[InstCombine] allow more narrowing of casted select

[spatel]

Summary

D47163 created a rule that we should not change the casted type of a select when we have matching types in its compare condition. That was intended to help vector codegen, but it also could create situations where we miss subsequent folds as shown in PR44545:
https://bugs.llvm.org/show_bug.cgi?id=44545

By using shouldChangeType(), we can continue to get the vector folds (because we always return false for vector types). But we also solve the motivating bug because it's ok to narrow the scalar select in that example.

Our canonicalization rules around select are a mess, but AFAICT, this will not induce any infinite looping from the reverse transform (but we'll need to watch for that possibility if committed).

Side note: there's a similar use of shouldChangeType() for phi ops just below this diff, and the source and destination types appear to be reversed.

[lebedev.ri]

I'm not sure what kind of fallout this will cause.
For example, won't this affect 'saturating math' peephole?
Won't this affect clamp (min-max) pattern?

[spatel]

In D72733#1820606, @lebedev.ri wrote:

I'm not sure what kind of fallout this will cause.
For example, won't this affect 'saturating math' peephole?
Won't this affect clamp (min-max) pattern?

As with PR44545, it's hard to see the fallout in advance, but keep in mind this patch is only restoring part of the behavior before D47163.
If we are detecting min/max patterns with ValueTracking's matchSelectPattern(), that can optionally include casts, so we should be able to peek through. If you have a pattern in mind that would be affected, let me know.

[bjope]

Thanks for looking into this. I did not know about shouldChangeType but that seems like a useful function here.

Slightly worried just like @lebedev.ri that we might swing back into getting a regression for some other code pattern. But it seems like it at least doesn't hit any regressions for existing lit-tests, so it doesn't regress anything that we have optimized (and added tests for) earlier at least.

Btw, D47163 landed 1.5 years ago, so we have had the regression from PR44545 for awhile. So even if it would be nice to fix it right away, probably no need to rush into something that we do not strongly believe in.

Anyway, the patch LGTM. But let's see if @lebedev.ri has any examples showing that this might a bad idea.

[spatel]

Ping @lebedev.ri - any examples that we can try to fix in advance?

[lebedev.ri]

In D72733#1842707, @spatel wrote:

Ping @lebedev.ri - any examples that we can try to fix in advance?

I don't have anything specific actionable in mind right now, feel free to proceed.

[spatel]

(marking as accepted here in Phab based on earlier LGTM)

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (Instruction::CastOps NewOpc = isEliminableCastPair(CSrc, &CI)) {
       // Point debug users of the dying cast to the new one.
       if (CSrc->hasOneUse())
         replaceAllDbgUsesWith(*CSrc, *Res, CI, DT);
       return Res;
     }
   }
 
   if (auto *Sel = dyn_cast<SelectInst>(Src)) {
-    // We are casting a select. Try to fold the cast into the select, but only
-    // if the select does not have a compare instruction with matching operand
-    // types. Creating a select with operands that are different sizes than its
+    // We are casting a select. Try to fold the cast into the select if the
+    // select does not have a compare instruction with matching operand types
+    // or the select is likely better done in a narrow type.
+    // Creating a select with operands that are different sizes than its
     // condition may inhibit other folds and lead to worse codegen.
     auto *Cmp = dyn_cast<CmpInst>(Sel->getCondition());
-    if (!Cmp || Cmp->getOperand(0)->getType() != Sel->getType())
+    if (!Cmp || Cmp->getOperand(0)->getType() != Sel->getType() ||
+        (CI.getOpcode() == Instruction::Trunc &&
+         shouldChangeType(CI.getSrcTy(), CI.getType()))) {
       if (Instruction *NV = FoldOpIntoSelect(CI, Sel)) {
         replaceAllDbgUsesWith(*Sel, *NV, CI, DT);
         return NV;
       }
     }
+  }
 
   // If we are casting a PHI, then fold the cast into the PHI.
   if (auto *PN = dyn_cast<PHINode>(Src)) {
     // Don't do this if it would create a PHI node with an illegal type from a
     // legal type.
     if (!Src->getType()->isIntegerTy() || !CI.getType()->isIntegerTy() ||
         shouldChangeType(CI.getType(), Src->getType()))
       if (Instruction *NV = foldOpIntoPhi(CI, PN))

llvm/test/Transforms/InstCombine/cast-select.ll

;
   %sel = select <2 x i1> %cmp, <2 x i8> %z, <2 x i8> <i8 42, i8 7>
   %r = sext <2 x i8> %sel to <2 x i32>
   ret <2 x i32> %r
 }
 
 define i16 @trunc(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: @trunc(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32 42, i32 [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = trunc i32 [[SEL]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[Z:%.*]] to i16
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[CMP]], i16 42, i16 [[TMP1]]
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %cmp = icmp ult i32 %x, %y
   %sel = select i1 %cmp, i32 42, i32 %z
   %r = trunc i32 %sel to i16
   ret i16 %r
 }
 

llvm/test/Transforms/InstCombine/select-imm-canon.ll

entry:
   %conv7 = trunc i32 %spec.select.i to i8
   ret i8 %conv7
 }
 
 define i8 @thisdoesnotloop(i32 %A, i32 %B) {
 ; CHECK-LABEL: @thisdoesnotloop(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[L1:%.*]] = icmp slt i32 [[A:%.*]], -128
-; CHECK-NEXT:    [[L2:%.*]] = select i1 [[L1]], i32 -128, i32 [[B:%.*]]
-; CHECK-NEXT:    [[CONV7:%.*]] = trunc i32 [[L2]] to i8
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i32 [[B:%.*]] to i8
+; CHECK-NEXT:    [[CONV7:%.*]] = select i1 [[L1]], i8 -128, i8 [[TMP0]]
 ; CHECK-NEXT:    ret i8 [[CONV7]]
 ;
 entry:
   %l1 = icmp slt i32 %A, -128
   %l2 = select i1 %l1, i32 128, i32 %B
   %conv7 = trunc i32 %l2 to i8
   ret i8 %conv7
 }

llvm/test/Transforms/InstCombine/trunc.ll

 ; CHECK-NEXT:    [[TR:%.*]] = sub <2 x i8> <i8 0, i8 127>, [[TMP1]]
 ; CHECK-NEXT:    ret <2 x i8> [[TR]]
 ;
   %sub = sub <2 x i32> <i32 256, i32 -129>, %x
   %tr = trunc <2 x i32> %sub to <2 x i8>
   ret <2 x i8> %tr
 }
 
-; FIXME: If the select is narrowed based on the target's datalayout, we allow more optimizations.
+; If the select is narrowed based on the target's datalayout, we allow more optimizations.
 
 define i16 @PR44545(i32 %t0, i32 %data) {
 ; CHECK-LABEL: @PR44545(
-; CHECK-NEXT:    [[T1:%.*]] = add nuw nsw i32 [[T0:%.*]], 1
 ; CHECK-NEXT:    [[ISZERO:%.*]] = icmp eq i32 [[DATA:%.*]], 0
-; CHECK-NEXT:    [[FFS:%.*]] = select i1 [[ISZERO]], i32 0, i32 [[T1]]
-; CHECK-NEXT:    [[CAST:%.*]] = trunc i32 [[FFS]] to i16
-; CHECK-NEXT:    [[SUB:%.*]] = add nsw i16 [[CAST]], -1
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[T0:%.*]] to i16
+; CHECK-NEXT:    [[SUB:%.*]] = select i1 [[ISZERO]], i16 -1, i16 [[TMP1]]
 ; CHECK-NEXT:    ret i16 [[SUB]]
 ;
   %t1 = add nuw nsw i32 %t0, 1
   %iszero = icmp eq i32 %data, 0
   %ffs = select i1 %iszero, i32 0, i32 %t1
   %cast = trunc i32 %ffs to i16
   %sub = add nsw i16 %cast, -1
   ret i16 %sub
 }