[DirectX] Validate registers are bound to root signature

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/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/DXILPostOptimizationValidation.cpp

@@ -130,6 +130,18 @@ static void reportOverlappingRegisters(
   M.getContext().diagnose(DiagnosticInfoGeneric(Message));
 }
 
+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) {
   switch (ET) {
@@ -154,7 +166,8 @@ tripleToVisibility(llvm::Triple::EnvironmentType ET) {
 
 static void validateRootSignature(Module &M,
                                   const mcdxbc::RootSignatureDesc &RSD,
-                                  dxil::ModuleMetadataInfo &MMI) {
+                                  dxil::ModuleMetadataInfo &MMI,
+                                  DXILResourceMap &DRM) {
 
   hlsl::BindingInfoBuilder Builder;
   dxbc::ShaderVisibility Visibility = tripleToVisibility(MMI.ShaderProfile);
@@ -213,14 +226,33 @@ static void validateRootSignature(Module &M,
     Builder.trackBinding(dxil::ResourceClass::Sampler, S.RegisterSpace,
                          S.ShaderRegister, S.ShaderRegister,
                          &IDs.emplace_back());
-
+  bool HasOverlap = false;
   hlsl::BindingInfo Info = Builder.calculateBindingInfo(
-      [&M](const llvm::hlsl::BindingInfoBuilder &Builder,
-           const llvm::hlsl::BindingInfoBuilder::Binding &ReportedBinding) {
+      [&M, &HasOverlap](
+          const llvm::hlsl::BindingInfoBuilder &Builder,
+          const llvm::hlsl::BindingInfoBuilder::Binding &ReportedBinding) {
+        HasOverlap = true;
         const llvm::hlsl::BindingInfoBuilder::Binding &Overlaping =
             Builder.findOverlapping(ReportedBinding);
         reportOverlappingRegisters(M, ReportedBinding, Overlaping);
       });
+  // Next checks require that the root signature definition is valid.
+  if (!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);
+
+        if (!Info.isBound(ResList.first, ResBinding.Space, ResRange))
+          reportRegNotBound(M, ResList.first, ResBinding);
+      }
+    }
+  }
 }
 
 static mcdxbc::RootSignatureDesc *
@@ -245,7 +277,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);
 }
 
 PreservedAnalyses

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

@@ -0,0 +1,34 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+
+; CHECK: error: register CBuffer (space=665, register=3) does not have a binding in the Root Signature
+
+; Root Signature(
+;   CBV(b3, space=666, visibility=SHADER_VISIBILITY_ALL)
+;   DescriptorTable(SRV(t0, space=0, numDescriptors=1), visibility=SHADER_VISIBILITY_ALL)
+;   DescriptorTable(Sampler(s0, numDescriptors=2), visibility=SHADER_VISIBILITY_VERTEX)
+;   DescriptorTable(UAV(u0, numDescriptors=unbounded), visibility=SHADER_VISIBILITY_ALL)
+
+%__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:
+; cbuffer CB : register(b3, space665) {
+;  float a;
+; }
+  %CB = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 665, i32 3, i32 1, i32 0, i1 false, ptr nonnull @CB.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}

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

@@ -0,0 +1,22 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; CHECK: error: register CBuffer (space=666, register=2) does not have a binding in the Root Signature
+; Root Signature(RootConstants(num32BitConstants=4, b2))
+
+%__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:
+; cbuffer CB : register(b2, space666) {
+;  float a;
+; }
+  %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-sampler.ll

@@ -0,0 +1,31 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+
+; CHECK: error: register Sampler (space=2, register=3) does not have a binding in the Root Signature
+
+; Root Signature(
+;   CBV(b3, space=666, visibility=SHADER_VISIBILITY_ALL)
+;   DescriptorTable(SRV(t0, space=0, numDescriptors=1), visibility=SHADER_VISIBILITY_VERTEX)
+;   DescriptorTable(Sampler(s0, numDescriptors=2), visibility=SHADER_VISIBILITY_ALL)
+;   DescriptorTable(UAV(u0, numDescriptors=unbounded), visibility=SHADER_VISIBILITY_ALL)
+
+@Smp.str = private unnamed_addr constant [4 x i8] c"Smp\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+; SamplerState S1 : register(s3, space2);
+  %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, !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}

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

@@ -0,0 +1,30 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+
+; CHECK: error: register SRV (space=0, register=0) does not have a binding in the Root Signature
+
+; Root Signature(
+;   CBV(b3, space=666, visibility=SHADER_VISIBILITY_ALL)
+;   DescriptorTable(SRV(t0, space=0, numDescriptors=1), visibility=SHADER_VISIBILITY_VERTEX)
+;   DescriptorTable(Sampler(s0, numDescriptors=2), visibility=SHADER_VISIBILITY_ALL)
+;   DescriptorTable(UAV(u0, numDescriptors=unbounded), visibility=SHADER_VISIBILITY_ALL)
+
+@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}

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

@@ -0,0 +1,30 @@
+; 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=4294967294) does not have a binding in the Root Signature
+
+; Root Signature(
+;   CBV(b3, space=666, visibility=SHADER_VISIBILITY_ALL)
+;   DescriptorTable(SRV(t0, space=0, numDescriptors=1), visibility=SHADER_VISIBILITY_VERTEX)
+;   DescriptorTable(Sampler(s0, numDescriptors=2), visibility=SHADER_VISIBILITY_ALL)
+;   DescriptorTable(UAV(u0, numDescriptors=unbounded), visibility=SHADER_VISIBILITY_ALL)
+
+@RWB.str = private unnamed_addr constant [4 x i8] c"RWB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+; 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
+}
+
+!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 10, i32 0, i32 0, i32 -1, i32 2}