Add ordered atomics

[VEZY]

Note: I used Codex with GPT5.5 High to implement this feature. For context, my original issue with this problem was about building a correct TLAS on the GPU with RayCore.jl (which failed): JuliaGeometry/Raycore.jl#19

Summary

This adds explicit memory-order support to Metal’s low-level atomic intrinsics and enables KernelAbstractions atomic support for MetalBackend.

Metal atomic operations previously exposed only relaxed ordering through the Julia API. That is fine for counters and reductions, but it is not enough for synchronization patterns where ordinary device-memory writes are published through an atomic operation and then read by another workitem.

Changes:

• Add Metal memory_order enum values for relaxed, acquire, release, acq_rel, and seq_cst.
• Add explicit ordering arguments to low-level atomic load/store/exchange/fetch/CAS functions.
• Preserve relaxed defaults for existing direct Metal atomic call sites.
• Add compare-exchange success/failure ordering support.
• Use Val-specialized ordering paths so GPU code receives compile-time ordering constants.
• Enable KernelAbstractions.supports_atomics(::MetalBackend).
• Add a regression test for a cross-workitem publish/consume pattern using device-memory writes plus an ordered atomic fetch_add.

Why

KernelAbstractions and Atomix can express atomic operations with acquire/release or sequentially-consistent semantics. Metal needs to receive those requested orderings for synchronization-sensitive GPU algorithms to be correct.

The new regression test models a common bottom-up refit pattern: each workitem writes ordinary device memory, then increments an atomic counter; the second arriving workitem reads both children’s ordinary writes and publishes the parent. This requires the atomic operation to provide ordering, while still allowing relaxed atomics for cases that do not need synchronization.

Validation

• Metal atomic tests passed.
• Metal KernelAbstractions integration test passed.
• Downstream RayCore TLAS/BLAS Metal repro passed with Atomix orderings wired through to Metal.
• ArchimedLight-derived static TLAS comparison passed: CPU-built/adapted and direct Metal-built TLAS match for nodes, instances, all BLAS buffers, descriptors, and root AABB.

[VEZY]

Of course I also kept the default memory order to relaxed.

[codecov]

Codecov Report

❌ Patch coverage is 91.07143% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.00%. Comparing base (bc55410) to head (ae30a4b).

Files with missing lines	Patch %	Lines
src/version.jl	50.00%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #870      +/-   ##
==========================================
+ Coverage   85.93%   86.00%   +0.06%     
==========================================
  Files          76       76              
  Lines        5246     5300      +54     
==========================================
+ Hits         4508     4558      +50     
- Misses        738      742       +4     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Metal Benchmarks

Details

Benchmark suite	Current: ae30a4b	Previous: bc55410	Ratio
array/accumulate/Float32/1d	405083 ns	403958 ns	1.00
array/accumulate/Float32/dims=1	383334 ns	374833 ns	1.02
array/accumulate/Float32/dims=1L	8805750 ns	8802042 ns	1.00
array/accumulate/Float32/dims=2	445042 ns	438667 ns	1.01
array/accumulate/Float32/dims=2L	2594375 ns	2654292 ns	0.98
array/accumulate/Int64/1d	850125 ns	848375 ns	1.00
array/accumulate/Int64/dims=1	945792 ns	939209 ns	1.01
array/accumulate/Int64/dims=1L	9506500 ns	9523000 ns	1.00
array/accumulate/Int64/dims=2	1251958 ns	932167 ns	1.34
array/accumulate/Int64/dims=2L	6516042 ns	6533834 ns	1.00
array/broadcast	227667 ns	218208 ns	1.04
array/construct	2375 ns	2291 ns	1.04
array/permutedims/2d	446166 ns	446417 ns	1.00
array/permutedims/3d	919917 ns	724583 ns	1.27
array/permutedims/4d	1174375 ns	1184500 ns	0.99
array/private/copy	237208 ns	230834 ns	1.03
array/private/copyto!/cpu_to_gpu	195333 ns	194042 ns	1.01
array/private/copyto!/gpu_to_cpu	197000 ns	192667 ns	1.02
array/private/copyto!/gpu_to_gpu	168541 ns	207209 ns	0.81
array/private/iteration/findall/bool	1080167 ns	1063666 ns	1.02
array/private/iteration/findall/int	1240792 ns	1243542 ns	1.00
array/private/iteration/findfirst/bool	1120833 ns	1125125 ns	1.00
array/private/iteration/findfirst/int	1138250 ns	1108708 ns	1.03
array/private/iteration/findmin/1d	1233375 ns	1220458 ns	1.01
array/private/iteration/findmin/2d	1072875 ns	1069125 ns	1.00
array/private/iteration/logical	1743208 ns	1698667 ns	1.03
array/private/iteration/scalar	1291667 ns	1211541 ns	1.07
array/random/rand/Float32	433875 ns	434041 ns	1.00
array/random/rand/Int64	522292 ns	473500 ns	1.10
array/random/rand!/Float32	382333 ns	392292 ns	0.97
array/random/rand!/Int64	404583 ns	324541 ns	1.25
array/random/randn/Float32	373167 ns	390416 ns	0.96
array/random/randn!/Float32	342375 ns	291667 ns	1.17
array/reductions/mapreduce/Float32/1d	402250 ns	409625 ns	0.98
array/reductions/mapreduce/Float32/dims=1	289541 ns	284083 ns	1.02
array/reductions/mapreduce/Float32/dims=1L	619167 ns	624042 ns	0.99
array/reductions/mapreduce/Float32/dims=2	328125 ns	278375 ns	1.18
array/reductions/mapreduce/Float32/dims=2L	996584 ns	998792 ns	1.00
array/reductions/mapreduce/Int64/1d	587125 ns	589208 ns	1.00
array/reductions/mapreduce/Int64/dims=1	625875 ns	621458 ns	1.01
array/reductions/mapreduce/Int64/dims=1L	1024791 ns	1030666 ns	0.99
array/reductions/mapreduce/Int64/dims=2	781834 ns	779208 ns	1.00
array/reductions/mapreduce/Int64/dims=2L	2179500 ns	2182458 ns	1.00
array/reductions/reduce/Float32/1d	403291 ns	420458 ns	0.96
array/reductions/reduce/Float32/dims=1	282625 ns	336833 ns	0.84
array/reductions/reduce/Float32/dims=1L	626250 ns	631834 ns	0.99
array/reductions/reduce/Float32/dims=2	221625 ns	216042 ns	1.03
array/reductions/reduce/Float32/dims=2L	445417 ns	447209 ns	1.00
array/reductions/reduce/Int64/1d	586041 ns	608791 ns	0.96
array/reductions/reduce/Int64/dims=1	630292 ns	629583 ns	1.00
array/reductions/reduce/Int64/dims=1L	1016334 ns	1026792 ns	0.99
array/reductions/reduce/Int64/dims=2	244125 ns	251792 ns	0.97
array/reductions/reduce/Int64/dims=2L	646209 ns	639833 ns	1.01
array/shared/copy	135792 ns	135417 ns	1.00
array/shared/copyto!/cpu_to_gpu	37750 ns	38333 ns	0.98
array/shared/copyto!/gpu_to_cpu	38250 ns	38708 ns	0.99
array/shared/copyto!/gpu_to_gpu	38541 ns	39250 ns	0.98
array/shared/iteration/findall/bool	1095208 ns	1070708 ns	1.02
array/shared/iteration/findall/int	1244750 ns	1242875 ns	1.00
array/shared/iteration/findfirst/bool	926500 ns	959833 ns	0.97
array/shared/iteration/findfirst/int	671916 ns	961875 ns	0.70
array/shared/iteration/findmin/1d	1094000 ns	1095791 ns	1.00
array/shared/iteration/findmin/2d	1074375 ns	1072458 ns	1.00
array/shared/iteration/logical	1582833 ns	1552541 ns	1.02
array/shared/iteration/scalar	4113.142857142857 ns	4047.5714285714284 ns	1.02
array/sorting/1d	1821875 ns	1788125 ns	1.02
array/sorting/2d	8304583 ns	8305084 ns	1.00
integration/byval/reference	1113834 ns	1102917 ns	1.01
integration/byval/slices=1	1105750 ns	1105833 ns	1.00
integration/byval/slices=2	2022625 ns	2025542 ns	1.00
integration/byval/slices=3	6592042 ns	6473959 ns	1.02
integration/metaldevrt	383583 ns	380750 ns	1.01
kernel/indexing	192458 ns	194000 ns	0.99
kernel/indexing_checked	379584 ns	382750 ns	0.99
kernel/launch	1862.5 ns	1887.5 ns	0.99
kernel/rand	373334 ns	312167 ns	1.20
latency/import	2038421792 ns	2034342500 ns	1.00
latency/precompile	39325503166 ns	38780551250 ns	1.01
latency/ttfp	2379460667 ns	2371294750 ns	1.00
metal/synchronization/context	590.7388888888889 ns	591.2055555555555 ns	1.00
metal/synchronization/stream	328.9070796460177 ns	334.8333333333333 ns	0.98

This comment was automatically generated by workflow using github-action-benchmark.

[VEZY]

I put supports_atomics to true:

KA.supports_atomics(::MetalBackend) = true

because I also opened a PR on Atomix here: JuliaConcurrent/Atomix.jl#79

[maleadt]

Even though these orderings seemed to work, both on CI and on my system, I prefer if we follow the Apple spec:

To make it easier to write code against different MSL versions, I added JuliaGPU/GPUCompiler.jl#880 and applied that here.

[christiangnrd]

FYI there seems to be a typo in Apple's table. memory_order_seq_cst was added in Metal 3.2 for atomic_thread_fence only and that doesn't seem to have changed in Metal 4

[christiangnrd]

This needs a minor rewrite. memory_order_relaxed is still the only supported order for all released versions of macOS, and with the unreleased Metal 4.1, atomics lower to a different form with an extra argument for the memory flag

[maleadt]

Can you test on macOS 27? I only have a virtualized environment.

[christiangnrd]

Tests seem to pass on macOS 27 except 1.10, but I can't help but feel like there's a much less messy way to do this. Considering the atomic operation changes really only apply to macOS 27, I'll split off the fence stuff into its own PR

[maleadt]

Okay, different approach, targeting a single IR form and downgrading during compilation. That's similar to what we do for simdgroup intrinsics. WDYT?

[christiangnrd]

targeting a single IR form and downgrading during compilation

That was one reason for my hesitation about the previous method. AIR 2.9 always lowers to the extra-arg version. I assume the old method is accepted for compatibility but I'd rather do it right since we never know when they'll get around to adding it to the validator.

My other gripe was about the sheer amount of code duplication which this new approach also seems to improve upon. I still think we should merge #875 separately first but at least now the diff is manageable for me to review

[christiangnrd]

Suggested change

	KA.supports_atomics(::MetalBackend) = true
	KA.supports_atomics(::MetalBackend) = metal_support() >= v"4.1"

[christiangnrd]

This or just remove? It doesn't feel like a very useful test

Suggested change

	@test KernelAbstractions.supports_atomics(MetalBackend())
	@test KernelAbstractions.supports_atomics(MetalBackend()) == Metal.metal_support() >= v"4.1"