JuliaSIMD / JuliaSIMD/LoopVectorization.jl

Add low-level tiling API

Offen
#96 0 Kommentare 4 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
enhancement
Vorherrschende Sprache
Julia
Sterne
789
Forks
73
PR-Merge-Kennzahlen
Keine gemergten PRs in 30 T.

Beschreibung

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.

Beitragsleitfaden

Für dieses Repository ist kein Beitragsleitfaden indexiert

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.