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

[dneilson]

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.

[mkazantsev]

I don't get why we need the last comparison. Can we just replace %cast with %val in your example?

[mkazantsev]

Ah, I see the type mismatch, but still. What prevents us from doing the same for less, greater etc. comparisons? Arithmetics?

[dneilson]

In D44997#1051341, @mkazantsev wrote:

Ah, I see the type mismatch, but still. What prevents us from doing the same for less, greater etc. comparisons? Arithmetics?

Nothing prevents us from doing so. I'll relax the restriction.

[dneilson]

• Generalize to predicates other than just (in)equality.

[efriedma]

test/Transforms/InstCombine/icmp-bc-vecsplat.ll

114 ↗

(On Diff #140258)

This can be optimized: https://rise4fun.com/Alive/wVf . Not sure if that's useful in practice.

[dneilson]

test/Transforms/InstCombine/icmp-bc-vecsplat.ll

114 ↗

(On Diff #140258)

Interesting. Not useful for the direct application that I have in mind for this peephole optimization (I really only need eq/ne), but good to keep in mind for the future. Thanks for pointing it out!

[mkazantsev]

lib/Transforms/InstCombine/InstCombineCompares.cpp

2488

I agree that it is true if C.isSplat, but how it will work in else branch?

[spatel]

The matching seems narrower than required. It won't catch cases where we splat some existing vector element (rather than a scalar value). Can we start the match from the shuffle instead:

%splat = shufflevector <M x iN> %x, <4 x iN> undef, <4 x i32> <i32 C, i32 C, i32 C, i32 C>
%cast = bitcast <M x iN> %splat to iMN
%cond = icmp Pred iMN %cast, SplatCmpC

And create this:

%ext = extractelement <4 x iN> %x, i32 C
%cond = icmp Pred iN %ext, CmpC

If there's an insertelement ahead of the extractelement, existing folds should take care of eliminating that.

[dneilson]

• Following spatel's suggestion -- making the pattern match a little more general by only matching the bitcast + shuffle.

[spatel]

lib/Transforms/InstCombine/InstCombineCompares.cpp

2470–2475

Early exit if the bitcast types aren't suitable:

if (!Bitcast->getType()->isIntegerTy() || !Bitcast->getSrcTy()->isIntOrIntVectorTy())
  return nullptr;

2478–2483

Constant::getSplatValue() ?

2497–2501

None of this is necessary AFAIK. A splat shuffle will be canonicalized so that the 2nd operand is undef and the mask constant is adjusted to choose from the 1st operand.

2509–2512

Probably better to handle this part independently in InstSimplify (if it really matters?).

test/Transforms/InstCombine/icmp-bc-vec.ll

194–202

This test isn't changing with this patch. Move it to instsimplify if there's no existing coverage?
It would be easier to see the transforms if you check in all of the tests first based on trunk. Then, just run the update script again with this patch in place locally to rebase.

[dneilson]

lib/Transforms/InstCombine/InstCombineCompares.cpp

2478–2483

Completely missed that interface. Thank you! I was not happy with the gymnastics I had to do here; I'm glad there's a better way.

2497–2501

Right you are. I had in my head, for some reason, that it wasn't getting simplified.

2509–2512

Seemed like a nice bonus; we've already done all of the heavy lifting doing the match and such here. Replicating that into inst-simplify doesn't seem worth the bother.

[dneilson]

• Simplify implementation
• Prune tests (remove redundant & unneccessary); all added tests have now been triple checked that they do not simplify/fold even a single instruction without this patch.

[spatel]

LGTM - see inline comments for more small improvements.

lib/Transforms/InstCombine/InstCombineCompares.cpp

2475–2476

No need to explicitly set these to nullptr; if the match fails, we don't use those variables.

2482

Slightly less generic if you name this 'EltTy' or similar to indicate that it's the element's type.

2492–2493

Less code if we let instcombine handle the replacement:

return new ICmpInst(Pred, Extract, NewC);

test/Transforms/InstCombine/icmp-bc-vec.ll

13–14

We should have a test that shows this most basic transform (the extractelt is in the output). I'd also include a test with a shuffle that changes the number of elements, and a test with weird types...so all-in-one:

define i1 @extending_shuffle_with_weird_types(<2 x i9> %v) {
  %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
}

[dneilson]

• Final updates suggested. Just updating the differential review to match what will be landed.

lib/Transforms/InstCombine/InstCombineCompares.cpp

if (Cmp.hasOneUse() &&
     if (TrueWhenGreaterThan)
       Cond = Builder.CreateOr(Cond, Builder.CreateICmp(ICmpInst::ICMP_SGT, OrigLHS, OrigRHS));
 
     return replaceInstUsesWith(Cmp, Cond);
   }
   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) {
   const APInt *C;
   if (!match(Cmp.getOperand(1), m_APInt(C)))
     return nullptr;
 
   if (auto *BO = dyn_cast<BinaryOperator>(Cmp.getOperand(0))) {
     switch (BO->getOpcode()) {
     case Instruction::Xor:
       if (Instruction *I = foldICmpXorConstant(Cmp, BO, *C))
         return I;
       break;
     case Instruction::And:
       if (Instruction *I = foldICmpAndConstant(Cmp, BO, *C))
         return I;
       break;
     case Instruction::Or:
       if (Instruction *I = foldICmpOrConstant(Cmp, BO, *C))
         return I;
       break;
     case Instruction::Mul:
if (ConstantInt *ConstRHS = dyn_cast<ConstantInt>(Cmp.getOperand(1)))
         return I;
   }
 
   if (auto *TI = dyn_cast<TruncInst>(Cmp.getOperand(0))) {
     if (Instruction *I = foldICmpTruncConstant(Cmp, TI, *C))
       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;
 
   return nullptr;
 }
 
 /// Fold an icmp equality instruction with binary operator LHS and constant RHS:
 /// icmp eq/ne BO, C.

lib/Transforms/InstCombine/InstCombineInternal.h

private:
   Instruction *foldICmpInstWithConstant(ICmpInst &Cmp);
   Instruction *foldICmpInstWithConstantNotInt(ICmpInst &Cmp);
   Instruction *foldICmpBinOp(ICmpInst &Cmp);
   Instruction *foldICmpEquality(ICmpInst &Cmp);
   Instruction *foldICmpWithZero(ICmpInst &Cmp);
 
   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,
                                    const APInt &C);
   Instruction *foldICmpXorConstant(ICmpInst &Cmp, BinaryOperator *Xor,
                                    const APInt &C);
   Instruction *foldICmpOrConstant(ICmpInst &Cmp, BinaryOperator *Or,
                                   const APInt &C);

test/Transforms/InstCombine/icmp-bc-vec.ll

+; 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
+}