[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 ↗

(On Diff #195144)

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

7538 ↗

(On Diff #197947)

Unrelated change, I'll undo before committing.

[dmgreen]

LGTM. Thanks for making the fix.

llvm/trunk/lib/Target/AArch64/AArch64ISelLowering.cpp

if ((unsigned)M[i] != 2 * i + WhichResult)
       return false;
   }
 
   return true;
 }
 
 static bool isTRNMask(ArrayRef<int> M, EVT VT, unsigned &WhichResult) {
   unsigned NumElts = VT.getVectorNumElements();
+  if (NumElts % 2 != 0)
+    return false;
   WhichResult = (M[0] == 0 ? 0 : 1);
   for (unsigned i = 0; i < NumElts; i += 2) {
     if ((M[i] >= 0 && (unsigned)M[i] != i + WhichResult) ||
         (M[i + 1] >= 0 && (unsigned)M[i + 1] != i + NumElts + WhichResult))
       return false;
   }
   return true;
 }
 
 /// isZIP_v_undef_Mask - Special case of isZIPMask for canonical form of
 /// "vector_shuffle v, v", i.e., "vector_shuffle v, undef".
 /// Mask is e.g., <0, 0, 1, 1> instead of <0, 4, 1, 5>.
 static bool isZIP_v_undef_Mask(ArrayRef<int> M, EVT VT, unsigned &WhichResult) {
   unsigned NumElts = VT.getVectorNumElements();
+  if (NumElts % 2 != 0)
+    return false;
   WhichResult = (M[0] == 0 ? 0 : 1);
   unsigned Idx = WhichResult * NumElts / 2;
   for (unsigned i = 0; i != NumElts; i += 2) {
     if ((M[i] >= 0 && (unsigned)M[i] != Idx) ||
         (M[i + 1] >= 0 && (unsigned)M[i + 1] != Idx))
       return false;
     Idx += 1;
   }
static bool isUZP_v_undef_Mask(ArrayRef<int> M, EVT VT, unsigned &WhichResult) {
   return true;
 }
 
 /// isTRN_v_undef_Mask - Special case of isTRNMask for canonical form of
 /// "vector_shuffle v, v", i.e., "vector_shuffle v, undef".
 /// Mask is e.g., <0, 0, 2, 2> instead of <0, 4, 2, 6>.
 static bool isTRN_v_undef_Mask(ArrayRef<int> M, EVT VT, unsigned &WhichResult) {
   unsigned NumElts = VT.getVectorNumElements();
+  if (NumElts % 2 != 0)
+    return false;
   WhichResult = (M[0] == 0 ? 0 : 1);
   for (unsigned i = 0; i < NumElts; i += 2) {
     if ((M[i] >= 0 && (unsigned)M[i] != i + WhichResult) ||
         (M[i + 1] >= 0 && (unsigned)M[i + 1] != i + WhichResult))
       return false;
   }
   return true;
 }

llvm/trunk/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:       // %bb.0:
+; 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
+}
+
+; Test case from PR41951
+define <4 x i32> @widen_shuffles_reduced(<3 x i32> %x, <3 x i32> %y) {
+; CHECK-LABEL: widen_shuffles_reduced:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    zip1 v2.4s, v0.4s, v1.4s
+; CHECK-NEXT:    zip1 v0.4s, v1.4s, v0.4s
+; CHECK-NEXT:    ext v0.16b, v0.16b, v2.16b, #8
+; CHECK-NEXT:    ret
+  %s3 = shufflevector <3 x i32> %y, <3 x i32> %x, <4 x i32> <i32 1, i32 4, i32 3, i32 0>
+  ret <4 x i32> %s3
+}