Skip to content

Complete the GPU extensions: atomic get/set, swap, float min/max #80

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vchuravy merged 1 commit into from
Sep 8, 2026
Merged

Complete the GPU extensions: atomic get/set, swap, float min/max #80

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 54 additions & 30 deletions ext/AtomixCUDAExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,20 @@
# TODO: respect ordering
module AtomixCUDAExt

using Atomix: Atomix, IndexableRef
using Atomix: Atomix, IndexableRef, right
using CUDA: CUDA, CuDeviceArray
using Core: LLVMPtr

const CuIndexableRef{Indexable<:CuDeviceArray} = IndexableRef{Indexable}

function Atomix.get(ref::CuIndexableRef, order)
error("not implemented")
end
# `Atomix.get` and `Atomix.set!` use Atomix's generic implementation: an LLVM atomic
# load/store on the device pointer (via UnsafeAtomicsLLVM), which honors the ordering.
#
# `Atomix.replace!` and `Atomix.modify!` go through CUDA.jl's intrinsics. These have no
# ordering parameter, so the requested ordering is ignored: every operation is performed
# with CUDA's default (relaxed, device-scope) semantics.

function Atomix.set!(ref::CuIndexableRef, v, order)
error("not implemented")
end
const NativeInt = Union{Int32,Int64,UInt32,UInt64}
const NativeFloat = Union{Float32,Float64}

@inline function Atomix.replace!(
ref::CuIndexableRef,
Expand All @@ -24,35 +26,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
old = CUDA.atomic_cas!(ptr, expected, desired)
return (; old = old, success = old === expected)
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
end
old = modify_native!(ptr, op, x)
return old => op(old, x)
end

# operations with a native atomic instruction
for (op, fn) in [(+) => :atomic_add!, (-) => :atomic_sub!, (&) => :atomic_and!,
(|) => :atomic_or!, xor => :atomic_xor!, min => :atomic_min!,
max => :atomic_max!]
@eval @inline modify_native!(ptr::LLVMPtr{<:NativeInt}, ::typeof($op), x) =
CUDA.$fn(ptr, x)
end
@inline modify_native!(ptr::LLVMPtr{Float32}, ::typeof(+), x) = CUDA.atomic_add!(ptr, x)
@inline modify_native!(ptr::LLVMPtr{Float32}, ::typeof(-), x) = CUDA.atomic_sub!(ptr, x)
# Float64 atomic add needs compute capability 6.0; use compare-and-swap below that.
@inline function modify_native!(ptr::LLVMPtr{Float64}, ::typeof(+), x)
if CUDA.compute_capability() >= v"6.0"
CUDA.atomic_add!(ptr, x)
else
modify_cas!(ptr, +, x)
end
end
@inline modify_native!(ptr::LLVMPtr{Float64}, ::typeof(-), x) = modify_native!(ptr, +, -x)

# swap: exchange floats through their integer representation
@inline modify_native!(ptr::LLVMPtr{<:NativeInt}, ::typeof(right), x) =
CUDA.atomic_xchg!(ptr, x)
@inline function modify_native!(ptr::LLVMPtr{T,A}, ::typeof(right), x) where {T<:NativeFloat,A}
I = CUDA.inttype(T)
old = CUDA.atomic_xchg!(reinterpret(LLVMPtr{I,A}, ptr), reinterpret(I, x))
return reinterpret(T, old)
end

# everything else (float min/max, arbitrary functions): compare-and-swap loop
@inline modify_native!(ptr::LLVMPtr, op, x) = modify_cas!(ptr, op, x)

@inline function modify_cas!(ptr::LLVMPtr{T}, op, x) where {T}
old = Base.unsafe_load(ptr)
while true
new = convert(T, op(old, x))
seen = CUDA.atomic_cas!(ptr, old, new)
# bitwise comparison: `==` would spin forever on NaN
seen === old && return old
old = seen
end
end

end # module AtomixCUDAExt
72 changes: 36 additions & 36 deletions ext/AtomixMetalExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,26 @@
# TODO: respect ordering
module AtomixMetalExt

using Atomix: Atomix, IndexableRef
using Atomix: Atomix, IndexableRef, right
using Metal: Metal, MtlDeviceArray
using Core: LLVMPtr

const MtlIndexableRef{Indexable<:MtlDeviceArray} = IndexableRef{Indexable}

function Atomix.get(ref::MtlIndexableRef, order)
error("not implemented")
# Everything goes through Metal.jl's `atomic_*_explicit` intrinsics, using their default
# (relaxed) memory order: the ordering Atomix is asked for is ignored.

const NativeInt = Union{Int32,UInt32}

@inline function Atomix.get(ref::MtlIndexableRef, order)
ptr = Atomix.pointer(ref)
return Metal.atomic_load_explicit(ptr)
end

function Atomix.set!(ref::MtlIndexableRef, v, order)
error("not implemented")
@inline function Atomix.set!(ref::MtlIndexableRef, v, order)
v = convert(eltype(ref), v)
ptr = Atomix.pointer(ref)
Metal.atomic_store_explicit(ptr, v)
return
end

@inline function Atomix.replace!(
Expand All @@ -24,44 +33,35 @@ 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
old = Metal.atomic_compare_exchange_weak_explicit(ptr, expected, desired)
return (; old = old, success = old === expected)
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}
@inline function Atomix.modify!(ref::MtlIndexableRef, op::OP, x, order) where {OP}
x = convert(eltype(ref), x)
ptr = Atomix.pointer(ref)
old = Metal.atomic_fetch_op_explicit(ptr, op, x)
old = modify_native!(ptr, op, x)
return old => op(old, x)
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
end
return old => op(old, x)
# operations with a native atomic instruction
for (op, fn) in [(+) => :atomic_fetch_add_explicit, (-) => :atomic_fetch_sub_explicit,
(&) => :atomic_fetch_and_explicit, (|) => :atomic_fetch_or_explicit,
xor => :atomic_fetch_xor_explicit, min => :atomic_fetch_min_explicit,
max => :atomic_fetch_max_explicit]
@eval @inline modify_native!(ptr::LLVMPtr{<:NativeInt}, ::typeof($op), x) =
Metal.$fn(ptr, x)
end
# Float32 add/sub is native on device memory only; threadgroup memory takes the
# compare-and-swap path below.
@inline modify_native!(ptr::LLVMPtr{Float32,Metal.AS.Device}, ::typeof(+), x) =
Metal.atomic_fetch_add_explicit(ptr, x)
@inline modify_native!(ptr::LLVMPtr{Float32,Metal.AS.Device}, ::typeof(-), x) =
Metal.atomic_fetch_sub_explicit(ptr, x)
@inline modify_native!(ptr::LLVMPtr{<:Union{NativeInt,Float32}}, ::typeof(right), x) =
Metal.atomic_exchange_explicit(ptr, x)

# everything else (float min/max, arbitrary functions): compare-and-swap loop
@inline modify_native!(ptr::LLVMPtr, op, x) = Metal.atomic_fetch_op_explicit(ptr, op, x)

end # module AtomixMetalExt
78 changes: 47 additions & 31 deletions ext/AtomixOpenCLExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,20 @@
# TODO: respect ordering
module AtomixOpenCLExt

using Atomix: Atomix, IndexableRef
using Atomix: Atomix, IndexableRef, right
using OpenCL: SPIRVIntrinsics, CLDeviceArray
using Core: LLVMPtr

const CLIndexableRef{Indexable<:CLDeviceArray} = IndexableRef{Indexable}

function Atomix.get(ref::CLIndexableRef, order)
error("not implemented")
end
# `Atomix.get` and `Atomix.set!` use Atomix's generic implementation: an LLVM atomic
# load/store on the device pointer (via UnsafeAtomicsLLVM), which honors the ordering.
#
# `Atomix.replace!` and `Atomix.modify!` go through the SPIR-V intrinsics. These have no
# ordering parameter, so the requested ordering is ignored: every operation is performed
# with the intrinsics' default (relaxed, workgroup-scope) semantics.

function Atomix.set!(ref::CLIndexableRef, v, order)
error("not implemented")
end
const NativeInt = Union{Int32,Int64,UInt32,UInt64}
const NativeFloat = Union{Float32,Float64}

@inline function Atomix.replace!(
ref::CLIndexableRef,
Expand All @@ -24,36 +26,50 @@ end
ptr = Atomix.pointer(ref)
expected = convert(eltype(ref), expected)
desired = convert(eltype(ref), desired)
begin
old = SPIRVIntrinsics.atomic_cmpxchg!(ptr, expected, desired)
end
old = SPIRVIntrinsics.atomic_cmpxchg!(ptr, expected, desired)
return (; old = old, success = old === expected)
end

@inline function Atomix.modify!(ref::CLIndexableRef, op::OP, x, order) where {OP}
x = convert(eltype(ref), x)
ptr = Atomix.pointer(ref)
begin
old = if op === (+)
SPIRVIntrinsics.atomic_add!(ptr, x)
elseif op === (-)
SPIRVIntrinsics.atomic_sub!(ptr, x)
elseif op === (&)
SPIRVIntrinsics.atomic_and!(ptr, x)
elseif op === (|)
SPIRVIntrinsics.atomic_or!(ptr, x)
elseif op === xor
SPIRVIntrinsics.atomic_xor!(ptr, x)
elseif op === min
SPIRVIntrinsics.atomic_min!(ptr, x)
elseif op === max
SPIRVIntrinsics.atomic_max!(ptr, x)
else
error("not implemented")
end
end
old = modify_native!(ptr, op, x)
return old => op(old, x)
end

end # module AtomixOpenCLExt
# operations with a native atomic instruction
for (op, fn) in [(+) => :atomic_add!, (-) => :atomic_sub!, (&) => :atomic_and!,
(|) => :atomic_or!, xor => :atomic_xor!, min => :atomic_min!,
max => :atomic_max!]
@eval @inline modify_native!(ptr::LLVMPtr{<:NativeInt}, ::typeof($op), x) =
SPIRVIntrinsics.$fn(ptr, x)
end
@inline modify_native!(ptr::LLVMPtr{<:NativeFloat}, ::typeof(+), x) =
SPIRVIntrinsics.atomic_add!(ptr, x)
@inline modify_native!(ptr::LLVMPtr{<:NativeFloat}, ::typeof(-), x) =
SPIRVIntrinsics.atomic_sub!(ptr, x)
# floating-point min/max: native or compare-and-swap, depending on SPIRVIntrinsics
if hasmethod(SPIRVIntrinsics.atomic_min!, Tuple{LLVMPtr{Float32,1},Float32})
@inline modify_native!(ptr::LLVMPtr{<:NativeFloat}, ::typeof(min), x) =
SPIRVIntrinsics.atomic_min!(ptr, x)
@inline modify_native!(ptr::LLVMPtr{<:NativeFloat}, ::typeof(max), x) =
SPIRVIntrinsics.atomic_max!(ptr, x)
end
@inline modify_native!(ptr::LLVMPtr{<:Union{NativeInt,NativeFloat}}, ::typeof(right), x) =
SPIRVIntrinsics.atomic_xchg!(ptr, x)

# everything else: compare-and-swap loop
@inline modify_native!(ptr::LLVMPtr, op, x) = modify_cas!(ptr, op, x)

@inline function modify_cas!(ptr::LLVMPtr{T}, op, x) where {T}
old = Base.unsafe_load(ptr)
while true
new = convert(T, op(old, x))
seen = SPIRVIntrinsics.atomic_cmpxchg!(ptr, old, new)
# bitwise comparison: `==` would spin forever on NaN
seen === old && return old
old = seen
end
end

end # module AtomixOpenCLExt
Loading
Loading