From a79eadd724c82698368627a55ce575afe441985c Mon Sep 17 00:00:00 2001 From: Swung0x48 Date: Fri, 21 Aug 2026 06:57:59 -0400 Subject: [PATCH] [Fix, Test] (ShaderTranspiler): bound the whole loop nest a resource-array index marks for unrolling --- .../LegalizeResourceArrayIndexTest.cpp | 40 ++++++++++ .../LegalizeResourceArrayIndexPass.cpp | 76 +++++++++++++++---- 2 files changed, 100 insertions(+), 16 deletions(-) diff --git a/MobileGL/MG_Test/ShaderTranspiler/LegalizeResourceArrayIndexTest.cpp b/MobileGL/MG_Test/ShaderTranspiler/LegalizeResourceArrayIndexTest.cpp index 19cfb6c6..1c9cc7c3 100644 --- a/MobileGL/MG_Test/ShaderTranspiler/LegalizeResourceArrayIndexTest.cpp +++ b/MobileGL/MG_Test/ShaderTranspiler/LegalizeResourceArrayIndexTest.cpp @@ -284,6 +284,23 @@ void main() { imageStore(g_image[i], ivec2(0), vec4(1.0)); } } +)"; + + // The same fold, buried in the loop nest a real image-writing shader has: a tile walk with + // the array walk innermost. Every level's trip count is inside the per-loop budget on its + // own, so nothing but a NEST budget stops the three from multiplying. + constexpr const char* kNestedLoopIndexedImageArray = R"(#version 450 core +layout(local_size_x = 1) in; +layout(rgba32f, binding = 0) uniform writeonly image2D g_image[4]; +void main() { + for (int y = 0; y < 64; ++y) { + for (int x = 0; x < 64; ++x) { + for (int i = 0; i < 4; ++i) { + imageStore(g_image[i], ivec2(x, y), vec4(1.0)); + } + } + } +} )"; // A uniform-sourced image index: nothing can fold it, so the switch/select lowering is what @@ -426,6 +443,29 @@ TEST(LegalizeResourceArrayIndexPass, FoldsALoopIndexedImageArray) { } } +// Marking a loop for unrolling means marking every loop enclosing it - SPIRV-Tools only unrolls +// innermost loops, so an outer one is unrollable only once its children are gone - and the copies +// those levels produce MULTIPLY. Bounding each loop on its own therefore bounds nothing: with the +// per-loop cap alone this nest (64 x 64 x 4, every level inside it) folded to 16384 OpImageWrite, +// a 3.68 MB module and 3.6 s of spirv-opt on desktop x86, from twelve lines of GLSL - before +// SPIRV-Cross or the device compiler saw any of it. Spending the budget as the walk climbs stops +// at the innermost level here, and the switch lowering - whose cost is the ARRAY LENGTH, not the +// trip counts - is what legalizes anything the unroll no longer reaches. +TEST(LegalizeResourceArrayIndexPass, BoundsTheWholeLoopNestAndNotEachLoopSeparately) { + const Vector input = CompileCompute(kNestedLoopIndexedImageArray); + ASSERT_FALSE(input.empty()); + EXPECT_TRUE(HasDynamicImageArrayIndex(input)); + + Vector output; + ASSERT_TRUE(ShaderCompiler::LegalizeResourceArrayIndexingForEssl(input, output, true)); + ASSERT_FALSE(output.empty()); + // Still legalized - that is not what is being traded away. + EXPECT_FALSE(HasDynamicImageArrayIndex(output)); + EXPECT_TRUE(Validates(output)); + // ...and paid for at the budget, not at its cube. 64 is kMaxUnrolledIterations. + EXPECT_LE(CountOpcode(output, spv::Op::OpImageWrite), 64u); +} + TEST(LegalizeResourceArrayIndexPass, LowersAUniformIndexedImageWriteToASwitchAndAReadToSelects) { const Vector input = CompileCompute(kUniformIndexedImageArray); ASSERT_FALSE(input.empty()); diff --git a/MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/LegalizeResourceArrayIndexPass.cpp b/MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/LegalizeResourceArrayIndexPass.cpp index 4d729fd6..fb0d5fbb 100644 --- a/MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/LegalizeResourceArrayIndexPass.cpp +++ b/MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/LegalizeResourceArrayIndexPass.cpp @@ -48,11 +48,15 @@ namespace MobileGL { // the round budget bounds the work on a pathological module. constexpr int kMaxLoweringRounds = 256; // Full unrolling copies the body once per iteration, and nothing in the stock - // unroller bounds that. Past this count the loop is left alone and the switch - // lowering, whose cost is the array length rather than the trip count, takes - // it instead. A loop over an array of storage blocks or of images iterates at - // most GL_MAX_*_SHADER_STORAGE_BLOCKS / GL_MAX_*_IMAGE_UNIFORMS times in any + // unroller bounds that. Past this many copies the nest is left alone and the + // switch lowering, whose cost is the array length rather than the trip count, + // takes it instead. A loop over an array of storage blocks or of images iterates + // at most GL_MAX_*_SHADER_STORAGE_BLOCKS / GL_MAX_*_IMAGE_UNIFORMS times in any // shader that is not already broken. + // + // This is a budget for the whole NEST, not for one loop: marking a loop for + // unrolling means marking its ancestors too (see MarkLoopsForUnroll), and the + // copies they produce multiply. constexpr size_t kMaxUnrolledIterations = 64; struct ResourceArray { @@ -286,12 +290,17 @@ namespace MobileGL { return false; } - // Whether fully unrolling |loop| is bounded work. The trip count is read the - // same way the stock unroller reads it, so a loop this declines to measure is - // one CanPerformUnroll would refuse anyway - the hint would be inert on it, - // and the fallback lowering is what handles it. Requires the induction - // variable to already be an OpPhi, which is why this runs after ssa-rewrite. - bool IsBoundedUnrollCandidate(spvtools::opt::Loop* loop) { + // |loop|'s trip count, when it has a measurable one, in *outIterations. The + // count is read the same way the stock unroller reads it, so a loop this + // declines to measure is one CanPerformUnroll would refuse anyway - the hint + // would be inert on it, and the fallback lowering is what handles it. Requires + // the induction variable to already be an OpPhi, which is why this runs after + // ssa-rewrite. + // + // A count of zero is reported as unmeasurable: it means nothing this pass can + // multiply a nest's budget by, and a loop that never runs is not one whose + // subscript needs folding. + bool TryGetUnrollTripCount(spvtools::opt::Loop* loop, size_t* outIterations) { const spvtools::opt::BasicBlock* condition = loop->FindConditionBlock(); if (condition == nullptr) { return false; @@ -301,10 +310,12 @@ namespace MobileGL { return false; } size_t iterations = 0; - if (!loop->FindNumberOfIterations(induction, &*condition->ctail(), &iterations)) { + if (!loop->FindNumberOfIterations(induction, &*condition->ctail(), &iterations) || + iterations == 0) { return false; } - return iterations <= kMaxUnrolledIterations; + *outIterations = iterations; + return true; } } // namespace @@ -345,25 +356,58 @@ namespace MobileGL { continue; } + // The offending chain's own loop AND every loop enclosing it, because + // SPIRV-Tools only ever unrolls an INNERMOST loop - an outer one is + // unrollable only once its children are gone - and because the index that + // has to become a literal may be an outer loop's induction variable. + // + // Marking a whole nest means the unrolled body count is the PRODUCT of its + // trip counts, not the largest of them, so the budget is spent as the walk + // climbs rather than tested loop by loop. Measured through + // LegalizeResourceArrayIndexingForEssl itself, on a twelve-line shader + // writing image2D g_image[4] from a 64/64/4 nest - every level individually + // inside the per-loop cap, which is all this used to test: 256 OpImageWrite + // with the per-loop cap alone against 4 with the nest budget, and the + // per-loop cap bounds nothing at all as the trip counts grow. The image half + // of this pass is what made such a nest reachable; a storage-block array + // rarely sits inside one. BoundsTheWholeLoopNestAndNotEachLoopSeparately is + // that measurement. + // + // Every exit is a BREAK rather than a skip-and-keep-climbing. A loop that + // cannot be marked - unmeasurable, out of budget, or carrying a control this + // pass will not overwrite - is a gap the unroller cannot cross, which makes + // every mark above it dead weight on a module that will not be unrolled + // anyway. Falling out of the unroll path costs nothing correctness-wise: + // LowerToConstantSwitch still legalizes the chain, at a cost proportional to + // the ARRAY LENGTH rather than to the trip counts. spvtools::opt::LoopDescriptor* loops = irContext->GetLoopDescriptor(function); + size_t nestIterations = 1; for (spvtools::opt::Loop* loop = (*loops)[block->id()]; loop != nullptr; loop = loop->GetParent()) { - if (!IsBoundedUnrollCandidate(loop)) { - continue; + size_t iterations = 0; + if (!TryGetUnrollTripCount(loop, &iterations)) { + break; + } + // Division, not multiplication, so the test itself cannot overflow. + if (iterations > kMaxUnrolledIterations / nestIterations) { + break; } Instruction* mergeInst = loop->GetHeaderBlock()->GetLoopMergeInst(); // Only a bare `None` control is promoted, and only when no extra // literal (PartialCount, PeelCount, ...) follows it: the unroller // tests the control word for equality with Unroll, so ORing the bit // into a control that already carries something - DontUnroll above - // all - would neither unroll nor mean what it says. + // all - would neither unroll nor mean what it says. An `Unroll` this + // pass itself already wrote for another chain in the same nest ends the + // walk too: everything above it was considered on that pass through. if (mergeInst == nullptr || mergeInst->NumOperands() != 3 || mergeInst->GetSingleWordOperand(2) != static_cast(spv::LoopControlMask::MaskNone)) { - continue; + break; } mergeInst->SetOperand( 2, {static_cast(spv::LoopControlMask::Unroll)}); + nestIterations *= iterations; modified = true; } }