Remove KA 0.9 compat and clean up generic_matmatmul implementation

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"
 Adapt = "4.7.3"
 GPUArraysCore = "= 0.2.1"
 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"

src/host/gemm.jl

@@ -1,5 +1,5 @@
 using LinearAlgebra: MulAddMul, wrapper_char, _unwrap
-
+import KernelInterface as KI
 ## Matmul kernel
 
 # A[i, j] of a symmetric/hermitian matrix with only one triangle stored
@@ -32,43 +32,37 @@ end
     end
 end
 
+# additive identity to start accumulating from: -0.0, not 0.0, for IEEE floats, so that a
+# sum of -0.0 products stays -0.0. other element types needn't support unary minus.
+@inline accumulator_zero(::Type{T}) where {T <: Number} = -zero(T)
+@inline accumulator_zero(::Type{T}) where {T} = zero(T)
+
 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)
-
-    # li, lj, _ = KI.get_local_id()
-    li, lj = @index(Local, NTuple)
+    TAT = eltype(A); TBT = eltype(B); R = eltype(C)
 
-    # 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)
+    # Revisit potentially padding if value of `TILE_DIM` ever changed
+    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 = accumulator_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 +79,21 @@ 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])
+        # skip the zero padding past K: adding its +0.0 products would turn a -0.0 into +0.0
+        @inbounds for k in 1:min(TILE, K - k0)
+            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 +131,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

@@ -1182,7 +1182,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]