[AArch64] Improve shuffles of i1 vectors (WIP)

[rjj]

Summary

If the operands of a shuffle are concatenations of i1 vectors and undef,
this patch tries to rewrite the operands without the concatenation by doing a
bitcast instead. This avoids BUILD_VECTORs during legalisation, which
ultimately leads to better codegen.

This is WIP because currently two tests are failing that need looking into, but
in the meantime I wanted to get your opinion @dmgreen on the approach.

[SjoerdMeijer]

Just for completeness, this is related to: https://reviews.llvm.org/D142359

It (partially) fixes a regression when we reduce the vector insertion costs:

this "distance" one from cmsisdsp: https://godbolt.org/z/1xz7GP3fM.

[dmgreen]

Thanks for uploading this. It does look very similar to the case I was looking at recently, but goes further than the optimization I was trying. (That was about altering anyext(buildvector(..)), not doing it earlier which looks like it allows the shuffles to lower more efficiently). I will try to take a look, to get my head around how this works.

[rjj]

In D152826#4420599, @dmgreen wrote:

Thanks for uploading this. It does look very similar to the case I was looking at recently, but goes further than the optimization I was trying. (That was about altering anyext(buildvector(..)), not doing it earlier which looks like it allows the shuffles to lower more efficiently). I will try to take a look, to get my head around how this works.

Thanks for taking a look! To give a bit of context, this was brought up by the following example (from cmsisdsp iirc):

define <4 x i1> @t1(<2 x i1> %a, <2 x i1> %b)  {
  %r = shufflevector <2 x i1> %a, <2 x i1> %b, <4 x i32> <i32 0, i32 3, i32 0, i32 3>
  ret <4 x i1> %r
}

This was getting lowered to:

t1:                                     // @t1
        fmov    d2, d0
        mov     w8, v1.s[1]
        mov     v0.16b, v2.16b
        mov     v0.h[1], w8
        mov     v0.h[2], v2.h[0]
        mov     v0.h[3], w8
        ret

which is essentially an element-by-element insert.

However, by taking advantage of the symmetry of the mask, <(0, 3), (0, 3)>, we could do better if we used a bigger element type to replicate the first half of the mask into the second, e.g.:

t1:
        mov     v0.h[1], v1.h[2]
        mov     v0.s[1], v0.s[0]

As I was working on this, I saw the compiler was actually able to do something similar if we used i32s (for example) instead of i1s, but was getting stuck with the latter because its initial SelectionDAG consisted of vector_shuffle (concat (tx, undef), concat (ty, undef)), which would get type-legalised into BUILD_VECTORs and eventually lead to a sequence of inserts.

Since during legalisation the CopyFromReg nodes were getting bitcast into smaller types (v4i16 in this case), I thought I could cut the middleman and cast them earlier in the pipeline (making the necessary adjustments) and thus avoid the undef concatenations. The rest of the optimisations fell through from this.

In the end, for the example above we got the following assembly:

t1:
        mov     v0.h[2], v1.h[2]
        uzp1    v0.4h, v0.4h, v0.4h
        ret

which, despite being a bit less obvious, is functionally equivalent to the assembly with two moves and should have the same performance on the V2.

You can check this example on Alive2 here https://alive2.llvm.org/ce/z/297iBB where I've essentially written the IR that corresponds to the initial SelectionDAGs before and after the patch.

[dmgreen]

Hello. Sorry for the delay. This seems to work OK for AArch64, but I'm not sure how much other architectures would like it. It might be worth possibly adding it to AArch64ISelLowering instead, as it is easier there to make assumptions about which types would be better.

I was trying it with the motivating example, and whilst it did help by reducing the instruction count by a lot, it wasn't enough for it to become profitable compared to scalar. There seems to just be too much shuffling going on. I am attempting to see if I can stop the vectorization in that case instead. We've seen this come up a few times though, so it would be good to get an improvement like this in if we can.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

24742

This creates a setcc with the result type set to the input type? What happens for targets which want to use a native vxi1 datatype, as they have native predicate registers?

There is a getSetCCResultType method in which the target specifies what the setcc result type should become. There is also a TargetLowering::BooleanContent that specifies whether the bits of the predicate from setcc should be 0/1 or 0/-1 or the top bits are undefined.

[SjoerdMeijer]

There seems to just be too much shuffling going on. I am attempting to see if I can stop the vectorization in that case instead.

Ok, so parking this, and instead focus on reigning in vectorisation? I guess it makes sense and is another way of doing it.
I can't remember why we went for fixing the codegen. Perhaps that was just a fun exercise, and avoids looking at cost-modelling again if we manage to get the vector code on par with the scalar code (that would be good enough for this case).

[dmgreen]

It certainly looks like the code is better, and that this would be an improvement. But the scores didn't move as much as expected when I tried it. (The example has some pretty long SLP vectorization chain that reaches over multiple loops).

This would be useful if you did want to continue it, but I don't think it should be a blocker for anything.

[SjoerdMeijer]

Ok, got it, cheers.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

for (int i = 0, e = (int)Mask.size(); i != e; ++i) {
     if (Mask0[Mask[i]] != Mask0[i])
       return SDValue();
   }
   // Every element of this shuffle is identical to the result of the previous
   // shuffle, so we can replace this value.
   return Shuf->getOperand(0);
 }
 
+/// Helper function for simplifyShuffleOfI1Concats. If the operand `OpIdx' of
+/// `SV' is a concatenation of an i1 vector and undef, see if we can rewrite it
+/// without the concatenation.
+static SDValue simplifyConcatOfI1AndUndef(SelectionDAG &DAG,
+                                          ShuffleVectorSDNode const *SV,
+                                          int const OpIdx,
+                                          SmallVector<int, 8> &Mask) {
+  EVT VT = SV->getValueType(0);
+  unsigned NumElts = VT.getVectorNumElements();
+
+  SDValue N = SV->getOperand(OpIdx);
+  if (N.getOpcode() != ISD::CONCAT_VECTORS || N.getNumOperands() != 2)
+    return SDValue();
+
+  // NOTE: As long as the second operand is not used (which we check below), it
+  // doesn't really matter if the second operand is undef or not.
+  SDValue N0 = N.getOperand(0);  // setcc|truncate
+  SDValue N1 = N.getOperand(1);  // undef/unused
+  if (N0.getOpcode() != ISD::SETCC && N0.getOpcode() != ISD::TRUNCATE)
+    return SDValue();
+  if (!N1.isUndef())
+    return SDValue();
+
+  // We are restricting ourselves to shuffles of i1 vectors, but we could apply
+  // the same idea to other types.
+  if (N0.getValueType().getVectorElementType() != MVT::i1)
+    return SDValue();
+
+  // Check if elements of the second operand are accessed.
+  unsigned MaskBegin = OpIdx*NumElts;
+  unsigned MaskEnd = (OpIdx+1)*NumElts;
+  unsigned NumEltsOfN0 = N0.getValueType().getVectorNumElements();
+  for (int i = 0, e = (int)Mask.size(); i != e; ++i) {
+    if (Mask[i] >= (int)(MaskBegin+NumEltsOfN0) && Mask[i] < (int)MaskEnd)
+      return SDValue();
+  }
+
+  // Check if we can do a bitcast from NumSrcBits to VT.
+  unsigned NumSrcBits = N0.getOperand(0).getValueSizeInBits();
+  if (NumSrcBits % NumElts != 0 || NumElts % NumEltsOfN0 != 0)
+    return SDValue();
+
+  EVT NewST = EVT::getIntegerVT(*DAG.getContext(), NumSrcBits/NumElts);
+  EVT NewVT = EVT::getVectorVT(*DAG.getContext(), NewST, NumElts);
+
+  SDValue NewN = N0.getOperand(0);
+  if (N0.getOpcode() == ISD::SETCC)
+    NewN = DAG.getSetCC(SDLoc(SV), N0.getOperand(0).getValueType(),
+                        N0.getOperand(0), N0.getOperand(1),
+                        cast<CondCodeSDNode>(N0.getOperand(2))->get());
+
+  // The bitcast essentially interleaves new elements in the operand. Here we
+  // adjust the mask to account for the interleaved elements.
+  unsigned MaskFactor = NumElts/NumEltsOfN0;
+  for (int i = 0, e = (int)Mask.size(); i != e; ++i) {
+    if (Mask[i] >= (int)MaskBegin && Mask[i] < (int)MaskEnd)
+      Mask[i] = (Mask[i]-MaskBegin)*MaskFactor+MaskBegin;
+  }
+
+  return DAG.getAnyExtOrTrunc(DAG.getBitcast(NewVT, NewN), SDLoc(SV), VT);
+}
+
+/// If the operands of a shuffle are concatenations of i1 vectors and undef,
+/// see if we can rewrite the operands without the concatenation by doing a
+/// bitcast instead. This avoids BUILD_VECTORs during legalisation, which
+/// ultimately leads to better codegen.
+static SDValue simplifyShuffleOfI1Concats(SelectionDAG &DAG,
+                                          ShuffleVectorSDNode const *SV) {
+  EVT VT = SV->getValueType(0);
+  SmallVector<int, 8> Mask(SV->getMask());
+
+  SDValue Op0 = simplifyConcatOfI1AndUndef(DAG, SV, 0, Mask);
+  SDValue Op1 = simplifyConcatOfI1AndUndef(DAG, SV, 1, Mask);
+
+  if (!Op0 && !Op1)
+    return SDValue();
+
+  if (!Op0) Op0 = SV->getOperand(0);
+  if (!Op1) Op1 = SV->getOperand(1);
+
+  return DAG.getVectorShuffle(VT, SDLoc(SV), Op0, Op1, Mask);
+}
+
 SDValue DAGCombiner::visitVECTOR_SHUFFLE(SDNode *N) {
   EVT VT = N->getValueType(0);
   unsigned NumElts = VT.getVectorNumElements();
 
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
 
   assert(N0.getValueType() == VT && "Vector shuffle must be normalized in DAG");
if (TLI.isBinOp(SrcOpcode) && N->isOnlyUserOf(N0.getNode()) &&
               VT, DL, RightSV0 ? RightSV0 : DAG.getUNDEF(VT),
               RightSV1 ? RightSV1 : DAG.getUNDEF(VT), RightMask);
           return DAG.getNode(SrcOpcode, DL, VT, LHS, RHS);
         }
       }
     }
   }
 
+  // Try to replace the concats in shuffle (concat (i1, undef), concat (i1,
+  // undef), Mask) with bitcasts
+  if (SDValue V = simplifyShuffleOfI1Concats(DAG, SVN))
+    return V;
+
   if (SDValue V = foldShuffleOfConcatUndefs(SVN, DAG))
     return V;
 
   // Match shuffles that can be converted to ISD::ZERO_EXTEND_VECTOR_INREG.
   // Perform this really late, because it could eliminate knowledge
   // of undef elements created by this shuffle.
   if (Level < AfterLegalizeTypes)
     if (SDValue V = combineShuffleToZeroExtendVectorInReg(SVN, DAG, TLI,

llvm/test/CodeGen/AArch64/aarch64-shuffle-of-i1s.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple aarch64 < %s | FileCheck %s
 
 ; https://alive2.llvm.org/ce/z/r_f2Zj
 define <4 x i1> @t11(<2 x i1> %a, <2 x i1> %b)  {
 ; CHECK-LABEL: t11:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    fmov d2, d0
-; CHECK-NEXT:    // kill: def $d1 killed $d1 def $q1
-; CHECK-NEXT:    mov w8, v1.s[1]
-; CHECK-NEXT:    mov v0.16b, v2.16b
-; CHECK-NEXT:    mov v0.h[1], w8
-; CHECK-NEXT:    mov v0.h[2], w8
-; CHECK-NEXT:    mov v0.h[3], v2.h[0]
-; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
+; CHECK-NEXT:    ext v0.8b, v1.8b, v0.8b, #2
+; CHECK-NEXT:    trn2 v0.4h, v0.4h, v0.4h
+; CHECK-NEXT:    ext v0.8b, v0.8b, v0.8b, #6
 ; CHECK-NEXT:    ret
   %r = shufflevector <2 x i1> %a, <2 x i1> %b, <4 x i32> <i32 0, i32 3, i32 3, i32 0>
   ret <4 x i1> %r
 }
 
 define <4 x i1> @t12(<2 x i32> %a, <2 x i32> %b)  {
 ; CHECK-LABEL: t12:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    fmov d2, d0
-; CHECK-NEXT:    // kill: def $d1 killed $d1 def $q1
-; CHECK-NEXT:    mov w8, v1.s[1]
-; CHECK-NEXT:    mov v0.16b, v2.16b
-; CHECK-NEXT:    mov v0.h[1], w8
-; CHECK-NEXT:    mov v0.h[2], w8
-; CHECK-NEXT:    mov v0.h[3], v2.h[0]
-; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
+; CHECK-NEXT:    ext v0.8b, v1.8b, v0.8b, #2
+; CHECK-NEXT:    trn2 v0.4h, v0.4h, v0.4h
+; CHECK-NEXT:    ext v0.8b, v0.8b, v0.8b, #6
 ; CHECK-NEXT:    ret
   %a2 = trunc <2 x i32> %a to <2 x i1>
   %b2 = trunc <2 x i32> %b to <2 x i1>
   %a3 = shufflevector <2 x i1> %a2, <2 x i1> undef, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %b3 = shufflevector <2 x i1> %b2, <2 x i1> undef, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %r  = shufflevector <4 x i1> %a3, <4 x i1> %b3, <4 x i32> <i32 0, i32 5, i32 5, i32 0>
   ret <4 x i1> %r
 }
 
 ; https://alive2.llvm.org/ce/z/297iBB
 define <4 x i1> @t21(<2 x i1> %a, <2 x i1> %b)  {
 ; CHECK-LABEL: t21:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    fmov d2, d0
+; CHECK-NEXT:    // kill: def $d0 killed $d0 def $q0
 ; CHECK-NEXT:    // kill: def $d1 killed $d1 def $q1
-; CHECK-NEXT:    mov w8, v1.s[1]
-; CHECK-NEXT:    mov v0.16b, v2.16b
-; CHECK-NEXT:    mov v0.h[1], w8
-; CHECK-NEXT:    mov v0.h[2], v2.h[0]
-; CHECK-NEXT:    mov v0.h[3], w8
-; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
+; CHECK-NEXT:    mov v0.h[2], v1.h[2]
+; CHECK-NEXT:    uzp1 v0.4h, v0.4h, v0.4h
 ; CHECK-NEXT:    ret
   %r = shufflevector <2 x i1> %a, <2 x i1> %b, <4 x i32> <i32 0, i32 3, i32 0, i32 3>
   ret <4 x i1> %r
 }
 
 define <4 x i1> @t22(<2 x i32> %a, <2 x i32> %b)  {
 ; CHECK-LABEL: t22:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    fmov d2, d0
+; CHECK-NEXT:    // kill: def $d0 killed $d0 def $q0
 ; CHECK-NEXT:    // kill: def $d1 killed $d1 def $q1
-; CHECK-NEXT:    mov w8, v1.s[1]
-; CHECK-NEXT:    mov v0.16b, v2.16b
-; CHECK-NEXT:    mov v0.h[1], w8
-; CHECK-NEXT:    mov v0.h[2], v2.h[0]
-; CHECK-NEXT:    mov v0.h[3], w8
-; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
+; CHECK-NEXT:    mov v0.h[2], v1.h[2]
+; CHECK-NEXT:    uzp1 v0.4h, v0.4h, v0.4h
 ; CHECK-NEXT:    ret
   %a2 = trunc <2 x i32> %a to <2 x i1>
   %b2 = trunc <2 x i32> %b to <2 x i1>
   %a3 = shufflevector <2 x i1> %a2, <2 x i1> undef, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %b3 = shufflevector <2 x i1> %b2, <2 x i1> undef, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %r  = shufflevector <4 x i1> %a3, <4 x i1> %b3, <4 x i32> <i32 0, i32 5, i32 0, i32 5>
   ret <4 x i1> %r
 }
 
 ; From cmsisdsp ( https://godbolt.org/z/1xz7GP3fM ), one of the regressions of
 ; https://reviews.llvm.org/D142359
 define <4 x i1> @t3(<2 x i32> %a, <2 x i32> %b)  {
 ; CHECK-LABEL: t3:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cmtst v2.2s, v0.2s, v0.2s
 ; CHECK-NEXT:    cmeq v0.2s, v0.2s, #0
-; CHECK-NEXT:    cmtst v3.2s, v1.2s, v1.2s
+; CHECK-NEXT:    ext v0.8b, v0.8b, v2.8b, #2
+; CHECK-NEXT:    cmtst v2.2s, v1.2s, v1.2s
 ; CHECK-NEXT:    cmeq v1.2s, v1.2s, #0
-; CHECK-NEXT:    mov w8, v0.s[1]
-; CHECK-NEXT:    mov w9, v1.s[1]
-; CHECK-NEXT:    mov v0.16b, v2.16b
-; CHECK-NEXT:    mov v1.16b, v3.16b
-; CHECK-NEXT:    mov v0.h[1], w8
-; CHECK-NEXT:    mov v1.h[1], w9
-; CHECK-NEXT:    mov v0.h[2], w8
-; CHECK-NEXT:    mov v1.h[2], v3.h[0]
-; CHECK-NEXT:    mov v0.h[3], v2.h[0]
-; CHECK-NEXT:    mov v1.h[3], w9
+; CHECK-NEXT:    mov v2.h[2], v1.h[2]
+; CHECK-NEXT:    trn2 v0.4h, v0.4h, v0.4h
+; CHECK-NEXT:    uzp1 v1.4h, v2.4h, v2.4h
+; CHECK-NEXT:    ext v0.8b, v0.8b, v0.8b, #6
 ; CHECK-NEXT:    and v0.8b, v0.8b, v1.8b
 ; CHECK-NEXT:    ret
   %a1 = icmp ne <2 x i32> %a, zeroinitializer
   %a2 = icmp eq <2 x i32> %a, zeroinitializer
   %a3 = shufflevector <2 x i1> %a1, <2 x i1> %a2, <4 x i32> <i32 0, i32 3, i32 3, i32 0>
   %b1 = icmp ne <2 x i32> %b, zeroinitializer
   %b2 = icmp eq <2 x i32> %b, zeroinitializer
   %b3 = shufflevector <2 x i1> %b1, <2 x i1> %b2, <4 x i32> <i32 0, i32 3, i32 0, i32 3>
   %r  = select <4 x i1> %a3, <4 x i1> %b3, <4 x i1> zeroinitializer
   ret <4 x i1> %r
 }