Write usage masks for mesh nodes in libraries into RDAT (#6522)

Previously, the step that is typically taken to set the usage mask for
signature elements was not executed, due to mesh nodes being defined in
library shaders. The mask was only written for non-lib shaders. Now that
mesh nodes exist in library shaders, we need to allow usage masks to be
set for mesh nodes, even though they are in library shaders.

This PR adds an alternative case for mesh nodes. In this case, usage
masks are written for all mesh nodes defined in a library shader, by
iterating over all functions declared in a library module. An
accompanying test was written that shows the RDAT now contains the usage
mask, and it is non-null.

Fixes #6521

---------

Co-authored-by: Tex Riddell <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 3 people

lib/HLSL/DxilPreparePasses.cpp

@@ -287,7 +287,7 @@ static bool GetUnsignedVal(Value *V, uint32_t *pValue) {
   return true;
 }
 
-static void MarkUsedSignatureElements(Function *F, DxilModule &DM) {
+static void MarkUsedSignatureElements(Function *F, DxilEntryProps &entryProps) {
   DXASSERT_NOMSG(F != nullptr);
   // For every loadInput/storeOutput, update the corresponding ReadWriteMask.
   // F is a pointer to a Function instance
@@ -309,54 +309,54 @@ static void MarkUsedSignatureElements(Function *F, DxilModule &DM) {
         continue;
       if (!GetUnsignedVal(LI.get_rowIndex(), &row))
         bDynIdx = true;
-      pSig = &DM.GetInputSignature();
+      pSig = &entryProps.sig.InputSignature;
     } else if (SO) {
       if (!GetUnsignedVal(SO.get_outputSigId(), &sigId))
         continue;
       if (!GetUnsignedVal(SO.get_colIndex(), &col))
         continue;
       if (!GetUnsignedVal(SO.get_rowIndex(), &row))
         bDynIdx = true;
-      pSig = &DM.GetOutputSignature();
+      pSig = &entryProps.sig.OutputSignature;
     } else if (SPC) {
       if (!GetUnsignedVal(SPC.get_outputSigID(), &sigId))
         continue;
       if (!GetUnsignedVal(SPC.get_col(), &col))
         continue;
       if (!GetUnsignedVal(SPC.get_row(), &row))
         bDynIdx = true;
-      pSig = &DM.GetPatchConstOrPrimSignature();
+      pSig = &entryProps.sig.PatchConstOrPrimSignature;
     } else if (LPC) {
       if (!GetUnsignedVal(LPC.get_inputSigId(), &sigId))
         continue;
       if (!GetUnsignedVal(LPC.get_col(), &col))
         continue;
       if (!GetUnsignedVal(LPC.get_row(), &row))
         bDynIdx = true;
-      pSig = &DM.GetPatchConstOrPrimSignature();
+      pSig = &entryProps.sig.PatchConstOrPrimSignature;
     } else if (SVO) {
       if (!GetUnsignedVal(SVO.get_outputSigId(), &sigId))
         continue;
       if (!GetUnsignedVal(SVO.get_colIndex(), &col))
         continue;
       if (!GetUnsignedVal(SVO.get_rowIndex(), &row))
         bDynIdx = true;
-      pSig = &DM.GetOutputSignature();
+      pSig = &entryProps.sig.OutputSignature;
     } else if (SPO) {
       if (!GetUnsignedVal(SPO.get_outputSigId(), &sigId))
         continue;
       if (!GetUnsignedVal(SPO.get_colIndex(), &col))
         continue;
       if (!GetUnsignedVal(SPO.get_rowIndex(), &row))
         bDynIdx = true;
-      pSig = &DM.GetPatchConstOrPrimSignature();
+      pSig = &entryProps.sig.PatchConstOrPrimSignature;
     } else {
       continue;
     }
 
     // Consider being more fine-grained about masks.
     // We report sometimes-read on input as always-read.
-    auto &El = pSig->GetElement(sigId);
+    hlsl::DxilSignatureElement &El = pSig->GetElement(sigId);
     unsigned UsageMask = El.GetUsageMask();
     unsigned colBit = 1 << col;
     if (!(colBit & UsageMask)) {
@@ -843,9 +843,25 @@ class DxilFinalizeModule : public ModulePass {
 
       if (!IsLib) {
         // Set used masks for signature elements
-        MarkUsedSignatureElements(DM.GetEntryFunction(), DM);
-        if (DM.GetShaderModel()->IsHS())
-          MarkUsedSignatureElements(DM.GetPatchConstantFunction(), DM);
+        Function *F = DM.GetEntryFunction();
+        // Assume entry props are available for entry function.
+        DxilEntryProps &entryProps = DM.GetDxilEntryProps(F);
+        MarkUsedSignatureElements(F, entryProps);
+        if (entryProps.props.IsHS() &&
+            entryProps.props.ShaderProps.HS.patchConstantFunc)
+          MarkUsedSignatureElements(
+              entryProps.props.ShaderProps.HS.patchConstantFunc, entryProps);
+      } else {
+        for (auto &function : M.getFunctionList()) {
+          if (!function.isDeclaration()) {
+            if (DM.HasDxilEntryProps(&function)) {
+              auto &entryProps = DM.GetDxilEntryProps(&function);
+              if (entryProps.props.IsMeshNode()) {
+                MarkUsedSignatureElements(&function, entryProps);
+              }
+            }
+          }
+        }
       }
 
       // Adding warning for pixel shader with unassigned target

tools/clang/test/CodeGenDXIL/hlsl/shaders/node/mesh-node.hlsl

@@ -104,7 +104,7 @@ void node_setmeshoutputcounts(DispatchNodeInputRecord<MeshPayload> mpl,
     op.layer[5] = mpl.Get().layer[5];
 
     // CHECK: %[[GSGEP2:[0-9]+]] = getelementptr [16 x float], [16 x float] addrspace(3)* @"\01?gsMem@@3PAMA", i32 0, i32 %[[GIDIV]]
-    // CHECK: store float %[[NRM]], float addrspace(3)* %[[GSGEP2]], align 4, !tbaa !29
+    // CHECK: store float %[[NRM]], float addrspace(3)* %[[GSGEP2]], align 4
     gsMem[tig / 3] = op.normal;
     // CHECK:  call void @dx.op.storePrimitiveOutput.f32(i32 172, i32 0, i32 0, i8 0, float %[[NRM]], i32 %[[GIDIV]])
     // CHECK:  call void @dx.op.storePrimitiveOutput.f32(i32 172, i32 1, i32 0, i8 0, float %[[GS]], i32 %[[GIDIV]])