Add GPU support for Complex atomics and refactor to shared utilities

- Extract shared Complex utilities to src/complex.jl
- Add ComplexF32/ComplexF64 support to CUDA extension (via CAS loops)
- Add ComplexF32 support to Metal extension (ComplexF64 not supported due to 128-bit limit)
- Refactor CPU implementation to use shared utilities
- Eliminates code duplication across backends

Author: albertomercurio

ext/AtomixCUDAExt.jl

@@ -2,6 +2,7 @@
 module AtomixCUDAExt
 
 using Atomix: Atomix, IndexableRef
+using Atomix.Internal: _int_type_for_complex, _to_int, _from_int, _with_int_repr, _generic_cas_loop_modify!
 using CUDA: CUDA, CuDeviceArray
 
 const CuIndexableRef{Indexable<:CuDeviceArray} = IndexableRef{Indexable}
@@ -24,35 +25,57 @@ end
     ptr = Atomix.pointer(ref)
     expected = convert(eltype(ref), expected)
     desired = convert(eltype(ref), desired)
-    begin
-        old = CUDA.atomic_cas!(ptr, expected, desired)
-    end
-    return (; old = old, success = old === expected)
+    _cuda_cas!(ptr, expected, desired)
+end
+
+# CUDA CAS - with Complex support via integer reinterpretation
+@inline function _cuda_cas!(ptr::Ptr{T}, expected::T, desired::T) where {T}
+    result = _with_int_repr(CUDA.atomic_cas!, ptr, expected, desired)
+    return _from_int(T, (; old = result, success = result === expected))
+end
+
+# CUDA load via CAS (for Complex support)
+@inline function _cuda_load(ptr::Ptr{T}) where {T}
+    _with_int_repr(_cuda_load_impl, ptr)
+end
+
+@inline function _cuda_load_impl(ptr::Ptr{T}) where {T}
+    # Load via CAS with same value
+    old = unsafe_load(ptr)
+    CUDA.atomic_cas!(ptr, old, old)
 end
 
 @inline function Atomix.modify!(ref::CuIndexableRef, op::OP, x, order) where {OP}
     x = convert(eltype(ref), x)
     ptr = Atomix.pointer(ref)
-    begin
-        old = if op === (+)
-            CUDA.atomic_add!(ptr, x)
-        elseif op === (-)
-            CUDA.atomic_sub!(ptr, x)
-        elseif op === (&)
-            CUDA.atomic_and!(ptr, x)
-        elseif op === (|)
-            CUDA.atomic_or!(ptr, x)
-        elseif op === xor
-            CUDA.atomic_xor!(ptr, x)
-        elseif op === min
-            CUDA.atomic_min!(ptr, x)
-        elseif op === max
-            CUDA.atomic_max!(ptr, x)
-        else
-            error("not implemented")
-        end
+    _cuda_modify!(ptr, op, x)
+end
+
+# CUDA modify! - native operations for non-Complex types
+@inline function _cuda_modify!(ptr::Ptr{T}, op::OP, x::T) where {T,OP}
+    old = if op === (+)
+        CUDA.atomic_add!(ptr, x)
+    elseif op === (-)
+        CUDA.atomic_sub!(ptr, x)
+    elseif op === (&)
+        CUDA.atomic_and!(ptr, x)
+    elseif op === (|)
+        CUDA.atomic_or!(ptr, x)
+    elseif op === xor
+        CUDA.atomic_xor!(ptr, x)
+    elseif op === min
+        CUDA.atomic_min!(ptr, x)
+    elseif op === max
+        CUDA.atomic_max!(ptr, x)
+    else
+        error("not implemented")
     end
     return old => op(old, x)
 end
 
+# CUDA modify! - CAS loop for Complex types
+@inline function _cuda_modify!(ptr::Ptr{Complex{T}}, op::OP, x::Complex{T}) where {T,OP}
+    _generic_cas_loop_modify!(_cuda_cas!, _cuda_load, ptr, op, x, nothing)
+end
+
 end  # module AtomixCUDAExt

ext/AtomixMetalExt.jl

@@ -2,6 +2,7 @@
 module AtomixMetalExt
 
 using Atomix: Atomix, IndexableRef
+using Atomix.Internal: _int_type_for_complex, _to_int, _from_int, _with_int_repr, _generic_cas_loop_modify!
 using Metal: Metal, MtlDeviceArray
 
 const MtlIndexableRef{Indexable<:MtlDeviceArray} = IndexableRef{Indexable}
@@ -24,12 +25,39 @@ end
     ptr = Atomix.pointer(ref)
     expected = convert(eltype(ref), expected)
     desired = convert(eltype(ref), desired)
-    begin
-        old = Metal.atomic_compare_exchange_weak_explicit(ptr, expected, desired)
-    end
+    _metal_cas!(ptr, expected, desired)
+end
+
+# Metal CAS - with ComplexF32 support (ComplexF64 not supported due to 128-bit limitation)
+@inline function _metal_cas!(ptr::Ptr{ComplexF32}, expected::ComplexF32, desired::ComplexF32)
+    IntType = UInt64
+    int_ptr = reinterpret(Ptr{IntType}, ptr)
+    int_expected = reinterpret(IntType, expected)
+    int_desired = reinterpret(IntType, desired)
+    old = Metal.atomic_compare_exchange_weak_explicit(int_ptr, int_expected, int_desired)
+    old_complex = reinterpret(ComplexF32, old)
+    return (; old = old_complex, success = old_complex === expected)
+end
+
+@inline function _metal_cas!(ptr::Ptr{T}, expected::T, desired::T) where {T}
+    old = Metal.atomic_compare_exchange_weak_explicit(ptr, expected, desired)
     return (; old = old, success = old === expected)
 end
 
+# Metal load via CAS (for Complex support)
+@inline function _metal_load(ptr::Ptr{ComplexF32})
+    IntType = UInt64
+    int_ptr = reinterpret(Ptr{IntType}, ptr)
+    old = unsafe_load(int_ptr)
+    result = Metal.atomic_compare_exchange_weak_explicit(int_ptr, old, old)
+    return reinterpret(ComplexF32, result)
+end
+
+@inline function _metal_load(ptr::Ptr{T}) where {T}
+    old = unsafe_load(ptr)
+    Metal.atomic_compare_exchange_weak_explicit(ptr, old, old)
+end
+
 
 # CAS is needed for FP ops on ThreadGroup memory
 @inline function Atomix.modify!(ref::IndexableRef{<:MtlDeviceArray{<:AbstractFloat, <:Any, Metal.AS.ThreadGroup}} , op::OP, x, order) where {OP}
@@ -42,26 +70,34 @@ end
 @inline function Atomix.modify!(ref::MtlIndexableRef, op::OP, x, order) where {OP}
     x = convert(eltype(ref), x)
     ptr = Atomix.pointer(ref)
-    begin
-        old = if op === (+)
-            Metal.atomic_fetch_add_explicit(ptr, x)
-        elseif op === (-)
-            Metal.atomic_fetch_sub_explicit(ptr, x)
-        elseif op === (&)
-            Metal.atomic_fetch_and_explicit(ptr, x)
-        elseif op === (|)
-            Metal.atomic_fetch_or_explicit(ptr, x)
-        elseif op === xor
-            Metal.atomic_fetch_xor_explicit(ptr, x)
-        elseif op === min
-            Metal.atomic_fetch_min_explicit(ptr, x)
-        elseif op === max
-            Metal.atomic_fetch_max_explicit(ptr, x)
-        else
-            error("not implemented")
-        end
+    _metal_modify!(ptr, op, x)
+end
+
+# Metal modify! - native operations for non-Complex types
+@inline function _metal_modify!(ptr::Ptr{T}, op::OP, x::T) where {T,OP}
+    old = if op === (+)
+        Metal.atomic_fetch_add_explicit(ptr, x)
+    elseif op === (-)
+        Metal.atomic_fetch_sub_explicit(ptr, x)
+    elseif op === (&)
+        Metal.atomic_fetch_and_explicit(ptr, x)
+    elseif op === (|)
+        Metal.atomic_fetch_or_explicit(ptr, x)
+    elseif op === xor
+        Metal.atomic_fetch_xor_explicit(ptr, x)
+    elseif op === min
+        Metal.atomic_fetch_min_explicit(ptr, x)
+    elseif op === max
+        Metal.atomic_fetch_max_explicit(ptr, x)
+    else
+        error("not implemented")
     end
     return old => op(old, x)
 end
 
+# Metal modify! - CAS loop for ComplexF32 only (ComplexF64 not supported)
+@inline function _metal_modify!(ptr::Ptr{ComplexF32}, op::OP, x::ComplexF32) where {OP}
+    _generic_cas_loop_modify!(_metal_cas!, _metal_load, ptr, op, x, nothing)
+end
+
 end  # module AtomixMetalExt

src/Atomix.jl

@@ -32,6 +32,7 @@ using UnsafeAtomics:
 
 include("utils.jl")
 include("references.jl")
+include("complex.jl")
 include("generic.jl")
 include("core.jl")
 include("sugar.jl")

src/complex.jl

@@ -0,0 +1,35 @@
+# Shared utilities for Complex number atomic operations across backends
+
+# Integer type mapping for Complex types
+_int_type_for_complex(::Type{Float32}) = UInt64
+_int_type_for_complex(::Type{Float64}) = UInt128
+
+# Convert to/from integer representation
+_to_int(::Type{I}, x::Complex{T}) where {I,T} = reinterpret(I, x)
+_to_int(::Type{I}, x) where {I} = x
+
+_from_int(::Type{Complex{T}}, x::Integer) where {T} = reinterpret(Complex{T}, x)
+_from_int(::Type{Complex{T}}, result::NamedTuple) where {T} = 
+    (old = reinterpret(Complex{T}, result.old), success = result.success)
+_from_int(::Type{T}, x) where {T} = x
+
+# Helper: apply atomic operation with integer reinterpretation for Complex types
+function _with_int_repr(f, ptr::Ptr{Complex{T}}, args...) where {T}
+    IntType = _int_type_for_complex(T)
+    int_ptr = reinterpret(Ptr{IntType}, ptr)
+    result = f(int_ptr, _to_int.(IntType, args)...)
+    return _from_int(Complex{T}, result)
+end
+
+_with_int_repr(f, ptr::Ptr{T}, args...) where {T} = f(ptr, args...)
+
+# Generic CAS loop for modify! operations (used when native atomics unavailable)
+function _generic_cas_loop_modify!(cas_fn, load_fn, ptr::Ptr{T}, op::OP, x::T, ord) where {T,OP}
+    old = load_fn(ptr, ord)
+    while true
+        new = op(old, x)
+        result = cas_fn(ptr, old, new, ord, ord)
+        result.success && return (old => new)
+        old = result.old
+    end
+end

src/core.jl

@@ -53,36 +53,7 @@ end
 
 # CAS loop fallback for Complex types (no native atomic modify!)
 function _atomic_modify!(ptr::Ptr{Complex{T}}, op::OP, x::Complex{T}, ord) where {T,OP}
-    old = _atomic_load(ptr, ord)
-    while true
-        new = op(old, x)
-        result = _atomic_cas!(ptr, old, new, ord, ord)
-        result.success && return (old => new)
-        old = result.old
-    end
-end
-
-# Helper: apply atomic operation with integer reinterpretation for Complex types
-function _with_int_repr(f, ptr::Ptr{Complex{T}}, args...) where {T}
-    IntType = _int_type_for_complex(T)
-    int_ptr = reinterpret(Ptr{IntType}, ptr)
-    result = f(int_ptr, _to_int.(IntType, args)...)
-    return _from_int(Complex{T}, result)
+    _generic_cas_loop_modify!(_atomic_cas!, _atomic_load, ptr, op, x, ord)
 end
 
-_with_int_repr(f, ptr::Ptr{T}, args...) where {T} = f(ptr, args...)
-
-# Integer type mapping for Complex types
-_int_type_for_complex(::Type{Float32}) = UInt64
-_int_type_for_complex(::Type{Float64}) = UInt128
-
-# Convert to/from integer representation
-_to_int(::Type{I}, x::Complex{T}) where {I,T} = reinterpret(I, x)
-_to_int(::Type{I}, x) where {I} = x
-
-_from_int(::Type{Complex{T}}, x::Integer) where {T} = reinterpret(Complex{T}, x)
-_from_int(::Type{Complex{T}}, result::NamedTuple) where {T} = 
-    (old = reinterpret(Complex{T}, result.old), success = result.success)
-_from_int(::Type{T}, x) where {T} = x
-
 Atomix.asstorable(ref, v) = convert(eltype(ref), v)