[llvm] [LV] Add initial legality checks for ee loops with stores (PR #145663)

[David Sherwood via llvm-commits]

================
@@ -1793,6 +1814,110 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
   return true;
 }
 
+bool LoopVectorizationLegality::canUncountedExitConditionLoadBeMoved(
+    BasicBlock *ExitingBlock) {
+  SmallVector<const SCEVPredicate *, 4> Predicates;
+  LoadInst *CriticalUncountedExitConditionLoad = nullptr;
+
+  // Try to find a load in the critical path for the uncounted exit condition.
+  // This is currently matching about the simplest form we can, expecting
+  // only one in-loop load, the result of which is directly compared against
+  // a loop-invariant value.
+  // FIXME: We're insisting on a single use for now, because otherwise we will
+  // need to make PHI nodes for other users. That can be done once the initial
+  // transform code lands.
+  if (BranchInst *Br = dyn_cast<BranchInst>(ExitingBlock->getTerminator())) {
+    // FIXME: Don't rely on operand ordering for the comparison.
+    ICmpInst *Cmp = dyn_cast<ICmpInst>(Br->getCondition());
+    if (Cmp && Cmp->hasOneUse() &&
+        TheLoop->isLoopInvariant(Cmp->getOperand(1))) {
+      LoadInst *Load = dyn_cast<LoadInst>(Cmp->getOperand(0));
+      if (Load && Load->hasOneUse() && !TheLoop->isLoopInvariant(Load)) {
+        // The following call also checks that the load address is either
+        // invariant or is an affine SCEVAddRecExpr with a constant step.
+        // In either case, we're not relying on another load.
+        // FIXME: Support gathers after first-faulting support lands.
+        if (isDereferenceableAndAlignedInLoop(Load, TheLoop, *PSE.getSE(), *DT,
+                                              AC, &Predicates)) {
+          ICFLoopSafetyInfo SafetyInfo;
+          SafetyInfo.computeLoopSafetyInfo(TheLoop);
+          // We need to know that load will be executed before we can hoist a
+          // copy out to run just before the first iteration.
+          if (SafetyInfo.isGuaranteedToExecute(*Load, DT, TheLoop))
+            CriticalUncountedExitConditionLoad = Load;
+          else
+            reportVectorizationFailure(
+                "Early exit condition load not guaranteed to execute",
+                "Cannot vectorize early exit loop when condition load is not "
+                "guaranteed to execute",
+                "EarlyExitLoadNotGuaranteed", ORE, TheLoop);
+        } else {
+          reportVectorizationFailure(
+              "Loop may fault",
+              "Cannot vectorize potentially faulting early exit loop",
+              "PotentiallyFaultingEarlyExitLoop", ORE, TheLoop);
+          return false;
+        }
+      }
+    }
+  } else {
+    reportVectorizationFailure(
+        "Unsupported control flow in early exit loop with side effects",
----------------
david-arm wrote:

I can't see this message in any of the tests, although you have a test for it - `@loop_contains_store_uncounted_exit_is_a_switch`. It suggests that we've caught the switch earlier on and already bailed out. Perhaps you can in fact completely remove the check for branches and simply assert it's a branch instead?

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