[InstCombine] Fold compare of int constant against a splatted vector of ints

Summary:
Folding patterns like:
  %vec = shufflevector <4 x i8> %insvec, <4 x i8> undef, <4 x i32> zeroinitializer
  %cast = bitcast <4 x i8> %vec to i32
  %cond = icmp eq i32 %cast, 0
into:
  %ext = extractelement <4 x i8> %insvec, i32 0
  %cond = icmp eq i32 %ext, 0

Combined with existing rules, this allows us to fold patterns like:
  %insvec = insertelement <4 x i8> undef, i8 %val, i32 0
  %vec = shufflevector <4 x i8> %insvec, <4 x i8> undef, <4 x i32> zeroinitializer
  %cast = bitcast <4 x i8> %vec to i32
  %cond = icmp eq i32 %cast, 0
into:
  %cond = icmp eq i8 %val, 0

When we construct a splat vector via a shuffle, and bitcast the vector into an integer type for comparison against an integer constant. Then we can simplify the the comparison to compare the splatted value against the integer constant.

Reviewers: spatel, anna, mkazantsev

Reviewed By: spatel

Subscribers: efriedma, rengolin, llvm-commits

Differential Revision: https://reviews.llvm.org/D44997

llvm-svn: 329087

3 files changed, 173 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index cea3a1d705b..246a335199c 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -2457,6 +2457,45 @@ Instruction *InstCombiner::foldICmpSelectConstant(ICmpInst &Cmp,
   return nullptr;
 }
 
+Instruction *InstCombiner::foldICmpBitCastConstant(ICmpInst &Cmp,
+                                                   BitCastInst *Bitcast,
+                                                   const APInt &C) {
+  // Folding: icmp <pred> iN X, C
+  //  where X = bitcast <M x iK> (shufflevector <M x iK> %vec, undef, SC)) to iN
+  //    and C is a splat of a K-bit pattern
+  //    and SC is a constant vector = <C', C', C', ..., C'>
+  // Into:
+  //   %E = extractelement <M x iK> %vec, i32 C'
+  //   icmp <pred> iK %E, trunc(C)
+  if (!Bitcast->getType()->isIntegerTy() ||
+      !Bitcast->getSrcTy()->isIntOrIntVectorTy())
+    return nullptr;
+
+  Value *BCIOp = Bitcast->getOperand(0);
+  Value *Vec = nullptr;     // 1st vector arg of the shufflevector
+  Constant *Mask = nullptr; // Mask arg of the shufflevector
+  if (match(BCIOp,
+            m_ShuffleVector(m_Value(Vec), m_Undef(), m_Constant(Mask)))) {
+    // Check whether every element of Mask is the same constant
+    if (auto *Elem = dyn_cast_or_null<ConstantInt>(Mask->getSplatValue())) {
+      auto *VecTy = cast<VectorType>(BCIOp->getType());
+      auto *EltTy = cast<IntegerType>(VecTy->getElementType());
+      auto Pred = Cmp.getPredicate();
+      if (C.isSplat(EltTy->getBitWidth())) {
+        // Fold the icmp based on the value of C
+        // If C is M copies of an iK sized bit pattern,
+        // then:
+        //   =>  %E = extractelement <N x iK> %vec, i32 Elem
+        //       icmp <pred> iK %SplatVal, <pattern>
+        Value *Extract = Builder.CreateExtractElement(Vec, Elem);
+        Value *NewC = ConstantInt::get(EltTy, C.trunc(EltTy->getBitWidth()));
+        return new ICmpInst(Pred, Extract, NewC);
+      }
+    }
+  }
+  return nullptr;
+}
+
 /// Try to fold integer comparisons with a constant operand: icmp Pred X, C
 /// where X is some kind of instruction.
 Instruction *InstCombiner::foldICmpInstWithConstant(ICmpInst &Cmp) {
@@ -2531,6 +2570,11 @@ Instruction *InstCombiner::foldICmpInstWithConstant(ICmpInst &Cmp) {
       return I;
   }
 
+  if (auto *BCI = dyn_cast<BitCastInst>(Cmp.getOperand(0))) {
+    if (Instruction *I = foldICmpBitCastConstant(Cmp, BCI, *C))
+      return I;
+  }
+
   if (Instruction *I = foldICmpIntrinsicWithConstant(Cmp, *C))
     return I;
 
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 5e934d542c0..40cfc3051c2 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -734,6 +734,8 @@ private:
 
   Instruction *foldICmpSelectConstant(ICmpInst &Cmp, SelectInst *Select,
                                       ConstantInt *C);
+  Instruction *foldICmpBitCastConstant(ICmpInst &Cmp, BitCastInst *Bitcast,
+                                       const APInt &C);
   Instruction *foldICmpTruncConstant(ICmpInst &Cmp, TruncInst *Trunc,
                                      const APInt &C);
   Instruction *foldICmpAndConstant(ICmpInst &Cmp, BinaryOperator *And,
diff --git a/llvm/test/Transforms/InstCombine/icmp-bc-vec.ll b/llvm/test/Transforms/InstCombine/icmp-bc-vec.ll
new file mode 100644
index 00000000000..26252d2b402
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/icmp-bc-vec.ll
@@ -0,0 +1,127 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+
+; Tests to verify proper functioning of the icmp folding implemented in
+;  InstCombiner::foldICmpBitCastConstant
+; Specifically, folding:
+;   icmp <pred> iN X, C
+;  where X = bitcast <M x iK> (shufflevector <M x iK> %vec, undef, SC)) to iN
+;    and C is a splat of a K-bit pattern
+;    and SC is a constant vector = <C', C', C', ..., C'>
+; Into:
+;  %E = extractelement <M x iK> %vec, i32 C'
+;  icmp <pred> iK %E, trunc(C)
+
+define i1 @test_i1_0(i1 %val) {
+; CHECK-LABEL: @test_i1_0(
+; CHECK-NEXT:    [[COND:%.*]] = xor i1 [[VAL:%.*]], true
+; CHECK-NEXT:    ret i1 [[COND]]
+;
+  %insvec = insertelement <4 x i1> undef, i1 %val, i32 0
+  %vec = shufflevector <4 x i1> %insvec, <4 x i1> undef, <4 x i32> zeroinitializer
+  %cast = bitcast <4 x i1> %vec to i4
+  %cond = icmp eq i4 %cast, 0
+  ret i1 %cond
+}
+
+define i1 @test_i1_0_2(i1 %val) {
+; CHECK-LABEL: @test_i1_0_2(
+; CHECK-NEXT:    [[COND:%.*]] = xor i1 [[VAL:%.*]], true
+; CHECK-NEXT:    ret i1 [[COND]]
+;
+  %insvec = insertelement <4 x i1> undef, i1 %val, i32 2
+  %vec = shufflevector <4 x i1> %insvec, <4 x i1> undef, <4 x i32> <i32 2, i32 2, i32 2, i32 2>
+  %cast = bitcast <4 x i1> %vec to i4
+  %cond = icmp eq i4 %cast, 0
+  ret i1 %cond
+}
+
+define i1 @test_i1_m1(i1 %val) {
+; CHECK-LABEL: @test_i1_m1(
+; CHECK-NEXT:    ret i1 [[VAL:%.*]]
+;
+  %insvec = insertelement <4 x i1> undef, i1 %val, i32 0
+  %vec = shufflevector <4 x i1> %insvec, <4 x i1> undef, <4 x i32> zeroinitializer
+  %cast = bitcast <4 x i1> %vec to i4
+  %cond = icmp eq i4 %cast, -1
+  ret i1 %cond
+}
+
+define i1 @test_i8_pattern(i8 %val) {
+; CHECK-LABEL: @test_i8_pattern(
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i8 [[VAL:%.*]], 72
+; CHECK-NEXT:    ret i1 [[COND]]
+;
+  %insvec = insertelement <4 x i8> undef, i8 %val, i32 0
+  %vec = shufflevector <4 x i8> %insvec, <4 x i8> undef, <4 x i32> zeroinitializer
+  %cast = bitcast <4 x i8> %vec to i32
+  %cond = icmp eq i32 %cast, 1212696648
+  ret i1 %cond
+}
+
+define i1 @test_i8_pattern_2(i8 %val) {
+; CHECK-LABEL: @test_i8_pattern_2(
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i8 [[VAL:%.*]], 72
+; CHECK-NEXT:    ret i1 [[COND]]
+;
+  %insvec = insertelement <4 x i8> undef, i8 %val, i32 2
+  %vec = shufflevector <4 x i8> %insvec, <4 x i8> undef, <4 x i32> <i32 2, i32 2, i32 2, i32 2>
+  %cast = bitcast <4 x i8> %vec to i32
+  %cond = icmp eq i32 %cast, 1212696648
+  ret i1 %cond
+}
+
+; Make sure we don't try to fold if the shufflemask has differing element values
+define i1 @test_i8_pattern_3(<4 x i8> %invec) {
+; CHECK-LABEL: @test_i8_pattern_3(
+; CHECK-NEXT:    [[VEC:%.*]] = shufflevector <4 x i8> [[INVEC:%.*]], <4 x i8> undef, <4 x i32> <i32 1, i32 0, i32 3, i32 2>
+; CHECK-NEXT:    [[CAST:%.*]] = bitcast <4 x i8> [[VEC]] to i32
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[CAST]], 1212696648
+; CHECK-NEXT:    ret i1 [[COND]]
+;
+  %vec = shufflevector <4 x i8> %invec, <4 x i8> undef, <4 x i32> <i32 1, i32 0, i32 3, i32 2>
+  %cast = bitcast <4 x i8> %vec to i32
+  %cond = icmp eq i32 %cast, 1212696648
+  ret i1 %cond
+}
+
+; Make sure we don't try to fold if the compared-to constant isn't a splatted value
+define i1 @test_i8_nopattern(i8 %val) {
+; CHECK-LABEL: @test_i8_nopattern(
+; CHECK-NEXT:    [[INSVEC:%.*]] = insertelement <4 x i8> undef, i8 [[VAL:%.*]], i32 0
+; CHECK-NEXT:    [[VEC:%.*]] = shufflevector <4 x i8> [[INSVEC]], <4 x i8> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[CAST:%.*]] = bitcast <4 x i8> [[VEC]] to i32
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[CAST]], 1212696647
+; CHECK-NEXT:    ret i1 [[COND]]
+;
+  %insvec = insertelement <4 x i8> undef, i8 %val, i32 0
+  %vec = shufflevector <4 x i8> %insvec, <4 x i8> undef, <4 x i32> zeroinitializer
+  %cast = bitcast <4 x i8> %vec to i32
+  %cond = icmp eq i32 %cast, 1212696647
+  ret i1 %cond
+}
+
+; Verify that we fold more than just the eq predicate
+define i1 @test_i8_ult_pattern(i8 %val) {
+; CHECK-LABEL: @test_i8_ult_pattern(
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i8 [[VAL:%.*]], 72
+; CHECK-NEXT:    ret i1 [[COND]]
+;
+  %insvec = insertelement <4 x i8> undef, i8 %val, i32 0
+  %vec = shufflevector <4 x i8> %insvec, <4 x i8> undef, <4 x i32> zeroinitializer
+  %cast = bitcast <4 x i8> %vec to i32
+  %cond = icmp ult i32 %cast, 1212696648
+  ret i1 %cond
+}
+
+define i1 @extending_shuffle_with_weird_types(<2 x i9> %v) {
+; CHECK-LABEL: @extending_shuffle_with_weird_types(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i9> [[V:%.*]], i32 0
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i9 [[TMP1]], 1
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %splat = shufflevector <2 x i9> %v, <2 x i9> undef, <3 x i32> zeroinitializer
+  %cast = bitcast <3 x i9> %splat to i27
+  %cmp = icmp slt i27 %cast, 262657 ; 0x040201
+  ret i1 %cmp
+}