From a1f442b2787ef63ba84c8627987cfebe3ba4f7f6 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Tue, 1 Mar 2022 15:46:25 +0100
Subject: [PATCH] [InstCombine] Support phi to cond fold with more than two
 preds

This transform can still be applied if there are more than two
phi inputs, as long as phi inputs with the same value are dominated
by the same idom edge.
---
 .../Transforms/InstCombine/InstCombinePHI.cpp | 66 +++++++++----------
 .../InstCombine/simple_phi_condition.ll       |  3 +-
 2 files changed, 33 insertions(+), 36 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp b/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp
index 13277e872759..f1a912822ff0 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp
@@ -1264,9 +1264,6 @@ static Value *simplifyUsingControlFlow(InstCombiner &Self, PHINode &PN,
   if (!PN.getType()->isIntegerTy(1))
     return nullptr;
 
-  if (PN.getNumOperands() != 2)
-    return nullptr;
-
   // Make sure all inputs are constants.
   if (!all_of(PN.operands(), [](Value *V) { return isa<ConstantInt>(V); }))
     return nullptr;
@@ -1276,22 +1273,7 @@ static Value *simplifyUsingControlFlow(InstCombiner &Self, PHINode &PN,
   if (!DT.isReachableFromEntry(BB))
     return nullptr;
 
-  // Same inputs.
-  if (PN.getOperand(0) == PN.getOperand(1))
-    return PN.getOperand(0);
-
-  BasicBlock *TruePred = nullptr, *FalsePred = nullptr;
-  for (auto *Pred : predecessors(BB)) {
-    auto *Input = cast<ConstantInt>(PN.getIncomingValueForBlock(Pred));
-    if (Input->isAllOnesValue())
-      TruePred = Pred;
-    else
-      FalsePred = Pred;
-  }
-  assert(TruePred && FalsePred && "Must be!");
-
-  // Check which edge of the dominator dominates the true input. If it is the
-  // false edge, we should invert the condition.
+  // Check that the immediate dominator has a conditional branch.
   auto *IDom = DT.getNode(BB)->getIDom()->getBlock();
   auto *BI = dyn_cast<BranchInst>(IDom->getTerminator());
   if (!BI || BI->isUnconditional())
@@ -1302,24 +1284,40 @@ static Value *simplifyUsingControlFlow(InstCombiner &Self, PHINode &PN,
   BasicBlockEdge TrueOutEdge(IDom, BI->getSuccessor(0));
   BasicBlockEdge FalseOutEdge(IDom, BI->getSuccessor(1));
 
-  BasicBlockEdge TrueIncEdge(TruePred, BB);
-  BasicBlockEdge FalseIncEdge(FalsePred, BB);
+  Optional<bool> Invert;
+  for (auto Pair : zip(PN.incoming_values(), PN.blocks())) {
+    auto *Input = cast<ConstantInt>(std::get<0>(Pair));
+    BasicBlock *Pred = std::get<1>(Pair);
+    BasicBlockEdge Edge(Pred, BB);
+
+    // The input needs to be dominated by one of the edges of the idom.
+    // Depending on the constant, the condition may need to be inverted.
+    bool NeedsInvert;
+    if (DT.dominates(TrueOutEdge, Edge))
+      NeedsInvert = Input->isZero();
+    else if (DT.dominates(FalseOutEdge, Edge))
+      NeedsInvert = Input->isOne();
+    else
+      return nullptr;
+
+    // Make sure the inversion requirement is always the same.
+    if (Invert && *Invert != NeedsInvert)
+      return nullptr;
+
+    Invert = NeedsInvert;
+  }
 
   auto *Cond = BI->getCondition();
-  if (DT.dominates(TrueOutEdge, TrueIncEdge) &&
-      DT.dominates(FalseOutEdge, FalseIncEdge))
-    // This Phi is actually equivalent to branching condition of IDom.
+  if (!*Invert)
     return Cond;
-  if (DT.dominates(TrueOutEdge, FalseIncEdge) &&
-      DT.dominates(FalseOutEdge, TrueIncEdge)) {
-    // This Phi is actually opposite to branching condition of IDom. We invert
-    // the condition that will potentially open up some opportunities for
-    // sinking.
-    auto InsertPt = BB->getFirstInsertionPt();
-    if (InsertPt != BB->end()) {
-      Self.Builder.SetInsertPoint(&*InsertPt);
-      return Self.Builder.CreateNot(Cond);
-    }
+
+  // This Phi is actually opposite to branching condition of IDom. We invert
+  // the condition that will potentially open up some opportunities for
+  // sinking.
+  auto InsertPt = BB->getFirstInsertionPt();
+  if (InsertPt != BB->end()) {
+    Self.Builder.SetInsertPoint(&*InsertPt);
+    return Self.Builder.CreateNot(Cond);
   }
 
   return nullptr;
diff --git a/llvm/test/Transforms/InstCombine/simple_phi_condition.ll b/llvm/test/Transforms/InstCombine/simple_phi_condition.ll
index 27e2a0142f97..87c864005784 100644
--- a/llvm/test/Transforms/InstCombine/simple_phi_condition.ll
+++ b/llvm/test/Transforms/InstCombine/simple_phi_condition.ll
@@ -168,8 +168,7 @@ define i1 @test_multiple_predecessors(i1 %cond, i1 %cond2) {
 ; CHECK:       if2.false:
 ; CHECK-NEXT:    br label [[MERGE]]
 ; CHECK:       merge:
-; CHECK-NEXT:    [[RET:%.*]] = phi i1 [ true, [[IF_TRUE]] ], [ false, [[IF2_TRUE]] ], [ false, [[IF2_FALSE]] ]
-; CHECK-NEXT:    ret i1 [[RET]]
+; CHECK-NEXT:    ret i1 [[COND]]
 ;
 entry:
   br i1 %cond, label %if.true, label %if.false