JuliaSIMD / JuliaSIMD/LoopVectorization.jl

Add low-level tiling API

Abierto
#96 0 comentarios 4 reacciones 0 asignados Ver en GitHub
enhancement
Lenguaje dominante
Julia
Estrellas
789
Forks
73
Métricas de merge de PR
Sin PR fusionados en 30 d

Descripción

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.

Guía de contribución

No hay ninguna guía de contribución indexada para este repositorio

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.