restructuring into an atomics.jl file

Author: leios

lib/CUDAKernels/src/CUDAKernels.jl

@@ -319,7 +319,8 @@ else
     const emit_shmem = CUDA._shmem
 end
 
-import KernelAbstractions: ConstAdaptor, SharedMemory, Scratchpad, __synchronize, __size, __atomic
+import KernelAbstractions: ConstAdaptor, SharedMemory, Scratchpad, __synchronize, __size
+import KernelAbstractions: __atomic, atomic_add!
 
 ###
 # GPU implementation of shared memory
@@ -386,7 +387,12 @@ end
 ###
 
 @inline function Cassette.overdub(::CUDACtx, ::typeof(__atomic), ex)
-    CUDA.@atomic(ex)
+    #CUDA.@atomic(ex)
+    KernelAbstractions.@print(ex)
+end
+
+@inline function Cassette.overdub(::CUDACtx, ::typeof(atomic_add!), ptr, b)
+    CUDA.atomic_add!(ptr, b)
 end
 
 end

src/KernelAbstractions.jl

@@ -1,7 +1,8 @@
 module KernelAbstractions
 
 export @kernel
-export @Const, @localmem, @private, @uniform, @synchronize, @index, @groupsize, @print, @atomic
+export @Const, @localmem, @private, @uniform, @synchronize, @index, @groupsize, @print
+export @atomic, atomic_add!
 export Device, GPU, CPU, Event, MultiEvent, NoneEvent
 export async_copy!
 
@@ -307,22 +308,6 @@ macro index(locale, args...)
     Expr(:call, GlobalRef(KernelAbstractions, index_function), esc(:__ctx__), map(esc, args)...)
 end
 
-"""
-    @atomic command
-
-This is a unified atomic interface
-
-# Platform differences
-  - `GPU`: This uses standard `@atomic` calls from CUDA.jl
-  - `CPU`: This reorganized the command to use atomic pointer logic
-"""
-
-macro atomic(ex)
-    quote
-        $__atomic($ex)
-    end
-end
-
 ###
 # Internal kernel functions
 ###
@@ -468,10 +453,6 @@ function __synchronize()
     error("@synchronize used outside kernel or not captured")
 end
 
-function __atomic(ex)
-    error("@atomic used outside kernel or not captured")
-end
-
 @generated function __print(items...)
     str = ""
     args = []
@@ -503,6 +484,10 @@ include("extras/extras.jl")
 
 include("reflection.jl")
 
+# Atomics
+
+include("atomics.jl")
+
 # CPU backend
 
 include("cpu.jl")

src/atomics.jl

@@ -0,0 +1,38 @@
+###
+# Atomics
+###
+
+"""
+    @atomic macro
+
+This is a unified atomic interface
+
+# Platform differences
+  - `GPU`: This uses standard `@atomic` calls from CUDA.jl
+  - `CPU`: This reorganized the command to use atomic pointer logic
+"""
+
+# NOTE: This only grabs the first symbol of the expression, not the entire expr
+macro atomic(ex)
+    quote
+        $__atomic($ex)
+    end
+end
+
+# Implement CPU macro here?
+function __atomic(ex)
+    error("@atomic macro incomplete!")
+end
+
+function atomic_add!(ptr, b)
+    Core.Intrinsics.atomic_pointermodify(ptr, +, b, :monotonic)
+end
+
+###
+# CPU implementation of atomic macro
+###
+
+#@inline function Cassette.overdub(::CPUCtx, ::typeof(__atomic), ex)
+#    println(ex)
+#end
+

src/cpu.jl

@@ -269,11 +269,3 @@ end
 
 # Argument conversion
 KernelAbstractions.argconvert(k::Kernel{CPU}, arg) = arg
-
-###
-# CPU implementation of atomic macro
-###
-
-@inline function Cassette.overdub(::CUDACtx, ::typeof(__atomic), ex)
-    println(ex)
-end