[InstCombine] allow non-splat folds of select cond (ext X), C

llvm-svn: 282906

2 files changed, 34 insertions, 40 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 7fb0659a79e..dca24d2eb91 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -590,8 +590,7 @@ private:
   Instruction *foldICmpIntrinsicWithConstant(ICmpInst &ICI, const APInt *C);
 
   // Helpers of visitSelectInst().
-  Instruction *foldSelectExtConst(SelectInst &Sel, Instruction *ExtInst,
-                                  const APInt &C);
+  Instruction *foldSelectExtConst(SelectInst &Sel);
   Instruction *foldSelectOpOp(SelectInst &SI, Instruction *TI, Instruction *FI);
   Instruction *foldSelectIntoOp(SelectInst &SI, Value *, Value *);
   Instruction *foldSPFofSPF(Instruction *Inner, SelectPatternFlavor SPF1,
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index c62598d4f2d..b8213426a63 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -916,32 +916,42 @@ static Instruction *foldAddSubSelect(SelectInst &SI,
   return nullptr;
 }
 
-/// If one of the operands is a sext/zext from i1 and the other is a constant,
-/// we may be able to create an i1 select which can be further folded to
-/// logical ops.
-Instruction *InstCombiner::foldSelectExtConst(SelectInst &Sel,
-                                              Instruction *ExtInst,
-                                              const APInt &C) {
+Instruction *InstCombiner::foldSelectExtConst(SelectInst &Sel) {
+  Instruction *ExtInst;
+  if (!match(Sel.getTrueValue(), m_Instruction(ExtInst)) &&
+      !match(Sel.getFalseValue(), m_Instruction(ExtInst)))
+    return nullptr;
+
+  auto ExtOpcode = ExtInst->getOpcode();
+  if (ExtOpcode != Instruction::ZExt && ExtOpcode != Instruction::SExt)
+    return nullptr;
+
   // TODO: Handle larger types? That requires adjusting FoldOpIntoSelect too.
-  Value *SmallVal = ExtInst->getOperand(0);
-  Type *SmallType = SmallVal->getType();
+  Value *X = ExtInst->getOperand(0);
+  Type *SmallType = X->getType();
   if (!SmallType->getScalarType()->isIntegerTy(1))
     return nullptr;
 
-  Value *Cond = Sel.getCondition();
-  bool IsExtTrueVal = Sel.getTrueValue() == ExtInst;
-  bool IsSext = ExtInst->getOpcode() == Instruction::SExt;
-  if (C == 0 || (!IsSext && C == 1) || (IsSext && C.isAllOnesValue())) {
-    Value *SmallConst = ConstantInt::get(SmallType, C.trunc(1));
-    Value *TrueVal = IsExtTrueVal ? SmallVal : SmallConst;
-    Value *FalseVal = IsExtTrueVal ? SmallConst : SmallVal;
-    Value *NewSel = Builder->CreateSelect(Cond, TrueVal, FalseVal,
-                                          "fold." + Sel.getName(), &Sel);
-
-    if (IsSext)
-      return new SExtInst(NewSel, Sel.getType());
+  Constant *C;
+  if (!match(Sel.getTrueValue(), m_Constant(C)) &&
+      !match(Sel.getFalseValue(), m_Constant(C)))
+    return nullptr;
 
-    return new ZExtInst(NewSel, Sel.getType());
+  // If the constant is the same after truncation to the smaller type and
+  // extension to the original type, we can narrow the select.
+  Type *SelType = Sel.getType();
+  Constant *TruncC = ConstantExpr::getTrunc(C, SmallType);
+  Constant *ExtC = ConstantExpr::getCast(ExtOpcode, TruncC, SelType);
+  if (ExtC == C) {
+    Value *Cond = Sel.getCondition();
+    Value *TruncCVal = cast<Value>(TruncC);
+    if (ExtInst == Sel.getFalseValue())
+      std::swap(X, TruncCVal);
+
+    // select Cond, (ext X), C --> ext(select Cond, X, C')
+    // select Cond, C, (ext X) --> ext(select Cond, C', X)
+    Value *NewSel = Builder->CreateSelect(Cond, X, TruncCVal, "narrow", &Sel);
+    return CastInst::Create(Instruction::CastOps(ExtOpcode), NewSel, SelType);
   }
 
   return nullptr;
@@ -1172,23 +1182,8 @@ Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
     if (Instruction *IV = foldSelectOpOp(SI, TI, FI))
       return IV;
 
-  // (select C, (ext X), const) -> (ext (select C, X, const')) and variations
-  // thereof when extending from i1, as that allows further folding into logic
-  // ops. When the sext is from a larger type, prefer to have it as an operand.
-  if (TI && (TI->getOpcode() == Instruction::ZExt ||
-             TI->getOpcode() == Instruction::SExt)) {
-    const APInt *C;
-    if (match(FalseVal, m_APInt(C)))
-      if (auto *I = foldSelectExtConst(SI, TI, *C))
-        return I;
-  }
-  if (FI && (FI->getOpcode() == Instruction::ZExt ||
-             FI->getOpcode() == Instruction::SExt)) {
-    const APInt *C;
-    if (match(TrueVal, m_APInt(C)))
-      if (auto *I = foldSelectExtConst(SI, FI, *C))
-        return I;
-  }
+  if (Instruction *I = foldSelectExtConst(SI))
+    return I;
 
   // See if we can fold the select into one of our operands.
   if (SelType->isIntOrIntVectorTy() || SelType->isFPOrFPVectorTy()) {