JuliaSIMD / JuliaSIMD/LoopVectorization.jl
Add low-level tiling API
- Langage dominant
- Julia
- Étoiles
- 789
- Forks
- 73
- Métriques de merge des PR
- Aucune PR mergée en 30 j
Description
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.
Guide de contribution
Aucun guide de contribution indexé pour ce dépôt
Piste de recherche
Commencez par l’exemple motivant dans benchmark/bench_gemm.jl, en particulier la transformation de fusetest1! en fusetest2!. Lisez les chemins de lowering existants de LoopVectorization afin de déterminer où une API de tiling de bas niveau pourrait exposer la génération finale d’expressions Julia. La réalisation devrait inclure une API capable d’exprimer cette transformation et de préserver l’équivalence numérique montrée par l’exemple.
Rédigé par le modèle d'indexation à partir du texte de l'issue.
Évaluation
- Stack technique
- julia
- Domaine
- performance
- Type d'issue
- Fonctionnalité
- Difficulté
- 5/5
- Temps estimé
- Plus d'une semaine
- Activité
- À l'abandon
- Clarté
- À clarifier
- Accessibilité débutants
- 25/100