Support SetMeshOutputCounts in mesh nodes (#6476)

Allow the usage of SetMeshOutputCounts() in mesh launch node shaders and
no other launch type nodes. This requires identifying the op initially
as permitted in node shaders, but restraining that permission later on.
Given the availability of launch type and opcode information, the
locations for doing that were limited. Future implementation might opt
to encode the launch type in the stage mask or else at least temporarily
use a different opcode for a SetMeshOutputCounts call within a node
shader.

Fixes #6473

Author: Greg Roth

lib/DXIL/DxilModule.cpp

@@ -2332,6 +2332,22 @@ void DxilModule::UpdateFunctionToShaderCompat(const llvm::Function *dxilFunc) {
       OP::GetMinShaderModelAndMask(CI, bWithTranslation, m_ValMajor, m_ValMinor,
                                    major, minor, mask);
       DXIL::UpdateToMaxOfVersions(info.minMajor, info.minMinor, major, minor);
+
+      // Fix up permitting SetMeshOutputCounts in node shaders
+
+      // This is a hack, but it is required to reject SetMeshOutputCounts from
+      // non-mesh node shaders. It is the least invasive place where the launch
+      // type and the opcode are both known. Within GetMinShaderModelAndMask,
+      // the launch type isn't known. After this point, the opcode that
+      // seemingly indicates support for all node shaders is lost.
+      // Ultimately, we'll probably need to encode launch type into the mask
+      if (DXIL::OpCode::SetMeshOutputCounts == OP::GetDxilOpFuncCallInst(CI) &&
+          HasDxilFunctionProps(F)) {
+        const DxilFunctionProps &props = GetDxilFunctionProps(F);
+        if (props.shaderKind != DXIL::ShaderKind::Node ||
+            props.Node.LaunchType != DXIL::NodeLaunchType::Mesh)
+          mask &= ~SFLAG(Node);
+      }
       info.mask &= mask;
     } else if (const llvm::LoadInst *LI = dyn_cast<LoadInst>(user)) {
       // If loading a groupshared variable, limit to CS/AS/MS/Node

lib/DXIL/DxilOperations.cpp

@@ -3251,14 +3251,21 @@ void OP::GetMinShaderModelAndMask(OpCode C, bool bWithTranslation,
     mask = SFLAG(Library) | SFLAG(Pixel);
     return;
   }
-  // Instructions: SetMeshOutputCounts=168, EmitIndices=169, GetMeshPayload=170,
-  // StoreVertexOutput=171, StorePrimitiveOutput=172
-  if ((168 <= op && op <= 172)) {
+  // Instructions: EmitIndices=169, GetMeshPayload=170, StoreVertexOutput=171,
+  // StorePrimitiveOutput=172
+  if ((169 <= op && op <= 172)) {
     major = 6;
     minor = 5;
     mask = SFLAG(Mesh);
     return;
   }
+  // Instructions: SetMeshOutputCounts=168
+  if (op == 168) {
+    major = 6;
+    minor = 5;
+    mask = SFLAG(Mesh) | SFLAG(Node);
+    return;
+  }
   // Instructions: CreateHandleFromHeap=218, Unpack4x8=219, Pack4x8=220,
   // IsHelperLane=221
   if ((218 <= op && op <= 221)) {

lib/HLSL/DxilValidation.cpp

@@ -1463,7 +1463,8 @@ static void ValidateSignatureDxilOp(CallInst *CI, DXIL::OpCode opcode,
     }
   } break;
   case DXIL::OpCode::SetMeshOutputCounts: {
-    if (!props.IsMS()) {
+    if (!props.IsMS() && (!props.IsNode() || props.Node.LaunchType !=
+                                                 DXIL::NodeLaunchType::Mesh)) {
       ValCtx.EmitInstrFormatError(CI, ValidationRule::SmOpcodeInInvalidFunction,
                                   {"SetMeshOutputCounts", "Mesh shader"});
     }
@@ -2795,9 +2796,11 @@ static void ValidateMsIntrinsics(Function *F, ValidationContext &ValCtx,
                                  CallInst *setMeshOutputCounts,
                                  CallInst *getMeshPayload) {
   if (ValCtx.DxilMod.HasDxilFunctionProps(F)) {
-    DXIL::ShaderKind shaderKind =
-        ValCtx.DxilMod.GetDxilFunctionProps(F).shaderKind;
-    if (shaderKind != DXIL::ShaderKind::Mesh)
+    DxilFunctionProps &props = ValCtx.DxilMod.GetDxilFunctionProps(F);
+    DXIL::ShaderKind shaderKind = props.shaderKind;
+    if (shaderKind != DXIL::ShaderKind::Mesh &&
+        (shaderKind != DXIL::ShaderKind::Node ||
+         props.Node.LaunchType != DXIL::NodeLaunchType::Mesh))
       return;
   } else {
     return;

tools/clang/test/CodeGenDXIL/hlsl/objects/NodeObjects/mesh-node-invalid-set-mesh-output-counts.hlsl

@@ -0,0 +1,34 @@
+// RUN: not %dxc -T lib_6_9 %s 2>&1 | FileCheck %s
+
+// REQUIRES: dxil-1-9
+ 
+// Ensure that setMeshOutputCounts will be appropriately rejected by non-mesh nodes
+
+RWBuffer<float> buf0;
+
+// CHECK:  14: error: Function uses features incompatible with the shader stage (node) of the entry function.
+[Shader("node")]
+[NumThreads(1024,1,1)]
+[NodeDispatchGrid(64,1,1)]
+[NodeLaunch("broadcasting")]
+void node_broadcasting() {
+  SetMeshOutputCounts(32, 16);
+  buf0[0] = 1.0;
+}
+
+// CHECK:  23: error: Function uses features incompatible with the shader stage (node) of the entry function.
+[Shader("node")]
+[NodeLaunch("coalescing")]
+[NumThreads(1024,1,1)]
+void node_coalescing() {
+  SetMeshOutputCounts(32, 16);
+  buf0[0] = 2.0;
+}
+
+// CHECK:  31: error: Function uses features incompatible with the shader stage (node) of the entry function.
+[Shader("node")]
+[NodeLaunch("thread")]
+void node_thread() {
+  SetMeshOutputCounts(32, 16);
+  buf0[0] = 4.0;
+}

tools/clang/test/CodeGenDXIL/hlsl/objects/NodeObjects/mesh-node-set-mesh-output-counts.hlsl

@@ -0,0 +1,24 @@
+// RUN: %dxc -T lib_6_9 %s | FileCheck %s
+
+// REQUIRES: dxil-1-9
+ 
+// Ensure that setMeshOutputCounts can be correctly lowered in a mesh node
+
+// CHECK: define void @node_setmeshoutputcounts()
+// CHECK: dx.op.setMeshOutputCounts(i32 168, i32 32, i32 16)
+// CHECK: ret void
+
+// CHECK: declare void @dx.op.setMeshOutputCounts(i32, i32, i32) #0
+
+RWBuffer<float> buf0;
+
+
+[Shader("node")]
+[NodeLaunch("mesh")]
+[outputtopology("triangle")]
+[numthreads(128, 1, 1)]
+[NodeDispatchGrid(64,1,1)]
+void node_setmeshoutputcounts() {
+    SetMeshOutputCounts(32, 16);
+    buf0[0] = 1.0;
+}

utils/hct/hctdb.py

@@ -594,9 +594,16 @@ def populate_categories_and_models(self):
         for i in "WaveMatch,WaveMultiPrefixOp,WaveMultiPrefixBitCount".split(","):
             self.name_idx[i].category = "Wave"
             self.name_idx[i].shader_model = 6, 5
+        for i in "SetMeshOutputCounts".split(","):
+            self.name_idx[i].category = "Mesh shader instructions"
+            self.name_idx[i].shader_stages = (
+                "mesh",
+                "node",
+            )
+            self.name_idx[i].shader_model = 6, 5
         for (
             i
-        ) in "SetMeshOutputCounts,EmitIndices,GetMeshPayload,StoreVertexOutput,StorePrimitiveOutput".split(
+        ) in "EmitIndices,GetMeshPayload,StoreVertexOutput,StorePrimitiveOutput".split(
             ","
         ):
             self.name_idx[i].category = "Mesh shader instructions"