[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

3399

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

3399

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

3399

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

3399

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

3437

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

3438–3439

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3437

Added an assert in the committed version, thanks!

3438–3439

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)).
+  Value *Step = nullptr;
   for (User *U : OrigPhi->users()) {
     auto *UI = cast<Instruction>(U);
     if (!OrigLoop->contains(UI)) {
       assert(isa<PHINode>(UI) && "Expected LCSSA form");
 
+      if (!Step) {
+        // Get the value corresponding to the expansion of the induction step.
+        // First try to directly convert the SCEV, otherwise look for a
+        // corresponding VPExpandSCEVRecipe.
+        if (auto *S = dyn_cast<SCEVConstant>(II.getStep()))
+          Step = S->getValue();
+        else if (auto *S = dyn_cast<SCEVUnknown>(II.getStep()))
+          Step = S->getValue();
+        else {
+          // All steps for induction should have been expanded earlier, so this
+          // only looks up existing VPValues.
+          VPBasicBlock *Preheader = Plan.getEntry()->getEntryBasicBlock();
+          for (auto &R : *Preheader) {
+            auto *StepR = dyn_cast<VPExpandSCEVRecipe>(&R);
+            if (StepR && StepR->getSCEV() == II.getStep()) {
+              Step = State.get(StepR, 0);
+              break;
+            }
+          }
+          assert(Step && "need a step here");
+        }
+      }
+
       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());
+
       Value *Escape =
           emitTransformedIndex(B, CountMinusOne, II.getStartValue(), Step, II);
       Escape->setName("ind.escape");
       MissingVals[UI] = Escape;
     }
   }
 
   for (auto &I : MissingVals) {
     PHINode *PHI = cast<PHINode>(I.first);
     // One corner case we have to handle is two IVs "chasing" each-other,
     // that is %IV2 = phi [...], [ %IV1, %latch ]
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);