[LV] Update incoming blocks in VPWidenPHIRecipe in reassociateBlocks

[lukel97]

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.

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: Luke Lau (lukel97)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/125481.diff

3 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/VPlan.h (+5)
• (modified) llvm/lib/Transforms/Vectorize/VPlanUtils.h (+12-1)
• (modified) llvm/unittests/Transforms/Vectorize/VPlanTest.cpp (+39)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 459222234bc37f5..cb58abec61c11c9 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 6ddb88308955f1a..ddbfc50ed961596 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 e7987a95f1ca26a..09f111ec31a4d4c 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, {});

[fhahn]

Suggested change

	for (auto &R : *cast<VPBasicBlock>(Succ))
	for (auto &R : cast<VPBasicBlock>(Succ)->phis())

[fhahn]

can you use Succ->getEntryBasicBlock()?

[lukel97]

Ah that's exactly the method I was looking for, thanks!

[fhahn]

Could you add a brief comment what this is testing?

[lukel97]

I thought I had responded to the comments + pushed my changes earlier, but turns out GitHub was stuck "processing updates" so it didn't show up in the PR, sorry about that. I had to push an empty commit to get things to refresh

[fhahn]

LGTM, thanks