AArch64/ARM64: improve spotting of EXT instructions from VECTOR_SHUFFLE.

We couldn't cope if the first mask element was UNDEF before, which
isn't ideal.

llvm-svn: 206588

1 files changed, 7 insertions, 5 deletions

diff --git a/llvm/lib/Target/ARM64/ARM64ISelLowering.cpp b/llvm/lib/Target/ARM64/ARM64ISelLowering.cpp
index efd979b5e2e..503e44b0dc0 100644
--- a/llvm/lib/Target/ARM64/ARM64ISelLowering.cpp
+++ b/llvm/lib/Target/ARM64/ARM64ISelLowering.cpp
@@ -3945,11 +3945,13 @@ static bool isEXTMask(ArrayRef<int> M, EVT VT, bool &ReverseEXT,
   unsigned NumElts = VT.getVectorNumElements();
   ReverseEXT = false;
 
-  // Assume that the first shuffle index is not UNDEF.  Fail if it is.
-  if (M[0] < 0)
-    return false;
-
-  Imm = M[0];
+  // 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());
 
   // 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