Dispatch to LLVM atomic OP (#21)

Co-authored-by: Valentin Churavy <v.churavy@gmail.com>

Author: pxl-th

Project.toml

@@ -17,6 +17,7 @@ julia = "1.10"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestFunctionRunner = "792026f5-ac9a-4a19-adcb-47b0ce2deb5d"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 
 [targets]
-test = ["Test", "TestFunctionRunner", "LLVM"]
+test = ["Test", "TestFunctionRunner", "LLVM", "InteractiveUtils"]

ext/UnsafeAtomicsLLVM/atomics.jl

@@ -257,6 +257,48 @@ end
     return old => x
 end
 
+const atomictypes = Any[
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Int128,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+    UInt128,
+    Float16,
+    Float32,
+    Float64,
+]
+
+for (opname, op, llvmop) in binoptable
+    opname === :xchg && continue
+    types = if opname in (:min, :max)
+        filter(t -> t <: Signed, atomictypes)
+    elseif opname in (:umin, :umax)
+        filter(t -> t <: Unsigned, atomictypes)
+    elseif opname in (:fadd, :fsub, :fmin, :fmax)
+        filter(t -> t <: AbstractFloat, atomictypes)
+    else
+        filter(t -> t <: Integer, atomictypes)
+    end
+    for T in types
+        @eval @inline function atomic_pointermodify(
+            ptr::LLVMPtr{$T},
+            ::$(typeof(op)),
+            x::$T,
+            order::AtomicOrdering,
+            sync::Val{S},
+        ) where {S}
+            old = llvm_atomic_op(
+                $(Val(llvmop)), ptr, x, llvm_from_julia_ordering(order), sync)
+            return old => $op(old, x)
+        end
+    end
+end
+
 # @inline atomic_pointerswap(pointer, new) = first(atomic_pointermodify(pointer, right, new))
 @inline atomic_pointerswap(pointer, new, order, sync) =
     first(atomic_pointermodify(pointer, right, new, order, sync))

test/UnsafeAtomicsLLVM.jl

@@ -1,4 +1,5 @@
 import LLVM
+import InteractiveUtils
 
 using UnsafeAtomics: UnsafeAtomics, acquire, release, acq_rel, seq_cst
 using UnsafeAtomics.Internal: OP_RMW_TABLE, inttypes
@@ -34,6 +35,14 @@ function check_default_ordering(xs::AbstractArray{T}, x1::T, x2::T) where T
             xs[1] = x1
             @test rmw(ptr, x2) === x1
             @test xs[1] === op(x1, x2)
+
+            # Check dispatch to LLVM atomic OP instead of CAS loop.
+            if (op == +) || (op == -)
+                IR = sprint(io->InteractiveUtils.code_llvm(io,
+                    UnsafeAtomics.modify!,
+                    typeof.((ptr, +, T(1)))))
+                @test occursin("atomicrmw", IR)
+            end
         end
     end
 end
@@ -80,6 +89,14 @@ function test_explicit_ordering(xs::AbstractArray{T}, x1::T, x2::T) where T
                       (x1 => op(x1, x2))
                 @test xs[1] === op(x1, x2)
             end
+
+            # Check dispatch to LLVM atomic OP instead of CAS loop.
+            if (op == +) || (op == -)
+                IR = sprint(io->InteractiveUtils.code_llvm(io,
+                    UnsafeAtomics.modify!,
+                    typeof.((ptr, +, T(1), seq_cst, UnsafeAtomics.singlethread))))
+                @test occursin("atomicrmw", IR)
+            end
         end
     end
 end