[AArch64] Fix zero element TBL indices

[dmgreen]

Summary

A TBL instruction will fill out-of-range values with 0's, something used in D121139 to turn tbl2 with a zero input into tbl1s. This works OK for v16i8, but for v8i8 the input is still treated as a v16i8, so "out-of-range" values (like a lane index of 8) would end up loading values from the top half of the input register. Clean this up by detecting the out of range values and making sure they really use out of range values. There is a fix for swapped indices of 64bit input vectors too, which could be incorrectly adjusted if the zerovector was the first operand.

Fixes #55545

[samtebbs]

LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static SDValue GenerateTBL(SDValue Op, ArrayRef<int> ShuffleMask,
   unsigned BytesPerElt = EltVT.getSizeInBits() / 8;
 
   bool Swap = false;
   if (V1.isUndef() || isZerosVector(V1.getNode())) {
     std::swap(V1, V2);
     Swap = true;
   }
 
+  // If the V2 source is undef or zero then we can use a tbl1, as tbl1 will fill
+  // out of range values with 0s. We do need to make sure that any out-of-range
+  // values are really out-of-range for a v16i8 vector.
+  bool IsUndefOrZero = V2.isUndef() || isZerosVector(V2.getNode());
+  MVT IndexVT = MVT::v8i8;
+  unsigned IndexLen = 8;
+  if (Op.getValueSizeInBits() == 128) {
+    IndexVT = MVT::v16i8;
+    IndexLen = 16;
+  }
+
   SmallVector<SDValue, 8> TBLMask;
   for (int Val : ShuffleMask) {
     for (unsigned Byte = 0; Byte < BytesPerElt; ++Byte) {
       unsigned Offset = Byte + Val * BytesPerElt;
       if (Swap)
-        Offset = Offset < 16 ? Offset + 16 : Offset - 16;
+        Offset = Offset < IndexLen ? Offset + IndexLen : Offset - IndexLen;
+      if (IsUndefOrZero && Offset >= IndexLen)
+        Offset = 255;
       TBLMask.push_back(DAG.getConstant(Offset, DL, MVT::i32));
     }
   }
 
-  MVT IndexVT = MVT::v8i8;
-  unsigned IndexLen = 8;
-  if (Op.getValueSizeInBits() == 128) {
-    IndexVT = MVT::v16i8;
-    IndexLen = 16;
-  }
-
   SDValue V1Cst = DAG.getNode(ISD::BITCAST, DL, IndexVT, V1);
   SDValue V2Cst = DAG.getNode(ISD::BITCAST, DL, IndexVT, V2);
 
   SDValue Shuffle;
-  // If the V2 source is undef or zero then we can use a tbl1, as tbl1 will fill
-  // out of range values with 0s.
-  if (V2.isUndef() || isZerosVector(V2.getNode())) {
+  if (IsUndefOrZero) {
     if (IndexLen == 8)
       V1Cst = DAG.getNode(ISD::CONCAT_VECTORS, DL, MVT::v16i8, V1Cst, V1Cst);
     Shuffle = DAG.getNode(
         ISD::INTRINSIC_WO_CHAIN, DL, IndexVT,
         DAG.getConstant(Intrinsic::aarch64_neon_tbl1, DL, MVT::i32), V1Cst,
         DAG.getBuildVector(IndexVT, DL,
                            makeArrayRef(TBLMask.data(), IndexLen)));
   } else {

llvm/test/CodeGen/AArch64/neon-bitwise-instructions.ll

 ; CHECK-NEXT:    tbl v0.8b, { v0.16b }, v1.8b
 ; CHECK-NEXT:    ret
   %c = shufflevector <8 x i8> %a, <8 x i8> %b, <8 x i32> <i32 0, i32 8, i32 2, i32 9, i32 4, i32 5, i32 6, i32 7>
   ret <8 x i8> %c
 }
 
 ; CHECK-LABEL: .LCPI90_0:
 ; CHECK-NEXT: .byte   0
-; CHECK-NEXT: .byte   8
+; CHECK-NEXT: .byte   255
 ; CHECK-NEXT: .byte   2
-; CHECK-NEXT: .byte   9
+; CHECK-NEXT: .byte   255
 ; CHECK-NEXT: .byte   4
 ; CHECK-NEXT: .byte   5
 ; CHECK-NEXT: .byte   6
 ; CHECK-NEXT: .byte   7
 define <8 x i8> @vselect_equivalent_shuffle_v8i8_zero(<8 x i8> %a) {
 ; CHECK-LABEL: vselect_equivalent_shuffle_v8i8_zero:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    adrp x8, .LCPI90_0
 ; CHECK-NEXT:    // kill: def $d0 killed $d0 def $q0
 ; CHECK-NEXT:    mov v0.d[1], v0.d[0]
 ; CHECK-NEXT:    ldr d1, [x8, :lo12:.LCPI90_0]
 ; CHECK-NEXT:    tbl v0.8b, { v0.16b }, v1.8b
 ; CHECK-NEXT:    ret
   %c = shufflevector <8 x i8> %a, <8 x i8> zeroinitializer, <8 x i32> <i32 0, i32 8, i32 2, i32 9, i32 4, i32 5, i32 6, i32 7>
   ret <8 x i8> %c
 }
 
 ; CHECK-LABEL: .LCPI91_0:
-; CHECK-NEXT: .byte   24
-; CHECK-NEXT: .byte   16
-; CHECK-NEXT: .byte   26
-; CHECK-NEXT: .byte   17
-; CHECK-NEXT: .byte   28
-; CHECK-NEXT: .byte   29
-; CHECK-NEXT: .byte   30
-; CHECK-NEXT: .byte   31
+; CHECK-NEXT: .byte   0
+; CHECK-NEXT: .byte   255
+; CHECK-NEXT: .byte   2
+; CHECK-NEXT: .byte   255
+; CHECK-NEXT: .byte   4
+; CHECK-NEXT: .byte   5
+; CHECK-NEXT: .byte   6
+; CHECK-NEXT: .byte   7
 define <8 x i8> @vselect_equivalent_shuffle_v8i8_zeroswap(<8 x i8> %a) {
 ; CHECK-LABEL: vselect_equivalent_shuffle_v8i8_zeroswap:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    adrp x8, .LCPI91_0
 ; CHECK-NEXT:    // kill: def $d0 killed $d0 def $q0
 ; CHECK-NEXT:    mov v0.d[1], v0.d[0]
 ; CHECK-NEXT:    ldr d1, [x8, :lo12:.LCPI91_0]
 ; CHECK-NEXT:    tbl v0.8b, { v0.16b }, v1.8b

 ; CHECK-NEXT:    ret
   %c = shufflevector <8 x i16> %a, <8 x i16> %b, <8 x i32> <i32 0, i32 8, i32 2, i32 9, i32 4, i32 5, i32 6, i32 7>
   ret <8 x i16> %c
 }
 
 ; CHECK-LABEL: .LCPI93_0:
 ; CHECK-NEXT: .byte   0
 ; CHECK-NEXT: .byte   1
-; CHECK-NEXT: .byte   16
-; CHECK-NEXT: .byte   17
+; CHECK-NEXT: .byte   255
+; CHECK-NEXT: .byte   255
 ; CHECK-NEXT: .byte   4
 ; CHECK-NEXT: .byte   5
-; CHECK-NEXT: .byte   18
-; CHECK-NEXT: .byte   19
+; CHECK-NEXT: .byte   255
+; CHECK-NEXT: .byte   255
 ; CHECK-NEXT: .byte   8
 ; CHECK-NEXT: .byte   9
 ; CHECK-NEXT: .byte   10
 ; CHECK-NEXT: .byte   11
 ; CHECK-NEXT: .byte   12
 ; CHECK-NEXT: .byte   13
 ; CHECK-NEXT: .byte   14
 ; CHECK-NEXT: .byte   15
 define <8 x i16> @vselect_equivalent_shuffle_v8i16_zero(<8 x i16> %a) {
 ; CHECK-LABEL: vselect_equivalent_shuffle_v8i16_zero:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    adrp x8, .LCPI93_0
 ; CHECK-NEXT:    ldr q1, [x8, :lo12:.LCPI93_0]
 ; CHECK-NEXT:    tbl v0.16b, { v0.16b }, v1.16b
 ; CHECK-NEXT:    ret
   %c = shufflevector <8 x i16> %a, <8 x i16> zeroinitializer, <8 x i32> <i32 0, i32 8, i32 2, i32 9, i32 4, i32 5, i32 6, i32 7>
   ret <8 x i16> %c
 }
 
 ; CHECK: .byte   0
 ; CHECK: .byte   1
-; CHECK: .byte   16
-; CHECK: .byte   17
+; CHECK: .byte   255
+; CHECK: .byte   255
 ; CHECK: .byte   4
 ; CHECK: .byte   5
-; CHECK: .byte   18
-; CHECK: .byte   19
+; CHECK: .byte   255
+; CHECK: .byte   255
 ; CHECK: .byte   8
 ; CHECK: .byte   9
 ; CHECK: .byte   10
 ; CHECK: .byte   11
 ; CHECK: .byte   12
 ; CHECK: .byte   13
 ; CHECK: .byte   14
 ; CHECK: .byte   15