[llvm] [BasicBlockUtils] Fix SplitBlockPredecessors incorrect dominator insert (PR #107190)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Joshua Cao (caojoshua)

<details>
<summary>Changes</summary>

When SplitBlockPredecessors uses DomTreeUpdater, it always assumes that the newly created BB dominates the old BB. This is not true if the consumer only passes in a subset of producers. For example:

```
bb0:
  br label %bb2
bb1:
  br label %bb2
bb2: ...
```

If the consumer calls SplitBlockPredecessors on bb2 with only bb0 as the predecessor:

```
bb0:
  br label %new_bb
new_bb:
  br label %bb2
bb1:
  br label %bb2
bb2: ...
```

In this case, new_bb does NOT dominate bb2 because there is an alternative path to bb2 from bb1. We fix this by only marking the new BB as dominating the old BB(bb2) if the old BB has a single predecessor.

SplitBlockPredecessors with DominatorTree is deprecated in favor of using the function with DomTreeUpdater. We move over a consumer in LICM::splitPredecessorsOfLoopExit() in this patch, which would have triggered the error.

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


3 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/LICM.cpp (+3-1) 
- (modified) llvm/lib/Transforms/Utils/BasicBlockUtils.cpp (+4-1) 
- (modified) llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp (+71) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index 86c7dceffc5245..d16a2a5fcad3a9 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -44,6 +44,7 @@
 #include "llvm/Analysis/AliasSetTracker.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CaptureTracking.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/GuardUtils.h"
 #include "llvm/Analysis/LazyBlockFrequencyInfo.h"
 #include "llvm/Analysis/Loads.h"
@@ -1603,13 +1604,14 @@ static void splitPredecessorsOfLoopExit(PHINode *PN, DominatorTree *DT,
   //
   const auto &BlockColors = SafetyInfo->getBlockColors();
   SmallSetVector<BasicBlock *, 8> PredBBs(pred_begin(ExitBB), pred_end(ExitBB));
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
   while (!PredBBs.empty()) {
     BasicBlock *PredBB = *PredBBs.begin();
     assert(CurLoop->contains(PredBB) &&
            "Expect all predecessors are in the loop");
     if (PN->getBasicBlockIndex(PredBB) >= 0) {
       BasicBlock *NewPred = SplitBlockPredecessors(
-          ExitBB, PredBB, ".split.loop.exit", DT, LI, MSSAU, true);
+          ExitBB, PredBB, ".split.loop.exit", &DTU, LI, MSSAU, true);
       // Since we do not allow splitting EH-block with BlockColors in
       // canSplitPredecessors(), we can simply assign predecessor's color to
       // the new block.
diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index 4144c7993b7e42..24d9d023211499 100644
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -1160,7 +1160,10 @@ static void UpdateAnalysisInformation(BasicBlock *OldBB, BasicBlock *NewBB,
       // Split block expects NewBB to have a non-empty set of predecessors.
       SmallVector<DominatorTree::UpdateType, 8> Updates;
       SmallPtrSet<BasicBlock *, 8> UniquePreds;
-      Updates.push_back({DominatorTree::Insert, NewBB, OldBB});
+      if (OldBB->getSinglePredecessor()) {
+        assert(OldBB->getSinglePredecessor() == NewBB);
+        Updates.push_back({DominatorTree::Insert, NewBB, OldBB});
+      }
       Updates.reserve(Updates.size() + 2 * Preds.size());
       for (auto *Pred : Preds)
         if (UniquePreds.insert(Pred).second) {
diff --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
index 56692cf25b7972..8b452e852f4565 100644
--- a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
+++ b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
@@ -414,6 +414,77 @@ define i32 @basic_func(i1 %cond) {
   EXPECT_TRUE(DT.verify());
 }
 
+TEST(BasicBlockUtils, SplitBlockPredecessorsLinear) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+define i32 @basic_func() {
+entry:
+  br label %bb0
+bb0:
+  ret i32 0
+}
+)IR");
+  Function *F = M->getFunction("basic_func");
+  DominatorTree DT(*F);
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
+
+  BasicBlock *Entry = &F->getEntryBlock();
+  BasicBlock *BB0 = getBasicBlockByName(*F, "bb0");
+  SplitBlockPredecessors(BB0, {Entry}, "split.entry", &DTU);
+  EXPECT_TRUE(DT.verify());
+}
+
+TEST(BasicBlockUtils, SplitBlockPredecessorsSubsetOfPreds) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+define i32 @basic_func(i1 %cond) {
+entry:
+  br i1 %cond, label %bb0, label %bb1
+bb0:
+  br label %bb2
+bb1:
+  br label %bb2
+bb2:
+  %phi = phi i32 [ 0, %bb0 ], [ 1, %bb1 ]
+  ret i32 %phi
+}
+)IR");
+  Function *F = M->getFunction("basic_func");
+  DominatorTree DT(*F);
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
+
+  BasicBlock *BB0 = getBasicBlockByName(*F, "bb0");
+  BasicBlock *BB2 = getBasicBlockByName(*F, "bb2");
+  SplitBlockPredecessors(BB2, {BB0}, "split.entry", &DTU);
+  EXPECT_TRUE(DT.verify());
+}
+
+TEST(BasicBlockUtils, SplitBlockPredecessorsAllPreds) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+define i32 @basic_func(i1 %cond) {
+entry:
+  br i1 %cond, label %bb0, label %bb1
+bb0:
+  br label %bb2
+bb1:
+  br label %bb2
+bb2:
+  %phi = phi i32 [ 0, %bb0 ], [ 1, %bb1 ]
+  ret i32 %phi
+}
+)IR");
+  Function *F = M->getFunction("basic_func");
+  DominatorTree DT(*F);
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
+
+  BasicBlock *BB0 = getBasicBlockByName(*F, "bb0");
+  BasicBlock *BB1 = getBasicBlockByName(*F, "bb1");
+  BasicBlock *BB2 = getBasicBlockByName(*F, "bb2");
+  SplitBlockPredecessors(BB2, {BB0, BB1}, "split.entry", &DTU);
+  EXPECT_TRUE(DT.verify());
+}
+
 TEST(BasicBlockUtils, SplitCriticalEdge) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"IR(

``````````

</details>


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