From 59f7219cb1b483afed0436e4f614e94bdfcc89a8 Mon Sep 17 00:00:00 2001 From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:41:06 -0300 Subject: [PATCH 1/4] Require KernelAbstractions 0.10 All back-ends, JLArrays included, have moved to KernelAbstractions 0.10. Update the back-end interface docs to define a KernelInterface back-end. --- Project.toml | 2 +- docs/src/interface.md | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index 9ecaec47..40f442b9 100644 --- a/Project.toml +++ b/Project.toml @@ -31,7 +31,7 @@ AcceleratedKernels = "0.5.8" Adapt = "4.7.3" GPUArraysCore = "= 0.2.2" JLD2 = "0.4, 0.5, 0.6" -KernelAbstractions = "0.9.43, 0.10" +KernelAbstractions = "0.10" LLVM = "3.9, 4, 5, 6, 7, 8, 9, 10" LinearAlgebra = "1" Printf = "1" diff --git a/docs/src/interface.md b/docs/src/interface.md index 37663494..b45ae976 100644 --- a/docs/src/interface.md +++ b/docs/src/interface.md @@ -21,12 +21,12 @@ end ``` This will allow your defined type (in this case `JLArray`) to use the GPUArrays interface where available. -To be able to actually use the functionality that is defined for `AbstractGPUArray`s, you need to define the backend, like so: +To be able to actually use the functionality that is defined for `AbstractGPUArray`s, you need to define a [KernelInterface](https://github.com/JuliaGPU/KernelAbstractions.jl/tree/main/lib/KernelInterface) back-end, which compiles and launches kernels, like so: ```julia -import KernelAbstractions: Backend -struct CustomBackend <: KernelAbstractions.GPU -KernelAbstractions.get_backend(a::CA) where CA <: CustomArray = CustomBackend() +import KernelInterface as KI +struct CustomBackend <: KI.Backend end +KI.get_backend(a::CA) where CA <: CustomArray = CustomBackend() ``` There are numerous examples of potential interfaces for GPUArrays, such as with [JLArrays](https://github.com/JuliaGPU/GPUArrays.jl/blob/main/lib/JLArrays/src/JLArrays.jl), [CuArrays](https://github.com/JuliaGPU/CUDA.jl/blob/main/CUDACore/src/array.jl), and [ROCArrays](https://github.com/JuliaGPU/AMDGPU.jl/blob/main/src/array.jl). From 1376941e13538b974f838e9f54e21c0f84c8c7ce Mon Sep 17 00:00:00 2001 From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:42:11 -0300 Subject: [PATCH 2/4] Clean up the matmul kernel for KA 0.10 KA 0.9's CPU back-end split kernels at every `@synchronize`, so values live across a barrier had to be `@uniform`, `@private` or recomputed. KA 0.10 runs kernels on POCL, which needs none of that: compute the global indices once and accumulate into a plain variable. --- src/host/gemm.jl | 42 ++++++++++-------------------------------- src/host/linalg.jl | 2 +- 2 files changed, 11 insertions(+), 33 deletions(-) diff --git a/src/host/gemm.jl b/src/host/gemm.jl index d018806d..55dcfe7d 100644 --- a/src/host/gemm.jl +++ b/src/host/gemm.jl @@ -38,37 +38,26 @@ const TILE_DIM = 16 M, N, K, ::Val{TA}, ::Val{TB}, ::Val{TILE}, ) where {TA, TB, TILE} - @uniform TAT = eltype(A); @uniform TBT = eltype(B); @uniform R = eltype(C) + TAT = eltype(A); TBT = eltype(B); R = eltype(C) - # li, lj, _ = KI.get_local_id() li, lj = @index(Local, NTuple) - - # grow, gcol, _ = KI.get_group_id() grow, gcol = @index(Group, NTuple) - # Uncomment once moved to KernelInterface - # gi = (grow - 1) * TILE + li - # gj = (gcol - 1) * TILE + lj + gi = (grow - 1) * TILE + li + gj = (gcol - 1) * TILE + lj # padded by one row to avoid bank conflicts on shared memory tile1 = @localmem TAT (TILE + 1, TILE) tile2 = @localmem TBT (TILE + 1, TILE) # private variable for tile C - @uniform Tacc = promote_type(R, typeof(zero(TAT) * zero(TBT) + zero(TAT) * zero(TBT))) - # acc = -zero(Tacc) #Use this when KernelInterface - acc = @private Tacc 1 - @inbounds acc[1] = zero(Tacc) + Tacc = promote_type(R, typeof(zero(TAT) * zero(TBT) + zero(TAT) * zero(TBT))) + acc = zero(Tacc) # number of tiles depends on inner dimension - # NUM_TILES = div(K + TILE - 1, TILE) - @uniform NUM_TILES = cld(K, TILE) + NUM_TILES = cld(K, TILE) # loop over all tiles needed for this calculation for t in 0:(NUM_TILES - 1) - # Remove once moved to KernelInterface - gi = (grow - 1) * TILE + li - gj = (gcol - 1) * TILE + lj - # load inputs into tiles, with bounds checking for non-square matrices k0 = t * TILE ac = k0 + lj @@ -85,28 +74,19 @@ const TILE_DIM = 16 end # wait for all tiles to be loaded - @synchronize() - - # Remove once moved to KernelInterface - gi = (grow - 1) * TILE + li - gj = (gcol - 1) * TILE + lj + @synchronize() @inbounds for k in 1:TILE - acc[1] = muladd(tile1[li, k], tile2[k, lj], acc[1]) + acc = muladd(tile1[li, k], tile2[k, lj], acc) end - @synchronize() + @synchronize() end - # Remove once moved to KernelInterface - gi = (grow - 1) * TILE + li - gj = (gcol - 1) * TILE + lj - # save if inbounds if gi <= M && gj <= N - @inbounds C[gi, gj] = add(acc[1], C[gi, gj]) + @inbounds C[gi, gj] = add(acc, C[gi, gj]) end - # return end @inline function check_matmul_shapes(C, A, B, cA='N', cB='N') @@ -144,8 +124,6 @@ function generic_matmatmul!(C::AbstractArray{R}, cA::AbstractChar, cB::AbstractC workgroupsize=(TILE_DIM, TILE_DIM) numworkgroups=(cld(M, TILE_DIM), cld(N, TILE_DIM)) - # KI.@kernel KI.get_backend(C) workgroupsize=workgroupsize numworkgroups=numworkgroups coalesced_matmul_kernel!( - # C, A, B, add, Int(M), Int(N), Int(K), Val(cA), Val(cB), Val(TILE_DIM)) coalesced_matmul_kernel!(get_backend(C), workgroupsize)( C, A, B, add, Int(M), Int(N), Int(K), Val(cA), Val(cB), Val(TILE_DIM); ndrange=numworkgroups .* workgroupsize) C diff --git a/src/host/linalg.jl b/src/host/linalg.jl index 6e162b7f..631c0ba7 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -1178,7 +1178,7 @@ end ## reflect function LinearAlgebra.reflect!(x::AbstractGPUArray, y::AbstractGPUArray, c::Number, s::Number) - @kernel function reflect_kernel!(x, y, c, s) + @kernel function reflect_kernel!(x, y, c, s) i = @index(Global, Linear) @inbounds xi = x[i] @inbounds yi = y[i] From 4fe32bfb2916eabb7e7b82ae641ee24a031a9df3 Mon Sep 17 00:00:00 2001 From: Christian <28689358+christiangnrd@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:14:19 -0300 Subject: [PATCH 3/4] Write the matmul kernel against KernelInterface The kernel always launches a 2D grid of fixed-size work-groups and indexes by work-group and local id, which is what KernelInterface exposes directly; KernelAbstractions' `@kernel` transformations add nothing here. Performance is unchanged. --- Project.toml | 2 ++ src/GPUArrays.jl | 1 + src/host/gemm.jl | 19 ++++++++++--------- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/Project.toml b/Project.toml index 40f442b9..8748f219 100644 --- a/Project.toml +++ b/Project.toml @@ -10,6 +10,7 @@ AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf" LLVM = "929cbde3-209d-540e-8aea-75f648917ca0" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" @@ -32,6 +33,7 @@ Adapt = "4.7.3" GPUArraysCore = "= 0.2.2" JLD2 = "0.4, 0.5, 0.6" KernelAbstractions = "0.10" +KernelInterface = "0.4.1" LLVM = "3.9, 4, 5, 6, 7, 8, 9, 10" LinearAlgebra = "1" Printf = "1" diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index d12d6762..3db7a050 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -16,6 +16,7 @@ using Reexport @reexport using GPUArraysCore using KernelAbstractions +import KernelInterface as KI import AcceleratedKernels as AK diff --git a/src/host/gemm.jl b/src/host/gemm.jl index 55dcfe7d..75ebac24 100644 --- a/src/host/gemm.jl +++ b/src/host/gemm.jl @@ -33,21 +33,21 @@ end end const TILE_DIM = 16 -@kernel unsafe_indices = true function coalesced_matmul_kernel!( +function coalesced_matmul_kernel!( C, A, B, add, M, N, K, ::Val{TA}, ::Val{TB}, ::Val{TILE}, ) where {TA, TB, TILE} TAT = eltype(A); TBT = eltype(B); R = eltype(C) - li, lj = @index(Local, NTuple) - grow, gcol = @index(Group, NTuple) + li, lj, _ = KI.get_local_id() + grow, gcol, _ = KI.get_group_id() gi = (grow - 1) * TILE + li gj = (gcol - 1) * TILE + lj # padded by one row to avoid bank conflicts on shared memory - tile1 = @localmem TAT (TILE + 1, TILE) - tile2 = @localmem TBT (TILE + 1, TILE) + tile1 = KI.localmemory(TAT, (TILE + 1, TILE)) + tile2 = KI.localmemory(TBT, (TILE + 1, TILE)) # private variable for tile C Tacc = promote_type(R, typeof(zero(TAT) * zero(TBT) + zero(TAT) * zero(TBT))) @@ -74,19 +74,20 @@ const TILE_DIM = 16 end # wait for all tiles to be loaded - @synchronize() + KI.barrier() @inbounds for k in 1:TILE acc = muladd(tile1[li, k], tile2[k, lj], acc) end - @synchronize() + KI.barrier() end # save if inbounds if gi <= M && gj <= N @inbounds C[gi, gj] = add(acc, C[gi, gj]) end + return end @inline function check_matmul_shapes(C, A, B, cA='N', cB='N') @@ -124,7 +125,7 @@ function generic_matmatmul!(C::AbstractArray{R}, cA::AbstractChar, cB::AbstractC workgroupsize=(TILE_DIM, TILE_DIM) numworkgroups=(cld(M, TILE_DIM), cld(N, TILE_DIM)) - coalesced_matmul_kernel!(get_backend(C), workgroupsize)( - C, A, B, add, Int(M), Int(N), Int(K), Val(cA), Val(cB), Val(TILE_DIM); ndrange=numworkgroups .* workgroupsize) + KI.@launch KI.get_backend(C) workgroupsize=workgroupsize numgroups=numworkgroups coalesced_matmul_kernel!( + C, A, B, add, Int(M), Int(N), Int(K), Val(cA), Val(cB), Val(TILE_DIM)) C end From 11da6e8a5d323a0a6d4e3bf9bf580afe374faa04 Mon Sep 17 00:00:00 2001 From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:17:48 -0300 Subject: [PATCH 4/4] Don't pad the matmul kernel's local memory tiles The padding avoided bank conflicts for row-major access, but work-items with consecutive local ids access consecutive elements of a column-major tile (or all read one element of `tile2`), which is conflict-free without it. Padding instead misaligned the columns of `tile2`, keeping each work-item from reading its column with vector loads. On an RTX 5080, 2048x2048 matmul takes 3.5 instead of 5.1 ms for Float32, 4.4 instead of 5.2 ms for ComplexF32, and 3.2 instead of 5.1 ms for Float16. Float64, bound by this GPU's FP64 throughput, stays within 5%. --- src/host/gemm.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/host/gemm.jl b/src/host/gemm.jl index 75ebac24..ee45f29b 100644 --- a/src/host/gemm.jl +++ b/src/host/gemm.jl @@ -45,9 +45,11 @@ function coalesced_matmul_kernel!( gi = (grow - 1) * TILE + li gj = (gcol - 1) * TILE + lj - # padded by one row to avoid bank conflicts on shared memory - tile1 = KI.localmemory(TAT, (TILE + 1, TILE)) - tile2 = KI.localmemory(TBT, (TILE + 1, TILE)) + # no padding: consecutive work-items (`li`) access consecutive elements of a column, or + # all the same element of `tile2`, so there are no bank conflicts to avoid. unpadded + # columns stay aligned, which lets each work-item read `tile2[:, lj]` with vector loads. + tile1 = KI.localmemory(TAT, (TILE, TILE)) + tile2 = KI.localmemory(TBT, (TILE, TILE)) # private variable for tile C Tacc = promote_type(R, typeof(zero(TAT) * zero(TBT) + zero(TAT) * zero(TBT)))