[SPIR-V] Add support for 64 bit atomics

[natevm]

Is your feature request related to a problem? Please describe.
When I attempt to use something like "InterlockedAdd64", dxc returns "error GCC929F7B: intrinsic 'InterlockedAdd64' method unimplemented" when targeting SPIR-V. (minimal reproducer: https://godbolt.org/z/oe1bG7T3d)

Describe the solution you'd like
We should instead request a -fspv-extension=SPV_KHR_shader_atomic_int64 in the compilation flags, then allow functions like InterlockedAdd64 to emit opcodes like OpAtomicIAdd with 64 bit types for inputs/output/comparands.

Describe alternatives you've considered
Perhaps it could theoretically be possible to emit custom SPIR-V, but OpAtomicIAdd requires a complex pointer type for the RWByteAddressBuffer that I'm trying to use, and I can't figure out how to actually write this with GL_EXT_spirv_intrinsics...

Additional context
64 bit atomic adds are necessary for certain operations when implementing a single pass parallel prefix sum, as described here: https://research.nvidia.com/sites/default/files/pubs/2016-03_Single-pass-Parallel-Prefix/nvr-2016-002.pdf. These prefix sums in turn are fundamental building blocks to many parallel algorithms relevant in ray tracing (eg, PLOC tree construction). Without 64 bit atomics, a single pass prefix sum is still possible, but requires some expensive memory fences.

[devshgraphicsprogramming]

and I can't figure out how to actually write this with GL_EXT_spirv_intrinsics...

Like this:
https://github.com/Devsh-Graphics-Programming/Nabla/blob/6b33b9730c65bd23569ffe977abcdbb426686d0b/include/nbl/builtin/hlsl/spirv_intrinsics/core.hlsl#L23-L30

In your case you'd need to make specialization for uint64_t and add an extra [[vk::extension("")]] atribute (or even capability) to the function specialization.

Then you can add one more layer of wrapping (to hide the memory semantics and scope) if you like playing with fire and relying on inout T being a T restrict & (which will break/no-longer-be-possible with the merging of #5249)

64 bit atomic adds are necessary for certain operations when implementing a single pass parallel prefix sum, as described here: https://research.nvidia.com/sites/default/files/pubs/2016-03_Single-pass-Parallel-Prefix/nvr-2016-002.pdf.

Damn, and I thought I independently invented that one 2 years ago.

[natevm]

Unfortunately the inline SPIR-V in the link above only works for RWStorageBuffers, and does not work for the RWByteAddressBuffers like I'm using.

@devshgraphicsprogramming Do you have any ideas on how to get your inline assembly to work with RWByteAddressBuffers?

(See the godbolt example I shared above with the reproducer using a RWByteAddressBuffer)

In retrospect, the single pass approach could also use 32 bit atomics, just limited to 30 bit unsigned integers.

[s-perron]

It is on our list of missing features to implement. No timeline on when we will get to it. I'll link this issue to our internal work list, which will increase its priority a little.

[natevm]

It is on our list of missing features to implement. No timeline on when we will get to it. I'll link this issue to our internal work list, which will increase its priority a little.

Awesome, thanks! Another motivating example for 64 bit atomics would probably be to support 64 bit keys in this "OneSweep" sorting algorithm which uses a similar atomic pattern for the lookback : https://arxiv.org/abs/2206.01784

This seems to be what's currently "state of the art" in GPU sorting these days. Today, we can implement a 2^30 bit key sort, but for tree construction algorithms, having more precise keys translates to more precise space filling curve codes, or being able to more efficiently sort key-value pairs.

[AntonTokarev]

Doesn't the InterlockedAdd on StructuredBuffer do the same already? https://godbolt.org/z/ed9nsqT93

[natevm]

On mobile atm, unfortunately godbolt doesn’t seem to work there. I was pretty sure that InterlockedAdd casts the input parameters to 32 bit integers.

InterlockedAdd64 would be the function to test.

@devshgraphicsprogramming ’s inline SPIR-V solution would work for 64 bit integers in structured buffers, but doesn’t work if you’re constrained to using RWByteAddressBuffers.

[devshgraphicsprogramming]

On mobile atm, unfortunately godbolt doesn’t seem to work there. I was pretty sure that InterlockedAdd casts the input parameters to 32 bit integers.

InterlockedAdd64 would be the function to test.

@devshgraphicsprogramming ’s inline SPIR-V solution would work for 64 bit integers in structured buffers, but doesn’t work if you’re constrained to using RWByteAddressBuffers.

Once proposal 0011 is in (AFAIK its being implemented right now), atomics might finally work with BDA so there'd be no reason to use ByteAddressBuffers on Vulkan.

[s-perron]

Unfortunately, it will not be possible to implement The Interlocked*64 member functions for RWByteAddressBuffers. The problem is that SPIR-V cannot express it in logical addressing mode.

The RWByteAddressBuffer is suppose to be unaware of its type. The type for the load or store is determined during the actual load or store operation. That does not directly translate to SPIR-V, which requires the declaration of the variable to have a type and all loads and stores have to match that type. In general, we get around this by declaring the variable to be an array of uint32. A load as a 64-bit int would be implemented as two 32-bit loads and have the two value merged together using bit operations.

When it comes to atomic operations, we cannot split the operations to act on 32-bit values because that would break atomicity for the the 64-bit memory location. Sorry this cannot be implemented in DXC until there is a change to SPIR-V. We will keep this in our list of missing features, and implement if SPIR-V get a change that will allow it.

FYI: @alan-baker