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

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;
                     }
                 }