adding overdubs for CUDA atomics

Author: leios

lib/CUDAKernels/src/CUDAKernels.jl

@@ -320,7 +320,7 @@ else
 end
 
 import KernelAbstractions: ConstAdaptor, SharedMemory, Scratchpad, __synchronize, __size
-import KernelAbstractions: __atomic, atomic_add!
+import KernelAbstractions: __atomic, atomic_add!, atomic_and!, atomic_cas!, atomic_dec!, atomic_inc!, atomic_max!, atomic_min!, atomic_op!, atomic_or!, atomic_sub!, atomic_xchg!, atomic_xor!
 
 ###
 # GPU implementation of shared memory
@@ -391,8 +391,12 @@ end
     KernelAbstractions.@print(ex)
 end
 
-@inline function Cassette.overdub(::CUDACtx, ::typeof(atomic_add!), ptr, b)
-    CUDA.atomic_add!(ptr, b)
+for afx in [atomic_add!, atomic_and!, atomic_cas!, atomic_dec!,
+            atomic_inc!, atomic_max!, atomic_min!, atomic_op!,
+            atomic_or!, atomic_sub!, atomic_xchg!, atomic_xor!]
+    @inline function Cassette.overdub(::CUDACtx, ::typeof(afx), args...)
+        CUDA.atomic_add!(args...)
+    end
 end
 
 end

src/atomics.jl

@@ -12,6 +12,11 @@ This is a unified atomic interface
   - `CPU`: This reorganized the command to use atomic pointer logic
 """
 
+# TODO: Remove?
+#const atomic_acquire = LLVM.API.LLVMAtomicOrderingAcquire
+#const atomic_release = LLVM.API.LLVMAtomicOrderingRelease
+#const atomic_acquire_release = LLVM.API.LLVMAtomicOrderingAcquireRelease
+
 # NOTE: This only grabs the first symbol of the expression, not the entire expr
 macro atomic(ex)
     quote
@@ -28,6 +33,58 @@ function atomic_add!(ptr, b)
     Core.Intrinsics.atomic_pointermodify(ptr, +, b, :monotonic)
 end
 
+function atomic_and!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, &, b, :monotonic)
+end
+
+# Not sure what this one is...
+function atomic_cas!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, &, b, :monotonic)
+end
+
+# TODO
+function atomic_dec!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, &, b, :monotonic)
+end
+
+# TODO
+function atomic_inc!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, &, b, :monotonic)
+end
+
+# TODO
+function atomic_max!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, &, b, :monotonic)
+end
+
+# TODO
+function atomic_min!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, &, b, :monotonic)
+end
+
+# TODO
+function atomic_op!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, &, b, :monotonic)
+end
+
+function atomic_or!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, |, b, :monotonic)
+end
+
+function atomic_sub!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, -, b, :monotonic)
+end
+
+# TODO
+function atomic_xchg!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, &, b, :monotonic)
+end
+
+# TODO
+function atomic_xor!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, &, b, :monotonic)
+end
+
 ###
 # CPU implementation of atomic macro
 ###