[HLSL] Adding support for Root Constants in LLVM Metadata

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -25,7 +25,7 @@ struct RootSignatureDesc {
 
   uint32_t Version = 2U;
   uint32_t Flags = 0U;
-  uint32_t RootParameterOffset = 0U;
+  uint32_t RootParameterOffset = 24U;
   uint32_t StaticSamplersOffset = 0u;
   uint32_t NumStaticSamplers = 0u;
   SmallVector<mcdxbc::RootParameter> Parameters;

llvm/lib/Target/DirectX/DXILRootSignature.cpp

@@ -40,14 +40,64 @@ static bool reportError(LLVMContext *Ctx, Twine Message,
   return true;
 }
 
+static bool reportValueError(LLVMContext *Ctx, Twine ParamName, uint32_t Value,
+                             DiagnosticSeverity Severity = DS_Error) {
+  Ctx->diagnose(DiagnosticInfoGeneric(
+      "Invalid value for " + ParamName + ": " + Twine(Value), Severity));
+  return true;
+}
+
+static bool extractMdIntValue(uint32_t &Value, MDNode *Node,
+                              unsigned int OpId) {
+  auto *CI = mdconst::dyn_extract<ConstantInt>(Node->getOperand(OpId).get());
+  if (CI == nullptr)
+    return true;
+
+  Value = CI->getZExtValue();
+  return false;
+}
+
 static bool parseRootFlags(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
                            MDNode *RootFlagNode) {
 
   if (RootFlagNode->getNumOperands() != 2)
     return reportError(Ctx, "Invalid format for RootFlag Element");
 
-  auto *Flag = mdconst::extract<ConstantInt>(RootFlagNode->getOperand(1));
-  RSD.Flags = Flag->getZExtValue();
+  if (extractMdIntValue(RSD.Flags, RootFlagNode, 1))
+    return reportError(Ctx, "Invalid value for RootFlag");
+
+  return false;
+}
+
+static bool parseRootConstants(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
+                               MDNode *RootConstantNode) {
+
+  if (RootConstantNode->getNumOperands() != 5)
+    return reportError(Ctx, "Invalid format for RootConstants Element");
+
+  mcdxbc::RootParameter NewParameter;
+  NewParameter.Header.ParameterType =
+      llvm::to_underlying(dxbc::RootParameterType::Constants32Bit);
+
+  uint32_t SV;
+  if (extractMdIntValue(SV, RootConstantNode, 1))
+    return reportError(Ctx, "Invalid value for ShaderVisibility");
+
+  NewParameter.Header.ShaderVisibility = SV;
+
+  if (extractMdIntValue(NewParameter.Constants.ShaderRegister, RootConstantNode,
+                        2))
+    return reportError(Ctx, "Invalid value for ShaderRegister");
+
+  if (extractMdIntValue(NewParameter.Constants.RegisterSpace, RootConstantNode,
+                        3))
+    return reportError(Ctx, "Invalid value for RegisterSpace");
+
+  if (extractMdIntValue(NewParameter.Constants.Num32BitValues, RootConstantNode,
+                        4))
+    return reportError(Ctx, "Invalid value for Num32BitValues");
+
+  RSD.Parameters.push_back(NewParameter);
 
   return false;
 }
@@ -62,12 +112,16 @@ static bool parseRootSignatureElement(LLVMContext *Ctx,
   RootSignatureElementKind ElementKind =
       StringSwitch<RootSignatureElementKind>(ElementText->getString())
           .Case("RootFlags", RootSignatureElementKind::RootFlags)
+          .Case("RootConstants", RootSignatureElementKind::RootConstants)
           .Default(RootSignatureElementKind::Error);
 
   switch (ElementKind) {
 
   case RootSignatureElementKind::RootFlags:
     return parseRootFlags(Ctx, RSD, Element);
+  case RootSignatureElementKind::RootConstants:
+    return parseRootConstants(Ctx, RSD, Element);
+    break;
   case RootSignatureElementKind::Error:
     return reportError(Ctx, "Invalid Root Signature Element: " +
                                 ElementText->getString());
@@ -94,10 +148,55 @@ static bool parse(LLVMContext *Ctx, mcdxbc::RootSignatureDesc &RSD,
 
 static bool verifyRootFlag(uint32_t Flags) { return (Flags & ~0xfff) == 0; }
 
+static bool verifyShaderVisibility(uint32_t Flags) {
+  switch (Flags) {
+
+  case llvm::to_underlying(dxbc::ShaderVisibility::All):
+  case llvm::to_underlying(dxbc::ShaderVisibility::Vertex):
+  case llvm::to_underlying(dxbc::ShaderVisibility::Hull):
+  case llvm::to_underlying(dxbc::ShaderVisibility::Domain):
+  case llvm::to_underlying(dxbc::ShaderVisibility::Geometry):
+  case llvm::to_underlying(dxbc::ShaderVisibility::Pixel):
+  case llvm::to_underlying(dxbc::ShaderVisibility::Amplification):
+  case llvm::to_underlying(dxbc::ShaderVisibility::Mesh):
+    return true;
+  }
+
+  return false;
+}
+
+static bool verifyParameterType(uint32_t Type) {
+  switch (Type) {
+  case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
+    return true;
+  }
+
+  return false;
+}
+
+static bool verifyVersion(uint32_t Version) {
+  return (Version == 1 || Version == 2);
+}
+
 static bool validate(LLVMContext *Ctx, const mcdxbc::RootSignatureDesc &RSD) {
+
+  if (!verifyVersion(RSD.Version)) {
+    return reportValueError(Ctx, "Version", RSD.Version);
+  }
+
   if (!verifyRootFlag(RSD.Flags)) {
-    return reportError(Ctx, "Invalid Root Signature flag value");
+    return reportValueError(Ctx, "RootFlags", RSD.Flags);
   }
+
+  for (const auto &P : RSD.Parameters) {
+    if (!verifyShaderVisibility(P.Header.ShaderVisibility))
+      return reportValueError(Ctx, "ShaderVisibility",
+                              (uint32_t)P.Header.ShaderVisibility);
+
+    assert(verifyParameterType(P.Header.ParameterType) &&
+           "Invalid value for ParameterType");
+  }
+
   return false;
 }
 
@@ -205,14 +304,35 @@ PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
 
     // start root signature header
     Space++;
-    OS << indent(Space) << "Flags: " << format_hex(RS.Flags, 8) << ":\n";
-    OS << indent(Space) << "Version: " << RS.Version << ":\n";
-    OS << indent(Space) << "NumParameters: " << RS.Parameters.size() << ":\n";
-    OS << indent(Space) << "RootParametersOffset: " << RSHSize << ":\n";
-    OS << indent(Space) << "NumStaticSamplers: " << 0 << ":\n";
+    OS << indent(Space) << "Flags: " << format_hex(RS.Flags, 8) << "\n";
+    OS << indent(Space) << "Version: " << RS.Version << "\n";
+    OS << indent(Space) << "NumParameters: " << RS.Parameters.size() << "\n";
+    OS << indent(Space) << "RootParametersOffset: " << RSHSize << "\n";
+    OS << indent(Space) << "NumStaticSamplers: " << 0 << "\n";
     OS << indent(Space)
        << "StaticSamplersOffset: " << RSHSize + RS.Parameters.size_in_bytes()
-       << ":\n";
+       << "\n";
+
+    Space++;
+    for (auto const &P : RS.Parameters) {
+      OS << indent(Space)
+         << "Parameter Type: " << (uint32_t)P.Header.ParameterType << "\n";
+      OS << indent(Space)
+         << "Shader Visibility: " << (uint32_t)P.Header.ShaderVisibility
+         << "\n";
+      switch (P.Header.ParameterType) {
+      case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
+        OS << indent(Space) << "Register Space: " << P.Constants.RegisterSpace
+           << "\n";
+        OS << indent(Space) << "Shader Register: " << P.Constants.ShaderRegister
+           << "\n";
+        OS << indent(Space)
+           << "Num 32 Bit Values: " << P.Constants.Num32BitValues << "\n";
+        break;
+      }
+    }
+    Space--;
+
     Space--;
     // end root signature header
   }

llvm/lib/Target/DirectX/DXILRootSignature.h

@@ -24,7 +24,11 @@
 namespace llvm {
 namespace dxil {
 
-enum class RootSignatureElementKind { Error = 0, RootFlags = 1 };
+enum class RootSignatureElementKind {
+  Error = 0,
+  RootFlags = 1,
+  RootConstants = 2
+};
 class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {
   friend AnalysisInfoMixin<RootSignatureAnalysis>;
   static AnalysisKey Key;

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags.ll

@@ -23,7 +23,7 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 ; DXC-NEXT:    RootSignature:
 ; DXC-NEXT:      Version:         2
 ; DXC-NEXT:      NumRootParameters: 0
-; DXC-NEXT:      RootParametersOffset: 0
+; DXC-NEXT:      RootParametersOffset: 24
 ; DXC-NEXT:      NumStaticSamplers: 0
 ; DXC-NEXT:      StaticSamplersOffset: 0
 ; DXC-NEXT:      Parameters: []

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Parameters-Validation-Error.ll

@@ -0,0 +1,20 @@
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+; CHECK: error: Invalid value for ShaderVisibility: 255
+; 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 = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"RootConstants", i32 255, i32 1, i32 2, i32 3 }

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Parameters.ll

@@ -0,0 +1,30 @@
+; RUN: opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
+
+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 = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !4, !5 } ; 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 }
+
+; CHECK-LABEL: Definition for 'main':
+; CHECK-NEXT: Flags: 0x000001
+; CHECK-NEXT: Version: 2
+; CHECK-NEXT: NumParameters: 1
+; CHECK-NEXT: RootParametersOffset: 24
+; CHECK-NEXT: NumStaticSamplers: 0
+; CHECK-NEXT: StaticSamplersOffset: 48
+; 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

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootConstants-Invalid-Num32BitValues.ll

@@ -0,0 +1,16 @@
+; 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 Num32BitValues
+; CHECK-NOT: Root Signature Definitions
+
+define void @main() {
+entry:
+  ret void
+}
+
+!dx.rootsignatures = !{!2} ; list of function/root signature pairs
+!2 = !{ ptr @main, !3 } ; function, root signature
+!3 = !{ !5 } ; list of root signature elements
+!5 = !{ !"RootConstants", i32 0, i32 1, i32 2, !"Invalid" }

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootConstants-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
+; 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 = !{ !"RootConstants", i32 0, i32 1, !"Invalid", i32 3 }

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootConstants-Invalid-ShaderRegister.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
+; 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 = !{ !"RootConstants", i32 0, !"Invalid", i32 2, i32 3 }

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootConstants.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 = !{ !"RootConstants", i32 0, i32 1, i32 2, i32 3 }
+
+; 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:   1
+; DXC-NEXT:          ShaderVisibility: 0
+; DXC-NEXT:          Constants:
+; DXC-NEXT:            Num32BitValues:  3
+; DXC-NEXT:            RegisterSpace:   2
+; DXC-NEXT:            ShaderRegister:  1

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags-Validation-Error.ll → llvm/test/CodeGen/DirectX/ContainerData/RootSignature-RootFlags-VisibilityValidationError.ll

@@ -1,6 +1,6 @@
 ; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
 
-; CHECK: error: Invalid Root Signature flag value
+; CHECK: error: Invalid value for RootFlags: 2147487744
 ; CHECK-NOT: Root Signature Definitions
 
 target triple = "dxil-unknown-shadermodel6.0-compute"