[AArch64] match TRN starting from undef elements

When the first element of a trn mask is undef, the isTRNMask function
assumes that the value for "WhichResult" should be 1. That has a 50%
chance of being wrong, so we fail to match some valid trn1/trn2.
This patch introduces a more precise test to determine the correct
value of "WhichResult", based on corresponding code in the isZIPMask
and isUZPMask functions.

Author: ginsbach

llvm/lib/Target/AArch64/AArch64PerfectShuffle.h

@@ -6685,15 +6685,30 @@ inline bool isUZPMask(ArrayRef<int> M, unsigned NumElts,
 ///  <0, 8, 2, 10, 4, 12, 6, 14> or
 ///  <1, 9, 3, 11, 5, 13, 7, 15>
 inline bool isTRNMask(ArrayRef<int> M, unsigned NumElts,
-                      unsigned &WhichResult) {
+                      unsigned &WhichResultOut) {
   if (NumElts % 2 != 0)
     return false;
-  WhichResult = (M[0] == 0 ? 0 : 1);
+  // Check the first non-undef element for trn1 vs trn2.
+  unsigned WhichResult = 2;
+  for (unsigned i = 0; i != NumElts; i += 2) {
+    if (M[i] >= 0) {
+      WhichResult = ((unsigned)M[i] == i ? 0 : 1);
+      break;
+    }
+    if (M[i + 1] >= 0) {
+      WhichResult = ((unsigned)M[i + 1] == i + NumElts ? 0 : 1);
+      break;
+    }
+  }
+  if (WhichResult == 2)
+    return false;
+
   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;
   }
+  WhichResultOut = WhichResult;
   return true;
 }
 

llvm/test/CodeGen/AArch64/insert-extend.ll

@@ -85,24 +85,24 @@ define i32 @large(ptr nocapture noundef readonly %p1, i32 noundef %st1, ptr noca
 ; CHECK-NEXT:    addp v2.4s, v3.4s, v2.4s
 ; CHECK-NEXT:    zip1 v16.4s, v5.4s, v4.4s
 ; CHECK-NEXT:    sub v7.4s, v3.4s, v7.4s
+; CHECK-NEXT:    trn1 v4.4s, v5.4s, v4.4s
 ; CHECK-NEXT:    zip2 v3.4s, v6.4s, v7.4s
 ; CHECK-NEXT:    mov v6.s[1], v7.s[0]
-; CHECK-NEXT:    ext v7.16b, v5.16b, v16.16b, #8
-; CHECK-NEXT:    mov v5.s[3], v4.s[2]
-; CHECK-NEXT:    ext v4.16b, v2.16b, v2.16b, #8
-; CHECK-NEXT:    mov v6.d[1], v7.d[1]
-; CHECK-NEXT:    mov v3.d[1], v5.d[1]
-; CHECK-NEXT:    uzp1 v1.4s, v4.4s, v0.4s
-; CHECK-NEXT:    uzp2 v4.4s, v4.4s, v0.4s
+; CHECK-NEXT:    ext v7.16b, v2.16b, v2.16b, #8
+; CHECK-NEXT:    ext v5.16b, v5.16b, v16.16b, #8
+; CHECK-NEXT:    mov v3.d[1], v4.d[1]
+; CHECK-NEXT:    uzp1 v1.4s, v7.4s, v0.4s
+; CHECK-NEXT:    uzp2 v4.4s, v7.4s, v0.4s
+; CHECK-NEXT:    mov v6.d[1], v5.d[1]
 ; CHECK-NEXT:    addp v0.4s, v2.4s, v0.4s
+; CHECK-NEXT:    sub v1.4s, v1.4s, v4.4s
+; CHECK-NEXT:    rev64 v7.4s, v0.4s
 ; CHECK-NEXT:    add v5.4s, v3.4s, v6.4s
 ; CHECK-NEXT:    sub v3.4s, v6.4s, v3.4s
-; CHECK-NEXT:    rev64 v7.4s, v0.4s
-; CHECK-NEXT:    sub v1.4s, v1.4s, v4.4s
+; CHECK-NEXT:    rev64 v2.4s, v1.4s
 ; CHECK-NEXT:    rev64 v4.4s, v5.4s
 ; CHECK-NEXT:    rev64 v6.4s, v3.4s
 ; CHECK-NEXT:    addp v16.4s, v0.4s, v5.4s
-; CHECK-NEXT:    rev64 v2.4s, v1.4s
 ; CHECK-NEXT:    sub v0.4s, v0.4s, v7.4s
 ; CHECK-NEXT:    zip1 v21.4s, v16.4s, v16.4s
 ; CHECK-NEXT:    sub v4.4s, v5.4s, v4.4s