Fold selects when the non-selected elements are undef from shufflevectors

[filcab]

Summary

If we have a select whose operands are two shufflevectors and we know the
non-selected elements are undef, replace the instruction with a
shufflevector from the source vectors of the shufflevectors.

Added test cases.

[bkramer]

lib/Transforms/InstCombine/InstCombineSelect.cpp

739

IMO DeMorgan's law makes this kind of condition more readable

if (!VecTy || !TrueSV || !FalseSV)

743

Why is it important that SI has only one use?

756–758

Since you never use any VectorType methods you can skip the cast<>s here.

778

If it's not a ConstantInt it could also be a ConstantExpr instead of an Undef. You probably want to bail out in that case.

[filcab]

Fixed bugs pointed by Benjamin Kramer.

[filcab]

Remove files committed by accident.

[bkramer]

lib/Transforms/InstCombine/InstCombineSelect.cpp

734

Argument "SI" is now unused :)

739

s/&&/||/g

750

I'd prefer

if (!isa<UndefValue>(TrueV2) || !isa<UndefValue>(FalseV2))

test/Transforms/InstCombine/select.ll

1037

Can we have a more positive test case? The mask generation code isn't entirely trivial, would be good to have it verified by a test.

[filcab]

More fixes pointed by Benjamin Kramer.
Make the test-cases more strict w.r. to shufflevector mask.

[filcab]

Remove unused argument

[bkramer]

This LGTM now.

[filcab]

Per discussion on the mailing list (e. g: http://lists.cs.uiuc.edu/pipermail/llvm-commits/Week-of-Mon-20140505/216033.html and following messages), I'm abandoning this revision.

lib/Transforms/InstCombine/InstCombineSelect.cpp

if (ShouldNotVal)
     V = Builder->CreateXor(V, ValC);
 
   // Apply an offset if needed.
   if (Offset)
     V = Builder->CreateAdd(V, Offset);
   return V;
 }
 
+static bool CanFoldSelectOfShuffleVectors(Value *CondVal, Value *TrueVal,
+                                          Value *FalseVal) {
+  VectorType *VecTy = dyn_cast<VectorType>(CondVal->getType());
+  ShuffleVectorInst *TrueSV = dyn_cast<ShuffleVectorInst>(TrueVal);
+  ShuffleVectorInst *FalseSV = dyn_cast<ShuffleVectorInst>(FalseVal);
+  if (!VecTy || !TrueSV || !FalseSV)
+    return false;
+
+  ConstantVector *CondV = dyn_cast<ConstantVector>(CondVal);
+  if (!CondV)
+    return false;
+
+  Value *TrueV2 = TrueSV->getOperand(1);
+  Value *FalseV2 = FalseSV->getOperand(1);
+  // We just check for *V2 being undef since instcombine will turn
+  // shufflevector(undef, v) into shufflevector(v, undef)
+  if (!isa<UndefValue>(TrueV2) || !isa<UndefValue>(FalseV2))
+    return false;
+
+  // The source vectors (not the mask) for the shuffle vectors have to
+  // have the same type. Otherwise we could end up trying to do a
+  // shufflevector <4 x i32> <2 x i32>
+  Type *TrueShuffleSrc = TrueV2->getType();
+  Type *FalseShuffleSrc = FalseV2->getType();
+  if (TrueShuffleSrc != FalseShuffleSrc)
+    return false;
+
+  return true;
+}
+
+// This is instruction is only safe to call if CanFoldSelectOfShuffleVectors is
+// true.
+static Instruction *FoldSelectOfShuffleVectors(SelectInst &SI) {
+  SmallVector<uint32_t, 16> ShuffleMask;
+  ShuffleVectorInst *Sources[] = { cast<ShuffleVectorInst>(SI.getFalseValue()),
+                                   cast<ShuffleVectorInst>(SI.getTrueValue()) };
+  ConstantVector *CondV = cast<ConstantVector>(SI.getCondition());
+
+  unsigned NumElems = cast<VectorType>(SI.getType())->getNumElements();
+  unsigned NumSourceElems =
+      cast<VectorType>(Sources[0]->getOperand(0)->getType())->getNumElements();
+  for (unsigned i = 0; i < NumElems; ++i) {
+    ConstantInt *Element = dyn_cast<ConstantInt>(CondV->getAggregateElement(i));
+    if (!Element)
+      return nullptr;
+
+    int Selector = Element->isOne();
+    int SourceIdx = Sources[Selector]->getMaskValue(i);
+    ShuffleMask.push_back(SourceIdx == -1 ? -1 : SourceIdx +
+                                                     NumSourceElems * Selector);
+  }
+
+  return new ShuffleVectorInst(
+      Sources[0]->getOperand(0), Sources[1]->getOperand(0),
+      ConstantDataVector::get(SI.getContext(), ShuffleMask));
+}
+
 Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();
 
   if (Value *V = SimplifySelectInst(CondVal, TrueVal, FalseVal, DL))
     return ReplaceInstUsesWith(SI, V);
 
if (Value *V = SimplifyDemandedVectorElts(&SI, AllOnesEltMask, UndefElts)) {
       if (V != &SI)
         return ReplaceInstUsesWith(SI, V);
       return &SI;
     }
 
     if (isa<ConstantAggregateZero>(CondVal)) {
       return ReplaceInstUsesWith(SI, FalseVal);
     }
+
+    // If all the elements with true in the mask correspond to undef on
+    // FalseV, and the reverse is true with false in the mask and TrueV, we
+    // can merge the shufflevectors and remove the select.
+    if (CanFoldSelectOfShuffleVectors(SI, CondVal, TrueVal, FalseVal))
+      if (Instruction *I = FoldSelectOfShuffleVectors(SI))
+        return I;
   }
 
   return nullptr;
 }

test/Transforms/InstCombine/select.ll

define i32 @test67(i16 %x) {
   ret i32 %3
 
 ; CHECK-LABEL: @test67(
 ; CHECK: and i16 %x, 4
 ; CHECK: zext i16 %1 to i32
 ; CHECK: lshr exact i32 %2, 1
 ; CHECK: xor i32 %3, 42
 }
+
+define <4 x float> @add2f_0(<2 x float> %a0, <2 x float> %b0, <2 x float> %a1, <2 x float> %b1) {
+; CHECK-LABEL: @add2f_0
+; CHECK-NOT: select
+; CHECK: shufflevector <2 x float> %a0, <2 x float> %a1, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK: shufflevector <2 x float> %b0, <2 x float> %b1, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK: ret
+  %1 = shufflevector <2 x float> %a0, <2 x float> undef, <4 x i32> <i32 0, i32 1, i32 undef, i32 undef>
+  %2 = shufflevector <2 x float> %a1, <2 x float> undef, <4 x i32> <i32 undef, i32 undef, i32 0, i32 1>
+  %3 = select <4 x i1> <i1 false, i1 false, i1 true, i1 true>, <4 x float> %2, <4 x float> %1
+  %4 = shufflevector <2 x float> %b0, <2 x float> undef, <4 x i32> <i32 0, i32 1, i32 undef, i32 undef>
+  %5 = shufflevector <2 x float> %b1, <2 x float> undef, <4 x i32> <i32 undef, i32 undef, i32 0, i32 1>
+  %6 = select <4 x i1> <i1 false, i1 false, i1 true, i1 true>, <4 x float> %5, <4 x float> %4
+  %7 = fadd <4 x float> %3, %6
+  ret <4 x float> %7
+}
+
+;; This test might need to change if we implement the transform
+;; (shuffle (add x y) (add z w)) -> (add (shuffle x z) (shuffle y w))
+define <4 x float> @add2f_1(<2 x float> %a0, <2 x float> %b0, <2 x float> %a1, <2 x float> %b1) {
+; CHECK-LABEL: @add2f_1
+; CHECK-NOT: select
+; CHECK: [[ADD0:%[a-z0-9]+]] = fadd <2 x float> %a0, %b0
+; CHECK: [[ADD1:%[a-z0-9]+]] = fadd <2 x float> %a1, %b1
+; CHECK: [[RES:%[a-z0-9]+]] = shufflevector <2 x float> [[ADD0]], <2 x float> [[ADD1]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK: ret
+  %1 = fadd <2 x float> %a0, %b0
+  %2 = shufflevector <2 x float> %1, <2 x float> undef, <4 x i32> <i32 0, i32 1, i32 undef, i32 undef>
+  %3 = fadd <2 x float> %a1, %b1
+  %4 = shufflevector <2 x float> %3, <2 x float> undef, <4 x i32> <i32 undef, i32 undef, i32 0, i32 1>
+  %5 = select <4 x i1> <i1 false, i1 false, i1 true, i1 true>, <4 x float> %4, <4 x float> %2
+  ret <4 x float> %5
+}