JIT: Validate loop invariant base case as part of IV analysis (#97122)

Move the validation of the base case of the IV analysis into
`FlowGraphNaturalLoop::AnalyzeIteration`. Previously the validation for
the base case was left up to the users:

- Loop cloning tried to accomplish this by checking whether loop
  inversion had run, but this is not sufficient (#96623) as nothing
  guarantees the loop is dominated by the inverted test. Loop cloning
  also skipped the check entirely for GDV, leading to #95315.

- Unrolling completely neglected to check the condition leading to
  #96591. Furthermore, unrolling made implicit assumptions that any
  `BBJ_COND` init block was an inverted test and downright removed the
  condition without any checks. This also led to another issue #97040
  where unrolling could uncover new natural loops that had not been
  canonicalized.

This change makes `FlowGraphNaturalLoop::AnalyzeIteration` try to prove
that the loop invariant it returns is true in the base case as well as
being maintained by the loop. If it cannot prove this then it fails.
This fixes all the issues, but sadly uncovers that we were doing a lot
of cloning in OSR methods without actually proving legality. We no
longer clone in these cases, but we can look into what to do about them
separately.

Fix #95315
Fix #96591
Fix #96623
Fix #97040

Author: jakobbotsch

src/coreclr/jit/compiler.h

@@ -2006,14 +2006,16 @@ struct NaturalLoopIterInfo
     // The local that is the induction variable.
     unsigned IterVar = BAD_VAR_NUM;
 
+#ifdef DEBUG
+    // Tree that initializes induction varaible outside the loop.
+    // Only valid if HasConstInit is true.
+    GenTree* InitTree = nullptr;
+#endif
+
     // Constant value that the induction variable is initialized with, outside
     // the loop. Only valid if HasConstInit is true.
     int ConstInitValue = 0;
 
-    // Block outside the loop that initializes the induction variable. Only
-    // valid if HasConstInit is true.
-    BasicBlock* InitBlock = nullptr;
-
     // Tree that has the loop test for the induction variable.
     GenTree* TestTree = nullptr;
 
@@ -2023,6 +2025,9 @@ struct NaturalLoopIterInfo
     // Tree that mutates the induction variable.
     GenTree* IterTree = nullptr;
 
+    // Is the loop exited when TestTree is true?
+    bool ExitedOnTrue : 1;
+
     // Whether or not we found an initialization of the induction variable.
     bool HasConstInit : 1;
 
@@ -2042,7 +2047,8 @@ struct NaturalLoopIterInfo
     bool HasArrayLengthLimit : 1;
 
     NaturalLoopIterInfo()
-        : HasConstInit(false)
+        : ExitedOnTrue(false)
+        , HasConstInit(false)
         , HasConstLimit(false)
         , HasSimdLimit(false)
         , HasInvariantLocalLimit(false)
@@ -2053,7 +2059,6 @@ struct NaturalLoopIterInfo
     int IterConst();
     genTreeOps IterOper();
     var_types IterOperType();
-    bool IsReversed();
     genTreeOps TestOper();
     bool IsIncreasingLoop();
     bool IsDecreasingLoop();
@@ -2062,6 +2067,9 @@ struct NaturalLoopIterInfo
     int ConstLimit();
     unsigned VarLimit();
     bool ArrLenLimit(Compiler* comp, ArrIndex* index);
+
+private:
+    bool IsReversed();
 };
 
 // Represents a natural loop in the flow graph. Natural loops are characterized
@@ -2136,7 +2144,9 @@ class FlowGraphNaturalLoop
 
     void MatchInit(NaturalLoopIterInfo* info, BasicBlock* initBlock, GenTree* init);
     bool MatchLimit(NaturalLoopIterInfo* info, GenTree* test);
-
+    bool CheckLoopConditionBaseCase(BasicBlock* initBlock, NaturalLoopIterInfo* info);
+    template<typename T>
+    static bool EvaluateRelop(T op1, T op2, genTreeOps oper);
 public:
     BasicBlock* GetHeader() const
     {
@@ -7201,7 +7211,6 @@ class Compiler
                            var_types  iterType,
                            genTreeOps testOper,
                            bool       unsignedTest,
-                           bool       dupCond,
                            unsigned*  iterCount);
 
 protected:

src/coreclr/jit/flowgraph.cpp

@@ -4933,15 +4933,31 @@ GenTreeLclVarCommon* FlowGraphNaturalLoop::FindDef(unsigned lclNum)
 //   otherwise false.
 //
 // Remarks:
-//   The function guarantees that at all definitions of the iteration local,
-//   the loop condition is reestablished before iteration continues. In other
-//   words, if this function returns true the loop invariant is guaranteed to
-//   be upheld in all blocks in the loop (but see below for establishing the
-//   base case).
-//
-//   Currently we do not validate that there is a zero-trip test ensuring the
-//   condition is true, so it is up to the user to validate that. This is
-//   normally done via GTF_RELOP_ZTT set by loop inversion.
+//   On a true return, the function guarantees that the loop invariant is true
+//   and maintained at all points within the loop, except possibly right after
+//   the update of the iterator variable (NaturalLoopIterInfo::IterTree). The
+//   function guarantees that the test (NaturalLoopIterInfo::TestTree) occurs
+//   immediately after the update, so no IR in the loop is executed without the
+//   loop invariant being true, except for the test.
+//
+//   The loop invariant is defined as the expression obtained by
+//   [info->IterVar] [info->TestOper()] [info->Limit()]. Note that
+//   [info->TestTree()] may not be of this form; it could for instance have the
+//   iterator variable as the second operand. However,
+//   [NaturalLoopIterInfo::TestOper()] will automatically normalize the test
+//   oper so that the invariant is equivalent to the returned form that has the
+//   iteration variable as op1 and the limit as op2.
+//
+//   The limit can be further decomposed via NaturalLoopIterInfo::ConstLimit,
+//   ::VarLimit and ::ArrLenLimit.
+//
+//   As an example, if info->IterVar == V02, info->TestOper() == GT_LT and
+//   info->ConstLimit() == 10, then the function guarantees that the value of
+//   the local V02 is less than 10 everywhere within the loop (except possibly
+//   at the test).
+//
+//   In some cases we also know the initial value on entry to the loop; see
+//   ::HasConstInit and ::ConstInitValue.
 //
 bool FlowGraphNaturalLoop::AnalyzeIteration(NaturalLoopIterInfo* info)
 {
@@ -4972,8 +4988,15 @@ bool FlowGraphNaturalLoop::AnalyzeIteration(NaturalLoopIterInfo* info)
         return false;
     }
 
-    info->TestBlock = cond;
-    info->IterVar   = comp->optIsLoopIncrTree(info->IterTree);
+    assert(ContainsBlock(cond->GetTrueTarget()) != ContainsBlock(cond->GetFalseTarget()));
+
+    info->TestBlock    = cond;
+    info->IterVar      = comp->optIsLoopIncrTree(info->IterTree);
+    info->ExitedOnTrue = !ContainsBlock(cond->GetTrueTarget());
+
+    // TODO-CQ: Currently, since we only look at the lexically bottom most block all loops have
+    // ExitedOnTrue == false. Once we recognize more structures this can be true.
+    assert(!info->ExitedOnTrue);
 
     assert(info->IterVar != BAD_VAR_NUM);
     LclVarDsc* const iterVarDsc = comp->lvaGetDesc(info->IterVar);
@@ -5027,13 +5050,20 @@ bool FlowGraphNaturalLoop::AnalyzeIteration(NaturalLoopIterInfo* info)
         return false;
     }
 
+    if (!CheckLoopConditionBaseCase(initBlock, info))
+    {
+        JITDUMP("  Loop condition is not always true\n");
+        return false;
+    }
+
 #ifdef DEBUG
     if (comp->verbose)
     {
         printf("  IterVar = V%02u\n", info->IterVar);
 
         if (info->HasConstInit)
-            printf("  Const init with value %d in " FMT_BB "\n", info->ConstInitValue, info->InitBlock->bbNum);
+            printf("  Const init with value %d (at [%06u])\n", info->ConstInitValue,
+                   Compiler::dspTreeID(info->InitTree));
 
         printf("  Test is [%06u] (", Compiler::dspTreeID(info->TestTree));
         if (info->HasConstLimit)
@@ -5075,7 +5105,7 @@ void FlowGraphNaturalLoop::MatchInit(NaturalLoopIterInfo* info, BasicBlock* init
 
     info->HasConstInit   = true;
     info->ConstInitValue = (int)initValue->AsIntCon()->IconValue();
-    info->InitBlock      = initBlock;
+    INDEBUG(info->InitTree = init);
 }
 
 //------------------------------------------------------------------------
@@ -5118,12 +5148,12 @@ bool FlowGraphNaturalLoop::MatchLimit(NaturalLoopIterInfo* info, GenTree* test)
     GenTree* limitOp;
 
     // Make sure op1 or op2 is the iterVar.
-    if (opr1->gtOper == GT_LCL_VAR && opr1->AsLclVarCommon()->GetLclNum() == info->IterVar)
+    if (opr1->OperIsScalarLocal() && (opr1->AsLclVarCommon()->GetLclNum() == info->IterVar))
     {
         iterOp  = opr1;
         limitOp = opr2;
     }
-    else if (opr2->gtOper == GT_LCL_VAR && opr2->AsLclVarCommon()->GetLclNum() == info->IterVar)
+    else if (opr2->OperIsScalarLocal() && (opr2->AsLclVarCommon()->GetLclNum() == info->IterVar))
     {
         iterOp  = opr2;
         limitOp = opr1;
@@ -5138,11 +5168,8 @@ bool FlowGraphNaturalLoop::MatchLimit(NaturalLoopIterInfo* info, GenTree* test)
         return false;
     }
 
-    // Mark the iterator node.
-    iterOp->gtFlags |= GTF_VAR_ITERATOR;
-
     // Check what type of limit we have - constant, variable or arr-len.
-    if (limitOp->gtOper == GT_CNS_INT)
+    if (limitOp->IsCnsIntOrI())
     {
         info->HasConstLimit = true;
         if ((limitOp->gtFlags & GTF_ICON_SIMD_COUNT) != 0)
@@ -5212,6 +5239,105 @@ bool FlowGraphNaturalLoop::MatchLimit(NaturalLoopIterInfo* info, GenTree* test)
     return true;
 }
 
+//------------------------------------------------------------------------
+// EvaluateRelop: Evaluate a relational operator with constant arguments.
+//
+// Parameters:
+//   op1  - First operand
+//   op2  - Second operand
+//   oper - Operator
+//
+// Returns:
+//   Result.
+//
+template <typename T>
+bool FlowGraphNaturalLoop::EvaluateRelop(T op1, T op2, genTreeOps oper)
+{
+    switch (oper)
+    {
+        case GT_EQ:
+            return op1 == op2;
+
+        case GT_NE:
+            return op1 != op2;
+
+        case GT_LT:
+            return op1 < op2;
+
+        case GT_LE:
+            return op1 <= op2;
+
+        case GT_GT:
+            return op1 > op2;
+
+        case GT_GE:
+            return op1 >= op2;
+
+        default:
+            unreached();
+    }
+}
+
+//------------------------------------------------------------------------
+// CheckLoopConditionBaseCase: Verify that the loop condition is true when the
+// loop is entered (or validated immediately on entry).
+//
+// Returns:
+//   True if we could prove that the condition is true at all interesting
+//   points in the loop.
+//
+// Remarks:
+//   Currently handles the following cases:
+//     * The condition being trivially true in the first iteration (e.g. for (int i = 0; i < 3; i++))
+//     * The condition is checked before entry (often due to loop inversion)
+//
+bool FlowGraphNaturalLoop::CheckLoopConditionBaseCase(BasicBlock* initBlock, NaturalLoopIterInfo* info)
+{
+    // TODO: A common loop idiom is to enter the loop at the test, with the
+    // unique in-loop predecessor of the header block being the increment. We
+    // currently do not handle these patterns in `optExtractInitTestIncr`.
+    // Instead we depend on loop inversion to put them into an
+    //   if (x) { do { ... } while (x) }
+    // form. Once we handle the pattern in `optExtractInitTestIncr` we can
+    // handle it here by checking for whether the test is the header and first
+    // thing in the header.
+
+    // Is it trivially true?
+    if (info->HasConstInit && info->HasConstLimit)
+    {
+        int initVal  = info->ConstInitValue;
+        int limitVal = info->ConstLimit();
+
+        assert(genActualType(info->TestTree->gtGetOp1()) == TYP_INT);
+
+        bool isTriviallyTrue = info->TestTree->IsUnsigned()
+                                   ? EvaluateRelop<uint32_t>((uint32_t)initVal, (uint32_t)limitVal, info->TestOper())
+                                   : EvaluateRelop<int32_t>(initVal, limitVal, info->TestOper());
+
+        if (isTriviallyTrue)
+        {
+            JITDUMP("  Condition is trivially true on entry (%d %s%s %d)\n", initVal,
+                    info->TestTree->IsUnsigned() ? "(uns)" : "", GenTree::OpName(info->TestOper()), limitVal);
+            return true;
+        }
+    }
+
+    // Do we have a zero-trip test?
+    if (initBlock->KindIs(BBJ_COND))
+    {
+        GenTree* enteringJTrue = initBlock->lastStmt()->GetRootNode();
+        assert(enteringJTrue->OperIs(GT_JTRUE));
+        if (enteringJTrue->gtGetOp1()->OperIsCompare() && ((enteringJTrue->gtGetOp1()->gtFlags & GTF_RELOP_ZTT) != 0))
+        {
+            JITDUMP("  Condition is established before entry at [%06u]\n",
+                    Compiler::dspTreeID(enteringJTrue->gtGetOp1()));
+            return true;
+        }
+    }
+
+    return false;
+}
+
 //------------------------------------------------------------------------
 // GetLexicallyTopMostBlock: Get the lexically top-most block contained within
 // the loop.
@@ -5340,7 +5466,7 @@ var_types NaturalLoopIterInfo::IterOperType()
 //
 bool NaturalLoopIterInfo::IsReversed()
 {
-    return TestTree->gtGetOp2()->OperIs(GT_LCL_VAR) && ((TestTree->gtGetOp2()->gtFlags & GTF_VAR_ITERATOR) != 0);
+    return TestTree->gtGetOp2()->OperIsScalarLocal() && (TestTree->gtGetOp2()->AsLclVar()->GetLclNum() == IterVar);
 }
 
 //------------------------------------------------------------------------
@@ -5353,7 +5479,15 @@ bool NaturalLoopIterInfo::IsReversed()
 genTreeOps NaturalLoopIterInfo::TestOper()
 {
     genTreeOps op = TestTree->OperGet();
-    return IsReversed() ? GenTree::SwapRelop(op) : op;
+    if (IsReversed())
+    {
+        op = GenTree::SwapRelop(op);
+    }
+    if (ExitedOnTrue)
+    {
+        op = GenTree::ReverseRelop(op);
+    }
+    return op;
 }
 
 //------------------------------------------------------------------------

src/coreclr/jit/loopcloning.cpp

@@ -1819,6 +1819,9 @@ bool Compiler::optIsLoopClonable(FlowGraphNaturalLoop* loop, LoopCloneContext* c
 
     // Is the entry block a handler or filter start?  If so, then if we cloned, we could create a jump
     // into the middle of a handler (to go to the cloned copy.)  Reject.
+    // TODO: This seems like it can be deleted. If the header is the beginning
+    // of a handler then the loop should be fully contained within the handler,
+    // and the cloned loop will also be in the handler.
     if (bbIsHandlerBeg(loop->GetHeader()))
     {
         JITDUMP("Loop cloning: rejecting loop " FMT_LP ". Header block is a handler start.\n", loop->GetIndex());
@@ -1873,6 +1876,8 @@ bool Compiler::optIsLoopClonable(FlowGraphNaturalLoop* loop, LoopCloneContext* c
             return false;
         }
 
+        // TODO-Quirk: Can be removed, loop invariant is validated by
+        // `FlowGraphNaturalLoop::AnalyzeIteration`.
         if (!iterInfo->TestTree->OperIsCompare() || ((iterInfo->TestTree->gtFlags & GTF_RELOP_ZTT) == 0))
         {
             JITDUMP("Loop cloning: rejecting loop " FMT_LP