[SPIR-V] Multiple variables with TaskPayloadWorkgroupEXT in task/amplification shaders

[godlikepanos]

Description
When the task payload is a struct the SPIR-V backend generates multiple variables with TaskPayloadWorkgroupEXT (one variable per member of the struct actually). The spec states that only a single TaskPayloadWorkgroupEXT variable can exist.

Each OpEntryPoint with the MeshEXT or TaskEXT Execution Models can have at most one global OpVariable of storage class TaskPayloadWorkgroupEXT.

Steps to Reproduce
Here is a simple task shader:

struct Payload
{
    uint index1;
    uint index2;
};

groupshared Payload s_payload;

[numthreads(16, 1, 1)]
void main(uint svGroupId : SV_GROUPID) 
{
    s_payload.index1 = svGroupId;
    s_payload.index2 = svGroupId * 100;

    DispatchMesh(svGroupId, 1, 1, s_payload);
}

Compiling with -T as_6_7 -E main -spirv -fspv-target-env=vulkan1.1spirv1.4 it gives the following SPIR-V:

               OpCapability MeshShadingEXT
               OpExtension "SPV_EXT_mesh_shader"
               OpMemoryModel Logical GLSL450
               OpEntryPoint TaskEXT %main "main" %gl_WorkGroupID %out_var_index1 %out_var_index2 %s_payload
               OpExecutionMode %main LocalSize 16 1 1
               OpMemberName %Payload 0 "index1"
               OpMemberName %Payload 1 "index2"
               OpModuleProcessed "dxc-commit-hash: 05094636"
               OpModuleProcessed "dxc-cl-option: /app/example.hlsl -E main -T as_6_7 -Zi -Qembed_debug -spirv -fspv-target-env=vulkan1.1spirv1.4"
               OpDecorate %gl_WorkGroupID BuiltIn WorkgroupId
        %int = OpTypeInt 32 1
      %int_0 = OpConstant %int 0
       %uint = OpTypeInt 32 0
   %uint_100 = OpConstant %uint 100
      %int_1 = OpConstant %int 1
     %uint_1 = OpConstant %uint 1
    %Payload = OpTypeStruct %uint %uint
%_ptr_Workgroup_Payload = OpTypePointer Workgroup %Payload
     %v3uint = OpTypeVector %uint 3
%_ptr_Input_v3uint = OpTypePointer Input %v3uint
%_ptr_TaskPayloadWorkgroupEXT_uint = OpTypePointer TaskPayloadWorkgroupEXT %uint
       %void = OpTypeVoid
         %20 = OpTypeFunction %void
%_ptr_Function_uint = OpTypePointer Function %uint
         %22 = OpTypeFunction %void %_ptr_Function_uint
%_ptr_Workgroup_uint = OpTypePointer Workgroup %uint
     %uint_2 = OpConstant %uint 2
   %uint_264 = OpConstant %uint 264
  %s_payload = OpVariable %_ptr_Workgroup_Payload Workgroup
%gl_WorkGroupID = OpVariable %_ptr_Input_v3uint Input
%out_var_index1 = OpVariable %_ptr_TaskPayloadWorkgroupEXT_uint TaskPayloadWorkgroupEXT
%out_var_index2 = OpVariable %_ptr_TaskPayloadWorkgroupEXT_uint TaskPayloadWorkgroupEXT
       %main = OpFunction %void None %20
         %26 = OpLabel
%param_var_svGroupId = OpVariable %_ptr_Function_uint Function
         %27 = OpLoad %v3uint %gl_WorkGroupID
         %28 = OpCompositeExtract %uint %27 0
               OpStore %param_var_svGroupId %28
         %29 = OpAccessChain %_ptr_Workgroup_uint %s_payload %int_0
               OpStore %29 %28
         %30 = OpIMul %uint %28 %uint_100
         %31 = OpAccessChain %_ptr_Workgroup_uint %s_payload %int_1
               OpStore %31 %30
               OpControlBarrier %uint_2 %uint_2 %uint_264
         %32 = OpLoad %Payload %s_payload
         %33 = OpCompositeExtract %uint %32 0
               OpStore %out_var_index1 %33
         %34 = OpCompositeExtract %uint %32 1
               OpStore %out_var_index2 %34
               OpEmitMeshTasksEXT %28 %uint_1 %uint_1
               OpFunctionEnd

You can play with the shader in compiler explorer: https://godbolt.org/z/hPEnx68E3

As you can see there are two TaskPayloadWorkgroupEXT variables (out_var_index1 and out_var_index2) which is wrong.

Note no1: It seems the code generated is not generally optimal. The s_payload is decorated as groupshared in SPIR-V and at DispatchMesh() the compiler will emit copies from s_payload to a new set of TaskPayloadWorkgroupEXT variables. It would have been preferable if s_payload was decorated as TaskPayloadWorkgroupEXT and the extra copies are removed.

Note no2: Similar problem exists in the mesh shaders where DXC will take the struct of a payload and create multiple global TaskPayloadWorkgroupEXT variables.

Actual Behavior
Ideally the s_payload is not decorated as groupshared but as TaskPayloadWorkgroupEXT and the copies from groupshared to TaskPayloadWorkgroupEXT are removed.

Environment

• DXC version: trunk & 1.7.2207
• Host Operating System: Linux & Windows

[godlikepanos]

cc @tywuAMD

[sudonatalie]

Thanks for identifying this @godlikepanos . @cassiebeckley Can you look into this?

Relevant spec: https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/EXT/SPV_EXT_mesh_shader.html

And it looks like adding this to spirv-val is already being tracked here: KhronosGroup/SPIRV-Tools#4919

[spencer-lunarg]

this seems like a related problem to KhronosGroup/SPIRV-Tools#4969 (comment) where there are more than 1 RayTracing StorageClass Variable

[godlikepanos]

Hi. Is there a chance this can be prioritized? On top of the problems it causes on RenderDoc (the mesh view in RenderDoc doesn't work when there are multiple TaskPayloadWorkgroupEXT) it also causes issues on AMD. AMD confirmed that their driver's compiler fails to link the task and mesh shaders because of the issue

[spencer-lunarg]

for notes, I think this was added already for Ray Tracing shaders
#6268

[sudonatalie]

@cassiebeckley feel free to re-assign this to me if you don't have bandwidth for it right now.