[ARM] Check the right order for halves of VZIP/VUZP if both parts are used

[mstorsjo]

Summary

This is the exact same fix as in SVN r247254. In that commit, the fix was applied only for isVTRNMask and isVTRN_v_undef_Mask, but the same issue is present for VZIP/VUZP as well.

This fixes PR33921.

[mstorsjo]

Understood the root cause better, with the hidden assumptions of isVZIPMask where both outputs are used.

[mstorsjo]

After digging in further, I found an even more elegant way of fixing it. And when expanding that fix to isVUZPMask and isVTRNMask, I noticed that the exact same fix as I was trying to provide already was present in isVTRNMask. So expanding that fix to all of them now, with the same code formatting style as in the existing check.

[rengolin]

Hi Martin,

This makes sense to me. Thanks for the great investigation and getting to the bottom of it with such a simple patch. LGTM.

cheers,
--renato

[compnerd]

Do you think it may be worthwhile to extract the operation into a function?

static unsigned SelectPairHalf(unsigned Elements, ArrayRef<int> Mask, unsigned Index) {
  if (Mask.size() == Elements * 2)
    return Index / Elements;
  return Mask[Index] == 0 ? 0 : 1;
}

[mstorsjo]

In D36899#846598, @compnerd wrote:

Do you think it may be worthwhile to extract the operation into a function?

static unsigned SelectPairHalf(unsigned Elements, ArrayRef<int> Mask, unsigned Index) {
  if (Mask.size() == Elements * 2)
    return Index / Elements;
  return Mask[Index] == 0 ? 0 : 1;
}

Hmm, maybe. I want this backported to 5.0 though, so I'd rather not to make too much factorization in the fix itself. Perhaps a separate factorization commit on top afterwards, that doesn't have to be backported?

[compnerd]

SGTM

llvm/trunk/lib/Target/ARM/ARMISelLowering.cpp

static bool isVUZPMask(ArrayRef<int> M, EVT VT, unsigned &WhichResult) {
   if (EltSz == 64)
     return false;
 
   unsigned NumElts = VT.getVectorNumElements();
   if (M.size() != NumElts && M.size() != NumElts*2)
     return false;
 
   for (unsigned i = 0; i < M.size(); i += NumElts) {
+    if (M.size() == NumElts * 2)
+      WhichResult = i / NumElts;
+    else
       WhichResult = M[i] == 0 ? 0 : 1;
     for (unsigned j = 0; j < NumElts; ++j) {
       if (M[i+j] >= 0 && (unsigned) M[i+j] != 2 * j + WhichResult)
         return false;
     }
   }
 
   if (M.size() == NumElts*2)
     WhichResult = 0;
if (EltSz == 64)
     return false;
 
   unsigned NumElts = VT.getVectorNumElements();
   if (M.size() != NumElts && M.size() != NumElts*2)
     return false;
 
   unsigned Half = NumElts / 2;
   for (unsigned i = 0; i < M.size(); i += NumElts) {
+    if (M.size() == NumElts * 2)
+      WhichResult = i / NumElts;
+    else
       WhichResult = M[i] == 0 ? 0 : 1;
     for (unsigned j = 0; j < NumElts; j += Half) {
       unsigned Idx = WhichResult;
       for (unsigned k = 0; k < Half; ++k) {
         int MIdx = M[i + j + k];
         if (MIdx >= 0 && (unsigned) MIdx != Idx)
           return false;
         Idx += 2;
       }
static bool isVZIPMask(ArrayRef<int> M, EVT VT, unsigned &WhichResult) {
   if (EltSz == 64)
     return false;
 
   unsigned NumElts = VT.getVectorNumElements();
   if (M.size() != NumElts && M.size() != NumElts*2)
     return false;
 
   for (unsigned i = 0; i < M.size(); i += NumElts) {
+    if (M.size() == NumElts * 2)
+      WhichResult = i / NumElts;
+    else
       WhichResult = M[i] == 0 ? 0 : 1;
     unsigned Idx = WhichResult * NumElts / 2;
     for (unsigned j = 0; j < NumElts; j += 2) {
       if ((M[i+j] >= 0 && (unsigned) M[i+j] != Idx) ||
           (M[i+j+1] >= 0 && (unsigned) M[i+j+1] != Idx + NumElts))
         return false;
       Idx += 1;
     }
   }
static bool isVZIP_v_undef_Mask(ArrayRef<int> M, EVT VT, unsigned &WhichResult){
   if (EltSz == 64)
     return false;
 
   unsigned NumElts = VT.getVectorNumElements();
   if (M.size() != NumElts && M.size() != NumElts*2)
     return false;
 
   for (unsigned i = 0; i < M.size(); i += NumElts) {
+    if (M.size() == NumElts * 2)
+      WhichResult = i / NumElts;
+    else
       WhichResult = M[i] == 0 ? 0 : 1;
     unsigned Idx = WhichResult * NumElts / 2;
     for (unsigned j = 0; j < NumElts; j += 2) {
       if ((M[i+j] >= 0 && (unsigned) M[i+j] != Idx) ||
           (M[i+j+1] >= 0 && (unsigned) M[i+j+1] != Idx))
         return false;
       Idx += 1;
     }
   }

llvm/trunk/test/CodeGen/ARM/vzip.ll

 ; CHECK-NEXT:    mov pc, lr
 entry:
   %tmp1 = load <4 x i16>, <4 x i16>* %A
   %tmp2 = load <4 x i16>, <4 x i16>* %B
   %0 = shufflevector <4 x i16> %tmp1, <4 x i16> %tmp2, <8 x i32> <i32 undef, i32 undef, i32 undef, i32 undef, i32 2, i32 6, i32 3, i32 7>
   ret <8 x i16> %0
 }
 
+; NOTE: The mask here looks like something that could be done with a vzip,
+; but which the current handling of two-result vzip can't do - thus ending up
+; as a vtrn.
+define <8 x i16> @vzip_lower_shufflemask_undef_rev(<4 x i16>* %A, <4 x i16>* %B) {
+; CHECK-LABEL: vzip_lower_shufflemask_undef_rev:
+; CHECK:       @ BB#0: @ %entry
+; CHECK-NEXT:    vldr d16, [r1]
+; CHECK-NEXT:    vldr d19, [r0]
+; CHECK-NEXT:    vtrn.16 d19, d16
+; CHECK-NEXT:    vmov r0, r1, d18
+; CHECK-NEXT:    vmov r2, r3, d19
+; CHECK-NEXT:    mov pc, lr
+entry:
+  %tmp1 = load <4 x i16>, <4 x i16>* %A
+  %tmp2 = load <4 x i16>, <4 x i16>* %B
+  %0 = shufflevector <4 x i16> %tmp1, <4 x i16> %tmp2, <8 x i32> <i32 undef, i32 undef, i32 undef, i32 undef, i32 0, i32 4, i32 undef, i32 undef>
+  ret <8 x i16> %0
+}
+
 define <4 x i32> @vzip_lower_shufflemask_zeroed(<2 x i32>* %A) {
 ; CHECK-LABEL: vzip_lower_shufflemask_zeroed:
 ; CHECK:       @ BB#0: @ %entry
 ; CHECK-NEXT:    vldr d16, [r0]
 ; CHECK-NEXT:    vdup.32 q9, d16[0]
 ; CHECK-NEXT:    vzip.32 q8, q9
 ; CHECK-NEXT:    vmov r0, r1, d16
 ; CHECK-NEXT:    vmov r2, r3, d17