[llvm] a1f442b - [InstCombine] Support phi to cond fold with more than two preds

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2022-03-01T16:31:49+01:00
New Revision: a1f442b2787ef63ba84c8627987cfebe3ba4f7f6

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

LOG: [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.

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp
    llvm/test/Transforms/InstCombine/simple_phi_condition.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp b/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp
index 13277e8727597..f1a912822ff06 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 27e2a0142f974..87c864005784a 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