[llvm] [VPlan] Make canonical IV part of the region (PR #156262)

[via llvm-commits]

================
@@ -9056,15 +9047,27 @@ static SmallVector<Instruction *> preparePlanForEpilogueVectorLoop(
          "ScalarIVSteps when resetting the start value");
   VPBuilder Builder(Header, Header->getFirstNonPhi());
   VPInstruction *Add = Builder.createAdd(IV, VPV);
-  IV->replaceAllUsesWith(Add);
-  Add->setOperand(0, IV);
+  // Replace all users of the canonical IV with the offset version, except for
+  // the Add itself and the canonical IV increment.
+  auto *Increment = vputils::findCanonicalIVIncrement(Plan);
+  IV->replaceUsesWithIf(Add, [Add, Increment](VPUser &U, unsigned) {
+    return &U != Add && &U != Increment;
+  });
+  // The exit condition compares the IV increment against the absolute trip
+  // count, so create an offset version of the increment for BranchOnCount.
+  auto *VectorLatch = VectorLoop->getExitingBasicBlock();
+  auto *BranchOnCount = cast<VPInstruction>(&VectorLatch->back());
+  assert(BranchOnCount->getOpcode() == VPInstruction::BranchOnCount &&
+         "expected BranchOnCount in exiting block");
+  Builder.setInsertPoint(BranchOnCount);
+  VPInstruction *OffsetIVInc = Builder.createAdd(Increment, VPV);
+  BranchOnCount->setOperand(0, OffsetIVInc);
----------------
ayalz wrote:

Better replace all uses of Increment with OffsetIVInc?
(Instead of bumping the start of IV, we keep it 0 and replace users of IV by IV+VPV, except for Increment.
Instead of bumping Increment to use IV+VPV, we keep it using IV and replace its users by Increment+VPV.)

https://github.com/llvm/llvm-project/pull/156262