Harden intrinsic compile-time validation

CUDACore/src/device/intrinsics/atomics.jl

@@ -93,7 +93,6 @@ for T in (:Float16, :Float32, :Float64)
         rmw = Symbol("f$op")
 
         fn = Symbol("atomic_$(op)!")
-        # XXX: cannot select
         @eval @inline $fn(ptr::Union{LLVMPtr{$T,AS.Generic},
                                      LLVMPtr{$T,AS.Global},
                                      LLVMPtr{$T,AS.Shared}}, val::$T) =
@@ -120,6 +119,7 @@ end
         atomic_add!(ptr, -val)
 end
 
+# cmpxchg is subject to the same address space restrictions as atomicrmw, above
 @generated function llvm_atomic_cas(ptr::LLVMPtr{T,A}, cmp::T, val::T) where {T, A}
     @dispose ctx=Context() begin
         T_val = convert(LLVMType, T)
@@ -149,11 +149,16 @@ end
 end
 
 for T in (:Int32, :Int64, :UInt32, :UInt64)
-    @eval @inline atomic_cas!(ptr::LLVMPtr{$T}, cmp::$T, val::$T) =
+    @eval @device_function @inline function atomic_cas!(ptr::LLVMPtr{$T,A}, cmp::$T,
+                                                        val::$T) where {A}
+        GPUCompiler.@static_assert(
+            A == AS.Generic || A == AS.Global || A == AS.Shared,
+            "atomics require a generic, global, or shared address space")
         llvm_atomic_cas(ptr, cmp, val)
+    end
 end
 
-# NVPTX doesn't support cmpxchg with i16 yet
+# LLVM expands i16 cmpxchg to a 32-bit CAS loop; use native PTX where available.
 for A in (AS.Generic, AS.Global, AS.Shared), T in (:Int16, :UInt16)
     if A == AS.Global
         scope = ".global"
@@ -165,9 +170,12 @@ for A in (AS.Generic, AS.Global, AS.Shared), T in (:Int16, :UInt16)
 
     intr = "atom$scope.cas.b16 \$0, [\$1], \$2, \$3;"
     @eval @device_function @inline function atomic_cas!(ptr::LLVMPtr{$T,$A}, cmp::$T, val::$T)
-        require_sm_70()
-        @asmcall($intr, "=h,l,h,h", true, $T,
-                 Tuple{Core.LLVMPtr{$T,$A},$T,$T}, ptr, cmp, val)
+        if compute_capability() >= sv"7.0"
+            @asmcall($intr, "=h,l,h,h", true, $T,
+                     Tuple{Core.LLVMPtr{$T,$A},$T,$T}, ptr, cmp, val)
+        else
+            llvm_atomic_cas(ptr, cmp, val)
+        end
     end
 end
 
@@ -237,9 +245,9 @@ Reads the value `old` located at address `ptr` and compare with `cmp`. If `old`
 `cmp`, stores `val` at the same address. Otherwise, doesn't change the value `old`. These
 operations are performed in one atomic transaction. The function returns `old`.
 
-This operation is supported for values of type Int32, Int64, UInt32 and UInt64.
-Additionally, on GPU hardware with compute capability 7.0+, values of type UInt16 are
-supported.
+This operation is supported for values of type Int16, Int32, Int64, UInt16, UInt32,
+UInt64, Float16, Float32, Float64, and BFloat16. 16-bit operations use a native
+instruction on GPU hardware with compute capability 7.0+, and are emulated otherwise.
 """
 atomic_cas!
 
@@ -260,9 +268,9 @@ Reads the value `old` located at address `ptr`, computes `old + val`, and stores
 back to memory at the same address. These operations are performed in one atomic
 transaction. The function returns `old`.
 
-This operation is supported for values of type Int32, Int64, UInt32, UInt64, and Float32.
-Additionally, on GPU hardware with compute capability 6.0+, values of type Float64 are
-supported.
+This operation is supported for values of type Int32, Int64, UInt32, UInt64, Float16,
+Float32, and Float64. BFloat16 is supported on Julia 1.11+. The back-end uses a native
+instruction where available and emulates the operation otherwise.
 """
 atomic_add!
 

CUDACore/src/device/intrinsics/cooperative_groups.jl

@@ -24,7 +24,7 @@ Noteworthy missing functionality:
 module CG
 
 using ..CUDACore
-using ..CUDACore: i32, Aligned, alignment, @device_function
+using ..CUDACore: i32, Aligned, alignment, GPUCompiler, @device_function
 
 import ..LLVM
 using ..LLVM.Interop
@@ -523,9 +523,10 @@ end
     wait_prior(group, stage)
 
 Make all threads in this group wait for all but `stage` previously submitted
-[`memcpy_async`](@ref) operations to complete.
+[`memcpy_async`](@ref) operations to complete. At most 8 stages can be kept in flight;
+larger values of `stage` are capped accordingly.
 """
-function wait_prior(group::memcpy_group, stage::Integer)
+@inline function wait_prior(group::memcpy_group, stage::Integer)
     if compute_capability() >= sv"8.0"
         pipeline_wait_prior(stage)
     end
@@ -564,9 +565,20 @@ end
 @device_function pipeline_commit() =
     ccall("llvm.nvvm.cp.async.commit.group", llvmcall, Cvoid, ())
 
-@device_function pipeline_wait_prior(n) =
+# the underlying `cp.async.wait_group` instruction takes an immediate operand, so the
+# number of stages has to be a compile-time constant
+@device_function @inline pipeline_wait_prior(::Val{n}) where {n} =
     ccall("llvm.nvvm.cp.async.wait.group", llvmcall, Cvoid, (Int32,), n)
 
+# ... which we cannot guarantee for a run-time value, so dispatch to a constant instead.
+# CUDA's `__pipeline_wait_prior` does the same, capping the number of stages at 8.
+@device_function @inline function pipeline_wait_prior(n::Integer)
+    Base.Cartesian.@nexprs 8 i -> begin
+        n <= i-1 && return pipeline_wait_prior(Val(i-1))
+    end
+    pipeline_wait_prior(Val(8))
+end
+
 @device_function @generated function pipeline_memcpy_async(dst::LLVMPtr{T}, src::LLVMPtr{T}) where T
     size_and_align = sizeof(T)
     size_and_align in (4, 8, 16) || :(return error($"Unsupported size $size_and_align"))
@@ -584,7 +596,9 @@ end
 @inline function _memcpy_async(group, dst::LLVMPtr, src::LLVMPtr,
                                bytes, ::Val{align_hint}) where {align_hint}
     align = min(16, align_hint)
-    ispow2(align) || throw(ArgumentError("Alignment must be a power of 2"))
+    GPUCompiler.@static_assert(
+        ispow2(align),
+        "memcpy_async alignment must be a power of 2")
     if compute_capability() >= sv"8.0"
         _memcpy_async_dispatch(group, Val{align}(), dst, src, bytes[])
         pipeline_commit()

test/core/device/intrinsics.jl

@@ -64,12 +64,26 @@ end
     nanosleep_kernel() = nanosleep(UInt32(1))
     @test_throws "requires compute capability 7.0" @cuda launch=false arch=sm"61" nanosleep_kernel()
 
-    out16 = CUDA.zeros(UInt16, 1)
-    function atomic_cas_kernel(out)
-        CUDA.atomic_cas!(pointer(out), UInt16(0), UInt16(1))
+    @test @filecheck CUDA.code_ptx((Core.LLVMPtr{UInt16,AS.Global},); arch=sm"61") do ptr
+        @check "{{atom(\\.sys)?\\.global\\.cas\\.b32}}"
+        @check_not "cas.b16"
+        CUDA.atomic_cas!(ptr, UInt16(0), UInt16(1))
+        return
+    end
+    @test @filecheck CUDA.code_ptx((Core.LLVMPtr{UInt16,AS.Global},); arch=sm"70") do ptr
+        @check "atom.global.cas.b16"
+        CUDA.atomic_cas!(ptr, UInt16(0), UInt16(1))
         return
     end
-    @test_throws "requires compute capability 7.0" @cuda launch=false arch=sm"61" atomic_cas_kernel(out16)
+
+    function invalid_atomic_address_space_kernel(ptr)
+        CUDA.atomic_cas!(ptr, UInt32(0), UInt32(1))
+        return
+    end
+    local_ptr_t = Core.LLVMPtr{UInt32,AS.Local}
+    @test_throws "atomics require a generic, global, or shared address space" begin
+        CUDA.cufunction(invalid_atomic_address_space_kernel, Tuple{local_ptr_t}; arch=sm"80")
+    end
 
     outf16 = CUDA.zeros(Float16, 16 * 16)
     function wmma_kernel(out)
@@ -108,6 +122,35 @@ end
         return
     end
     @test_throws "requires compute capability 9.0" @cuda launch=false arch=sm"80" distributed_shared_kernel()
+
+    # `cp.async.wait_group` takes an immediate operand, so the number of stages needs to
+    # be materialized as a constant, capping it at 8 like CUDA's `__pipeline_wait_prior`
+    @test @filecheck CUDA.code_ptx(Tuple{}; arch=sm"80", kernel=true) do
+        @check "cp.async.wait_group {{0;}}"
+        CG.wait_prior(CG.this_thread_block(), 0)
+        return
+    end
+    @test @filecheck CUDA.code_ptx(Tuple{}; arch=sm"80", kernel=true) do
+        @check "cp.async.wait_group {{8;}}"
+        CG.wait_prior(CG.this_thread_block(), 20)
+        return
+    end
+    # a run-time number of stages is dispatched to those constants
+    @test @filecheck CUDA.code_ptx((Int32,); arch=sm"80", kernel=true) do stage
+        @check "cp.async.wait_group"
+        CG.wait_prior(CG.this_thread_block(), stage)
+        return
+    end
+
+    function invalid_memcpy_alignment_kernel(dst, src)
+        CG.memcpy_async(CG.this_thread_block(), dst, src, 4)
+        return
+    end
+    dst_t = CUDACore.Aligned{Core.LLVMPtr{UInt32,AS.Shared},3}
+    src_t = CUDACore.Aligned{Core.LLVMPtr{UInt32,AS.Global},3}
+    @test_throws "memcpy_async alignment must be a power of 2" begin
+        CUDA.cufunction(invalid_memcpy_alignment_kernel, Tuple{dst_t,src_t}; arch=sm"80")
+    end
 end