[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:

• https://github.com/microsoft/DirectXShaderCompiler/issues/2191
• https://github.com/microsoft/DirectXShaderCompiler/issues/2188