[DirectX] Validate registers are bound to root signature

clang/lib/Sema/SemaHLSL.cpp

@@ -1146,12 +1146,12 @@ struct PerVisibilityBindingChecker {
     using llvm::hlsl::BindingInfoBuilder;
     auto ReportOverlap = [this, &HadOverlap](
                              const BindingInfoBuilder &Builder,
-                             const BindingInfoBuilder::Binding &Reported) {
+                             const llvm::hlsl::Binding &Reported) {
       HadOverlap = true;
 
       const auto *Elem =
           static_cast<const hlsl::RootSignatureElement *>(Reported.Cookie);
-      const BindingInfoBuilder::Binding &Previous =
+      const llvm::hlsl::Binding &Previous =
           Builder.findOverlapping(Reported);
       const auto *PrevElem =
           static_cast<const hlsl::RootSignatureElement *>(Previous.Cookie);

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

@@ -13,6 +13,7 @@
 #ifndef LLVM_FRONTEND_HLSL_HLSLBINDING_H
 #define LLVM_FRONTEND_HLSL_HLSLBINDING_H
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Compiler.h"
@@ -98,37 +99,52 @@ class BindingInfo {
 
   friend class BindingInfoBuilder;
 };
+struct Binding {
+  dxil::ResourceClass RC;
+  uint32_t Space;
+  uint32_t LowerBound;
+  uint32_t UpperBound;
+  const void *Cookie;
+
+  Binding(dxil::ResourceClass RC, uint32_t Space, uint32_t LowerBound,
+          uint32_t UpperBound, const void *Cookie)
+      : RC(RC), Space(Space), LowerBound(LowerBound), UpperBound(UpperBound),
+        Cookie(Cookie) {}
+
+  bool isUnbounded() const { return UpperBound == ~0U; }
+
+  bool operator==(const Binding &RHS) const {
+    return std::tie(RC, Space, LowerBound, UpperBound, Cookie) ==
+           std::tie(RHS.RC, RHS.Space, RHS.LowerBound, RHS.UpperBound,
+                    RHS.Cookie);
+  }
+  bool operator!=(const Binding &RHS) const { return !(*this == RHS); }
 
-/// Builder class for creating a /c BindingInfo.
-class BindingInfoBuilder {
-public:
-  struct Binding {
-    dxil::ResourceClass RC;
-    uint32_t Space;
-    uint32_t LowerBound;
-    uint32_t UpperBound;
-    const void *Cookie;
-
-    Binding(dxil::ResourceClass RC, uint32_t Space, uint32_t LowerBound,
-            uint32_t UpperBound, const void *Cookie)
-        : RC(RC), Space(Space), LowerBound(LowerBound), UpperBound(UpperBound),
-          Cookie(Cookie) {}
-
-    bool isUnbounded() const { return UpperBound == ~0U; }
-
-    bool operator==(const Binding &RHS) const {
-      return std::tie(RC, Space, LowerBound, UpperBound, Cookie) ==
-             std::tie(RHS.RC, RHS.Space, RHS.LowerBound, RHS.UpperBound,
-                      RHS.Cookie);
-    }
-    bool operator!=(const Binding &RHS) const { return !(*this == RHS); }
+  bool operator<(const Binding &RHS) const {
+    return std::tie(RC, Space, LowerBound) <
+           std::tie(RHS.RC, RHS.Space, RHS.LowerBound);
+  }
+};
+class BoundRegs {
+  SmallVector<Binding> Bindings;
 
-    bool operator<(const Binding &RHS) const {
-      return std::tie(RC, Space, LowerBound) <
-             std::tie(RHS.RC, RHS.Space, RHS.LowerBound);
-    }
-  };
+public:
+  BoundRegs(SmallVector<Binding> &&Bindings) : Bindings(std::move(Bindings)) {}
+
+  bool isBound(dxil::ResourceClass RC, uint32_t Space, uint32_t LowerBound,
+               uint32_t UpperBound) const {
+    const Binding *It =
+        llvm::upper_bound(Bindings, Binding{RC, Space, LowerBound, 0, nullptr});
+    if (It == Bindings.begin())
+      return false;
+    --It;
+    return It->RC == RC && It->Space == Space && It->LowerBound <= LowerBound &&
+           It->UpperBound >= UpperBound;
+  }
+};
 
+/// Builder class for creating a /c BindingInfo.
+class BindingInfoBuilder {
 private:
   SmallVector<Binding> Bindings;
 
@@ -152,6 +168,11 @@ class BindingInfoBuilder {
         [&HasOverlap](auto, auto) { HasOverlap = true; });
   }
 
+  LLVM_ABI BoundRegs calculateBoundRegs(
+      llvm::function_ref<void(const BindingInfoBuilder &Builder,
+                              const Binding &Overlapping)>
+            ReportOverlap);
+
   /// For use in the \c ReportOverlap callback of \c calculateBindingInfo -
   /// finds a binding that the \c ReportedBinding overlaps with.
   LLVM_ABI const Binding &findOverlapping(const Binding &ReportedBinding) const;

llvm/lib/Frontend/HLSL/HLSLBinding.cpp

@@ -131,9 +131,40 @@ BindingInfo BindingInfoBuilder::calculateBindingInfo(
   return Info;
 }
 
-const BindingInfoBuilder::Binding &BindingInfoBuilder::findOverlapping(
-    const BindingInfoBuilder::Binding &ReportedBinding) const {
-  for (const BindingInfoBuilder::Binding &Other : Bindings)
+BoundRegs BindingInfoBuilder::calculateBoundRegs(
+    llvm::function_ref<void(const BindingInfoBuilder &Builder,
+                            const Binding &Overlapping)>
+        ReportOverlap) {
+  // sort all the collected bindings
+  llvm::stable_sort(Bindings);
+
+  // remove duplicates
+  Binding *NewEnd = llvm::unique(Bindings);
+  if (NewEnd != Bindings.end())
+    Bindings.erase(NewEnd, Bindings.end());
+
+  if (Bindings.size() < 2)
+    return BoundRegs(std::move(Bindings));
+
+  for (auto Curr = Bindings.begin() + 1, End = Bindings.end(); Curr != End;
+       ++Curr) {
+    const Binding *Prev = Curr - 1;
+    if (Curr->Space != Prev->Space || Curr->RC != Prev->RC)
+      continue;
+
+    if (std::max(Curr->LowerBound, Prev->LowerBound) <=
+        std::min(Curr->UpperBound, Prev->UpperBound)) {
+      ReportOverlap(*this, *Curr);
+      continue;
+    }
+  }
+
+  return BoundRegs(std::move(Bindings));
+}
+
+const Binding &
+BindingInfoBuilder::findOverlapping(const Binding &ReportedBinding) const {
+  for (const Binding &Other : Bindings)
     if (ReportedBinding.LowerBound <= Other.UpperBound &&
         Other.LowerBound <= ReportedBinding.UpperBound)
       return Other;

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -118,10 +118,8 @@ static void reportOverlappingBinding(Module &M, DXILResourceMap &DRM) {
                        "true, yet no overlapping binding was found");
 }
 
-static void
-reportOverlappingRegisters(Module &M,
-                           const llvm::hlsl::BindingInfoBuilder::Binding &R1,
-                           const llvm::hlsl::BindingInfoBuilder::Binding &R2) {
+static void reportOverlappingRegisters(Module &M, const llvm::hlsl::Binding &R1,
+                                       const llvm::hlsl::Binding &R2) {
   SmallString<128> Message;
 
   raw_svector_ostream OS(Message);
@@ -133,6 +131,17 @@ reportOverlappingRegisters(Module &M,
   M.getContext().diagnose(DiagnosticInfoGeneric(Message));
 }
 
+static void
+reportRegNotBound(Module &M, ResourceClass Class,
+                  const llvm::dxil::ResourceInfo::ResourceBinding &Unbound) {
+  SmallString<128> Message;
+  raw_svector_ostream OS(Message);
+  OS << getResourceClassName(Class) << " register " << Unbound.LowerBound
+     << " in space " << Unbound.Space
+     << " does not have a binding in the Root Signature";
+  M.getContext().diagnose(DiagnosticInfoGeneric(Message));
+}
+
 static dxbc::ShaderVisibility
 tripleToVisibility(llvm::Triple::EnvironmentType ET) {
   switch (ET) {
@@ -157,7 +166,9 @@ tripleToVisibility(llvm::Triple::EnvironmentType ET) {
 
 static void validateRootSignature(Module &M,
                                   const mcdxbc::RootSignatureDesc &RSD,
-                                  dxil::ModuleMetadataInfo &MMI) {
+                                  dxil::ModuleMetadataInfo &MMI,
+                                  DXILResourceMap &DRM,
+                                  DXILResourceTypeMap &DRTM) {
 
   hlsl::BindingInfoBuilder Builder;
   dxbc::ShaderVisibility Visibility = tripleToVisibility(MMI.ShaderProfile);
@@ -214,13 +225,20 @@ static void validateRootSignature(Module &M,
     Builder.trackBinding(dxil::ResourceClass::Sampler, S.RegisterSpace,
                          S.ShaderRegister, S.ShaderRegister, &S);
 
-  Builder.calculateBindingInfo(
+  hlsl::BoundRegs Info = Builder.calculateBoundRegs(
       [&M](const llvm::hlsl::BindingInfoBuilder &Builder,
-           const llvm::hlsl::BindingInfoBuilder::Binding &ReportedBinding) {
-        const llvm::hlsl::BindingInfoBuilder::Binding &Overlaping =
+           const llvm::hlsl::Binding &ReportedBinding) {
+        const llvm::hlsl::Binding &Overlaping =
             Builder.findOverlapping(ReportedBinding);
         reportOverlappingRegisters(M, ReportedBinding, Overlaping);
       });
+  for (const ResourceInfo &RI : DRM) {
+    const ResourceInfo::ResourceBinding &Binding = RI.getBinding();
+    ResourceClass RC = DRTM[RI.getHandleTy()].getResourceClass();
+    if (!Info.isBound(RC, Binding.Space, Binding.LowerBound,
+                      Binding.LowerBound + Binding.Size - 1))
+      reportRegNotBound(M, RC, Binding);
+  }
 }
 
 static mcdxbc::RootSignatureDesc *
@@ -234,7 +252,8 @@ getRootSignature(RootSignatureBindingInfo &RSBI,
 static void reportErrors(Module &M, DXILResourceMap &DRM,
                          DXILResourceBindingInfo &DRBI,
                          RootSignatureBindingInfo &RSBI,
-                         dxil::ModuleMetadataInfo &MMI) {
+                         dxil::ModuleMetadataInfo &MMI,
+                         DXILResourceTypeMap &DRTM) {
   if (DRM.hasInvalidCounterDirection())
     reportInvalidDirection(M, DRM);
 
@@ -245,7 +264,7 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
                                        "DXILResourceImplicitBinding pass");
 
   if (mcdxbc::RootSignatureDesc *RSD = getRootSignature(RSBI, MMI))
-    validateRootSignature(M, *RSD, MMI);
+    validateRootSignature(M, *RSD, MMI, DRM, DRTM);
 }
 
 PreservedAnalyses
@@ -254,8 +273,9 @@ DXILPostOptimizationValidation::run(Module &M, ModuleAnalysisManager &MAM) {
   DXILResourceBindingInfo &DRBI = MAM.getResult<DXILResourceBindingAnalysis>(M);
   RootSignatureBindingInfo &RSBI = MAM.getResult<RootSignatureAnalysis>(M);
   ModuleMetadataInfo &MMI = MAM.getResult<DXILMetadataAnalysis>(M);
+  DXILResourceTypeMap &DRTM = MAM.getResult<DXILResourceTypeAnalysis>(M);
 
-  reportErrors(M, DRM, DRBI, RSBI, MMI);
+  reportErrors(M, DRM, DRBI, RSBI, MMI, DRTM);
   return PreservedAnalyses::all();
 }
 
@@ -271,8 +291,10 @@ class DXILPostOptimizationValidationLegacy : public ModulePass {
         getAnalysis<RootSignatureAnalysisWrapper>().getRSInfo();
     dxil::ModuleMetadataInfo &MMI =
         getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
+    DXILResourceTypeMap &DRTM =
+        getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
 
-    reportErrors(M, DRM, DRBI, RSBI, MMI);
+    reportErrors(M, DRM, DRBI, RSBI, MMI, DRTM);
     return false;
   }
   StringRef getPassName() const override {
@@ -286,11 +308,13 @@ class DXILPostOptimizationValidationLegacy : public ModulePass {
     AU.addRequired<DXILResourceBindingWrapperPass>();
     AU.addRequired<DXILMetadataAnalysisWrapperPass>();
     AU.addRequired<RootSignatureAnalysisWrapper>();
+    AU.addRequired<DXILResourceTypeWrapperPass>();
     AU.addPreserved<DXILResourceWrapperPass>();
     AU.addPreserved<DXILResourceBindingWrapperPass>();
     AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
     AU.addPreserved<ShaderFlagsAnalysisWrapper>();
     AU.addPreserved<RootSignatureAnalysisWrapper>();
+    AU.addPreserved<DXILResourceTypeWrapperPass>();
   }
 };
 char DXILPostOptimizationValidationLegacy::ID = 0;
@@ -303,6 +327,7 @@ INITIALIZE_PASS_DEPENDENCY(DXILResourceTypeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DXILResourceWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DXILMetadataAnalysisWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(RootSignatureAnalysisWrapper)
+INITIALIZE_PASS_DEPENDENCY(DXILResourceTypeWrapperPass)
 INITIALIZE_PASS_END(DXILPostOptimizationValidationLegacy, DEBUG_TYPE,
                     "DXIL Post Optimization Validation", false, false)
 

llvm/test/CodeGen/DirectX/rootsignature-validation-binding-limits-upperbound.ll

@@ -0,0 +1,20 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; This is a valid code, it checks the limits of a binding space
+; CHECK-NOT: error: .*
+
+%__cblayout_CB = type <{ float }>
+
+@CB.str = private unnamed_addr constant [3 x i8] c"CB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+  %CB = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 5, i32 0, i1 false, ptr nonnull @CB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2}
+!2 = !{!"DescriptorTable", i32 0, !3}
+!3 = !{!"CBV", i32 5, i32 0, i32 0, i32 0, i32 4}

llvm/test/CodeGen/DirectX/rootsignature-validation-binding-limits.ll

@@ -0,0 +1,22 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; This is a valid code, it checks the limits of a binding space
+
+; CHECK-NOT: error: .*
+
+%__cblayout_CB = type <{ float }>
+
+@CB.str = private unnamed_addr constant [3 x i8] c"CB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+  %CB = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 0, i32 4294967294, i32 1, i32 0, i1 false, ptr nonnull @CB.str)
+  %CB1 = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false, ptr nonnull @CB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !3}
+!2 = !{!"RootCBV", i32 0, i32 4294967294, i32 0, i32 4}
+!3 = !{!"RootCBV", i32 0, i32 0, i32 0, i32 4}

llvm/test/CodeGen/DirectX/rootsignature-validation-fail-cbv-binding.ll

@@ -0,0 +1,18 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; CHECK: error: CBV register 2 in space 666 does not have a binding in the Root Signature
+
+%__cblayout_CB = type <{ float }>
+
+@CB.str = private unnamed_addr constant [3 x i8] c"CB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+  %CB = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 666, i32 2, i32 1, i32 0, i1 false, ptr nonnull @CB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2}
+!2 = !{!"RootConstants", i32 0, i32 2, i32 0, i32 4}

llvm/test/CodeGen/DirectX/rootsignature-validation-fail-consecutive-ranges.ll

@@ -0,0 +1,19 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; CHECK: CBV register 3 in space 0 does not have a binding in the Root Signature
+%__cblayout_CB = type <{ float }>
+
+@CB.str = private unnamed_addr constant [3 x i8] c"CB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+  %CB = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 0, i32 3, i32 6, i32 0, i1 false, ptr nonnull @CB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2}
+!2 = !{!"DescriptorTable", i32 0, !3, !4}
+!3 = !{!"CBV", i32 5, i32 2, i32 0, i32 0, i32 4}
+!4 = !{!"CBV", i32 4, i32 7, i32 0, i32 10, i32 4}

llvm/test/CodeGen/DirectX/rootsignature-validation-fail-sampler-binding.ll

@@ -0,0 +1,18 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; CHECK: error: Sampler register 3 in space 2 does not have a binding in the Root Signature 
+
+@Smp.str = private unnamed_addr constant [4 x i8] c"Smp\00", align 1
+
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+  %Sampler = call target("dx.Sampler", 0) @llvm.dx.resource.handlefrombinding(i32 2, i32 3, i32 1, i32 0, i1 false, ptr nonnull @Smp.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2}
+!2 = !{!"DescriptorTable", i32 0, !3}
+!3 = !{!"Sampler", i32 1, i32 42, i32 0, i32 -1, i32 0}