JuliaSIMD / JuliaSIMD/LoopVectorization.jl

Add low-level tiling API

Aperta
#96 0 commenti 4 reazioni 0 assegnatari Vedi su GitHub
enhancement
Lingua principale
Julia
Stelle
789
Fork
73
Metriche di merge delle PR
Nessuna PR unita negli ultimi 30g

Descrizione

Having LoopVectorization automatically tile for cache efficiency could be broken down into two major components:
1. A cost model determining profitability.
2. Lowering, generating the final Julia expression.

"2." would be useful on its own as a low-level API, and allow quick experimentation to inform development of "1." (although "1." should be largely theory-driven).

tkf wanted this as a feature, and [provided this](https://github.com/tkf/Transducers.jl/blob/2217a84a7a382d9611a7750ae6628681ebb5c667/benchmark/bench_gemm.jl#L171) as a motivating example:
```julia
using LoopVectorization, BenchmarkTools, Transducers, ArgCheck, Test
using Referenceables: referenceable

using Transducers: @simd_if, @next, complete, maybe_usesimd, BottomRF, SideEffect

_size(x) = size.(axes(x), 1)
_size(x, i) = size(axes(x)[i], 1)

@inline function ij_ik_kj_foldable(C, A, B)
@argcheck _size(C) === (_size(A, 1), _size(B, 2))
@argcheck _size(A, 2) === _size(B, 1)
@argcheck !any(Base.has_offset_axes, (A, B, C))
return AdHocFoldable() do rf, acc, _
Base.@_inline_meta
for j in 1:size(C, 2), k in 1:size(A, 2)
b = @inbounds B[k, j]
@simd_if rf for i in 1:size(A, 1)
c = @inbounds C[i, j]
a = @inbounds A[i, k]
acc = @next(rf, acc, (c, a, b))
end
end
return complete(rf, acc)
end
end

function fusedxfmul!(C, A, B1, B2)
fill!(C, 0)

B = Broadcast.instantiate(Broadcast.broadcasted(tuple, B1, B2))
CAB = ij_ik_kj_foldable(referenceable(C), A, B)

foreach(CAB; simd=Val(:ivdep)) do (c, a, (b1, b2))
c0 = c[]
c1 = muladd(a, b1, c0)
c2 = a^2 * b2
c[] = c1 + c2
return # for type stability
end

return C
end

function fusetest1!(C, A, B₁, B₂)
N, M = size(C)
@avx for nᵣ ∈ 1:N, m ∈ 1:M
Cnm = zero(eltype(C))
for nc ∈ 1:N
Cnm += A[nᵣ, nc] * B₁[nc, m] + A[nᵣ, nc]^2 * B₂[nc, m]
end
C[nᵣ, m] = Cnm
end
end
function fusetest2!(C, A, B₁, B₂)
N, M = size(C)
fill!(C, 0)
Nblock = 20
Nblocks, rem = divrem(N, Nblock)
for n in 0:Nblocks-1
Ncoffset = n * Nblock
@avx for nᵣ ∈ 1:N, m ∈ 1:M, nc ∈ 1:20
temp = A[nᵣ, nc + Ncoffset] * B₁[nc + Ncoffset, m] + A[nᵣ, nc + Ncoffset]^2 * B₂[nc + Ncoffset, m]
C[nᵣ, m] += temp
end
end
if rem > 0
Ncoffset = Nblocks * Nblock
@avx for nᵣ ∈ 1:N, m ∈ 1:M, nc ∈ 1 + Ncoffset:rem + Ncoffset
temp = A[nᵣ, nc] * B₁[nc, m] + A[nᵣ, nc]^2 * B₂[nc, m]
C[nᵣ, m] += temp
end
end
end

n = 1024; m = 8;
A1 = randn(n, n);
A2 = A1 .^ 2;
B1 = randn(n, m);
B2 = randn(n, m);
C1 = zero(A1 * B2);
C2 = similar(C1); C3 = similar(C1);

fusetest1!(C1, A1, B1, B2)
fusetest2!(C2, A1, B1, B2)
fusedxfmul!(C3, A1, B1, B2);
@test A1 * B1 + A2 * B2 ≈ C1 ≈ C2 ≈ C3

@benchmark fusetest1!($C1, $A1, $B1, $B2)
@benchmark fusetest2!($C2, $A1, $B1, $B2)
@benchmark fusedxfmul!($C3, $A1, $B1, $B2)
```

This yields:
```julia
julia> fusetest1!(C1, A1, B1, B2)

julia> fusetest2!(C2, A1, B1, B2)

julia> fusedxfmul!(C3, A1, B1, B2);

julia> @test A1 * B1 + A2 * B2 ≈ C1 ≈ C2 ≈ C3
Test Passed

julia> @benchmark fusetest1!($C1, $A1, $B1, $B2)
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 8.936 ms (0.00% GC)
median time: 9.449 ms (0.00% GC)
mean time: 9.410 ms (0.00% GC)
maximum time: 12.198 ms (0.00% GC)
--------------
samples: 531
evals/sample: 1

julia> @benchmark fusetest2!($C2, $A1, $B1, $B2)
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 2.816 ms (0.00% GC)
median time: 2.927 ms (0.00% GC)
mean time: 2.969 ms (0.00% GC)
maximum time: 5.903 ms (0.00% GC)
--------------
samples: 1676
evals/sample: 1

julia> @benchmark fusedxfmul!($C3, $A1, $B1, $B2)
BenchmarkTools.Trial:
memory estimate: 176 bytes
allocs estimate: 7
--------------
minimum time: 6.635 ms (0.00% GC)
median time: 7.052 ms (0.00% GC)
mean time: 7.100 ms (0.00% GC)
maximum time: 9.614 ms (0.00% GC)
--------------
samples: 703
evals/sample: 1
```
The transformation to go from `fusetest1!` to `fusetest2!` is very straightforward, and yields substantial performance benefits.

Guida per i contributori

Nessuna guida per i contributori indicizzata per questo repository

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.