[SPIR-V] Allow thread group size to be specified with specialization constants

[SaschaWillems]

In glsl it's possible to set the thread group / work group size for compute shaders via specialization constants. This allows an application to use different sizes depending e.g. on the hardware limits, without the need for different shaders.

The GL_ARB_gl_spirv extension added explicit qualifiers for this:

"The built-in constant vector gl_WorkGroupSize can be specialized using
the local_size_{xyz}_id qualifiers, to individually give the components
an id. For example:

     layout(local_size_x_id = 18, local_size_z_id = 19) in;

There doesn't seem to be an equivalent of this for HLSL to SPIR-V right now.

[ehsannas]

Thanks for reporting @SaschaWillems

[shangjiaxuan]

Related: #4032 #2191
Seems like specialization constant logic in dxc doesn't propagate constants that can be evaluated early-on?
(earlier shader models seems to recognize #2191)

[shangjiaxuan]

One suggestion would be detecting [[xxx]] attributes in [numthreads(4, 4, 1)], and use the actual numbers as the default number (no current code will be affected). Then the following may be ok:

    [numthreads([[vk::constant_id(1)]] 4, [[vk::constant_id(2)]] 4, [[vk::constant_id(3)]] 1)]

Seems a bit hacky though...

[shangjiaxuan]

Sample assembly from glslc compilation of following code:

layout(local_size_x_id = 1, local_size_y_id = 2, local_size_z = 3) in;
void main(void)
{
}

Assembly:

; SPIR-V
; Version: 1.5
; Generator: Google Shaderc over Glslang; 10
; Bound: 12
; Schema: 0
               OpCapability Shader
          %1 = OpExtInstImport "GLSL.std.450"
               OpMemoryModel Logical GLSL450
               OpEntryPoint GLCompute %main "main"
               OpExecutionMode %main LocalSize 1 1 3
               OpSource GLSL 450
               OpSourceExtension "GL_GOOGLE_cpp_style_line_directive"
               OpSourceExtension "GL_GOOGLE_include_directive"
               OpName %main "main"
               OpDecorate %7 SpecId 1
               OpDecorate %8 SpecId 2
               OpDecorate %gl_WorkGroupSize BuiltIn WorkgroupSize
       %void = OpTypeVoid
          %3 = OpTypeFunction %void
       %uint = OpTypeInt 32 0
          %7 = OpSpecConstant %uint 1
          %8 = OpSpecConstant %uint 1
     %uint_3 = OpConstant %uint 3
     %v3uint = OpTypeVector %uint 3
%gl_WorkGroupSize = OpSpecConstantComposite %v3uint %7 %8 %uint_3
       %main = OpFunction %void None %3
          %5 = OpLabel
               OpReturn
               OpFunctionEnd

Don't know if OpSpecConstantComposite is currently implemented in dxc, also it seems to reference OpDecorate %gl_WorkGroupSize BuiltIn WorkgroupSize, things may be easier with a post-processor in this case.

[SaschaWillems]

Any updates on this?

[s-perron]

My preferred syntax would be

[[vk::constant_id(13)]]
const int X = 10;

[numthreads(X,1,1)]
void CSMain()
{
}

We can already define spec constants, so I would want to reuse the same syntax, just allow them in the numthreads attribute.

@llvm-beanz Is this something that would need to go through hlsl specs?

[llvm-beanz]

@s-perron, no I think this is just a bug. We should allow any compile-time constant in that attribute, but I think the way it is implemented we don’t correctly support that. We have a few related issues:

• Assert when a static const uint is used with [numthreads] #2191
• fxc.exe vs dxc.exe: "static const int" use as compile time constant #2188

[SaschaWillems]

Any updates on this? I'm working on translating some complex glsl shaders to hlsl, and this is somewhat of a blocker.

[s-perron]

I have a draft PR that implements this: #7084. However, this cannot go into DXC yet. We would need an HLSL spec update. We could add a new attribute or update numthreads. If we update numthreads, the implementation of it will have to completely change. That would need to be discussed.

Also in the compute derivatives feature, the quads are determined by the values in the numthreads attribute. This assumes the values are all known at compile time, so we can add the correct SPIR-V capability/executionmode. We will need to update the HLSL in this respect for the new attribute.

There are solutions to these, but they have to be handled in the HLSL spec first.

The draft PR also includes a partial refactoring of the OpExecutionModeId instruction in the SPIR-V backend. The current implementation is not expressive enough for what we want to do. This refactoring should be pulled out into its own PR.

• Add a new attribute vk::LocalSizeId or update the spec for numthreads to allow spec constant.
• Update the compute derivatives spec to account for an unknown dimension size.
• Refactor the DXC implementation of OpExecutionModeId.

[s-perron]

I do not have any timeline on when we can get to this.