[DirectX] Validate registers are bound to root signature

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

@@ -17,6 +17,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/DXILABI.h"
 #include "llvm/Support/ErrorHandling.h"
+#include <cstdint>
 
 namespace llvm {
 namespace hlsl {
@@ -57,6 +58,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 +97,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/include/llvm/Frontend/HLSL/RootSignatureValidations.h

@@ -148,51 +148,6 @@ struct OverlappingRanges {
 ///      B: Check for overlap with any overlapping Visibility ResourceRange
 LLVM_ABI llvm::SmallVector<OverlappingRanges>
 findOverlappingRanges(ArrayRef<RangeInfo> Infos);
-
-class RootSignatureBindingValidation {
-private:
-  llvm::SmallVector<RangeInfo, 16> Bindings;
-  struct TypeRange {
-    uint32_t Start;
-    uint32_t End;
-  };
-  std::unordered_map<dxil::ResourceClass, TypeRange> Ranges;
-
-public:
-  void addBinding(dxil::ResourceClass Type, const RangeInfo &Binding) {
-    auto It = Ranges.find(Type);
-
-    if (It == Ranges.end()) {
-      uint32_t InsertPos = Bindings.size();
-      Bindings.push_back(Binding);
-      Ranges[Type] = {InsertPos, InsertPos + 1};
-      return;
-    } 
-    uint32_t InsertPos = It->second.End;
-    Bindings.insert(Bindings.begin() + InsertPos, Binding);
-
-    It->second.End++;
-
-    for (auto &[Type, Range] : Ranges) {
-      if (Range.Start > InsertPos) {
-        Range.Start++;
-        Range.End++;
-      }
-    }
-  }
-
-  llvm::ArrayRef<RangeInfo>
-  getBindingsOfType(const dxil::ResourceClass &Type) const {
-    auto It = Ranges.find(Type);
-    if (It == Ranges.end()) 
-      return {};
-    return llvm::ArrayRef<RangeInfo>(Bindings.data() + It->second.Start,
-                                     It->second.End - It->second.Start);
-  }
-};
-llvm::SmallVector<RangeInfo>
-findUnboundRanges(const llvm::ArrayRef<RangeInfo> &Ranges,
-                  const llvm::ArrayRef<RangeInfo> &Bindings);
 } // namespace rootsig
 } // namespace hlsl
 } // namespace llvm

llvm/lib/Frontend/HLSL/HLSLBinding.cpp

@@ -8,6 +8,7 @@
 
 #include "llvm/Frontend/HLSL/HLSLBinding.h"
 #include "llvm/ADT/STLExtras.h"
+#include <cstdint>
 
 using namespace llvm;
 using namespace hlsl;
@@ -66,6 +67,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/DXILPostOptimizationValidation.cpp

@@ -14,12 +14,16 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/DXILMetadataAnalysis.h"
 #include "llvm/Analysis/DXILResource.h"
+#include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/Frontend/HLSL/HLSLBinding.h"
 #include "llvm/Frontend/HLSL/RootSignatureValidations.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicsDirectX.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
+#include "llvm/Support/DXILABI.h"
+#include <utility>
 
 #define DEBUG_TYPE "dxil-post-optimization-validation"
 
@@ -116,11 +120,12 @@ static void reportOverlappingBinding(Module &M, DXILResourceMap &DRM) {
   }
 }
 
-static void reportRegNotBound(Module &M,
-                              llvm::hlsl::rootsig::RangeInfo Unbound) {
+static void
+reportRegNotBound(Module &M, ResourceClass Class,
+                  llvm::dxil::ResourceInfo::ResourceBinding Unbound) {
   SmallString<128> Message;
   raw_svector_ostream OS(Message);
-  OS << "register " << getResourceClassName(Unbound.Class)
+  OS << "register " << getResourceClassName(Class)
      << " (space=" << Unbound.Space << ", register=" << Unbound.LowerBound
      << ")"
      << " does not have a binding in the Root Signature";
@@ -155,12 +160,9 @@ tripleToVisibility(llvm::Triple::EnvironmentType ET) {
   }
 }
 
-static hlsl::rootsig::RootSignatureBindingValidation
-initRSBindingValidation(const mcdxbc::RootSignatureDesc &RSD,
-                        dxbc::ShaderVisibility Visibility) {
-
-  hlsl::rootsig::RootSignatureBindingValidation Validation;
-
+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);
@@ -175,15 +177,10 @@ initRSBindingValidation(const mcdxbc::RootSignatureDesc &RSD,
     case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
       dxbc::RTS0::v1::RootConstants Const =
           RSD.ParametersContainer.getConstant(Loc);
-
-      hlsl::rootsig::RangeInfo Binding;
-      Binding.LowerBound = Const.ShaderRegister;
-      Binding.Space = Const.RegisterSpace;
-      Binding.UpperBound = Binding.LowerBound;
-
-      // Root Constants Bind to CBuffers
-      Validation.addBinding(ResourceClass::CBuffer, Binding);
-
+      Builder.trackBinding(dxil::ResourceClass::CBuffer, Const.RegisterSpace,
+                           Const.ShaderRegister,
+                           Const.ShaderRegister + Const.Num32BitValues,
+                           nullptr);
       break;
     }
 
@@ -192,32 +189,25 @@ initRSBindingValidation(const mcdxbc::RootSignatureDesc &RSD,
     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);
 
-      hlsl::rootsig::RangeInfo Binding;
-      Binding.LowerBound = Desc.ShaderRegister;
-      Binding.Space = Desc.RegisterSpace;
-      Binding.UpperBound = Binding.LowerBound;
-
-      Validation.addBinding(ParameterToResourceClass(Type), Binding);
       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) {
-        hlsl::rootsig::RangeInfo Binding;
-        Binding.LowerBound = Range.BaseShaderRegister;
-        Binding.Space = Range.RegisterSpace;
-        Binding.UpperBound = Binding.LowerBound + Range.NumDescriptors - 1;
-        Validation.addBinding(RangeToResourceClass(Range.RangeType), Binding);
+        Builder.trackBinding(RangeToResourceClass(Range.RangeType),
+                             Range.RegisterSpace, Range.BaseShaderRegister,
+                             Range.BaseShaderRegister + Range.NumDescriptors,
+                             nullptr);
       }
       break;
     }
     }
   }
-
-  return Validation;
 }
 
 std::optional<mcdxbc::RootSignatureDesc>
@@ -232,30 +222,6 @@ getRootSignature(RootSignatureBindingInfo &RSBI,
   return RootSigDesc;
 }
 
-static void reportUnboundRegisters(
-    Module &M,
-    const llvm::hlsl::rootsig::RootSignatureBindingValidation &Validation,
-    ResourceClass Class,
-    const iterator_range<SmallVectorImpl<dxil::ResourceInfo>::iterator>
-        &Resources) {
-  SmallVector<hlsl::rootsig::RangeInfo> Ranges;
-  for (auto Res = Resources.begin(), End = Resources.end(); Res != End; Res++) {
-    ResourceInfo::ResourceBinding ResBinding = Res->getBinding();
-    hlsl::rootsig::RangeInfo Range;
-    Range.Space = ResBinding.Space;
-    Range.LowerBound = ResBinding.LowerBound;
-    Range.UpperBound = Range.LowerBound + ResBinding.Size - 1;
-    Range.Class = Class;
-    Ranges.push_back(Range);
-  }
-
-  SmallVector<hlsl::rootsig::RangeInfo> Unbounds =
-      hlsl::rootsig::findUnboundRanges(Ranges,
-                                       Validation.getBindingsOfType(Class));
-  for (const auto &Unbound : Unbounds)
-    reportRegNotBound(M, Unbound);
-}
-
 static void reportErrors(Module &M, DXILResourceMap &DRM,
                          DXILResourceBindingInfo &DRBI,
                          RootSignatureBindingInfo &RSBI,
@@ -271,15 +237,29 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
 
   if (auto RSD = getRootSignature(RSBI, MMI)) {
 
-    llvm::hlsl::rootsig::RootSignatureBindingValidation Validation =
-        initRSBindingValidation(*RSD, tripleToVisibility(MMI.ShaderProfile));
-
-    reportUnboundRegisters(M, Validation, ResourceClass::CBuffer,
-                           DRM.cbuffers());
-    reportUnboundRegisters(M, Validation, ResourceClass::UAV, DRM.uavs());
-    reportUnboundRegisters(M, Validation, ResourceClass::Sampler,
-                           DRM.samplers());
-    reportUnboundRegisters(M, Validation, ResourceClass::SRV, DRM.srvs());
+    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

llvm/test/CodeGen/DirectX/rootsignature-validation-fail-uav.ll

@@ -1,6 +1,6 @@
 ; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
 
-; CHECK: error: register UAV (space=0, register=4294967295) does not have a binding in the Root Signature
+; CHECK: error: register UAV (space=0, register=4294967294) does not have a binding in the Root Signature
 
 ; Root Signature(
 ;   CBV(b3, space=666, visibility=SHADER_VISIBILITY_ALL)
@@ -12,8 +12,8 @@
 
 define void @CSMain() "hlsl.shader"="compute" {
 entry:
-; RWBuffer<float> UAV : register(u4294967295);
-  %RWB = tail call target("dx.TypedBuffer", float, 1, 0, 0) @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0_0t(i32 0, i32 -1, i32 1, i32 0, i1 false, ptr nonnull @RWB.str)
+; RWBuffer<float> UAV : register(4294967294);
+  %RWB = tail call target("dx.TypedBuffer", float, 1, 0, 0) @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0_0t(i32 0, i32 4294967294, i32 1, i32 0, i1 false, ptr nonnull @RWB.str)
   ret void
 }
 
@@ -27,4 +27,4 @@ entry:
 !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}
+!8 = !{!"UAV", i32 10, i32 0, i32 0, i32 -1, i32 2}