[DirectX] Validate registers are bound to root signature

llvm/include/llvm/Analysis/DXILResource.h

@@ -33,6 +33,20 @@ class DXILResourceTypeMap;
 
 namespace dxil {
 
+inline StringRef getResourceClassName(ResourceClass RC) {
+  switch (RC) {
+  case ResourceClass::SRV:
+    return "SRV";
+  case ResourceClass::UAV:
+    return "UAV";
+  case ResourceClass::CBuffer:
+    return "CBuffer";
+  case ResourceClass::Sampler:
+    return "Sampler";
+  }
+  llvm_unreachable("Unhandled ResourceClass");
+}
+
 // Returns the resource name from dx_resource_handlefrombinding or
 // dx_resource_handlefromimplicitbinding call
 LLVM_ABI StringRef getResourceNameFromBindingCall(CallInst *CI);

llvm/include/llvm/Frontend/HLSL/HLSLBinding.h

@@ -57,6 +57,7 @@ class BindingInfo {
     }
     // Size == -1 means unbounded array
     LLVM_ABI std::optional<uint32_t> findAvailableBinding(int32_t Size);
+    LLVM_ABI bool isBound(BindingRange B);
   };
 
   struct BindingSpaces {
@@ -95,6 +96,8 @@ class BindingInfo {
   LLVM_ABI std::optional<uint32_t>
   findAvailableBinding(dxil::ResourceClass RC, uint32_t Space, int32_t Size);
 
+  LLVM_ABI bool isBound(dxil::ResourceClass RC, uint32_t Space, BindingRange B);
+
   friend class BindingInfoBuilder;
 };
 

llvm/lib/Analysis/DXILResource.cpp

@@ -29,20 +29,6 @@
 using namespace llvm;
 using namespace dxil;
 
-static StringRef getResourceClassName(ResourceClass RC) {
-  switch (RC) {
-  case ResourceClass::SRV:
-    return "SRV";
-  case ResourceClass::UAV:
-    return "UAV";
-  case ResourceClass::CBuffer:
-    return "CBuffer";
-  case ResourceClass::Sampler:
-    return "Sampler";
-  }
-  llvm_unreachable("Unhandled ResourceClass");
-}
-
 static StringRef getResourceKindName(ResourceKind RK) {
   switch (RK) {
   case ResourceKind::Texture1D:

llvm/lib/Frontend/HLSL/HLSLBinding.cpp

@@ -66,6 +66,22 @@ BindingInfo::RegisterSpace::findAvailableBinding(int32_t Size) {
   return std::nullopt;
 }
 
+bool BindingInfo::RegisterSpace::isBound(BindingRange B) {
+  for (BindingRange &R : FreeRanges) {
+    if (B.LowerBound >= R.LowerBound && B.LowerBound < R.UpperBound &&
+        B.UpperBound > R.LowerBound && B.UpperBound <= R.UpperBound)
+      return false;
+  }
+  return true;
+}
+
+bool BindingInfo::isBound(dxil::ResourceClass RC, uint32_t Space,
+                          BindingRange B) {
+  BindingSpaces &BS = getBindingSpaces(RC);
+  RegisterSpace &RS = BS.getOrInsertSpace(Space);
+  return RS.isBound(B);
+}
+
 BindingInfo BindingInfoBuilder::calculateBindingInfo(
     llvm::function_ref<void(const BindingInfoBuilder &Builder,
                             const Binding &Overlapping)>

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -9,6 +9,7 @@
 #include "DXILOpLowering.h"
 #include "DXILConstants.h"
 #include "DXILOpBuilder.h"
+#include "DXILRootSignature.h"
 #include "DXILShaderFlags.h"
 #include "DirectX.h"
 #include "llvm/ADT/SmallVector.h"
@@ -945,6 +946,7 @@ class DXILOpLoweringLegacy : public ModulePass {
     AU.addPreserved<DXILResourceWrapperPass>();
     AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
     AU.addPreserved<ShaderFlagsAnalysisWrapper>();
+    AU.addPreserved<RootSignatureAnalysisWrapper>();
   }
 };
 char DXILOpLoweringLegacy::ID = 0;

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "DXILPostOptimizationValidation.h"
+#include "DXILRootSignature.h"
 #include "DXILShaderFlags.h"
 #include "DirectX.h"
 #include "llvm/ADT/SmallString.h"
@@ -24,6 +25,35 @@ using namespace llvm;
 using namespace llvm::dxil;
 
 namespace {
+static ResourceClass RangeToResourceClass(uint32_t RangeType) {
+  using namespace dxbc;
+  switch (static_cast<DescriptorRangeType>(RangeType)) {
+  case DescriptorRangeType::SRV:
+    return ResourceClass::SRV;
+  case DescriptorRangeType::UAV:
+    return ResourceClass::UAV;
+  case DescriptorRangeType::CBV:
+    return ResourceClass::CBuffer;
+  case DescriptorRangeType::Sampler:
+    return ResourceClass::Sampler;
+  }
+}
+
+ResourceClass ParameterToResourceClass(uint32_t Type) {
+  using namespace dxbc;
+  switch (Type) {
+  case llvm::to_underlying(RootParameterType::Constants32Bit):
+    return ResourceClass::CBuffer;
+  case llvm::to_underlying(RootParameterType::SRV):
+    return ResourceClass::SRV;
+  case llvm::to_underlying(RootParameterType::UAV):
+    return ResourceClass::UAV;
+  case llvm::to_underlying(RootParameterType::CBV):
+    return ResourceClass::CBuffer;
+  default:
+    llvm_unreachable("Unknown RootParameterType");
+  }
+}
 
 static void reportInvalidDirection(Module &M, DXILResourceMap &DRM) {
   for (const auto &UAV : DRM.uavs()) {
@@ -84,8 +114,112 @@ static void reportOverlappingBinding(Module &M, DXILResourceMap &DRM) {
   }
 }
 
+static void
+reportRegNotBound(Module &M, ResourceClass Class,
+                  llvm::dxil::ResourceInfo::ResourceBinding Unbound) {
+  SmallString<128> Message;
+  raw_svector_ostream OS(Message);
+  OS << "register " << getResourceClassName(Class)
+     << " (space=" << Unbound.Space << ", register=" << Unbound.LowerBound
+     << ")"
+     << " does not have a binding in the Root Signature";
+  M.getContext().diagnose(DiagnosticInfoGeneric(Message));
+}
+
+static dxbc::ShaderVisibility
+tripleToVisibility(llvm::Triple::EnvironmentType ET) {
+  assert((ET == Triple::Pixel || ET == Triple::Vertex ||
+          ET == Triple::Geometry || ET == Triple::Hull ||
+          ET == Triple::Domain || ET == Triple::Mesh ||
+          ET == Triple::Compute) &&
+         "Invalid Triple to shader stage conversion");
+
+  switch (ET) {
+  case Triple::Pixel:
+    return dxbc::ShaderVisibility::Pixel;
+  case Triple::Vertex:
+    return dxbc::ShaderVisibility::Vertex;
+  case Triple::Geometry:
+    return dxbc::ShaderVisibility::Geometry;
+  case Triple::Hull:
+    return dxbc::ShaderVisibility::Hull;
+  case Triple::Domain:
+    return dxbc::ShaderVisibility::Domain;
+  case Triple::Mesh:
+    return dxbc::ShaderVisibility::Mesh;
+  case Triple::Compute:
+    return dxbc::ShaderVisibility::All;
+  default:
+    llvm_unreachable("Invalid triple to shader stage conversion");
+  }
+}
+
+static void trackRootSigDescBinding(hlsl::BindingInfoBuilder &Builder,
+                                    const mcdxbc::RootSignatureDesc &RSD,
+                                    dxbc::ShaderVisibility Visibility) {
+  for (size_t I = 0; I < RSD.ParametersContainer.size(); I++) {
+    const auto &[Type, Loc] =
+        RSD.ParametersContainer.getTypeAndLocForParameter(I);
+
+    const auto &Header = RSD.ParametersContainer.getHeader(I);
+    if (Header.ShaderVisibility !=
+            llvm::to_underlying(dxbc::ShaderVisibility::All) &&
+        Header.ShaderVisibility != llvm::to_underlying(Visibility))
+      continue;
+
+    switch (Type) {
+    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
+      dxbc::RTS0::v1::RootConstants Const =
+          RSD.ParametersContainer.getConstant(Loc);
+      Builder.trackBinding(dxil::ResourceClass::CBuffer, Const.RegisterSpace,
+                           Const.ShaderRegister,
+                           Const.ShaderRegister + Const.Num32BitValues,
+                           nullptr);
+      break;
+    }
+
+    case llvm::to_underlying(dxbc::RootParameterType::SRV):
+    case llvm::to_underlying(dxbc::RootParameterType::UAV):
+    case llvm::to_underlying(dxbc::RootParameterType::CBV): {
+      dxbc::RTS0::v2::RootDescriptor Desc =
+          RSD.ParametersContainer.getRootDescriptor(Loc);
+      Builder.trackBinding(ParameterToResourceClass(Type), Desc.RegisterSpace,
+                           Desc.ShaderRegister, Desc.ShaderRegister, nullptr);
+
+      break;
+    }
+    case llvm::to_underlying(dxbc::RootParameterType::DescriptorTable): {
+      const mcdxbc::DescriptorTable &Table =
+          RSD.ParametersContainer.getDescriptorTable(Loc);
+
+      for (const dxbc::RTS0::v2::DescriptorRange &Range : Table.Ranges) {
+        Builder.trackBinding(RangeToResourceClass(Range.RangeType),
+                             Range.RegisterSpace, Range.BaseShaderRegister,
+                             Range.BaseShaderRegister + Range.NumDescriptors,
+                             nullptr);
+      }
+      break;
+    }
+    }
+  }
+}
+
+std::optional<mcdxbc::RootSignatureDesc>
+getRootSignature(RootSignatureBindingInfo &RSBI,
+                 dxil::ModuleMetadataInfo &MMI) {
+  if (MMI.EntryPropertyVec.size() == 0)
+    return std::nullopt;
+  std::optional<mcdxbc::RootSignatureDesc> RootSigDesc =
+      RSBI.getDescForFunction(MMI.EntryPropertyVec[0].Entry);
+  if (!RootSigDesc)
+    return std::nullopt;
+  return RootSigDesc;
+}
+
 static void reportErrors(Module &M, DXILResourceMap &DRM,
-                         DXILResourceBindingInfo &DRBI) {
+                         DXILResourceBindingInfo &DRBI,
+                         RootSignatureBindingInfo &RSBI,
+                         dxil::ModuleMetadataInfo &MMI) {
   if (DRM.hasInvalidCounterDirection())
     reportInvalidDirection(M, DRM);
 
@@ -94,14 +228,44 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
 
   assert(!DRBI.hasImplicitBinding() && "implicit bindings should be handled in "
                                        "DXILResourceImplicitBinding pass");
+
+  if (auto RSD = getRootSignature(RSBI, MMI)) {
+
+    hlsl::BindingInfoBuilder Builder;
+    dxbc::ShaderVisibility Visibility = tripleToVisibility(MMI.ShaderProfile);
+    trackRootSigDescBinding(Builder, *RSD, Visibility);
+
+    bool HasOverlap;
+    hlsl::BindingInfo Info = Builder.calculateBindingInfo(HasOverlap);
+
+    for (const auto &ResList :
+         {std::make_pair(ResourceClass::SRV, DRM.srvs()),
+          std::make_pair(ResourceClass::UAV, DRM.uavs()),
+          std::make_pair(ResourceClass::CBuffer, DRM.cbuffers()),
+          std::make_pair(ResourceClass::Sampler, DRM.samplers())}) {
+      for (auto Res : ResList.second) {
+        llvm::dxil::ResourceInfo::ResourceBinding ResBinding = Res.getBinding();
+        llvm::hlsl::BindingInfo::BindingRange ResRange(
+            ResBinding.LowerBound, ResBinding.LowerBound + ResBinding.Size);
+
+        auto IsBound = Info.isBound(ResList.first, ResBinding.Space, ResRange);
+        if (!IsBound) {
+          reportRegNotBound(M, ResList.first, ResBinding);
+        }
+      }
+    }
+  }
 }
 } // namespace
 
 PreservedAnalyses
 DXILPostOptimizationValidation::run(Module &M, ModuleAnalysisManager &MAM) {
   DXILResourceMap &DRM = MAM.getResult<DXILResourceAnalysis>(M);
   DXILResourceBindingInfo &DRBI = MAM.getResult<DXILResourceBindingAnalysis>(M);
-  reportErrors(M, DRM, DRBI);
+  RootSignatureBindingInfo &RSBI = MAM.getResult<RootSignatureAnalysis>(M);
+  ModuleMetadataInfo &MMI = MAM.getResult<DXILMetadataAnalysis>(M);
+
+  reportErrors(M, DRM, DRBI, RSBI, MMI);
   return PreservedAnalyses::all();
 }
 
@@ -113,7 +277,12 @@ class DXILPostOptimizationValidationLegacy : public ModulePass {
         getAnalysis<DXILResourceWrapperPass>().getResourceMap();
     DXILResourceBindingInfo &DRBI =
         getAnalysis<DXILResourceBindingWrapperPass>().getBindingInfo();
-    reportErrors(M, DRM, DRBI);
+    RootSignatureBindingInfo &RSBI =
+        getAnalysis<RootSignatureAnalysisWrapper>().getRSInfo();
+    dxil::ModuleMetadataInfo &MMI =
+        getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
+
+    reportErrors(M, DRM, DRBI, RSBI, MMI);
     return false;
   }
   StringRef getPassName() const override {
@@ -125,10 +294,13 @@ class DXILPostOptimizationValidationLegacy : public ModulePass {
   void getAnalysisUsage(llvm::AnalysisUsage &AU) const override {
     AU.addRequired<DXILResourceWrapperPass>();
     AU.addRequired<DXILResourceBindingWrapperPass>();
+    AU.addRequired<DXILMetadataAnalysisWrapperPass>();
+    AU.addRequired<RootSignatureAnalysisWrapper>();
     AU.addPreserved<DXILResourceWrapperPass>();
     AU.addPreserved<DXILResourceBindingWrapperPass>();
     AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
     AU.addPreserved<ShaderFlagsAnalysisWrapper>();
+    AU.addPreserved<RootSignatureAnalysisWrapper>();
   }
 };
 char DXILPostOptimizationValidationLegacy::ID = 0;
@@ -139,6 +311,8 @@ INITIALIZE_PASS_BEGIN(DXILPostOptimizationValidationLegacy, DEBUG_TYPE,
 INITIALIZE_PASS_DEPENDENCY(DXILResourceBindingWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DXILResourceTypeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DXILResourceWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DXILMetadataAnalysisWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(RootSignatureAnalysisWrapper)
 INITIALIZE_PASS_END(DXILPostOptimizationValidationLegacy, DEBUG_TYPE,
                     "DXIL Post Optimization Validation", false, false)
 

llvm/lib/Target/DirectX/DXILPrepare.cpp

@@ -300,7 +300,6 @@ class DXILPrepareModule : public ModulePass {
   DXILPrepareModule() : ModulePass(ID) {}
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<DXILMetadataAnalysisWrapperPass>();
-    AU.addRequired<RootSignatureAnalysisWrapper>();
     AU.addPreserved<RootSignatureAnalysisWrapper>();
     AU.addPreserved<ShaderFlagsAnalysisWrapper>();
     AU.addPreserved<DXILMetadataAnalysisWrapperPass>();

llvm/lib/Target/DirectX/DXILRootSignature.cpp

@@ -229,14 +229,15 @@ bool RootSignatureAnalysisWrapper::runOnModule(Module &M) {
 
 void RootSignatureAnalysisWrapper::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
-  AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
+  AU.addRequired<DXILMetadataAnalysisWrapperPass>();
 }
 
 char RootSignatureAnalysisWrapper::ID = 0;
 
 INITIALIZE_PASS_BEGIN(RootSignatureAnalysisWrapper,
                       "dxil-root-signature-analysis",
                       "DXIL Root Signature Analysis", true, true)
+INITIALIZE_PASS_DEPENDENCY(DXILMetadataAnalysisWrapperPass)
 INITIALIZE_PASS_END(RootSignatureAnalysisWrapper,
                     "dxil-root-signature-analysis",
                     "DXIL Root Signature Analysis", true, true)

llvm/lib/Target/DirectX/DXILRootSignature.h

@@ -96,4 +96,4 @@ class RootSignatureAnalysisPrinter
 
 } // namespace dxil
 } // namespace llvm
-#endif
+#endif // LLVM_LIB_TARGET_DIRECTX_DXILROOTSIGNATURE_H

llvm/test/CodeGen/DirectX/llc-pipeline.ll

@@ -33,9 +33,9 @@
 ; CHECK-NEXT:   DXIL Module Metadata analysis
 ; CHECK-NEXT:   DXIL Shader Flag Analysis
 ; CHECK-NEXT:   DXIL Translate Metadata
+; CHECK-NEXT:   DXIL Root Signature Analysis
 ; CHECK-NEXT:   DXIL Post Optimization Validation
 ; CHECK-NEXT:   DXIL Op Lowering
-; CHECK-NEXT:   DXIL Root Signature Analysis
 ; CHECK-NEXT:   DXIL Prepare Module
 
 ; CHECK-ASM-NEXT: DXIL Metadata Pretty Printer