[SemaHLSL] Correct descriptor range overflow validation

[inbelic]

This pr corrects the validation behaviour to allow valid root signatures of the form:
DescriptorTable(CBV(b0, offset = 4294967294), CBV(b1))
which will append a range onto the location of UINT_MAX, which is valid.

Resolves: #159478.

[llvmbot]

@llvm/pr-subscribers-hlsl

@llvm/pr-subscribers-clang

Author: Finn Plummer (inbelic)

Changes

This pr corrects the validation behaviour to allow valid root signatures of the form:
DescriptorTable(CBV(b0, offset = 4294967294), CBV(b1))
which will append a range onto the location of UINT_MAX, which is valid.

---

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

5 Files Affected:

• (modified) clang/lib/Sema/SemaHLSL.cpp (+8-4)
• (modified) clang/test/SemaHLSL/RootSignature-resource-ranges-err.hlsl (+5-1)
• (modified) clang/test/SemaHLSL/RootSignature-resource-ranges.hlsl (+3)
• (modified) llvm/include/llvm/Frontend/HLSL/RootSignatureValidations.h (-1)
• (modified) llvm/lib/Frontend/HLSL/RootSignatureValidations.cpp (-4)

diff --git a/clang/lib/Sema/SemaHLSL.cpp b/clang/lib/Sema/SemaHLSL.cpp
index 0af38472b0fec..33edd7f9a916a 100644
--- a/clang/lib/Sema/SemaHLSL.cpp
+++ b/clang/lib/Sema/SemaHLSL.cpp
@@ -1361,6 +1361,7 @@ bool SemaHLSL::handleRootSignatureElements(
       bool HasAnySampler = false;
       bool HasAnyNonSampler = false;
       uint32_t Offset = 0;
+      bool Unbound = false;
       for (const auto &[Clause, ClauseElem] : UnboundClauses) {
         SourceLocation Loc = ClauseElem->getLocation();
         if (Clause->Type == llvm::dxil::ResourceClass::Sampler)
@@ -1381,15 +1382,16 @@ bool SemaHLSL::handleRootSignatureElements(
             llvm::hlsl::rootsig::DescriptorTableOffsetAppend) {
           // Manually specified the offset
           Offset = Clause->Offset;
+          Unbound = false;
+        } else if (Unbound) {
+          // Trying to append onto unbound offset
+          Diag(Loc, diag::err_hlsl_appending_onto_unbound);
         }
 
         uint64_t RangeBound = llvm::hlsl::rootsig::computeRangeBound(
             Offset, Clause->NumDescriptors);
 
-        if (!llvm::hlsl::rootsig::verifyBoundOffset(Offset)) {
-          // Trying to append onto unbound offset
-          Diag(Loc, diag::err_hlsl_appending_onto_unbound);
-        } else if (!llvm::hlsl::rootsig::verifyNoOverflowedOffset(RangeBound)) {
+        if (!llvm::hlsl::rootsig::verifyNoOverflowedOffset(RangeBound)) {
           // Upper bound overflows maximum offset
           Diag(Loc, diag::err_hlsl_offset_overflow) << Offset << RangeBound;
         }
@@ -1397,6 +1399,8 @@ bool SemaHLSL::handleRootSignatureElements(
         Offset = RangeBound == llvm::hlsl::rootsig::NumDescriptorsUnbounded
                      ? uint32_t(RangeBound)
                      : uint32_t(RangeBound + 1);
+        Unbound = Clause->NumDescriptors ==
+                  llvm::hlsl::rootsig::NumDescriptorsUnbounded;
 
         // Compute the register bounds and track resource binding
         uint32_t LowerBound(Clause->Reg.Number);
diff --git a/clang/test/SemaHLSL/RootSignature-resource-ranges-err.hlsl b/clang/test/SemaHLSL/RootSignature-resource-ranges-err.hlsl
index 2d025d0e6e5ce..5ce4419149592 100644
--- a/clang/test/SemaHLSL/RootSignature-resource-ranges-err.hlsl
+++ b/clang/test/SemaHLSL/RootSignature-resource-ranges-err.hlsl
@@ -141,4 +141,8 @@ void append_offset_overflow_signature() {}
 
 // expected-error@+1 {{descriptor range offset overflows [4294967292, 4294967296]}}
 [RootSignature("DescriptorTable(CBV(b0, offset = 4294967292, numDescriptors = 5))")]
-void offset_() {}
+void offset_overflow() {}
+
+// expected-error@+1 {{descriptor range offset overflows [4294967295, 4294967296]}}
+[RootSignature("DescriptorTable(CBV(b0, offset = 4294967294), CBV(b1, numDescriptors = 2))")]
+void appended_offset_overflow() {}
diff --git a/clang/test/SemaHLSL/RootSignature-resource-ranges.hlsl b/clang/test/SemaHLSL/RootSignature-resource-ranges.hlsl
index 10e7215eccf6e..37bb4f173aac4 100644
--- a/clang/test/SemaHLSL/RootSignature-resource-ranges.hlsl
+++ b/clang/test/SemaHLSL/RootSignature-resource-ranges.hlsl
@@ -25,3 +25,6 @@ void valid_root_signature_6() {}
 
 [RootSignature("DescriptorTable(CBV(b0, offset = 4294967292), CBV(b1, numDescriptors = 3))")]
 void valid_root_signature_7() {}
+
+[RootSignature("DescriptorTable(CBV(b0, offset = 4294967294), CBV(b1))")]
+void valid_root_signature_8() {}
diff --git a/llvm/include/llvm/Frontend/HLSL/RootSignatureValidations.h b/llvm/include/llvm/Frontend/HLSL/RootSignatureValidations.h
index ea96094b18300..45353142267a6 100644
--- a/llvm/include/llvm/Frontend/HLSL/RootSignatureValidations.h
+++ b/llvm/include/llvm/Frontend/HLSL/RootSignatureValidations.h
@@ -38,7 +38,6 @@ LLVM_ABI bool verifyMipLODBias(float MipLODBias);
 LLVM_ABI bool verifyMaxAnisotropy(uint32_t MaxAnisotropy);
 LLVM_ABI bool verifyLOD(float LOD);
 
-LLVM_ABI bool verifyBoundOffset(uint32_t Offset);
 LLVM_ABI bool verifyNoOverflowedOffset(uint64_t Offset);
 LLVM_ABI uint64_t computeRangeBound(uint32_t Offset, uint32_t Size);
 
diff --git a/llvm/lib/Frontend/HLSL/RootSignatureValidations.cpp b/llvm/lib/Frontend/HLSL/RootSignatureValidations.cpp
index 0970977b5064f..f3ea1698f3bf8 100644
--- a/llvm/lib/Frontend/HLSL/RootSignatureValidations.cpp
+++ b/llvm/lib/Frontend/HLSL/RootSignatureValidations.cpp
@@ -125,10 +125,6 @@ bool verifyMaxAnisotropy(uint32_t MaxAnisotropy) {
 
 bool verifyLOD(float LOD) { return !std::isnan(LOD); }
 
-bool verifyBoundOffset(uint32_t Offset) {
-  return Offset != NumDescriptorsUnbounded;
-}
-
 bool verifyNoOverflowedOffset(uint64_t Offset) {
   return Offset <= std::numeric_limits<uint32_t>::max();
 }

[tex3d]

I think there may be a gap in this implementation allowing binding of overlapping size 1 range at UINT_MAX Offset; plus some other nits; see comments.

[tex3d]

The code in this vicinity does seem to have a few minor nits that make it harder to reason about:

• Offset will be set to the next available offset, unless that's UINT_MAX, so could there be a bug where a prior not-unbounded Clause ends at UINT_MAX, and now the next Clause of size 1 appears to fit starting at UINT_MAX? I think Offset should be a uint64_t and always point to the next location. Otherwise, every value for a uint32_t Offset is a valid location for one descriptor to fit. If that's larger than UINT_MAX, the next one won't fit.
• Inconsistent use of llvm::hlsl::rootsig::NumDescriptorsUnbounded vs. ~0u when comparing Clause->NumDescriptors. We should use the same llvm::hlsl::rootsig::NumDescriptorsUnbounded everywhere.
• You use llvm::hlsl::rootsig::computeRangeBound to compute upper bound elsewhere, but you duplicate that logic to initialize UpperBound below. Why not use computeRangeBound for UpperBound too?

[tex3d]

You'll want another test for the gap I think I noticed:

// expected-error@+1 {{descriptor range offset overflows [4294967296, 4294967296]}}
[RootSignature("DescriptorTable(CBV(b0, offset = 4294967294), CBV(b1), CBV(b2))")]
void appended_offset_overflow2() {}

[joaosaffran]

nit: Since that reference to the previous clause, maybe something like: IsPrevUnbound, would make it clearer, IMHO

[joaosaffran]

I think you might be able to remove the need to track Unbound, if you check if the offset is overflowing, after updating the offset and before calculating rangeBound, something like:

if(Clause->Offset != llvm::hlsl::rootsig::DescriptorTableOffsetAppend)
    Offset = Clause->Offset;

if (!llvm::hlsl::rootsig::verifyNoOverflowedOffset(Offset))
    Diag(Loc, diag::err_hlsl_appending_onto_unbound);
// From this point onwards, Offset is not overflowing

uint64_t RangeBound = llvm::hlsl::rootsig::computeRangeBound(
            Offset, Clause->NumDescriptors);

if (!llvm::hlsl::rootsig::verifyNoOverflowedOffset(RangeBound))
    Diag(Loc, diag::err_hlsl_offset_overflow) << Offset << RangeBound;

Offset = RangeBound + 1;