From 4c4f1ae93ea7477ccb4772007fc78313f5a0644f Mon Sep 17 00:00:00 2001
From: Max Kazantsev <mkazantsev@azul.com>
Date: Tue, 22 Jun 2021 12:21:38 +0700
Subject: [PATCH] Re-land "[LoopDeletion] Handle Phis with similar inputs from
 different blocks"

Patch was reverted due to a bug that existed before it and was exposed
by it. Returning after the underlying bug has been fixed.

Differential Revision: https://reviews.llvm.org/D103959
---
 llvm/lib/Transforms/Scalar/LoopDeletion.cpp        | 47 +++++++++++-----------
 .../LoopDeletion/eval_first_iteration.ll           | 11 ++---
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/LoopDeletion.cpp b/llvm/lib/Transforms/Scalar/LoopDeletion.cpp
index 33572e3..2dd36ae 100644
--- a/llvm/lib/Transforms/Scalar/LoopDeletion.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopDeletion.cpp
@@ -246,23 +246,25 @@ static bool canProveExitOnFirstIteration(Loop *L, DominatorTree &DT,
       MarkLiveEdge(BB, Succ);
   };
 
-  // Check if there is only one predecessor on 1st iteration. Note that because
-  // we iterate in RPOT, we have already visited all its (non-latch)
-  // predecessors.
-  auto GetSolePredecessorOnFirstIteration = [&](BasicBlock * BB)->BasicBlock * {
+  // Check if there is only one value coming from all live predecessor blocks.
+  // Note that because we iterate in RPOT, we have already visited all its
+  // (non-latch) predecessors.
+  auto GetSoleInputOnFirstIteration = [&](PHINode & PN)->Value * {
+    BasicBlock *BB = PN.getParent();
     if (BB == Header)
-      return L->getLoopPredecessor();
-    BasicBlock *OnlyPred = nullptr;
+      return PN.getIncomingValueForBlock(Predecessor);
+    Value *OnlyInput = nullptr;
     for (auto *Pred : predecessors(BB))
-      if (OnlyPred != Pred && LiveEdges.count({ Pred, BB })) {
-        // 2 live preds.
-        if (OnlyPred)
+      if (LiveEdges.count({ Pred, BB })) {
+        Value *Incoming = PN.getIncomingValueForBlock(Pred);
+        // Two inputs.
+        if (OnlyInput && OnlyInput != Incoming)
           return nullptr;
-        OnlyPred = Pred;
+        OnlyInput = Incoming;
       }
 
-    assert(OnlyPred && "No live predecessors?");
-    return OnlyPred;
+    assert(OnlyInput && "No live predecessors?");
+    return OnlyInput;
   };
   DenseMap<Value *, Value *> FirstIterValue;
 
@@ -291,18 +293,17 @@ static bool canProveExitOnFirstIteration(Loop *L, DominatorTree &DT,
       continue;
     }
 
-    // If this block has only one live pred, map its phis onto their SCEVs.
-    if (auto *OnlyPred = GetSolePredecessorOnFirstIteration(BB))
-      for (auto &PN : BB->phis()) {
-        if (!PN.getType()->isIntegerTy())
-          continue;
-        auto *Incoming = PN.getIncomingValueForBlock(OnlyPred);
-        if (DT.dominates(Incoming, BB->getTerminator())) {
-          Value *FirstIterV =
-              getValueOnFirstIteration(Incoming, FirstIterValue, SQ);
-          FirstIterValue[&PN] = FirstIterV;
-        }
+    // If Phi has only one input from all live input blocks, use it.
+    for (auto &PN : BB->phis()) {
+      if (!PN.getType()->isIntegerTy())
+        continue;
+      auto *Incoming = GetSoleInputOnFirstIteration(PN);
+      if (Incoming && DT.dominates(Incoming, BB->getTerminator())) {
+        Value *FirstIterV =
+            getValueOnFirstIteration(Incoming, FirstIterValue, SQ);
+        FirstIterValue[&PN] = FirstIterV;
       }
+    }
 
     using namespace PatternMatch;
     ICmpInst::Predicate Pred;
diff --git a/llvm/test/Transforms/LoopDeletion/eval_first_iteration.ll b/llvm/test/Transforms/LoopDeletion/eval_first_iteration.ll
index f2543e2..15ebd7d 100644
--- a/llvm/test/Transforms/LoopDeletion/eval_first_iteration.ll
+++ b/llvm/test/Transforms/LoopDeletion/eval_first_iteration.ll
@@ -655,20 +655,19 @@ failure:
   unreachable
 }
 
-; TODO: We can break the backedge here.
 define i32 @test_multiple_pred_2() {
 ; CHECK-LABEL: @test_multiple_pred_2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[SUM_NEXT:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[SUB:%.*]] = sub i32 4, [[SUM]]
 ; CHECK-NEXT:    [[IS_POSITIVE:%.*]] = icmp sgt i32 [[SUB]], 0
 ; CHECK-NEXT:    br i1 [[IS_POSITIVE]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
 ; CHECK:       if.true:
 ; CHECK-NEXT:    br i1 undef, label [[IF_TRUE_1:%.*]], label [[IF_TRUE_2:%.*]]
 ; CHECK:       if.true.1:
-; CHECK-NEXT:    br label [[BACKEDGE]]
+; CHECK-NEXT:    br label [[BACKEDGE:%.*]]
 ; CHECK:       if.true.2:
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       if.false:
@@ -679,9 +678,11 @@ define i32 @test_multiple_pred_2() {
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[MERGE_PHI:%.*]] = phi i32 [ 0, [[IF_FALSE_1]] ], [ 0, [[IF_FALSE_2]] ], [ [[SUB]], [[IF_TRUE_1]] ], [ [[SUB]], [[IF_TRUE_2]] ]
-; CHECK-NEXT:    [[SUM_NEXT]] = add i32 [[SUM]], [[MERGE_PHI]]
+; CHECK-NEXT:    [[SUM_NEXT:%.*]] = add i32 [[SUM]], [[MERGE_PHI]]
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp ne i32 [[SUM_NEXT]], 4
-; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[DONE:%.*]]
+; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[BACKEDGE_LOOP_CRIT_EDGE:%.*]], label [[DONE:%.*]]
+; CHECK:       backedge.loop_crit_edge:
+; CHECK-NEXT:    unreachable
 ; CHECK:       done:
 ; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT]], [[BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
-- 
2.7.4