[DirectX] Add Range Overlap validation

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>
@@ -29,20 +30,6 @@
 using namespace llvm;
 using namespace dxil;
 
-static 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");
-}
-
 static StringRef getResourceKindName(ResourceKind RK) {
   switch (RK) {
   case ResourceKind::Texture1D:

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

@@ -9,6 +9,7 @@
 #include "DXILOpLowering.h"
 #include "DXILConstants.h"
 #include "DXILOpBuilder.h"
+#include "DXILRootSignature.h"
 #include "DXILShaderFlags.h"
 #include "DirectX.h"
 #include "llvm/ADT/SmallVector.h"
@@ -918,6 +919,7 @@ PreservedAnalyses DXILOpLowering::run(Module &M, ModuleAnalysisManager &MAM) {
   PA.preserve<DXILResourceAnalysis>();
   PA.preserve<DXILMetadataAnalysis>();
   PA.preserve<ShaderFlagsAnalysis>();
+  PA.preserve<RootSignatureAnalysis>();
   return PA;
 }
 
@@ -945,6 +947,7 @@ class DXILOpLoweringLegacy : public ModulePass {
     AU.addPreserved<DXILResourceWrapperPass>();
     AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
     AU.addPreserved<ShaderFlagsAnalysisWrapper>();
+    AU.addPreserved<RootSignatureAnalysisWrapper>();
   }
 };
 char DXILOpLoweringLegacy::ID = 0;

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "DXILPostOptimizationValidation.h"
+#include "DXILRootSignature.h"
 #include "DXILShaderFlags.h"
 #include "DirectX.h"
 #include "llvm/ADT/SmallString.h"
@@ -17,13 +18,43 @@
 #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 toResourceClass(dxbc::DescriptorRangeType RangeType) {
+  using namespace dxbc;
+  switch (RangeType) {
+  case DescriptorRangeType::SRV:
+    return ResourceClass::SRV;
+  case DescriptorRangeType::UAV:
+    return ResourceClass::UAV;
+  case DescriptorRangeType::CBV:
+    return ResourceClass::CBuffer;
+  case DescriptorRangeType::Sampler:
+    return ResourceClass::Sampler;
+  }
+}
+
+static ResourceClass toResourceClass(dxbc::RootParameterType Type) {
+  using namespace dxbc;
+  switch (Type) {
+  case RootParameterType::Constants32Bit:
+    return ResourceClass::CBuffer;
+  case RootParameterType::SRV:
+    return ResourceClass::SRV;
+  case RootParameterType::UAV:
+    return ResourceClass::UAV;
+  case RootParameterType::CBV:
+    return ResourceClass::CBuffer;
+  case dxbc::RootParameterType::DescriptorTable:
+    break;
+  }
+  llvm_unreachable("Unconvertible RootParameterType");
+}
 
 static void reportInvalidDirection(Module &M, DXILResourceMap &DRM) {
   for (const auto &UAV : DRM.uavs()) {
@@ -84,8 +115,123 @@ static void reportOverlappingBinding(Module &M, DXILResourceMap &DRM) {
   }
 }
 
+static void reportOverlappingRegisters(
+    Module &M, const llvm::hlsl::BindingInfoBuilder::Binding &Reported,
+    const llvm::hlsl::BindingInfoBuilder::Binding &Overlaping) {
+  SmallString<128> Message;
+  raw_svector_ostream OS(Message);
+  OS << "register " << getResourceClassName(Reported.RC)
+     << " (space=" << Reported.Space << ", register=" << Reported.LowerBound
+     << ")" << " is overlapping with" << " register "
+     << getResourceClassName(Overlaping.RC) << " (space=" << Overlaping.Space
+     << ", register=" << Overlaping.LowerBound << ")"
+     << ", verify your root signature definition.";
+
+  M.getContext().diagnose(DiagnosticInfoGeneric(Message));
+}
+
+static dxbc::ShaderVisibility
+tripleToVisibility(llvm::Triple::EnvironmentType ET) {
+  switch (ET) {
+  case Triple::Pixel:
+    return dxbc::ShaderVisibility::Pixel;
+  case Triple::Vertex:
+    return dxbc::ShaderVisibility::Vertex;
+  case Triple::Geometry:
+    return dxbc::ShaderVisibility::Geometry;
+  case Triple::Hull:
+    return dxbc::ShaderVisibility::Hull;
+  case Triple::Domain:
+    return dxbc::ShaderVisibility::Domain;
+  case Triple::Mesh:
+    return dxbc::ShaderVisibility::Mesh;
+  case Triple::Compute:
+    return dxbc::ShaderVisibility::All;
+  default:
+    llvm_unreachable("Invalid triple to shader stage conversion");
+  }
+}
+
+static void validateRootSignature(Module &M,
+                                  const mcdxbc::RootSignatureDesc &RSD,
+                                  dxil::ModuleMetadataInfo &MMI) {
+
+  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(ParamInfo.Location);
+      Builder.trackBinding(dxil::ResourceClass::CBuffer, Const.RegisterSpace,
+                           Const.ShaderRegister, Const.ShaderRegister, nullptr);
+      break;
+    }
+
+    case dxbc::RootParameterType::SRV:
+    case dxbc::RootParameterType::UAV:
+    case dxbc::RootParameterType::CBV: {
+      dxbc::RTS0::v2::RootDescriptor 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 dxbc::RootParameterType::DescriptorTable: {
+      const mcdxbc::DescriptorTable &Table =
+          RSD.ParametersContainer.getDescriptorTable(ParamInfo.Location);
+
+      for (const dxbc::RTS0::v2::DescriptorRange &Range : Table.Ranges) {
+        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 (const dxbc::RTS0::v1::StaticSampler &S : RSD.StaticSamplers)
+    Builder.trackBinding(dxil::ResourceClass::Sampler, S.RegisterSpace,
+                         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);
+      });
+}
+
+static mcdxbc::RootSignatureDesc *
+getRootSignature(RootSignatureBindingInfo &RSBI,
+                 dxil::ModuleMetadataInfo &MMI) {
+  if (MMI.EntryPropertyVec.size() == 0)
+    return nullptr;
+  return RSBI.getDescForFunction(MMI.EntryPropertyVec[0].Entry);
+}
+
 static void reportErrors(Module &M, DXILResourceMap &DRM,
-                         DXILResourceBindingInfo &DRBI) {
+                         DXILResourceBindingInfo &DRBI,
+                         RootSignatureBindingInfo &RSBI,
+                         dxil::ModuleMetadataInfo &MMI) {
   if (DRM.hasInvalidCounterDirection())
     reportInvalidDirection(M, DRM);
 
@@ -94,14 +240,19 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
 
   assert(!DRBI.hasImplicitBinding() && "implicit bindings should be handled in "
                                        "DXILResourceImplicitBinding pass");
+
+  if (mcdxbc::RootSignatureDesc *RSD = getRootSignature(RSBI, MMI))
+    validateRootSignature(M, *RSD, MMI);
 }
-} // namespace
 
 PreservedAnalyses
 DXILPostOptimizationValidation::run(Module &M, ModuleAnalysisManager &MAM) {
   DXILResourceMap &DRM = MAM.getResult<DXILResourceAnalysis>(M);
   DXILResourceBindingInfo &DRBI = MAM.getResult<DXILResourceBindingAnalysis>(M);
-  reportErrors(M, DRM, DRBI);
+  RootSignatureBindingInfo &RSBI = MAM.getResult<RootSignatureAnalysis>(M);
+  ModuleMetadataInfo &MMI = MAM.getResult<DXILMetadataAnalysis>(M);
+
+  reportErrors(M, DRM, DRBI, RSBI, MMI);
   return PreservedAnalyses::all();
 }
 
@@ -113,7 +264,12 @@ class DXILPostOptimizationValidationLegacy : public ModulePass {
         getAnalysis<DXILResourceWrapperPass>().getResourceMap();
     DXILResourceBindingInfo &DRBI =
         getAnalysis<DXILResourceBindingWrapperPass>().getBindingInfo();
-    reportErrors(M, DRM, DRBI);
+    RootSignatureBindingInfo &RSBI =
+        getAnalysis<RootSignatureAnalysisWrapper>().getRSInfo();
+    dxil::ModuleMetadataInfo &MMI =
+        getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
+
+    reportErrors(M, DRM, DRBI, RSBI, MMI);
     return false;
   }
   StringRef getPassName() const override {
@@ -125,10 +281,13 @@ class DXILPostOptimizationValidationLegacy : public ModulePass {
   void getAnalysisUsage(llvm::AnalysisUsage &AU) const override {
     AU.addRequired<DXILResourceWrapperPass>();
     AU.addRequired<DXILResourceBindingWrapperPass>();
+    AU.addRequired<DXILMetadataAnalysisWrapperPass>();
+    AU.addRequired<RootSignatureAnalysisWrapper>();
     AU.addPreserved<DXILResourceWrapperPass>();
     AU.addPreserved<DXILResourceBindingWrapperPass>();
     AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
     AU.addPreserved<ShaderFlagsAnalysisWrapper>();
+    AU.addPreserved<RootSignatureAnalysisWrapper>();
   }
 };
 char DXILPostOptimizationValidationLegacy::ID = 0;
@@ -139,6 +298,8 @@ INITIALIZE_PASS_BEGIN(DXILPostOptimizationValidationLegacy, DEBUG_TYPE,
 INITIALIZE_PASS_DEPENDENCY(DXILResourceBindingWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DXILResourceTypeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DXILResourceWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DXILMetadataAnalysisWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(RootSignatureAnalysisWrapper)
 INITIALIZE_PASS_END(DXILPostOptimizationValidationLegacy, DEBUG_TYPE,
                     "DXIL Post Optimization Validation", false, false)
 

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/llc-pipeline.ll

@@ -33,9 +33,9 @@
 ; CHECK-NEXT:   DXIL Module Metadata analysis
 ; CHECK-NEXT:   DXIL Shader Flag Analysis
 ; CHECK-NEXT:   DXIL Translate Metadata
+; CHECK-NEXT:   DXIL Root Signature Analysis
 ; CHECK-NEXT:   DXIL Post Optimization Validation
 ; CHECK-NEXT:   DXIL Op Lowering
-; CHECK-NEXT:   DXIL Root Signature Analysis
 ; CHECK-NEXT:   DXIL Prepare Module
 
 ; CHECK-ASM-NEXT: DXIL Metadata Pretty Printer

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

@@ -0,0 +1,22 @@
+; RUN: opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 
+; expected-no-diagnostics
+; Root Signature(RootConstants(num32BitConstants=4, b2))
+
+%__cblayout_CB = type <{ float }>
+
+@CB.str = private unnamed_addr constant [3 x i8] c"CB\00", align 1
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+; cbuffer CB : register(b2, space0) {
+;  float a;
+; }
+  %CB = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 4, 0)) @llvm.dx.resource.handlefrombinding(i32 0, i32 2, i32 1, i32 0, i1 false, ptr nonnull @CB.str)
+  ret void
+}
+
+!dx.rootsignatures = !{!0}
+
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2}
+!2 = !{!"RootConstants", i32 0, i32 2, i32 0, i32 4}