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

[David Green via llvm-commits]

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

>From 225b677bb17901c9944b86c0c14f5b2ea899079e Mon Sep 17 00:00:00 2001
From: David Green <david.green at arm.com>
Date: Mon, 26 Jan 2026 18:10:58 +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..b27139ca7ac18 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
+}