Add dynamic shared memory bounds checking.

Author: maleadt

src/device/intrinsics/memory_shared.jl

@@ -40,8 +40,10 @@ shared memory; in the case of a homogeneous multi-part buffer it is preferred to
 """
 @inline function CuDynamicSharedArray(::Type{T}, dims, offset=0) where {T}
     len = prod(dims)
+    @boundscheck if offset+len > dynamic_smem_size()
+        throw(BoundsError())
+    end
     ptr = emit_shmem(T) + offset
-    # TODO: boundscheck against %dynamic_smem_size (currently unsupported by LLVM)
     CuDeviceArray(dims, ptr)
 end
 
@@ -52,6 +54,8 @@ macro cuDynamicSharedMem(T, dims, offset=0)
     end
 end
 
+dynamic_smem_size() = @asmcall("mov.u32 \$0, %dynamic_smem_size;", "=r", true, UInt32, Tuple{})
+
 # get a pointer to shared memory, with known (static) or zero length (dynamic shared memory)
 @generated function emit_shmem(::Type{T}, ::Val{len}=Val(0)) where {T,len}
     Context() do ctx

test/device/intrinsics/memory.jl

@@ -19,20 +19,20 @@ n = 256
     @on_device CuStaticSharedArray(Tuple{RGB{Float32}, UInt32}, (1,2))
 
     # dynamic
-    @on_device CuDynamicSharedArray(Float32, 1)
-    @on_device CuDynamicSharedArray(Float32, (1, 2))
-    @on_device CuDynamicSharedArray(Tuple{Float32, Float32}, 1)
-    @on_device CuDynamicSharedArray(Tuple{Float32, Float32}, (1,2))
-    @on_device CuDynamicSharedArray(Tuple{RGB{Float32}, UInt32}, 1)
-    @on_device CuDynamicSharedArray(Tuple{RGB{Float32}, UInt32}, (1,2))
+    @on_device shmem=sizeof(Float32) CuDynamicSharedArray(Float32, 1)
+    @on_device shmem=sizeof(Float32) CuDynamicSharedArray(Float32, (1, 2))
+    @on_device shmem=sizeof(Tuple{Float32, Float32}) CuDynamicSharedArray(Tuple{Float32, Float32}, 1)
+    @on_device shmem=sizeof(Tuple{Float32, Float32}) CuDynamicSharedArray(Tuple{Float32, Float32}, (1,2))
+    @on_device shmem=sizeof(Tuple{RGB{Float32}, UInt32}) CuDynamicSharedArray(Tuple{RGB{Float32}, UInt32}, 1)
+    @on_device shmem=sizeof(Tuple{RGB{Float32}, UInt32}) CuDynamicSharedArray(Tuple{RGB{Float32}, UInt32}, (1,2))
 
     # dynamic with offset
-    @on_device CuDynamicSharedArray(Float32, 1, 8)
-    @on_device CuDynamicSharedArray(Float32, (1,2), 8)
-    @on_device CuDynamicSharedArray(Tuple{Float32, Float32}, 1, 8)
-    @on_device CuDynamicSharedArray(Tuple{Float32, Float32}, (1,2), 8)
-    @on_device CuDynamicSharedArray(Tuple{RGB{Float32}, UInt32}, 1, 8)
-    @on_device CuDynamicSharedArray(Tuple{RGB{Float32}, UInt32}, (1,2), 8)
+    @on_device shmem=sizeof(Float32)+8 CuDynamicSharedArray(Float32, 1, 8)
+    @on_device shmem=sizeof(Float32)+8 CuDynamicSharedArray(Float32, (1,2), 8)
+    @on_device shmem=sizeof(Tuple{Float32, Float32})+8 CuDynamicSharedArray(Tuple{Float32, Float32}, 1, 8)
+    @on_device shmem=sizeof(Tuple{Float32, Float32})+8 CuDynamicSharedArray(Tuple{Float32, Float32}, (1,2), 8)
+    @on_device shmem=sizeof(Tuple{RGB{Float32}, UInt32})+8 CuDynamicSharedArray(Tuple{RGB{Float32}, UInt32}, 1, 8)
+    @on_device shmem=sizeof(Tuple{RGB{Float32}, UInt32})+8 CuDynamicSharedArray(Tuple{RGB{Float32}, UInt32}, (1,2), 8)
 end
 
 

test/setup.jl

@@ -196,16 +196,19 @@ macro test_throws_macro(ty, ex)
 end
 
 # Run some code on-device
-macro on_device(ex)
+macro on_device(ex...)
+    code = ex[end]
+    kwargs = ex[1:end-1]
+
     @gensym kernel
     esc(quote
         let
             function $kernel()
-                $ex
+                $code
                 return
             end
 
-            CUDA.@sync @cuda $kernel()
+            CUDA.@sync @cuda $(kwargs...) $kernel()
         end
     end)
 end