[llvm] [LV] Vectorize uncountable early exit store loops with combined conditions (PR #205109)

[Florian Hahn via llvm-commits]

================
@@ -1667,16 +1668,32 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
   }
 
   // The latch block must have a countable exit.
-  if (isa<SCEVCouldNotCompute>(
-          PSE.getSE()->getPredicatedExitCount(TheLoop, LatchBB, &Predicates))) {
+  if (isa<SCEVCouldNotCompute>(PSE.getSE()->getPredicatedExitCount(
+          TheLoop, LatchBB, &Predicates, ScalarEvolution::SymbolicMaximum))) {
     reportVectorizationFailure(
         "Cannot determine exact exit count for latch block",
         "Cannot vectorize early exit loop",
         "UnknownLatchExitCountEarlyExitLoop", ORE, TheLoop);
     return false;
   }
-  assert(llvm::is_contained(CountableExitingBlocks, LatchBB) &&
-         "Latch block not found in list of countable exits!");
+
+  if (!is_contained(CountableExitingBlocks, LatchBB)) {
+    // If not a separate counted exit in the latch, then check for a combined
+    // countable and uncountable exit.
+    BasicBlock *TrueBB, *FalseBB;
+    if (!match(LatchBB->getTerminator(),
+               m_Br(m_c_LogicalOr(
+                        m_Value(),
+                        m_c_ICmp(m_c_Add(m_Specific(getPrimaryInduction()),
----------------
fhahn wrote:

Do we need to limit this to the primary induction? If so, would be good to add a TODO to remove the restriction. And add test cases where we reject because the OR is of an induction that is not the primary induction

https://github.com/llvm/llvm-project/pull/205109