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

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

@@ -133,17 +133,19 @@ class BoundRegs {
 public:
   BoundRegs(SmallVector<Binding> &&Bindings) : Bindings(std::move(Bindings)) {}
 
-  bool isBound(dxil::ResourceClass RC, uint32_t Space, uint32_t LowerBound,
-               uint32_t UpperBound) const {
+  const Binding *findBoundReg(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;
+      return nullptr;
     --It;
-    return It->RC == RC && It->Space == Space && It->LowerBound <= LowerBound &&
-           It->UpperBound >= UpperBound;
+    if (It->RC == RC && It->Space == Space && It->LowerBound <= LowerBound &&
+        It->UpperBound >= UpperBound)
+      return It;
+    return nullptr;
   }
 };
 

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -103,6 +103,17 @@ static void reportOverlappingBinding(Module &M, DXILResourceMap &DRM) {
                        "true, yet no overlapping binding was found");
 }
 
+static void reportInvalidHandleTyError(Module &M, ResourceClass RC,
+                                       ResourceInfo::ResourceBinding Binding) {
+  SmallString<160> Message;
+  raw_svector_ostream OS(Message);
+  StringRef RCName = getResourceClassName(RC);
+  OS << RCName << " at register " << Binding.LowerBound << " and space "
+     << Binding.Space << " is bound to a texture or typed buffer. " << RCName
+     << " root descriptors can only be Raw or Structured buffers.";
+  M.getContext().diagnose(DiagnosticInfoGeneric(Message));
+}
+
 static void reportOverlappingRegisters(Module &M, const llvm::hlsl::Binding &R1,
                                        const llvm::hlsl::Binding &R2) {
   SmallString<128> Message;
@@ -217,10 +228,29 @@ static void validateRootSignature(Module &M,
   const hlsl::BoundRegs &BoundRegs = Builder.takeBoundRegs();
   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))
+    const dxil::ResourceTypeInfo &RTI = DRTM[RI.getHandleTy()];
+    dxil::ResourceClass RC = RTI.getResourceClass();
+    dxil::ResourceKind RK = RTI.getResourceKind();
+
+    const llvm::hlsl::Binding *Reg =
+        BoundRegs.findBoundReg(RC, Binding.Space, Binding.LowerBound,
+                               Binding.LowerBound + Binding.Size - 1);
+
+    if (Reg != nullptr) {
+      const auto *ParamInfo =
+          static_cast<const mcdxbc::RootParameterInfo *>(Reg->Cookie);
+
+      if (RC != ResourceClass::SRV && RC != ResourceClass::UAV)
+        continue;
+
+      if (ParamInfo->Type == dxbc::RootParameterType::DescriptorTable)
+        continue;
+
+      if (RK != ResourceKind::RawBuffer && RK != ResourceKind::StructuredBuffer)
+        reportInvalidHandleTyError(M, RC, Binding);
+    } else {
       reportRegNotBound(M, RC, Binding);
+    }
   }
 }
 

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

@@ -0,0 +1,17 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 
+; "This is a valid root signature with a texture/typed buffer resource"
+
+@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", <4 x float>, 1, 0, 0, 4) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, 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-valid-typedbuffer.ll

@@ -0,0 +1,17 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 
+; This is a valid root signature with a texture/typed buffer resource
+
+@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(i32 0, i32 0, i32 1, i32 0, 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-textures-fail.ll

@@ -0,0 +1,16 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; CHECK: error: UAV at register 0 and space 0 is bound to a texture or typed buffer.
+
+@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, 4) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, 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-fail.ll

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