[llvm] [BranchFolding] Fix assertion failure in HoistCommonCodeInSuccs (PR #141028)

[Orlando Cazalet-Hyams via llvm-commits]

https://github.com/OCHyams updated https://github.com/llvm/llvm-project/pull/141028

>From e9de2d886dc5aeb3ea1da6ef9d3e8cd38388d8f5 Mon Sep 17 00:00:00 2001
From: Orlando Cazalet-Hyams <orlando.hyams at sony.com>
Date: Thu, 22 May 2025 10:34:44 +0100
Subject: [PATCH 1/2] [BranchFolding] Fix assertion failure in
 HoistCommonCodeInSuccs

Assertion failure introduced in #140063, which didn't account for TBB and FBB
being the same block.
---
 llvm/lib/CodeGen/BranchFolding.cpp            | 10 ++-
 .../ARM/branch-folder-single-bb-crash.mir     | 80 +++++++++++++++++++
 2 files changed, 86 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/CodeGen/ARM/branch-folder-single-bb-crash.mir

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index d10e6722f4548..41e66bb82650e 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -2072,9 +2072,11 @@ bool BranchFolder::HoistCommonCodeInSuccs(MachineBasicBlock *MBB) {
     return false;
 
   // Hoist the instructions from [T.begin, TIB) and then delete [F.begin, FIB).
-  // Merge the debug locations. FIXME: We should do something with the
-  // debug instructions too (from BOTH branches).
-  {
+  // If we're hoisting from a single block then just splice. Else step through
+  // and merge the debug locations.
+  if (TBB == FBB) {
+    MBB->splice(Loc, TBB, TBB->begin(), TIB);
+  } else {
     // TIB and FIB point to the end of the regions to hoist/merge in TBB and
     // FBB.
     MachineBasicBlock::iterator FE = FIB;
@@ -2083,7 +2085,7 @@ bool BranchFolder::HoistCommonCodeInSuccs(MachineBasicBlock *MBB) {
          make_early_inc_range(make_range(TBB->begin(), TIB))) {
       // Move debug instructions and pseudo probes without modifying them.
       // FIXME: This is the wrong thing to do for debug locations, which
-      // should at least be killed.
+      // should at least be killed (and hoisted from BOTH blocks).
       if (TI->isDebugOrPseudoInstr()) {
         TI->moveBefore(&*Loc);
         continue;
diff --git a/llvm/test/CodeGen/ARM/branch-folder-single-bb-crash.mir b/llvm/test/CodeGen/ARM/branch-folder-single-bb-crash.mir
new file mode 100644
index 0000000000000..56bcd3d73d616
--- /dev/null
+++ b/llvm/test/CodeGen/ARM/branch-folder-single-bb-crash.mir
@@ -0,0 +1,80 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+# RUN: llc -mtriple=armv7-apple-ios --run-pass=branch-folder %s -o - | FileCheck %s
+
+## Don't crash in branch-folder's HoistCommonCodeInSuccs when TBB and FBB are
+## the same block.
+
+...
+---
+name:            f_1418_2054
+tracksRegLiveness: true
+noPhis:          true
+isSSA:           false
+noVRegs:         true
+body:             |
+  ; CHECK-LABEL: name: f_1418_2054
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   $r0 = MOVr undef $r0, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK-NEXT:   Bcc %bb.1, 1 /* CC::ne */, undef $cpsr
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   successors: %bb.3(0x40000000), %bb.4(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   $r0 = MOVr undef $r0, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK-NEXT:   Bcc %bb.4, 1 /* CC::ne */, undef $cpsr
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.3:
+  ; CHECK-NEXT:   BX_RET 14 /* CC::al */, $noreg
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.4:
+  ; CHECK-NEXT:   TRAP
+  bb.0:
+    successors: %bb.1
+
+  bb.1:
+    successors: %bb.3, %bb.2
+
+    Bcc %bb.3, 1, undef $cpsr
+    B %bb.2
+
+  bb.2:
+    successors: %bb.5, %bb.6
+
+    Bcc %bb.5, 1, undef $cpsr
+    B %bb.6
+
+  bb.3:
+    successors: %bb.4
+
+    $r0 = MOVr undef $r0, 14, $noreg, $noreg
+
+  bb.4:
+    successors: %bb.1
+
+    B %bb.1
+
+  bb.5:
+    successors: %bb.6
+
+  bb.6:
+    successors: %bb.8, %bb.7
+
+    $r0 = MOVr undef $r0, 14, $noreg, $noreg
+    $r0 = MOVr undef $r0, 14, $noreg, $noreg
+
+    Bcc %bb.7, 1, undef $cpsr
+    B %bb.8
+
+  bb.7:
+    successors:
+
+    TRAP
+
+  bb.8:
+    BX_RET 14, _
+...

>From 7ea4874533588146891abf524a85ea6dd465521c Mon Sep 17 00:00:00 2001
From: Orlando Cazalet-Hyams <orlando.hyams at sony.com>
Date: Thu, 22 May 2025 13:17:16 +0100
Subject: [PATCH 2/2] better comment

---
 .../test/CodeGen/ARM/branch-folder-single-bb-crash.mir | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/llvm/test/CodeGen/ARM/branch-folder-single-bb-crash.mir b/llvm/test/CodeGen/ARM/branch-folder-single-bb-crash.mir
index 56bcd3d73d616..0cb5f997784b6 100644
--- a/llvm/test/CodeGen/ARM/branch-folder-single-bb-crash.mir
+++ b/llvm/test/CodeGen/ARM/branch-folder-single-bb-crash.mir
@@ -1,8 +1,14 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
 # RUN: llc -mtriple=armv7-apple-ios --run-pass=branch-folder %s -o - | FileCheck %s
 
-## Don't crash in branch-folder's HoistCommonCodeInSuccs when TBB and FBB are
-## the same block.
+## We hoist code from successor blocks, and in abnormal circumstances we can
+## have control flow that branches to the same block. Test we handle that
+## correctly.
+##
+## bb.5 (empty fall-through) folds into bb.6, making bb.2 branch to bb.6
+## unconditionally on a conditional branch (both edges go to bb.6).
+## Instructions are hosited from bb.2's successors, which is bb.6 twice,
+## along both edges.
 
 ...
 ---