[InstCombine] Support vector splats in transformZExtICmp

[craig.topper]

Summary

This patch adds splat support to transformZExtICmp. The test cases are vector versions of tests that failed when commenting out parts of the existing scalar code.

One test didn't vectorize optimize properly due to another bug so a TODO has been added.

[spatel]

LGTM.

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

Instruction *InstCombiner::visitTrunc(TruncInst &CI) {
   return nullptr;
 }
 
 Instruction *InstCombiner::transformZExtICmp(ICmpInst *ICI, ZExtInst &CI,
                                              bool DoTransform) {
   // If we are just checking for a icmp eq of a single bit and zext'ing it
   // to an integer, then shift the bit to the appropriate place and then
   // cast to integer to avoid the comparison.
-  if (ConstantInt *Op1C = dyn_cast<ConstantInt>(ICI->getOperand(1))) {
-    const APInt &Op1CV = Op1C->getValue();
+  const APInt *Op1CV;
+  if (match(ICI->getOperand(1), m_APInt(Op1CV))) {
 
     // zext (x <s  0) to i32 --> x>>u31      true if signbit set.
     // zext (x >s -1) to i32 --> (x>>u31)^1  true if signbit clear.
-    if ((ICI->getPredicate() == ICmpInst::ICMP_SLT && Op1CV.isNullValue()) ||
-        (ICI->getPredicate() == ICmpInst::ICMP_SGT && Op1CV.isAllOnesValue())) {
+    if ((ICI->getPredicate() == ICmpInst::ICMP_SLT && Op1CV->isNullValue()) ||
+        (ICI->getPredicate() == ICmpInst::ICMP_SGT && Op1CV->isAllOnesValue())) {
       if (!DoTransform) return ICI;
 
       Value *In = ICI->getOperand(0);
       Value *Sh = ConstantInt::get(In->getType(),
                                    In->getType()->getScalarSizeInBits() - 1);
       In = Builder.CreateLShr(In, Sh, In->getName() + ".lobit");
       if (In->getType() != CI.getType())
         In = Builder.CreateIntCast(In, CI.getType(), false /*ZExt*/);
if (match(ICI->getOperand(1), m_APInt(Op1CV))) {
     // zext (X == 0) to i32 --> X^1      iff X has only the low bit set.
     // zext (X == 0) to i32 --> (X>>1)^1 iff X has only the 2nd bit set.
     // zext (X == 1) to i32 --> X        iff X has only the low bit set.
     // zext (X == 2) to i32 --> X>>1     iff X has only the 2nd bit set.
     // zext (X != 0) to i32 --> X        iff X has only the low bit set.
     // zext (X != 0) to i32 --> X>>1     iff X has only the 2nd bit set.
     // zext (X != 1) to i32 --> X^1      iff X has only the low bit set.
     // zext (X != 2) to i32 --> (X>>1)^1 iff X has only the 2nd bit set.
-    if ((Op1CV.isNullValue() || Op1CV.isPowerOf2()) &&
+    if ((Op1CV->isNullValue() || Op1CV->isPowerOf2()) &&
         // This only works for EQ and NE
         ICI->isEquality()) {
       // If Op1C some other power of two, convert:
       KnownBits Known = computeKnownBits(ICI->getOperand(0), 0, &CI);
 
       APInt KnownZeroMask(~Known.Zero);
       if (KnownZeroMask.isPowerOf2()) { // Exactly 1 possible 1?
         if (!DoTransform) return ICI;
 
         bool isNE = ICI->getPredicate() == ICmpInst::ICMP_NE;
-        if (!Op1CV.isNullValue() && (Op1CV != KnownZeroMask)) {
+        if (!Op1CV->isNullValue() && (*Op1CV != KnownZeroMask)) {
           // (X&4) == 2 --> false
           // (X&4) != 2 --> true
           Constant *Res = ConstantInt::get(Type::getInt1Ty(CI.getContext()),
                                            isNE);
           Res = ConstantExpr::getZExt(Res, CI.getType());
           return replaceInstUsesWith(CI, Res);
         }
 
         uint32_t ShAmt = KnownZeroMask.logBase2();
         Value *In = ICI->getOperand(0);
         if (ShAmt) {
           // Perform a logical shr by shiftamt.
           // Insert the shift to put the result in the low bit.
           In = Builder.CreateLShr(In, ConstantInt::get(In->getType(), ShAmt),
                                   In->getName() + ".lobit");
         }
 
-        if (!Op1CV.isNullValue() == isNE) { // Toggle the low bit.
+        if (!Op1CV->isNullValue() == isNE) { // Toggle the low bit.
           Constant *One = ConstantInt::get(In->getType(), 1);
           In = Builder.CreateXor(In, One);
         }
 
         if (CI.getType() == In->getType())
           return replaceInstUsesWith(CI, In);
 
         Value *IntCast = Builder.CreateIntCast(In, CI.getType(), false);

llvm/trunk/test/Transforms/InstCombine/compare-signs.ll

 ;
   %t0 = lshr i32 %a, 31
   %t1 = lshr i32 %b, 31
   %t2 = icmp eq i32 %t0, %t1
   %t3 = zext i1 %t2 to i32
   ret i32 %t3
 }
 
+; TODO this should optimize but doesn't due to missing vector support in InstCombiner::foldICmpEquality.
+define <2 x i32> @test3vec(<2 x i32> %a, <2 x i32> %b) nounwind readnone {
+; CHECK-LABEL: @test3vec(
+; CHECK-NEXT:    [[T0:%.*]] = lshr <2 x i32> [[A:%.*]], <i32 31, i32 31>
+; CHECK-NEXT:    [[T1:%.*]] = lshr <2 x i32> [[B:%.*]], <i32 31, i32 31>
+; CHECK-NEXT:    [[T2:%.*]] = icmp eq <2 x i32> [[T0]], [[T1]]
+; CHECK-NEXT:    [[T3:%.*]] = zext <2 x i1> [[T2]] to <2 x i32>
+; CHECK-NEXT:    ret <2 x i32> [[T3]]
+;
+  %t0 = lshr <2 x i32> %a, <i32 31, i32 31>
+  %t1 = lshr <2 x i32> %b, <i32 31, i32 31>
+  %t2 = icmp eq <2 x i32> %t0, %t1
+  %t3 = zext <2 x i1> %t2 to <2 x i32>
+  ret <2 x i32> %t3
+}
+
 ; Variation on @test3: checking the 2nd bit in a situation where the 5th bit
 ; is one, not zero.
 define i32 @test3i(i32 %a, i32 %b) nounwind readnone {
 ; CHECK-LABEL: @test3i(
 ; CHECK-NEXT:    [[T01:%.*]] = xor i32 %a, %b
 ; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[T01]], 31
 ; CHECK-NEXT:    [[T4:%.*]] = xor i32 [[TMP1]], 1
 ; CHECK-NEXT:    ret i32 [[T4]]

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 target datalayout = "e-p:64:64:64-p1:16:16:16-p2:32:32:32-p3:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64"
 
 define i32 @test1(i32 %X) {
 ; CHECK-LABEL: @test1(
 ; CHECK-NEXT:    [[X_LOBIT:%.*]] = lshr i32 %X, 31
 ; CHECK-NEXT:    ret i32 [[X_LOBIT]]
 ;
   %a = icmp slt i32 %X, 0
   %b = zext i1 %a to i32
   ret i32 %b
 }
 
+define <2 x i32> @test1vec(<2 x i32> %X) {
+; CHECK-LABEL: @test1vec(
+; CHECK-NEXT:    [[X_LOBIT:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 31, i32 31>
+; CHECK-NEXT:    ret <2 x i32> [[X_LOBIT]]
+;
+  %a = icmp slt <2 x i32> %X, zeroinitializer
+  %b = zext <2 x i1> %a to <2 x i32>
+  ret <2 x i32> %b
+}
+
 define i32 @test2(i32 %X) {
 ; CHECK-LABEL: @test2(
 ; CHECK-NEXT:    [[X_LOBIT:%.*]] = lshr i32 %X, 31
 ; CHECK-NEXT:    [[X_LOBIT_NOT:%.*]] = xor i32 [[X_LOBIT]], 1
 ; CHECK-NEXT:    ret i32 [[X_LOBIT_NOT]]
 ;
   %a = icmp ult i32 %X, -2147483648
   %b = zext i1 %a to i32
   ret i32 %b
 }
 
+define <2 x i32> @test2vec(<2 x i32> %X) {
+; CHECK-LABEL: @test2vec(
+; CHECK-NEXT:    [[X_LOBIT:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 31, i32 31>
+; CHECK-NEXT:    [[X_LOBIT_NOT:%.*]] = xor <2 x i32> [[X_LOBIT]], <i32 1, i32 1>
+; CHECK-NEXT:    ret <2 x i32> [[X_LOBIT_NOT]]
+;
+  %a = icmp ult <2 x i32> %X, <i32 -2147483648, i32 -2147483648>
+  %b = zext <2 x i1> %a to <2 x i32>
+  ret <2 x i32> %b
+}
+
 define i32 @test3(i32 %X) {
 ; CHECK-LABEL: @test3(
 ; CHECK-NEXT:    [[X_LOBIT:%.*]] = ashr i32 %X, 31
 ; CHECK-NEXT:    ret i32 [[X_LOBIT]]
 ;
   %a = icmp slt i32 %X, 0
   %b = sext i1 %a to i32
   ret i32 %b

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

 ; CHECK-NEXT:    ret i32 [[B]]
 ;
   %B = and i32 %A, 1
   %C = icmp ne i32 %B, 0
   %D = zext i1 %C to i32
   ret i32 %D
 }
 
+define <2 x i32> @test20vec(<2 x i32> %A) {
+; CHECK-LABEL: @test20vec(
+; CHECK-NEXT:    [[B:%.*]] = and <2 x i32> [[A:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    ret <2 x i32> [[B]]
+;
+  %B = and <2 x i32> %A, <i32 1, i32 1>
+  %C = icmp ne <2 x i32> %B, zeroinitializer
+  %D = zext <2 x i1> %C to <2 x i32>
+  ret <2 x i32> %D
+}
+
 define i32 @test21(i32 %a) {
 ; CHECK-LABEL: @test21(
 ; CHECK-NEXT:    [[TMP_6:%.*]] = lshr i32 %a, 2
 ; CHECK-NEXT:    [[TMP_6_LOBIT:%.*]] = and i32 [[TMP_6]], 1
 ; CHECK-NEXT:    ret i32 [[TMP_6_LOBIT]]
 ;
   %tmp.6 = and i32 %a, 4
   %not.tmp.7 = icmp ne i32 %tmp.6, 0
   %retval = zext i1 %not.tmp.7 to i32
   ret i32 %retval
 }
 
+define <2 x i32> @test21vec(<2 x i32> %a) {
+; CHECK-LABEL: @test21vec(
+; CHECK-NEXT:    [[TMP_6:%.*]] = lshr <2 x i32> [[A:%.*]], <i32 2, i32 2>
+; CHECK-NEXT:    [[TMP_6_LOBIT:%.*]] = and <2 x i32> [[TMP_6]], <i32 1, i32 1>
+; CHECK-NEXT:    ret <2 x i32> [[TMP_6_LOBIT]]
+;
+  %tmp.6 = and <2 x i32> %a, <i32 4, i32 4>
+  %not.tmp.7 = icmp ne <2 x i32> %tmp.6, zeroinitializer
+  %retval = zext <2 x i1> %not.tmp.7 to <2 x i32>
+  ret <2 x i32> %retval
+}
+
 define i1 @test22(i32 %A, i32 %X) {
 ; CHECK-LABEL: @test22(
 ; CHECK-NEXT:    ret i1 true
 ;
   %B = and i32 %A, 100663295
   %C = icmp ult i32 %B, 268435456
   %Y = and i32 %X, 7
   %Z = icmp sgt i32 %Y, -1
   %R = or i1 %C, %Z
   ret i1 %R
 }
 
 define i32 @test23(i32 %a) {
 ; CHECK-LABEL: @test23(
 ; CHECK-NEXT:    [[TMP_1:%.*]] = and i32 %a, 1
 ; CHECK-NEXT:    [[TMP1:%.*]] = xor i32 [[TMP_1]], 1
 ; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %tmp.1 = and i32 %a, 1
   %tmp.2 = icmp eq i32 %tmp.1, 0
   %tmp.3 = zext i1 %tmp.2 to i32
   ret i32 %tmp.3
 }
 
+define <2 x i32> @test23vec(<2 x i32> %a) {
+; CHECK-LABEL: @test23vec(
+; CHECK-NEXT:    [[TMP_1:%.*]] = and <2 x i32> [[A:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[TMP1:%.*]] = xor <2 x i32> [[TMP_1]], <i32 1, i32 1>
+; CHECK-NEXT:    ret <2 x i32> [[TMP1]]
+;
+  %tmp.1 = and <2 x i32> %a, <i32 1, i32 1>
+  %tmp.2 = icmp eq <2 x i32> %tmp.1, zeroinitializer
+  %tmp.3 = zext <2 x i1> %tmp.2 to <2 x i32>
+  ret <2 x i32> %tmp.3
+}
+
 define i32 @test24(i32 %a) {
 ; CHECK-LABEL: @test24(
 ; CHECK-NEXT:    [[TMP_1:%.*]] = lshr i32 %a, 2
 ; CHECK-NEXT:    [[TMP_1_LOBIT:%.*]] = and i32 [[TMP_1]], 1
 ; CHECK-NEXT:    [[TMP1:%.*]] = xor i32 [[TMP_1_LOBIT]], 1
 ; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %tmp1 = and i32 %a, 4
   %tmp.1 = lshr i32 %tmp1, 2
   %tmp.2 = icmp eq i32 %tmp.1, 0
   %tmp.3 = zext i1 %tmp.2 to i32
   ret i32 %tmp.3
 }
 
+define <2 x i32> @test24vec(<2 x i32> %a) {
+; CHECK-LABEL: @test24vec(
+; CHECK-NEXT:    [[TMP_1:%.*]] = lshr <2 x i32> [[A:%.*]], <i32 2, i32 2>
+; CHECK-NEXT:    [[TMP_1_LOBIT:%.*]] = and <2 x i32> [[TMP_1]], <i32 1, i32 1>
+; CHECK-NEXT:    [[TMP1:%.*]] = xor <2 x i32> [[TMP_1_LOBIT]], <i32 1, i32 1>
+; CHECK-NEXT:    ret <2 x i32> [[TMP1]]
+;
+  %tmp1 = and <2 x i32> %a, <i32 4, i32 4>
+  %tmp.1 = lshr <2 x i32> %tmp1, <i32 2, i32 2>
+  %tmp.2 = icmp eq <2 x i32> %tmp.1, zeroinitializer
+  %tmp.3 = zext <2 x i1> %tmp.2 to <2 x i32>
+  ret <2 x i32> %tmp.3
+}
+
 define i1 @test25(i32 %A) {
 ; CHECK-LABEL: @test25(
 ; CHECK-NEXT:    ret i1 false
 ;
   %B = and i32 %A, 2
   %C = icmp ugt i32 %B, 2
   ret i1 %C
 }