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

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -13,11 +13,13 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Analysis/DXILMetadataAnalysis.h"
 #include "llvm/Analysis/DXILResource.h"
+#include "llvm/BinaryFormat/DXContainer.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/MC/DXContainerRootSignature.h"
 #include "llvm/Support/DXILABI.h"
 
 #define DEBUG_TYPE "dxil-post-optimization-validation"
@@ -115,6 +117,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 reportOverlappingRegisters(
     Module &M, const llvm::hlsl::BindingInfoBuilder::Binding &Reported,
     const llvm::hlsl::BindingInfoBuilder::Binding &Overlaping) {
@@ -142,6 +155,25 @@ reportRegNotBound(Module &M, ResourceClass Class,
   M.getContext().diagnose(DiagnosticInfoGeneric(Message));
 }
 
+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();
+    for (const auto &RD : RDs) {
+      if (Binding.overlapsWith(RD)) {
+        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 dxbc::ShaderVisibility
 tripleToVisibility(llvm::Triple::EnvironmentType ET) {
   switch (ET) {
@@ -164,6 +196,46 @@ tripleToVisibility(llvm::Triple::EnvironmentType ET) {
   }
 }
 
+static SmallVector<ResourceInfo::ResourceBinding>
+getRootDescriptorsBindingInfo(const mcdxbc::RootSignatureDesc &RSD,
+                              dxbc::ShaderVisibility Visibility) {
+
+  SmallVector<ResourceInfo::ResourceBinding> RDs;
+
+  for (const mcdxbc::RootParameterInfo &Info : RSD.ParametersContainer) {
+    dxbc::ShaderVisibility ParamVisibility =
+        static_cast<dxbc::ShaderVisibility>(Info.Header.ShaderVisibility);
+    if (ParamVisibility != dxbc::ShaderVisibility::All &&
+        ParamVisibility != Visibility)
+      continue;
+
+    dxbc::RootParameterType ParamType =
+        static_cast<dxbc::RootParameterType>(Info.Header.ParameterType);
+    switch (ParamType) {
+
+    case dxbc::RootParameterType::SRV:
+    case dxbc::RootParameterType::UAV:
+    case dxbc::RootParameterType::CBV: {
+      dxbc::RTS0::v2::RootDescriptor Desc =
+          RSD.ParametersContainer.getRootDescriptor(Info.Location);
+
+      ResourceInfo::ResourceBinding Binding;
+      Binding.LowerBound = Desc.ShaderRegister;
+      Binding.Space = Desc.RegisterSpace;
+      Binding.Size = 1;
+
+      RDs.push_back(Binding);
+      break;
+    }
+    case dxbc::RootParameterType::DescriptorTable:
+    case dxbc::RootParameterType::Constants32Bit:
+      break;
+    }
+  }
+
+  return RDs;
+}
+
 static void validateRootSignature(Module &M,
                                   const mcdxbc::RootSignatureDesc &RSD,
                                   dxil::ModuleMetadataInfo &MMI,
@@ -238,6 +310,8 @@ static void validateRootSignature(Module &M,
       });
   // Next checks require that the root signature definition is valid.
   if (!HasOverlap) {
+    SmallVector<ResourceInfo::ResourceBinding> RDs =
+        getRootDescriptorsBindingInfo(RSD, Visibility);
     for (const auto &ResList :
          {std::make_pair(ResourceClass::SRV, DRM.srvs()),
           std::make_pair(ResourceClass::UAV, DRM.uavs()),
@@ -251,13 +325,13 @@ static void validateRootSignature(Module &M,
         if (!Info.isBound(ResList.first, ResBinding.Space, ResRange))
           reportRegNotBound(M, ResList.first, ResBinding);
       }
+      reportInvalidHandleTy(M, RDs, ResList.second);
     }
   }
 }
 
-static mcdxbc::RootSignatureDesc *
-getRootSignature(RootSignatureBindingInfo &RSBI,
-                 dxil::ModuleMetadataInfo &MMI) {
+mcdxbc::RootSignatureDesc *getRootSignature(RootSignatureBindingInfo &RSBI,
+                                            dxil::ModuleMetadataInfo &MMI) {
   if (MMI.EntryPropertyVec.size() == 0)
     return nullptr;
   return RSBI.getDescForFunction(MMI.EntryPropertyVec[0].Entry);

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
 }