Add WorkgroupCount function (#7734)

davli-nv · web-flow · commit d39590228241 · 2025-07-12T00:46:22.000Z
Fixes #7733 Copy gl_NumWorkGroups into hlsl.meta.slang as WorkgroupCount function so that it can be used for GLSL and SPIR-V targets without GLSL syntax. Also change WorkgroupSize function to allow use with mesh shading capability. Update pipeline/rasterization/mesh/task-simple.slang to test it in task and mesh stages.
diff --git a/source/slang/glsl.meta.slang b/source/slang/glsl.meta.slang
@@ -141,25 +141,18 @@ public property int gl_DeviceIndex
 public property uint3 gl_NumWorkGroups
 {
     [require(glsl_spirv, GLSL_430_SPIRV_1_0_compute)]
+    [require(glsl_spirv, meshshading)]
     get 
     {
-        __target_switch
-        {
-        case glsl:
-            __intrinsic_asm "(gl_NumWorkGroups)";
-        case spirv:
-            return spirv_asm {
-                    result:$$uint3 = OpLoad builtin(NumWorkgroups:uint3);
-                };
-        }
+        return WorkgroupCount();
     }
 }
 
-[require(compute)]
 public property uint3 gl_WorkGroupSize
 {
     [__unsafeForceInlineEarly]
     [require(compute)]
+    [require(meshshading)]
     get
     {
         return WorkgroupSize();
diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
@@ -6619,9 +6619,26 @@ void AllMemoryBarrierWithGroupSync()
 
 // Returns the workgroup size of the calling entry point.
 [require(compute)]
+[require(meshshading)]
 __intrinsic_op($(kIROp_GetWorkGroupSize))
 int3 WorkgroupSize();
 
+// Returns number of workgroups that have been dispatched to a GLSL or SPIR-V compute shader
+[require(glsl_spirv, GLSL_430_SPIRV_1_0_compute)]
+[require(glsl_spirv, meshshading)]
+uint3 WorkgroupCount()
+{
+    __target_switch
+    {
+    case glsl:
+        __intrinsic_asm "(gl_NumWorkGroups)";
+    case spirv:
+        return spirv_asm {
+            result:$$uint3 = OpLoad builtin(NumWorkgroups:uint3);
+        };
+    }
+}
+
 // Test if any components is non-zero.
 
 __generic<T : __BuiltinType>
diff --git a/tests/pipeline/rasterization/mesh/task-simple.slang b/tests/pipeline/rasterization/mesh/task-simple.slang
@@ -34,12 +34,14 @@ struct MeshPayload
     int exponent;
 };
 
-[numthreads(1, 1, 1)]
+const static uint AMPLIFICATION_NUM_THREADS_X = 1;
+
+[numthreads(AMPLIFICATION_NUM_THREADS_X, 1, 1)]
 [shader("amplification")]
 void taskMain(in uint tig : SV_GroupIndex)
 {
     MeshPayload p;
-    p.exponent = 3;
+    p.exponent = select(AMPLIFICATION_NUM_THREADS_X == WorkgroupSize().x, 3, 0);
     DispatchMesh(1,1,1,p);
 }
 
@@ -71,8 +73,10 @@ struct Vertex
 const static uint MAX_VERTS = 12;
 const static uint MAX_PRIMS = 4;
 
+const static uint MESH_NUM_THREADS_X = 12;
+
 [outputtopology("triangle")]
-[numthreads(12, 1, 1)]
+[numthreads(MESH_NUM_THREADS_X, 1, 1)]
 void meshMain(
     in uint tig : SV_GroupIndex,
     in payload MeshPayload meshPayload,
@@ -88,7 +92,7 @@ void meshMain(
 
     if(tig < numVertices)
     {
-        const int tri = tig / 3;
+        const int tri = select(WorkgroupSize().x == MESH_NUM_THREADS_X, tig / 3, -1);
         verts[tig] = {float4(positions[tig % 3], 0, 1), colors[tig % 3], tri, int(pow(tri, meshPayload.exponent))};
     }
 

Original file line number	Diff line number	Diff line change
`@@ -141,25 +141,18 @@ public property int gl_DeviceIndex`
`141`	`141`	`public property uint3 gl_NumWorkGroups`
`142`	`142`	`{`
`143`	`143`	`[require(glsl_spirv, GLSL_430_SPIRV_1_0_compute)]`
	`144`	`+ [require(glsl_spirv, meshshading)]`
`144`	`145`	`get`
`145`	`146`	`{`
`146`		`- __target_switch`
`147`		`- {`
`148`		`- case glsl:`
`149`		`- __intrinsic_asm "(gl_NumWorkGroups)";`
`150`		`- case spirv:`
`151`		`- return spirv_asm {`
`152`		`- result:$$uint3 = OpLoad builtin(NumWorkgroups:uint3);`
`153`		`- };`
`154`		`- }`
	`147`	`+ return WorkgroupCount();`
`155`	`148`	`}`
`156`	`149`	`}`
`157`	`150`
`158`		`-[require(compute)]`
`159`	`151`	`public property uint3 gl_WorkGroupSize`
`160`	`152`	`{`
`161`	`153`	`[__unsafeForceInlineEarly]`
`162`	`154`	`[require(compute)]`
	`155`	`+ [require(meshshading)]`
`163`	`156`	`get`
`164`	`157`	`{`
`165`	`158`	`return WorkgroupSize();`