[HLSL] Adding support for root descriptors in root signature metadata representation

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 = cast<MDString>(Node->getOperand(OpId));
+  if (NodeText == nullptr)
+    return std::nullopt;
+  return NodeText->getString();
+}
+
 static bool parseRootFlags(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
                            MDNode *RootFlagNode) {
 
@@ -105,6 +113,56 @@ static bool parseRootConstants(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
   return false;
 }
 
+static bool parseRootDescriptors(LLVMContext *Ctx,
+                                 mcdxbc::RootSignatureDesc &RSD,
+                                 MDNode *RootDescriptorNode) {
+
+  if (RootDescriptorNode->getNumOperands() != 5)
+    return reportError(Ctx, "Invalid format for RootConstants Element");
+
+  std::optional<StringRef> ElementText =
+      extractMdStringValue(RootDescriptorNode, 0);
+  assert(!ElementText->empty());
+
+  dxbc::RootParameterHeader Header;
+  Header.ParameterType =
+      StringSwitch<uint32_t>(*ElementText)
+          .Case("RootCBV", llvm::to_underlying(dxbc::RootParameterType::CBV))
+          .Case("RootSRV", llvm::to_underlying(dxbc::RootParameterType::SRV))
+          .Case("RootUAV", llvm::to_underlying(dxbc::RootParameterType::UAV));
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(RootDescriptorNode, 1))
+    Header.ShaderVisibility = *Val;
+  else
+    return reportError(Ctx, "Invalid value for ShaderVisibility");
+
+  dxbc::RTS0::v1::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);
+  dxbc::RTS0::v2::RootDescriptor DescriptorV2(Descriptor);
+
+  if (std::optional<uint32_t> Val = extractMdIntValue(RootDescriptorNode, 4))
+    DescriptorV2.Flags = *Val;
+  else
+    return reportError(Ctx, "Invalid value for Root Descriptor Flags");
+
+  RSD.ParametersContainer.addParameter(Header, DescriptorV2);
+  return false;
+}
+
 static bool parseRootSignatureElement(LLVMContext *Ctx,
                                       mcdxbc::RootSignatureDesc &RSD,
                                       MDNode *Element) {
@@ -116,6 +174,9 @@ static bool parseRootSignatureElement(LLVMContext *Ctx,
       StringSwitch<RootSignatureElementKind>(ElementText->getString())
           .Case("RootFlags", RootSignatureElementKind::RootFlags)
           .Case("RootConstants", RootSignatureElementKind::RootConstants)
+          .Case("RootCBV", RootSignatureElementKind::RootDescriptors)
+          .Case("RootSRV", RootSignatureElementKind::RootDescriptors)
+          .Case("RootUAV", RootSignatureElementKind::RootDescriptors)
           .Default(RootSignatureElementKind::Error);
 
   switch (ElementKind) {
@@ -124,7 +185,8 @@ static bool parseRootSignatureElement(LLVMContext *Ctx,
     return parseRootFlags(Ctx, RSD, Element);
   case RootSignatureElementKind::RootConstants:
     return parseRootConstants(Ctx, RSD, Element);
-    break;
+  case RootSignatureElementKind::RootDescriptors:
+    return parseRootDescriptors(Ctx, RSD, Element);
   case RootSignatureElementKind::Error:
     return reportError(Ctx, "Invalid Root Signature Element: " +
                                 ElementText->getString());
@@ -155,6 +217,16 @@ static bool verifyVersion(uint32_t Version) {
   return (Version == 1 || Version == 2);
 }
 
+static bool verifyRegisterValue(uint32_t RegisterValue) {
+  return !(RegisterValue == 0xFFFFFFFF);
+}
+
+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)) {
@@ -172,6 +244,39 @@ static bool validate(LLVMContext *Ctx, const mcdxbc::RootSignatureDesc &RSD) {
 
     assert(dxbc::isValidParameterType(Info.Header.ParameterType) &&
            "Invalid value for ParameterType");
+
+    auto P = RSD.ParametersContainer.getParameter(&Info);
+    if (!P)
+      return reportError(Ctx, "Cannot locate parameter from Header Info");
+
+    if (std::holds_alternative<const dxbc::RTS0::v1::RootDescriptor *>(*P)) {
+      auto *Descriptor =
+          std::get<const dxbc::RTS0::v1::RootDescriptor *>(P.value());
+
+      if (!verifyRegisterValue(Descriptor->ShaderRegister))
+        return reportValueError(Ctx, "ShaderRegister",
+                                Descriptor->ShaderRegister);
+
+      if (!verifyRegisterSpace(Descriptor->RegisterSpace))
+        return reportValueError(Ctx, "RegisterSpace",
+                                Descriptor->RegisterSpace);
+
+    } else if (std::holds_alternative<const dxbc::RTS0::v2::RootDescriptor *>(
+                   *P)) {
+      auto *Descriptor =
+          std::get<const dxbc::RTS0::v2::RootDescriptor *>(P.value());
+
+      if (!verifyRegisterValue(Descriptor->ShaderRegister))
+        return reportValueError(Ctx, "ShaderRegister",
+                                Descriptor->ShaderRegister);
+
+      if (!verifyRegisterSpace(Descriptor->RegisterSpace))
+        return reportValueError(Ctx, "RegisterSpace",
+                                Descriptor->RegisterSpace);
+
+      if (!verifyDescriptorFlag(Descriptor->Flags))
+        return reportValueError(Ctx, "DescriptorFlag", Descriptor->Flags);
+    }
   }
 
   return false;
@@ -308,6 +413,21 @@ PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
            << "Shader Register: " << Constants->ShaderRegister << "\n";
         OS << indent(Space + 2)
            << "Num 32 Bit Values: " << Constants->Num32BitValues << "\n";
+      } else if (std::holds_alternative<const dxbc::RTS0::v1::RootDescriptor *>(
+                     *P)) {
+        auto *Constants = std::get<const dxbc::RTS0::v1::RootDescriptor *>(*P);
+        OS << indent(Space + 2)
+           << "Register Space: " << Constants->RegisterSpace << "\n";
+        OS << indent(Space + 2)
+           << "Shader Register: " << Constants->ShaderRegister << "\n";
+      } else if (std::holds_alternative<const dxbc::RTS0::v2::RootDescriptor *>(
+                     *P)) {
+        auto *Constants = std::get<const dxbc::RTS0::v2::RootDescriptor *>(*P);
+        OS << indent(Space + 2)
+           << "Register Space: " << Constants->RegisterSpace << "\n";
+        OS << indent(Space + 2)
+           << "Shader Register: " << Constants->ShaderRegister << "\n";
+        OS << indent(Space + 2) << "Flags: " << Constants->Flags << "\n";
       }
     }
     Space--;

llvm/lib/Target/DirectX/DXILRootSignature.h

@@ -27,7 +27,8 @@ namespace dxil {
 enum class RootSignatureElementKind {
   Error = 0,
   RootFlags = 1,
-  RootConstants = 2
+  RootConstants = 2,
+  RootDescriptors = 3
 };
 class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {
   friend AnalysisInfoMixin<RootSignatureAnalysis>;

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  }

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  }

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  }

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  }

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