[llvm] [LoopFusion] Non-loop block must be the immediate successor of exit (PR #175034)

[Alireza Torabian via llvm-commits]

https://github.com/1997alireza updated https://github.com/llvm/llvm-project/pull/175034

>From a46c288267f72932dfa8c0a14d3fab62f74ad2e2 Mon Sep 17 00:00:00 2001
From: Alireza Torabian <alireza.torabian at huawei.com>
Date: Wed, 7 Jan 2026 13:26:34 -0500
Subject: [PATCH] [LoopFusion] Non-loop block must be the immediate successor
 of exit

Loop fusion assumes the non-loop block of a guarded adjacent loop
is the immediate successor of its exit block. This patch ensures
this condition is hold and fixes the crash #166356.
---
 llvm/lib/Transforms/Scalar/LoopFuse.cpp     | 11 ++---
 llvm/test/Transforms/LoopFusion/pr166356.ll | 50 +++++++++++++++++++++
 2 files changed, 56 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/LoopFusion/pr166356.ll

diff --git a/llvm/lib/Transforms/Scalar/LoopFuse.cpp b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
index 3a06c3f00fa02..d7284c98b48c3 100644
--- a/llvm/lib/Transforms/Scalar/LoopFuse.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
@@ -1399,10 +1399,11 @@ struct LoopFuser {
   /// This method will determine if there are additional basic blocks in the CFG
   /// between the exit of \p FC0 and the entry of \p FC1.
   /// If the two candidates are guarded loops, then it checks whether the
-  /// non-loop successor of the \p FC0 guard branch is the entry block of \p
-  /// FC1. If not, then the loops are not adjacent. If the two candidates are
-  /// not guarded loops, then it checks whether the exit block of \p FC0 is the
-  /// preheader of \p FC1.
+  /// exit block of the \p FC0 is the predecessor of the \p FC1 preheader. This
+  /// implicitly ensures that the non-loop successor of the \p FC0 guard branch
+  /// is the entry block of \p FC1. If not, then the loops are not adjacent. If
+  /// the two candidates are not guarded loops, then it checks whether the exit
+  /// block of \p FC0 is the preheader of \p FC1.
   /// Strictly means there is no predecessor for FC1 unless it is from FC0,
   /// i.e., FC0 dominates FC1.
   bool isStrictlyAdjacent(const FusionCandidate &FC0,
@@ -1410,7 +1411,7 @@ struct LoopFuser {
     // If the successor of the guard branch is FC1, then the loops are adjacent
     if (FC0.GuardBranch)
       return DT.dominates(FC0.getEntryBlock(), FC1.getEntryBlock()) &&
-             FC0.getNonLoopBlock() == FC1.getEntryBlock();
+             FC0.ExitBlock->getSingleSuccessor() == FC1.getEntryBlock();
     else
       return FC0.ExitBlock == FC1.getEntryBlock();
   }
diff --git a/llvm/test/Transforms/LoopFusion/pr166356.ll b/llvm/test/Transforms/LoopFusion/pr166356.ll
new file mode 100644
index 0000000000000..764f0ce3ffc72
--- /dev/null
+++ b/llvm/test/Transforms/LoopFusion/pr166356.ll
@@ -0,0 +1,50 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -passes=loop-fusion -S < %s 2>&1 | FileCheck %s
+
+; The non-loop block of loop for.cond.cleanup13 is not the immediate successor
+; of its exit block, then it should not be eligible for fusion.
+
+define void @non_immediate_exit() {
+; CHECK-LABEL: define void @non_immediate_exit() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br i1 false, label %[[IF_ELSE_1:.*]], label %[[FOR_COND_CLEANUP13_PREHEADER:.*]]
+; CHECK:       [[FOR_COND_CLEANUP13_PREHEADER]]:
+; CHECK-NEXT:    br label %[[FOR_COND_CLEANUP13:.*]]
+; CHECK:       [[FOR_COND_CLEANUP13]]:
+; CHECK-NEXT:    br i1 true, label %[[FOR_INC21:.*]], label %[[FOR_COND_CLEANUP13]]
+; CHECK:       [[FOR_INC21]]:
+; CHECK-NEXT:    br label %[[IF_THEN_1:.*]]
+; CHECK:       [[IF_THEN_1]]:
+; CHECK-NEXT:    br label %[[IF_ELSE_1]]
+; CHECK:       [[IF_ELSE_1]]:
+; CHECK-NEXT:    br i1 false, label %[[IF_ELSE_2:.*]], label %[[FOR_COND_CLEANUP13_1_PREHEADER:.*]]
+; CHECK:       [[FOR_COND_CLEANUP13_1_PREHEADER]]:
+; CHECK-NEXT:    br label %[[FOR_COND_CLEANUP13_1:.*]]
+; CHECK:       [[FOR_COND_CLEANUP13_1]]:
+; CHECK-NEXT:    br i1 true, label %[[IF_ELSE_2_LOOPEXIT:.*]], label %[[FOR_COND_CLEANUP13_1]]
+; CHECK:       [[IF_ELSE_2_LOOPEXIT]]:
+; CHECK-NEXT:    br label %[[IF_ELSE_2]]
+; CHECK:       [[IF_ELSE_2]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 false, label %if.else.1, label %for.cond.cleanup13
+
+for.cond.cleanup13:                               ; preds = %for.cond.cleanup13, %entry
+  br i1 true, label %for.inc21, label %for.cond.cleanup13
+
+for.inc21:                                        ; preds = %for.cond.cleanup13
+  br label %if.then.1
+
+if.then.1:                                        ; preds = %for.inc21
+  br label %if.else.1
+
+if.else.1:                                        ; preds = %if.then.1, %entry
+  br i1 false, label %if.else.2, label %for.cond.cleanup13.1
+
+for.cond.cleanup13.1:                             ; preds = %for.cond.cleanup13.1, %if.else.1
+  br i1 true, label %if.else.2, label %for.cond.cleanup13.1
+
+if.else.2:                                        ; preds = %for.cond.cleanup13.1, %if.else.1
+  ret void
+}