[X86][SSE] Detect zeroable shuffle elements from different value types

[RKSimon]

Summary

Improve computeZeroableShuffleElements to be able to peek through bitcasts to extract zero/undef values from BUILD_VECTOR nodes of different element sizes to the shuffle mask.

[andreadb]

Hi Simon,

lib/Target/X86/X86ISelLowering.cpp

6730–6732

Do you think it would make sense to limit the recursion level?

6747–6750

I don't think you need to compute V1IsZero and V2IsZero anymore. The two calls to 'isBuildVectorAllZeros' are now made redundant by the calls to 'GetSubZeroable'.

You can also simplify the if statement at line 6754. In particular, you would only need to check for the presence of an undef index. All the remaining cases should be taken care by the checks after line 6760.

6781–6791

At line 6777, you don't need to check if 'Size < SubSize'. If control reaches line 6777, then Size can never be bigger than or equal to SubSize.

The motivation is that the check for (Size < SubSize) is dominated by the checks for (Size == SubSize) and the check for (Size > SubSize). You can also remove the llvm_unreachable at line 6787 as it is not needed.

[RKSimon]

Added recursion limit and other recommendations from Andrea.

[chandlerc]

I'm really not convinced this is the correct approach.

Instead, I think we should combine SHUFFLE_VECTOR nodes into a single node so that we don't need this kind of recursive logic. If that doesn't work for some reason, I think that needs to be pretty clearly explained.

[RKSimon]

In D14261#289720, @chandlerc wrote:

I'm really not convinced this is the correct approach.

Is it mainly the recursion that is concerning you? There is more that we could do at the DAGCombiner or X86 level to merge shuffles, but it would still be recursion. Other targets have shuffle instructions that set elements to constants (zero and allbits being the most common) as well as as input permutations, but I don't see much that makes use of this at all.

I've been trying to think of ways to canonicalize shuffles with zeros/constants inputs to reduce this but can't think of anything that would really help.

Instead, I think we should combine SHUFFLE_VECTOR nodes into a single node so that we don't need this kind of recursive logic. If that doesn't work for some reason, I think that needs to be pretty clearly explained.

The X86 shuffle combining could do some of this, but it would still be recursive and it would be repeating a lot of what is done at lowering already.

insertps is a curious instruction in that it really takes 3 inputs (well, 2 inputs + zero), which makes combining rather tricky with our existing methods - as another approach I could disable recursion in computeZeroableShuffleElements by default and just have the insertps lowering use it?

[chandlerc]

In D14261#289721, @RKSimon wrote:

In D14261#289720, @chandlerc wrote:

I'm really not convinced this is the correct approach.

Is it mainly the recursion that is concerning you? There is more that we could do at the DAGCombiner or X86 level to merge shuffles, but it would still be recursion.

It's not the same at all though. When we combine N shuffles into 1 shuffle, we can do that in O(N) because we can combine as we go. When we instead recursively match during lowering here we can do O(N^2) work.

Other targets have shuffle instructions that set elements to constants (zero and allbits being the most common) as well as as input permutations, but I don't see much that makes use of this at all.

I'm not really concerned about helping other targets here, I'm concerned about making this use the correct model.

I've been trying to think of ways to canonicalize shuffles with zeros/constants inputs to reduce this but can't think of anything that would really help.

Can you describe why not? I'm not seeing it.

Specifically, I can't come up with any way that we would need to look through more than N-1 VECTOR_SHUFFLE nodes to find N inputs.

Instead, I think we should combine SHUFFLE_VECTOR nodes into a single node so that we don't need this kind of recursive logic. If that doesn't work for some reason, I think that needs to be pretty clearly explained.

The X86 shuffle combining could do some of this, but it would still be recursive and it would be repeating a lot of what is done at lowering already.

As above, the point is to do *less* of this at lowering time. Even if we still need some of this logic during lowering in order to handle iterative decomposition, we should minimize it and always avoid unbounded recursion.

insertps is a curious instruction in that it really takes 3 inputs (well, 2 inputs + zero), which makes combining rather tricky with our existing methods - as another approach I could disable recursion in computeZeroableShuffleElements by default and just have the insertps lowering use it?

Ah, I think this is really the critical thing to realize.

Because of the simplification provided by it, we should essentially re-associate shuffles to cause all constant inputs to come from a single build_vector, and that a build_vector of constants should be the last input shuffled:

For example:

S1 = (shuffle A, (build constant, constant, ...))
S2 = (shuffle B, S1)

should combine to:

S1 = (shuffle A, B)
S2 = (shuffle S1, (build constant, constant, ...))

So that we can always identify constant inputs at the bottom of the shuffle tree. We can still indicate *unused* inputs in S1 with undef shuffle indices.

Does this make sense to you? Is there some *other* thing broken by this style of reassociation?

[RKSimon]

In D14261#289722, @chandlerc wrote:

In D14261#289721, @RKSimon wrote:

insertps is a curious instruction in that it really takes 3 inputs (well, 2 inputs + zero), which makes combining rather tricky with our existing methods - as another approach I could disable recursion in computeZeroableShuffleElements by default and just have the insertps lowering use it?

Ah, I think this is really the critical thing to realize.

Because of the simplification provided by it, we should essentially re-associate shuffles to cause all constant inputs to come from a single build_vector, and that a build_vector of constants should be the last input shuffled:

For example:

S1 = (shuffle A, (build constant, constant, ...))
S2 = (shuffle B, S1)

should combine to:

S1 = (shuffle A, B)
S2 = (shuffle S1, (build constant, constant, ...))

So that we can always identify constant inputs at the bottom of the shuffle tree. We can still indicate *unused* inputs in S1 with undef shuffle indices.

Does this make sense to you? Is there some *other* thing broken by this style of reassociation?

The approach you propose is the canonicalization step I mentioned. My main concern with it is making sure our existing shuffle lowering can be adapted to 'peek' through a bottom shuffle with constants (in this case zeroable vectors) - this will affect insertps and the perm2f128 instructions and possibly others - but its certainly doable, but may involve reordering some of the lowerings so that they are attempted earlier. This will probably affect build_vector lowering as well - but I won't know the scope until I investigate this more.

Do you think this should be limited to the X86? Large parts of the code, such as combining multiple constant build_vector shuffle inputs into one, would almost certainly be useful in the DAGCombiner and I can't think of reasons why pushing down the shuffles with constants would cause any notable regressions in other targets' shuffle codegen.

That would leave this patch without the recursion element - I think reducing it to just allowing it to check for zeroable build_vectors of different element counts would still be useful, I'll see if I can update this patch.

Cheers, Simon.

[qcolombet]

Hi Simon,

Is this patch still alive?
(I.e., does it need review.)

Thanks,
-Quentin

[RKSimon]

It needs redoing (and reducing in scope) now that more of the target shuffle combine work is completed. We should be able to reduce it to a test for zero on BUILD_VECTORs with different number of elements to the mask (same as to what has been done in setTargetShuffleZeroElements).

It just hasn't been high enough on my todo list......

[RKSimon]

As recommended by Chandler, I've adjusted this patch to avoid recursion (we now do this in setTargetShuffleZeroElements for target shuffle combining) and instead focus on improving computeZeroableShuffleElements to be able to peek through bitcasts to extract zero/undef values from BUILD_VECTOR nodes of different element sizes to the shuffle mask.

[chandlerc]

FWIW, this looks really nice now. Nits below, LGTM with suggested changes.

lib/Target/X86/X86ISelLowering.cpp

43–44

Please just use 'int' unless you need modular arithmetic...

45

int here as well.

66–68

int here, and j < Scale.

lib/Target/X86/X86ISelLowering.cpp

 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/WinEHFuncInfo.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Target/TargetOptions.h"
 #include "X86IntrinsicsInfo.h"
 #include <bitset>
 #include <numeric>
 #include <cctype>
 using namespace llvm;
 
 #define DEBUG_TYPE "x86-isel"
 
 STATISTIC(NumTailCalls, "Number of tail calls");
 
 static cl::opt<bool> ExperimentalVectorWideningLegalization(
     "x86-experimental-vector-widening-legalization", cl::init(false),
     cl::desc("Enable an experimental vector type legalization through widening "
              "rather than promotion."),
     cl::Hidden);
 
 X86TargetLowering::X86TargetLowering(const X86TargetMachine &TM,
                                      const X86Subtarget &STI)
     : TargetLowering(TM), Subtarget(&STI) {
   X86ScalarSSEf64 = Subtarget->hasSSE2();
   X86ScalarSSEf32 = Subtarget->hasSSE1();
   MVT PtrVT = MVT::getIntegerVT(8 * TM.getPointerSize());

 ///
 /// A "zeroable" vector shuffle element is one which can be lowered to zero.
 /// Either it is an undef element in the shuffle mask, the element of the input
 /// referenced is undef, or the element of the input referenced is known to be
 /// zero. Many x86 shuffles can zero lanes cheaply and we often want to handle
 /// as many lanes with this technique as possible to simplify the remaining
 /// shuffle.
 static SmallBitVector computeZeroableShuffleElements(ArrayRef<int> Mask,
-                                                     SDValue V1, SDValue V2) {
+                                                     SDValue V1, SDValue V2,
+                                                     unsigned Depth = 0) {
+  // Determine zeroable elements of each input.
+  auto GetSubZeroable = [&](SDValue V) {
+    MVT VT = V.getSimpleValueType();
+    unsigned NumElts = VT.isVector() ? VT.getVectorNumElements() : 1;
+
+    // BUILD_VECTOR - check each operand.
+    if (V.getOpcode() == ISD::BUILD_VECTOR) {
+      SmallBitVector SubZeroable(NumElts, false);
+      for (unsigned i = 0; i != NumElts; ++i) {
+        SDValue Op = V.getOperand(i);
+        SubZeroable[i] = (Op.getOpcode() == ISD::UNDEF || X86::isZeroNode(Op));
+      }
+      return SubZeroable;
+    }
+
+    // SHUFFLE_VECTOR - recursive call to computeZeroableShuffleElements.
+    else if (ShuffleVectorSDNode *S = dyn_cast<ShuffleVectorSDNode>(V))
+      return computeZeroableShuffleElements(S->getMask(), S->getOperand(0),
+                                            S->getOperand(1), Depth + 1);
+
+    return SmallBitVector(NumElts, false);
+  };
+
   SmallBitVector Zeroable(Mask.size(), false);
 
+  if (Depth == 6)
+    return Zeroable; // Limit search depth.
+
   while (V1.getOpcode() == ISD::BITCAST)
     V1 = V1->getOperand(0);
   while (V2.getOpcode() == ISD::BITCAST)
     V2 = V2->getOperand(0);
 
-  bool V1IsZero = ISD::isBuildVectorAllZeros(V1.getNode());
-  bool V2IsZero = ISD::isBuildVectorAllZeros(V2.getNode());
+  SmallBitVector V1Zeroables = GetSubZeroable(V1);
+  SmallBitVector V2Zeroables = GetSubZeroable(V2);
 
   for (int i = 0, Size = Mask.size(); i < Size; ++i) {
     int M = Mask[i];
-    // Handle the easy cases.
-    if (M < 0 || (M >= 0 && M < Size && V1IsZero) || (M >= Size && V2IsZero)) {
+
+    // Handle the easy case.
+    if (M < 0) {
       Zeroable[i] = true;
       continue;
     }
 
-    // If this is an index into a build_vector node (which has the same number
-    // of elements), dig out the input value and use it.
-    SDValue V = M < Size ? V1 : V2;
-    if (V.getOpcode() != ISD::BUILD_VECTOR || Size != (int)V.getNumOperands())
+    int Idx = M % Size;
+    SmallBitVector SubZeroable = M < Size ? V1Zeroables : V2Zeroables;
+    int SubSize = SubZeroable.size();
+
+    // If fewer (wider) elements - the entire element must be zero.
+    if (Size > SubSize) {
+      assert(0 == (Size % SubSize) && "Bad scale");
+      unsigned Scale = Size / SubSize;
+      Zeroable[i] = SubZeroable[Idx / Scale];
       continue;
+    }
 
-    SDValue Input = V.getOperand(M % Size);
-    // The UNDEF opcode check really should be dead code here, but not quite
-    // worth asserting on (it isn't invalid, just unexpected).
-    if (Input.getOpcode() == ISD::UNDEF || X86::isZeroNode(Input))
-      Zeroable[i] = true;
+    // If more (narrower) elements - all aliased 'sub-elements' must be zero.
+    if (Size < SubSize) {
+      assert(0 == (SubSize % Size) && "Bad scale");
+      unsigned Scale = SubSize / Size;
+      bool Zero = true;
+      for (unsigned j = 0; j != Scale; ++j)
+        Zero &= SubZeroable[(Idx * Scale) + j];
+      Zeroable[i] = Zero;
+      continue;
+    }
+
+    assert(Size == SubSize && "Unexpected mask size");
+    Zeroable[i] = SubZeroable[Idx];
   }
 
   return Zeroable;
 }
 
 // X86 has dedicated unpack instructions that can handle specific blend
 // operations: UNPCKH and UNPCKL.
 static SDValue lowerVectorShuffleWithUNPCK(SDLoc DL, MVT VT, ArrayRef<int> Mask,
                                            SDValue V1, SDValue V2,
                                            SelectionDAG &DAG) {
   int NumElts = VT.getVectorNumElements();
   int NumEltsInLane = 128 / VT.getScalarSizeInBits();
   SmallVector<int, 8> Unpckl;
   SmallVector<int, 8> Unpckh;
 

test/CodeGen/X86/sse41.ll

; X64-NEXT:    retq
   ret <4 x float> %vecinit4
 }
 
 define <4 x float> @shuf_X00A(<4 x float> %x, <4 x float> %a) {
 ; X32-LABEL: shuf_X00A:
 ; X32:       ## BB#0:
 ; X32-NEXT:    xorps %xmm2, %xmm2
 ; X32-NEXT:    blendps {{.*#+}} xmm0 = xmm0[0],xmm2[1,2,3]
-; X32-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0,1,2],xmm1[0]
+; X32-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0],zero,zero,xmm1[0]
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: shuf_X00A:
 ; X64:       ## BB#0:
 ; X64-NEXT:    xorps %xmm2, %xmm2
 ; X64-NEXT:    blendps {{.*#+}} xmm0 = xmm0[0],xmm2[1,2,3]
-; X64-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0,1,2],xmm1[0]
+; X64-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0],zero,zero,xmm1[0]
 ; X64-NEXT:    retq
   %vecext = extractelement <4 x float> %x, i32 0
   %vecinit = insertelement <4 x float> undef, float %vecext, i32 0
   %vecinit1 = insertelement <4 x float> %vecinit, float 0.0, i32 1
   %vecinit2 = insertelement <4 x float> %vecinit1, float 0.0, i32 2
   %vecinit4 = shufflevector <4 x float> %vecinit2, <4 x float> %a, <4 x i32> <i32 0, i32 1, i32 2, i32 4>
   ret <4 x float> %vecinit4
 }
; X64-NEXT:    retq
   %vecinit4 = shufflevector <4 x float> %vecinit2, <4 x float> %x, <4 x i32> <i32 0, i32 1, i32 2, i32 4>
   ret <4 x float> %vecinit4
 }
 
 define <4 x float> @shuf_X0YC(<4 x float> %x, <4 x float> %a) {
 ; X32-LABEL: shuf_X0YC:
 ; X32:       ## BB#0:
 ; X32-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0],zero,xmm0[1],zero
-; X32-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0,1,2],xmm1[2]
+; X32-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0],zero,xmm0[2],xmm1[2]
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: shuf_X0YC:
 ; X64:       ## BB#0:
 ; X64-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0],zero,xmm0[1],zero
-; X64-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0,1,2],xmm1[2]
+; X64-NEXT:    insertps {{.*#+}} xmm0 = xmm0[0],zero,xmm0[2],xmm1[2]
 ; X64-NEXT:    retq
   %vecext = extractelement <4 x float> %x, i32 0
   %vecinit = insertelement <4 x float> undef, float %vecext, i32 0
   %vecinit1 = insertelement <4 x float> %vecinit, float 0.0, i32 1
   %vecinit3 = shufflevector <4 x float> %vecinit1, <4 x float> %x, <4 x i32> <i32 0, i32 1, i32 5, i32 undef>
   %vecinit5 = shufflevector <4 x float> %vecinit3, <4 x float> %a, <4 x i32> <i32 0, i32 1, i32 2, i32 6>
   ret <4 x float> %vecinit5
 }

test/CodeGen/X86/widen_load-2.ll

; CHECK-NEXT:    movdqa  %[[R0]], (%{{.*}})
   store %i8vec31 %x, %i8vec31* %ret, align 16
   ret void
 }
 
 
 %i8vec3pack = type { <3 x i8>, i8 }
 define void @rot(%i8vec3pack* nocapture sret %result, %i8vec3pack* %X, %i8vec3pack* %rot) nounwind {
 ; CHECK-LABEL: rot:
-; CHECK:         movdqa  {{.*}}, %[[CONSTANT0:xmm[0-9]+]]
-; CHECK-NEXT:    movdqa  {{.*}}, %[[SHUFFLE_MASK:xmm[0-9]+]]
+; CHECK:         movdqa  {{.*}}, %[[SHUFFLE_MASK:xmm[0-9]+]]
+; CHECK-NEXT:    movdqa  {{.*}}, %[[CONSTANT0:xmm[0-9]+]]
 ; CHECK-NEXT:    pshufb  %[[SHUFFLE_MASK]], %[[CONSTANT0]]
 ; CHECK-NEXT:    pmovzxwq %[[CONSTANT0]], %[[CONSTANT0]]
 ; CHECK-NEXT:    movd    %[[CONSTANT0]], %e[[R0:[abcd]]]x
 ; CHECK-NEXT:    movw    %[[R0]]x, (%[[PTR0:.*]])
 ; CHECK-NEXT:    movb    $-98, 2(%[[PTR0]])
 ; CHECK-NEXT:    movdqa  {{.*}}, %[[CONSTANT1:xmm[0-9]+]]
 ; CHECK-NEXT:    pshufb  %[[SHUFFLE_MASK]], %[[CONSTANT1]]
-; CHECK-NEXT:    pmovzxwq %[[CONSTANT1]], %[[CONSTANT1]]
-; CHECK-NEXT:    movd    %[[CONSTANT1]], %e[[R1:[abcd]]]x
+; CHECK-NEXT:    pmovzxwq %[[CONSTANT1]], %[[SHUFFLE_MASK]]
+; CHECK-NEXT:    movd    %[[SHUFFLE_MASK]], %e[[R1:[abcd]]]x
 ; CHECK-NEXT:    movw    %[[R1]]x, (%[[PTR1:.*]])
 ; CHECK-NEXT:    movb    $1, 2(%[[PTR1]])
 ; CHECK-NEXT:    pmovzxbd (%[[PTR0]]), %[[X0:xmm[0-9]+]]
 ; CHECK-NEXT:    movdqa  %[[X0]], %[[X1:xmm[0-9]+]]
 ; CHECK-NEXT:    psrld   $1, %[[X1]]
 ; CHECK-NEXT:    pblendw $192, %[[X0]], %[[X1]]
 ; CHECK-NEXT:    pextrb  $8, %[[X1]], 2(%{{.*}})
-; CHECK-NEXT:    pshufb  %[[SHUFFLE_MASK]], %[[X1]]
+; CHECK-NEXT:    pshufb  {{.*}}, %[[X1]]
 ; CHECK-NEXT:    pmovzxwq %[[X1]], %[[X3:xmm[0-9]+]]
 ; CHECK-NEXT:    movd    %[[X3]], %e[[R0:[abcd]]]x
 ; CHECK-NEXT:    movw    %[[R0]]x, (%{{.*}})
 
 entry:
   %storetmp = bitcast %i8vec3pack* %X to <3 x i8>*
   store <3 x i8> <i8 -98, i8 -98, i8 -98>, <3 x i8>* %storetmp
   %storetmp1 = bitcast %i8vec3pack* %rot to <3 x i8>*