[DirectX] Validate if Textures/TypedBuffers are being bound in Root Signatures

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -10,6 +10,7 @@
 #include "DXILRootSignature.h"
 #include "DXILShaderFlags.h"
 #include "DirectX.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLForwardCompat.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
@@ -248,6 +249,44 @@ initRSBindingValidation(const mcdxbc::RootSignatureDesc &RSD,
   return Validation;
 }
 
+static SmallVector<ResourceInfo::ResourceBinding>
+getRootDescriptorsBindingInfo(const mcdxbc::RootSignatureDesc &RSD,
+                              dxbc::ShaderVisibility Visibility) {
+
+  SmallVector<ResourceInfo::ResourceBinding> RDs;
+
+  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::SRV):
+    case llvm::to_underlying(dxbc::RootParameterType::UAV):
+    case llvm::to_underlying(dxbc::RootParameterType::CBV): {
+      dxbc::RTS0::v2::RootDescriptor Desc =
+          RSD.ParametersContainer.getRootDescriptor(Loc);
+
+      ResourceInfo::ResourceBinding Binding;
+      Binding.LowerBound = Desc.ShaderRegister;
+      Binding.Space = Desc.RegisterSpace;
+      Binding.Size = 1;
+
+      RDs.push_back(Binding);
+      break;
+    }
+    }
+  }
+
+  return RDs;
+}
+
 std::optional<mcdxbc::RootSignatureDesc>
 getRootSignature(RootSignatureBindingInfo &RSBI,
                  dxil::ModuleMetadataInfo &MMI) {
@@ -261,10 +300,22 @@ getRootSignature(RootSignatureBindingInfo &RSBI,
 }
 
 static void reportInvalidHandleTy(
-    Module &M,
+    Module &M, const llvm::ArrayRef<dxil::ResourceInfo::ResourceBinding> &RDs,
     const iterator_range<SmallVectorImpl<dxil::ResourceInfo>::iterator>
         &Resources) {
   for (auto Res = Resources.begin(), End = Resources.end(); Res != End; Res++) {
+    llvm::dxil::ResourceInfo::ResourceBinding Binding = Res->getBinding();
+    bool IsBound = false;
+    for (const auto &RD : RDs) {
+      if (Binding.overlapsWith(RD)) {
+        IsBound = true;
+        break;
+      }
+    }
+
+    if (!IsBound)
+      continue;
+
     TargetExtType *Handle = Res->getHandleTy();
     auto *TypedBuffer = dyn_cast_or_null<TypedBufferExtType>(Handle);
     auto *Texture = dyn_cast_or_null<TextureExtType>(Handle);
@@ -312,9 +363,9 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
                                        "DXILResourceImplicitBinding pass");
 
   if (auto RSD = getRootSignature(RSBI, MMI)) {
-
+    dxbc::ShaderVisibility Visibility = tripleToVisibility(MMI.ShaderProfile);
     llvm::hlsl::rootsig::RootSignatureBindingValidation Validation =
-        initRSBindingValidation(*RSD, tripleToVisibility(MMI.ShaderProfile));
+        initRSBindingValidation(*RSD, Visibility);
 
     reportUnboundRegisters(M, Validation, ResourceClass::CBuffer,
                            DRM.cbuffers());
@@ -323,10 +374,13 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
                            DRM.samplers());
     reportUnboundRegisters(M, Validation, ResourceClass::SRV, DRM.srvs());
 
-    reportInvalidHandleTy(M, DRM.cbuffers());
-    reportInvalidHandleTy(M, DRM.srvs());
-    reportInvalidHandleTy(M, DRM.uavs());
-    reportInvalidHandleTy(M, DRM.samplers());
+    SmallVector<ResourceInfo::ResourceBinding> RDs =
+        getRootDescriptorsBindingInfo(*RSD, Visibility);
+
+    reportInvalidHandleTy(M, RDs, DRM.cbuffers());
+    reportInvalidHandleTy(M, RDs, DRM.srvs());
+    reportInvalidHandleTy(M, RDs, DRM.uavs());
+    reportInvalidHandleTy(M, RDs, DRM.samplers());
   }
 }
 } // namespace

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

@@ -3,10 +3,7 @@
 
 
 ; Root Signature(
-;   CBV(b3, space=1, 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)
+;   UAV(b0, space=0, visibility=SHADER_VISIBILITY_ALL)
 
 @TB.str = private unnamed_addr constant [3 x i8] c"TB\00", align 1
 
@@ -21,11 +18,5 @@ entry:
 !dx.rootsignatures = !{!0}
 
 !0 = !{ptr @CSMain, !1, i32 2}
-!1 = !{!2, !3, !5, !7}
-!2 = !{!"RootCBV", i32 0, i32 3, i32 1, i32 4}
-!3 = !{!"DescriptorTable", i32 0, !4}
-!4 = !{!"SRV", i32 1, i32 0, i32 0, i32 -1, i32 0}
-!5 = !{!"DescriptorTable", i32 0, !6}
-!6 = !{!"Sampler", i32 5, i32 3, i32 2, i32 -1, i32 0}
-!7 = !{!"DescriptorTable", i32 0, !8}
-!8 = !{!"UAV", i32 -1, i32 0, i32 0, i32 -1, i32 2}
+!1 = !{!2}
+!2 = !{!"RootUAV", i32 0, i32 0, i32 0, i32 4}

llvm/test/CodeGen/DirectX/rootsignature-validation-textures.ll

@@ -0,0 +1,23 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 
+
+
+
+; Root Signature(
+;   DescriptorTable(UAV(b0, space=0, visibility=SHADER_VISIBILITY_ALL))
+
+@TB.str = private unnamed_addr constant [3 x i8] c"TB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+
+  %TB =  tail call target("dx.Texture", float, 1, 0, 0) @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0_0t(i32 0, i32 0, i32 1, i32 0, i1 false, ptr nonnull @TB.str)
+
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!3}
+!3 = !{!"DescriptorTable", i32 0, !4}
+!4 = !{!"UAV", i32 1, i32 0, i32 0, i32 -1, i32 0}

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

@@ -13,19 +13,13 @@
 define void @CSMain() "hlsl.shader"="compute" {
 entry:
 
-  %TB =  tail call target("dx.Texture", float, 1, 0, 0) @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0_0t(i32 0, i32 0, i32 1, i32 0, i1 false, ptr nonnull @TB.str)
+  %TB =  tail call target("dx.TypedBuffer", float, 1, 0, 0) @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0_0t(i32 0, i32 0, i32 1, i32 0, i1 false, ptr nonnull @TB.str)
 
   ret void
 }
 
 !dx.rootsignatures = !{!0}
 
 !0 = !{ptr @CSMain, !1, i32 2}
-!1 = !{!2, !3, !5, !7}
-!2 = !{!"RootCBV", i32 0, i32 3, i32 1, i32 4}
-!3 = !{!"DescriptorTable", i32 0, !4}
-!4 = !{!"SRV", i32 1, i32 0, i32 0, i32 -1, i32 0}
-!5 = !{!"DescriptorTable", i32 0, !6}
-!6 = !{!"Sampler", i32 5, i32 3, i32 2, i32 -1, i32 0}
-!7 = !{!"DescriptorTable", i32 0, !8}
-!8 = !{!"UAV", i32 -1, i32 0, i32 0, i32 -1, i32 2}
+!1 = !{!2}
+!2 = !{!"RootUAV", i32 0, i32 0, i32 0, i32 4}

llvm/test/CodeGen/DirectX/rootsignature-validation-typedbuffer.ll

@@ -0,0 +1,23 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 
+
+
+
+; Root Signature(
+;   DescriptorTable(UAV(b0, space=0, visibility=SHADER_VISIBILITY_ALL))
+
+@TB.str = private unnamed_addr constant [3 x i8] c"TB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+
+  %TB =  tail call target("dx.TypedBuffer", float, 1, 0, 0) @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0_0t(i32 0, i32 0, i32 1, i32 0, i1 false, ptr nonnull @TB.str)
+
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!3}
+!3 = !{!"DescriptorTable", i32 0, !4}
+!4 = !{!"UAV", i32 1, i32 0, i32 0, i32 -1, i32 0}