From c60cdb44f7ecb4b02edf8b3ada707e8a74e79211 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo@fhahn.com>
Date: Mon, 7 Mar 2022 12:01:15 +0000
Subject: [PATCH] [ConstraintElimination] Only add cond from assume to succs if
 valid.

Add missing CanAdd check before adding a condition from an assume
to the successor blocks. When adding information from assume to
successor blocks we need to perform the same CanAdd as we do for adding
a condition from a branch.

Fixes https://github.com/llvm/llvm-project/issues/54217
---
 .../Transforms/Scalar/ConstraintElimination.cpp    | 30 ++++++++++++----------
 .../Transforms/ConstraintElimination/assumes.ll    |  2 +-
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index af90043..c094b5c 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -413,6 +413,19 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT) {
       continue;
     WorkList.emplace_back(DT.getNode(&BB));
 
+    // Returns true if we can add a known condition from BB to its successor
+    // block Succ. Each predecessor of Succ can either be BB or be dominated by
+    // Succ (e.g. the case when adding a condition from a pre-header to a loop
+    // header).
+    auto CanAdd = [&BB, &DT](BasicBlock *Succ) {
+      assert(isa<BranchInst>(BB.getTerminator()));
+      return any_of(successors(&BB),
+                    [Succ](const BasicBlock *S) { return S != Succ; }) &&
+             all_of(predecessors(Succ), [&BB, &DT, Succ](BasicBlock *Pred) {
+               return Pred == &BB || DT.dominates(Succ, Pred);
+             });
+    };
+
     // True as long as long as the current instruction is guaranteed to execute.
     bool GuaranteedToExecute = true;
     // Scan BB for assume calls.
@@ -431,9 +444,12 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT) {
           WorkList.emplace_back(DT.getNode(&BB), cast<ICmpInst>(Cond), false);
         } else {
           // Otherwise the condition only holds in the successors.
-          for (BasicBlock *Succ : successors(&BB))
+          for (BasicBlock *Succ : successors(&BB)) {
+            if (!CanAdd(Succ))
+              continue;
             WorkList.emplace_back(DT.getNode(Succ), cast<ICmpInst>(Cond),
                                   false);
+          }
         }
       }
       GuaranteedToExecute &= isGuaranteedToTransferExecutionToSuccessor(&I);
@@ -443,18 +459,6 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT) {
     if (!Br || !Br->isConditional())
       continue;
 
-    // Returns true if we can add a known condition from BB to its successor
-    // block Succ. Each predecessor of Succ can either be BB or be dominated by
-    // Succ (e.g. the case when adding a condition from a pre-header to a loop
-    // header).
-    auto CanAdd = [&BB, &DT](BasicBlock *Succ) {
-      assert(isa<BranchInst>(BB.getTerminator()));
-      return any_of(successors(&BB),
-                    [Succ](const BasicBlock *S) { return S != Succ; }) &&
-             all_of(predecessors(Succ), [&BB, &DT, Succ](BasicBlock *Pred) {
-               return Pred == &BB || DT.dominates(Succ, Pred);
-             });
-    };
     // If the condition is an OR of 2 compares and the false successor only has
     // the current block as predecessor, queue both negated conditions for the
     // false successor.
diff --git a/llvm/test/Transforms/ConstraintElimination/assumes.ll b/llvm/test/Transforms/ConstraintElimination/assumes.ll
index 73a24dd..d51098e 100644
--- a/llvm/test/Transforms/ConstraintElimination/assumes.ll
+++ b/llvm/test/Transforms/ConstraintElimination/assumes.ll
@@ -164,7 +164,7 @@ define i1 @assume_does_not_dominates_successor_with_may_unwind_call_before_assum
 ; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    [[C_2:%.*]] = icmp eq i16 [[A]], 0
-; CHECK-NEXT:    ret i1 true
+; CHECK-NEXT:    ret i1 [[C_2]]
 ;
 entry:
   br i1 %i.0, label %exit, label %if.then
-- 
2.7.4