[InstCombine] try to fold (select C, (sext A), B) into logical ops

Summary:
Turn (select C, (sext A), B) into (sext (select C, A, B')) when A is i1 and
B is a compatible constant, also for zext instead of sext. This will then be
further folded into logical operations.

The transformation would be valid for non-i1 types as well, but other parts of
InstCombine prefer to have sext from non-i1 as an operand of select.

Motivated by the shader compiler frontend in Mesa for AMDGPU, which emits i32
for boolean operations. With this change, the boolean logic is fully
recovered.

Reviewers: majnemer, spatel, tstellarAMD

Subscribers: llvm-commits

Differential Revision: https://reviews.llvm.org/D22747

llvm-svn: 277801

1 files changed, 56 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 018cdcc24d5..b76241a70a9 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -912,6 +912,37 @@ 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.
+static Instruction *foldSelectExtConst(InstCombiner::BuilderTy &Builder,
+                                       SelectInst &SI, Instruction *EI,
+                                       const APInt &C, bool isExtTrueVal,
+                                       bool isSigned) {
+  Value *SmallVal = EI->getOperand(0);
+  Type *SmallType = SmallVal->getType();
+
+  // TODO Handle larger types as well? Note this requires adjusting
+  // FoldOpIntoSelect as well.
+  if (!SmallType->getScalarType()->isIntegerTy(1))
+    return nullptr;
+
+  if (C != 0 && (isSigned || C != 1) &&
+      (!isSigned || !C.isAllOnesValue()))
+    return nullptr;
+
+  Value *SmallConst = ConstantInt::get(SmallType, C.trunc(1));
+  Value *TrueVal = isExtTrueVal ? SmallVal : SmallConst;
+  Value *FalseVal = isExtTrueVal ? SmallConst : SmallVal;
+  Value *Select = Builder.CreateSelect(SI.getOperand(0), TrueVal, FalseVal,
+                                       "fold." + SI.getName());
+
+  if (isSigned)
+    return new SExtInst(Select, SI.getType());
+
+  return new ZExtInst(Select, SI.getType());
+}
+
 Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
@@ -1098,6 +1129,31 @@ Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
     if (Instruction *IV = FoldSelectOpOp(SI, TI, FI))
       return IV;
 
+  // (select C, (sext X), const) -> (sext (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, we prefer to have it
+  // as an operand.
+  if (TI &&
+      (TI->getOpcode() == Instruction::ZExt || TI->getOpcode() == Instruction::SExt)) {
+    bool IsSExt = TI->getOpcode() == Instruction::SExt;
+    const APInt *C;
+    if (match(FalseVal, m_APInt(C))) {
+      if (Instruction *IV =
+              foldSelectExtConst(*Builder, SI, TI, *C, true, IsSExt))
+        return IV;
+    }
+  }
+  if (FI &&
+      (FI->getOpcode() == Instruction::ZExt || FI->getOpcode() == Instruction::SExt)) {
+    bool IsSExt = FI->getOpcode() == Instruction::SExt;
+    const APInt *C;
+    if (match(TrueVal, m_APInt(C))) {
+      if (Instruction *IV =
+              foldSelectExtConst(*Builder, SI, FI, *C, false, IsSExt))
+          return IV;
+    }
+  }
+
   // See if we can fold the select into one of our operands.
   if (SelType->isIntOrIntVectorTy() || SelType->isFPOrFPVectorTy()) {
     if (Instruction *FoldI = FoldSelectIntoOp(SI, TrueVal, FalseVal))