[LV] Use VPValue for SCEV expansion in fixupIVUsers.

[fhahn]

Summary

The step is already expanded in the VPlan. Use this expansion instead.
This is a step towards modeling fixing up IV users in VPlan.

It also fixes a crash casued by SCEV-expanding the Step expression in
fixupIVUsers, where the IR is in an incomplete state

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3370

Would it make sense to look for the IV recipe that corresponds to OrigPhi (with overridden getVPValue() or by scanning header phi recipes) and then reach its Step, instead of looking for StepR in preheader and relying on II.getStep()?
(Traversing IV recipes (or live-outs) instead of original phi's could probably follow.)

[fhahn]

Add comment explaining step expansion.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3370

Unfortunately we cannot rely on mapping IV recipes to OrigPhi here, as an original VPWidenIntOrFpInductionRecipe ( which maps to OrigPhi) could be transformed to a VPScalarIVStepsRecipe, which doesn't store OrigPhi.

This will only be cleaned up once we completely model fixing up inductions in VPlan.

[fhahn]

Added test case that shows this patch is not an NFC. It fixes a crash casued by SCEV-expanding the Step expression in fixupIVUsers, where the IR is in an incomplete state. I will adjust the title accordingly.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3370

ok, deserves a TODO note emphasizing it's a temporary solution.

Worth recording a SCEVtoExpandSCEVRecipe mapping in VPlan (with a note that it too is temporary)?

[fhahn]

Update to add temporary mapping from SCEV -> VPValue expansions.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3370

Updated the code to have a temporary mapping. I've not added a TODO, as moving IV fix up to VPlan is a general TODO :)

[fhahn]

rebase and ping :)

[Ayal]

LGTM, with some comments, also provided in D147964.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3384

nit: worth asserting that step SCEV is expected to have a VPValue?

3385–3386

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3384

Added an assert in the committed version, thanks!

3385–3386

Updated in the committed version, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

protected:
   /// in the new unrolled loop, where UF is the unroll factor and VF is the
   /// vectorization factor.
   using ScalarParts = SmallVector<SmallVector<Value *, 4>, 2>;
 
   /// Set up the values of the IVs correctly when exiting the vector loop.
   void fixupIVUsers(PHINode *OrigPhi, const InductionDescriptor &II,
                     Value *VectorTripCount, Value *EndValue,
                     BasicBlock *MiddleBlock, BasicBlock *VectorHeader,
-                    VPlan &Plan);
+                    VPlan &Plan, VPTransformState &State);
 
   /// Handle all cross-iteration phis in the header.
   void fixCrossIterationPHIs(VPTransformState &State);
 
   /// Create the exit value of first order recurrences in the middle block and
   /// update their users.
   void fixFixedOrderRecurrence(VPFirstOrderRecurrencePHIRecipe *PhiR,
                                VPTransformState &State);

 // Fix up external users of the induction variable. At this point, we are
 // in LCSSA form, with all external PHIs that use the IV having one input value,
 // coming from the remainder loop. We need those PHIs to also have a correct
 // value for the IV when arriving directly from the middle block.
 void InnerLoopVectorizer::fixupIVUsers(PHINode *OrigPhi,
                                        const InductionDescriptor &II,
                                        Value *VectorTripCount, Value *EndValue,
                                        BasicBlock *MiddleBlock,
-                                       BasicBlock *VectorHeader, VPlan &Plan) {
+                                       BasicBlock *VectorHeader, VPlan &Plan,
+                                       VPTransformState &State) {
   // There are two kinds of external IV usages - those that use the value
   // computed in the last iteration (the PHI) and those that use the penultimate
   // value (the value that feeds into the phi from the loop latch).
   // We allow both, but they, obviously, have different values.
 
   assert(OrigLoop->getUniqueExitBlock() && "Expected a single exit block");
 
   DenseMap<Value *, Value *> MissingVals;
 
   // An external user of the last iteration's value should see the value that
   // the remainder loop uses to initialize its own IV.
   Value *PostInc = OrigPhi->getIncomingValueForBlock(OrigLoop->getLoopLatch());
   for (User *U : PostInc->users()) {
     Instruction *UI = cast<Instruction>(U);
     if (!OrigLoop->contains(UI)) {
       assert(isa<PHINode>(UI) && "Expected LCSSA form");
       MissingVals[UI] = EndValue;
     }
   }
 
   // An external user of the penultimate value need to see EndValue - Step.
   // The simplest way to get this is to recompute it from the constituent SCEVs,
   // that is Start + (Step * (CRD - 1)).
   for (User *U : OrigPhi->users()) {
     auto *UI = cast<Instruction>(U);
     if (!OrigLoop->contains(UI)) {
       assert(isa<PHINode>(UI) && "Expected LCSSA form");
-
       IRBuilder<> B(MiddleBlock->getTerminator());
 
       // Fast-math-flags propagate from the original induction instruction.
       if (II.getInductionBinOp() && isa<FPMathOperator>(II.getInductionBinOp()))
         B.setFastMathFlags(II.getInductionBinOp()->getFastMathFlags());
 
       Value *CountMinusOne = B.CreateSub(
           VectorTripCount, ConstantInt::get(VectorTripCount->getType(), 1));
       CountMinusOne->setName("cmo");
-      Value *Step = CreateStepValue(II.getStep(), *PSE.getSE(),
-                                    VectorHeader->getTerminator());
+
+      VPValue *StepVPV = Plan.getSCEVExpansion(II.getStep());
+      assert(StepVPV && "step must have been expanded during VPlan execution");
+      Value *Step = StepVPV->isLiveIn() ? StepVPV->getLiveInIRValue()
+                                        : State.get(StepVPV, 0);
       Value *Escape =
           emitTransformedIndex(B, CountMinusOne, II.getStartValue(), Step, II);
       Escape->setName("ind.escape");
       MissingVals[UI] = Escape;
     }
   }
 
   for (auto &I : MissingVals) {
if (Cost->requiresScalarEpilogue(VF)) {
     // update uses outside the loop (phis) to account for the newly inserted
     // edge.
 
     // Fix-up external users of the induction variables.
     for (const auto &Entry : Legal->getInductionVars())
       fixupIVUsers(Entry.first, Entry.second,
                    getOrCreateVectorTripCount(VectorLoop->getLoopPreheader()),
                    IVEndValues[Entry.first], LoopMiddleBlock,
-                   VectorLoop->getHeader(), Plan);
+                   VectorLoop->getHeader(), Plan, State);
   }
 
   // Fix LCSSA phis not already fixed earlier. Extracts may need to be generated
   // in the exit block, so update the builder.
   State.Builder.SetInsertPoint(State.CFG.ExitBB->getFirstNonPHI());
   for (const auto &KV : Plan.getLiveOuts())
     KV.second->fixPhi(Plan, State);
 

llvm/lib/Transforms/Vectorize/VPlan.h

class VPlan {
 
   /// Indicates whether it is safe use the Value2VPValue mapping or if the
   /// mapping cannot be used any longer, because it is stale.
   bool Value2VPValueEnabled = true;
 
   /// Values used outside the plan.
   MapVector<PHINode *, VPLiveOut *> LiveOuts;
 
+  /// Mapping from SCEVs to the VPValues representing their expansions.
+  /// NOTE: This mapping is temporary and will be removed once all users have
+  /// been modeled in VPlan directly.
+  DenseMap<const SCEV *, VPValue *> SCEVToExpansion;
+
 public:
   VPlan(VPBasicBlock *Entry = nullptr) : Entry(Entry) {
     if (Entry)
       Entry->setPlan(this);
   }
 
   ~VPlan();
 
void removeLiveOut(PHINode *PN) {
     delete LiveOuts[PN];
     LiveOuts.erase(PN);
   }
 
   const MapVector<PHINode *, VPLiveOut *> &getLiveOuts() const {
     return LiveOuts;
   }
 
+  VPValue *getSCEVExpansion(const SCEV *S) const {
+    auto I = SCEVToExpansion.find(S);
+    if (I == SCEVToExpansion.end())
+      return nullptr;
+    return I->second;
+  }
+
+  void addSCEVExpansion(const SCEV *S, VPValue *V) {
+    assert(SCEVToExpansion.find(S) == SCEVToExpansion.end() &&
+           "SCEV already expanded");
+    SCEVToExpansion[S] = V;
+  }
+
 private:
   /// Add to the given dominator tree the header block and every new basic block
   /// that was created between it and the latch block, inclusive.
   static void updateDominatorTree(DominatorTree *DT, BasicBlock *LoopLatchBB,
                                   BasicBlock *LoopPreHeaderBB,
                                   BasicBlock *LoopExitBB);
 };
 

llvm/lib/Transforms/Vectorize/VPlan.cpp

 
 bool vputils::onlyFirstLaneUsed(VPValue *Def) {
   return all_of(Def->users(),
                 [Def](VPUser *U) { return U->onlyFirstLaneUsed(Def); });
 }
 
 VPValue *vputils::getOrCreateVPValueForSCEVExpr(VPlan &Plan, const SCEV *Expr,
                                                 ScalarEvolution &SE) {
+  if (auto *Expanded = Plan.getSCEVExpansion(Expr))
+    return Expanded;
+  VPValue *Expanded = nullptr;
   if (auto *E = dyn_cast<SCEVConstant>(Expr))
-    return Plan.getVPValueOrAddLiveIn(E->getValue());
-  if (auto *E = dyn_cast<SCEVUnknown>(Expr))
-    return Plan.getVPValueOrAddLiveIn(E->getValue());
+    Expanded = Plan.getVPValueOrAddLiveIn(E->getValue());
+  else if (auto *E = dyn_cast<SCEVUnknown>(Expr))
+    Expanded = Plan.getVPValueOrAddLiveIn(E->getValue());
+  else {
 
     VPBasicBlock *Preheader = Plan.getEntry();
-  VPExpandSCEVRecipe *Step = new VPExpandSCEVRecipe(Expr, SE);
-  Preheader->appendRecipe(Step);
-  return Step;
+    Expanded = new VPExpandSCEVRecipe(Expr, SE);
+    Preheader->appendRecipe(Expanded->getDefiningRecipe());
+  }
+  Plan.addSCEVExpansion(Expr, Expanded);
+  return Expanded;
 }

llvm/test/Transforms/LoopVectorize/pr58811-scev-expansion.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes=loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -S %s | FileCheck %s
 
-; REQUIRES: asserts
-; XFAIL: *
-
 define void @test1_pr58811() {
 ; CHECK-LABEL: @test1_pr58811(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP_1_PREHEADER:%.*]]
 ; CHECK:       loop.1.preheader:
 ; CHECK-NEXT:    [[IV_1_PH:%.*]] = phi i32 [ [[SUB93_2:%.*]], [[UNREACHABLE_BB:%.*]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[TMP0:%.*]] = sub i32 0, [[IV_1_PH]]
 ; CHECK-NEXT:    br label [[LOOP_1:%.*]]