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

[nhaehnle]

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.

[arsenm]

lib/Transforms/InstCombine/InstCombineSelect.cpp

918–921

I would swap the parameter order so that the Builder is first and all bools are at the end

919

Make C a ConstantInt and dyn_cast at the call site?

938–941

No return after else

[nhaehnle]

Implement review suggestions.

[spatel]

test/Transforms/InstCombine/select-bitext.ll

1

Please:

1. Remove all CHECK lines from this file.
2. Auto-generate all CHECK lines using utils/update_test_checks.py against trunk.
3. Commit the test file to trunk to show the current behavior.
4. Apply your patch locally.
5. Regenerate the CHECK lines again using the script.
6. Update this patch.

This will allow us to see the exact diffs caused by this patch.

[nhaehnle]

I have just committed the tests baseline to trunk, this is the updated
patch.

[spatel]

In D22747#504723, @nhaehnle wrote:

I have just committed the tests baseline to trunk, this is the updated
patch.

Thanks!

Sorry I didn't notice this earlier, but can you reduce the first 8 tests by making them have i1 parameters instead of having icmp instructions in the tests?

Is there some reason we shouldn't do this transform for vector selects with splatted constants? If that's possible, it could actually simplify the patch by having foldSelectExtConst() take an APInt rather than a ConstantInt parameter with the caller code looking something like this:

if (TI && (TI->getOpcode() == Instruction::ZExt || TI->getOpcode() == Instruction::SExt)) {
  const APInt *C;
  if (match(FalseVal, m_APInt(C))) {
    bool IsSext = TI->getOpcode() == Instruction::SExt;
    if (auto *Res = foldSelectExtConst(*Builder, SI, TI, C, true, IsSext))
      return Res;
  }
}
if (FI... )

[nhaehnle]

Simplified the test cases (baseline is already in SVN), and use the APInt
hint, thanks!

About the vector case: As you can see, it requires a change to
FoldOpIntoSelect (to avoid an infinite loop) and a change to one other test
case. Now the change looks correct, but I don't know if there may be
unintended optimization regressions in some backend...

Reverting the behavior for vectors is easy enough, though, just remove the
->getScalarType() in foldSelectExtConst and FoldOpIntoSelect.

[spatel]

In D22747#505014, @nhaehnle wrote:

Simplified the test cases (baseline is already in SVN), and use the APInt
hint, thanks!

About the vector case: As you can see, it requires a change to
FoldOpIntoSelect (to avoid an infinite loop) and a change to one other test
case. Now the change looks correct, but I don't know if there may be
unintended optimization regressions in some backend...

Can you add a test case like:

define <2 x i32> @scalar_select_of_vectors(<2 x i1> %cca, i1 %ccb) {
  %ccax = zext <2 x i1> %cca to <2 x i32>
  %r = select i1 %ccb, <2 x i32> %ccax, <2 x i32> <i32 0, i32 0>   ; scalar condition
  ret <2 x i32> %r
}

Also, I'm curious why this wouldn't be good for non-i1 types. Did you see a case where it caused a problem? Particularly in the case of vectors, I think we should be performing ops in smaller types as much as possible before extending to a wider type (ref: https://llvm.org/bugs/show_bug.cgi?id=28160 ). Add a 'TODO' comment about that?

[nhaehnle]

That type of test was actually already there, but I added a copy with zext
instead of sext for good measure.

I added a TODO about handling larger types as well.

[nhaehnle]

To clarify, I didn't see a case where extending to non-i1 types caused a problem (other than the necessary FoldOpIntoSelect change), but I also didn't see cases where it helped, and in any case I was only looking at the AMDGPU assembly results. I expect that the more common CPU targets could be sensitive to a change there, that's why I kept it conservative.

[spatel]

LGTM. Thanks!

lib/Transforms/InstCombine/InstCombineSelect.cpp

if (OtherAddOp) {
         return RI;
       } else
         return BinaryOperator::CreateAdd(SubOp->getOperand(0), NewSel);
     }
   }
   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();
+
+  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();
   Value *FalseVal = SI.getFalseValue();
   Type *SelType = SI.getType();
 
   if (Value *V =
           SimplifySelectInst(CondVal, TrueVal, FalseVal, DL, TLI, DT, AC))
Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
 
   // Turn (select C, (op X, Y), (op X, Z)) -> (op X, (select C, Y, Z))
   auto *TI = dyn_cast<Instruction>(TrueVal);
   auto *FI = dyn_cast<Instruction>(FalseVal);
   if (TI && FI && TI->getOpcode() == FI->getOpcode())
     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))
       return FoldI;
 
     Value *LHS, *RHS, *LHS2, *RHS2;
     Instruction::CastOps CastOp;
     SelectPatternResult SPR = matchSelectPattern(&SI, LHS, RHS, &CastOp);

lib/Transforms/InstCombine/InstructionCombining.cpp

 Instruction *InstCombiner::FoldOpIntoSelect(Instruction &Op, SelectInst *SI) {
   // Don't modify shared select instructions
   if (!SI->hasOneUse()) return nullptr;
   Value *TV = SI->getOperand(1);
   Value *FV = SI->getOperand(2);
 
   if (isa<Constant>(TV) || isa<Constant>(FV)) {
     // Bool selects with constant operands can be folded to logical ops.
-    if (SI->getType()->isIntegerTy(1)) return nullptr;
+    if (SI->getType()->getScalarType()->isIntegerTy(1)) return nullptr;
 
     // If it's a bitcast involving vectors, make sure it has the same number of
     // elements on both sides.
     if (BitCastInst *BC = dyn_cast<BitCastInst>(&Op)) {
       VectorType *DestTy = dyn_cast<VectorType>(BC->getDestTy());
       VectorType *SrcTy = dyn_cast<VectorType>(BC->getSrcTy());
 
       // Verify that either both or neither are vectors.

test/Transforms/InstCombine/select-bitext.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 define i32 @test_sext1(i1 %cca, i1 %ccb) {
 ; CHECK-LABEL: @test_sext1(
-; CHECK-NEXT:    [[CCAX:%.*]] = sext i1 %cca to i32
-; CHECK-NEXT:    [[R:%.*]] = select i1 %ccb, i32 [[CCAX]], i32 0
+; CHECK-NEXT:    [[FOLD_R:%.*]] = and i1 %ccb, %cca
+; CHECK-NEXT:    [[R:%.*]] = sext i1 [[FOLD_R]] to i32
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %ccax = sext i1 %cca to i32
   %r = select i1 %ccb, i32 %ccax, i32 0
   ret i32 %r
 }
 
 define i32 @test_sext2(i1 %cca, i1 %ccb) {
 ; CHECK-LABEL: @test_sext2(
-; CHECK-NEXT:    [[CCAX:%.*]] = sext i1 %cca to i32
-; CHECK-NEXT:    [[R:%.*]] = select i1 %ccb, i32 -1, i32 [[CCAX]]
+; CHECK-NEXT:    [[FOLD_R:%.*]] = or i1 %ccb, %cca
+; CHECK-NEXT:    [[R:%.*]] = sext i1 [[FOLD_R]] to i32
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %ccax = sext i1 %cca to i32
   %r = select i1 %ccb, i32 -1, i32 %ccax
   ret i32 %r
 }
 
 define i32 @test_sext3(i1 %cca, i1 %ccb) {
 ; CHECK-LABEL: @test_sext3(
-; CHECK-NEXT:    [[CCAX:%.*]] = sext i1 %cca to i32
-; CHECK-NEXT:    [[R:%.*]] = select i1 %ccb, i32 0, i32 [[CCAX]]
+; CHECK-NEXT:    [[NOT_CCB:%.*]] = xor i1 %ccb, true
+; CHECK-NEXT:    [[FOLD_R:%.*]] = and i1 [[NOT_CCB]], %cca
+; CHECK-NEXT:    [[R:%.*]] = sext i1 [[FOLD_R]] to i32
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %ccax = sext i1 %cca to i32
   %r = select i1 %ccb, i32 0, i32 %ccax
   ret i32 %r
 }
 
 define i32 @test_sext4(i1 %cca, i1 %ccb) {
 ; CHECK-LABEL: @test_sext4(
-; CHECK-NEXT:    [[CCAX:%.*]] = sext i1 %cca to i32
-; CHECK-NEXT:    [[R:%.*]] = select i1 %ccb, i32 [[CCAX]], i32 -1
+; CHECK-NEXT:    [[NOT_CCB:%.*]] = xor i1 %ccb, true
+; CHECK-NEXT:    [[FOLD_R:%.*]] = or i1 [[NOT_CCB]], %cca
+; CHECK-NEXT:    [[R:%.*]] = sext i1 [[FOLD_R]] to i32
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %ccax = sext i1 %cca to i32
   %r = select i1 %ccb, i32 %ccax, i32 -1
   ret i32 %r
 }
 
 define i32 @test_zext1(i1 %cca, i1 %ccb) {
 ; CHECK-LABEL: @test_zext1(
-; CHECK-NEXT:    [[CCAX:%.*]] = zext i1 %cca to i32
-; CHECK-NEXT:    [[R:%.*]] = select i1 %ccb, i32 [[CCAX]], i32 0
+; CHECK-NEXT:    [[FOLD_R:%.*]] = and i1 %ccb, %cca
+; CHECK-NEXT:    [[R:%.*]] = zext i1 [[FOLD_R]] to i32
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %ccax = zext i1 %cca to i32
   %r = select i1 %ccb, i32 %ccax, i32 0
   ret i32 %r
 }
 
 define i32 @test_zext2(i1 %cca, i1 %ccb) {
 ; CHECK-LABEL: @test_zext2(
-; CHECK-NEXT:    [[CCAX:%.*]] = zext i1 %cca to i32
-; CHECK-NEXT:    [[R:%.*]] = select i1 %ccb, i32 1, i32 [[CCAX]]
+; CHECK-NEXT:    [[FOLD_R:%.*]] = or i1 %ccb, %cca
+; CHECK-NEXT:    [[R:%.*]] = zext i1 [[FOLD_R]] to i32
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %ccax = zext i1 %cca to i32
   %r = select i1 %ccb, i32 1, i32 %ccax
   ret i32 %r
 }
 
 define i32 @test_zext3(i1 %cca, i1 %ccb) {
 ; CHECK-LABEL: @test_zext3(
-; CHECK-NEXT:    [[CCAX:%.*]] = zext i1 %cca to i32
-; CHECK-NEXT:    [[R:%.*]] = select i1 %ccb, i32 0, i32 [[CCAX]]
+; CHECK-NEXT:    [[NOT_CCB:%.*]] = xor i1 %ccb, true
+; CHECK-NEXT:    [[FOLD_R:%.*]] = and i1 [[NOT_CCB]], %cca
+; CHECK-NEXT:    [[R:%.*]] = zext i1 [[FOLD_R]] to i32
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %ccax = zext i1 %cca to i32
   %r = select i1 %ccb, i32 0, i32 %ccax
   ret i32 %r
 }
 
 define i32 @test_zext4(i1 %cca, i1 %ccb) {
 ; CHECK-LABEL: @test_zext4(
-; CHECK-NEXT:    [[CCAX:%.*]] = zext i1 %cca to i32
-; CHECK-NEXT:    [[R:%.*]] = select i1 %ccb, i32 [[CCAX]], i32 1
+; CHECK-NEXT:    [[NOT_CCB:%.*]] = xor i1 %ccb, true
+; CHECK-NEXT:    [[FOLD_R:%.*]] = or i1 [[NOT_CCB]], %cca
+; CHECK-NEXT:    [[R:%.*]] = zext i1 [[FOLD_R]] to i32
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %ccax = zext i1 %cca to i32
   %r = select i1 %ccb, i32 %ccax, i32 1
   ret i32 %r
 }
 
 define i32 @test_negative_sext(i1 %a, i1 %cc) {
;
   %ccbx = sext i1 %ccb to i32
   %ccc = icmp sgt i32 %c, 0
   %r = select i1 %ccc, i32 %ccax, i32 %ccbx
   ret i32 %r
 }
 
 define <2 x i32> @test_vectors1(<2 x i1> %cca, <2 x i1> %ccb) {
 ; CHECK-LABEL: @test_vectors1(
-; CHECK-NEXT:    [[CCAX:%.*]] = sext <2 x i1> %cca to <2 x i32>
-; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> %ccb, <2 x i32> [[CCAX]], <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[FOLD_R:%.*]] = and <2 x i1> %ccb, %cca
+; CHECK-NEXT:    [[R:%.*]] = sext <2 x i1> [[FOLD_R]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %ccax = sext <2 x i1> %cca to <2 x i32>
   %r = select <2 x i1> %ccb, <2 x i32> %ccax, <2 x i32> <i32 0, i32 0>
   ret <2 x i32> %r
 }
 
 define <2 x i32> @test_vectors2(<2 x i1> %cca, i1 %ccb) {
 ; CHECK-LABEL: @test_vectors2(
-; CHECK-NEXT:    [[CCAX:%.*]] = sext <2 x i1> %cca to <2 x i32>
-; CHECK-NEXT:    [[R:%.*]] = select i1 %ccb, <2 x i32> [[CCAX]], <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[FOLD_R:%.*]] = select i1 %ccb, <2 x i1> %cca, <2 x i1> zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = sext <2 x i1> [[FOLD_R]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %ccax = sext <2 x i1> %cca to <2 x i32>
   %r = select i1 %ccb, <2 x i32> %ccax, <2 x i32> <i32 0, i32 0>
   ret <2 x i32> %r
 }

test/Transforms/InstCombine/vector-casts.ll

entry:
   %sext = sext <4 x i1> %cmp to <4 x i32>  
   %cmp4 = fcmp ult <4 x float> %b, zeroinitializer
   %sext5 = sext <4 x i1> %cmp4 to <4 x i32>
   %and = and <4 x i32> %sext, %sext5
   %conv = bitcast <4 x i32> %and to <2 x i64>
   ret <2 x i64> %conv
         
 ; CHECK-LABEL: @test5(
-; CHECK:   sext <4 x i1> %cmp to <4 x i32>  
-; The sext-and pair is canonicalized to a select.
-; CHECK:   select <4 x i1> %cmp4, <4 x i32> %sext, <4 x i32> zeroinitializer
+; CHECK:   %fold.and = and <4 x i1> %cmp4, %cmp
+; CHECK:   sext <4 x i1> %fold.and to <4 x i32>
 }
 
 
 define void @convert(<2 x i32>* %dst.addr, <2 x i64> %src) nounwind {
 entry:
   %val = trunc <2 x i64> %src to <2 x i32>
   %add = add <2 x i32> %val, <i32 1, i32 1>
   store <2 x i32> %add, <2 x i32>* %dst.addr