JuliaSIMD / JuliaSIMD/LoopVectorization.jl

Add low-level tiling API

Open
#96 0 comments 4 reactions 0 assignees View on GitHub
enhancement
Dominant language
Julia
Stars
789
Forks
73
PR merge metrics
No merged PRs in 30d

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.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.