finish implementing && fix tests

joaosaffran · joaosaffran · commit 8f40e83ab0db · 2025-07-04T19:46:24.000Z
diff --git a/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp b/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp
@@ -153,23 +153,29 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
     RootSignatureBindingValidation Validation;
     Validation.addRsBindingInfo(*RSD, tripleToVisibility(MMI.ShaderProfile));
 
-    for (const auto &CBuf : DRM.cbuffers()) {
+    for (const ResourceInfo &CBuf : DRM.cbuffers()) {
       ResourceInfo::ResourceBinding Binding = CBuf.getBinding();
-      if (!Validation.checkCregBinding(Binding))
+      if (!Validation.checkCRegBinding(Binding))
         reportRegNotBound(M, "cbuffer", Binding);
     }
 
-    for (const auto &CBuf : DRM.srvs()) {
-      ResourceInfo::ResourceBinding Binding = CBuf.getBinding();
+    for (const ResourceInfo &SRV : DRM.srvs()) {
+      ResourceInfo::ResourceBinding Binding = SRV.getBinding();
       if (!Validation.checkTRegBinding(Binding))
         reportRegNotBound(M, "srv", Binding);
     }
 
-    for (const auto &CBuf : DRM.uavs()) {
-      ResourceInfo::ResourceBinding Binding = CBuf.getBinding();
+    for (const ResourceInfo &UAV : DRM.uavs()) {
+      ResourceInfo::ResourceBinding Binding = UAV.getBinding();
       if (!Validation.checkURegBinding(Binding))
         reportRegNotBound(M, "uav", Binding);
     }
+
+    for (const ResourceInfo &Sampler : DRM.samplers()) {
+      ResourceInfo::ResourceBinding Binding = Sampler.getBinding();
+      if (!Validation.checkSamplerBinding(Binding))
+        reportRegNotBound(M, "sampler", Binding);
+    }
   }
 }
 } // namespace
@@ -199,10 +205,6 @@ void RootSignatureBindingValidation::addRsBindingInfo(
           RSD.ParametersContainer.getDescriptorTable(Loc);
 
       for (const dxbc::RTS0::v2::DescriptorRange &Range : Table.Ranges) {
-        if (Range.RangeType ==
-            llvm::to_underlying(dxbc::DescriptorRangeType::Sampler))
-          continue;
-
         if (Header.ShaderVisibility ==
                 llvm::to_underlying(dxbc::ShaderVisibility::All) ||
             Header.ShaderVisibility == llvm::to_underlying(Visibility))
@@ -214,29 +216,6 @@ void RootSignatureBindingValidation::addRsBindingInfo(
   }
 }
 
-bool RootSignatureBindingValidation::checkCregBinding(
-    ResourceInfo::ResourceBinding Binding) {
-  return CRegBindingsMap.overlaps(
-      combineUint32ToUint64(Binding.Space, Binding.LowerBound),
-      combineUint32ToUint64(Binding.Space,
-                            Binding.LowerBound + Binding.Size - 1));
-}
-
-bool RootSignatureBindingValidation::checkTRegBinding(
-    ResourceInfo::ResourceBinding Binding) {
-  return TRegBindingsMap.overlaps(
-      combineUint32ToUint64(Binding.Space, Binding.LowerBound),
-      combineUint32ToUint64(Binding.Space, Binding.LowerBound + Binding.Size));
-}
-
-bool RootSignatureBindingValidation::checkURegBinding(
-    ResourceInfo::ResourceBinding Binding) {
-  return URegBindingsMap.overlaps(
-      combineUint32ToUint64(Binding.Space, Binding.LowerBound),
-      combineUint32ToUint64(Binding.Space,
-                            Binding.LowerBound + Binding.Size - 1));
-}
-
 PreservedAnalyses
 DXILPostOptimizationValidation::run(Module &M, ModuleAnalysisManager &MAM) {
   DXILResourceMap &DRM = MAM.getResult<DXILResourceAnalysis>(M);
diff --git a/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.h b/llvm/lib/Target/DirectX/DXILPostOptimizationValidation.h
@@ -36,12 +36,13 @@ class RootSignatureBindingValidation {
   MapT CRegBindingsMap;
   MapT TRegBindingsMap;
   MapT URegBindingsMap;
+  MapT SamplersBindingsMap;
 
   void addRange(const dxbc::RTS0::v2::RootDescriptor &Desc, uint32_t Type) {
     assert((Type == llvm::to_underlying(dxbc::RootParameterType::CBV) ||
             Type == llvm::to_underlying(dxbc::RootParameterType::SRV) ||
             Type == llvm::to_underlying(dxbc::RootParameterType::UAV)) &&
-           "Invalid Type");
+           "Invalid Type in add Range Method");
 
     llvm::dxil::ResourceInfo::ResourceBinding Binding;
     Binding.LowerBound = Desc.ShaderRegister;
@@ -53,19 +54,20 @@ class RootSignatureBindingValidation {
     uint64_t HighRange = combineUint32ToUint64(
         Binding.Space, Binding.LowerBound + Binding.Size - 1);
 
+    assert(LowRange <= HighRange && "Invalid range configuration");
+
     switch (Type) {
 
     case llvm::to_underlying(dxbc::RootParameterType::CBV):
       CRegBindingsMap.insert(LowRange, HighRange, Binding);
-      return;
+      break;
     case llvm::to_underlying(dxbc::RootParameterType::SRV):
       TRegBindingsMap.insert(LowRange, HighRange, Binding);
-      return;
+      break;
     case llvm::to_underlying(dxbc::RootParameterType::UAV):
       URegBindingsMap.insert(LowRange, HighRange, Binding);
-      return;
+      break;
     }
-    llvm_unreachable("Invalid Type in add Range Method");
   }
 
   void addRange(const dxbc::RTS0::v2::DescriptorRange &Range) {
@@ -80,33 +82,59 @@ class RootSignatureBindingValidation {
     uint64_t HighRange = combineUint32ToUint64(
         Binding.Space, Binding.LowerBound + Binding.Size - 1);
 
+    assert(LowRange <= HighRange && "Invalid range configuration");
+
     switch (Range.RangeType) {
     case llvm::to_underlying(dxbc::DescriptorRangeType::CBV):
       CRegBindingsMap.insert(LowRange, HighRange, Binding);
-      return;
+      break;
     case llvm::to_underlying(dxbc::DescriptorRangeType::SRV):
       TRegBindingsMap.insert(LowRange, HighRange, Binding);
-      return;
+      break;
     case llvm::to_underlying(dxbc::DescriptorRangeType::UAV):
       URegBindingsMap.insert(LowRange, HighRange, Binding);
-      return;
+      break;
+    case llvm::to_underlying(dxbc::DescriptorRangeType::Sampler):
+      SamplersBindingsMap.insert(LowRange, HighRange, Binding);
+      break;
     }
-    llvm_unreachable("Invalid Type in add Range Method");
   }
 
 public:
   RootSignatureBindingValidation()
       : Allocator(), CRegBindingsMap(Allocator), TRegBindingsMap(Allocator),
-        URegBindingsMap(Allocator) {}
+        URegBindingsMap(Allocator), SamplersBindingsMap(Allocator) {}
 
   void addRsBindingInfo(mcdxbc::RootSignatureDesc &RSD,
                         dxbc::ShaderVisibility Visibility);
 
-  bool checkCregBinding(dxil::ResourceInfo::ResourceBinding Binding);
+  bool checkCRegBinding(dxil::ResourceInfo::ResourceBinding Binding) {
+    return CRegBindingsMap.overlaps(
+        combineUint32ToUint64(Binding.Space, Binding.LowerBound),
+        combineUint32ToUint64(Binding.Space,
+                              Binding.LowerBound + Binding.Size - 1));
+  }
 
-  bool checkTRegBinding(dxil::ResourceInfo::ResourceBinding Binding);
+  bool checkTRegBinding(dxil::ResourceInfo::ResourceBinding Binding) {
+    return TRegBindingsMap.overlaps(
+        combineUint32ToUint64(Binding.Space, Binding.LowerBound),
+        combineUint32ToUint64(Binding.Space,
+                              Binding.LowerBound + Binding.Size - 1));
+  }
 
-  bool checkURegBinding(dxil::ResourceInfo::ResourceBinding Binding);
+  bool checkURegBinding(dxil::ResourceInfo::ResourceBinding Binding) {
+    return URegBindingsMap.overlaps(
+        combineUint32ToUint64(Binding.Space, Binding.LowerBound),
+        combineUint32ToUint64(Binding.Space,
+                              Binding.LowerBound + Binding.Size - 1));
+  }
+
+  bool checkSamplerBinding(dxil::ResourceInfo::ResourceBinding Binding) {
+    return SamplersBindingsMap.overlaps(
+        combineUint32ToUint64(Binding.Space, Binding.LowerBound),
+        combineUint32ToUint64(Binding.Space,
+                              Binding.LowerBound + Binding.Size - 1));
+  }
 };
 
 class DXILPostOptimizationValidation
diff --git a/llvm/lib/Target/DirectX/DXILRootSignature.cpp b/llvm/lib/Target/DirectX/DXILRootSignature.cpp
@@ -635,8 +635,9 @@ PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
 
 //===----------------------------------------------------------------------===//
 bool RootSignatureAnalysisWrapper::runOnModule(Module &M) {
-  FuncToRsMap = std::make_unique<RootSignatureBindingInfo>(
-      RootSignatureBindingInfo(analyzeModule(M)));
+  if (!FuncToRsMap)
+    FuncToRsMap = std::make_unique<RootSignatureBindingInfo>(
+        RootSignatureBindingInfo(analyzeModule(M)));
   return false;
 }
 
diff --git a/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-DescriptorTable-AllValidFlagCombinationsV1.ll b/llvm/test/CodeGen/DirectX/ContainerData/RootSignature-DescriptorTable-AllValidFlagCombinationsV1.ll
@@ -13,7 +13,7 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 !2 = !{ ptr @main, !3, i32 1 } ; function, root signature
 !3 = !{ !5 } ; list of root signature elements
 !5 = !{ !"DescriptorTable", i32 0, !6, !7 }
-!6 = !{ !"Sampler", i32 0, i32 1, i32 0, i32 -1, i32 1 }
+!6 = !{ !"Sampler", i32 1, i32 1, i32 0, i32 -1, i32 1 }
 !7 = !{ !"UAV", i32 5, i32 1, i32 10, i32 5, i32 3 }
 
 
@@ -33,7 +33,7 @@ attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
 ; DXC-NEXT:             RangesOffset:    44
 ; DXC-NEXT:             Ranges:
 ; DXC-NEXT:               - RangeType:       3
-; DXC-NEXT:                 NumDescriptors:  0
+; DXC-NEXT:                 NumDescriptors:  1
 ; DXC-NEXT:                 BaseShaderRegister: 1
 ; DXC-NEXT:                 RegisterSpace:   0
 ; DXC-NEXT:                 OffsetInDescriptorsFromTableStart: 4294967295
diff --git a/llvm/test/CodeGen/DirectX/llc-pipeline.ll b/llvm/test/CodeGen/DirectX/llc-pipeline.ll
@@ -31,6 +31,7 @@
 ; 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 Prepare Module
