From bb1dc876ebb8a2eef38d5183d00c2db1437f1c91 Mon Sep 17 00:00:00 2001
From: Max Kazantsev <mkazantsev@azul.com>
Date: Mon, 21 Jun 2021 11:37:06 +0700
Subject: [PATCH] [LoopDeletion] Handle Phis with similar inputs from different
 blocks

This patch lifts the requirement to have the only incoming live block
for Phis. There can be multiple live blocks if the same value comes to
phi from all of them.

Differential Revision: https://reviews.llvm.org/D103959
Reviewed By: nikic, lebedev.ri
---
 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 4d7925e..a287f68 100644
--- a/llvm/lib/Transforms/Scalar/LoopDeletion.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopDeletion.cpp
@@ -245,23 +245,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(L->getLoopPredecessor());
+    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;
 
@@ -290,18 +292,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 612228a..744853e 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