[llvm] [LoopFusion] Validate loop structure before creating LoopCandidates (PR #192280)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Madhur Amilkanthwar (madhur13490)

<details>
<summary>Changes</summary>

This patch add early check in LoopFuse.cpp
to skip loops missing preheader, header, exiting,
exit block or latch. It also makes `reportInvalidCandidate` gracefully return false instead of asserting when Preheader is null.

In other words this patch skips loops which lack the required structure.

Fixes #<!-- -->156670.

---
Full diff: https://github.com/llvm/llvm-project/pull/192280.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/LoopFuse.cpp (+11-1) 
- (added) llvm/test/Transforms/LoopFusion/indirect-br.ll (+59) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/LoopFuse.cpp b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
index b577bf6584615..9c7500b0a290e 100644
--- a/llvm/lib/Transforms/Scalar/LoopFuse.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
@@ -363,7 +363,9 @@ struct FusionCandidate {
 
   bool reportInvalidCandidate(Statistic &Stat) const {
     using namespace ore;
-    assert(L && Preheader && "Fusion candidate not initialized properly!");
+    assert(L && "Fusion candidate has null loop!");
+    if (!Preheader)
+      return false;
 #if LLVM_ENABLE_STATS
     ++Stat;
     ORE.emit(OptimizationRemarkAnalysis(DEBUG_TYPE, Stat.getName(),
@@ -573,6 +575,14 @@ struct LoopFuser {
   /// Flow Equivalent sets, sorted by dominance.
   void collectFusionCandidates(const LoopVector &LV) {
     for (Loop *L : LV) {
+      // Skip loops that do not have the structure required for fusion (e.g.
+      // loops with indirectbr may lack a preheader). Creating a FusionCandidate
+      // for such loops can lead to reportInvalidCandidate being called with
+      // a null Preheader.
+      if (!L->getLoopPreheader() || !L->getHeader() || !L->getExitingBlock() ||
+          !L->getExitBlock() || !L->getLoopLatch())
+        continue;
+
       TTI::PeelingPreferences PP =
           gatherPeelingPreferences(L, SE, TTI, std::nullopt, std::nullopt);
       FusionCandidate CurrCand(L, DT, &PDT, ORE, PP);
diff --git a/llvm/test/Transforms/LoopFusion/indirect-br.ll b/llvm/test/Transforms/LoopFusion/indirect-br.ll
new file mode 100644
index 0000000000000..2f95a3a6c4b6f
--- /dev/null
+++ b/llvm/test/Transforms/LoopFusion/indirect-br.ll
@@ -0,0 +1,59 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes=loop-fusion < %s | FileCheck %s
+; Test that loop-fusion does not crash on functions with indirectbr and
+; loops that lack a preheader (or other required loop structure).
+
+define i32 @foo(ptr nocapture readonly %p) {
+; CHECK-LABEL: define i32 @foo(
+; CHECK-SAME: ptr readonly captures(none) [[P:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[TBL:%.*]] = alloca <2 x ptr>, align 16
+; CHECK-NEXT:    store <2 x ptr> <ptr blockaddress(@foo, %[[INDIRECT:.*]]), ptr blockaddress(@foo, %[[END:.*]])>, ptr [[TBL]], align 16
+; CHECK-NEXT:    br label %[[FOR_COND:.*]]
+; CHECK:       [[FOR_COND]]:
+; CHECK-NEXT:    [[CUR:%.*]] = phi ptr [ [[P]], %[[ENTRY]] ], [ [[NEXT:%.*]], %[[FOR_COND]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i8, ptr [[CUR]], align 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[TMP0]], 7
+; CHECK-NEXT:    [[NEXT]] = getelementptr inbounds i8, ptr [[CUR]], i64 1
+; CHECK-NEXT:    br i1 [[CMP]], label %[[INDIRECT_PREHEADER:.*]], label %[[FOR_COND]]
+; CHECK:       [[INDIRECT_PREHEADER]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = load i8, ptr [[NEXT]], align 1
+; CHECK-NEXT:    [[IDX:%.*]] = sext i8 [[TMP1]] to i64
+; CHECK-NEXT:    [[SLOT:%.*]] = getelementptr inbounds <2 x ptr>, ptr [[TBL]], i64 0, i64 [[IDX]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load ptr, ptr [[SLOT]], align 8
+; CHECK-NEXT:    br label %[[INDIRECT]]
+; CHECK:       [[INDIRECT]]:
+; CHECK-NEXT:    indirectbr ptr [[TMP2]], [label %[[INDIRECT]], label %[[END]]]
+; CHECK:       [[INDIRECT2:.*:]]
+; CHECK-NEXT:    indirectbr ptr [[TMP2]], [label %[[INDIRECT]], label %[[END]]]
+; CHECK:       [[END]]:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %tbl = alloca <2 x ptr>, align 16
+  store <2 x ptr> <ptr blockaddress(@foo, %indirect), ptr blockaddress(@foo, %exit)>, ptr %tbl, align 16
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.cond, %entry
+  %cur = phi ptr [ %p, %entry ], [ %next, %for.cond ]
+  %val = load i8, ptr %cur, align 1
+  %cmp = icmp eq i8 %val, 7
+  %next = getelementptr inbounds i8, ptr %cur, i64 1
+  br i1 %cmp, label %indirect.preheader, label %for.cond
+
+indirect.preheader:                               ; preds = %for.cond
+  %raw = load i8, ptr %next, align 1
+  %idx = sext i8 %raw to i64
+  %slot = getelementptr inbounds <2 x ptr>, ptr %tbl, i64 0, i64 %idx
+  %dest = load ptr, ptr %slot, align 8
+  br label %indirect
+
+indirect:
+  indirectbr ptr %dest, [label %indirect, label %exit]
+
+indirect2:
+  indirectbr ptr %dest, [label %indirect, label %exit]
+
+exit:
+  ret i32 0
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/192280