[llvm] [BasicBlockUtils] Fix loopinfo header in split block before (PR #179408)

[Yevgeny Rouban via llvm-commits]

https://github.com/yrouban created https://github.com/llvm/llvm-project/pull/179408


SplitBlock() breaks LoopInfo. Splitting a basic block BB into a pair of blocks NewBB->BB the call SplitBlock(... DT, &LI,... /* Before */ true) does not change NewBB to be the new head of the loop if BB was the loop head before the split.

This PR fixes the loop header in splitBlockBefore(). If BB was the loop header before NewBB = splitBlockBefore(BB, ...) then the NewBB is set to be the new header of the loop.



>From c0e312713d7e99bb44a0eb75a27cd096decc934f Mon Sep 17 00:00:00 2001
From: Yevgeny Rouban <yrouban at gmail.com>
Date: Tue, 3 Feb 2026 09:02:27 +0300
Subject: [PATCH 1/2] [BasicBlockUtilsTests] SplitBlock breaks LoopInfo. NFC

Splitting a basic block BB into a pair of blocks NewBB->BB
SplitBlock(... DT, &LI,... /* Before */ true) does not change
NewBB to be the new head of the loop if BB was the loop head
before the split.

This change just introduces a new unit test that demonstrates
the bug.
---
 .../Transforms/Utils/BasicBlockUtilsTest.cpp  | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
index 00d9e9ff81e05..37b20d00d18d9 100644
--- a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
+++ b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
@@ -26,6 +26,8 @@
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
+#define DEBUG_TYPE "basic-block-utils-tests"
+
 using namespace llvm;
 
 static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
@@ -354,6 +356,38 @@ define void @foo() {
 }
 #endif
 
+#ifndef NDEBUG
+TEST(BasicBlockUtils, SplitBlockBefore) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+define void @split-block-before-test(i1 %flag) {
+entry:
+  br label %loop
+
+loop:
+  br i1 %flag, label %loop, label %exit
+
+exit:
+  ret void
+}
+)IR");
+  Function *F = M->getFunction("split-block-before-test");
+  DominatorTree DT(*F);
+  LoopInfo LI(DT);
+
+  EXPECT_TRUE(DT.verify());
+  LI.verify(DT);
+  auto *LoopBB = getBasicBlockByName(*F, "loop");
+  SplitBlock(LoopBB, LoopBB->getFirstInsertionPt(), &DT, &LI,
+              /* MemorySSAUpdater */ nullptr, LoopBB->getName() + ".split",
+              /* Before */ true);
+
+  EXPECT_TRUE(DT.verify());
+  LLVM_DEBUG(F->dump(); LI.print(dbgs()));
+  EXPECT_DEATH({LI.verify(DT);}, "Assertion.*Loop is unreachable!");
+}
+#endif
+
 TEST(BasicBlockUtils, NoUnreachableBlocksToEliminate) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"IR(

>From ef0bf05db95cb097cba0e94f0c732cd1019f718a Mon Sep 17 00:00:00 2001
From: Yevgeny Rouban <yrouban at gmail.com>
Date: Tue, 3 Feb 2026 10:38:06 +0300
Subject: [PATCH 2/2] [BasicBlockUtils] Fix loop header in splitBlockBefore()

If BB was the loop header before NewBB = splitBlockBefore(BB, ...)
then the NewBB must become the new header of the loop.
---
 llvm/lib/Transforms/Utils/BasicBlockUtils.cpp          |  5 ++++-
 .../unittests/Transforms/Utils/BasicBlockUtilsTest.cpp | 10 ++++++----
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index b0c04086f5e84..381cb7170c0f6 100644
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -1080,8 +1080,11 @@ BasicBlock *llvm::splitBlockBefore(BasicBlock *Old, BasicBlock::iterator SplitPt
   // The new block lives in whichever loop the old one did. This preserves
   // LCSSA as well, because we force the split point to be after any PHI nodes.
   if (LI)
-    if (Loop *L = LI->getLoopFor(Old))
+    if (Loop *L = LI->getLoopFor(Old)) {
       L->addBasicBlockToLoop(New, *LI);
+      if (L->getHeader() == Old)
+        L->moveToHeader(New);
+    }
 
   if (DTU) {
     SmallVector<DominatorTree::UpdateType, 8> DTUpdates;
diff --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
index 37b20d00d18d9..198518e130551 100644
--- a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
+++ b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
@@ -378,13 +378,15 @@ define void @split-block-before-test(i1 %flag) {
   EXPECT_TRUE(DT.verify());
   LI.verify(DT);
   auto *LoopBB = getBasicBlockByName(*F, "loop");
-  SplitBlock(LoopBB, LoopBB->getFirstInsertionPt(), &DT, &LI,
-              /* MemorySSAUpdater */ nullptr, LoopBB->getName() + ".split",
-              /* Before */ true);
+  auto *New =
+      SplitBlock(LoopBB, LoopBB->getFirstInsertionPt(), &DT, &LI,
+                 /* MemorySSAUpdater */ nullptr, LoopBB->getName() + ".split",
+                 /* Before */ true);
 
   EXPECT_TRUE(DT.verify());
   LLVM_DEBUG(F->dump(); LI.print(dbgs()));
-  EXPECT_DEATH({LI.verify(DT);}, "Assertion.*Loop is unreachable!");
+  LI.verify(DT);
+  EXPECT_EQ(LI.getLoopFor(New)->getHeader(), New);
 }
 #endif