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

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -10,6 +10,8 @@
 #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/Analysis/DXILMetadataAnalysis.h"
 #include "llvm/Analysis/DXILResource.h"
@@ -128,6 +130,17 @@ static void reportOverlappingBinding(Module &M, DXILResourceMap &DRM) {
   }
 }
 
+static void
+reportInvalidHandleTyBoundInRs(Module &M, Twine Type,
+                               ResourceInfo::ResourceBinding Binding) {
+  SmallString<128> Message;
+  raw_svector_ostream OS(Message);
+  OS << "resource " << Type << " at register (space=" << Binding.Space
+     << ", register=" << Binding.LowerBound << ")"
+     << " is bound to a texture or typed buffer.";
+  M.getContext().diagnose(DiagnosticInfoGeneric(Message));
+}
+
 static void reportRegNotBound(Module &M,
                               llvm::hlsl::rootsig::RangeInfo Unbound) {
   SmallString<128> Message;
@@ -232,6 +245,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) {
@@ -244,6 +295,32 @@ getRootSignature(RootSignatureBindingInfo &RSBI,
   return RootSigDesc;
 }
 
+static void reportInvalidHandleTy(
+    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);
+
+    if (TypedBuffer != nullptr || Texture != nullptr)
+      reportInvalidHandleTyBoundInRs(M, Res->getName(), Res->getBinding());
+  }
+}
+
 static void reportUnboundRegisters(
     Module &M,
     const llvm::hlsl::rootsig::RootSignatureBindingValidation &Validation,
@@ -282,16 +359,24 @@ 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());
     reportUnboundRegisters(M, Validation, ResourceClass::UAV, DRM.uavs());
     reportUnboundRegisters(M, Validation, ResourceClass::Sampler,
                            DRM.samplers());
     reportUnboundRegisters(M, Validation, ResourceClass::SRV, DRM.srvs());
+
+    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

@@ -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: resource TB at register (space=0, register=0) is bound to a texture or typed buffer.
+
+
+; Root Signature(
+;   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 = !{!2}
+!2 = !{!"RootUAV", i32 0, i32 0, i32 0, i32 4}

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

@@ -0,0 +1,21 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 
+; expected-no-diagnostics
+; 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

@@ -0,0 +1,25 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; CHECK: error: resource TB at register (space=0, register=0) is bound to a texture or typed buffer.
+
+
+; 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)
+
+@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 = !{!2}
+!2 = !{!"RootUAV", i32 0, i32 0, i32 0, i32 4}

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

@@ -0,0 +1,21 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 
+; expected-no-diagnostics
+; 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}

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

@@ -21,7 +21,7 @@ entry:
   %CB = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 1, i32 3, i32 1, i32 0, i1 false, ptr nonnull @CB.str)
   %Sampler = call target("dx.Sampler", 0) @llvm.dx.resource.handlefrombinding(i32 2, i32 3, i32 1, i32 0, i1 false, ptr nonnull @Smp.str)
   %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)
-  %RWB =  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 @RWB.str)
+  %RWB =  tail call target("dx.TypedBuffer", float, 1, 0, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0_0t(i32 0, i32 0, i32 1, i32 0, i1 false, ptr nonnull @RWB.str)
 
   ret void
 }