[HLSL][RootSignature] Plug into the thing

Author: inbelic

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -12967,6 +12967,11 @@ def err_hlsl_expect_arg_const_int_one_or_neg_one: Error<
 def err_invalid_hlsl_resource_type: Error<
   "invalid __hlsl_resource_t type attributes">;
 
+def err_hlsl_resource_range_overlap: Error<
+  "resource ranges %select{t|u|b|s}0[%1;%2] and %select{t|u|b|s}3[%4;%5] "
+  "overlap within space = %6 and visibility = "
+  "%select{All|Vertex|Hull|Domain|Geometry|Pixel|Amplification|Mesh}7">;
+
 // Layout randomization diagnostics.
 def err_non_designated_init_used : Error<
   "a randomized struct can only be initialized with a designated initializer">;

clang/include/clang/Sema/SemaHLSL.h

@@ -119,6 +119,8 @@ class SemaHLSL : public SemaBase {
                                        bool IsCompAssign);
   void emitLogicalOperatorFixIt(Expr *LHS, Expr *RHS, BinaryOperatorKind Opc);
 
+  // Returns true when D is invalid and a diagnostic was produced
+  bool handleRootSignatureDecl(HLSLRootSignatureDecl *D, SourceLocation Loc);
   void handleRootSignatureAttr(Decl *D, const ParsedAttr &AL);
   void handleNumThreadsAttr(Decl *D, const ParsedAttr &AL);
   void handleWaveSizeAttr(Decl *D, const ParsedAttr &AL);

clang/lib/Sema/SemaHLSL.cpp

@@ -39,6 +39,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/Frontend/HLSL/HLSLRootSignature.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/DXILABI.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -951,6 +952,108 @@ void SemaHLSL::emitLogicalOperatorFixIt(Expr *LHS, Expr *RHS,
       << NewFnName << FixItHint::CreateReplacement(FullRange, OS.str());
 }
 
+namespace {
+
+// A resource range overlaps with another resource range if they have:
+// - equivalent ResourceClass (SRV, UAV, CBuffer, Sampler)
+// - equivalent resource space
+// - overlapping visbility
+class ResourceRanges {
+public:
+  // KeyT: 32-lsb denotes resource space, and 32-msb denotes resource type enum
+  using KeyT = uint64_t;
+
+  static const unsigned NumVisEnums =
+      (unsigned)llvm::hlsl::rootsig::ShaderVisibility::NumEnums;
+
+private:
+  llvm::hlsl::rootsig::ResourceRange::IMap::Allocator Allocator;
+
+  // Denotes a mapping of a unique combination of ResourceClass and register
+  // space to a ResourceRange
+  using MapT = llvm::SmallDenseMap<KeyT, llvm::hlsl::rootsig::ResourceRange>;
+
+  // Denotes a mapping for each unique visibility
+  MapT RangeMaps[NumVisEnums];
+
+  constexpr static KeyT getKey(const llvm::hlsl::rootsig::RangeInfo &Info) {
+    uint64_t SpacePacked = (uint64_t)Info.Space;
+    uint64_t ClassPacked = (uint64_t)llvm::to_underlying(Info.Class);
+    return (ClassPacked << 32) | SpacePacked;
+  }
+
+public:
+  // Returns std::nullopt if there was no collision. Otherwise, it will
+  // return the RangeInfo of the collision
+  std::optional<const llvm::hlsl::rootsig::RangeInfo *>
+  addRange(const llvm::hlsl::rootsig::RangeInfo &Info) {
+    MapT &VisRangeMap = RangeMaps[llvm::to_underlying(Info.Vis)];
+    auto [It, _] = VisRangeMap.insert(
+        {getKey(Info), llvm::hlsl::rootsig::ResourceRange(Allocator)});
+    auto Res = It->second.insert(Info);
+    if (Res.has_value())
+      return Res;
+
+    MutableArrayRef<MapT> Maps =
+        Info.Vis == llvm::hlsl::rootsig::ShaderVisibility::All
+            ? MutableArrayRef<MapT>{RangeMaps}.drop_front()
+            : MutableArrayRef<MapT>{RangeMaps}.take_front();
+
+    for (MapT &CurMap : Maps) {
+      auto CurIt = CurMap.find(getKey(Info));
+      if (CurIt != CurMap.end())
+        if (auto Overlapping = CurIt->second.getOverlapping(Info))
+          return Overlapping;
+    }
+
+    return std::nullopt;
+  }
+};
+
+} // namespace
+
+bool SemaHLSL::handleRootSignatureDecl(HLSLRootSignatureDecl *D,
+                                       SourceLocation Loc) {
+  auto Elements = D->getRootElements();
+
+  // First we will go through and collect our range info
+  llvm::SmallVector<llvm::hlsl::rootsig::RangeInfo> Infos;
+  for (const auto &Elem : Elements) {
+    if (const auto *Param =
+            std::get_if<llvm::hlsl::rootsig::RootParam>(&Elem)) {
+      llvm::hlsl::rootsig::RangeInfo Info;
+      Info.LowerBound = Param->Reg.Number;
+      Info.UpperBound = Info.LowerBound; // use inclusive ranges []
+
+      Info.Class = Param->Type;
+      Info.Space = Param->Space;
+      Info.Vis = Param->Visibility;
+      Infos.push_back(Info);
+    }
+  }
+
+  // Iterate through info and attempt to insert corresponding range
+  ResourceRanges Ranges;
+  bool HadOverlap = false;
+  for (const llvm::hlsl::rootsig::RangeInfo &Info : Infos)
+    if (auto MaybeOverlappingInfo = Ranges.addRange(Info)) {
+      const llvm::hlsl::rootsig::RangeInfo *OInfo =
+          MaybeOverlappingInfo.value();
+      auto CommonVis = Info.Vis == llvm::hlsl::rootsig::ShaderVisibility::All
+                           ? OInfo->Vis
+                           : Info.Vis;
+
+      Diag(Loc, diag::err_hlsl_resource_range_overlap)
+          << llvm::to_underlying(Info.Class) << Info.LowerBound
+          << Info.UpperBound << llvm::to_underlying(OInfo->Class)
+          << OInfo->LowerBound << OInfo->UpperBound << Info.Space << CommonVis;
+
+      HadOverlap = true;
+    }
+
+  return HadOverlap;
+}
+
 void SemaHLSL::handleRootSignatureAttr(Decl *D, const ParsedAttr &AL) {
   if (AL.getNumArgs() != 1) {
     Diag(AL.getLoc(), diag::err_attribute_wrong_number_arguments) << AL << 1;
@@ -973,6 +1076,8 @@ void SemaHLSL::handleRootSignatureAttr(Decl *D, const ParsedAttr &AL) {
     if (auto *SignatureDecl =
             dyn_cast<HLSLRootSignatureDecl>(R.getFoundDecl())) {
       // Perform validation of constructs here
+      if (handleRootSignatureDecl(SignatureDecl, AL.getLoc()))
+        return;
       D->addAttr(::new (getASTContext()) RootSignatureAttr(
           getASTContext(), AL, Ident, SignatureDecl));
     }

clang/test/SemaHLSL/RootSignature-resource-ranges-err.hlsl

@@ -0,0 +1,26 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -o - %s -verify
+
+#define Overlap0 "CBV(b42), CBV(b42)"
+
+[RootSignature(Overlap0)] // expected-error {{resource ranges b[42;42] and b[42;42] overlap within space = 0 and visibility = All}}
+void bad_root_signature_0() {}
+
+#define Overlap1 "SRV(t0, space = 3), SRV(t0, space = 3)"
+
+[RootSignature(Overlap1)] // expected-error {{resource ranges t[0;0] and t[0;0] overlap within space = 3 and visibility = All}}
+void bad_root_signature_1() {}
+
+#define Overlap2 "UAV(u0, visibility = SHADER_VISIBILITY_PIXEL), UAV(u0, visibility = SHADER_VISIBILITY_PIXEL)"
+
+[RootSignature(Overlap2)] // expected-error {{resource ranges u[0;0] and u[0;0] overlap within space = 0 and visibility = Pixel}}
+void bad_root_signature_2() {}
+
+#define Overlap3 "UAV(u0, visibility = SHADER_VISIBILITY_ALL), UAV(u0, visibility = SHADER_VISIBILITY_PIXEL)"
+
+[RootSignature(Overlap3)] // expected-error {{resource ranges u[0;0] and u[0;0] overlap within space = 0 and visibility = Pixel}}
+void bad_root_signature_3() {}
+
+#define Overlap4 "UAV(u0, visibility = SHADER_VISIBILITY_PIXEL), UAV(u0, visibility = SHADER_VISIBILITY_ALL)"
+
+[RootSignature(Overlap4)] // expected-error {{resource ranges u[0;0] and u[0;0] overlap within space = 0 and visibility = Pixel}}
+void bad_root_signature_4() {}

clang/test/SemaHLSL/RootSignature-resource-ranges.hlsl

@@ -0,0 +1,22 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -o - %s -verify
+// expected-no-diagnostics
+
+#define NoOverlap0 "CBV(b0), CBV(b1)"
+
+[RootSignature(NoOverlap0)]
+void valid_root_signature_0() {}
+
+#define NoOverlap1 "CBV(b0, visibility = SHADER_VISIBILITY_DOMAIN), CBV(b0, visibility = SHADER_VISIBILITY_PIXEL)"
+
+[RootSignature(NoOverlap1)]
+void valid_root_signature_1() {}
+
+#define NoOverlap2 "CBV(b0, space = 1), CBV(b0, space = 2)"
+
+[RootSignature(NoOverlap2)]
+void valid_root_signature_2() {}
+
+#define NoOverlap3 "CBV(b0), SRV(t0)"
+
+[RootSignature(NoOverlap3)]
+void valid_root_signature_3() {}

llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h

@@ -75,6 +75,7 @@ enum class ShaderVisibility {
   Pixel = 5,
   Amplification = 6,
   Mesh = 7,
+  NumEnums = 8,
 };
 
 // Definitions of the in-memory data layout structures
@@ -199,13 +200,17 @@ class MetadataBuilder {
   SmallVector<Metadata *> GeneratedMetadata;
 };
 
-// RangeInfo holds the information to correctly construct a ResourceRange
-// and retains this information to be used for displaying a better diagnostic
 struct RangeInfo {
   const static uint32_t Unbounded = static_cast<uint32_t>(-1);
 
+  // Interval information
   uint32_t LowerBound;
   uint32_t UpperBound;
+
+  // Information retained for diagnostics
+  llvm::dxil::ResourceClass Class;
+  uint32_t Space;
+  ShaderVisibility Vis;
 };
 
 class ResourceRange {

llvm/lib/Frontend/HLSL/HLSLRootSignature.cpp

@@ -66,6 +66,8 @@ static raw_ostream &operator<<(raw_ostream &OS,
   case ShaderVisibility::Mesh:
     OS << "Mesh";
     break;
+  case ShaderVisibility::NumEnums:
+    break;
   }
 
   return OS;