[Fix, Test] (ShaderTranspiler): bound the whole loop nest a resource-array index marks for unrolling

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<Uint32> input = CompileCompute(kNestedLoopIndexedImageArray);
+    ASSERT_FALSE(input.empty());
+    EXPECT_TRUE(HasDynamicImageArrayIndex(input));
+
+    Vector<Uint32> 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<Uint32> input = CompileCompute(kUniformIndexedImageArray);
     ASSERT_FALSE(input.empty());

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<uint32_t>(spv::LoopControlMask::MaskNone)) {
-                            continue;
+                            break;
                         }
                         mergeInst->SetOperand(
                             2, {static_cast<uint32_t>(spv::LoopControlMask::Unroll)});
+                        nestIterations *= iterations;
                         modified = true;
                     }
                 }