[llvm] [Utils] Fix duplicate DomTree updates in SplitIndirectBrCriticalEdges (PR #196475)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Maosu Zhao (zhaomaosu)

<details>
<summary>Changes</summary>

SplitIndirectBrCriticalEdges generates DomTree Insert/Delete pairs for each predecessor in OtherPreds. However, OtherPreds can contain duplicate entries when a conditional branch has both targets pointing to the same block (e.g., `br i1 %c, label %X, label %X`). This produces duplicate DomTree updates for the same edge, triggering the assertion `std::abs(NumInsertions) <= 1 && "Unbalanced operations!"` in LegalizeUpdates.

Fix by tracking which source blocks have already had DomTree updates emitted, and skipping duplicates.

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp (+5-1) 
- (added) llvm/test/Transforms/CodeGenPrepare/X86/split-indirectbr-duplicate-pred.ll (+25) 


``````````diff
diff --git a/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp b/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
index fd6b2562bee54..65bff2b878bda 100644
--- a/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
+++ b/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
@@ -416,6 +416,7 @@ bool llvm::SplitIndirectBrCriticalEdges(Function &F,
 
     BlockFrequency BlockFreqForDirectSucc;
     SmallVector<DominatorTree::UpdateType, 8> DTUpdates;
+    SmallPtrSet<BasicBlock *, 8> SeenSrcs;
     if (DTU)
       DTUpdates.reserve(OtherPreds.size() * 2 + 1);
     for (BasicBlock *Pred : OtherPreds) {
@@ -426,7 +427,10 @@ bool llvm::SplitIndirectBrCriticalEdges(Function &F,
       if (ShouldUpdateAnalysis)
         BlockFreqForDirectSucc += BFI->getBlockFreq(Src) *
             BPI->getEdgeProbability(Src, DirectSucc);
-      if (DTU) {
+      // A predecessor may appear multiple times in OtherPreds (e.g., a CondBr
+      // with both targets pointing to the same block). Only emit one pair of
+      // DomTree updates per unique source.
+      if (DTU && SeenSrcs.insert(Src).second) {
         DTUpdates.push_back({DominatorTree::Insert, Src, DirectSucc});
         DTUpdates.push_back({DominatorTree::Delete, Src, Target});
       }
diff --git a/llvm/test/Transforms/CodeGenPrepare/X86/split-indirectbr-duplicate-pred.ll b/llvm/test/Transforms/CodeGenPrepare/X86/split-indirectbr-duplicate-pred.ll
new file mode 100644
index 0000000000000..978684fa61073
--- /dev/null
+++ b/llvm/test/Transforms/CodeGenPrepare/X86/split-indirectbr-duplicate-pred.ll
@@ -0,0 +1,25 @@
+; RUN: opt -passes='require<profile-summary>,function(codegenprepare)' -S -mtriple=x86_64 < %s | FileCheck %s
+
+; Test that SplitIndirectBrCriticalEdges does not crash when a predecessor
+; block has a conditional branch with both targets pointing to the same
+; destination (producing duplicate entries in the predecessor list).
+
+; CHECK-LABEL: @duplicate_pred_condbr
+; CHECK: indirectbr ptr %addr, [label %target, label %cond.bb]
+; CHECK: cond.bb:
+; CHECK-NEXT: br label %.split
+; CHECK: .split:
+; CHECK-NEXT: %merge = phi i32
+; CHECK-NEXT: ret i32 %merge
+
+define i32 @duplicate_pred_condbr(ptr %addr, i1 %cond) {
+entry:
+  indirectbr ptr %addr, [label %target, label %cond.bb]
+
+cond.bb:
+  br i1 %cond, label %target, label %target
+
+target:
+  %result = phi i32 [ 0, %entry ], [ 1, %cond.bb ], [ 1, %cond.bb ]
+  ret i32 %result
+}

``````````

</details>


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