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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: None (Camsyn)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp (+10-1) 
- (modified) llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp (+57) 


``````````diff
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(

``````````

</details>


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