small test fixes

Author: leios

lib/CUDAKernels/src/CUDAKernels.jl

@@ -391,11 +391,24 @@ end
     KernelAbstractions.@print(ex)
 end
 
-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!]
+afxs = Dict(
+    atomic_add! => CUDA.atomic_add!,
+    atomic_and! => CUDA.atomic_and!,
+    atomic_cas! => CUDA.atomic_cas!,
+    atomic_dec! => CUDA.atomic_dec!,
+    atomic_inc! => CUDA.atomic_inc!,
+    atomic_max! => CUDA.atomic_max!,
+    atomic_min! => CUDA.atomic_min!,
+    atomic_op! => CUDA.atomic_op!,
+    atomic_or! => CUDA.atomic_or!,
+    atomic_sub! => CUDA.atomic_sub!,
+    atomic_xchg! => CUDA.atomic_xchg!,
+    atomic_xor! => CUDA.atomic_xor!
+)
+
+for (afx, cfx) in afxs
     @inline function Cassette.overdub(::CUDACtx, ::typeof(afx), args...)
-        CUDA.atomic_add!(args...)
+        cfx(args...)
     end
 end
 

src/atomics.jl

@@ -6,6 +6,16 @@ export @atomic, atomic_add!, atomic_sub!, atomic_and!, atomic_or!, atomic_xor!,
        atomic_min!, atomic_max!, atomic_inc!, atomic_dec!, atomic_xchg!,
        atomic_op!, atomic_cas!
 
+# helper functions for inc and dec
+
+function dec(a::T,b::T) where T
+    ((a == 0) | (a > b)) ? b : (old-T(1))
+end
+
+function inc(a::T,b::T) where T
+    (a >= b) ? T(0) : (a+T(1))
+end
+
 """
     @atomic macro
 
@@ -58,23 +68,13 @@ function atomic_cas!(ptr::Ptr{T}, old::T, new::T) where T
 end
 
 # Implementation of: (((old == 0) | (old > b)) ? b : (old-1)), returns old
-# Currently broken
 function atomic_dec!(ptr::Ptr{T}, b::T) where T
-    Core.Intrinsics.atomic_fence(:monotonic)
-    if (unsafe_load(ptr) == 0 | unsafe_load(ptr) > b)
-        Core.Intrinsics.atomic_pointerswap(ptr, b, :monotonic)
-    else
-        Core.Intrinsics.atomic_pointermodify(ptr, -, T(1), :monotonic)
-    end
+    Core.Intrinsics.atomic_pointermodify(ptr::Ptr{T}, dec, b::T, :monotonic)
 end
 
 # implementation of: ((old >= b) ? 0 : (old+1)), returns old
 function atomic_inc!(ptr::Ptr{T}, b::T) where T
-    if unsafe_load(ptr) >= b
-        Core.Intrinsics.atomic_pointerswap(ptr, T(0), :monotonic)
-    else
-        Core.Intrinsics.atomic_pointermodify(ptr, +, T(1), :monotonic)
-    end
+    Core.Intrinsics.atomic_pointermodify(ptr::Ptr{T}, inc, b::T, :monotonic)
 end
 
 function atomic_xchg!(ptr::Ptr{T}, b::T) where T

test/atomic_test.jl

@@ -22,7 +22,7 @@ function atomics_testsuite(backend)
 
     @testset "atomic addition tests" begin
         types = [Int32, Int64, UInt32, UInt64, Float32]
-        if ArrayT != CuArray
+        if ArrayT == CuArray
             CUDA.capability(CUDA.device()) >= v"6.0" && push!(types, Float64)
             CUDA.capability(CUDA.device()) >= v"7.0" && push!(types, Float16)
         else
@@ -82,9 +82,8 @@ function atomics_testsuite(backend)
             kernel = atomic_inc_kernel(backend(), 4)
             wait(kernel(A, T(512), ndrange=(256)))
 
-            @test Array(A)[2] == 257
+            @test Array(A)[2] == 255
         end
     end
 
-
 end