[InstCombine] don't change the size of a select if it would mismatch its condition operands' sizes

[spatel]

Summary

Don't always:
cast (select (cmp x, y), z, C) --> select (cmp x, y), (cast z), C'

This is something that came up as far back as D26556, and I lost track of it. I suspect that this transform is part of the underlying problem that is inspiring many of the recent proposals that seek to match larger patterns that include a cast op.

If this looks ok, then I can also add the canonicalization in the other direction, so we actively try to get the compare and select to match sizes.

[lebedev.ri]

Looks sane to me, but please do wait for someone else to review this too.

lib/Transforms/InstCombine/InstCombineCasts.cpp

289–290 ↗

(On Diff #147872)

Hm, indeed.
https://godbolt.org/g/9zMbm8

[spatel]

@opaparo can you review this patch? I'm wondering if this will help any of the motivating cases for D46380.

[opaparo]

In D47163#1108032, @spatel wrote:

@opaparo can you review this patch? I'm wondering if this will help any of the motivating cases for D46380.

It's actually not very helpful for the motivating cases for D46380.
Regardless, the proposed functionality seems optimization-wise correct.

llvm/trunk/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (Instruction::CastOps NewOpc = isEliminableCastPair(CSrc, &CI)) {
           DIB.insertDbgValueIntrinsic(
               Res, DII->getVariable(), DII->getExpression(),
               DII->getDebugLoc().get(), &*std::next(CI.getIterator()));
       }
       return Res;
     }
   }
 
-  // If we are casting a select, then fold the cast into the select.
-  if (auto *SI = dyn_cast<SelectInst>(Src))
-    if (Instruction *NV = FoldOpIntoSelect(CI, SI))
+  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
+    // 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 (Instruction *NV = FoldOpIntoSelect(CI, Sel))
         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/trunk/test/Transforms/InstCombine/cast-select.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 define i64 @zext(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: @zext(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[Z:%.*]] to i64
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[CMP]], i64 0, i64 [[TMP1]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32 0, i32 [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = zext i32 [[SEL]] to i64
 ; CHECK-NEXT:    ret i64 [[R]]
 ;
   %cmp = icmp eq i32 %x, %y
   %sel = select i1 %cmp, i32 0, i32 %z
   %r = zext i32 %sel to i64
   ret i64 %r
 }
 
 define <2 x i32> @zext_vec(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
 ; CHECK-LABEL: @zext_vec(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt <2 x i8> [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = zext <2 x i8> [[Z:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[CMP]], <2 x i32> [[TMP1]], <2 x i32> <i32 42, i32 7>
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x i8> [[Z:%.*]], <2 x i8> <i8 42, i8 7>
+; CHECK-NEXT:    [[R:%.*]] = zext <2 x i8> [[SEL]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %cmp = icmp ugt <2 x i8> %x, %y
   %sel = select <2 x i1> %cmp, <2 x i8> %z, <2 x i8> <i8 42, i8 7>
   %r = zext <2 x i8> %sel to <2 x i32>
   ret <2 x i32> %r
 }
 
 define i64 @sext(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @sext(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = sext i8 [[Z:%.*]] to i64
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[CMP]], i64 42, i64 [[TMP1]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i8 42, i8 [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = sext i8 [[SEL]] to i64
 ; CHECK-NEXT:    ret i64 [[R]]
 ;
   %cmp = icmp ult i8 %x, %y
   %sel = select i1 %cmp, i8 42, i8 %z
   %r = sext i8 %sel to i64
   ret i64 %r
 }
 
 define <2 x i32> @sext_vec(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
 ; CHECK-LABEL: @sext_vec(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt <2 x i8> [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = sext <2 x i8> [[Z:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[CMP]], <2 x i32> [[TMP1]], <2 x i32> <i32 42, i32 7>
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x i8> [[Z:%.*]], <2 x i8> <i8 42, i8 7>
+; CHECK-NEXT:    [[R:%.*]] = sext <2 x i8> [[SEL]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %cmp = icmp ugt <2 x i8> %x, %y
   %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:    [[TMP1:%.*]] = trunc i32 [[Z:%.*]] to i16
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[CMP]], i16 42, i16 [[TMP1]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32 42, i32 [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = trunc i32 [[SEL]] to i16
 ; 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
 }
 
 define <2 x i32> @trunc_vec(<2 x i64> %x, <2 x i64> %y, <2 x i64> %z) {
 ; CHECK-LABEL: @trunc_vec(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt <2 x i64> [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc <2 x i64> [[Z:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[CMP]], <2 x i32> [[TMP1]], <2 x i32> <i32 42, i32 7>
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x i64> [[Z:%.*]], <2 x i64> <i64 42, i64 7>
+; CHECK-NEXT:    [[R:%.*]] = trunc <2 x i64> [[SEL]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %cmp = icmp ugt <2 x i64> %x, %y
   %sel = select <2 x i1> %cmp, <2 x i64> %z, <2 x i64> <i64 42, i64 7>
   %r = trunc <2 x i64> %sel to <2 x i32>
   ret <2 x i32> %r
 }
 
 define double @fpext(float %x, float %y, float %z) {
 ; CHECK-LABEL: @fpext(
 ; CHECK-NEXT:    [[CMP:%.*]] = fcmp oeq float [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = fpext float [[Z:%.*]] to double
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[CMP]], double 1.700000e+01, double [[TMP1]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], float 1.700000e+01, float [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = fpext float [[SEL]] to double
 ; CHECK-NEXT:    ret double [[R]]
 ;
   %cmp = fcmp oeq float %x, %y
   %sel = select i1 %cmp, float 17.0, float %z
   %r = fpext float %sel to double
   ret double %r
 }
 
 define <2 x double> @fpext_vec(<2 x float> %x, <2 x float> %y, <2 x float> %z) {
 ; CHECK-LABEL: @fpext_vec(
 ; CHECK-NEXT:    [[CMP:%.*]] = fcmp ugt <2 x float> [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = fpext <2 x float> [[Z:%.*]] to <2 x double>
-; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[CMP]], <2 x double> [[TMP1]], <2 x double> <double 4.200000e+01, double -2.000000e+00>
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x float> [[Z:%.*]], <2 x float> <float 4.200000e+01, float -2.000000e+00>
+; CHECK-NEXT:    [[R:%.*]] = fpext <2 x float> [[SEL]] to <2 x double>
 ; CHECK-NEXT:    ret <2 x double> [[R]]
 ;
   %cmp = fcmp ugt <2 x float> %x, %y
   %sel = select <2 x i1> %cmp, <2 x float> %z, <2 x float> <float 42.0, float -2.0>
   %r = fpext <2 x float> %sel to <2 x double>
   ret <2 x double> %r
 }
 
 define float @fptrunc(double %x, double %y, double %z) {
 ; CHECK-LABEL: @fptrunc(
 ; CHECK-NEXT:    [[CMP:%.*]] = fcmp ult double [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = fptrunc double [[Z:%.*]] to float
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[CMP]], float 4.200000e+01, float [[TMP1]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], double 4.200000e+01, double [[Z:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = fptrunc double [[SEL]] to float
 ; CHECK-NEXT:    ret float [[R]]
 ;
   %cmp = fcmp ult double %x, %y
   %sel = select i1 %cmp, double 42.0, double %z
   %r = fptrunc double %sel to float
   ret float %r
 }
 
 define <2 x float> @fptrunc_vec(<2 x double> %x, <2 x double> %y, <2 x double> %z) {
 ; CHECK-LABEL: @fptrunc_vec(
 ; CHECK-NEXT:    [[CMP:%.*]] = fcmp oge <2 x double> [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = fptrunc <2 x double> [[Z:%.*]] to <2 x float>
-; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[CMP]], <2 x float> [[TMP1]], <2 x float> <float -4.200000e+01, float 1.200000e+01>
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[CMP]], <2 x double> [[Z:%.*]], <2 x double> <double -4.200000e+01, double 1.200000e+01>
+; CHECK-NEXT:    [[R:%.*]] = fptrunc <2 x double> [[SEL]] to <2 x float>
 ; CHECK-NEXT:    ret <2 x float> [[R]]
 ;
   %cmp = fcmp oge <2 x double> %x, %y
   %sel = select <2 x i1> %cmp, <2 x double> %z, <2 x double> <double -42.0, double 12.0>
   %r = fptrunc <2 x double> %sel to <2 x float>
   ret <2 x float> %r
 }
 

llvm/trunk/test/Transforms/InstCombine/fpcast.ll

;
   %b = fsub fast float -0.0, %a
   %c = fptrunc float %b to half
   ret half %c
 }
 
 define half @test5(float %a, float %b, float %c) {
 ; CHECK-LABEL: @test5(
 ; CHECK-NEXT:    [[D:%.*]] = fcmp ogt float [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = fptrunc float [[C:%.*]] to half
-; CHECK-NEXT:    [[F:%.*]] = select i1 [[D]], half [[TMP1]], half 0xH3C00
+; CHECK-NEXT:    [[E:%.*]] = select i1 [[D]], float [[C:%.*]], float 1.000000e+00
+; CHECK-NEXT:    [[F:%.*]] = fptrunc float [[E]] to half
 ; CHECK-NEXT:    ret half [[F]]
 ;
   %d = fcmp ogt float %a, %b
   %e = select i1 %d, float %c, float 1.0
   %f = fptrunc float %e to half
   ret half %f
 }