User defined WorkGroup sizes in Vulkan Compute shaders

Hi everyone,

I’ve been playing with Vulkan for a while and so far so good, but I recently stumble upon a problem that I am unable to solve.

I’m experimenting with Vulkan Compute Shaders and I managed to run a simple vector add. Then I wanted to enhance it by having the possibility to more or less dynamically setting the work group sizes.
In plain core OpenGL that is easy using the “GL_ARB_compute_variable_group_size” extension. In Vulkan I know I have to use specialization constants, as reported here.

So I modified both host and shader code, in order to do that and I did this (glsl shader compiled into SPIR-V with LUNARG glslValidator)

#version 450

layout (constant_id = 0) const int TX = 512; //we push 1024 from code using spec. constants.
layout (constant_id = 1) const int TY = 1;
layout (constant_id = 2) const int TZ = 1; 

layout( local_size_x = TX, local_size_y = TY, local_size_z = TZ ) in;
//rest of the shader

The code compiles into SPIR-V and survives pipeline creation without crashes.
TX,TY and TZ values are correctly passed into the shader as I am able to assign those values to an output buffer and read those values from host.
However, seems like TX, TY and TZ did not manage to change the workgroup size, as I got one workitem per workgroup (some kind of default value.)

So I thought I was missing something. After doing some research I saw this.
Looks like I can assign constants to built-in variables. So I changed my code and now looks like this:

#version 450

//commented out previous definitions
/*layout (constant_id = 0) const int TX = 512; 
layout (constant_id = 1) const int TY = 1;
layout (constant_id = 2) const int TZ = 1; */

 layout(local_size_x_id = 0) in;               
 layout(local_size_y_id = 1) in;
 layout(local_size_z_id = 2) in;

//rest of the shader

This code translates into SPIR-V with no probs, but crashes during pipeline creation. By activating the debug layer I got this:

Internal error: assembly compile error for compute shader at offset 82:
-- error message --
line 4, column 12:  error: GROUP_SIZE declaration must be followed by integers
-- internal assembly text --
OPTION NV_shader_storage_buffer;
OPTION NV_bindless_texture;
# cgc version 3.4.0001, build date Sep 16 2016
# command line args:
#vendor NVIDIA Corporation
#version COP Build Date Sep 16 2016
#profile gp5cp
#program main
#semantic Output.output_data : SBO_BUFFER[2]
#semantic Input0.input_data0 : SBO_BUFFER[0]
#semantic Input1.input_data1 : SBO_BUFFER[1]
#var uint3 gl_GlobalInvocationID : $vin.GBLID : GBLID[3] : -1 : 1
#var uint3 gl_LocalInvocationID : $vin.LCLID : LCLID[2] : -1 : 1
#var int output_data.elements[0] :  : sbo_buffer[2][0] : -1 : 1
#var int input_data0.elements[0] :  : sbo_buffer[0][0] : -1 : 0
#var int input_data1.elements[0] :  : sbo_buffer[1][0] : -1 : 0
STORAGE sbo_buf0[] = {[0] };
STORAGE sbo_buf1[] = {[1] };
STORAGE sbo_buf2[] = {[2] };
MUL.S R0.x, invocation.globalid, {4, 0, 0, 0};
MOV.U R0.x, R0;
STB.S32 invocation.localid, sbo_buf2[R0.x];
# 3 instructions, 1 R-regs

vkCreateComputePipelines: returned VK_ERROR_INVALID_SHADER_NV, indicating that one or more shaders failed to compile or link

So, how can I fix that?


That’s a bug in the current NVIDIA drivers, see this thread for details.

Until this has been fixed in the current drivers, start with constant_id 1 instead of zero:

 layout(local_size_x_id = 1) in;               
 layout(local_size_y_id = 2) in;
 layout(local_size_z_id = 3) in;

And offset accordingly at pipeline creation time.

Yep. That was it.
Thank you very much. BTW, thanks for your really useful and complete tutorials :slight_smile: