[VPlan] Use VPlan predecessors in VPWidenPHIRecipe (NFC). (#126388)

Update VPWidenPHIRecipe to use the predecessors in VPlan to determine
the incoming blocks instead of tracking them separately. This brings
VPWidenPHIRecipe in line with the other phi recipes.

PR: https://github.com/llvm/llvm-project/pull/126388

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1958,13 +1958,12 @@ public:
 #endif
 };
 
-/// A recipe for handling phis that are widened in the vector loop.
-/// In the VPlan native path, all incoming VPValues & VPBasicBlock pairs are
-/// managed in the recipe directly.
+/// A recipe for widened phis. Incoming values are operands of the recipe and
+/// their operand index corresponds to the incoming predecessor block. If the
+/// recipe is placed in an entry block to a (non-replicate) region, it must have
+/// exactly 2 incoming values, the first from the predecessor of the region and
+/// the second from the exiting block of the region.
 class VPWidenPHIRecipe : public VPSingleDefRecipe {
-  /// List of incoming blocks. Only used in the VPlan native path.
-  SmallVector<VPBasicBlock *, 2> IncomingBlocks;
-
 public:
   /// Create a new VPWidenPHIRecipe for \p Phi with start value \p Start and
   /// debug location \p DL.
@@ -1991,19 +1990,8 @@ public:
              VPSlotTracker &SlotTracker) const override;
 #endif
 
-  /// Adds a pair (\p IncomingV, \p IncomingBlock) to the phi.
-  void addIncoming(VPValue *IncomingV, VPBasicBlock *IncomingBlock) {
-    addOperand(IncomingV);
-    IncomingBlocks.push_back(IncomingBlock);
-  }
-
   /// Returns the \p I th incoming VPBasicBlock.
-  VPBasicBlock *getIncomingBlock(unsigned I) { return IncomingBlocks[I]; }
-
-  /// Set the \p I th incoming VPBasicBlock to \p IncomingBlock.
-  void setIncomingBlock(unsigned I, VPBasicBlock *IncomingBlock) {
-    IncomingBlocks[I] = IncomingBlock;
-  }
+  VPBasicBlock *getIncomingBlock(unsigned I);
 
   /// Returns the \p I th incoming VPValue.
   VPValue *getIncomingValue(unsigned I) { return getOperand(I); }

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

@@ -136,19 +136,23 @@ void PlainCFGBuilder::fixPhiNodes() {
       // predecessor is the first operand of the recipe.
       assert(Phi->getNumOperands() == 2);
       BasicBlock *LoopPred = L->getLoopPredecessor();
-      VPPhi->addIncoming(
-          getOrCreateVPOperand(Phi->getIncomingValueForBlock(LoopPred)),
-          BB2VPBB[LoopPred]);
+      VPPhi->addOperand(
+          getOrCreateVPOperand(Phi->getIncomingValueForBlock(LoopPred)));
       BasicBlock *LoopLatch = L->getLoopLatch();
-      VPPhi->addIncoming(
-          getOrCreateVPOperand(Phi->getIncomingValueForBlock(LoopLatch)),
-          BB2VPBB[LoopLatch]);
+      VPPhi->addOperand(
+          getOrCreateVPOperand(Phi->getIncomingValueForBlock(LoopLatch)));
       continue;
     }
 
-    for (unsigned I = 0; I != Phi->getNumOperands(); ++I)
-      VPPhi->addIncoming(getOrCreateVPOperand(Phi->getIncomingValue(I)),
-                         BB2VPBB[Phi->getIncomingBlock(I)]);
+    // Add operands for VPPhi in the order matching its predecessors in VPlan.
+    DenseMap<const VPBasicBlock *, VPValue *> VPPredToIncomingValue;
+    for (unsigned I = 0; I != Phi->getNumOperands(); ++I) {
+      VPPredToIncomingValue[BB2VPBB[Phi->getIncomingBlock(I)]] =
+          getOrCreateVPOperand(Phi->getIncomingValue(I));
+    }
+    for (VPBlockBase *Pred : VPPhi->getParent()->getPredecessors())
+      VPPhi->addOperand(
+          VPPredToIncomingValue.lookup(Pred->getExitingBasicBlock()));
   }
 }
 

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -3621,6 +3621,27 @@ void VPReductionPHIRecipe::print(raw_ostream &O, const Twine &Indent,
 }
 #endif
 
+VPBasicBlock *VPWidenPHIRecipe::getIncomingBlock(unsigned I) {
+  VPBasicBlock *Parent = getParent();
+  VPBlockBase *Pred = nullptr;
+  if (Parent->getNumPredecessors() > 0) {
+    Pred = Parent->getPredecessors()[I];
+  } else {
+    auto *Region = Parent->getParent();
+    assert(Region && !Region->isReplicator() && Region->getEntry() == Parent &&
+           "must be in the entry block of a non-replicate region");
+    assert(
+        I < 2 && getNumOperands() == 2 &&
+        "when placed in an entry block, only 2 incoming blocks are available");
+
+    // I ==  0 selects the predecessor of the region, I == 1 selects the region
+    // itself whose exiting block feeds the phi across the backedge.
+    Pred = I == 0 ? Region->getSinglePredecessor() : Region;
+  }
+
+  return Pred->getExitingBasicBlock();
+}
+
 void VPWidenPHIRecipe::execute(VPTransformState &State) {
   assert(EnableVPlanNativePath &&
          "Non-native vplans are not expected to have VPWidenPHIRecipes.");

llvm/lib/Transforms/Vectorize/VPlanUtils.h

@@ -169,16 +169,8 @@ public:
   static void reassociateBlocks(VPBlockBase *Old, VPBlockBase *New) {
     for (auto *Pred : to_vector(Old->getPredecessors()))
       Pred->replaceSuccessor(Old, New);
-    for (auto *Succ : to_vector(Old->getSuccessors())) {
+    for (auto *Succ : to_vector(Old->getSuccessors()))
       Succ->replacePredecessor(Old, New);
-
-      // Replace any references to Old in widened phi incoming blocks.
-      for (auto &R : Succ->getEntryBasicBlock()->phis())
-        if (auto *WidenPhiR = dyn_cast<VPWidenPHIRecipe>(&R))
-          for (unsigned I = 0; I < WidenPhiR->getNumOperands(); I++)
-            if (WidenPhiR->getIncomingBlock(I) == Old)
-              WidenPhiR->setIncomingBlock(I, cast<VPBasicBlock>(New));
-    }
     New->setPredecessors(Old->getPredecessors());
     New->setSuccessors(Old->getSuccessors());
     Old->clearPredecessors();

llvm/test/Transforms/LoopVectorize/outer-loop-wide-phis.ll

@@ -134,7 +134,7 @@ define void @wide_phi_2_predecessors_phi_ops_swapped(ptr noalias %A, ptr noalias
 ; CHECK-NEXT:    [[WIDE_MASKED_GATHER:%.*]] = call <4 x i64> @llvm.masked.gather.v4i64.v4p0(<4 x ptr> [[TMP1]], i32 8, <4 x i1> splat (i1 true), <4 x i64> poison)
 ; CHECK-NEXT:    br label %[[INNER_LATCH4]]
 ; CHECK:       [[INNER_LATCH4]]:
-; CHECK-NEXT:    [[VEC_PHI5:%.*]] = phi <4 x i64> [ zeroinitializer, %[[INNER_HEADER1]] ], [ [[WIDE_MASKED_GATHER]], %[[THEN3]] ]
+; CHECK-NEXT:    [[VEC_PHI5:%.*]] = phi <4 x i64> [ [[WIDE_MASKED_GATHER]], %[[THEN3]] ], [ zeroinitializer, %[[INNER_HEADER1]] ]
 ; CHECK-NEXT:    [[TMP2:%.*]] = add nsw <4 x i64> [[VEC_PHI5]], [[VEC_IND]]
 ; CHECK-NEXT:    [[TMP3]] = add nsw <4 x i64> [[TMP2]], [[VEC_PHI2]]
 ; CHECK-NEXT:    [[TMP4]] = add nuw nsw <4 x i64> [[VEC_PHI]], splat (i64 1)

llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

@@ -672,7 +672,7 @@ TEST_F(VPBasicBlockTest, reassociateBlocks) {
     auto *WidenPhi = new VPWidenPHIRecipe(nullptr);
     IntegerType *Int32 = IntegerType::get(C, 32);
     VPValue *Val = Plan.getOrAddLiveIn(ConstantInt::get(Int32, 1));
-    WidenPhi->addIncoming(Val, VPBB1);
+    WidenPhi->addOperand(Val);
     VPBB2->appendRecipe(WidenPhi);
 
     VPBasicBlock *VPBBNew = Plan.createVPBasicBlock("VPBBNew");
@@ -693,7 +693,8 @@ TEST_F(VPBasicBlockTest, reassociateBlocks) {
     auto *WidenPhi = new VPWidenPHIRecipe(nullptr);
     IntegerType *Int32 = IntegerType::get(C, 32);
     VPValue *Val = Plan.getOrAddLiveIn(ConstantInt::get(Int32, 1));
-    WidenPhi->addIncoming(Val, VPBB1);
+    WidenPhi->addOperand(Val);
+    WidenPhi->addOperand(Val);
     VPBB2->appendRecipe(WidenPhi);
 
     VPBasicBlock *VPBBNew = Plan.createVPBasicBlock("VPBBNew");