From d1d3563278589b2e19077e6a81cd9455d0195577 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo@fhahn.com>
Date: Tue, 29 Mar 2022 20:27:59 +0100
Subject: [PATCH] [LV] Move code to place pointer induction increment to VPlan
 post-processing.

This patch moves the code to set the correct incoming block for the
backedge value to VPlan::execute.

When generating the phi node, the backedge value is temporarily added
using the pre-header as incoming block. The invalid phi node will be
fixed up during VPlan::execute after main VPlan code generation.
At the same time, the backedge value is also moved to the latch.

This change removes the requirement to create the latch block up-front
for VPWidenInductionPHIRecipe::execute, which in turn will enable
modeling the pre-header in VPlan.

Depends on D121617.

Reviewed By: Ayal

Differential Revision: https://reviews.llvm.org/D121618
---
 .../Transforms/Vectorize/LoopVectorize.cpp    | 13 ++++++----
 llvm/lib/Transforms/Vectorize/VPlan.cpp       | 25 ++++++++++++++-----
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 8a4d8f2c39f6..efebf8efdec3 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -9599,11 +9599,9 @@ void VPWidenPointerInductionRecipe::execute(VPTransformState &State) {
   NewPointerPhi->addIncoming(ScalarStartValue, State.CFG.VectorPreHeader);
 
   // A pointer induction, performed by using a gep
-  BasicBlock *LoopLatch =
-      State.LI->getLoopFor(State.CFG.PrevBB)->getLoopLatch();
+  const DataLayout &DL = NewPointerPhi->getModule()->getDataLayout();
+  Instruction *InductionLoc = &*State.Builder.GetInsertPoint();
 
-  const DataLayout &DL = LoopLatch->getModule()->getDataLayout();
-  Instruction *InductionLoc = LoopLatch->getTerminator();
   const SCEV *ScalarStep = IndDesc.getStep();
   SCEVExpander Exp(SE, DL, "induction");
   Value *ScalarStepValue = Exp.expandCodeFor(ScalarStep, PhiType, InductionLoc);
@@ -9614,7 +9612,12 @@ void VPWidenPointerInductionRecipe::execute(VPTransformState &State) {
       IndDesc.getElementType(), NewPointerPhi,
       State.Builder.CreateMul(ScalarStepValue, NumUnrolledElems), "ptr.ind",
       InductionLoc);
-  NewPointerPhi->addIncoming(InductionGEP, LoopLatch);
+  // Add induction update using an incorrect block temporarily. The phi node
+  // will be fixed after VPlan execution. Note that at this point the latch
+  // block cannot be used, as it does not exist yet.
+  // TODO: Model increment value in VPlan, by turning the recipe into a
+  // multi-def and a subclass of VPHeaderPHIRecipe.
+  NewPointerPhi->addIncoming(InductionGEP, State.CFG.VectorPreHeader);
 
   // Create UF many actual address geps that use the pointer
   // phi as base and a vectorized version of the step value
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index 6e7c12a9cda2..c5bd434b209c 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -981,14 +981,27 @@ void VPlan::execute(VPTransformState *State) {
   }
   for (VPRecipeBase &R : Header->phis()) {
     // Skip phi-like recipes that generate their backedege values themselves.
-    // TODO: Model their backedge values explicitly.
-    if (isa<VPWidenPHIRecipe>(&R) || isa<VPWidenPointerInductionRecipe>(&R))
+    if (isa<VPWidenPHIRecipe>(&R))
       continue;
 
-    // Set the correct incoming block for backedge values and move induction to
-    // latch.
-    if (auto *IndR = dyn_cast<VPWidenIntOrFpInductionRecipe>(&R)) {
-      auto *Phi = cast<PHINode>(State->get(IndR, 0));
+    if (isa<VPWidenPointerInductionRecipe>(&R) ||
+        isa<VPWidenIntOrFpInductionRecipe>(&R)) {
+      PHINode *Phi = nullptr;
+      if (isa<VPWidenIntOrFpInductionRecipe>(&R)) {
+        Phi = cast<PHINode>(State->get(R.getVPSingleValue(), 0));
+      } else {
+        auto *WidenPhi = cast<VPWidenPointerInductionRecipe>(&R);
+        // TODO: Split off the case that all users of a pointer phi are scalar
+        // from the VPWidenPointerInductionRecipe.
+        if (all_of(WidenPhi->users(), [WidenPhi](const VPUser *U) {
+              return cast<VPRecipeBase>(U)->usesScalars(WidenPhi);
+            }))
+          continue;
+
+        auto *GEP = cast<GetElementPtrInst>(State->get(WidenPhi, 0));
+        Phi = cast<PHINode>(GEP->getPointerOperand());
+      }
+
       Phi->setIncomingBlock(1, VectorLatchBB);
 
       // Move the last step to the end of the latch block. This ensures