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 0% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 85.93%. Comparing base (998b6cb) to head (7837754).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/MetalKernels.jl	0.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #870   +/-   ##
=======================================
  Coverage   85.93%   85.93%           
=======================================
  Files          76       76           
  Lines        5246     5246           
=======================================
  Hits         4508     4508           
  Misses        738      738           

☔ 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: 7837754	Previous: 998b6cb	Ratio
array/accumulate/Float32/1d	404709 ns	406500 ns	1.00
array/accumulate/Float32/dims=1	374125 ns	376292 ns	0.99
array/accumulate/Float32/dims=1L	8788583 ns	8806959 ns	1.00
array/accumulate/Float32/dims=2	444042 ns	443333 ns	1.00
array/accumulate/Float32/dims=2L	2181875 ns	2641417 ns	0.83
array/accumulate/Int64/1d	851875 ns	862708 ns	0.99
array/accumulate/Int64/dims=1	939208 ns	940333 ns	1.00
array/accumulate/Int64/dims=1L	9504000 ns	9511750 ns	1.00
array/accumulate/Int64/dims=2	1242000 ns	1247291 ns	1.00
array/accumulate/Int64/dims=2L	6532417 ns	6530125 ns	1.00
array/broadcast	221792 ns	230750 ns	0.96
array/construct	2375 ns	2375 ns	1
array/permutedims/2d	455292 ns	405292 ns	1.12
array/permutedims/3d	1012542 ns	1017958 ns	0.99
array/permutedims/4d	1156167 ns	1200542 ns	0.96
array/private/copy	235334 ns	238625 ns	0.99
array/private/copyto!/cpu_to_gpu	201333 ns	196792 ns	1.02
array/private/copyto!/gpu_to_cpu	195208 ns	198500 ns	0.98
array/private/copyto!/gpu_to_gpu	189458 ns	197708 ns	0.96
array/private/iteration/findall/bool	1082000 ns	937417 ns	1.15
array/private/iteration/findall/int	1237250 ns	1234500 ns	1.00
array/private/iteration/findfirst/bool	1111000 ns	1105500 ns	1.00
array/private/iteration/findfirst/int	1127583 ns	1141833 ns	0.99
array/private/iteration/findmin/1d	1233000 ns	1229875 ns	1.00
array/private/iteration/findmin/2d	1075542 ns	1065458 ns	1.01
array/private/iteration/logical	1764334 ns	1689875 ns	1.04
array/private/iteration/scalar	1079542 ns	1236041 ns	0.87
array/random/rand/Float32	389750 ns	433084 ns	0.90
array/random/rand/Int64	521584 ns	464000 ns	1.12
array/random/rand!/Float32	384125 ns	382333 ns	1.00
array/random/rand!/Int64	327417 ns	413458 ns	0.79
array/random/randn/Float32	367792 ns	388208 ns	0.95
array/random/randn!/Float32	347250 ns	347667 ns	1.00
array/reductions/mapreduce/Float32/1d	381250 ns	405375 ns	0.94
array/reductions/mapreduce/Float32/dims=1	271709 ns	340084 ns	0.80
array/reductions/mapreduce/Float32/dims=1L	619000 ns	627500 ns	0.99
array/reductions/mapreduce/Float32/dims=2	339500 ns	335375 ns	1.01
array/reductions/mapreduce/Float32/dims=2L	1012292 ns	810334 ns	1.25
array/reductions/mapreduce/Int64/1d	601542 ns	603333 ns	1.00
array/reductions/mapreduce/Int64/dims=1	444833 ns	441291 ns	1.01
array/reductions/mapreduce/Int64/dims=1L	1026584 ns	1036208 ns	0.99
array/reductions/mapreduce/Int64/dims=2	655917 ns	773334 ns	0.85
array/reductions/mapreduce/Int64/dims=2L	2183791 ns	2182583 ns	1.00
array/reductions/reduce/Float32/1d	402125 ns	416334 ns	0.97
array/reductions/reduce/Float32/dims=1	338958 ns	346875 ns	0.98
array/reductions/reduce/Float32/dims=1L	628125 ns	636500 ns	0.99
array/reductions/reduce/Float32/dims=2	237709 ns	231875 ns	1.03
array/reductions/reduce/Float32/dims=2L	450750 ns	449541 ns	1.00
array/reductions/reduce/Int64/1d	532625 ns	582667 ns	0.91
array/reductions/reduce/Int64/dims=1	547959 ns	634417 ns	0.86
array/reductions/reduce/Int64/dims=1L	1019958 ns	1027458 ns	0.99
array/reductions/reduce/Int64/dims=2	239458 ns	243584 ns	0.98
array/reductions/reduce/Int64/dims=2L	642000 ns	649584 ns	0.99
array/shared/copy	135708 ns	136208 ns	1.00
array/shared/copyto!/cpu_to_gpu	37292 ns	38583 ns	0.97
array/shared/copyto!/gpu_to_cpu	38000 ns	38459 ns	0.99
array/shared/copyto!/gpu_to_gpu	38708 ns	39084 ns	0.99
array/shared/iteration/findall/bool	1097167 ns	1062125 ns	1.03
array/shared/iteration/findall/int	1245250 ns	1222417 ns	1.02
array/shared/iteration/findfirst/bool	955125 ns	654709 ns	1.46
array/shared/iteration/findfirst/int	948000 ns	670708 ns	1.41
array/shared/iteration/findmin/1d	1101750 ns	1098167 ns	1.00
array/shared/iteration/findmin/2d	1069792 ns	1070833 ns	1.00
array/shared/iteration/logical	1612792 ns	1565834 ns	1.03
array/shared/iteration/scalar	4113.142857142857 ns	4226.142857142857 ns	0.97
array/sorting/1d	1819792 ns	1786375 ns	1.02
array/sorting/2d	8276458 ns	8281333 ns	1.00
integration/byval/reference	1101459 ns	1113750 ns	0.99
integration/byval/slices=1	1102625 ns	1113459 ns	0.99
integration/byval/slices=2	2022375 ns	2029333 ns	1.00
integration/byval/slices=3	6790791 ns	6488667 ns	1.05
integration/metaldevrt	382417 ns	387792 ns	0.99
kernel/indexing	198833 ns	200208 ns	0.99
kernel/indexing_checked	382875 ns	376125 ns	1.02
kernel/launch	1895.9 ns	1891.6 ns	1.00
kernel/rand	387833 ns	384666 ns	1.01
latency/import	2041220500 ns	2035107500 ns	1.00
latency/precompile	39253324666 ns	38873630084 ns	1.01
latency/ttfp	2374825334 ns	2371256000 ns	1.00
metal/synchronization/context	602.1843575418994 ns	601.6685714285715 ns	1.00
metal/synchronization/stream	338.6834862385321 ns	340.1826484018265 ns	1.00

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]

Is 5598ff2 still here on purpose?

[maleadt]

Is 5598ff2 still here on purpose?

No, that's something I reverted on main, and apparently snuck in here again.

[christiangnrd]

I lost the more detailed comment I had typed up, but I think I confused the robot when merging #875 and force-pushing to update this PR and it's somehow back to the messy version that made me comment #870 (comment)

This PR shouldn't need any code or test changes to atomic fences unless we decide we want the atomic check functions in synchronization. But I had removed check_atomic_thread_fence_order in #875 since it was only used once and didn't clutter up the function too much so the other ones could probably be moved to atomics.jl

[maleadt]

Github is being confused here... I pushed a clean-up commit, but it doesn't appear. I'll re-open on a first-party branch.

[maleadt]

#889

[VEZY]

Thanks @maleadt and @christiangnrd for pushing this forward ! And sorry for not having participated more in the discussion; that was way out of my capabilities.