[DirectX] Add Range Overlap validation to DXILPostOptimizationValidation.cpp

[joaosaffran]

According to Root Signature spec, we need to validate if ranges defined in root signature are overlapping. This PR adds such check.
Closes: #126645

[llvmbot]

@llvm/pr-subscribers-backend-directx

Author: None (joaosaffran)

Changes

According to Root Signature spec, we need to validate if ranges defined in root signature are overlapping. This PR adds such check.
Closes: #126645

---

Full diff: https://github.com/llvm/llvm-project/pull/148919.diff

3 Files Affected:

• (modified) llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp (+153-15)
• (added) llvm/test/CodeGen/DirectX/rootsignature-validation-fail-cbuffer-range.ll (+15)
• (added) llvm/test/CodeGen/DirectX/rootsignature-validation-fail-overlap-range.ll (+16)

diff --git a/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp b/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp
index f95a1d5913fd9..f22ae36df98b8 100644
--- a/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp
+++ b/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp
@@ -15,6 +15,7 @@
 #include "llvm/Analysis/DXILMetadataAnalysis.h"
 #include "llvm/Analysis/DXILResource.h"
 #include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/Frontend/HLSL/RootSignatureValidations.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicsDirectX.h"
@@ -227,6 +228,139 @@ getRootSignature(RootSignatureBindingInfo &RSBI,
   return RootSigDesc;
 }
 
+static void
+reportOverlappingRegisters(Module &M,
+                           llvm::hlsl::rootsig::OverlappingRanges Overlap) {
+  const llvm::hlsl::rootsig::RangeInfo *Info = Overlap.A;
+  const llvm::hlsl::rootsig::RangeInfo *OInfo = Overlap.B;
+  SmallString<128> Message;
+  raw_svector_ostream OS(Message);
+  auto ResourceClassToString =
+      [](llvm::dxil::ResourceClass Class) -> const char * {
+    switch (Class) {
+
+    case ResourceClass::SRV:
+      return "SRV";
+    case ResourceClass::UAV:
+      return "UAV";
+    case ResourceClass::CBuffer:
+      return "CBuffer";
+    case ResourceClass::Sampler:
+      return "Sampler";
+      break;
+    }
+  };
+  OS << "register " << ResourceClassToString(Info->Class)
+     << " (space=" << Info->Space << ", register=" << Info->LowerBound << ")"
+     << " is overlapping with"
+     << " register " << ResourceClassToString(OInfo->Class)
+     << " (space=" << OInfo->Space << ", register=" << OInfo->LowerBound << ")"
+     << ", verify your root signature definition";
+
+  M.getContext().diagnose(DiagnosticInfoGeneric(Message));
+}
+
+static bool reportOverlappingRanges(Module &M,
+                                    const mcdxbc::RootSignatureDesc &RSD) {
+  using namespace llvm::hlsl::rootsig;
+
+  llvm::SmallVector<RangeInfo> Infos;
+  // Helper to map DescriptorRangeType to ResourceClass
+  auto RangeToResourceClass = [](uint32_t RangeType) -> ResourceClass {
+    using namespace dxbc;
+    switch (static_cast<DescriptorRangeType>(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;
+    }
+  };
+
+  // Helper to map RootParameterType to ResourceClass
+  auto ParameterToResourceClass = [](uint32_t Type) -> ResourceClass {
+    using namespace dxbc;
+    switch (Type) {
+    case llvm::to_underlying(RootParameterType::Constants32Bit):
+      return ResourceClass::CBuffer;
+    case llvm::to_underlying(RootParameterType::SRV):
+      return ResourceClass::SRV;
+    case llvm::to_underlying(RootParameterType::UAV):
+      return ResourceClass::UAV;
+    case llvm::to_underlying(RootParameterType::CBV):
+      return ResourceClass::CBuffer;
+    default:
+      llvm_unreachable("Unknown RootParameterType");
+    }
+  };
+
+  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);
+    switch (Type) {
+    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
+      dxbc::RTS0::v1::RootConstants Const =
+          RSD.ParametersContainer.getConstant(Loc);
+
+      RangeInfo Info;
+      Info.Space = Const.RegisterSpace;
+      Info.LowerBound = Const.ShaderRegister;
+      Info.UpperBound = Info.LowerBound;
+      Info.Class = ParameterToResourceClass(Type);
+      Info.Visibility = (dxbc::ShaderVisibility)Header.ShaderVisibility;
+
+      Infos.push_back(Info);
+      break;
+    }
+    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);
+
+      RangeInfo Info;
+      Info.Space = Desc.RegisterSpace;
+      Info.LowerBound = Desc.ShaderRegister;
+      Info.UpperBound = Info.LowerBound;
+      Info.Class = ParameterToResourceClass(Type);
+      Info.Visibility = (dxbc::ShaderVisibility)Header.ShaderVisibility;
+
+      Infos.push_back(Info);
+      break;
+    }
+    case llvm::to_underlying(dxbc::RootParameterType::DescriptorTable): {
+      const mcdxbc::DescriptorTable &Table =
+          RSD.ParametersContainer.getDescriptorTable(Loc);
+
+      for (const dxbc::RTS0::v2::DescriptorRange &Range : Table.Ranges) {
+        RangeInfo Info;
+        Info.Space = Range.RegisterSpace;
+        Info.LowerBound = Range.BaseShaderRegister;
+        Info.UpperBound = Info.LowerBound + ((Range.NumDescriptors == ~0U)
+                                                 ? Range.NumDescriptors
+                                                 : Range.NumDescriptors - 1);
+        Info.Visibility = (dxbc::ShaderVisibility)Header.ShaderVisibility;
+        Info.Class = RangeToResourceClass(Range.RangeType);
+
+        Infos.push_back(Info);
+      }
+      break;
+    }
+    }
+  }
+
+  llvm::SmallVector<OverlappingRanges> Overlaps =
+      llvm::hlsl::rootsig::findOverlappingRanges(Infos);
+  for (OverlappingRanges Overlap : Overlaps)
+    reportOverlappingRegisters(M, Overlap);
+
+  return Overlaps.size() > 0;
+}
+
 static void reportInvalidRegistersBinding(
     Module &M,
     const llvm::ArrayRef<llvm::dxil::ResourceInfo::ResourceBinding> &Bindings,
@@ -272,21 +406,25 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
 
   if (auto RSD = getRootSignature(RSBI, MMI)) {
 
-    RootSignatureBindingValidation Validation =
-        initRSBindingValidation(*RSD, tripleToVisibility(MMI.ShaderProfile));
-
-    reportInvalidRegistersBinding(
-        M, Validation.getBindingsOfType(dxbc::DescriptorRangeType::CBV),
-        DRM.cbuffers());
-    reportInvalidRegistersBinding(
-        M, Validation.getBindingsOfType(dxbc::DescriptorRangeType::UAV),
-        DRM.uavs());
-    reportInvalidRegistersBinding(
-        M, Validation.getBindingsOfType(dxbc::DescriptorRangeType::Sampler),
-        DRM.samplers());
-    reportInvalidRegistersBinding(
-        M, Validation.getBindingsOfType(dxbc::DescriptorRangeType::SRV),
-        DRM.srvs());
+    if (!reportOverlappingRanges(M, *RSD)) {
+      // Those checks require that no range is overlapping to provide correct
+      // diagnostic.
+      RootSignatureBindingValidation Validation =
+          initRSBindingValidation(*RSD, tripleToVisibility(MMI.ShaderProfile));
+
+      reportInvalidRegistersBinding(
+          M, Validation.getBindingsOfType(dxbc::DescriptorRangeType::CBV),
+          DRM.cbuffers());
+      reportInvalidRegistersBinding(
+          M, Validation.getBindingsOfType(dxbc::DescriptorRangeType::UAV),
+          DRM.uavs());
+      reportInvalidRegistersBinding(
+          M, Validation.getBindingsOfType(dxbc::DescriptorRangeType::Sampler),
+          DRM.samplers());
+      reportInvalidRegistersBinding(
+          M, Validation.getBindingsOfType(dxbc::DescriptorRangeType::SRV),
+          DRM.srvs());
+    }
   }
 }
 } // namespace
diff --git a/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-cbuffer-range.ll b/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-cbuffer-range.ll
new file mode 100644
index 0000000000000..431ec577efc5f
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-cbuffer-range.ll
@@ -0,0 +1,15 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.6-compute %s 2>&1 | FileCheck %s
+; CHECK: error: register CBuffer (space=0, register=0) is overlapping with register CBuffer (space=0, register=2), verify your root signature definition
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+  ret void
+}
+
+; RootConstants(num32BitConstants=4, b2), DescriptorTable(CBV(b10, numDescriptors=3))
+!dx.rootsignatures = !{!0}
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !3}
+!2 = !{!"RootConstants", i32 0, i32 2, i32 0, i32 4}
+!3 = !{!"DescriptorTable", i32 0, !4}
+!4 = !{!"CBV", i32 3, i32 0, i32 0, i32 -1, i32 4}
diff --git a/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-overlap-range.ll b/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-overlap-range.ll
new file mode 100644
index 0000000000000..17b22b62ff8f5
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/rootsignature-validation-fail-overlap-range.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: register UAV (space=0, register=2) is overlapping with register UAV (space=0, register=0), verify your root signature definition
+
+define void @CSMain() "hlsl.shader"="compute" {
+entry:
+  ret void
+}
+
+; DescriptorTable(UAV(u0, numDescriptors=unbounded), visibility = SHADER_VISIBILITY_HULL), DescriptorTable(UAV(u2, numDescriptors=4))
+!dx.rootsignatures = !{!0}
+!0 = !{ptr @CSMain, !1, i32 2}
+!1 = !{!2, !4}
+!2 = !{!"DescriptorTable", i32 2, !3}
+!3 = !{!"UAV", i32 -1, i32 0, i32 0, i32 -1, i32 2}
+!4 = !{!"DescriptorTable", i32 0, !5}
+!5 = !{!"UAV", i32 4, i32 2, i32 0, i32 -1, i32 2}

[bob80905]

I see your test case for UAV, and if I understand correctly, the second test file you added might be testing the Constants32Bit.
Is there an existing test for the Descriptor table overlapping case?

[llvm-beanz]

Rather than using llvm::to_underlying if you convert the int to the enum early (and handle the case where it doesn't match), the compiler will make sure you have full switch coverage.

That's definitely the better pattern to use for cases like this.

[llvm-beanz]

I filed an issue to clean up this pattern across the DX backend. We really shouldn't be writing switch statements like this. The issue is #150676.

[joaosaffran]

I will abandon this PR, since parent PR has changed significantly.