[llvm] [AArch64] Refactor creation of a shuffle mask for TBL (NFC) (PR #92529)

[Momchil Velikov via llvm-commits]

================
@@ -15710,48 +15710,51 @@ bool AArch64TargetLowering::shouldSinkOperands(
   return false;
 }
 
-static bool createTblShuffleForZExt(ZExtInst *ZExt, FixedVectorType *DstTy,
-                                    bool IsLittleEndian) {
-  Value *Op = ZExt->getOperand(0);
-  auto *SrcTy = cast<FixedVectorType>(Op->getType());
-  auto SrcWidth = cast<IntegerType>(SrcTy->getElementType())->getBitWidth();
-  auto DstWidth = cast<IntegerType>(DstTy->getElementType())->getBitWidth();
+static bool createTblShuffleMask(unsigned SrcWidth, unsigned DstWidth,
+                                 unsigned NumElts, bool IsLittleEndian,
+                                 SmallVectorImpl<int> &Mask) {
   if (DstWidth % 8 != 0 || DstWidth <= 16 || DstWidth >= 64)
     return false;
 
-  assert(DstWidth % SrcWidth == 0 &&
-         "TBL lowering is not supported for a ZExt instruction with this "
-         "source & destination element type.");
-  unsigned ZExtFactor = DstWidth / SrcWidth;
+  if (DstWidth % SrcWidth != 0)
+    return false;
+
+  unsigned Factor = DstWidth / SrcWidth;
+  unsigned MaskLen = NumElts * Factor;
+
+  Mask.clear();
+  Mask.resize(MaskLen, NumElts);
+
+  unsigned SrcIndex = 0;
+  for (unsigned I = 0; I < MaskLen; I += Factor)
+    Mask[I] = SrcIndex++;
+
+  if (!IsLittleEndian)
+    std::rotate(Mask.rbegin(), Mask.rbegin() + Factor - 1, Mask.rend());
----------------
momchil-velikov wrote:

Eh, this could be a rotate by 1 in the opposite direction, will fix later.

https://github.com/llvm/llvm-project/pull/92529