[VPlan] Move IV predicate handling to VPlan. (llvm#192876)

Author: fhahn

llvm/include/llvm/Analysis/IVDescriptors.h

@@ -28,6 +28,7 @@ class Loop;
 class PredicatedScalarEvolution;
 class ScalarEvolution;
 class SCEV;
+class SCEVPredicate;
 class StoreInst;
 
 /// These are the kinds of recurrences that we support.
@@ -405,9 +406,13 @@ class InductionDescriptor {
   /// analysis, it can be passed through \p Expr. If the def-use chain
   /// associated with the phi includes casts (that we know we can ignore
   /// under proper runtime checks), they are passed through \p CastsToIgnore.
+  /// SCEV predicates checking potential overflow for \p Phi to be an induction,
+  /// if any, are passed via \p NoWrapPreds and recorded.
   LLVM_ABI static bool
   isInductionPHI(PHINode *Phi, const Loop *L, ScalarEvolution *SE,
-                 InductionDescriptor &D, const SCEV *Expr = nullptr,
+                 InductionDescriptor &D,
+                 ArrayRef<const SCEVPredicate *> NoWrapPreds = {},
+                 const SCEV *Expr = nullptr,
                  SmallVectorImpl<Instruction *> *CastsToIgnore = nullptr);
 
   /// Returns true if \p Phi is a floating point induction in the loop \p L.
@@ -449,12 +454,18 @@ class InductionDescriptor {
   /// SCEV overflow check.
   ArrayRef<Instruction *> getCastInsts() const { return RedundantCasts; }
 
+  /// Returns the SCEV predicates associated with this induction.
+  ArrayRef<const SCEVPredicate *> getNoWrapPredicates() const {
+    return NoWrapPredicates;
+  }
+
 private:
   /// Private constructor - used by \c isInductionPHI and
   /// \c getCanonicalIntInduction.
   InductionDescriptor(Value *Start, InductionKind K, const SCEV *Step,
                       BinaryOperator *InductionBinOp = nullptr,
-                      SmallVectorImpl<Instruction *> *Casts = nullptr);
+                      SmallVectorImpl<Instruction *> *Casts = nullptr,
+                      ArrayRef<const SCEVPredicate *> NoWrapPreds = {});
 
   /// Start value.
   TrackingVH<Value> StartValue;
@@ -467,6 +478,8 @@ class InductionDescriptor {
   // Instructions used for type-casts of the induction variable,
   // that are redundant when guarded with a runtime SCEV overflow check.
   SmallVector<Instruction *, 2> RedundantCasts;
+  // SCEV predicates checking overflow needed for this induction.
+  SmallVector<const SCEVPredicate *, 2> NoWrapPredicates;
 };
 
 } // end namespace llvm

llvm/include/llvm/Analysis/ScalarEvolution.h

@@ -2657,8 +2657,11 @@ class PredicatedScalarEvolution {
   /// Attempts to produce an AddRecExpr for V by adding additional SCEV
   /// predicates. If we can't transform the expression into an AddRecExpr we
   /// return nullptr and not add additional SCEV predicates to the current
-  /// context.
-  LLVM_ABI const SCEVAddRecExpr *getAsAddRec(Value *V);
+  /// context. If \p WrapPredsAdded is non-null, the required predicates are
+  /// collected there instead of being added to this context.
+  LLVM_ABI const SCEVAddRecExpr *
+  getAsAddRec(Value *V,
+              SmallVectorImpl<const SCEVPredicate *> *WrapPredsAdded = nullptr);
 
   /// Proves that V doesn't overflow by adding SCEV predicate.
   LLVM_ABI void setNoOverflow(Value *V,
@@ -2680,9 +2683,10 @@ class PredicatedScalarEvolution {
   LLVM_ABI void print(raw_ostream &OS, unsigned Depth) const;
 
   /// Check if \p AR1 and \p AR2 are equal, while taking into account
-  /// Equal predicates in Preds.
-  LLVM_ABI bool areAddRecsEqualWithPreds(const SCEVAddRecExpr *AR1,
-                                         const SCEVAddRecExpr *AR2) const;
+  /// Equal predicates in Preds and \p ExtraPreds.
+  LLVM_ABI bool areAddRecsEqualWithPreds(
+      const SCEVAddRecExpr *AR1, const SCEVAddRecExpr *AR2,
+      ArrayRef<const SCEVPredicate *> ExtraPreds = {}) const;
 
 private:
   /// Increments the version number of the predicate.  This needs to be called

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

@@ -659,8 +659,7 @@ class LoopVectorizationLegality {
   /// Updates the vectorization state by adding \p Phi to the inductions list.
   /// This can set \p Phi as the main induction of the loop if \p Phi is a
   /// better choice for the main induction than the existing one.
-  void addInductionPhi(PHINode *Phi, const InductionDescriptor &ID,
-                       SmallPtrSetImpl<Value *> &AllowedExit);
+  void addInductionPhi(PHINode *Phi, const InductionDescriptor &ID);
 
   /// The loop that we evaluate.
   Loop *TheLoop;
@@ -718,10 +717,6 @@ class LoopVectorizationLegality {
   /// Holds the widest induction type encountered.
   IntegerType *WidestIndTy = nullptr;
 
-  /// Allowed outside users. This holds the variables that can be accessed from
-  /// outside the loop.
-  SmallPtrSet<Value *, 4> AllowedExit;
-
   /// Vectorization requirements that will go through late-evaluation.
   LoopVectorizationRequirements *Requirements;
 

llvm/lib/Analysis/IVDescriptors.cpp

@@ -1381,9 +1381,10 @@ RecurrenceDescriptor::getReductionOpChain(PHINode *Phi, Loop *L) const {
   return ReductionOperations;
 }
 
-InductionDescriptor::InductionDescriptor(Value *Start, InductionKind K,
-                                         const SCEV *Step, BinaryOperator *BOp,
-                                         SmallVectorImpl<Instruction *> *Casts)
+InductionDescriptor::InductionDescriptor(
+    Value *Start, InductionKind K, const SCEV *Step, BinaryOperator *BOp,
+    SmallVectorImpl<Instruction *> *Casts,
+    ArrayRef<const SCEVPredicate *> NoWrapPreds)
     : StartValue(Start), IK(K), Step(Step), InductionBinOp(BOp) {
   assert(IK != IK_NoInduction && "Not an induction");
 
@@ -1412,6 +1413,7 @@ InductionDescriptor::InductionDescriptor(Value *Start, InductionKind K,
 
   if (Casts)
     llvm::append_range(RedundantCasts, *Casts);
+  llvm::append_range(NoWrapPredicates, NoWrapPreds);
 }
 
 InductionDescriptor
@@ -1511,15 +1513,22 @@ bool InductionDescriptor::isFPInductionPHI(PHINode *Phi, const Loop *TheLoop,
 /// If we are able to find such sequence, we return the instructions
 /// we found, namely %casted_phi and the instructions on its use-def chain up
 /// to the phi (not including the phi).
-static bool getCastsForInductionPHI(PredicatedScalarEvolution &PSE,
-                                    const SCEVUnknown *PhiScev,
-                                    const SCEVAddRecExpr *AR,
-                                    SmallVectorImpl<Instruction *> &CastInsts) {
+static bool
+getCastsForInductionPHI(PredicatedScalarEvolution &PSE,
+                        const SCEVUnknown *PhiScev, const SCEVAddRecExpr *AR,
+                        SmallVectorImpl<Instruction *> &CastInsts,
+                        ArrayRef<const SCEVPredicate *> NoWrapPreds) {
 
   assert(CastInsts.empty() && "CastInsts is expected to be empty.");
   auto *PN = cast<PHINode>(PhiScev->getValue());
-  assert(PSE.getSCEV(PN) == AR && "Unexpected phi node SCEV expression");
+
+  // Build a predicate to rewrite SCEVs of values in the cast chain using the
+  // predicates needed for this induction.
+  ScalarEvolution &SE = *PSE.getSE();
+  SCEVUnionPredicate NoWrapUnionPred(NoWrapPreds, SE);
   const Loop *L = AR->getLoop();
+  assert(SE.rewriteUsingPredicate(SE.getSCEV(PN), L, NoWrapUnionPred) == AR &&
+         "Unexpected phi node SCEV expression");
 
   // Find any cast instructions that participate in the def-use chain of
   // PhiScev in the loop.
@@ -1564,8 +1573,10 @@ static bool getCastsForInductionPHI(PredicatedScalarEvolution &PSE,
     if (!Inst || !L->contains(Inst)) {
       return false;
     }
-    auto *AddRec = dyn_cast<SCEVAddRecExpr>(PSE.getSCEV(Val));
-    if (AddRec && PSE.areAddRecsEqualWithPreds(AddRec, AR))
+    // Create AddRec with NoWrapPredicates applied.
+    auto *AddRec = dyn_cast<SCEVAddRecExpr>(
+        SE.rewriteUsingPredicate(SE.getSCEV(Val), L, NoWrapUnionPred));
+    if (AddRec && PSE.areAddRecsEqualWithPreds(AddRec, AR, NoWrapPreds))
       InCastSequence = true;
     if (InCastSequence) {
       // Only the last instruction in the cast sequence is expected to have
@@ -1603,9 +1614,12 @@ bool InductionDescriptor::isInductionPHI(PHINode *Phi, const Loop *TheLoop,
   const SCEV *PhiScev = PSE.getSCEV(Phi);
   const auto *AR = dyn_cast<SCEVAddRecExpr>(PhiScev);
 
+  // Collect predicates needed to force the SCEV into an AddRecExpr.
+  SmallVector<const SCEVPredicate *, 2> Preds;
+
   // We need this expression to be an AddRecExpr.
   if (Assume && !AR)
-    AR = PSE.getAsAddRec(Phi);
+    AR = PSE.getAsAddRec(Phi, &Preds);
 
   if (!AR) {
     LLVM_DEBUG(dbgs() << "LV: PHI is not a poly recurrence.\n");
@@ -1621,17 +1635,17 @@ bool InductionDescriptor::isInductionPHI(PHINode *Phi, const Loop *TheLoop,
   // induction.
   if (PhiScev != AR && SymbolicPhi) {
     SmallVector<Instruction *, 2> Casts;
-    if (getCastsForInductionPHI(PSE, SymbolicPhi, AR, Casts))
-      return isInductionPHI(Phi, TheLoop, PSE.getSE(), D, AR, &Casts);
+    if (getCastsForInductionPHI(PSE, SymbolicPhi, AR, Casts, Preds))
+      return isInductionPHI(Phi, TheLoop, PSE.getSE(), D, Preds, AR, &Casts);
   }
 
-  return isInductionPHI(Phi, TheLoop, PSE.getSE(), D, AR);
+  return isInductionPHI(Phi, TheLoop, PSE.getSE(), D, Preds, AR);
 }
 
 bool InductionDescriptor::isInductionPHI(
     PHINode *Phi, const Loop *TheLoop, ScalarEvolution *SE,
-    InductionDescriptor &D, const SCEV *Expr,
-    SmallVectorImpl<Instruction *> *CastsToIgnore) {
+    InductionDescriptor &D, ArrayRef<const SCEVPredicate *> Preds,
+    const SCEV *Expr, SmallVectorImpl<Instruction *> *CastsToIgnore) {
   Type *PhiTy = Phi->getType();
   // isSCEVable returns true for integer and pointer types.
   if (!SE->isSCEVable(PhiTy))
@@ -1671,13 +1685,14 @@ bool InductionDescriptor::isInductionPHI(
     BinaryOperator *BOp =
         dyn_cast<BinaryOperator>(Phi->getIncomingValueForBlock(Latch));
     D = InductionDescriptor(StartValue, IK_IntInduction, Step, BOp,
-                            CastsToIgnore);
+                            CastsToIgnore, Preds);
     return true;
   }
 
   assert(PhiTy->isPointerTy() && "The PHI must be a pointer");
 
   // This allows induction variables w/non-constant steps.
-  D = InductionDescriptor(StartValue, IK_PtrInduction, Step);
+  D = InductionDescriptor(StartValue, IK_PtrInduction, Step,
+                          /*InductionBinOp=*/nullptr, /*Casts=*/nullptr, Preds);
   return true;
 }

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -5884,14 +5884,17 @@ ScalarEvolution::createAddRecFromPHIWithCasts(const SCEVUnknown *SymbolicPHI) {
 // even when the following Equal predicate exists:
 // "%step == (sext ix (trunc iy to ix) to iy)".
 bool PredicatedScalarEvolution::areAddRecsEqualWithPreds(
-    const SCEVAddRecExpr *AR1, const SCEVAddRecExpr *AR2) const {
+    const SCEVAddRecExpr *AR1, const SCEVAddRecExpr *AR2,
+    ArrayRef<const SCEVPredicate *> NoWrapPreds) const {
   if (AR1 == AR2)
     return true;
 
+  SCEVUnionPredicate NoWrapUnionPred(NoWrapPreds, SE);
+  SCEVUnionPredicate AllPreds = Preds->getUnionWith(&NoWrapUnionPred, SE);
   auto areExprsEqual = [&](const SCEV *Expr1, const SCEV *Expr2) -> bool {
     if (Expr1 != Expr2 &&
-        !Preds->implies(SE.getEqualPredicate(Expr1, Expr2), SE) &&
-        !Preds->implies(SE.getEqualPredicate(Expr2, Expr1), SE))
+        !AllPreds.implies(SE.getEqualPredicate(Expr1, Expr2), SE) &&
+        !AllPreds.implies(SE.getEqualPredicate(Expr2, Expr1), SE))
       return false;
     return true;
   };
@@ -15625,14 +15628,20 @@ bool PredicatedScalarEvolution::hasNoOverflow(
   return Flags == SCEVWrapPredicate::IncrementAnyWrap;
 }
 
-const SCEVAddRecExpr *PredicatedScalarEvolution::getAsAddRec(Value *V) {
+const SCEVAddRecExpr *PredicatedScalarEvolution::getAsAddRec(
+    Value *V, SmallVectorImpl<const SCEVPredicate *> *ExtraPreds) {
   const SCEV *Expr = this->getSCEV(V);
   SmallVector<const SCEVPredicate *, 4> NewPreds;
   auto *New = SE.convertSCEVToAddRecWithPredicates(Expr, &L, NewPreds);
 
   if (!New)
     return nullptr;
 
+  if (ExtraPreds) {
+    ExtraPreds->append(NewPreds);
+    return New;
+  }
+
   for (const auto *P : NewPreds)
     addPredicate(*P);