2 files changed, 39 insertions, 18 deletions
diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index 9fc10af9668..ddebcfaad06 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -3924,6 +3924,42 @@ static Value *simplifySelectWithFCmp(Value *Cond, Value *T, Value *F) {
   return nullptr;
 }
 
+/// Try to determine the result of a select based on a dominating condition.
+static Value *foldSelectWithDominatingCond(Value *Cond, Value *TV, Value *FV,
+                                           const SimplifyQuery &Q) {
+  // First, make sure that we have a select in a basic block.
+  // We don't know if we are called from some incomplete state.
+  if (!Q.CxtI || !Q.CxtI->getParent())
+    return nullptr;
+
+  // TODO: This is a poor/cheap way to determine dominance. Should we use the
+  // dominator tree in the SimplifyQuery instead?
+  const BasicBlock *SelectBB = Q.CxtI->getParent();
+  const BasicBlock *PredBB = SelectBB->getSinglePredecessor();
+  if (!PredBB)
+    return nullptr;
+
+  // We need a conditional branch in the predecessor.
+  Value *PredCond;
+  BasicBlock *TrueBB, *FalseBB;
+  if (!match(PredBB->getTerminator(), m_Br(m_Value(PredCond), TrueBB, FalseBB)))
+    return nullptr;
+
+  // The branch should get simplified. Don't bother simplifying the select.
+  if (TrueBB == FalseBB)
+    return nullptr;
+
+  assert((TrueBB == SelectBB || FalseBB == SelectBB) &&
+         "Predecessor block does not point to successor?");
+
+  // Is the select condition implied by the predecessor condition?
+  bool CondIsTrue = TrueBB == SelectBB;
+  Optional<bool> Implied = isImpliedCondition(PredCond, Cond, Q.DL, CondIsTrue);
+  if (!Implied)
+    return nullptr;
+  return *Implied ? TV : FV;
+}
+
 /// Given operands for a SelectInst, see if we can fold the result.
 /// If not, this returns null.
 static Value *SimplifySelectInst(Value *Cond, Value *TrueVal, Value *FalseVal,
@@ -3966,6 +4002,9 @@ static Value *SimplifySelectInst(Value *Cond, Value *TrueVal, Value *FalseVal,
   if (Value *V = foldSelectWithBinaryOp(Cond, TrueVal, FalseVal))
     return V;
 
+  if (Value *V = foldSelectWithDominatingCond(Cond, TrueVal, FalseVal, Q))
+    return V;
+
   return nullptr;
 }
 
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index db75c9383a1..19858ae149a 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -2021,24 +2021,6 @@ Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
     }
   }
 
-  // See if we can determine the result of this select based on a dominating
-  // condition.
-  BasicBlock *Parent = SI.getParent();
-  if (BasicBlock *Dom = Parent->getSinglePredecessor()) {
-    auto *PBI = dyn_cast_or_null<BranchInst>(Dom->getTerminator());
-    if (PBI && PBI->isConditional() &&
-        PBI->getSuccessor(0) != PBI->getSuccessor(1) &&
-        (PBI->getSuccessor(0) == Parent || PBI->getSuccessor(1) == Parent)) {
-      bool CondIsTrue = PBI->getSuccessor(0) == Parent;
-      Optional<bool> Implication = isImpliedCondition(
-          PBI->getCondition(), SI.getCondition(), DL, CondIsTrue);
-      if (Implication) {
-        Value *V = *Implication ? TrueVal : FalseVal;
-        return replaceInstUsesWith(SI, V);
-      }
-    }
-  }
-
   // If we can compute the condition, there's no need for a select.
   // Like the above fold, we are attempting to reduce compile-time cost by
   // putting this fold here with limitations rather than in InstSimplify.