[llvm] [AArch64] isTBLMask(M, VT) as part of the shuffle mask check (PR #79058)

[via llvm-commits]

https://github.com/AtariDreams updated https://github.com/llvm/llvm-project/pull/79058

>From ae9b9337dc81d1425f08567a0443c42cde95e53d Mon Sep 17 00:00:00 2001
From: Rose <83477269+AtariDreams at users.noreply.github.com>
Date: Sat, 20 Jan 2024 22:05:15 -0500
Subject: [PATCH] [AArch64] isTBLMask(M, VT) as part of the shuffle mask check

---
 .../Target/AArch64/AArch64ISelLowering.cpp    | 19 ++++++++++++++-----
 llvm/lib/Target/ARM/ARMTargetTransformInfo.h  |  3 ++-
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 332fb37655288c..8ae758603089aa 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -11471,6 +11471,8 @@ static bool isWideDUPMask(ArrayRef<int> M, EVT VT, unsigned BlockSize,
 // vector sources of the shuffle are different.
 static bool isEXTMask(ArrayRef<int> M, EVT VT, bool &ReverseEXT,
                       unsigned &Imm) {
+  ReverseEXT = false;
+
   // Look for the first non-undef element.
   const int *FirstRealElt = find_if(M, [](int Elt) { return Elt >= 0; });
 
@@ -11508,6 +11510,14 @@ static bool isEXTMask(ArrayRef<int> M, EVT VT, bool &ReverseEXT,
   return true;
 }
 
+static bool isTBLMask(ArrayRef<int> M, EVT VT) {
+  // We can handle <16 x i8> and <8 x i8> vector shuffles. If the index in the
+  // mask is out of range, then 0 is placed into the resulting vector. So pretty
+  // much any mask of 16 or 8 elements can work here.
+  return (VT == MVT::v8i8 && M.size() == 8) ||
+         (VT == MVT::v16i8 && M.size() == 16);
+}
+
 /// isREVMask - Check if a vector shuffle corresponds to a REV
 /// instruction with the specified blocksize.  (The order of the elements
 /// within each block of the vector is reversed.)
@@ -11542,7 +11552,7 @@ static bool isZIPMask(ArrayRef<int> M, EVT VT, unsigned &WhichResult) {
     return false;
   WhichResult = (M[0] == 0 ? 0 : 1);
   unsigned Idx = WhichResult * NumElts / 2;
-  for (unsigned i = 0; i != NumElts; i += 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 + NumElts))
       return false;
@@ -12257,7 +12267,7 @@ SDValue AArch64TargetLowering::LowerVECTOR_SHUFFLE(SDValue Op,
                        DAG.getConstant(8, dl, MVT::i32));
   }
 
-  bool ReverseEXT = false;
+  bool ReverseEXT;
   unsigned Imm;
   if (isEXTMask(ShuffleMask, VT, ReverseEXT, Imm)) {
     if (ReverseEXT)
@@ -13755,8 +13765,7 @@ bool AArch64TargetLowering::isShuffleMaskLegal(ArrayRef<int> M, EVT VT) const {
 
   return (ShuffleVectorSDNode::isSplatMask(&M[0], VT) || isREVMask(M, VT, 64) ||
           isREVMask(M, VT, 32) || isREVMask(M, VT, 16) ||
-          isEXTMask(M, VT, DummyBool, DummyUnsigned) ||
-          // isTBLMask(M, VT) || // FIXME: Port TBL support from ARM.
+          isEXTMask(M, VT, DummyBool, DummyUnsigned) || isTBLMask(M, VT) ||
           isTRNMask(M, VT, DummyUnsigned) || isUZPMask(M, VT, DummyUnsigned) ||
           isZIPMask(M, VT, DummyUnsigned) ||
           isTRN_v_undef_Mask(M, VT, DummyUnsigned) ||
@@ -26641,7 +26650,7 @@ SDValue AArch64TargetLowering::LowerFixedLengthVECTOR_SHUFFLEToSVE(
     return convertFromScalableVector(DAG, VT, Op);
   }
 
-  bool ReverseEXT = false;
+  bool ReverseEXT;
   unsigned Imm;
   if (isEXTMask(ShuffleMask, VT, ReverseEXT, Imm) &&
       Imm == VT.getVectorNumElements() - 1) {
diff --git a/llvm/lib/Target/ARM/ARMTargetTransformInfo.h b/llvm/lib/Target/ARM/ARMTargetTransformInfo.h
index bb4b321b530091..e83d8b830a43cd 100644
--- a/llvm/lib/Target/ARM/ARMTargetTransformInfo.h
+++ b/llvm/lib/Target/ARM/ARMTargetTransformInfo.h
@@ -347,6 +347,7 @@ inline bool isVREVMask(ArrayRef<int> M, EVT VT, unsigned BlockSize) {
   if (EltSz != 8 && EltSz != 16 && EltSz != 32)
     return false;
 
+  unsigned NumElts = VT.getVectorNumElements();
   unsigned BlockElts = M[0] + 1;
   // If the first shuffle index is UNDEF, be optimistic.
   if (M[0] < 0)
@@ -355,7 +356,7 @@ inline bool isVREVMask(ArrayRef<int> M, EVT VT, unsigned BlockSize) {
   if (BlockSize <= EltSz || BlockSize != BlockElts * EltSz)
     return false;
 
-  for (unsigned i = 0, e = M.size(); i < e; ++i) {
+  for (unsigned i = 0; i < NumElts; ++i) {
     if (M[i] < 0)
       continue; // ignore UNDEF indices
     if ((unsigned)M[i] != (i - i % BlockElts) + (BlockElts - 1 - i % BlockElts))