[llvm] 0856592 - Ensure `collectTransitivePredecessors` returns Pred only from the Loop. (#113831)

[via llvm-commits]

Author: Manish Kausik H
Date: 2024-10-31T11:08:15-07:00
New Revision: 0856592f6f8d2bdaf2d66017d073b8d55fc27552

URL: https://github.com/llvm/llvm-project/commit/0856592f6f8d2bdaf2d66017d073b8d55fc27552
DIFF: https://github.com/llvm/llvm-project/commit/0856592f6f8d2bdaf2d66017d073b8d55fc27552.diff

LOG: Ensure `collectTransitivePredecessors` returns Pred only from the Loop. (#113831)

It's possible that we encounter Irreducible control flow, due to which,
we may find that a few predecessors of BB are not a part of the CurLoop.
Currently we crash in the function for such cases. This patch ensures
that we only return Predecessors that are a part of CurLoop and
gracefully ignore other Predecessors.

For example, consider Irreducible IR of this form:
```
define i64 @baz() {
bb:
  br label %bb1

bb1:                                              ; preds = %bb3, %bb
  br label %bb3

bb2:                                              ; No predecessors!
  br label %bb3

bb3:                                              ; preds = %bb2, %bb1
  %load = load ptr addrspace(1), ptr addrspace(1) null, align 8
  br label %bb1
}
```

This crashes when `collectTransitivePredecessors` is called on the
`%bb1<Header>, %bb3<latch>` loop, because the loop body has a
predecessor `%bb2` which is not a part of the loop.

See https://godbolt.org/z/E9fM1q3cT for the crash

Added: 
    llvm/test/Analysis/MustExecute/irreducible-cfg.ll

Modified: 
    llvm/lib/Analysis/MustExecute.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/MustExecute.cpp b/llvm/lib/Analysis/MustExecute.cpp
index caed62679a683c..8e233a5c8d7a54 100644
--- a/llvm/lib/Analysis/MustExecute.cpp
+++ b/llvm/lib/Analysis/MustExecute.cpp
@@ -162,6 +162,9 @@ static bool CanProveNotTakenFirstIteration(const BasicBlock *ExitBlock,
 /// Collect all blocks from \p CurLoop which lie on all possible paths from
 /// the header of \p CurLoop (inclusive) to BB (exclusive) into the set
 /// \p Predecessors. If \p BB is the header, \p Predecessors will be empty.
+/// Note: It's possible that we encounter Irreducible control flow, due to
+/// which, we may find that a few predecessors of \p BB are not a part of the
+/// \p CurLoop. We only return Predecessors that are a part of \p CurLoop.
 static void collectTransitivePredecessors(
     const Loop *CurLoop, const BasicBlock *BB,
     SmallPtrSetImpl<const BasicBlock *> &Predecessors) {
@@ -171,6 +174,8 @@ static void collectTransitivePredecessors(
     return;
   SmallVector<const BasicBlock *, 4> WorkList;
   for (const auto *Pred : predecessors(BB)) {
+    if (!CurLoop->contains(Pred))
+      continue;
     Predecessors.insert(Pred);
     WorkList.push_back(Pred);
   }
@@ -187,7 +192,7 @@ static void collectTransitivePredecessors(
     // We can ignore backedge of all loops containing BB to get a sligtly more
     // optimistic result.
     for (const auto *PredPred : predecessors(Pred))
-      if (Predecessors.insert(PredPred).second)
+      if (CurLoop->contains(PredPred) && Predecessors.insert(PredPred).second)
         WorkList.push_back(PredPred);
   }
 }

diff  --git a/llvm/test/Analysis/MustExecute/irreducible-cfg.ll b/llvm/test/Analysis/MustExecute/irreducible-cfg.ll
new file mode 100644
index 00000000000000..a452761ab33566
--- /dev/null
+++ b/llvm/test/Analysis/MustExecute/irreducible-cfg.ll
@@ -0,0 +1,29 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -disable-output -passes=print-mustexecute %s 2>&1 | FileCheck %s
+
+; The loop body has two predecessors, %header and %side-entry. This leads to irreducible-cfg
+define i64 @baz() {
+; CHECK-LABEL: define i64 @baz() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[HEADER:.*]]
+; CHECK:       [[HEADER]]:
+; CHECK-NEXT:    br label %[[BODY:.*]] ; (mustexec in: header)
+; CHECK:       [[SIDE_ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[BODY]]
+; CHECK:       [[BODY]]:
+; CHECK-NEXT:    [[LOAD:%.*]] = load ptr addrspace(1), ptr addrspace(1) null, align 8 ; (mustexec in: header)
+; CHECK-NEXT:    br label %[[HEADER]] ; (mustexec in: header)
+;
+entry:
+  br label %header
+
+header:
+  br label %body
+
+side-entry:
+  br label %body
+
+body:
+  %load = load ptr addrspace(1), ptr addrspace(1) null, align 8
+  br label %header
+}