[AArch64] Skip mask checks for masks with an odd number of elements.

[fhahn]

Summary

Some checks in isShuffleMaskLegal expect an even number of elements,
e.g. isTRN_v_undef_Mask or isUZP_v_undef_Mask, otherwise they access
invalid elements and crash. This patch adds a check to
isShuffleMaskLegal to ensure there is an even number of elements or
exactly one. The case with one element is relevant for matching some dup
cases. Masks with one element will still cause a crash in isTRN_v_undef_Mask and
isUZP_v_undef_Mask, but I am not sure if that is a realistic scenario. I
can add a separate check or an assertion, if required.

I am not entirely sure if the check is overly restrictive, as I am not
too familiar with the patterns here. We could always push it down to the
problematic functions.

[Gerolf]

Could you separate size one checks are OK vs. checks that require even powers of 2? That smells like an inconsistency worth calling out explicitly.

Thanks
Gerolf

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7553

This check could be the early exit at the start of the function.

[fhahn]

Push even element checks to the impacted functions.

[fhahn]

In D60690#1469460, @Gerolf wrote:

Could you separate size one checks are OK vs. checks that require even powers of 2? That smells like an inconsistency worth calling out explicitly.

Thanks Gerolf. I pushed the checks to the impacted functions, which seems to better limit the impact of the checks. What do you think?

[fhahn]

ping

[fhahn]

Ping. @Gerolf, are you happy with the update?

This also independently popped up as PR41951

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7532–7533

Unrelated change, I'll undo before committing.

[dmgreen]

LGTM. Thanks for making the fix.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

SDValue AArch64TargetLowering::LowerEXTRACT_SUBVECTOR(SDValue Op,
   // that directly.
   if (Size == 64 && Val * VT.getScalarSizeInBits() == 64)
     return Op;
 
   return SDValue();
 }
 
 bool AArch64TargetLowering::isShuffleMaskLegal(ArrayRef<int> M, EVT VT) const {
-  if (VT.getVectorNumElements() == 4 &&
+  unsigned NumElts = VT.getVectorNumElements();
+  if (NumElts == 4 &&
       (VT.is128BitVector() || VT.is64BitVector())) {
     unsigned PFIndexes[4];
     for (unsigned i = 0; i != 4; ++i) {
       if (M[i] < 0)
         PFIndexes[i] = 8;
       else
         PFIndexes[i] = M[i];
     }
 
     // Compute the index in the perfect shuffle table.
     unsigned PFTableIndex = PFIndexes[0] * 9 * 9 * 9 + PFIndexes[1] * 9 * 9 +
                             PFIndexes[2] * 9 + PFIndexes[3];
     unsigned PFEntry = PerfectShuffleTable[PFTableIndex];
     unsigned Cost = (PFEntry >> 30);
 
     if (Cost <= 4)
       return true;
   }
 
+  if (NumElts > 1 && NumElts % 2 != 0)
+    return false;
+
   bool DummyBool;
   int DummyInt;
   unsigned DummyUnsigned;
 
   return (ShuffleVectorSDNode::isSplatMask(&M[0], VT) || isREVMask(M, VT, 64) ||
           isREVMask(M, VT, 32) || isREVMask(M, VT, 16) ||
           isEXTMask(M, VT, DummyBool, DummyUnsigned) ||
           // isTBLMask(M, VT) || // FIXME: Port TBL support from ARM.

llvm/test/CodeGen/AArch64/arm64-neon-vector-shuffle-extract.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -verify-machineinstrs -mtriple=aarch64-unknown-linux -o - | FileCheck %s
+
+define void @test(i32* %p1, i32* %p2) {
+; CHECK-LABEL: test:
+; CHECK-NEXT:    mov w8, #3
+; CHECK-NEXT:    mov w9, #1
+; CHECK-NEXT:    str w8, [x0]
+; CHECK-NEXT:    str w9, [x1]
+; CHECK-NEXT:    ret
+  %tmp = shufflevector <1 x i32> <i32 1>, <1 x i32> undef, <3 x i32> <i32 0, i32 undef, i32 undef>
+  %tmp2 = shufflevector <3 x i32> <i32 2, i32 3, i32 4>, <3 x i32> %tmp, <3 x i32> <i32 0, i32 1, i32 3>
+  %tmp3 = shufflevector <3 x i32> %tmp2, <3 x i32> undef, <6 x i32> <i32 0, i32 1, i32 2, i32 undef, i32 undef, i32 undef>
+  %tmp4 = shufflevector <6 x i32> undef, <6 x i32> %tmp3, <9 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8>
+  %tmp6 = extractelement <9 x i32> %tmp4, i32 7
+  %tmp8 = extractelement <9 x i32> %tmp4, i32 8
+  store i32 %tmp6, i32* %p1, align 4
+  store i32 %tmp8, i32* %p2, align 4
+  ret void
+}