[SPIRV] Feature Parity Atomics & Shared Array

[duburcqa]

Brief Summary

1. Fix float atomics on shared memory for Metal and Vulkan:
   • allocate shared float arrays as uint, bitcast at load/store
   • only retype arrays targeted by atomics (pre-scan via scan_shared_atomic_allocs)
2. Add f16 shared memory float atomics support:
   • back f16 arrays with u32 (Metal/Vulkan lack 16-bit atomics), with width conversion at load/store/CAS boundaries
3. Add support of shared memory of arbitrary dtype on Metal and Vulkan:
   • flatten nested tensor types (vec3 etc.), handle OpPtrAccessChain for component access
4. Add official support of multiple shared arrays on Metal:
   • enable existing test on Metal
5. Fix shared memory float atomics on GPUs with native float atomic support:
   • handle dest_is_ptr before at_buffer, disable native float atomics for uint-retyped shared arrays

┌─────┬────────────────────────────────────────────────────────────────────┬──────────┐
│  #  │                               Commit                               │  Delta   │
├─────┼────────────────────────────────────────────────────────────────────┼──────────┤
│ 1   │ [Test] Expand unit test coverage.                                  │ +96/-6   │
├─────┼────────────────────────────────────────────────────────────────────┼──────────┤
│ 2   │ [SPIRV] Fix float atomic bugs: f16/f64 min/max, CAS type mismatch. │ +85/-7   │
├─────┼────────────────────────────────────────────────────────────────────┼──────────┤
│ 3   │ [SPIRV] Add Int64Atomics and shared float atomic cap detection.    │ +32/-0   │
├─────┼────────────────────────────────────────────────────────────────────┼──────────┤
│ 4   │ [SPIRV] Add CAS-based float atomic emulation for shared memory.    │ +444/-10 │
├─────┼────────────────────────────────────────────────────────────────────┼──────────┤
│ 5   │ [SPIRV] Add nested tensor type flattening.                         │ +40/-25  │
└─────┴────────────────────────────────────────────────────────────────────┴──────────┘

[duburcqa]

I was assisted by Claude Opus to write this PR. I have read every line added in this PR, and reviewed the lines. I take full responsibility for the lines added and removed in this PR. I won't blame any issue on Claude Opus.

[hughperkins]

elem_dt and elem_type are very confusing. Could we either give more intuitive names, or at least add a comment on what is the difference between them?

[duburcqa]

It should be better now.

[hughperkins]

it's not clear to me from the name what get_quadrants_uint_type does. specifically around nubmer of bits. Could we add a comment to clarify what is happening in this line, specifically around nubmer of bits?

[duburcqa]

I should be better now.

[hughperkins]

this seems very similar to what happens above. Could this be factorized into a helper function?

[duburcqa]

done.

[hughperkins]

same here as for the AllocaStmt: lets have a minimal footprint in this function, that peels off thinknig about uint-backed float shared arrays into a separate function in aseparate file pelase.

[duburcqa]

done.

[hughperkins]

ditto for questoin about helper function.

[duburcqa]

done.

[hughperkins]

can we add a comment about what this if statement is checking for intuitively

[duburcqa]

done.

[hughperkins]

can we add a comment about what this new else if block is checking for

[duburcqa]

done.

[hughperkins]

lots more new code. lets mov to the uint backed shared array file please.

[duburcqa]

Before this PR, shared arrays never reached the float-atomic path; now they can, so all branches guard at_buffer() with dest_is_ptr ternaries.

[hughperkins]

this is just refactorizing right? seems like a nice refactorization, if I've undrestood correctly.

[duburcqa]

done.

[hughperkins]

why are we refactorziing by the way? It's just extra codde for me to read. Is this reafctoziing necessary?

[duburcqa]

updated.

[hughperkins]

I'm not sure I follow. I thought the purpose of changing the backing type to uint was to enable the spirv atomics? Could you give a little more clarification (in the comments) about this point please.

Also, how do we know if we ar edealing with a shared array here?

[duburcqa]

Added some comment on the PR itself to clarify this.

[hughperkins]

total what? total_threads? Why are we dividing by 2.0? Oh, perhaps we are doing some kind of arithmetic progression or similar, and this is the expected_sum of that progression? Could we update the name to make the meaning more intuitive please. By the way, this calculation could be done using ints. Could we make this something that needs actual floats? Like, e.g. multiply each term in the progression by 0.333, which is pretty incompatible with binary representation.

[duburcqa]

fixed.

[hughperkins]

Opus review:

Thoughts

What's good:
• The approach of retyping float shared arrays to uint and bitcasting at the boundaries is a well-known pattern for working around missing float atomics. It's clean and
localized.
• The taint-tracking via shared_float_retyped_ is a reasonable way to propagate the retyping through the IR without modifying the IR itself.
• Test coverage looks solid — the new test_shared_array_float_atomics is parametrized over all four atomic ops.

Things I'd flag:

1. The shared_float_retyped_ set is a std::unordered_set<const Stmt*> tracking raw pointers. This is fine as long as stmts outlive the codegen pass, but it's worth
   confirming there's no stmt lifetime issue. Given this is a visitor pass over existing IR, it should be safe.
2. The retyping applies to all real types on shared memory, not just f32. If someone uses f16 or f64 shared arrays, the bitcast to the corresponding uint type should work,
   but there's no explicit test for those. The get_quadrants_uint_type utility presumably handles the width mapping, but worth verifying for f16 especially.
3. The early dest_is_ptr branch in AtomicOpStmt unconditionally uses dest_val for shared memory, which means it never goes through at_buffer. This is correct (shared
   memory isn't in a buffer), but the original code had some interleaving of the dest_is_ptr check inside each float type branch. The refactor is cleaner, but it's a behavior
   change — previously a shared f32 array on a GPU that does report spirv_has_atomic_float_add would have used dest_val directly (correct), and now it still does (correct),
   but the native-atomics flag is also explicitly suppressed. The comment explains why (shared arrays are uint-retyped so native float atomics can't work), which is good.
4. The test_shared_array_float_atomics test initializes sharr[tid] = qd.f32(tid) then atomically accumulates into sharr[0]. The initial value of sharr[0] is 0.0 (from
   tid=0), then all threads do atomic_op(sharr[0], tid). For sub, the expected result is -total which assumes the initial value 0.0 stays as-is and only the atomic subtracts
   are applied. That's correct since atomic_sub(sharr[0], 0) is a no-op for tid=0.
5. Minor: test_shared_array_matrix dropped debug=True — intentional? Could be that debug mode was causing issues on Vulkan/Metal or just wasn't needed.

Overall, this is a well-structured branch. The Metal/Vulkan shared memory work is the kind of backend plumbing that's easy to get wrong, but the approach here is principled and
the test coverage backs it up.

[hughperkins]

from the AI review, pelase could we address at least:
2. The retyping applies to all real types on shared memory, not just f32. If someone uses f16 or f64 shared arrays, the bitcast to the corresponding uint type should work,
but there's no explicit test for those. The get_quadrants_uint_type utility presumably handles the width mapping, but worth verifying for f16 especially.
5. Minor: test_shared_array_matrix dropped debug=True — intentional? Could be that debug mode was causing issues on Vulkan/Metal or just wasn't needed.

[hughperkins]

(so AI and myself concur about the ambiguity over what get_quadrants_uint_type does)

[duburcqa]

I was assisted by Claude Opus to write this PR. I have read and review every changes in this PR. I take full responsibility for the lines added and removed in this PR. I won't blame any issue on Claude Opus.

[duburcqa]

I have reviewed this PR, and approve it.

[hughperkins]

Approving, as a codeowner, on the basis that Alexis has reviewed and approved this PR."