[InstCombine] Teach foldSelectICmpAndOr to handle vector splats

[craig.topper]

Summary

This was pretty close to working already. While I was here I went ahead and passed the ICmpInst pointer from the caller instead of doing a dyn_cast that can never fail.

[spatel]

I didn't test it, but I don't see a hasOneUse() check on the cmp, so this might be creating more instructions than we expect. That would be a follow-up change though, so LGTM. See inline for a small potential improvement.

lib/Transforms/InstCombine/InstCombineSelect.cpp

322 ↗

(On Diff #112980)

Could use IC->getType() and make this a compound predicate in the first line?

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

 ///   C1 and C2 are both powers of 2
 /// where:
 ///   C3 = Log(C2) - Log(C1)
 ///
 /// This transform handles cases where:
 /// 1. The icmp predicate is inverted
 /// 2. The select operands are reversed
 /// 3. The magnitude of C2 and C1 are flipped
-static Value *foldSelectICmpAndOr(const SelectInst &SI, Value *TrueVal,
+static Value *foldSelectICmpAndOr(const ICmpInst *IC, Value *TrueVal,
                                   Value *FalseVal,
                                   InstCombiner::BuilderTy &Builder) {
-  const ICmpInst *IC = dyn_cast<ICmpInst>(SI.getCondition());
-  if (!IC || !SI.getType()->isIntegerTy())
+  // Only handle integer compares. Also, if this is a vector select, we need a
+  // vector compare.
+  if (!TrueVal->getType()->isIntOrIntVectorTy() ||
+      TrueVal->getType()->isVectorTy() != IC->getType()->isVectorTy())
     return nullptr;
 
   Value *CmpLHS = IC->getOperand(0);
   Value *CmpRHS = IC->getOperand(1);
 
   Value *V;
   unsigned C1Log;
   bool IsEqualZero;
if (!OrOnFalseVal && !OrOnTrueVal)
     return nullptr;
 
   Value *Y = OrOnFalseVal ? TrueVal : FalseVal;
 
   unsigned C2Log = C2->logBase2();
 
   bool NeedXor = (!IsEqualZero && OrOnFalseVal) || (IsEqualZero && OrOnTrueVal);
   bool NeedShift = C1Log != C2Log;
-  bool NeedZExtTrunc = Y->getType()->getIntegerBitWidth() !=
-                       V->getType()->getIntegerBitWidth();
+  bool NeedZExtTrunc = Y->getType()->getScalarSizeInBits() !=
+                       V->getType()->getScalarSizeInBits();
 
   // Make sure we don't create more instructions than we save.
   Value *Or = OrOnFalseVal ? FalseVal : TrueVal;
   if ((NeedShift + NeedXor + NeedZExtTrunc) >
       (IC->hasOneUse() + Or->hasOneUse()))
     return nullptr;
 
   if (NeedAnd) {
if (IsBitTest) {
                match(FalseVal, m_Xor(m_Specific(X), m_APInt(C))) && *Y == *C)
         V = Builder.CreateOr(X, *Y);
 
       if (V)
         return replaceInstUsesWith(SI, V);
     }
   }
 
-  if (Value *V = foldSelectICmpAndOr(SI, TrueVal, FalseVal, Builder))
+  if (Value *V = foldSelectICmpAndOr(ICI, TrueVal, FalseVal, Builder))
     return replaceInstUsesWith(SI, V);
 
   if (Value *V = foldSelectCttzCtlz(ICI, TrueVal, FalseVal, Builder))
     return replaceInstUsesWith(SI, V);
 
   return Changed ? &SI : nullptr;
 }
 

llvm/trunk/test/Transforms/InstCombine/select-with-bitwise-ops.ll

 ;
   %and = and i32 %x, 1
   %cmp = icmp eq i32 %and, 0
   %or = or i32 %y, 2
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }
 
+define <2 x i32> @select_icmp_eq_and_1_0_or_2_vec(<2 x i32> %x, <2 x i32> %y) {
+; CHECK-LABEL: @select_icmp_eq_and_1_0_or_2_vec(
+; CHECK-NEXT:    [[AND:%.*]] = shl <2 x i32> [[X:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i32> [[AND]], <i32 2, i32 2>
+; CHECK-NEXT:    [[TMP2:%.*]] = or <2 x i32> [[TMP1]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[TMP2]]
+;
+  %and = and <2 x i32> %x, <i32 1, i32 1>
+  %cmp = icmp eq <2 x i32> %and, zeroinitializer
+  %or = or <2 x i32> %y, <i32 2, i32 2>
+  %select = select <2 x i1> %cmp, <2 x i32> %y, <2 x i32> %or
+  ret <2 x i32> %select
+}
+
 define i32 @select_icmp_eq_and_32_0_or_8(i32 %x, i32 %y) {
 ; CHECK-LABEL: @select_icmp_eq_and_32_0_or_8(
 ; CHECK-NEXT:    [[AND:%.*]] = lshr i32 %x, 2
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[AND]], 8
 ; CHECK-NEXT:    [[TMP2:%.*]] = or i32 [[TMP1]], %y
 ; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   %and = and i32 %x, 32
   %cmp = icmp eq i32 %and, 0
   %or = or i32 %y, 8
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }
 
+define <2 x i32> @select_icmp_eq_and_32_0_or_8_vec(<2 x i32> %x, <2 x i32> %y) {
+; CHECK-LABEL: @select_icmp_eq_and_32_0_or_8_vec(
+; CHECK-NEXT:    [[AND:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 2, i32 2>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i32> [[AND]], <i32 8, i32 8>
+; CHECK-NEXT:    [[TMP2:%.*]] = or <2 x i32> [[TMP1]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[TMP2]]
+;
+  %and = and <2 x i32> %x, <i32 32, i32 32>
+  %cmp = icmp eq <2 x i32> %and, zeroinitializer
+  %or = or <2 x i32> %y, <i32 8, i32 8>
+  %select = select <2 x i1> %cmp, <2 x i32> %y, <2 x i32> %or
+  ret <2 x i32> %select
+}
+
 define i32 @select_icmp_ne_0_and_4096_or_4096(i32 %x, i32 %y) {
 ; CHECK-LABEL: @select_icmp_ne_0_and_4096_or_4096(
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 %x, 4096
 ; CHECK-NEXT:    [[TMP1:%.*]] = xor i32 [[AND]], 4096
 ; CHECK-NEXT:    [[TMP2:%.*]] = or i32 [[TMP1]], %y
 ; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   %and = and i32 %x, 4096
   %cmp = icmp ne i32 0, %and
   %or = or i32 %y, 4096
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }
 
+define <2 x i32> @select_icmp_ne_0_and_4096_or_4096_vec(<2 x i32> %x, <2 x i32> %y) {
+; CHECK-LABEL: @select_icmp_ne_0_and_4096_or_4096_vec(
+; CHECK-NEXT:    [[AND:%.*]] = and <2 x i32> [[X:%.*]], <i32 4096, i32 4096>
+; CHECK-NEXT:    [[TMP1:%.*]] = xor <2 x i32> [[AND]], <i32 4096, i32 4096>
+; CHECK-NEXT:    [[TMP2:%.*]] = or <2 x i32> [[TMP1]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[TMP2]]
+;
+  %and = and <2 x i32> %x, <i32 4096, i32 4096>
+  %cmp = icmp ne <2 x i32> zeroinitializer, %and
+  %or = or <2 x i32> %y, <i32 4096, i32 4096>
+  %select = select <2 x i1> %cmp, <2 x i32> %y, <2 x i32> %or
+  ret <2 x i32> %select
+}
+
 define i32 @select_icmp_eq_and_4096_0_or_4096(i32 %x, i32 %y) {
 ; CHECK-LABEL: @select_icmp_eq_and_4096_0_or_4096(
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 %x, 4096
 ; CHECK-NEXT:    [[TMP1:%.*]] = or i32 [[AND]], %y
 ; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %and = and i32 %x, 4096
   %cmp = icmp eq i32 %and, 0
   %or = or i32 %y, 4096
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }
 
+define <2 x i32> @select_icmp_eq_and_4096_0_or_4096_vec(<2 x i32> %x, <2 x i32> %y) {
+; CHECK-LABEL: @select_icmp_eq_and_4096_0_or_4096_vec(
+; CHECK-NEXT:    [[AND:%.*]] = and <2 x i32> [[X:%.*]], <i32 4096, i32 4096>
+; CHECK-NEXT:    [[TMP1:%.*]] = or <2 x i32> [[AND]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[TMP1]]
+;
+  %and = and <2 x i32> %x, <i32 4096, i32 4096>
+  %cmp = icmp eq <2 x i32> %and, zeroinitializer
+  %or = or <2 x i32> %y, <i32 4096, i32 4096>
+  %select = select <2 x i1> %cmp, <2 x i32> %y, <2 x i32> %or
+  ret <2 x i32> %select
+}
+
 define i32 @select_icmp_eq_0_and_1_or_1(i64 %x, i32 %y) {
 ; CHECK-LABEL: @select_icmp_eq_0_and_1_or_1(
 ; CHECK-NEXT:    [[X_TR:%.*]] = trunc i64 %x to i32
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X_TR]], 1
 ; CHECK-NEXT:    [[TMP2:%.*]] = or i32 [[TMP1]], %y
 ; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   %and = and i64 %x, 1
   %cmp = icmp eq i64 %and, 0
   %or = or i32 %y, 1
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }
 
+define <2 x i32> @select_icmp_eq_0_and_1_or_1_vec(<2 x i64> %x, <2 x i32> %y) {
+; CHECK-LABEL: @select_icmp_eq_0_and_1_or_1_vec(
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc <2 x i64> [[X:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[TMP2:%.*]] = and <2 x i32> [[TMP1]], <i32 1, i32 1>
+; CHECK-NEXT:    [[TMP3:%.*]] = or <2 x i32> [[TMP2]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[TMP3]]
+;
+  %and = and <2 x i64> %x, <i64 1, i64 1>
+  %cmp = icmp eq <2 x i64> %and, zeroinitializer
+  %or = or <2 x i32> %y, <i32 1, i32 1>
+  %select = select <2 x i1> %cmp, <2 x i32> %y, <2 x i32> %or
+  ret <2 x i32> %select
+}
+
 define i32 @select_icmp_ne_0_and_4096_or_32(i32 %x, i32 %y) {
 ; CHECK-LABEL: @select_icmp_ne_0_and_4096_or_32(
 ; CHECK-NEXT:    [[AND:%.*]] = lshr i32 %x, 7
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[AND]], 32
 ; CHECK-NEXT:    [[TMP2:%.*]] = xor i32 [[TMP1]], 32
 ; CHECK-NEXT:    [[TMP3:%.*]] = or i32 [[TMP2]], %y
 ; CHECK-NEXT:    ret i32 [[TMP3]]
 ;

 ;
   %and = and i32 %x, 32
   %cmp = icmp ne i32 0, %and
   %or = or i32 %y, 4096
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }
 
+define <2 x i32> @select_icmp_ne_0_and_32_or_4096_vec(<2 x i32> %x, <2 x i32> %y) {
+; CHECK-LABEL: @select_icmp_ne_0_and_32_or_4096_vec(
+; CHECK-NEXT:    [[AND:%.*]] = shl <2 x i32> [[X:%.*]], <i32 7, i32 7>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i32> [[AND]], <i32 4096, i32 4096>
+; CHECK-NEXT:    [[TMP2:%.*]] = xor <2 x i32> [[TMP1]], <i32 4096, i32 4096>
+; CHECK-NEXT:    [[TMP3:%.*]] = or <2 x i32> [[TMP2]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[TMP3]]
+;
+  %and = and <2 x i32> %x, <i32 32, i32 32>
+  %cmp = icmp ne <2 x i32> zeroinitializer, %and
+  %or = or <2 x i32> %y, <i32 4096, i32 4096>
+  %select = select <2 x i1> %cmp, <2 x i32> %y, <2 x i32> %or
+  ret <2 x i32> %select
+}
+
 define i8 @select_icmp_ne_0_and_1073741824_or_8(i32 %x, i8 %y) {
 ; CHECK-LABEL: @select_icmp_ne_0_and_1073741824_or_8(
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 %x, 1073741824
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND]], 0
 ; CHECK-NEXT:    [[OR:%.*]] = or i8 %y, 8
 ; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[CMP]], i8 [[OR]], i8 %y
 ; CHECK-NEXT:    ret i8 [[SELECT]]
 ;

 ;
   %and = and i32 %x, 8
   %cmp = icmp eq i32 %and, 0
   %or = or i64 %y, 8
   %or.y = select i1 %cmp, i64 %or, i64 %y
   ret i64 %or.y
 }
 
+define <2 x i64> @select_icmp_x_and_8_ne_0_y_or_8_vec(<2 x i32> %x, <2 x i64> %y) {
+; CHECK-LABEL: @select_icmp_x_and_8_ne_0_y_or_8_vec(
+; CHECK-NEXT:    [[AND:%.*]] = and <2 x i32> [[X:%.*]], <i32 8, i32 8>
+; CHECK-NEXT:    [[TMP1:%.*]] = xor <2 x i32> [[AND]], <i32 8, i32 8>
+; CHECK-NEXT:    [[TMP2:%.*]] = zext <2 x i32> [[TMP1]] to <2 x i64>
+; CHECK-NEXT:    [[TMP3:%.*]] = or <2 x i64> [[TMP2]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i64> [[TMP3]]
+;
+  %and = and <2 x i32> %x, <i32 8, i32 8>
+  %cmp = icmp eq <2 x i32> %and, zeroinitializer
+  %or = or <2 x i64> %y, <i64 8, i64 8>
+  %or.y = select <2 x i1> %cmp, <2 x i64> %or, <2 x i64> %y
+  ret <2 x i64> %or.y
+}
+
 define i32 @select_icmp_and_2147483648_ne_0_xor_2147483648(i32 %x) {
 ; CHECK-LABEL: @select_icmp_and_2147483648_ne_0_xor_2147483648(
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 %x, 2147483647
 ; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %and = and i32 %x, 2147483648
   %cmp = icmp eq i32 %and, 0
   %xor = xor i32 %x, 2147483648

 ;
   %and = and i32 %x, 128
   %cmp = icmp eq i32 %and, 0
   %or = or i32 %y, 2
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }
 
+define <2 x i32> @test68vec(<2 x i32> %x, <2 x i32> %y) {
+; CHECK-LABEL: @test68vec(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 6, i32 6>
+; CHECK-NEXT:    [[TMP2:%.*]] = and <2 x i32> [[TMP1]], <i32 2, i32 2>
+; CHECK-NEXT:    [[TMP3:%.*]] = or <2 x i32> [[TMP2]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[TMP3]]
+;
+  %and = and <2 x i32> %x, <i32 128, i32 128>
+  %cmp = icmp eq <2 x i32> %and, zeroinitializer
+  %or = or <2 x i32> %y, <i32 2, i32 2>
+  %select = select <2 x i1> %cmp, <2 x i32> %y, <2 x i32> %or
+  ret <2 x i32> %select
+}
+
 define i32 @test69(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test69(
 ; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[X:%.*]], 6
 ; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 2
 ; CHECK-NEXT:    [[TMP3:%.*]] = xor i32 [[TMP2]], 2
 ; CHECK-NEXT:    [[TMP4:%.*]] = or i32 [[TMP3]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i32 [[TMP4]]
 ;
   %and = and i32 %x, 128
   %cmp = icmp ne i32 %and, 0
   %or = or i32 %y, 2
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }
 
+define <2 x i32> @test69vec(<2 x i32> %x, <2 x i32> %y) {
+; CHECK-LABEL: @test69vec(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 6, i32 6>
+; CHECK-NEXT:    [[TMP2:%.*]] = and <2 x i32> [[TMP1]], <i32 2, i32 2>
+; CHECK-NEXT:    [[TMP3:%.*]] = xor <2 x i32> [[TMP2]], <i32 2, i32 2>
+; CHECK-NEXT:    [[TMP4:%.*]] = or <2 x i32> [[TMP3]], [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[TMP4]]
+;
+  %and = and <2 x i32> %x, <i32 128, i32 128>
+  %cmp = icmp ne <2 x i32> %and, zeroinitializer
+  %or = or <2 x i32> %y, <i32 2, i32 2>
+  %select = select <2 x i1> %cmp, <2 x i32> %y, <2 x i32> %or
+  ret <2 x i32> %select
+}
+
 ; TODO: we should be able to remove this select
 define i8 @test70(i8 %x, i8 %y) {
 ; CHECK-LABEL: @test70(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i8 [[X:%.*]], 0
 ; CHECK-NEXT:    [[OR:%.*]] = or i8 [[Y:%.*]], 2
 ; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[CMP]], i8 [[OR]], i8 [[Y]]
 ; CHECK-NEXT:    ret i8 [[SELECT]]
 ;