[LV] Identify more induction PHIs by coercing expressions to AddRecExprs

[sbaranga]

Summary

Some PHIs can have expressions that are not AddRecExprs due to the presence
of sext/zext instructions. In order to prevent the Loop Vectorizer from
bailing out when encountering these PHIs, we now coerce the SCEV
expressions to AddRecExprs using SCEV predicates (when possible).

We only do this when the alternative would be to not vectorize.

[sbaranga]

Perform a rebase of the patch.

Only add SCEV predicates after we know that we couldn't do anything else.

However, we still need to identify induction variables before trying to find
first-order recurrences. Therefore, we will now classify in Phis in the following
order:

1. induction
2. first-order recurrence
3. induction, with the help of SCEV predicates.
4. give up

[sbaranga]

Ping?

[anemet]

lib/Transforms/Vectorize/LoopVectorize.cpp

4671–4672

Why did you choose a lambda rather than a helper for this? This is a pretty big function and the lambda is not small either. Unless there is a strong reason I would prefer factoring this part out into a helper.

test/Transforms/LoopVectorize/induction.ll

171

Please comment which induction PHI you mean here? Same on the other testcase.

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

4671–4672

There's no strong reason for this (it just avoided passing a lot of arguments). I'll convert it to a helper function.

test/Transforms/LoopVectorize/induction.ll

171

Sure! (it's %sphi in both tests)

[sbaranga]

Convert the lambda into a helper function.

[sbaranga]

Ping?

[sbaranga]

Ping?

-Silviu

[anemet]

Sorry about the delay, this LGTM.

include/llvm/Transforms/Utils/LoopUtils.h

296–297

Please add function comment. Would be nice to also add comment for the original function. Thanks!

lib/Transforms/Vectorize/LoopVectorize.cpp

1417

"main main"

4582–4586

Please peel off the NFC part of splitting out the helper into a separate post-commit review change.

[sbaranga]

Update the code after the helper creation part was committed in r268632.

[sbaranga]

Thanks, committed in r268633.

-Silviu

include/llvm/Transforms/Utils/LoopUtils.h

 class AliasSetTracker;
 class AssumptionCache;
 class BasicBlock;
 class DataLayout;
 class DominatorTree;
 class Loop;
 class LoopInfo;
 class Pass;
+class PredicatedScalarEvolution;
 class PredIteratorCache;
 class ScalarEvolution;
 class TargetLibraryInfo;
 
 /// \brief Captures loop safety information.
 /// It keep information for loop & its header may throw exception.
 struct LICMSafetyInfo {
   bool MayThrow;           // The current loop contains an instruction which
public:
   /// flags, which can be found from the original scalar operations.
   Value *transform(IRBuilder<> &B, Value *Index) const;
 
   Value *getStartValue() const { return StartValue; }
   InductionKind getKind() const { return IK; }
   ConstantInt *getStepValue() const { return StepValue; }
 
   static bool isInductionPHI(PHINode *Phi, ScalarEvolution *SE,
-                             InductionDescriptor &D);
+                             InductionDescriptor &D,
+                             const SCEV *Expr = nullptr);
+
+  static bool isInductionPHI(PHINode *Phi, PredicatedScalarEvolution &PSE,
+                             InductionDescriptor &D, bool Assume = false);
 
 private:
   /// Private constructor - used by \c isInductionPHI.
   InductionDescriptor(Value *Start, InductionKind K, ConstantInt *Step);
 
   /// Start value.
   TrackingVH<Value> StartValue;
   /// Induction kind.

lib/Transforms/Utils/LoopUtils.cpp

case IK_PtrInduction:
     return B.CreateGEP(nullptr, StartValue, Index);
 
   case IK_NoInduction:
     return nullptr;
   }
   llvm_unreachable("invalid enum");
 }
 
-bool InductionDescriptor::isInductionPHI(PHINode *Phi, ScalarEvolution *SE,
-                                         InductionDescriptor &D) {
+bool InductionDescriptor::isInductionPHI(PHINode *Phi,
+                                         PredicatedScalarEvolution &PSE,
+                                         InductionDescriptor &D,
+                                         bool Assume) {
+  Type *PhiTy = Phi->getType();
+  // We only handle integer and pointer inductions variables.
+  if (!PhiTy->isIntegerTy() && !PhiTy->isPointerTy())
+    return false;
+
+  const SCEV *PhiScev = PSE.getSCEV(Phi);
+  const auto *AR = dyn_cast<SCEVAddRecExpr>(PhiScev);
+
+  // We need this expression to be an AddRecExpr.
+  if (Assume && !AR)
+    AR = PSE.getAsAddRec(Phi);
+
+  if (!AR) {
+    DEBUG(dbgs() << "LV: PHI is not a poly recurrence.\n");
+    return false;
+  }
+
+  return isInductionPHI(Phi, PSE.getSE(), D, AR);
+}
+
+bool InductionDescriptor::isInductionPHI(PHINode *Phi,
+                                         ScalarEvolution *SE,
+                                         InductionDescriptor &D,
+                                         const SCEV *Expr) {
   Type *PhiTy = Phi->getType();
   // We only handle integer and pointer inductions variables.
   if (!PhiTy->isIntegerTy() && !PhiTy->isPointerTy())
     return false;
 
   // Check that the PHI is consecutive.
-  const SCEV *PhiScev = SE->getSCEV(Phi);
+  const SCEV *PhiScev = Expr ? Expr : SE->getSCEV(Phi);
   const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(PhiScev);
+
   if (!AR) {
     DEBUG(dbgs() << "LV: PHI is not a poly recurrence.\n");
     return false;
   }
 
   assert(AR->getLoop()->getHeader() == Phi->getParent() &&
          "PHI is an AddRec for a different loop?!");
   Value *StartValue =

lib/Transforms/Vectorize/LoopVectorize.cpp

private:
   bool blockCanBePredicated(BasicBlock *BB, SmallPtrSetImpl<Value *> &SafePtrs);
 
   /// \brief Collect memory access with loop invariant strides.
   ///
   /// Looks for accesses like "a[i * StrideA]" where "StrideA" is loop
   /// invariant.
   void collectStridedAccess(Value *LoadOrStoreInst);
 
+  /// \brief Returns true if we can vectorize using this PHI node as an
+  /// induction.
+  ///
+  /// Updates the vectorization state by adding \p Phi to the inductions
+  /// list. This can set \p Phi as the main main induction of the loop if
+  /// \p Phi is a better choice for the main induction than the existing one.
+  bool addInductionPhi(PHINode *Phi, InductionDescriptor ID);
+
   /// Report an analysis message to assist the user in diagnosing loops that are
   /// not vectorized.  These are handled as LoopAccessReport rather than
   /// VectorizationReport because the << operator of VectorizationReport returns
   /// LoopAccessReport.
   void emitAnalysis(const LoopAccessReport &Message) const {
     emitAnalysisDiag(TheFunction, TheLoop, *Hints, Message);
   }
 
for (User *U : Inst->users()) {
       if (!TheLoop->contains(UI)) {
         DEBUG(dbgs() << "LV: Found an outside user for : " << *UI << '\n');
         return true;
       }
     }
   return false;
 }
 
+bool LoopVectorizationLegality::addInductionPhi(PHINode *Phi,
+                                                InductionDescriptor ID) {
+  Inductions[Phi] = ID;
+  Type *PhiTy = Phi->getType();
+  const DataLayout &DL = Phi->getModule()->getDataLayout();
+
+  // Get the widest type.
+  if (!WidestIndTy)
+    WidestIndTy = convertPointerToIntegerType(DL, PhiTy);
+  else
+    WidestIndTy = getWiderType(DL, PhiTy, WidestIndTy);
+
+  // Int inductions are special because we only allow one IV.
+  if (ID.getKind() == InductionDescriptor::IK_IntInduction &&
+      ID.getStepValue()->isOne() &&
+      isa<Constant>(ID.getStartValue()) &&
+        cast<Constant>(ID.getStartValue())->isNullValue()) {
+    // Use the phi node with the widest type as induction. Use the last
+    // one if there are multiple (no good reason for doing this other
+    // than it is expedient). We've checked that it begins at zero and
+    // steps by one, so this is a canonical induction variable.
+    if (!Induction || PhiTy == WidestIndTy)
+      Induction = Phi;
+  }
+
+  DEBUG(dbgs() << "LV: Found an induction variable.\n");
+
+  // Until we explicitly handle the case of an induction variable with
+  // an outside loop user we have to give up vectorizing this loop.
+  if (hasOutsideLoopUser(TheLoop, Phi, AllowedExit)) {
+    emitAnalysis(VectorizationReport(Phi) <<
+                 "use of induction value outside of the "
+                 "loop is not handled by vectorizer");
+    return false;
+  }
+
+  return true;
+}
+
 bool LoopVectorizationLegality::canVectorizeInstrs() {
   BasicBlock *Header = TheLoop->getHeader();
 
   // Look for the attribute signaling the absence of NaNs.
   Function &F = *Header->getParent();
   const DataLayout &DL = F.getParent()->getDataLayout();
   HasFunNoNaNAttr =
       F.getFnAttribute("no-nans-fp-math").getValueAsString() == "true";
for (BasicBlock::iterator it = (*bb)->begin(), e = (*bb)->end(); it != e;
 
         // We only allow if-converted PHIs with exactly two incoming values.
         if (Phi->getNumIncomingValues() != 2) {
           emitAnalysis(VectorizationReport(&*it)
                        << "control flow not understood by vectorizer");
           DEBUG(dbgs() << "LV: Found an invalid PHI.\n");
           return false;
         }
 
-        InductionDescriptor ID;
-        if (InductionDescriptor::isInductionPHI(Phi, PSE.getSE(), ID)) {
-          Inductions[Phi] = ID;
-          // Get the widest type.
-          if (!WidestIndTy)
-            WidestIndTy = convertPointerToIntegerType(DL, PhiTy);
-          else
-            WidestIndTy = getWiderType(DL, PhiTy, WidestIndTy);
-
-          // Int inductions are special because we only allow one IV.
-          if (ID.getKind() == InductionDescriptor::IK_IntInduction &&
-              ID.getStepValue()->isOne() &&
-              isa<Constant>(ID.getStartValue()) &&
-                cast<Constant>(ID.getStartValue())->isNullValue()) {
-            // Use the phi node with the widest type as induction. Use the last
-            // one if there are multiple (no good reason for doing this other
-            // than it is expedient). We've checked that it begins at zero and
-            // steps by one, so this is a canonical induction variable.
-            if (!Induction || PhiTy == WidestIndTy)
-              Induction = Phi;
-          }
-
-          DEBUG(dbgs() << "LV: Found an induction variable.\n");
-
-          // Until we explicitly handle the case of an induction variable with
-          // an outside loop user we have to give up vectorizing this loop.
-          if (hasOutsideLoopUser(TheLoop, &*it, AllowedExit)) {
-            emitAnalysis(VectorizationReport(&*it) <<
-                         "use of induction value outside of the "
-                         "loop is not handled by vectorizer");
-            return false;
-          }
-
-          continue;
-        }
-
         RecurrenceDescriptor RedDes;
         if (RecurrenceDescriptor::isReductionPHI(Phi, TheLoop, RedDes)) {
           if (RedDes.hasUnsafeAlgebra())
             Requirements->addUnsafeAlgebraInst(RedDes.getUnsafeAlgebraInst());
           AllowedExit.insert(RedDes.getLoopExitInstr());
           Reductions[Phi] = RedDes;
           continue;
         }
 
+        InductionDescriptor ID;
+        if (InductionDescriptor::isInductionPHI(Phi, PSE, ID)) {
+          if (!addInductionPhi(Phi, ID))
+            return false;
+          continue;
+        }
+
         if (RecurrenceDescriptor::isFirstOrderRecurrence(Phi, TheLoop, DT)) {
           FirstOrderRecurrences.insert(Phi);
           continue;
         }
 
+        // As a last resort, coerce the PHI to a AddRec expression
+        // and re-try classifying it a an induction PHI.
+        if (InductionDescriptor::isInductionPHI(Phi, PSE, ID, true)) {
+          if (!addInductionPhi(Phi, ID))
+            return false;
+          continue;
+        }
+
         emitAnalysis(VectorizationReport(&*it) <<
                      "value that could not be identified as "
                      "reduction is used outside the loop");
         DEBUG(dbgs() << "LV: Found an unidentified PHI."<< *Phi <<"\n");
         return false;
       }// end of PHI handling
 
       // We handle calls that:

test/Transforms/LoopVectorize/induction.ll

cond.end.i:
   %and.i = and i32 %0, %and3.i
   %inc.i = add i8 %inc4.i, 1
   %tobool.i = icmp eq i8 %inc.i, 0
   br i1 %tobool.i, label %loopexit, label %cond.end.i
 
 loopexit:
   ret i32 %and.i
 }
+
+; The SCEV expression of %sphi is (zext i8 {%t,+,1}<%loop> to i32)
+; In order to recognize %sphi as an induction PHI and vectorize this loop,
+; we need to convert the SCEV expression into an AddRecExpr.
+; The expression gets converted to {zext i8 %t to i32,+,1}.
+
+; CHECK-LABEL: wrappingindvars1
+; CHECK-LABEL: vector.scevcheck
+; CHECK-LABEL: vector.body
+; CHECK: add <2 x i32> {{%[^ ]*}}, <i32 0, i32 1>
+define void @wrappingindvars1(i8 %t, i32 %len, i32 *%A) {
+ entry:
+  %st = zext i8 %t to i16
+  %ext = zext i8 %t to i32
+  %ecmp = icmp ult i16 %st, 42
+  br i1 %ecmp, label %loop, label %exit
+
+ loop:
+
+  %idx = phi i8 [ %t, %entry ], [ %idx.inc, %loop ]
+  %idx.b = phi i32 [ 0, %entry ], [ %idx.b.inc, %loop ]
+  %sphi = phi i32 [ %ext, %entry ], [%idx.inc.ext, %loop]
+
+  %ptr = getelementptr inbounds i32, i32* %A, i8 %idx
+  store i32 %sphi, i32* %ptr
+
+  %idx.inc = add i8 %idx, 1
+  %idx.inc.ext = zext i8 %idx.inc to i32
+  %idx.b.inc = add nuw nsw i32 %idx.b, 1
+
+  %c = icmp ult i32 %idx.b, %len
+  br i1 %c, label %loop, label %exit
+
+ exit:
+  ret void
+}
+
+; The SCEV expression of %sphi is (4 * (zext i8 {%t,+,1}<%loop> to i32))
+; In order to recognize %sphi as an induction PHI and vectorize this loop,
+; we need to convert the SCEV expression into an AddRecExpr.
+; The expression gets converted to ({4 * (zext %t to i32),+,4}).
+; CHECK-LABEL: wrappingindvars2
+; CHECK-LABEL: vector.scevcheck
+; CHECK-LABEL: vector.body
+; CHECK: add <2 x i32> {{%[^ ]*}}, <i32 0, i32 4>
+define void @wrappingindvars2(i8 %t, i32 %len, i32 *%A) {
+
+entry:
+  %st = zext i8 %t to i16
+  %ext = zext i8 %t to i32
+  %ext.mul = mul i32 %ext, 4
+
+  %ecmp = icmp ult i16 %st, 42
+  br i1 %ecmp, label %loop, label %exit
+
+ loop:
+
+  %idx = phi i8 [ %t, %entry ], [ %idx.inc, %loop ]
+  %sphi = phi i32 [ %ext.mul, %entry ], [%mul, %loop]
+  %idx.b = phi i32 [ 0, %entry ], [ %idx.b.inc, %loop ]
+
+  %ptr = getelementptr inbounds i32, i32* %A, i8 %idx
+  store i32 %sphi, i32* %ptr
+
+  %idx.inc = add i8 %idx, 1
+  %idx.inc.ext = zext i8 %idx.inc to i32
+  %mul = mul i32 %idx.inc.ext, 4
+  %idx.b.inc = add nuw nsw i32 %idx.b, 1
+
+  %c = icmp ult i32 %idx.b, %len
+  br i1 %c, label %loop, label %exit
+
+ exit:
+  ret void
+}