[AArch64] Try to fold shuffle (tbl2, tbl2) to tbl4.

[fhahn]

Summary

shuffle (tbl2, tbl2) can be folded into a single tbl4 if the shuffle selects
the first 8 elements of each tbl2 and the tbl2 masks can be merged.

[t.p.northover]

At first glance this seems like a hyper-specific optimization, I take it there's some reasonably common idiom that motivates us even bothering?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10702–10703

Why do we care about this? It looks like we've already checked that lanes being filled by this check are discarded by the shuffle.

10707

Won't this overflow if it's a tbl2 produding an <8 x i8>?

10716

Maybe default fill with SDValue()? We just overwrite all of them immediately afterwards anyway so that'd signal early that the reader doesn't have to care about this line.

10719

Have we checked anywhere that the lower 8 operands are actually constant?

[fhahn]

Address latest comments, thanks!

In D133491#3791661, @t.p.northover wrote:

At first glance this seems like a hyper-specific optimization, I take it there's some reasonably common idiom that motivates us even bothering?

It is quite specific but this kind of pattern can be produced by loop-vectorization, in combination with the recent changes using tbl instructions for extends/truncates.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10702–10703

Good point, I removed the value checks and re-purposed the helper to check if the first 8 elements are constants.

10707

Yes, added extra tests in 9f2c39418b85 and limit this to v16i8 for now, to be extended further in a follow-up.

10716

Done, thanks!

10719

That's checked in the latest version and I added extra test cases in 9f2c39418b85.

[t.p.northover]

It is quite specific but this kind of pattern can be produced by loop-vectorization, in combination with the recent changes using tbl instructions for extends/truncates.

I think it's more restrictive than it needs to be, and we should drop the isExtractLowerHalfMask check entirely. Any constant shuffle of two constant tbl2 operations ought to be representable with a constant tbl4, and that check's only there because our index-mapping is too naive.

Instead I think we want something like this in the loop where we generate the new tbl indices (now running through all 16 lanes):

if (ShuffleMask[I] < 16)
  TblMaskParts[I] = Mask1->getOperand(ShuffleMask[I]);
else {
  auto *C = cast<ConstantSDNode>(Mask2->getOperand(ShuffleMask[I] - 16));
  TblMaskParts[I] = DAG.getConstant(C->getSExtValue() + 32, dl, MVT::i32);
}

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10706

This doesn't appear to capture anything, so maybe consider making it a static function instead? The difference from isExtractLowerHalfMask in the same patch seems a bit unmotivated.

Not insisting on one style or the other (or even consistency), just making sure it's an active decision.

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10709

Would also need to extend this check. Or perhaps better only check the indices that are actually used by the shuffle, so move this into the main loop later.

[fhahn]

In D133491#3802133, @t.p.northover wrote:

It is quite specific but this kind of pattern can be produced by loop-vectorization, in combination with the recent changes using tbl instructions for extends/truncates.

I think it's more restrictive than it needs to be, and we should drop the isExtractLowerHalfMask check entirely. Any constant shuffle of two constant tbl2 operations ought to be representable with a constant tbl4, and that check's only there because our index-mapping is too naive.

Instead I think we want something like this in the loop where we generate the new tbl indices (now running through all 16 lanes):

Thanks Tim! Generalized as suggested.

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

10706

This is not needed any longer, because it is checked in the loop below now.

10709

Done in the loop below, thanks!

[t.p.northover]

Thanks Florian. LGTM!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (DAG.getTargetLoweringInfo().isTypeLegal(NewVT)) {
       return DAG.getBitcast(VT,
                             DAG.getVectorShuffle(NewVT, DL, V0, V1, NewMask));
     }
   }
 
   return SDValue();
 }
 
+static bool isExtractLowerHalfMask(ArrayRef<int> Mask) {
+  if (Mask.size() != 16)
+    return false;
+
+  for (int I = 0; I < 8; ++I) {
+    if (Mask[I] != I)
+      return false;
+    if (Mask[I + 8] != 16 + I)
+      return false;
+  }
+  return true;
+}
+
+// Try to fold shuffle (tbl2, tbl2) into a single tbl4 if the shuffle selects
+// the first 8 elements of each tbl2 and the tbl2 masks can be merged.
+static SDValue tryToConvertShuffleOfTbl2ToTbl4(SDValue Op,
+                                               ArrayRef<int> ShuffleMask,
+                                               SelectionDAG &DAG) {
+  if (!isExtractLowerHalfMask(ShuffleMask))
+    return SDValue();
+
+  SDValue Tbl1 = Op->getOperand(0);
+  SDValue Tbl2 = Op->getOperand(1);
+  SDLoc dl(Op);
+  SDValue Tbl2ID =
+      DAG.getTargetConstant(Intrinsic::aarch64_neon_tbl2, dl, MVT::i64);
+
+  EVT VT = Op.getValueType();
+  if (Tbl1->getOpcode() != ISD::INTRINSIC_WO_CHAIN ||
+      Tbl1->getOperand(0) != Tbl2ID ||
+      Tbl2->getOpcode() != ISD::INTRINSIC_WO_CHAIN ||
+      Tbl2->getOperand(0) != Tbl2ID)
+    return SDValue();
+
+  SDValue Mask1 = Tbl1->getOperand(3);
+  SDValue Mask2 = Tbl2->getOperand(3);
+  // Make sure the tbl2 mask only selects values in the first 8 lanes (i.e. the
+  // last 8 lanes all have an index of -1).
+  auto IsLowerExtractMask = [](SDValue Mask) {
+    if (Mask->getOpcode() != ISD::BUILD_VECTOR)
+      return false;
+    for (unsigned I = 8; I < 16; I++) {
+      auto *C = dyn_cast<ConstantSDNode>(Mask->getOperand(I));
+      if (!C || C->getSExtValue() != -1)
+        return false;
+    }
+    return true;
+  };
+  if (!IsLowerExtractMask(Mask1) || !IsLowerExtractMask(Mask2))
+    return SDValue();
+  SmallVector<SDValue, 16> TBLMaskParts(16, Mask1->getOperand(0));
+  for (unsigned I = 0; I < 8; I++) {
+    TBLMaskParts[I] = Mask1->getOperand(I);
+    auto *C = cast<ConstantSDNode>(Mask2->getOperand(I));
+    TBLMaskParts[I + 8] = DAG.getConstant(C->getSExtValue() + 32, dl, MVT::i32);
+  }
+
+  SDValue TBLMask = DAG.getBuildVector(VT, dl, TBLMaskParts);
+  SDValue ID =
+      DAG.getTargetConstant(Intrinsic::aarch64_neon_tbl4, dl, MVT::i64);
+
+  return DAG.getNode(ISD::INTRINSIC_WO_CHAIN, dl, MVT::v16i8,
+                     {ID, Tbl1->getOperand(1), Tbl1->getOperand(2),
+                      Tbl2->getOperand(1), Tbl2->getOperand(2), TBLMask});
+}
+
 SDValue AArch64TargetLowering::LowerVECTOR_SHUFFLE(SDValue Op,
                                                    SelectionDAG &DAG) const {
   SDLoc dl(Op);
   EVT VT = Op.getValueType();
 
   ShuffleVectorSDNode *SVN = cast<ShuffleVectorSDNode>(Op.getNode());
 
   if (useSVEForFixedLengthVectorVT(VT))
     return LowerFixedLengthVECTOR_SHUFFLEToSVE(Op, DAG);
 
   // Convert shuffles that are directly supported on NEON to target-specific
   // DAG nodes, instead of keeping them as shuffles and matching them again
   // during code selection.  This is more efficient and avoids the possibility
   // of inconsistencies between legalization and selection.
   ArrayRef<int> ShuffleMask = SVN->getMask();
 
   SDValue V1 = Op.getOperand(0);
   SDValue V2 = Op.getOperand(1);
 
   assert(V1.getValueType() == VT && "Unexpected VECTOR_SHUFFLE type!");
   assert(ShuffleMask.size() == VT.getVectorNumElements() &&
          "Unexpected VECTOR_SHUFFLE mask size!");
 
+  if (SDValue Res = tryToConvertShuffleOfTbl2ToTbl4(Op, ShuffleMask, DAG))
+    return Res;
+
   if (SVN->isSplat()) {
     int Lane = SVN->getSplatIndex();
     // If this is undef splat, generate it via "just" vdup, if possible.
     if (Lane == -1)
       Lane = 0;
 
     if (Lane == 0 && V1.getOpcode() == ISD::SCALAR_TO_VECTOR)
       return DAG.getNode(AArch64ISD::DUP, dl, V1.getValueType(),

llvm/test/CodeGen/AArch64/arm64-tbl.ll

 ; CHECK-NEXT:     .byte    0                               // 0x0
 ; CHECK-NEXT:     .byte    4                               // 0x4
 ; CHECK-NEXT:     .byte    8                               // 0x8
 ; CHECK-NEXT:     .byte    12                              // 0xc
 ; CHECK-NEXT:     .byte    16                              // 0x10
 ; CHECK-NEXT:     .byte    20                              // 0x14
 ; CHECK-NEXT:     .byte    24                              // 0x18
 ; CHECK-NEXT:     .byte    28                              // 0x1c
-; CHECK-NEXT:     .byte    255                             // 0xff
-; CHECK-NEXT:     .byte    255                             // 0xff
-; CHECK-NEXT:     .byte    255                             // 0xff
-; CHECK-NEXT:     .byte    255                             // 0xff
-; CHECK-NEXT:     .byte    255                             // 0xff
-; CHECK-NEXT:     .byte    255                             // 0xff
-; CHECK-NEXT:     .byte    255                             // 0xff
-; CHECK-NEXT:     .byte    255                             // 0xff
+; CHECK-NEXT:    .byte   32                              // 0x20
+; CHECK-NEXT:    .byte   36                              // 0x24
+; CHECK-NEXT:    .byte   40                              // 0x28
+; CHECK-NEXT:    .byte   44                              // 0x2c
+; CHECK-NEXT:    .byte   48                              // 0x30
+; CHECK-NEXT:    .byte   52                              // 0x34
+; CHECK-NEXT:    .byte   56                              // 0x38
+; CHECK-NEXT:    .byte   60                              // 0x3c
 
 define <16 x i8> @shuffled_tbl2_to_tbl4(<16 x i8> %a, <16 x i8> %b, <16 x i8> %c, <16 x i8> %d) {
 ; CHECK-LABEL: shuffled_tbl2_to_tbl4:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    adrp x8, .LCPI8_0
-; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $q0_q1 def $q0_q1
-; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $q2_q3 def $q2_q3
-; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $q0_q1 def $q0_q1
-; CHECK-NEXT:    // kill: def $q2 killed $q2 killed $q2_q3 def $q2_q3
+; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $q0_q1_q2_q3 def $q0_q1_q2_q3
+; CHECK-NEXT:    // kill: def $q2 killed $q2 killed $q0_q1_q2_q3 def $q0_q1_q2_q3
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $q0_q1_q2_q3 def $q0_q1_q2_q3
 ; CHECK-NEXT:    ldr q4, [x8, :lo12:.LCPI8_0]
-; CHECK-NEXT:    tbl.16b v0, { v0, v1 }, v4
-; CHECK-NEXT:    tbl.16b v1, { v2, v3 }, v4
-; CHECK-NEXT:    mov.d v0[1], v1[0]
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $q0_q1_q2_q3 def $q0_q1_q2_q3
+; CHECK-NEXT:    tbl.16b v0, { v0, v1, v2, v3 }, v4
 ; CHECK-NEXT:    ret
   %t1 = call <16 x i8> @llvm.aarch64.neon.tbl2.v16i8(<16 x i8> %a, <16 x i8> %b, <16 x i8> <i8 0, i8 4, i8 8, i8 12, i8 16, i8 20, i8 24, i8 28, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1>)
   %t2 = call <16 x i8> @llvm.aarch64.neon.tbl2.v16i8(<16 x i8> %c, <16 x i8> %d, <16 x i8> <i8 0, i8 4, i8 8, i8 12, i8 16, i8 20, i8 24, i8 28, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1>)
   %s = shufflevector <16 x i8> %t1, <16 x i8> %t2, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 16, i32 17, i32 18, i32 19, i32 20, i32 21, i32 22, i32 23>
   ret <16 x i8> %s
 }
 
 define <16 x i8> @shuffled_tbl2_to_tbl4_incompatible_shuffle(<16 x i8> %a, <16 x i8> %b, <16 x i8> %c, <16 x i8> %d) {