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

[via llvm-commits]

https://github.com/Camsyn created https://github.com/llvm/llvm-project/pull/131744

#### Description  
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.  

---

#### Issue Details  
When splitting edges from blocks with >1 identical edges to a loop exit block (e.g., via `switch`  
instructions), the interaction between these options leads to invalid PHI nodes:  
1. **MergeIdenticalEdges** merges identical edges into a single edge to the new split block.  
2. **PreserveLCSSA** would only insert one PHI entry for the merged edge, ignoring the original  
   number of identical edges.  

**Example** (simplified from function `Parse` in `lemon.c` from SQLite3 in FuzzBench):  
```llvm  
; Original IR  
exiting:  
  switch i8 %cond, label %while.body [  
      i8 0, label %exit  
      i8 10, label %exit  
  ]  
exit:  
  %phi = phi i32 [%val, %exiting], [%val, %exiting]  ; Requires two predecessors  
```  

After splitting with `PreserveLCSSA` and `MergeIdenticalEdges`, the PHI node in the split block  
**incorrectly** became:  
```llvm  
split:  
  %lcssa = phi i32 [%val, %exiting]   ; Missing one predecessor!  
  br label %exit  
```  
instead of the correct:  
```llvm  
split:  
  %lcssa = phi i32 [%val, %exiting], [%val, %exiting]  ; Both edges preserved  
  br label %exit  
```  

---

#### Fix  
Modify `createPHIsForSplitLoopExit` to accept all predecessor entries corresponding to the number  
of merged edges. Instead of passing a single predecessor block, we now pass `NumSplittedIdenticalEdges`  
copies of the original block, ensuring PHI nodes have correct incoming entries.  


---

#### Impact  
Fixes miscompilations in edge cases where critical edge splitting interacts with LCSSA and  
edge merging. No known regressions introduced.  

---

#### References  
- Original bug context: `BreakCriticalEdges.cpp` (lines 291)  
- Example derived from FuzzBench/SQLite3’s `Parse` function in `lemon.c`.  

>From 1563edcf81023e6d370df48ae1cf44ecfb1761f0 Mon Sep 17 00:00:00 2001
From: Camsyn <65994555+Camsyn at users.noreply.github.com>
Date: Tue, 18 Mar 2025 15:07:01 +0800
Subject: [PATCH] [Utils] Fix incorrect LCSSA PHI nodes when splitting critical
 edges with MergeIdenticalEdges

When splitting a critical edge from a block with multiple identical
edges to an exit block while both `PreserveLCSSA` and
`MergeIdenticalEdges` are enabled, the generated LCSSA PHI
nodes in the split block would miss incoming values from merged edges.

This occurs because:
1. `MergeIdenticalEdges` merges multiple identical edges into a single
    edge to the new split block.
2. `PreserveLCSSA` (in `createPHIsForSplitLoopExit`) previously assumed
    only one incoming edge from the original block, creating PHI nodes
    with incomplete predecessors.

The fix modifies `createPHIsForSplitLoopExit` to account for merged
edges by passing all original predecessor entries (matching the number
of merged edges) when creating PHI nodes.

This ensures all incoming values are properly reflected in the split
block's PHI nodes.

Add unittest case in `BasicBlockUtilsTest.cpp` to verify the correction
of PHI node generation in this scenario.
---
 .../Transforms/Utils/BreakCriticalEdges.cpp   | 11 +++-
 .../Transforms/Utils/BasicBlockUtilsTest.cpp  | 57 +++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp b/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
index 62b4b545f29bb..075a1afb4f6d9 100644
--- a/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
+++ b/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
@@ -205,6 +205,8 @@ llvm::SplitKnownCriticalEdge(Instruction *TI, unsigned SuccNum,
     }
   }
 
+  unsigned NumSplittedIdenticalEdges = 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).
@@ -217,6 +219,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 splitted identical edges to DestBB.
+      NumSplittedIdenticalEdges++;
     }
   }
 
@@ -288,7 +293,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 splitted, we need to introduce 
+          // the incoming blocks of the same number for the new PHINode.
+          createPHIsForSplitLoopExit(
+            SmallVector<BasicBlock *, 4>(NumSplittedIdenticalEdges, TIBB),
+            NewBB, DestBB);
         }
 
         if (!LoopPreds.empty()) {
diff --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
index 56692cf25b797..a2f17ea672833 100644
--- a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
+++ b/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 dso_local i1 @predicate(ptr noundef %p)
+
+define dso_local ptr @Parse(ptr noundef %gp) {
+entry:
+  br label %for.inc
+
+for.inc:
+  %phi = phi ptr [ %gp, %entry ], [ %cp, %while.cond ], [ %cp, %while.cond ]
+  %cond = call i1 @predicate(ptr noundef %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
+  ASSERT_EQ(1u, SplitBB->getTerminator()->getNumSuccessors());
+  // MergeIdenticalEdges: SplitBB has two identical predecessors, %while.cond.
+  ASSERT_EQ(WhileBB, SplitBB->getUniquePredecessor());
+  ASSERT_EQ(true, SplitBB->hasNPredecessors(2));
+
+  PHINode *PN = dyn_cast<PHINode>(&(SplitBB->front()));
+  // PreserveLCSSA: should insert a PHI node in front of SplitBB
+  ASSERT_NE(nullptr, PN);
+  // The PHI node should have 2 identical incoming blocks.
+  ASSERT_EQ(2u, PN->getNumIncomingValues());
+  ASSERT_EQ(PN->getIncomingBlock(0), PN->getIncomingBlock(1));
+}
+
 TEST(BasicBlockUtils, SplitIndirectBrCriticalEdgesIgnorePHIs) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"IR(