[llvm] [LV] Update incoming blocks in VPWidenPHIRecipe in reassociateBlocks (PR #125481)

[Luke Lau via llvm-commits]

https://github.com/lukel97 created https://github.com/llvm/llvm-project/pull/125481

This is extracted from #118638

After c7ebe4f we will crash in fixNonInductionPHIs if we use a VPWidenPHIRecipe with the vector preheader as an incoming block, because the phi will reference the old non-IRBB vector preheader.

This fixes this by updating VPBlockUtils::reassociateBlocks to update any VPWidenPHIRecipes's incoming blocks.

This assumes that if the VPWidenPHIRecipe is in a VPRegionBlock, it's in the entry block, and that we are replacing a VPBasicBlock with another VPBasicBlock.


>From 8cfe2f2ea9c827ffb6c7619615c7bf52b80cc785 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Mon, 3 Feb 2025 19:13:33 +0800
Subject: [PATCH] [LV] Update incoming blocks in VPWidenPHIRecipe in
 reassociateBlocks

This is extracted from #118638

After c7ebe4f we will crash in fixNonInductionPHIs if we use a VPWidenPHIRecipe with the vector preheader as an incoming block, because the phi will reference the old non-IRBB vector preheader.

This fixes this by updating VPBlockUtils::reassociateBlocks to update any VPWidenPHIRecipes's incoming blocks.

This assumes that if the VPWidenPHIRecipe is in a VPRegionBlock, it's in the entry block, and that we are replacing a VPBasicBlock with another VPBasicBlock.
---
 llvm/lib/Transforms/Vectorize/VPlan.h         |  5 +++
 llvm/lib/Transforms/Vectorize/VPlanUtils.h    | 13 ++++++-
 .../Transforms/Vectorize/VPlanTest.cpp        | 39 +++++++++++++++++++
 3 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 459222234bc37f..cb58abec61c11c 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -2347,6 +2347,11 @@ class VPWidenPHIRecipe : public VPSingleDefRecipe {
   /// 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;
+  }
+
   /// Returns the \p I th incoming VPValue.
   VPValue *getIncomingValue(unsigned I) { return getOperand(I); }
 };
diff --git a/llvm/lib/Transforms/Vectorize/VPlanUtils.h b/llvm/lib/Transforms/Vectorize/VPlanUtils.h
index 6ddb88308955f1..ddbfc50ed96159 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanUtils.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanUtils.h
@@ -169,8 +169,19 @@ class VPBlockUtils {
   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.
+      while (auto *Region = dyn_cast<VPRegionBlock>(Succ))
+        Succ = Region->getEntry();
+
+      for (auto &R : *cast<VPBasicBlock>(Succ))
+        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();
diff --git a/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp b/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp
index e7987a95f1ca26..09f111ec31a4d4 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp
@@ -659,6 +659,45 @@ TEST_F(VPBasicBlockTest, TraversingIteratorTest) {
   }
 }
 
+TEST_F(VPBasicBlockTest, reassociateBlocks) {
+  {
+    VPlan &Plan = getPlan();
+    VPBasicBlock *VPBB1 = Plan.createVPBasicBlock("VPBB1");
+    VPBasicBlock *VPBB2 = Plan.createVPBasicBlock("VPBB2");
+    VPBlockUtils::connectBlocks(VPBB1, VPBB2);
+
+    auto *WidenPhi = new VPWidenPHIRecipe(nullptr);
+    IntegerType *Int32 = IntegerType::get(C, 32);
+    VPValue *Val = Plan.getOrAddLiveIn(ConstantInt::get(Int32, 1));
+    WidenPhi->addIncoming(Val, VPBB1);
+    VPBB2->appendRecipe(WidenPhi);
+
+    VPBasicBlock *VPBBNew = Plan.createVPBasicBlock("VPBBNew");
+    VPBlockUtils::reassociateBlocks(VPBB1, VPBBNew);
+    EXPECT_EQ(VPBB2->getSinglePredecessor(), VPBBNew);
+    EXPECT_EQ(WidenPhi->getIncomingBlock(0), VPBBNew);
+  }
+
+  {
+    VPlan &Plan = getPlan();
+    VPBasicBlock *VPBB1 = Plan.createVPBasicBlock("VPBB1");
+    VPBasicBlock *VPBB2 = Plan.createVPBasicBlock("VPBB2");
+    VPRegionBlock *R1 = Plan.createVPRegionBlock(VPBB2, VPBB2, "R1");
+    VPBlockUtils::connectBlocks(VPBB1, R1);
+
+    auto *WidenPhi = new VPWidenPHIRecipe(nullptr);
+    IntegerType *Int32 = IntegerType::get(C, 32);
+    VPValue *Val = Plan.getOrAddLiveIn(ConstantInt::get(Int32, 1));
+    WidenPhi->addIncoming(Val, VPBB1);
+    VPBB2->appendRecipe(WidenPhi);
+
+    VPBasicBlock *VPBBNew = Plan.createVPBasicBlock("VPBBNew");
+    VPBlockUtils::reassociateBlocks(VPBB1, VPBBNew);
+    EXPECT_EQ(R1->getSinglePredecessor(), VPBBNew);
+    EXPECT_EQ(WidenPhi->getIncomingBlock(0), VPBBNew);
+  }
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 TEST_F(VPBasicBlockTest, print) {
   VPInstruction *TC = new VPInstruction(Instruction::Add, {});