diff --git a/ext/AtomixMetalExt.jl b/ext/AtomixMetalExt.jl index af98afb..c0eca13 100644 --- a/ext/AtomixMetalExt.jl +++ b/ext/AtomixMetalExt.jl @@ -1,17 +1,27 @@ -# TODO: respect ordering module AtomixMetalExt using Atomix: Atomix, IndexableRef using Metal: Metal, MtlDeviceArray +using UnsafeAtomics: UnsafeAtomics const MtlIndexableRef{Indexable<:MtlDeviceArray} = IndexableRef{Indexable} -function Atomix.get(ref::MtlIndexableRef, order) - error("not implemented") +@inline metal_memory_order(::typeof(UnsafeAtomics.unordered)) = Val(Metal.memory_order_relaxed) +@inline metal_memory_order(::typeof(UnsafeAtomics.monotonic)) = Val(Metal.memory_order_relaxed) +@inline metal_memory_order(::typeof(UnsafeAtomics.acquire)) = Val(Metal.memory_order_acquire) +@inline metal_memory_order(::typeof(UnsafeAtomics.release)) = Val(Metal.memory_order_release) +@inline metal_memory_order(::typeof(UnsafeAtomics.acq_rel)) = Val(Metal.memory_order_acq_rel) +@inline metal_memory_order(::typeof(UnsafeAtomics.seq_cst)) = Val(Metal.memory_order_seq_cst) + +@inline function Atomix.get(ref::MtlIndexableRef, order) + ptr = Atomix.pointer(ref) + return Metal.atomic_load_explicit(ptr, metal_memory_order(order)) 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) + return Metal.atomic_store_explicit(ptr, v, metal_memory_order(order)) end @inline function Atomix.replace!( @@ -25,38 +35,48 @@ end expected = convert(eltype(ref), expected) desired = convert(eltype(ref), desired) begin - old = Metal.atomic_compare_exchange_weak_explicit(ptr, expected, desired) + old = Metal.atomic_compare_exchange_weak_explicit( + ptr, expected, desired, + metal_memory_order(success_ordering), + metal_memory_order(failure_ordering), + ) end 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::IndexableRef{<:MtlDeviceArray{<:AbstractFloat, <:Any, Metal.AS.ThreadGroup}}, + 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 = Metal.atomic_fetch_op_explicit(ptr, op, x, metal_memory_order(order)) 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) + metal_order = metal_memory_order(order) begin old = if op === (+) - Metal.atomic_fetch_add_explicit(ptr, x) + Metal.atomic_fetch_add_explicit(ptr, x, metal_order) elseif op === (-) - Metal.atomic_fetch_sub_explicit(ptr, x) + Metal.atomic_fetch_sub_explicit(ptr, x, metal_order) elseif op === (&) - Metal.atomic_fetch_and_explicit(ptr, x) + Metal.atomic_fetch_and_explicit(ptr, x, metal_order) elseif op === (|) - Metal.atomic_fetch_or_explicit(ptr, x) + Metal.atomic_fetch_or_explicit(ptr, x, metal_order) elseif op === xor - Metal.atomic_fetch_xor_explicit(ptr, x) + Metal.atomic_fetch_xor_explicit(ptr, x, metal_order) elseif op === min - Metal.atomic_fetch_min_explicit(ptr, x) + Metal.atomic_fetch_min_explicit(ptr, x, metal_order) elseif op === max - Metal.atomic_fetch_max_explicit(ptr, x) + Metal.atomic_fetch_max_explicit(ptr, x, metal_order) else error("not implemented") end diff --git a/test/test_atomix_metal.jl b/test/test_atomix_metal.jl index 051c9e0..bcf34ad 100644 --- a/test/test_atomix_metal.jl +++ b/test/test_atomix_metal.jl @@ -1,5 +1,6 @@ using Metal using Metal: @allowscalar +using UnsafeAtomics: UnsafeAtomics @testset "AtomixMetalExt:extension_found" begin @@ -16,20 +17,19 @@ function metal(f) end -# Not implemented: -#= -function test_get_set() - A = CUDA.ones(Int, 3) - cuda() do +@testset "AtomixMetalExt:test_get_set" begin + A = Metal.ones(Int32, 3) + metal() do GC.@preserve A begin ref = Atomix.IndexableRef(A, (1,)) - x = Atomix.get(ref) - Atomix.set!(ref, -x) + Atomix.set!(ref, -Atomix.get(ref, Atomix.acquire), Atomix.release) + A[2] = Atomix.get(ref, Atomix.acquire) + Atomix.set!(ref, Int32(7), UnsafeAtomics.unordered) + A[3] = Atomix.get(ref, Atomix.monotonic) end end - @test collect(A) == [-1, 1, 1] + @test collect(A) == Int32[7, -1, 7] end -=# @testset "AtomixMetalExt:test_cas" begin @@ -86,6 +86,39 @@ end end +@testset "AtomixMetalExt:test_orderings" begin + A = Metal.zeros(Int32, 4) + metal() do + GC.@preserve A begin + ref = Atomix.IndexableRef(A, (1,)) + Atomix.modify!(ref, +, 1, UnsafeAtomics.unordered) + Atomix.modify!(ref, +, 1, Atomix.monotonic) + Atomix.modify!(ref, +, 1, Atomix.acquire) + Atomix.modify!(ref, +, 1, Atomix.release) + Atomix.modify!(ref, +, 1, Atomix.acquire_release) + old, success = Atomix.replace!( + ref, + 5, + 10, + Atomix.sequentially_consistent, + Atomix.acquire, + ) + A[2] = old + A[3] = success + old, success = Atomix.replace!( + ref, + 5, + 11, + Atomix.acq_rel, + Atomix.sequentially_consistent, + ) + A[4] = old + success + end + end + @test collect(A) == Int32[10, 5, 1, 10] +end + + @testset "AtomixMetalExt:test_inc_sugar" begin A = Metal.ones(Int32, 3) metal() do