[InstCombine] move vector select ahead of select-shuffle

[spatel]

Summary

select Cond, (shuf_sel X, Y), X --> shuf_sel X, (select Cond, Y, X)

A select of a select-shuffle ("blend" in x86 lingo) can be reversed so that the select is done first.
This is a more limited version of what I was trying in D80658, but it enables existing demanded bits transforms to catch some of the motivating cases. The tricky bit in that seems to be that by moving the shuffle later, we can always guarantee that poison is correctly inhibited by the shuffle mask in the final value.

Alive2 checks for the basic tests:
http://volta.cs.utah.edu:8080/z/Qqd3RK
http://volta.cs.utah.edu:8080/z/S4wchM
http://volta.cs.utah.edu:8080/z/wf9zPL
http://volta.cs.utah.edu:8080/z/wJeEGk

[lebedev.ri]

LG, alive2 is happy with test changes.

[srj]

This seems to have injected breakages into Halide when used for x86-64 targets with 'narrow' SIMD (eg., using only SSE2 or SSE4.1); many of our tests now fail in that configuration with variations of

/path/tollvm-project/llvm/lib/IR/Instructions.cpp:2049: bool isSingleSourceMaskImpl(llvm::ArrayRef<int>, int): Assertion `Mask[i] >= 0 && Mask[i] < (NumOpElts * 2) && "Out-of-bounds shuffle mask element"' failed.

[lebedev.ri]

In D81013#2074535, @srj wrote:

This seems to have injected breakages into Halide when used for x86-64 targets with 'narrow' SIMD (eg., using only SSE2 or SSE4.1); many of our tests now fail in that configuration with variations of

/path/tollvm-project/llvm/lib/IR/Instructions.cpp:2049: bool isSingleSourceMaskImpl(llvm::ArrayRef<int>, int): Assertion `Mask[i] >= 0 && Mask[i] < (NumOpElts * 2) && "Out-of-bounds shuffle mask element"' failed.

Looks like we are missing !increasesLength() check somewhere.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2416–2417

The problem is that we call ShuffleVectorInst::isSelectMask() before we ensure that select and shuffle have common operand, so by now types of their ops might be non-matching, and shuffle is lenght-changing.

[lebedev.ri]

That being said, i do not understand how we can end up with non--1 widening element,
https://godbolt.org/z/nBhmko <- impossible as per ShuffleVectorInst::isValidOperands().

So i'd like to see a reproducer.

[lebedev.ri]

Ah, i guess i see it.
The problem is that ShuffleVectorInst::isSingleSourceMask(ArrayRef<int> Mask) is kinda broken:
it can't tell isSingleSourceMaskImpl() how many src elts there are,
so naturally it will immediately assert as long as it selects from non-first source..

[spatel]

In D81013#2074708, @lebedev.ri wrote:

Ah, i guess i see it.
The problem is that ShuffleVectorInst::isSingleSourceMask(ArrayRef<int> Mask) is kinda broken:
it can't tell isSingleSourceMaskImpl() how many src elts there are,
so naturally it will immediately assert as long as it selects from non-first source..

Would this be fixed with "cast<ShuffleVectorInst>(TVal)->isSelect()" instead of the mask check?
Let me know if I should revert or if we have a test case to work on.

[lebedev.ri]

In D81013#2074721, @spatel wrote:

In D81013#2074708, @lebedev.ri wrote:

Ah, i guess i see it.
The problem is that ShuffleVectorInst::isSingleSourceMask(ArrayRef<int> Mask) is kinda broken:
it can't tell isSingleSourceMaskImpl() how many src elts there are,
so naturally it will immediately assert as long as it selects from non-first source..

Would this be fixed with "cast<ShuffleVectorInst>(TVal)->isSelect()" instead of the mask check?

I think so.

Let me know if I should revert or if we have a test case to work on.

This should do the trick, but my opt is out-of-date to check

define <5 x i8> @widening(<4 x i8> %x, <4 x i8> %y, <5 x i8> %x2, <5 x i1> %cmp) {
  %blend = shufflevector <4 x i8> %x, <4 x i8> %y, <5 x i32> <i32 0, i32 5, i32 2, i32 7, i32 7>
  %r = select <5 x i1> %cmp, <5 x i8> %blend, <5 x i8> %x2
  ret <5 x i8> %r
}

[spatel]

In D81013#2074725, @lebedev.ri wrote:

This should do the trick, but my opt is out-of-date to check

define <5 x i8> @widening(<4 x i8> %x, <4 x i8> %y, <5 x i8> %x2, <5 x i1> %cmp) {
  %blend = shufflevector <4 x i8> %x, <4 x i8> %y, <5 x i32> <i32 0, i32 5, i32 2, i32 7, i32 7>
  %r = select <5 x i1> %cmp, <5 x i8> %blend, <5 x i8> %x2
  ret <5 x i8> %r
}

I couldn't find a widening version that fails, but the narrowing variant does it:
rG192cb718361dbd7be082bc0893f43bbc9782288f

@srj - sorry about the crashing. Let me know if there are still problems.

[srj]

@spatel thanks for the quick fix!

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

Instruction *InstCombiner::foldVectorSelect(SelectInst &Sel) {
   APInt UndefElts(NumElts, 0);
   APInt AllOnesEltMask(APInt::getAllOnesValue(NumElts));
   if (Value *V = SimplifyDemandedVectorElts(&Sel, AllOnesEltMask, UndefElts)) {
     if (V != &Sel)
       return replaceInstUsesWith(Sel, V);
     return &Sel;
   }
 
+  // A select of a "select shuffle" with a common operand can be rearranged
+  // to select followed by "select shuffle". Because of poison, this only works
+  // in the case of a shuffle with no undefined mask elements.
+  Value *Cond = Sel.getCondition();
+  Value *TVal = Sel.getTrueValue();
+  Value *FVal = Sel.getFalseValue();
+  Value *X, *Y;
+  ArrayRef<int> Mask;
+  if (match(TVal, m_OneUse(m_Shuffle(m_Value(X), m_Value(Y), m_Mask(Mask)))) &&
+      !is_contained(Mask, UndefMaskElem) &&
+      ShuffleVectorInst::isSelectMask(Mask)) {
+    if (X == FVal) {
+      // select Cond, (shuf_sel X, Y), X --> shuf_sel X, (select Cond, Y, X)
+      Value *NewSel = Builder.CreateSelect(Cond, Y, X, "sel", &Sel);
+      return new ShuffleVectorInst(X, NewSel, Mask);
+    }
+    if (Y == FVal) {
+      // select Cond, (shuf_sel X, Y), Y --> shuf_sel (select Cond, X, Y), Y
+      Value *NewSel = Builder.CreateSelect(Cond, X, Y, "sel", &Sel);
+      return new ShuffleVectorInst(NewSel, Y, Mask);
+    }
+  }
+  if (match(FVal, m_OneUse(m_Shuffle(m_Value(X), m_Value(Y), m_Mask(Mask)))) &&
+      !is_contained(Mask, UndefMaskElem) &&
+      ShuffleVectorInst::isSelectMask(Mask)) {
+    if (X == TVal) {
+      // select Cond, X, (shuf_sel X, Y) --> shuf_sel X, (select Cond, X, Y)
+      Value *NewSel = Builder.CreateSelect(Cond, X, Y, "sel", &Sel);
+      return new ShuffleVectorInst(X, NewSel, Mask);
+    }
+    if (Y == TVal) {
+      // select Cond, Y, (shuf_sel X, Y) --> shuf_sel (select Cond, Y, X), Y
+      Value *NewSel = Builder.CreateSelect(Cond, Y, X, "sel", &Sel);
+      return new ShuffleVectorInst(NewSel, Y, Mask);
+    }
+  }
+
   return nullptr;
 }
 
 Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();
   Type *SelType = SI.getType();

llvm/test/Transforms/InstCombine/select-select.ll

 ; CHECK-NEXT:    [[SEL1:%.*]] = select i1 [[BOOL:%.*]], <2 x i32> [[SEL0]], <2 x i32> [[V]]
 ; CHECK-NEXT:    ret <2 x i32> [[SEL1]]
 ;
   %sel0 = select <2 x i1> %vec_bool, <2 x i32> zeroinitializer, <2 x i32> %V
   %sel1 = select i1 %bool, <2 x i32> %sel0, <2 x i32> %V
   ret <2 x i32> %sel1
 }
 
+; Four variations of (select (select-shuffle)) with a common operand.
+
 define <4 x i8> @sel_shuf_commute0(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @sel_shuf_commute0(
-; CHECK-NEXT:    [[BLEND:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 2, i32 7>
-; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x i8> [[BLEND]], <4 x i8> [[X]]
+; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x i8> [[Y:%.*]], <4 x i8> [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x i8> [[X]], <4 x i8> [[SEL]], <4 x i32> <i32 0, i32 5, i32 2, i32 7>
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %blend = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
   %r = select <4 x i1> %cmp, <4 x i8> %blend, <4 x i8> %x
   ret <4 x i8> %r
 }
 
+; Weird types are ok.
+
 define <5 x i9> @sel_shuf_commute1(<5 x i9> %x, <5 x i9> %y, <5 x i1> %cmp) {
 ; CHECK-LABEL: @sel_shuf_commute1(
-; CHECK-NEXT:    [[BLEND:%.*]] = shufflevector <5 x i9> [[X:%.*]], <5 x i9> [[Y:%.*]], <5 x i32> <i32 0, i32 6, i32 2, i32 8, i32 9>
-; CHECK-NEXT:    [[R:%.*]] = select <5 x i1> [[CMP:%.*]], <5 x i9> [[BLEND]], <5 x i9> [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select <5 x i1> [[CMP:%.*]], <5 x i9> [[X:%.*]], <5 x i9> [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <5 x i9> [[SEL]], <5 x i9> [[Y]], <5 x i32> <i32 0, i32 6, i32 2, i32 8, i32 9>
 ; CHECK-NEXT:    ret <5 x i9> [[R]]
 ;
   %blend = shufflevector <5 x i9> %x, <5 x i9> %y, <5 x i32> <i32 0, i32 6, i32 2, i32 8, i32 9>
   %r = select <5 x i1> %cmp, <5 x i9> %blend, <5 x i9> %y
   ret <5 x i9> %r
 }
 
 define <4 x float> @sel_shuf_commute2(<4 x float> %x, <4 x float> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @sel_shuf_commute2(
-; CHECK-NEXT:    [[BLEND:%.*]] = shufflevector <4 x float> [[X:%.*]], <4 x float> [[Y:%.*]], <4 x i32> <i32 0, i32 1, i32 2, i32 7>
-; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x float> [[X]], <4 x float> [[BLEND]]
+; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x float> [[X:%.*]], <4 x float> [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x float> [[X]], <4 x float> [[SEL]], <4 x i32> <i32 0, i32 1, i32 2, i32 7>
 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %blend = shufflevector <4 x float> %x, <4 x float> %y, <4 x i32> <i32 0, i32 1, i32 2, i32 7>
   %r = select <4 x i1> %cmp, <4 x float> %x, <4 x float> %blend
   ret <4 x float> %r
 }
 
+; Scalar condition is ok.
+
 define <4 x i8> @sel_shuf_commute3(<4 x i8> %x, <4 x i8> %y, i1 %cmp) {
 ; CHECK-LABEL: @sel_shuf_commute3(
-; CHECK-NEXT:    [[BLEND:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 2, i32 3>
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[CMP:%.*]], <4 x i8> [[Y]], <4 x i8> [[BLEND]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP:%.*]], <4 x i8> [[Y:%.*]], <4 x i8> [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x i8> [[SEL]], <4 x i8> [[Y]], <4 x i32> <i32 0, i32 5, i32 2, i32 3>
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %blend = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 2, i32 3>
   %r = select i1 %cmp, <4 x i8> %y, <4 x i8> %blend
   ret <4 x i8> %r
 }
 
 declare void @use(<4 x i8>)
 
+; Negative test - extra use would require another instruction.
+
 define <4 x i8> @sel_shuf_use(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @sel_shuf_use(
 ; CHECK-NEXT:    [[BLEND:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 2, i32 7>
 ; CHECK-NEXT:    call void @use(<4 x i8> [[BLEND]])
 ; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x i8> [[BLEND]], <4 x i8> [[X]]
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %blend = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
   call void @use(<4 x i8> %blend)
   %r = select <4 x i1> %cmp, <4 x i8> %blend, <4 x i8> %x
   ret <4 x i8> %r
 }
 
+; Negative test - undef in shuffle mask prevents transform.
+
 define <4 x i8> @sel_shuf_undef(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @sel_shuf_undef(
 ; CHECK-NEXT:    [[BLEND:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 2, i32 undef>
 ; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x i8> [[BLEND]], <4 x i8> [[Y]]
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %blend = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 2, i32 undef>
   %r = select <4 x i1> %cmp, <4 x i8> %blend, <4 x i8> %y
   ret <4 x i8> %r
 }
 
+; Negative test - not a "select shuffle"
+
 define <4 x i8> @sel_shuf_not(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @sel_shuf_not(
 ; CHECK-NEXT:    [[NOTBLEND:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 2, i32 6>
 ; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x i8> [[NOTBLEND]], <4 x i8> [[Y]]
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %notblend = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 2, i32 6>
   %r = select <4 x i1> %cmp, <4 x i8> %notblend, <4 x i8> %y
   ret <4 x i8> %r
 }
 
+; Negative test - must shuffle one of the select operands
+
 define <4 x i8> @sel_shuf_no_common_operand(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp, <4 x i8> %z) {
 ; CHECK-LABEL: @sel_shuf_no_common_operand(
 ; CHECK-NEXT:    [[BLEND:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 2, i32 3>
 ; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x i8> [[Z:%.*]], <4 x i8> [[BLEND]]
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %blend = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 2, i32 3>
   %r = select <4 x i1> %cmp, <4 x i8> %z, <4 x i8> %blend
   ret <4 x i8> %r
 }

llvm/test/Transforms/InstCombine/vec_demanded_elts.ll

bb:
   %t2 = extractelement <4 x i32*> %t, i64 0
   ret i32* %t2
 }
 
 ; The non-zero elements of the result are always 'y', so the splat is unnecessary.
 
 define <4 x i8> @select_cond_with_eq_true_false_elts(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @select_cond_with_eq_true_false_elts(
-; CHECK-NEXT:    [[TVAL:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[SPLAT]], <4 x i8> [[TVAL]], <4 x i8> [[Y]]
+; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x i8> [[SEL]], <4 x i8> [[Y]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %tval = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 6, i32 7>
   %splat = shufflevector <4 x i1> %cmp, <4 x i1> undef, <4 x i32> zeroinitializer
   %r = select <4 x i1> %splat, <4 x i8> %tval, <4 x i8> %y
   ret <4 x i8> %r
 }
 
 ; First element of the result is always x[0], so first element of select condition is unnecessary.
 
 define <4 x i8> @select_cond_with_eq_true_false_elts2(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @select_cond_with_eq_true_false_elts2(
-; CHECK-NEXT:    [[TVAL:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
-; CHECK-NEXT:    [[COND:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
-; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[COND]], <4 x i8> [[TVAL]], <4 x i8> [[X]]
+; CHECK-NEXT:    [[COND:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <4 x i32> <i32 undef, i32 1, i32 0, i32 1>
+; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> [[COND]], <4 x i8> [[Y:%.*]], <4 x i8> [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x i8> [[X]], <4 x i8> [[SEL]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %tval = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 6, i32 7>
   %cond = shufflevector <4 x i1> %cmp, <4 x i1> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
   %r = select <4 x i1> %cond, <4 x i8> %tval, <4 x i8> %x
   ret <4 x i8> %r
 }
 

 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %tval = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 undef, i32 5, i32 6, i32 7>
   %cond = shufflevector <4 x i1> %cmp, <4 x i1> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
   %r = select <4 x i1> %cond, <4 x i8> %tval, <4 x i8> %x
   ret <4 x i8> %r
 }
 
-; The insert can not be safely eliminated because cmp[0] might be poison.
+; The insert can be safely eliminated because the shuffle blocks poison from cmp[0].
 
 define <4 x i8> @select_cond_(<4 x i8> %x, <4 x i8> %min, <4 x i1> %cmp, i1 %poison_blocker) {
 ; CHECK-LABEL: @select_cond_(
-; CHECK-NEXT:    [[INS:%.*]] = insertelement <4 x i1> [[CMP:%.*]], i1 [[POISON_BLOCKER:%.*]], i32 0
-; CHECK-NEXT:    [[VECINS:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[MIN:%.*]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
-; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[INS]], <4 x i8> [[VECINS]], <4 x i8> [[X]]
+; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x i8> [[MIN:%.*]], <4 x i8> [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x i8> [[X]], <4 x i8> [[SEL]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %ins = insertelement <4 x i1> %cmp, i1 %poison_blocker, i32 0
   %vecins = shufflevector <4 x i8> %x, <4 x i8> %min, <4 x i32> <i32 0, i32 5, i32 6, i32 7>
   %r = select <4 x i1> %ins, <4 x i8> %vecins, <4 x i8> %x
   ret <4 x i8> %r
 }