[llvm] [LV] Add initial legality checks for early exit loops with side effects (PR #145663)

[Graham Hunter via llvm-commits]

================
@@ -1797,6 +1822,114 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
                        "backedge taken count: "
                     << *SymbolicMaxBTC << '\n');
   UncountableExitingBB = SingleUncountableExitingBlock;
+  UncountableExitWithSideEffects = HasSideEffects;
+  return true;
+}
+
+bool LoopVectorizationLegality::canUncountableExitConditionLoadBeMoved(
+    BasicBlock *ExitingBlock) {
+  LoadInst *CriticalUncountableExitConditionLoad = nullptr;
+
+  // Try to find a load in the critical path for the uncountable 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.
+  auto *Br = cast<BranchInst>(ExitingBlock->getTerminator());
+
+  using namespace llvm::PatternMatch;
+  Value *L = nullptr;
+  Value *R = nullptr;
+  if (!match(Br->getCondition(),
+             m_OneUse(m_ICmp(m_OneUse(m_Value(L)), (m_Value(R)))))) {
+    reportVectorizationFailure(
+        "Early exit loop with store but no supported condition load",
+        "NoConditionLoadForEarlyExitLoop", ORE, TheLoop);
+    return false;
+  }
+
+  // FIXME: Don't rely on operand ordering for the comparison.
+  if (!TheLoop->isLoopInvariant(R)) {
+    reportVectorizationFailure(
+        "Early exit loop with store but no supported condition load",
+        "NoConditionLoadForEarlyExitLoop", ORE, TheLoop);
+    return false;
+  }
+
+  if (auto *Load = dyn_cast<LoadInst>(L)) {
+    // Make sure that the load address is not loop invariant; we want an
+    // address calculation that we can rotate to the next vector iteration.
+    const SCEV *PtrScev = PSE.getSE()->getSCEV(Load->getPointerOperand());
+    if (PSE.getSE()->isLoopInvariant(PtrScev, TheLoop)) {
+      reportVectorizationFailure(
+          "Uncountable exit condition depends on load from invariant address",
+          "EarlyExitLoadInvariantAddress", ORE, TheLoop);
+      return false;
+    }
+
+    // The following call also checks that the load address is either
+    // invariant (which we've just ruled out) or is an affine SCEVAddRecExpr
+    // with a constant step. In either case, we're not relying on another
+    // load within the loop.
+    // FIXME: Support gathers after first-faulting load support lands.
+    SmallVector<const SCEVPredicate *, 4> Predicates;
+    if (!isDereferenceableAndAlignedInLoop(Load, TheLoop, *PSE.getSE(), *DT, AC,
+                                           &Predicates)) {
+      reportVectorizationFailure(
+          "Loop may fault",
+          "Cannot vectorize potentially faulting early exit loop",
+          "PotentiallyFaultingEarlyExitLoop", ORE, TheLoop);
+      return false;
+    }
+
+    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.
+    // FIXME: Currently, other restrictions prevent us from reaching this point
+    //        with a loop where the uncountable exit condition is determined
+    //        by a conditional load.
+    assert(SafetyInfo.isGuaranteedToExecute(*Load, DT, TheLoop) &&
+           "Unhandled control flow in uncountable exit loop with side effects");
+
+    CriticalUncountableExitConditionLoad = Load;
+  }
+
+  if (!CriticalUncountableExitConditionLoad) {
+    reportVectorizationFailure(
+        "Early exit loop with store but no supported condition load",
+        "NoConditionLoadForEarlyExitLoop", ORE, TheLoop);
+    return false;
+  }
+
+  // Prohibit any potential aliasing with any instruction in the loop which
+  // might store to memory.
+  // FIXME: Relax this constraint where possible.
+  AAResults *AA = LAIs.getAAResults();
----------------
huntergr-arm wrote:

AAResults are obtained from the analysis manager, at least in all cases I've looked at in the codebase; LVL doesn't have access to that, so I've switched it to LoopVectorize obtaining the results and passing to LVL in the constructor. I guess I could also pass in the analysis manager and instantiate here, but either way we'll be passing more into the constructor. Is this ok? Or do you have an example of obtaining AAResults without the analysis manager I can use?

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