[llvm] [ShrinkWrap] Ensure we do not crash on unreachable blocks. (PR #178009)

[David Green via llvm-commits]

https://github.com/davemgreen created https://github.com/llvm/llvm-project/pull/178009

Since we started optimizating always-true branches in the AArch64 backend (like cbz wzr), shrink wrap has been exposed to some block structures that it does not handle correctly, usually with unreachable blocks. This prevents the call to FindIDom/findNearestCommonDominator from failing when looking at the predecessors of a loop if one is unreachable.

Fixes: #177866

>From 2036b32e0af387c84db188bebdf337e64308036a Mon Sep 17 00:00:00 2001
From: David Green <david.green at arm.com>
Date: Mon, 26 Jan 2026 17:56:20 +0000
Subject: [PATCH] [ShrinkWrap] Ensure we do not crash on unreachable blocks.

Since we started optimizating always-true branches in the AArch64 backend (like
cbz wzr), shrink wrap has been exposed to some block structures that it does
not handle correctly, usually with unreachable blocks. This prevents the call
to FindIDom/findNearestCommonDominator from failing when looking at the
predecessors of a loop if one is unreachable.

Fixes: #177866
---
 llvm/lib/CodeGen/ShrinkWrap.cpp               |  6 +++-
 .../AArch64/shrink-wrap-unreachable.ll        | 36 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AArch64/shrink-wrap-unreachable.ll

diff --git a/llvm/lib/CodeGen/ShrinkWrap.cpp b/llvm/lib/CodeGen/ShrinkWrap.cpp
index 0cf21b261fdbb..8285d4fd60622 100644
--- a/llvm/lib/CodeGen/ShrinkWrap.cpp
+++ b/llvm/lib/CodeGen/ShrinkWrap.cpp
@@ -782,7 +782,11 @@ void ShrinkWrapImpl::updateSaveRestorePoints(MachineBasicBlock &MBB,
       if (MLI->getLoopDepth(Save) > MLI->getLoopDepth(Restore)) {
         // Push Save outside of this loop if immediate dominator is different
         // from save block. If immediate dominator is not different, bail out.
-        Save = FindIDom<>(*Save, Save->predecessors(), *MDT);
+        SmallVector<MachineBasicBlock *> Preds;
+        for(auto *PBB : Save->predecessors())
+          if (MDT->isReachableFromEntry(PBB))
+            Preds.push_back(PBB);
+        Save = FindIDom<>(*Save, Preds, *MDT);
         if (!Save)
           break;
       } else {
diff --git a/llvm/test/CodeGen/AArch64/shrink-wrap-unreachable.ll b/llvm/test/CodeGen/AArch64/shrink-wrap-unreachable.ll
new file mode 100644
index 0000000000000..414045591b691
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/shrink-wrap-unreachable.ll
@@ -0,0 +1,36 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=aarch64-linux-gnu -o - %s | FileCheck %s
+
+; Ensure we do not crash on unreachable blocks.
+
+define void @func_44(i1 %tobool3.not, ptr %g_530) {
+; CHECK-LABEL: func_44:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:  .LBB0_1: // %for.cond2
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    ldr wzr, [x8]
+; CHECK-NEXT:    b .LBB0_1
+entry:
+  br label %for.cond2.outer
+
+for.cond2.outer:                                  ; preds = %for.inc, %entry
+  %add23.ph = phi i8 [ 0, %for.inc ], [ 0, %entry ]
+  br label %for.cond2
+
+for.cond2:                                        ; preds = %cleanup12.split, %for.cond2, %for.cond2.outer
+  %load = load volatile i32, ptr null, align 4
+  br i1 %tobool3.not, label %if.end, label %for.cond2
+
+if.end:                                           ; preds = %for.cond2
+  %tobool7.not = icmp eq i64 0, 0
+  br i1 %tobool7.not, label %for.inc, label %cleanup12.split
+
+for.inc:                                          ; preds = %if.end
+  %add = or i8 %add23.ph, 0
+  br label %for.cond2.outer
+
+cleanup12.split:                                  ; preds = %if.end
+  store i64 0, ptr %g_530, align 8
+  br label %for.cond2
+}