Skip to content

Add atomic memory orderings for Metal #79

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

Open
VEZY wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add atomic memory orderings for Metal #79

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
50 changes: 35 additions & 15 deletions ext/AtomixMetalExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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!(
Expand All @@ -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
Expand Down
51 changes: 42 additions & 9 deletions test/test_atomix_metal.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
using Metal
using Metal: @allowscalar
using UnsafeAtomics: UnsafeAtomics


@testset "AtomixMetalExt:extension_found" begin
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading