Merge pull request #14 from JuliaGPU/vc/shmem

Fix and test local memory

Author: vchuravy

src/KernelAbstractions.jl

@@ -79,16 +79,15 @@ Query the workgroupsize on the device.
 """
 function groupsize end
 
-const shmem_id = Ref(0)
-
 """
    @localmem T dims
 """
 macro localmem(T, dims)
-    id = (shmem_id[]+= 1)
+    # Stay in sync with CUDAnative
+    id = gensym("static_shmem")
 
     quote
-        $SharedMemory($(esc(T)), Val($(esc(dims))), Val($id))
+        $SharedMemory($(esc(T)), Val($(esc(dims))), Val($(QuoteNode(id))))
     end
 end
 
@@ -281,11 +280,11 @@ include("macros.jl")
 ###
 
 function Scratchpad(::Type{T}, ::Val{Dims}) where {T, Dims}
-    throw(MethodError(ScratchArray, (T, Val(Dims))))
+    throw(MethodError(Scratchpad, (T, Val(Dims))))
 end
 
 function SharedMemory(::Type{T}, ::Val{Dims}, ::Val{Id}) where {T, Dims, Id}
-    throw(MethodError(ScratchArray, (T, Val(Dims), Val(Id))))
+    throw(MethodError(SharedMemory, (T, Val(Dims), Val(Id))))
 end
 
 function __synchronize()

src/backends/cuda.jl

@@ -203,7 +203,7 @@ end
 ###
 @inline function Cassette.overdub(ctx::CUDACtx, ::typeof(SharedMemory), ::Type{T}, ::Val{Dims}, ::Val{Id}) where {T, Dims, Id}
     ptr = CUDAnative._shmem(Val(Id), T, Val(prod(Dims)))
-    CUDAnative.CuDeviceArray(Dims, CUDAnative.DevicePtr{T, CUDAnative.AS.Shared}(ptr))
+    CUDAnative.CuDeviceArray(Dims, ptr)
 end
 
 ###

test/localmem.jl

@@ -0,0 +1,32 @@
+using KernelAbstractions
+using Test
+using CUDAapi
+if has_cuda_gpu()
+    using CuArrays
+    CuArrays.allowscalar(false)
+end
+
+@kernel function localmem(A)
+    I = @index(Global, Linear)
+    i = @index(Local, Linear)
+    lmem = @localmem Int (groupsize(),) # Ok iff groupsize is static 
+    lmem[i] = i
+    @synchronize
+    A[I] = lmem[groupsize() - i + 1]
+end
+
+function harness(backend, ArrayT)
+    A = ArrayT{Int}(undef, 64)
+    wait(localmem(backend, 16)(A, ndrange=size(A)))
+    @test all(A[1:16] .== 16:-1:1)
+    @test all(A[17:32] .== 16:-1:1)
+    @test all(A[33:48] .== 16:-1:1)
+    @test all(A[49:64] .== 16:-1:1)
+end
+
+@testset "kernels" begin
+    harness(CPU(), Array)
+    if has_cuda_gpu()
+        harness(CUDA(), CuArray)
+    end
+end

test/runtests.jl

@@ -5,4 +5,8 @@ using Test
     include("test.jl")
 end
 
+@testset "Localmem" begin
+    include("localmem.jl")
+end
+
 include("examples.jl")