[DirectX] Validate registers are bound to root signature

clang/lib/Sema/SemaHLSL.cpp

@@ -1144,15 +1144,14 @@ struct PerVisibilityBindingChecker {
     bool HadOverlap = false;
 
     using llvm::hlsl::BindingInfoBuilder;
-    auto ReportOverlap = [this, &HadOverlap](
-                             const BindingInfoBuilder &Builder,
-                             const BindingInfoBuilder::Binding &Reported) {
+    auto ReportOverlap = [this,
+                          &HadOverlap](const BindingInfoBuilder &Builder,
+                                       const llvm::hlsl::Binding &Reported) {
       HadOverlap = true;
 
       const auto *Elem =
           static_cast<const hlsl::RootSignatureElement *>(Reported.Cookie);
-      const BindingInfoBuilder::Binding &Previous =
-          Builder.findOverlapping(Reported);
+      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"
@@ -99,36 +100,55 @@ class BindingInfo {
   friend class BindingInfoBuilder;
 };
 
-/// 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;
+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) {}
+  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 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, 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);
-    }
-  };
+  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;
 
+public:
+  BoundRegs(SmallVector<Binding> &&Bindings) : Bindings(std::move(Bindings)) {}
+
+  bool isBound(dxil::ResourceClass RC, uint32_t Space, uint32_t LowerBound,
+               uint32_t UpperBound) const {
+    // UpperBound and Cookie are given dummy values, since they aren't
+    // interesting for operator<
+    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 +172,12 @@ class BindingInfoBuilder {
         [&HasOverlap](auto, auto) { HasOverlap = true; });
   }
 
+  LLVM_ABI BoundRegs getBoundRegs() {
+    assert(std::is_sorted(Bindings.begin(), Bindings.end()) &&
+           "Bindings must be sorted");
+    return BoundRegs(std::move(Bindings));
+  }
+
   /// 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,9 @@ BindingInfo BindingInfoBuilder::calculateBindingInfo(
   return Info;
 }
 
-const BindingInfoBuilder::Binding &BindingInfoBuilder::findOverlapping(
-    const BindingInfoBuilder::Binding &ReportedBinding) const {
-  for (const BindingInfoBuilder::Binding &Other : 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);
@@ -216,11 +227,19 @@ static void validateRootSignature(Module &M,
 
   Builder.calculateBindingInfo(
       [&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);
       });
+  const hlsl::BoundRegs &BoundRegs = Builder.getBoundRegs();
+  for (const ResourceInfo &RI : DRM) {
+    const ResourceInfo::ResourceBinding &Binding = RI.getBinding();
+    ResourceClass RC = DRTM[RI.getHandleTy()].getResourceClass();
+    if (!BoundRegs.isBound(RC, Binding.Space, Binding.LowerBound,
+                           Binding.LowerBound + Binding.Size - 1))
+      reportRegNotBound(M, RC, Binding);
+  }
 }
 
 static mcdxbc::RootSignatureDesc *
@@ -234,7 +253,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 +265,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 +274,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 +292,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,6 +309,7 @@ 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>();
@@ -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}

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

@@ -0,0 +1,23 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; CHECK: error: SRV register 0 in space 0 does not have a binding in the Root Signature
+
+@SB.str = private unnamed_addr constant [3 x i8] c"SB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+; StructuredBuffer<int> In : register(t0, space0);
+  %SB = tail call target("dx.RawBuffer", i32, 0, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i32_0_0t(i32 0, i32 0, i32 1, i32 0, i1 false, ptr nonnull @SB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !3, !5, !7}
+!2 = !{!"RootCBV", i32 0, i32 3, i32 666, i32 4}
+!3 = !{!"DescriptorTable", i32 1, !4}
+!4 = !{!"SRV", i32 1, i32 0, i32 0, i32 -1, i32 4}
+!5 = !{!"DescriptorTable", i32 0, !6}
+!6 = !{!"Sampler", i32 2, i32 0, i32 0, i32 -1, i32 0}
+!7 = !{!"DescriptorTable", i32 0, !8}
+!8 = !{!"UAV", i32 -1, i32 0, i32 0, i32 -1, i32 2}