[DirectX] Add Range Overlap validation

llvm/include/llvm/Analysis/DXILResource.h

@@ -33,20 +33,6 @@ class DXILResourceTypeMap;
 
 namespace dxil {
 
-inline StringRef getResourceClassName(ResourceClass RC) {
-  switch (RC) {
-  case ResourceClass::SRV:
-    return "SRV";
-  case ResourceClass::UAV:
-    return "UAV";
-  case ResourceClass::CBuffer:
-    return "CBuffer";
-  case ResourceClass::Sampler:
-    return "Sampler";
-  }
-  llvm_unreachable("Unhandled ResourceClass");
-}
-
 // Returns the resource name from dx_resource_handlefrombinding or
 // dx_resource_handlefromimplicitbinding call
 LLVM_ABI StringRef getResourceNameFromBindingCall(CallInst *CI);

llvm/include/llvm/Support/DXILABI.h

@@ -17,6 +17,7 @@
 #ifndef LLVM_SUPPORT_DXILABI_H
 #define LLVM_SUPPORT_DXILABI_H
 
+#include "llvm/ADT/StringRef.h"
 #include <cstdint>
 
 namespace llvm {
@@ -99,6 +100,8 @@ enum class SamplerFeedbackType : uint32_t {
 const unsigned MinWaveSize = 4;
 const unsigned MaxWaveSize = 128;
 
+StringRef getResourceClassName(ResourceClass RC);
+
 } // namespace dxil
 } // namespace llvm
 

llvm/lib/Analysis/DXILResource.cpp

@@ -20,6 +20,7 @@
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
+#include "llvm/Support/DXILABI.h"
 #include "llvm/Support/FormatVariadic.h"
 #include <cstdint>
 #include <optional>

llvm/lib/Support/CMakeLists.txt

@@ -182,6 +182,7 @@ add_llvm_component_library(LLVMSupport
   DivisionByConstantInfo.cpp
   DAGDeltaAlgorithm.cpp
   DJB.cpp
+  DXILABI.cpp
   DynamicAPInt.cpp
   ELFAttributes.cpp
   ELFAttrParserCompact.cpp

llvm/lib/Support/DXILABI.cpp

@@ -0,0 +1,22 @@
+
+#include "llvm/Support/DXILABI.h"
+#include "llvm/Support/ErrorHandling.h"
+
+using namespace llvm;
+namespace llvm {
+namespace dxil {
+StringRef getResourceClassName(dxil::ResourceClass RC) {
+  switch (RC) {
+  case dxil::ResourceClass::SRV:
+    return "SRV";
+  case dxil::ResourceClass::UAV:
+    return "UAV";
+  case dxil::ResourceClass::CBuffer:
+    return "CBuffer";
+  case dxil::ResourceClass::Sampler:
+    return "Sampler";
+  }
+  llvm_unreachable("Unhandled ResourceClass");
+}
+} // namespace dxil
+} // namespace llvm

llvm/lib/Target/DirectX/DXContainerGlobals.cpp

@@ -162,8 +162,7 @@ void DXContainerGlobals::addRootSignature(Module &M,
 
   auto &RSA = getAnalysis<RootSignatureAnalysisWrapper>().getRSInfo();
   const Function *EntryFunction = MMI.EntryPropertyVec[0].Entry;
-  const std::optional<mcdxbc::RootSignatureDesc> &RS =
-      RSA.getDescForFunction(EntryFunction);
+  const mcdxbc::RootSignatureDesc *RS = RSA.getDescForFunction(EntryFunction);
 
   if (!RS)
     return;

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -919,6 +919,7 @@ PreservedAnalyses DXILOpLowering::run(Module &M, ModuleAnalysisManager &MAM) {
   PA.preserve<DXILResourceAnalysis>();
   PA.preserve<DXILMetadataAnalysis>();
   PA.preserve<ShaderFlagsAnalysis>();
+  PA.preserve<RootSignatureAnalysis>();
   return PA;
 }
 

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -18,16 +18,16 @@
 #include "llvm/IR/IntrinsicsDirectX.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
+#include "llvm/Support/DXILABI.h"
 
 #define DEBUG_TYPE "dxil-post-optimization-validation"
 
 using namespace llvm;
 using namespace llvm::dxil;
 
-namespace {
-static ResourceClass RangeToResourceClass(uint32_t RangeType) {
+static ResourceClass toResourceClass(dxbc::DescriptorRangeType RangeType) {
   using namespace dxbc;
-  switch (static_cast<DescriptorRangeType>(RangeType)) {
+  switch (RangeType) {
   case DescriptorRangeType::SRV:
     return ResourceClass::SRV;
   case DescriptorRangeType::UAV:
@@ -39,20 +39,21 @@ static ResourceClass RangeToResourceClass(uint32_t RangeType) {
   }
 }
 
-ResourceClass ParameterToResourceClass(uint32_t Type) {
+static ResourceClass toResourceClass(dxbc::RootParameterType Type) {
   using namespace dxbc;
   switch (Type) {
-  case llvm::to_underlying(RootParameterType::Constants32Bit):
+  case RootParameterType::Constants32Bit:
     return ResourceClass::CBuffer;
-  case llvm::to_underlying(RootParameterType::SRV):
+  case RootParameterType::SRV:
     return ResourceClass::SRV;
-  case llvm::to_underlying(RootParameterType::UAV):
+  case RootParameterType::UAV:
     return ResourceClass::UAV;
-  case llvm::to_underlying(RootParameterType::CBV):
+  case RootParameterType::CBV:
     return ResourceClass::CBuffer;
-  default:
-    llvm_unreachable("Unknown RootParameterType");
+  case dxbc::RootParameterType::DescriptorTable:
+    break;
   }
+  llvm_unreachable("Unconvertible RootParameterType");
 }
 
 static void reportInvalidDirection(Module &M, DXILResourceMap &DRM) {
@@ -131,12 +132,6 @@ static void reportOverlappingRegisters(
 
 static dxbc::ShaderVisibility
 tripleToVisibility(llvm::Triple::EnvironmentType ET) {
-  assert((ET == Triple::Pixel || ET == Triple::Vertex ||
-          ET == Triple::Geometry || ET == Triple::Hull ||
-          ET == Triple::Domain || ET == Triple::Mesh ||
-          ET == Triple::Compute) &&
-         "Invalid Triple to shader stage conversion");
-
   switch (ET) {
   case Triple::Pixel:
     return dxbc::ShaderVisibility::Pixel;
@@ -157,73 +152,80 @@ tripleToVisibility(llvm::Triple::EnvironmentType ET) {
   }
 }
 
-static void trackRootSigDescBinding(hlsl::BindingInfoBuilder &Builder,
-                                    const mcdxbc::RootSignatureDesc &RSD,
-                                    dxbc::ShaderVisibility Visibility) {
-  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;
+static void validateRootSignature(Module &M,
+                                  const mcdxbc::RootSignatureDesc &RSD,
+                                  dxil::ModuleMetadataInfo &MMI) {
 
-    switch (Type) {
-    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
+  hlsl::BindingInfoBuilder Builder;
+  dxbc::ShaderVisibility Visibility = tripleToVisibility(MMI.ShaderProfile);
+
+  for (const mcdxbc::RootParameterInfo &ParamInfo : RSD.ParametersContainer) {
+    dxbc::ShaderVisibility ParamVisibility =
+        static_cast<dxbc::ShaderVisibility>(ParamInfo.Header.ShaderVisibility);
+    if (ParamVisibility != dxbc::ShaderVisibility::All &&
+        ParamVisibility != Visibility)
+      continue;
+    dxbc::RootParameterType ParamType =
+        static_cast<dxbc::RootParameterType>(ParamInfo.Header.ParameterType);
+    switch (ParamType) {
+    case dxbc::RootParameterType::Constants32Bit: {
       dxbc::RTS0::v1::RootConstants Const =
-          RSD.ParametersContainer.getConstant(Loc);
+          RSD.ParametersContainer.getConstant(ParamInfo.Location);
       Builder.trackBinding(dxil::ResourceClass::CBuffer, Const.RegisterSpace,
-                           Const.ShaderRegister,
-                           Const.ShaderRegister + Const.Num32BitValues, &Const);
+                           Const.ShaderRegister, Const.ShaderRegister, nullptr);
       break;
     }
 
-    case llvm::to_underlying(dxbc::RootParameterType::SRV):
-    case llvm::to_underlying(dxbc::RootParameterType::UAV):
-    case llvm::to_underlying(dxbc::RootParameterType::CBV): {
+    case dxbc::RootParameterType::SRV:
+    case dxbc::RootParameterType::UAV:
+    case dxbc::RootParameterType::CBV: {
       dxbc::RTS0::v2::RootDescriptor Desc =
-          RSD.ParametersContainer.getRootDescriptor(Loc);
-      Builder.trackBinding(ParameterToResourceClass(Type), Desc.RegisterSpace,
-                           Desc.ShaderRegister, Desc.ShaderRegister, &Desc);
+          RSD.ParametersContainer.getRootDescriptor(ParamInfo.Location);
+      Builder.trackBinding(toResourceClass(static_cast<dxbc::RootParameterType>(
+                               ParamInfo.Header.ParameterType)),
+                           Desc.RegisterSpace, Desc.ShaderRegister,
+                           Desc.ShaderRegister, nullptr);
 
       break;
     }
-    case llvm::to_underlying(dxbc::RootParameterType::DescriptorTable): {
+    case dxbc::RootParameterType::DescriptorTable: {
       const mcdxbc::DescriptorTable &Table =
-          RSD.ParametersContainer.getDescriptorTable(Loc);
+          RSD.ParametersContainer.getDescriptorTable(ParamInfo.Location);
 
       for (const dxbc::RTS0::v2::DescriptorRange &Range : Table.Ranges) {
-        Builder.trackBinding(RangeToResourceClass(Range.RangeType),
-                             Range.RegisterSpace, Range.BaseShaderRegister,
-                             Range.NumDescriptors == ~0U
-                                 ? Range.NumDescriptors
-                                 : Range.BaseShaderRegister +
-                                       Range.NumDescriptors,
-                             &Range);
+        uint32_t UpperBound =
+            Range.NumDescriptors == ~0U
+                ? Range.BaseShaderRegister
+                : Range.BaseShaderRegister + Range.NumDescriptors - 1;
+        Builder.trackBinding(
+            toResourceClass(
+                static_cast<dxbc::DescriptorRangeType>(Range.RangeType)),
+            Range.RegisterSpace, Range.BaseShaderRegister, UpperBound, nullptr);
       }
       break;
     }
     }
   }
 
-  for (auto &S : RSD.StaticSamplers) {
+  for (const dxbc::RTS0::v1::StaticSampler &S : RSD.StaticSamplers)
     Builder.trackBinding(dxil::ResourceClass::Sampler, S.RegisterSpace,
-                         S.ShaderRegister, S.ShaderRegister, &S);
-  }
+                         S.ShaderRegister, S.ShaderRegister, nullptr);
+
+  hlsl::BindingInfo Info = Builder.calculateBindingInfo(
+      [&M](const llvm::hlsl::BindingInfoBuilder &Builder,
+           const llvm::hlsl::BindingInfoBuilder::Binding &ReportedBinding) {
+        const llvm::hlsl::BindingInfoBuilder::Binding &Overlaping =
+            Builder.findOverlapping(ReportedBinding);
+        reportOverlappingRegisters(M, ReportedBinding, Overlaping);
+      });
 }
 
-std::optional<mcdxbc::RootSignatureDesc>
+static mcdxbc::RootSignatureDesc *
 getRootSignature(RootSignatureBindingInfo &RSBI,
                  dxil::ModuleMetadataInfo &MMI) {
   if (MMI.EntryPropertyVec.size() == 0)
-    return std::nullopt;
-  std::optional<mcdxbc::RootSignatureDesc> RootSigDesc =
-      RSBI.getDescForFunction(MMI.EntryPropertyVec[0].Entry);
-  if (!RootSigDesc)
-    return std::nullopt;
-  return RootSigDesc;
+    return nullptr;
+  return RSBI.getDescForFunction(MMI.EntryPropertyVec[0].Entry);
 }
 
 static void reportErrors(Module &M, DXILResourceMap &DRM,
@@ -239,21 +241,9 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
   assert(!DRBI.hasImplicitBinding() && "implicit bindings should be handled in "
                                        "DXILResourceImplicitBinding pass");
 
-  if (auto RSD = getRootSignature(RSBI, MMI)) {
-
-    hlsl::BindingInfoBuilder Builder;
-    dxbc::ShaderVisibility Visibility = tripleToVisibility(MMI.ShaderProfile);
-    trackRootSigDescBinding(Builder, *RSD, Visibility);
-    hlsl::BindingInfo Info = Builder.calculateBindingInfo(
-        [&M](const llvm::hlsl::BindingInfoBuilder &Builder,
-             const llvm::hlsl::BindingInfoBuilder::Binding &ReportedBinding) {
-          const llvm::hlsl::BindingInfoBuilder::Binding &Overlaping =
-              Builder.findOverlapping(ReportedBinding);
-          reportOverlappingRegisters(M, ReportedBinding, Overlaping);
-        });
-  }
+  if (mcdxbc::RootSignatureDesc *RSD = getRootSignature(RSBI, MMI))
+    validateRootSignature(M, *RSD, MMI);
 }
-} // namespace
 
 PreservedAnalyses
 DXILPostOptimizationValidation::run(Module &M, ModuleAnalysisManager &MAM) {

llvm/lib/Target/DirectX/DXILRootSignature.h

@@ -43,13 +43,11 @@ class RootSignatureBindingInfo {
 
   iterator end() { return FuncToRsMap.end(); }
 
-  std::optional<mcdxbc::RootSignatureDesc>
-  getDescForFunction(const Function *F) {
+  mcdxbc::RootSignatureDesc *getDescForFunction(const Function *F) {
     const auto FuncRs = find(F);
     if (FuncRs == end())
-      return std::nullopt;
-
-    return FuncRs->second;
+      return nullptr;
+    return &FuncRs->second;
   }
 };
 

llvm/test/CodeGen/DirectX/rootsignature-validation-fail-root-descriptor-range.ll

@@ -6,10 +6,9 @@ entry:
   ret void
 }
 
-; DescriptorTable(UAV(u0, numDescriptors=unbounded), visibility = SHADER_VISIBILITY_HULL), UAV(u3, space=1)
 !dx.rootsignatures = !{!0}
 !0 = !{ptr @CSMain, !1, i32 2}
 !1 = !{!2, !4}
 !2 = !{!"RootUAV", i32 0, i32 3, i32 1, i32 4}
 !4 = !{!"DescriptorTable", i32 0, !5}
-!5 = !{!"UAV", i32 3, i32 0, i32 1, i32 -1, i32 2}
+!5 = !{!"UAV", i32 4, i32 0, i32 1, i32 -1, i32 2}