groupshared improvement discussion

[jeremyong]

This is a general issue I wanted to file to encourage open discussion about future improvements to groupshared LDS/shmem usage and allocation.

Problem Discussion

This is unlikely exhaustive, but here are the problems that are likely to crop up for users of groupshared memory (I will use groupshared, LDS, and shmem interchangeably).

• LDS allocation usually scales linearly with the number of waves in a thread group, however, because WaveGetLaneCount is not a compile-time constant, the wave size cannot readily be used in a groupshared declaration. The current workaround for this is to either make (bad) assumptions about the wave size, or produce multiple specializations of each shader and dispatching the correct one with a matching WaveSize at runtime. Neither option is ideal, with the former resulting in brittle hardware-specific code, and the latter resulting in build and runtime complexity.
• Currently, groupshared data must be declared in the global declaration context. This impedes code composability -- for example, if we wanted to include a header to use a function defined in that header, we may hurt occupancy by inadvertently dragging along groupshared declarations. This is a real footgun in larger (and sometimes smaller) codebases, and it is difficult to detect when it occurs (or at least, it takes some work to understand why occupancy is lower than expected).
• Along the same vein of code composability, LDS memory usage must interact directly with the groupshared variable as declared, due to the lack of a user-accessible ref function parameter qualifier. A function that operates over some input data and exports output data cannot always rely on "copy-in" and "copy-out" semantics, because the in and out semantics do not permit usage of the various Interlocked* functions (which internally are modeled using ref-qualified parameters)

Possible Solutions

If possible, I think there are a few things that would immediately improve quality-of-life for compute shader authors. These suggestions are written from the perspective of an ISV (the shader writer), and it's understood that other solutions may end up being preferable due to practicality, performance, ease-of-implementation, or all of the above from the perspective of a hardware vendor or DXC compiler implementer.

1. Permit the use of WaveGetLaneCount in the declaration type for groupshared storage.
2. Permit local variables to be declared as groupshared.
3. Add a ref keyword that would permit use of Interlocked* intrinsics for ref-qualified parameters

The first item would allow developers to conceptually treat WaveGetLaneCount as a constexpr function, whose value is realized only when a PSO is actually created at runtime. This has implications beyond LDS allocation, but would be a very useful tool in the toolbox for other use cases.

For the second item, because functions are still fully unrolled currently, the total storage needed per-thread-group for a given compute shader should still be statically known, although DXIL may require modifications to properly alias types allocated from the virtual shmem pool. The idea here is that a static analysis pass would determine the amount of LDS memory needed in the "middle swell" of the program, accounting for all possible branches taken where groupshared variables are declared.

The counterargument to the second item is that statically knowing how much LDS is used precludes future HLSL code in a world where function calling is possible. At this point, one option would be to permit functions to allocate LDS (similar to alloca) using the same semantics as locally declared groupshared variables. The driver would need to be able to suspend thread groups if LDS isn't available, or possibly demote allocated LDS to slower vram (possibly from a fixed size pool of reserved memory).

The last item addresses the ability to perform operations on memory in LDS, regardless of where or how that LDS memory was allocated.

All that said, my main goal is to encourage discussion, and not attempt to be overly prescriptive about the solutions. I think starting from a well-defined problem statement is likely step one.

[jeremyong]

@jongreenberg-ea mentioned a related issue filed against DXC w.r.t. groupshared usage and WorkGraphs, with similar suggestions: https://github.com/microsoft/DirectXShaderCompiler/issues/5618

I think this raises a good point, code composability issues are exacerbated in a world with work graphs.

[devshgraphicsprogramming]

Along the same vein of code composability, LDS memory usage must interact directly with the groupshared variable as declared, due to the lack of a user-accessible ref function parameter qualifier. A function that operates over some input data and exports output data cannot always rely on "copy-in" and "copy-out" semantics, because the in and out semantics do not permit usage of the various Interlocked* functions (which internally are modeled using ref-qualified parameters)

the way we solve that is via an Accessor pattern, basically every algorithm we have takes a stateful-lambda to access whatever memory it needs.

There's a funny bug in codegen in DXC which causes inout to work exactly like a reference, we've tried it in SPIR-V for atomics, an issue was opened about it at one point.

[devshgraphicsprogramming]

Permit local variables to be declared as groupshared.

in spir-v land there's a capability/extension SPV_KHR_workgroup_memory_explicit_layout which allows you to alias multiple shared variable declarations (or just generally set the offset).

[jeremyong]

the way we solve that is via an Accessor pattern, basically every algorithm we have takes a stateful-lambda to access whatever memory it needs.

I don't think this relevant to this proposal, which I would rather keep to a minimal surface area if possible (lambdas feel like an entirely separate matter, which would require some form of reference semantics anyways).

in spir-v land there's a capability/extension SPV_KHR_workgroup_memory_explicit_layout which allows you to alias multiple shared variable declarations (or just generally set the offset).

This is probably more relevant to the proposal related to union support.

[devshgraphicsprogramming]

I don't think this relevant to this proposal, which I would rather keep to a minimal surface area if possible (lambdas feel like an entirely separate matter, which would require some form of reference semantics anyways).

Functor structs, a.k.a. the precursors to C++11 lambdas, work just fine for us.

[jeremyong]

Please show a snippet of code where in the function body, you are able to perform an interlocked compare exchange on data passed into the function that may either be LDS memory or memory in a structured buffer or some other buffer type. I can imagine code contortions one might need to make this possible, but restructuring all code to become templates or perform interlocked operations through proxy objects is pretty undesirable in a general case.

[devshgraphicsprogramming]

Voila: https://godbolt.org/z/9qso3Y7eW

see it emits:

        %107 = OpAtomicIAdd %uint %particleCount %uint_1 %uint_0 %uint_1

its a codegen bug, but it works... for now

[jeremyong]

That isn't what I asked for unfortunately -- I thought it was pretty obvious I was replying to your somewhat snarky "work just fine for us" comment. As for the codegen bug, obviously this should not be relied on, and is better discussed as an actual DXC issue?

[devshgraphicsprogramming]

That isn't what I asked for unfortunately -- I thought it was pretty obvious I was replying to your somewhat snarky "work just fine for us" comment.

Well its not my fault if you can't ask the question plainly and clearly.

As for the codegen bug, obviously this should not be relied on, and is better discussed as an actual DXC issue?

There's already a DXC issue, but unlikely to get "fixed" until references are in.

[devshgraphicsprogramming]

Also look, same function takes either shared or RWStructuredBuffer

https://godbolt.org/z/1j3hczKcT

[jeremyong]

Well its not my fault if you can't ask the question plainly and clearly.

If you scroll up, I am replying directly to this claim:

Functor structs, a.k.a. the precursors to C++11 lambdas, work just fine for us.

I'm not particularly interested in engaging in further discussion with you but with respect to:

Also look, same function takes either shared or RWStructuredBuffer

https://godbolt.org/z/1j3hczKcT

Yes, obviously you can modify the snippet to exploit the same codegen bug, but this is not the same as demonstrating a working snippet working within intended language semantics using your functor struct method.

[devshgraphicsprogramming]

working snippet working within intended language semantics using your functor struct method

I have a whole library of them.

Write exactly what you want and I can give you an example.

[llvm-beanz]

There's a funny bug in codegen in DXC which causes inout to work exactly like a reference, we've tried it in SPIR-V for atomics, an issue was opened about it at one point.

DO NOT rely on this. If you do your code will break in the near future. I have a draft PR to fix this issue in the SPIR-V code generator and fix a bunch of related issues in the DXIL code generator (microsoft/DirectXShaderCompiler#5249). The PR is currently blocked on resolving some related issues with matrix orientation annotations, but we are hoping to get both of those fixes into a DXC release this fall.