[InstCombine] canonicalize vector select with constant vector condition to shuffle

[spatel]

Summary

As discussed on llvm-dev ( http://lists.llvm.org/pipermail/llvm-dev/2016-August/104210.html ): turn a vector select into a shuffle when possible as a canonicalization step. Shuffles may be easier to reason about in conjunction with other shuffles and insert/extract.

[majnemer]

lib/Transforms/InstCombine/InstCombineSelect.cpp

965

I'd just reserve NumElts and push_back in the loop.

966

Might be shorter to use Builder->getInt32Ty()

967

I think we'd just use unsigned.

968

Can't this return null?

[spatel]

lib/Transforms/InstCombine/InstCombineSelect.cpp

968

Argh...yes. Let me concoct some tests from the last time I forgot about constant expressions. :)

[spatel]

Patch updated:

1. If the select condition is a constant expression - getAggregateElement returns nullptr - bail; add test.
2. If any element of the select condition is a constant expression - ie, not {0,1,undef} - bail; add test.
3. The code grew just enough to warrant being a helper function IMO, so...
4. I didn't use Builder->getInt32Ty() to avoid needing the builder as a param.
5. Fixed to use 'reserve' and 'push_back'.
6. Fixed 'unsigned int'.

[spatel]

Ping.

[mkuper]

Thanks Sanjay, as I wrote on the email thread, I think this is probably the right canonicalization.
Did you happen to check that we don't regress codegen for the IR tests that changed? At least X86? (This should probably go in even if it does regress some cases, but a PR would be good.)

test/Transforms/InstCombine/select.ll

1794

Yikes.

[spatel]

In D24279#541376, @mkuper wrote:

Did you happen to check that we don't regress codegen for the IR tests that changed? At least X86? (This should probably go in even if it does regress some cases, but a PR would be good.)

Yes - x86 is the easy target because the x86 DAG combining/lowering already turns vselect into shuffle which then gets matched to blend or whatever shuffle instruction scraps are available before SSE4.1.

There is one existing problem for x86 that becomes slightly worse with this change. It manifests in vec_demanded_elts.ll:test_select(), and I just filed PR30371 (https://llvm.org/bugs/show_bug.cgi?id=30371) for that. That test also demonstrates a missing InstCombine fold that could fuse an insertelement with a later shuffle, but even if we had that fold, the x86 backend would screw up the codegen. :)

I filed https://llvm.org/bugs/show_bug.cgi?id=28530 for AArch64.

[mkuper]

Thanks, Sanjay.
Anyway, this LGTM, but I'm not an InstCombine expert,so take this with the appropriate amount of salt...

[spatel]

In D24279#545080, @mkuper wrote:

Thanks, Sanjay.
Anyway, this LGTM, but I'm not an InstCombine expert,so take this with the appropriate amount of salt...

Thanks, Michael! I haven't heard any philosophical objections, so I'll commit this soon, and then I'll post back on the llvm-dev thread so people can keep an eye open for any regressions.

lib/Transforms/InstCombine/InstCombineSelect.cpp

Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();
   Type *SelType = SI.getType();
 
   if (Value *V =
           SimplifySelectInst(CondVal, TrueVal, FalseVal, DL, &TLI, &DT, &AC))
     return replaceInstUsesWith(SI, V);
 
+  // A vector select with a constant condition vector is canonicalized to a
+  // shuffle for easier combining with other shuffles and insert/extract.
+  Constant *CondC;
+  if (CondVal->getType()->isVectorTy() && match(CondVal, m_Constant(CondC))) {
+    unsigned NumElts = CondVal->getType()->getVectorNumElements();
+    SmallVector<Constant *, 16> Mask(NumElts);
+    Type *Int32Ty = Type::getInt32Ty(CondVal->getContext());
+    for (unsigned int i = 0; i != NumElts; ++i) {
+      Constant *Elt = CondC->getAggregateElement(i);
+      if (Elt->isOneValue()) {
+        // If the select condition element is true, choose from the 1st vector.
+        Mask[i] = ConstantInt::get(Int32Ty, i);
+      } else if (Elt->isNullValue()) {
+        // If the select condition element is false, choose from the 2nd vector.
+        Mask[i] = ConstantInt::get(Int32Ty, i + NumElts);
+      } else {
+        // If the select condition element is undef, the shuffle mask is undef.
+        Mask[i] = UndefValue::get(Int32Ty);
+      }
+    }
+    return new ShuffleVectorInst(TrueVal, FalseVal, ConstantVector::get(Mask));
+  }
+
   if (SelType->getScalarType()->isIntegerTy(1) &&
       TrueVal->getType() == CondVal->getType()) {
     if (match(TrueVal, m_One())) {
       // Change: A = select B, true, C --> A = or B, C
       return BinaryOperator::CreateOr(CondVal, FalseVal);
     }
     if (match(TrueVal, m_Zero())) {
       // Change: A = select B, false, C --> A = and !B, C

test/Transforms/InstCombine/blend_x86.ll

 ; RUN: opt < %s -instcombine -mtriple=x86_64-apple-macosx -mcpu=core-avx2 -S | FileCheck %s
 
 define <2 x double> @constant_blendvpd(<2 x double> %xy, <2 x double> %ab) {
-; CHECK-LABEL: @constant_blendvpd
-; CHECK-NEXT: %1 = select <2 x i1> <i1 true, i1 false>, <2 x double> %ab, <2 x double> %xy
-; CHECK-NEXT: ret <2 x double> %1
+; CHECK-LABEL: @constant_blendvpd(
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <2 x double> %ab, <2 x double> %xy, <2 x i32> <i32 0, i32 3>
+; CHECK-NEXT:    ret <2 x double> [[TMP1]]
+;
   %1 = tail call <2 x double> @llvm.x86.sse41.blendvpd(<2 x double> %xy, <2 x double> %ab, <2 x double> <double 0xFFFFFFFFE0000000, double 0.000000e+00>)
   ret <2 x double> %1
 }
 
 define <2 x double> @constant_blendvpd_zero(<2 x double> %xy, <2 x double> %ab) {
 ; CHECK-LABEL: @constant_blendvpd_zero
 ; CHECK-NEXT: ret <2 x double> %xy
   %1 = tail call <2 x double> @llvm.x86.sse41.blendvpd(<2 x double> %xy, <2 x double> %ab, <2 x double> zeroinitializer)
   ret <2 x double> %1
 }
 
 define <2 x double> @constant_blendvpd_dup(<2 x double> %xy, <2 x double> %sel) {
 ; CHECK-LABEL: @constant_blendvpd_dup
 ; CHECK-NEXT: ret <2 x double> %xy
   %1 = tail call <2 x double> @llvm.x86.sse41.blendvpd(<2 x double> %xy, <2 x double> %xy, <2 x double> %sel)
   ret <2 x double> %1
 }
 
 define <4 x float> @constant_blendvps(<4 x float> %xyzw, <4 x float> %abcd) {
-; CHECK-LABEL: @constant_blendvps
-; CHECK-NEXT: %1 = select <4 x i1> <i1 false, i1 false, i1 false, i1 true>, <4 x float> %abcd, <4 x float> %xyzw
-; CHECK-NEXT: ret <4 x float> %1
+; CHECK-LABEL: @constant_blendvps(
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x float> %abcd, <4 x float> %xyzw, <4 x i32> <i32 4, i32 5, i32 6, i32 3>
+; CHECK-NEXT:    ret <4 x float> [[TMP1]]
+;
   %1 = tail call <4 x float> @llvm.x86.sse41.blendvps(<4 x float> %xyzw, <4 x float> %abcd, <4 x float> <float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0xFFFFFFFFE0000000>)
   ret <4 x float> %1
 }
 
 define <4 x float> @constant_blendvps_zero(<4 x float> %xyzw, <4 x float> %abcd) {
 ; CHECK-LABEL: @constant_blendvps_zero
 ; CHECK-NEXT: ret <4 x float> %xyzw
   %1 = tail call <4 x float> @llvm.x86.sse41.blendvps(<4 x float> %xyzw, <4 x float> %abcd, <4 x float> zeroinitializer)
   ret <4 x float> %1
 }
 
 define <4 x float> @constant_blendvps_dup(<4 x float> %xyzw, <4 x float> %sel) {
 ; CHECK-LABEL: @constant_blendvps_dup
 ; CHECK-NEXT: ret <4 x float> %xyzw
   %1 = tail call <4 x float> @llvm.x86.sse41.blendvps(<4 x float> %xyzw, <4 x float> %xyzw, <4 x float> %sel)
   ret <4 x float> %1
 }
 
 define <16 x i8> @constant_pblendvb(<16 x i8> %xyzw, <16 x i8> %abcd) {
-; CHECK-LABEL: @constant_pblendvb
-; CHECK-NEXT: %1 = select <16 x i1> <i1 false, i1 false, i1 true, i1 false, i1 true, i1 true, i1 true, i1 false, i1 false, i1 false, i1 true, i1 false, i1 true, i1 true, i1 true, i1 false>, <16 x i8> %abcd, <16 x i8> %xyzw
-; CHECK-NEXT: ret <16 x i8> %1
+; CHECK-LABEL: @constant_pblendvb(
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <16 x i8> %abcd, <16 x i8> %xyzw, <16 x i32> <i32 16, i32 17, i32 2, i32 19, i32 4, i32 5, i32 6, i32 23, i32 24, i32 25, i32 10, i32 27, i32 12, i32 13, i32 14, i32 31>
+; CHECK-NEXT:    ret <16 x i8> [[TMP1]]
+;
   %1 = tail call <16 x i8> @llvm.x86.sse41.pblendvb(<16 x i8> %xyzw, <16 x i8> %abcd, <16 x i8> <i8 0, i8 0, i8 255, i8 0, i8 255, i8 255, i8 255, i8 0, i8 0, i8 0, i8 255, i8 0, i8 255, i8 255, i8 255, i8 0>)
   ret <16 x i8> %1
 }
 
 define <16 x i8> @constant_pblendvb_zero(<16 x i8> %xyzw, <16 x i8> %abcd) {
 ; CHECK-LABEL: @constant_pblendvb_zero
 ; CHECK-NEXT: ret <16 x i8> %xyzw
   %1 = tail call <16 x i8> @llvm.x86.sse41.pblendvb(<16 x i8> %xyzw, <16 x i8> %abcd, <16 x i8> zeroinitializer)
   ret <16 x i8> %1
 }
 
 define <16 x i8> @constant_pblendvb_dup(<16 x i8> %xyzw, <16 x i8> %sel) {
 ; CHECK-LABEL: @constant_pblendvb_dup
 ; CHECK-NEXT: ret <16 x i8> %xyzw
   %1 = tail call <16 x i8> @llvm.x86.sse41.pblendvb(<16 x i8> %xyzw, <16 x i8> %xyzw, <16 x i8> %sel)
   ret <16 x i8> %1
 }
 
 define <4 x double> @constant_blendvpd_avx(<4 x double> %xy, <4 x double> %ab) {
-; CHECK-LABEL: @constant_blendvpd_avx
-; CHECK-NEXT: %1 = select <4 x i1> <i1 true, i1 false, i1 true, i1 false>, <4 x double> %ab, <4 x double> %xy
-; CHECK-NEXT: ret <4 x double> %1
+; CHECK-LABEL: @constant_blendvpd_avx(
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x double> %ab, <4 x double> %xy, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    ret <4 x double> [[TMP1]]
+;
   %1 = tail call <4 x double> @llvm.x86.avx.blendv.pd.256(<4 x double> %xy, <4 x double> %ab, <4 x double> <double 0xFFFFFFFFE0000000, double 0.000000e+00, double 0xFFFFFFFFE0000000, double 0.000000e+00>)
   ret <4 x double> %1
 }
 
 define <4 x double> @constant_blendvpd_avx_zero(<4 x double> %xy, <4 x double> %ab) {
 ; CHECK-LABEL: @constant_blendvpd_avx_zero
 ; CHECK-NEXT: ret <4 x double> %xy
   %1 = tail call <4 x double> @llvm.x86.avx.blendv.pd.256(<4 x double> %xy, <4 x double> %ab, <4 x double> zeroinitializer)
   ret <4 x double> %1
 }
 
 define <4 x double> @constant_blendvpd_avx_dup(<4 x double> %xy, <4 x double> %sel) {
 ; CHECK-LABEL: @constant_blendvpd_avx_dup
 ; CHECK-NEXT: ret <4 x double> %xy
   %1 = tail call <4 x double> @llvm.x86.avx.blendv.pd.256(<4 x double> %xy, <4 x double> %xy, <4 x double> %sel)
   ret <4 x double> %1
 }
 
 define <8 x float> @constant_blendvps_avx(<8 x float> %xyzw, <8 x float> %abcd) {
-; CHECK-LABEL: @constant_blendvps_avx
-; CHECK-NEXT: %1 = select <8 x i1> <i1 false, i1 false, i1 false, i1 true, i1 false, i1 false, i1 false, i1 true>, <8 x float> %abcd, <8 x float> %xyzw
-; CHECK-NEXT: ret <8 x float> %1
+; CHECK-LABEL: @constant_blendvps_avx(
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <8 x float> %abcd, <8 x float> %xyzw, <8 x i32> <i32 8, i32 9, i32 10, i32 3, i32 12, i32 13, i32 14, i32 7>
+; CHECK-NEXT:    ret <8 x float> [[TMP1]]
+;
   %1 = tail call <8 x float> @llvm.x86.avx.blendv.ps.256(<8 x float> %xyzw, <8 x float> %abcd, <8 x float> <float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0xFFFFFFFFE0000000, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0xFFFFFFFFE0000000>)
   ret <8 x float> %1
 }
 
 define <8 x float> @constant_blendvps_avx_zero(<8 x float> %xyzw, <8 x float> %abcd) {
 ; CHECK-LABEL: @constant_blendvps_avx_zero
 ; CHECK-NEXT: ret <8 x float> %xyzw
   %1 = tail call <8 x float> @llvm.x86.avx.blendv.ps.256(<8 x float> %xyzw, <8 x float> %abcd, <8 x float> zeroinitializer)
   ret <8 x float> %1
 }
 
 define <8 x float> @constant_blendvps_avx_dup(<8 x float> %xyzw, <8 x float> %sel) {
 ; CHECK-LABEL: @constant_blendvps_avx_dup
 ; CHECK-NEXT: ret <8 x float> %xyzw
   %1 = tail call <8 x float> @llvm.x86.avx.blendv.ps.256(<8 x float> %xyzw, <8 x float> %xyzw, <8 x float> %sel)
   ret <8 x float> %1
 }
 
 define <32 x i8> @constant_pblendvb_avx2(<32 x i8> %xyzw, <32 x i8> %abcd) {
-; CHECK-LABEL: @constant_pblendvb_avx2
-; CHECK-NEXT: %1 = select <32 x i1> <i1 false, i1 false, i1 true, i1 false, i1 true, i1 true, i1 true, i1 false, i1 false, i1 false, i1 true, i1 false, i1 true, i1 true, i1 true, i1 false, i1 false, i1 false, i1 true, i1 false, i1 true, i1 true, i1 true, i1 false, i1 false, i1 false, i1 true, i1 false, i1 true, i1 true, i1 true, i1 false>, <32 x i8> %abcd, <32 x i8> %xyzw
-; CHECK-NEXT: ret <32 x i8> %1
+; CHECK-LABEL: @constant_pblendvb_avx2(
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <32 x i8> %abcd, <32 x i8> %xyzw, <32 x i32> <i32 32, i32 33, i32 2, i32 35, i32 4, i32 5, i32 6, i32 39, i32 40, i32 41, i32 10, i32 43, i32 12, i32 13, i32 14, i32 47, i32 48, i32 49, i32 18, i32 51, i32 20, i32 21, i32 22, i32 55, i32 56, i32 57, i32 26, i32 59, i32 28, i32 29, i32 30, i32 63>
+; CHECK-NEXT:    ret <32 x i8> [[TMP1]]
+;
   %1 = tail call <32 x i8> @llvm.x86.avx2.pblendvb(<32 x i8> %xyzw, <32 x i8> %abcd,
         <32 x i8> <i8 0, i8 0, i8 255, i8 0, i8 255, i8 255, i8 255, i8 0,
                    i8 0, i8 0, i8 255, i8 0, i8 255, i8 255, i8 255, i8 0,
                    i8 0, i8 0, i8 255, i8 0, i8 255, i8 255, i8 255, i8 0,
                    i8 0, i8 0, i8 255, i8 0, i8 255, i8 255, i8 255, i8 0>)
   ret <32 x i8> %1
 }
 

test/Transforms/InstCombine/logical-select.ll

;
   %bc1 = bitcast <4 x i1> %not to i4
   %bc2 = bitcast <4 x i1> %c to i4
   %and1 = and i4 %a, %bc1
   %and2 = and i4 %bc2, %b
   %or = or i4 %and1, %and2
   ret i4 %or
 }
 
-; Inverted 'and' constants mean this is a select.
+; Inverted 'and' constants mean this is a select which is canonicalized to a shuffle.
 
 define <4 x i32> @vec_sel_consts(<4 x i32> %a, <4 x i32> %b) {
 ; CHECK-LABEL: @vec_sel_consts(
-; CHECK-NEXT:    [[TMP1:%.*]] = select <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x i32> %a, <4 x i32> %b
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> %a, <4 x i32> %b, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
 ; CHECK-NEXT:    ret <4 x i32> [[TMP1]]
 ;
   %and1 = and <4 x i32> %a, <i32 -1, i32 0, i32 0, i32 -1>
   %and2 = and <4 x i32> %b, <i32 0, i32 -1, i32 -1, i32 0>
   %or = or <4 x i32> %and1, %and2
   ret <4 x i32> %or
 }
 
-; The select condition constant is always derived from the first operand of the 'or'.
-
 define <3 x i129> @vec_sel_consts_weird(<3 x i129> %a, <3 x i129> %b) {
 ; CHECK-LABEL: @vec_sel_consts_weird(
-; CHECK-NEXT:    [[TMP1:%.*]] = select <3 x i1> <i1 false, i1 true, i1 false>, <3 x i129> %b, <3 x i129> %a
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <3 x i129> %b, <3 x i129> %a, <3 x i32> <i32 3, i32 1, i32 5>
 ; CHECK-NEXT:    ret <3 x i129> [[TMP1]]
 ;
   %and1 = and <3 x i129> %a, <i129 -1, i129 0, i129 -1>
   %and2 = and <3 x i129> %b, <i129 0, i129 -1, i129 0>
   %or = or <3 x i129> %and2, %and1
   ret <3 x i129> %or
 }
 

test/Transforms/InstCombine/select.ll

 ; CHECK-NEXT:    ret <2 x i32> [[TMP1]]
 ;
   %cmp = icmp slt <2 x i32> %x, zeroinitializer
   %xor = xor <2 x i32> %x, <i32 2147483648, i32 2147483648>
   %x.xor = select <2 x i1> %cmp, <2 x i32> %x, <2 x i32> %xor
   ret <2 x i32> %x.xor
 }
 
+; Make sure that undef elements of the select condition are
+; translated to undef elements of the shuffle mask.
+
+define <4 x i32> @canonicalize_to_shuffle(<4 x i32> %a, <4 x i32> %b) {
+; CHECK-LABEL: @canonicalize_to_shuffle(
+; CHECK-NEXT:    [[SEL:%.*]] = shufflevector <4 x i32> %a, <4 x i32> %b, <4 x i32> <i32 0, i32 undef, i32 6, i32 undef>
+; CHECK-NEXT:    ret <4 x i32> [[SEL]]
+;
+  %sel = select <4 x i1> <i1 true, i1 undef, i1 false, i1 undef>, <4 x i32> %a, <4 x i32> %b
+  ret <4 x i32> %sel
+}
+

test/Transforms/InstCombine/vec_demanded_elts.ll

;
   %ret = shufflevector <4 x double> %tmp5, <4 x double> undef, <2 x i32> <i32 0, i32 1>
   ret <2 x double> %ret
 }
 
 define <4 x float> @test_select(float %f, float %g) {
 ; CHECK-LABEL: @test_select(
 ; CHECK-NEXT:    [[A0:%.*]] = insertelement <4 x float> undef, float %f, i32 0
 ; CHECK-NEXT:    [[A3:%.*]] = insertelement <4 x float> [[A0]], float 3.000000e+00, i32 3
-; CHECK-NEXT:    [[RET:%.*]] = select <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x float> [[A3]], <4 x float> <float undef, float 4.000000e+00, float 5.000000e+00, float undef>
+; CHECK-NEXT:    [[RET:%.*]] = shufflevector <4 x float> [[A3]], <4 x float> <float undef, float 4.000000e+00, float 5.000000e+00, float undef>, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
 ; CHECK-NEXT:    ret <4 x float> [[RET]]
 ;
   %a0 = insertelement <4 x float> undef, float %f, i32 0
   %a1 = insertelement <4 x float> %a0, float 1.000000e+00, i32 1
   %a2 = insertelement <4 x float> %a1, float 2.000000e+00, i32 2
   %a3 = insertelement <4 x float> %a2, float 3.000000e+00, i32 3
   %b0 = insertelement <4 x float> undef, float %g, i32 0
   %b1 = insertelement <4 x float> %b0, float 4.000000e+00, i32 1
   %b2 = insertelement <4 x float> %b1, float 5.000000e+00, i32 2
   %b3 = insertelement <4 x float> %b2, float 6.000000e+00, i32 3
   %ret = select <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x float> %a3, <4 x float> %b3
   ret <4 x float> %ret
 }
 
-; Check that instcombine doesn't wrongly fold the select statement into a ret <2 x i64> %v
+; Check that instcombine doesn't wrongly fold away the select completely.
+; TODO: Should this be an insertelement rather than a shuffle?
+
 define <2 x i64> @PR24922(<2 x i64> %v) {
 ; CHECK-LABEL: @PR24922(
-; CHECK-NEXT:    [[RESULT:%.*]] = select <2 x i1> <i1 false, i1 true>, <2 x i64> %v, <2 x i64> <i64 0, i64 undef>
+; CHECK-NEXT:    [[RESULT:%.*]] = shufflevector <2 x i64> %v, <2 x i64> <i64 0, i64 undef>, <2 x i32> <i32 2, i32 1>
 ; CHECK-NEXT:    ret <2 x i64> [[RESULT]]
 ;
   %result = select <2 x i1> <i1 icmp eq (i64 extractelement (<2 x i64> bitcast (<4 x i32> <i32 15, i32 15, i32 15, i32 15> to <2 x i64>), i64 0), i64 0), i1 true>, <2 x i64> %v, <2 x i64> zeroinitializer
   ret <2 x i64> %result
 }