adding framework to build off of

Author: leios

lib/CUDAKernels/src/CUDAKernels.jl

@@ -319,7 +319,7 @@ else
     const emit_shmem = CUDA._shmem
 end
 
-import KernelAbstractions: ConstAdaptor, SharedMemory, Scratchpad, __synchronize, __size
+import KernelAbstractions: ConstAdaptor, SharedMemory, Scratchpad, __synchronize, __size, __atomic
 
 ###
 # GPU implementation of shared memory
@@ -381,4 +381,12 @@ end
     CUDA.ptx_isa_version(args...)
 end
 
+###
+# GPU implementation of atomics
+###
+
+@inline function Cassette.overdub(::CUDACtx, ::typeof(__atomic), ex)
+    CUDA.@atomic(ex)
+end
+
 end

src/KernelAbstractions.jl

@@ -1,7 +1,7 @@
 module KernelAbstractions
 
 export @kernel
-export @Const, @localmem, @private, @uniform, @synchronize, @index, @groupsize, @print
+export @Const, @localmem, @private, @uniform, @synchronize, @index, @groupsize, @print, @atomic
 export Device, GPU, CPU, Event, MultiEvent, NoneEvent
 export async_copy!
 
@@ -28,6 +28,7 @@ and then invoked on the arguments.
 - [`@uniform`](@ref)
 - [`@synchronize`](@ref)
 - [`@print`](@ref)
+- [`@atomic`](@ref)
 
 # Example:
 
@@ -306,6 +307,22 @@ 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
+        return $__atomic(ex)
+    end
+end
+
 ###
 # Internal kernel functions
 ###
@@ -451,6 +468,10 @@ 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 = []