[ARM64]Fix a bug when lowing shuffle vector to EXT instruction

[• HaoLiu]

Summary

Hi Tim and other reviewers,

When lowing shuffle vector, there is a function called isEXTMask to check whether we can optimize it into a EXT. But that function misses on situation, as a result it will generate incorrect indexed EXT.
It considers about the mask starts with UNDEFs as following:

<-1, -1, 3, ...> means that an EXT must start at 3 - 2 = 1,

But if there are too many UNDEFs, such as <-1, -1, -1, 1, ...>, the index counter will be overflowed into a very large unsigned int (i.e. unsigned Imm = 1 - 3).

This patch fixes this bug. To make it looks more simpler, it also adjust the position of logic about operands traverse. Review please.

Thanks,
-Hao

[t.p.northover]

Hi Hao,

This logic looks horribly convoluted, given the problem it's trying to solve. I think it may be time to step back and rewrite it with UNDEF support from the beginning.

One possible idea: the width is going to be a power of 2, so if we made ExpectedElt an APInt of the correct width, any overflow might well do precisely what we *want*. Further, we then ReverseExt if and only if ExpectedElt *starts* in the second half (I think).

So the total logic might be:

FirstRealElt = ...
ExpectedElt = APInt(...);
FirstWrongElt = std::find_if(M.begin(), M.end(), [&](int Elt) { return Elt != ExpectedElt++ && Elt != -1; })
if (FirstWrongElt != M.end()) return false;
// N.b. ExpectedElt is now back to its original value via wrapping
if (ExpectedElt >= NumElts) {
  ReverseExt = true;
  ExpectedElt -= NumElts;
}
Imm = ExpectedElt;
return true;

How does that look?

Cheers.

Tim.

[• HaoLiu]

Hi Tim,

Thanks for your suggestion on using APInt and std::find_if. It is very simple and easy to be understood. I've refactored the patch.

BTW, there is a minor issue for this line

"Imm = ExpectedElt;"

It is only true when there is UNDEF in the beginning. If the mask starts from non-undef index, it should be

"Imm = M[0]"

For example, "shuffle V0, V1, <6, 7, 0, 1>" (V0, V1 are vectors of 4 elements), Imm should be "6" not "2".

Thanks,
-Hao

[t.p.northover]

Hi Hao,

BTW, there is a minor issue for this line

"Imm = ExpectedElt;"

It is only true when there is UNDEF in the beginning. If the mask starts from non-undef index, it should be

"Imm = M[0]"

For example, "shuffle V0, V1, <6, 7, 0, 1>" (V0, V1 are vectors of 4 elements), Imm should be "6" not "2".

Ah yes, I see. Well spotted! Actually, doesn't that apply even if M[0]
*is* -1: <-1 7, 0, 1> say? I think I'd been assuming the mask was 2 *
NumElts long for some reason.

So the logic we really want is to count back from *FirstRealElt.
Either of these look like they'd work:

1. *FirstRealElt - (FirstRealElt - M.begin())
2. ExpectedElt - NumElts

Cheers.

Tim.

[• HaoLiu]

Hi Tim,

Yes. As "(1). *FirstRealElt - (FirstRealElt - M.begin())" may be overflow, this patch implements in "(2). ExpectedElt - NumElts".

I think you mean this patch is OK to be committed. Committed in http://llvm.org/viewvc/llvm-project?view=revision&revision=207485.

Thanks,
-Hao

lib/Target/ARM64/ARM64ISelLowering.cpp

 // vector sources of the shuffle are different.
 static bool isEXTMask(ArrayRef<int> M, EVT VT, bool &ReverseEXT,
                       unsigned &Imm) {
-  unsigned NumElts = VT.getVectorNumElements();
-  ReverseEXT = false;
-
-  // Look for the first non-undef choice and count backwards from
-  // that. E.g. <-1, -1, 3, ...> means that an EXT must start at 3 - 2 = 1. This
-  // guarantees that at least one index is correct.
-  const int *FirstRealElt =
-      std::find_if(M.begin(), M.end(), [](int Elt) { return Elt >= 0; });
-  assert(FirstRealElt != M.end() && "Completely UNDEF shuffle? Why bother?");
-  Imm = *FirstRealElt - (FirstRealElt - M.begin());
+  // Look for the first non-undef element.
+  const int *FirstRealElt = std::find_if(M.begin(), M.end(),
+      [](int Elt) {return Elt >= 0;});
 
-  // If this is a VEXT shuffle, the immediate value is the index of the first
-  // element.  The other shuffle indices must be the successive elements after
-  // the first one.
-  unsigned ExpectedElt = Imm;
-  for (unsigned i = 1; i < NumElts; ++i) {
-    // Increment the expected index.  If it wraps around, it may still be
-    // a VEXT but the source vectors must be swapped.
-    ExpectedElt += 1;
-    if (ExpectedElt == NumElts * 2) {
-      ExpectedElt = 0;
-      ReverseEXT = true;
-    }
+  // Benefit form APInt to avoid overflow when calculating the expected element.
+  unsigned NumElts = VT.getVectorNumElements();
+  unsigned MaskBits = APInt(32, NumElts * 2).logBase2();
+  APInt ExpectedElt = APInt(MaskBits, *FirstRealElt + 1);
+  // The following shuffle indices must be the successive elements after the
+  // first real element.
+  const int *FirstWrongElt = std::find_if(FirstRealElt + 1, M.end(),
+      [&](int Elt) {return Elt != ExpectedElt++ && Elt != -1;});
+  if (FirstWrongElt != M.end())
+    return false;
 
-    if (M[i] < 0)
-      continue; // ignore UNDEF indices
-    if (ExpectedElt != static_cast<unsigned>(M[i]))
-      return false;
-  }
+  // The index of an EXT is the first element if it is not UNDEF.
+  // Watch out for the beginning UNDEFs. The EXT index should be the expected
+  // value of the first element.
+  // E.g. <-1, -1, 3, ...> is treated as <1, 2, 3, ...>.
+  //      <-1, -1, 0, 1, ...> is treated as <IDX, IDX+1, 0, 1, ...>. IDX is
+  // equal to the ExpectedElt. For this case, ExpectedElt is (NumElts*2 - 2).
+  Imm = (M[0] >= 0) ? static_cast<unsigned>(M[0]) : ExpectedElt.getZExtValue();
 
-  // Adjust the index value if the source operands will be swapped.
-  if (ReverseEXT)
+  if (Imm >= NumElts) {
+    ReverseEXT = true;
     Imm -= NumElts;
+  }
 
   return true;
 }

test/CodeGen/ARM64/ext.ll

   ret <8 x i8> %tmp3
 }
 
+define <8 x i8> @test_vextd_undef2(<8 x i8>* %A, <8 x i8>* %B) nounwind {
+;CHECK-LABEL: test_vextd_undef2:
+;CHECK: {{ext.8b.*#6}}
+  %tmp1 = load <8 x i8>* %A
+  %tmp2 = load <8 x i8>* %B
+  %tmp3 = shufflevector <8 x i8> %tmp1, <8 x i8> %tmp2, <8 x i32> <i32 undef, i32 undef, i32 undef, i32 undef, i32 2, i32 3, i32 4, i32 5>
+  ret <8 x i8> %tmp3
+}
+
 define <16 x i8> @test_vextRq_undef(<16 x i8>* %A, <16 x i8>* %B) nounwind {
 ;CHECK-LABEL: test_vextRq_undef:
 ;CHECK: {{ext.16b.*#7}}

   ret <16 x i8> %tmp3
 }
 
+define <8 x i16> @test_vextRq_undef2(<8 x i16>* %A) nounwind {
+;CHECK-LABEL: test_vextRq_undef2:
+;CHECK: {{ext.16b.*#10}}
+  %tmp1 = load <8 x i16>* %A
+  %vext = shufflevector <8 x i16> %tmp1, <8 x i16> undef, <8 x i32> <i32 undef, i32 undef, i32 undef, i32 undef, i32 1, i32 2, i32 3, i32 4>
+  ret <8 x i16> %vext;
+}
+
 ; Tests for ReconstructShuffle function. Indices have to be carefully
 ; chosen to reach lowering phase as a BUILD_VECTOR.