JIT: only skip loop inversion when the IV test is the bottom test

src/coreclr/jit/optimizer.cpp

@@ -1903,14 +1903,16 @@ bool Compiler::optTryInvertWhileLoop(FlowGraphNaturalLoop* loop)
     }
 
     // There may be multiple exits, and one of the other exits may also be a
-    // latch. That latch could be preferable to leave (for example because it
-    // is an IV test). The general BBJ_COND-latch-that-exits check below
-    // subsumes the IV-test case.
-
-    // If the loop is already bottom-tested (has a BBJ_COND latch that exits the loop),
-    // there is no need to invert. Also handle the canonical multi-backedge case, where
-    // optCanonicalizeBackedges has installed a BBJ_ALWAYS latch whose in-loop predecessors
-    // are the original bottom tests.
+    // latch. So avoid inversion if the loop is already bottom-tested.
+    //
+    // We check for two cases:
+    //  1. The latch itself is a BBJ_COND that exits the loop, AND the latch is
+    //     the recognized IV test (so the latch BBJ_COND determines loop
+    //     continuation, not a body-internal early exit that happens to feed the
+    //     back-edge).
+    //  2. The latch is a BBJ_ALWAYS (installed by optCanonicalizeBackedges for
+    //     loops with originally multiple backedges) and one of its in-loop
+    //     predecessors is the recognized IV test BBJ_COND that exits the loop.
     //
     auto isExitingCondLatch = [&](BasicBlock* block) -> bool {
         if ((block == condBlock) || !block->KindIs(BBJ_COND))
@@ -1927,13 +1929,33 @@ bool Compiler::optTryInvertWhileLoop(FlowGraphNaturalLoop* loop)
         return false;
     };
 
+    // Only run AnalyzeIteration when we see a back-edge that could
+    // qualify for the bail-out (i.e. its source block, or an in-loop predecessor
+    // of a BBJ_ALWAYS canonical latch, is a candidate for being the IV test).
+    //
+    NaturalLoopIterInfo iterInfo;
+    bool                analyzedIteration = false;
+    BasicBlock*         ivTestBlock       = nullptr;
+
+    auto isIvTest = [&](BasicBlock* candidate) -> bool {
+        if (!analyzedIteration)
+        {
+            analyzedIteration = true;
+            if (loop->AnalyzeIteration(&iterInfo))
+            {
+                ivTestBlock = iterInfo.TestBlock;
+            }
+        }
+        return (ivTestBlock != nullptr) && (candidate == ivTestBlock);
+    };
+
     for (FlowEdge* const backEdge : loop->BackEdges())
     {
         BasicBlock* const latch = backEdge->getSourceBlock();
 
-        if (isExitingCondLatch(latch))
+        if (isExitingCondLatch(latch) && isIvTest(latch))
         {
-            JITDUMP("No loop-inversion for " FMT_LP "; latch " FMT_BB " already makes it bottom-tested\n",
+            JITDUMP("No loop-inversion for " FMT_LP "; IV-test latch " FMT_BB " already makes it bottom-tested\n",
                     loop->GetIndex(), latch->bbNum);
             return false;
         }
@@ -1943,10 +1965,10 @@ bool Compiler::optTryInvertWhileLoop(FlowGraphNaturalLoop* loop)
             for (FlowEdge* const predEdge : latch->PredEdges())
             {
                 BasicBlock* const pred = predEdge->getSourceBlock();
-                if (loop->ContainsBlock(pred) && isExitingCondLatch(pred))
+                if (loop->ContainsBlock(pred) && isExitingCondLatch(pred) && isIvTest(pred))
                 {
-                    JITDUMP("No loop-inversion for " FMT_LP "; predecessor " FMT_BB " of canonical latch " FMT_BB
-                            " already makes it bottom-tested\n",
+                    JITDUMP("No loop-inversion for " FMT_LP "; IV-test predecessor " FMT_BB
+                            " of canonical latch " FMT_BB " already makes it bottom-tested\n",
                             loop->GetIndex(), pred->bbNum, latch->bbNum);
                     return false;
                 }