[HLSL] Adding support for root descriptors in root signature metadata representation (#139781)

- adds parsing from metadata into dxcontainer binary - adds validations as described in the spec - adds testing scenarios closes: [#126638](https://github.com/llvm/llvm-project/issues/126638) --------- Co-authored-by: joaosaffran <joao.saffran@microsoft.com>
2025-06-04 11:21:22 -07:00 · 2025-06-04 11:21:22 -07:00 · ad6575f2da
commit ad6575f2da
parent b194cf1e40
9 changed files with 263 additions and 9 deletions
--- a/llvm/lib/Target/DirectX/DXILRootSignature.cpp
+++ b/llvm/lib/Target/DirectX/DXILRootSignature.cpp
@ -55,6 +55,14 @@ static std::optional<uint32_t> extractMdIntValue(MDNode *Node,
  return std::nullopt;
 }

+static std::optional<StringRef> extractMdStringValue(MDNode *Node,
+                                                     unsigned int OpId) {
+  MDString *NodeText = dyn_cast<MDString>(Node->getOperand(OpId));
+  if (NodeText == nullptr)
+    return std::nullopt;
+  return NodeText->getString();
+}
+
 static bool parseRootFlags(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
                           MDNode *RootFlagNode) {

@ -107,17 +115,79 @@ static bool parseRootConstants(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
  return false;
 }

+static bool parseRootDescriptors(LLVMContext *Ctx,
+                                 mcdxbc::RootSignatureDesc &RSD,
+                                 MDNode *RootDescriptorNode,
+                                 RootSignatureElementKind ElementKind) {
+  assert(ElementKind == RootSignatureElementKind::SRV ||
+         ElementKind == RootSignatureElementKind::UAV ||
+         ElementKind == RootSignatureElementKind::CBV &&
+             "parseRootDescriptors should only be called with RootDescriptor "
+             "element kind.");
+  if (RootDescriptorNode->getNumOperands() != 5)
+    return reportError(Ctx, "Invalid format for Root Descriptor Element");
+
+  dxbc::RTS0::v1::RootParameterHeader Header;
+  switch (ElementKind) {
+  case RootSignatureElementKind::SRV:
+    Header.ParameterType = llvm::to_underlying(dxbc::RootParameterType::SRV);
+    break;
+  case RootSignatureElementKind::UAV:
+    Header.ParameterType = llvm::to_underlying(dxbc::RootParameterType::UAV);
+    break;
+  case RootSignatureElementKind::CBV:
+    Header.ParameterType = llvm::to_underlying(dxbc::RootParameterType::CBV);
+    break;
+  default:
+    llvm_unreachable("invalid Root Descriptor kind");
+    break;
+  }
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(RootDescriptorNode, 1))
+    Header.ShaderVisibility = *Val;
+  else
+    return reportError(Ctx, "Invalid value for ShaderVisibility");
+
+  dxbc::RTS0::v2::RootDescriptor Descriptor;
+  if (std::optional<uint32_t> Val = extractMdIntValue(RootDescriptorNode, 2))
+    Descriptor.ShaderRegister = *Val;
+  else
+    return reportError(Ctx, "Invalid value for ShaderRegister");
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(RootDescriptorNode, 3))
+    Descriptor.RegisterSpace = *Val;
+  else
+    return reportError(Ctx, "Invalid value for RegisterSpace");
+
+  if (RSD.Version == 1) {
+    RSD.ParametersContainer.addParameter(Header, Descriptor);
+    return false;
+  }
+  assert(RSD.Version > 1);
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(RootDescriptorNode, 4))
+    Descriptor.Flags = *Val;
+  else
+    return reportError(Ctx, "Invalid value for Root Descriptor Flags");
+
+  RSD.ParametersContainer.addParameter(Header, Descriptor);
+  return false;
+}
+
 static bool parseRootSignatureElement(LLVMContext *Ctx,
                                      mcdxbc::RootSignatureDesc &RSD,
                                      MDNode *Element) {
-  MDString *ElementText = cast<MDString>(Element->getOperand(0));
-  if (ElementText == nullptr)
+  std::optional<StringRef> ElementText = extractMdStringValue(Element, 0);
+  if (!ElementText.has_value())
    return reportError(Ctx, "Invalid format for Root Element");

  RootSignatureElementKind ElementKind =
-      StringSwitch<RootSignatureElementKind>(ElementText->getString())
+      StringSwitch<RootSignatureElementKind>(*ElementText)
          .Case("RootFlags", RootSignatureElementKind::RootFlags)
          .Case("RootConstants", RootSignatureElementKind::RootConstants)
+          .Case("RootCBV", RootSignatureElementKind::CBV)
+          .Case("RootSRV", RootSignatureElementKind::SRV)
+          .Case("RootUAV", RootSignatureElementKind::UAV)
          .Default(RootSignatureElementKind::Error);

  switch (ElementKind) {
@ -126,10 +196,12 @@ static bool parseRootSignatureElement(LLVMContext *Ctx,
    return parseRootFlags(Ctx, RSD, Element);
  case RootSignatureElementKind::RootConstants:
    return parseRootConstants(Ctx, RSD, Element);
-    break;
+  case RootSignatureElementKind::CBV:
+  case RootSignatureElementKind::SRV:
+  case RootSignatureElementKind::UAV:
+    return parseRootDescriptors(Ctx, RSD, Element, ElementKind);
  case RootSignatureElementKind::Error:
-    return reportError(Ctx, "Invalid Root Signature Element: " +
-                                ElementText->getString());
+    return reportError(Ctx, "Invalid Root Signature Element: " + *ElementText);
  }

  llvm_unreachable("Unhandled RootSignatureElementKind enum.");
@ -157,6 +229,18 @@ static bool verifyVersion(uint32_t Version) {
  return (Version == 1 || Version == 2);
 }

+static bool verifyRegisterValue(uint32_t RegisterValue) {
+  return RegisterValue != ~0U;
+}
+
+// This Range is reserverved, therefore invalid, according to the spec
+// https://github.com/llvm/wg-hlsl/blob/main/proposals/0002-root-signature-in-clang.md#all-the-values-should-be-legal
+static bool verifyRegisterSpace(uint32_t RegisterSpace) {
+  return !(RegisterSpace >= 0xFFFFFFF0 && RegisterSpace <= 0xFFFFFFFF);
+}
+
+static bool verifyDescriptorFlag(uint32_t Flags) { return (Flags & ~0xE) == 0; }
+
 static bool validate(LLVMContext *Ctx, const mcdxbc::RootSignatureDesc &RSD) {

  if (!verifyVersion(RSD.Version)) {
@ -174,6 +258,28 @@ static bool validate(LLVMContext *Ctx, const mcdxbc::RootSignatureDesc &RSD) {

    assert(dxbc::isValidParameterType(Info.Header.ParameterType) &&
           "Invalid value for ParameterType");
+
+    switch (Info.Header.ParameterType) {
+
+    case llvm::to_underlying(dxbc::RootParameterType::CBV):
+    case llvm::to_underlying(dxbc::RootParameterType::UAV):
+    case llvm::to_underlying(dxbc::RootParameterType::SRV): {
+      const dxbc::RTS0::v2::RootDescriptor &Descriptor =
+          RSD.ParametersContainer.getRootDescriptor(Info.Location);
+      if (!verifyRegisterValue(Descriptor.ShaderRegister))
+        return reportValueError(Ctx, "ShaderRegister",
+                                Descriptor.ShaderRegister);
+
+      if (!verifyRegisterSpace(Descriptor.RegisterSpace))
+        return reportValueError(Ctx, "RegisterSpace", Descriptor.RegisterSpace);
+
+      if (RSD.Version > 1) {
+        if (!verifyDescriptorFlag(Descriptor.Flags))
+          return reportValueError(Ctx, "DescriptorFlag", Descriptor.Flags);
+      }
+      break;
+    }
+    }
  }

  return false;
@ -313,6 +419,20 @@ PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
           << "Shader Register: " << Constants.ShaderRegister << "\n";
        OS << indent(Space + 2)
           << "Num 32 Bit Values: " << Constants.Num32BitValues << "\n";
+        break;
+      }
+      case llvm::to_underlying(dxbc::RootParameterType::CBV):
+      case llvm::to_underlying(dxbc::RootParameterType::UAV):
+      case llvm::to_underlying(dxbc::RootParameterType::SRV): {
+        const dxbc::RTS0::v2::RootDescriptor &Descriptor =
+            RS.ParametersContainer.getRootDescriptor(Loc);
+        OS << indent(Space + 2)
+           << "Register Space: " << Descriptor.RegisterSpace << "\n";
+        OS << indent(Space + 2)
+           << "Shader Register: " << Descriptor.ShaderRegister << "\n";
+        if (RS.Version > 1)
+          OS << indent(Space + 2) << "Flags: " << Descriptor.Flags << "\n";
+        break;
      }
      }
      Space--;
--- a/llvm/lib/Target/DirectX/DXILRootSignature.h
+++ b/llvm/lib/Target/DirectX/DXILRootSignature.h
@ -27,7 +27,10 @@ namespace dxil {
 enum class RootSignatureElementKind {
  Error = 0,
  RootFlags = 1,
-  RootConstants = 2
+  RootConstants = 2,
+  SRV = 3,
+  UAV = 4,
+  CBV = 5,
 };
 class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {
  friend AnalysisInfoMixin<RootSignatureAnalysis>;
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Parameters-Invalid-ParameterIsNotString.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Parameters-Invalid-ParameterIsNotString.ll
@ -0,0 +1,19 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+; CHECK: error: Invalid format for Root Element
+; CHECK-NOT: Root Signature Definitions
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+
+define void @main() #0 {
+entry:
+  ret void
+}
+
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+!dx.rootsignatures = !{!0}
+!0 = !{ ptr @main, !1 }
+!1 = !{ !2 }
+!2 = !{ i32 0 }
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Parameters.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Parameters.ll
@ -12,19 +12,25 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }

 !dx.rootsignatures = !{!2} ; list of function/root signature pairs
 !2 = !{ ptr @main, !3 } ; function, root signature
-!3 = !{ !4, !5 } ; list of root signature elements
+!3 = !{ !4, !5, !6 } ; list of root signature elements
 !4 = !{ !"RootFlags", i32 1 } ; 1 = allow_input_assembler_input_layout
 !5 = !{ !"RootConstants", i32 0, i32 1, i32 2, i32 3 }
+!6 = !{ !"RootSRV", i32 1, i32 4, i32 5, i32 6 }

 ;CHECK-LABEL: Definition for 'main':
 ;CHECK-NEXT:  Flags: 0x000001
 ;CHECK-NEXT:  Version: 2
 ;CHECK-NEXT:  RootParametersOffset: 24
-;CHECK-NEXT:  NumParameters: 1
+;CHECK-NEXT:  NumParameters: 2
 ;CHECK-NEXT:   - Parameter Type: 1
 ;CHECK-NEXT:     Shader Visibility: 0
 ;CHECK-NEXT:     Register Space: 2
 ;CHECK-NEXT:     Shader Register: 1
 ;CHECK-NEXT:     Num 32 Bit Values: 3
+;CHECK-NEXT:   - Parameter Type: 3
+;CHECK-NEXT:     Shader Visibility: 1
+;CHECK-NEXT:     Register Space: 5
+;CHECK-NEXT:     Shader Register: 4
+;CHECK-NEXT:     Flags: 6
 ;CHECK-NEXT:  NumStaticSamplers: 0
 ;CHECK-NEXT:  StaticSamplersOffset: 0
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor-Invalid-Flags.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor-Invalid-Flags.ll
@ -0,0 +1,18 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+
+; CHECK: error: Invalid value for DescriptorFlag: 3
+; CHECK-NOT: Root Signature Definitions
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"RootCBV", i32 0, i32 1, i32 2, i32 3  }
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor-Invalid-RegisterKind.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor-Invalid-RegisterKind.ll
@ -0,0 +1,18 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+
+; CHECK: error:  Invalid Root Signature Element: Invalid 
+; CHECK-NOT: Root Signature Definitions
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"Invalid", i32 0, i32 1, i32 2, i32 3  }
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor-Invalid-RegisterSpace.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor-Invalid-RegisterSpace.ll
@ -0,0 +1,18 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+
+; CHECK: error:  Invalid value for RegisterSpace: 4294967280
+; CHECK-NOT: Root Signature Definitions
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"RootCBV", i32 0, i32 1, i32 4294967280, i32 0  }
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor-Invalid-RegisterValue.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor-Invalid-RegisterValue.ll
@ -0,0 +1,18 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+
+; CHECK: error: Invalid value for ShaderRegister: 4294967295 
+; CHECK-NOT: Root Signature Definitions
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"RootCBV", i32 0, i32 4294967295, i32 2, i32 3  }
--- a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor.ll
+++ b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootDescriptor.ll
@ -0,0 +1,34 @@
+; RUN: opt %s -dxil-embed -dxil-globals -S -o - | FileCheck %s
+; RUN: llc %s --filetype=obj -o - | obj2yaml | FileCheck %s --check-prefix=DXC
+
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+; CHECK: @dx.rts0 = private constant [48 x i8]  c"{{.*}}", section "RTS0", align 4
+
+define void @main() #0 {
+entry:
+  ret void
+}
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"RootCBV", i32 0, i32 1, i32 2, i32 8  }
+
+; DXC:  - Name:            RTS0
+; DXC-NEXT:    Size:            48
+; DXC-NEXT:    RootSignature:
+; DXC-NEXT:      Version:         2
+; DXC-NEXT:      NumRootParameters: 1 
+; DXC-NEXT:      RootParametersOffset: 24 
+; DXC-NEXT:      NumStaticSamplers: 0
+; DXC-NEXT:      StaticSamplersOffset: 0
+; DXC-NEXT:      Parameters:
+; DXC-NEXT:        - ParameterType:   2
+; DXC-NEXT:          ShaderVisibility: 0
+; DXC-NEXT:          Descriptor:
+; DXC-NEXT:            RegisterSpace: 2
+; DXC-NEXT:            ShaderRegister: 1
+; DXC-NEXT:            DATA_STATIC: true