Expose syncscope parameter for atomic add/sub (#14)

Author: pxl-th

.github/workflows/ci.yml

@@ -13,8 +13,7 @@ jobs:
     strategy:
       matrix:
         julia-version:
-          - '1'
-          - '1.6'
+          - '1.10'
           - 'nightly'
       fail-fast: false
     name: Test Julia ${{ matrix.julia-version }}
@@ -36,8 +35,7 @@ jobs:
     strategy:
       matrix:
         julia-version:
-          - '1'
-          - '1.6'
+          - '1.10'
           - 'nightly'
       fail-fast: false
     steps:

Project.toml

@@ -1,13 +1,13 @@
 name = "UnsafeAtomicsLLVM"
 uuid = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
 authors = ["Takafumi Arakaki <[email protected]> and contributors"]
-version = "0.1.5"
+version = "0.2.0"
 
 [deps]
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 
 [compat]
-LLVM = "6, 7, 8"
+LLVM = "8.1"
 UnsafeAtomics = "0.2"
-julia = "1.6"
+julia = "1.10"

src/atomics.jl

@@ -248,14 +248,15 @@ const binoptable = [
     (:umin, min, LLVM.API.LLVMAtomicRMWBinOpUMin),
     (:fadd, +, LLVM.API.LLVMAtomicRMWBinOpFAdd),
     (:fsub, -, LLVM.API.LLVMAtomicRMWBinOpFSub),
+    (:fmax, max, LLVM.API.LLVMAtomicRMWBinOpFMax),
+    (:fmin, min, LLVM.API.LLVMAtomicRMWBinOpFMin),
 ]
-if VERSION ≥ v"1.10-"
-    push!(binoptable, (:fmax, max, LLVM.API.LLVMAtomicRMWBinOpFMax))
-    push!(binoptable, (:fmin, min, LLVM.API.LLVMAtomicRMWBinOpFMin))
-end
 
 const AtomicRMWBinOpVal = Union{(Val{binop} for (_, _, binop) in binoptable)...}
 
+# LLVM API accepts string literal as a syncscope argument.
+@inline syncscope_to_string(::Type{Val{S}}) where {S} = string(S)
+
 @generated function llvm_atomic_op(
     binop::AtomicRMWBinOpVal,
     ptr::LLVMPtr{T,A},
@@ -293,6 +294,40 @@ const AtomicRMWBinOpVal = Union{(Val{binop} for (_, _, binop) in binoptable)...}
     end
 end
 
+@generated function llvm_atomic_op(
+    binop::AtomicRMWBinOpVal,
+    ptr::LLVMPtr{T,A},
+    val::T,
+    order::LLVMOrderingVal,
+    syncscope::Val{S},
+) where {T,A,S}
+    @dispose ctx = Context() begin
+        T_val = convert(LLVMType, T)
+        T_ptr = convert(LLVMType, ptr)
+
+        T_typed_ptr = LLVM.PointerType(T_val, A)
+        llvm_f, _ = create_function(T_val, [T_ptr, T_val])
+
+        @dispose builder = IRBuilder() begin
+            entry = BasicBlock(llvm_f, "entry")
+            position!(builder, entry)
+
+            typed_ptr = bitcast!(builder, parameters(llvm_f)[1], T_typed_ptr)
+            rv = atomic_rmw!(
+                builder,
+                _valueof(binop()),
+                typed_ptr,
+                parameters(llvm_f)[2],
+                _valueof(order()),
+                syncscope_to_string(syncscope),
+            )
+
+            ret!(builder, rv)
+        end
+        call_function(llvm_f, T, Tuple{LLVMPtr{T,A},T}, :ptr, :val)
+    end
+end
+
 @inline function atomic_pointermodify(pointer, op::OP, x, order::Symbol) where {OP}
     @dynamic_order(order) do order
         atomic_pointermodify(pointer, op, x, order)
@@ -359,8 +394,18 @@ for (opname, op, llvmop) in binoptable
             ::$(typeof(op)),
             x::$T,
             order::AtomicOrdering,
-        )
-            old = llvm_atomic_op($(Val(llvmop)), ptr, x, llvm_from_julia_ordering(order))
+            syncscope::Val{S} = Val{:system}(),
+        ) where {S}
+            old =
+                syncscope isa Val{:system} ?
+                llvm_atomic_op($(Val(llvmop)), ptr, x, llvm_from_julia_ordering(order)) :
+                llvm_atomic_op(
+                    $(Val(llvmop)),
+                    ptr,
+                    x,
+                    llvm_from_julia_ordering(order),
+                    syncscope,
+                )
             return old => $op(old, x)
         end
     end

src/internal.jl

@@ -20,6 +20,15 @@ mapop(::typeof(UnsafeAtomics.right)) = right
 @inline UnsafeAtomics.modify!(ptr::LLVMPtr, op::OP, x, order::Ordering) where {OP} =
     atomic_pointermodify(ptr, mapop(op), x, Val{julia_ordering_name(order)}())
 
+@inline UnsafeAtomics.modify!(
+    ptr::LLVMPtr,
+    op::OP,
+    x,
+    order::Ordering,
+    syncscope::Val{S} = Val(:system),
+) where {OP<:Union{typeof(+),typeof(-)},S} =
+    atomic_pointermodify(ptr, mapop(op), x, Val{julia_ordering_name(order)}(), syncscope)
+
 @inline UnsafeAtomics.cas!(
     ptr::LLVMPtr,
     expected,

test/runtests.jl

@@ -1,6 +1,6 @@
 import UnsafeAtomicsLLVM
 
-using UnsafeAtomics: UnsafeAtomics, acquire, release, acq_rel
+using UnsafeAtomics: UnsafeAtomics, acquire, release, acq_rel, seq_cst
 using UnsafeAtomics.Internal: OP_RMW_TABLE, inttypes
 using Test
 
@@ -57,6 +57,19 @@ function test_explicit_ordering(T::Type = UInt)
             xs[1] = x1
             @test rmw(ptr, x2, acquire) === x1
             @test xs[1] === op(x1, x2)
+
+            # Test syncscopes.
+            if (op == +) || (op == -)
+                xs[1] = x1
+                @test UnsafeAtomics.modify!(ptr, op, x2, seq_cst, Val(:system)) ===
+                      (x1 => op(x1, x2))
+                @test xs[1] === op(x1, x2)
+
+                xs[1] = x1
+                @test UnsafeAtomics.modify!(ptr, op, x2, seq_cst, Val(:singlethread)) ===
+                      (x1 => op(x1, x2))
+                @test xs[1] === op(x1, x2)
+            end
         end
     end
 end