Require KernelAbstractions 0.10 and clean up the matmul kernel

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"
@@ -31,7 +32,8 @@ 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"
+KernelInterface = "0.4.1"
 LLVM = "3.9, 4, 5, 6, 7, 8, 9, 10"
 LinearAlgebra = "1"
 Printf = "1"

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).

src/GPUArrays.jl

@@ -16,6 +16,7 @@ using Reexport
 @reexport using GPUArraysCore
 
 using KernelAbstractions
+import KernelInterface as KI
 
 import AcceleratedKernels as AK
 

src/host/gemm.jl

@@ -33,42 +33,33 @@ 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}
-    @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
+    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)
+    # 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
-    @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 +76,20 @@ 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
+        KI.barrier()
 
         @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()
+        KI.barrier()
     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
+    return
 end
 
 @inline function check_matmul_shapes(C, A, B, cA='N', cB='N')
@@ -144,9 +127,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))
 
-    # 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)
+    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

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]