[Utils] Fix incorrect LCSSA PHI nodes when splitting critical edges with MergeIdenticalEdges (#131744)

This PR fixes incorrect LCSSA PHI node generation when splitting
critical edges with both
`PreserveLCSSA` and `MergeIdenticalEdges` enabled. The bug caused PHI
nodes in the split block
to miss predecessors when multiple identical edges were merged.

Author: Camsyn

llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp

@@ -207,6 +207,8 @@ llvm::SplitKnownCriticalEdge(Instruction *TI, unsigned SuccNum,
     }
   }
 
+  unsigned NumSplitIdenticalEdges = 1;
+
   // If there are any other edges from TIBB to DestBB, update those to go
   // through the split block, making those edges non-critical as well (and
   // reducing the number of phi entries in the DestBB if relevant).
@@ -219,6 +221,9 @@ llvm::SplitKnownCriticalEdge(Instruction *TI, unsigned SuccNum,
 
       // We found another edge to DestBB, go to NewBB instead.
       TI->setSuccessor(i, NewBB);
+
+      // Record the number of split identical edges to DestBB.
+      NumSplitIdenticalEdges++;
     }
   }
 
@@ -290,7 +295,11 @@ llvm::SplitKnownCriticalEdge(Instruction *TI, unsigned SuccNum,
 
         // Update LCSSA form in the newly created exit block.
         if (Options.PreserveLCSSA) {
-          createPHIsForSplitLoopExit(TIBB, NewBB, DestBB);
+          // If > 1 identical edges to be split, we need to introduce the same
+          // number of the incoming blocks for the new PHINode.
+          createPHIsForSplitLoopExit(
+              SmallVector<BasicBlock *, 4>(NumSplitIdenticalEdges, TIBB), NewBB,
+              DestBB);
         }
 
         if (!LoopPreds.empty()) {

llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp

@@ -438,6 +438,63 @@ define void @crit_edge(i1 %cond0, i1 %cond1) {
   EXPECT_TRUE(PDT.verify());
 }
 
+TEST(BasicBlockUtils, SplitLoopCriticalEdge) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+declare i1 @predicate(ptr %p)
+
+define ptr @Parse(ptr %gp) {
+entry:
+  br label %for.inc
+
+for.inc:
+  %phi = phi ptr [ %gp, %entry ], [ %cp, %while.cond ], [ %cp, %while.cond ]
+  %cond = call i1 @predicate(ptr %phi)
+  %inc= getelementptr inbounds i8, ptr %phi, i64 1
+  br i1 %cond, label %while.cond, label %exit
+
+while.cond:
+  %cp = phi ptr [ %inc, %for.inc ], [ %incdec, %while.body ]
+  %val = load i8, ptr %cp, align 1
+  switch i8 %val, label %while.body [
+    i8 10, label %for.inc
+    i8 0, label %for.inc
+  ]
+
+while.body:
+  %incdec = getelementptr inbounds i8, ptr %cp, i64 1
+  br label %while.cond
+
+exit:
+  ret ptr %phi
+}
+)IR");
+  Function *F = M->getFunction("Parse");
+  DominatorTree DT(*F);
+  LoopInfo LI(DT);
+
+  CriticalEdgeSplittingOptions CESO =
+      CriticalEdgeSplittingOptions(nullptr, &LI, nullptr)
+          .setMergeIdenticalEdges()
+          .setPreserveLCSSA();
+  EXPECT_EQ(2u, SplitAllCriticalEdges(*F, CESO));
+
+  BasicBlock *WhileBB = getBasicBlockByName(*F, "while.cond");
+  BasicBlock *SplitBB = WhileBB->getTerminator()->getSuccessor(1);
+  // The only 1 successor of SplitBB is %for.inc
+  EXPECT_EQ(1u, SplitBB->getTerminator()->getNumSuccessors());
+  // MergeIdenticalEdges: SplitBB has two identical predecessors, %while.cond.
+  EXPECT_EQ(WhileBB, SplitBB->getUniquePredecessor());
+  EXPECT_TRUE(SplitBB->hasNPredecessors(2));
+
+  auto *PN = dyn_cast<PHINode>(&SplitBB->front());
+  // PreserveLCSSA: should insert a PHI node in front of SplitBB
+  EXPECT_NE(nullptr, PN);
+  // The PHI node should have 2 identical incoming blocks.
+  EXPECT_EQ(2u, PN->getNumIncomingValues());
+  EXPECT_EQ(PN->getIncomingBlock(0), PN->getIncomingBlock(1));
+}
+
 TEST(BasicBlockUtils, SplitIndirectBrCriticalEdgesIgnorePHIs) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"IR(