AArch64PerfectShuffle.h

Go to the documentation of this file.

//===-- AArch64PerfectShuffle.h - AdvSIMD Perfect Shuffle Table -----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file declares data for the optimal way to build a perfect shuffle using
// AdvSIMD instructions. The data is generated by llvm-PerfectShuffle.
//
//===----------------------------------------------------------------------===//
 
#ifndef LLVM_LIB_TARGET_AARCH64_AARCH64PERFECTSHUFFLE_H
#define LLVM_LIB_TARGET_AARCH64_AARCH64PERFECTSHUFFLE_H
 
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
 
namespace llvm {
 
extern const unsigned PerfectShuffleTable[6561 + 1];
 
inline unsigned getPerfectShuffleCost(llvm::ArrayRef<int> M) {
  assert(M.size() == 4 && "Expected a 4 entry perfect shuffle");
 
  // Special case zero-cost nop copies, from either LHS or RHS.
  if (llvm::all_of(llvm::enumerate(M), [](const auto &E) {
        return E.value() < 0 || E.value() == (int)E.index();
      }))
    return 0;
  if (llvm::all_of(llvm::enumerate(M), [](const auto &E) {
        return E.value() < 0 || E.value() == (int)E.index() + 4;
      }))
    return 0;
 
  // Get the four mask elementd from the 2 inputs. Perfect shuffles encode undef
  // elements with value 8.
  unsigned PFIndexes[4];
  for (unsigned i = 0; i != 4; ++i) {
    assert(M[i] < 8 && "Expected a maximum entry of 8 for shuffle mask");
    if (M[i] < 0)
      PFIndexes[i] = 8;
    else
      PFIndexes[i] = M[i];
  }
 
  // Compute the index in the perfect shuffle table.
  unsigned PFTableIndex = PFIndexes[0] * 9 * 9 * 9 + PFIndexes[1] * 9 * 9 +
                          PFIndexes[2] * 9 + PFIndexes[3];
  unsigned PFEntry = PerfectShuffleTable[PFTableIndex];
  // And extract the cost from the upper bits. The cost is encoded as Cost-1.
  return (PFEntry >> 30) + 1;
}
 
/// Return true for zip1 or zip2 masks of the form:
///  <0,  8, 1,  9, 2, 10, 3, 11> (WhichResultOut = 0, OperandOrderOut = 0) or
///  <4, 12, 5, 13, 6, 14, 7, 15> (WhichResultOut = 1, OperandOrderOut = 0) or
///  <8,  0, 9,  1, 10, 2, 11, 3> (WhichResultOut = 0, OperandOrderOut = 1) or
///  <12, 4, 13, 5, 14, 6, 15, 7> (WhichResultOut = 1, OperandOrderOut = 1)
inline bool isZIPMask(ArrayRef<int> M, unsigned NumElts,
                      unsigned &WhichResultOut, unsigned &OperandOrderOut) {
  if (NumElts % 2 != 0)
    return false;
 
  // "Result" corresponds to "WhichResultOut", selecting between zip1 and zip2.
  // "Order" corresponds to "OperandOrderOut", selecting the order of operands
  // for the instruction (flipped or not).
  bool Result0Order0 = true; // WhichResultOut = 0, OperandOrderOut = 0
  bool Result1Order0 = true; // WhichResultOut = 1, OperandOrderOut = 0
  bool Result0Order1 = true; // WhichResultOut = 0, OperandOrderOut = 1
  bool Result1Order1 = true; // WhichResultOut = 1, OperandOrderOut = 1
  // Check all elements match.
  for (unsigned i = 0; i != NumElts; i += 2) {
    if (M[i] >= 0) {
      unsigned EvenElt = (unsigned)M[i];
      if (EvenElt != i / 2)
        Result0Order0 = false;
      if (EvenElt != NumElts / 2 + i / 2)
        Result1Order0 = false;
      if (EvenElt != NumElts + i / 2)
        Result0Order1 = false;
      if (EvenElt != NumElts + NumElts / 2 + i / 2)
        Result1Order1 = false;
    }
    if (M[i + 1] >= 0) {
      unsigned OddElt = (unsigned)M[i + 1];
      if (OddElt != NumElts + i / 2)
        Result0Order0 = false;
      if (OddElt != NumElts + NumElts / 2 + i / 2)
        Result1Order0 = false;
      if (OddElt != i / 2)
        Result0Order1 = false;
      if (OddElt != NumElts / 2 + i / 2)
        Result1Order1 = false;
    }
  }
 
  if (Result0Order0 + Result1Order0 + Result0Order1 + Result1Order1 != 1)
    return false;
 
  WhichResultOut = (Result0Order0 || Result0Order1) ? 0 : 1;
  OperandOrderOut = (Result0Order0 || Result1Order0) ? 0 : 1;
  return true;
}
 
/// isZIP_v_undef_Mask - Special case of isZIPMask for canonical form of
/// "vector_shuffle v, v", i.e., "vector_shuffle v, undef".
/// Mask is e.g., <0, 0, 1, 1> instead of <0, 4, 1, 5>.
inline bool isZIP_v_undef_Mask(ArrayRef<int> M, unsigned NumElts,
                               unsigned &WhichResult) {
  if (NumElts % 2 != 0)
    return false;
  WhichResult = (M[0] == 0 ? 0 : 1);
  unsigned Idx = WhichResult * NumElts / 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))
      return false;
    Idx += 1;
  }
 
  return true;
}
 
/// Return true for uzp1 or uzp2 masks of the form:
///  <0, 2, 4, 6, 8, 10, 12, 14> or
///  <1, 3, 5, 7, 9, 11, 13, 15>
inline bool isUZPMask(ArrayRef<int> M, unsigned NumElts,
                      unsigned &WhichResultOut) {
  // Check the first non-undef element for which half to use.
  unsigned WhichResult = 2;
  for (unsigned i = 0; i != NumElts; i++) {
    if (M[i] >= 0) {
      WhichResult = ((unsigned)M[i] == i * 2 ? 0 : 1);
      break;
    }
  }
  if (WhichResult == 2)
    return false;
 
  // Check all elements match.
  for (unsigned i = 0; i != NumElts; ++i) {
    if (M[i] < 0)
      continue; // ignore UNDEF indices
    if ((unsigned)M[i] != 2 * i + WhichResult)
      return false;
  }
  WhichResultOut = WhichResult;
  return true;
}
 
/// isUZP_v_undef_Mask - Special case of isUZPMask for canonical form of
/// "vector_shuffle v, v", i.e., "vector_shuffle v, undef".
/// Mask is e.g., <0, 2, 0, 2> instead of <0, 2, 4, 6>,
inline bool isUZP_v_undef_Mask(ArrayRef<int> M, unsigned NumElts,
                               unsigned &WhichResult) {
  unsigned Half = NumElts / 2;
  WhichResult = (M[0] == 0 ? 0 : 1);
  for (unsigned j = 0; j != 2; ++j) {
    unsigned Idx = WhichResult;
    for (unsigned i = 0; i != Half; ++i) {
      int MIdx = M[i + j * Half];
      if (MIdx >= 0 && (unsigned)MIdx != Idx)
        return false;
      Idx += 2;
    }
  }
 
  return true;
}
 
/// Return true for trn1 or trn2 masks of the form:
///  <0, 8, 2, 10, 4, 12, 6, 14> (WhichResultOut = 0, OperandOrderOut = 0) or
///  <1, 9, 3, 11, 5, 13, 7, 15> (WhichResultOut = 1, OperandOrderOut = 0) or
///  <8, 0, 10, 2, 12, 4, 14, 6> (WhichResultOut = 0, OperandOrderOut = 1) or
///  <9, 1, 11, 3, 13, 5, 15, 7> (WhichResultOut = 1, OperandOrderOut = 1) or
inline bool isTRNMask(ArrayRef<int> M, unsigned NumElts,
                      unsigned &WhichResultOut, unsigned &OperandOrderOut) {
  if (NumElts % 2 != 0)
    return false;
 
  // "Result" corresponds to "WhichResultOut", selecting between trn1 and trn2.
  // "Order" corresponds to "OperandOrderOut", selecting the order of operands
  // for the instruction (flipped or not).
  bool Result0Order0 = true; // WhichResultOut = 0, OperandOrderOut = 0
  bool Result1Order0 = true; // WhichResultOut = 1, OperandOrderOut = 0
  bool Result0Order1 = true; // WhichResultOut = 0, OperandOrderOut = 1
  bool Result1Order1 = true; // WhichResultOut = 1, OperandOrderOut = 1
  // Check all elements match.
  for (unsigned i = 0; i != NumElts; i += 2) {
    if (M[i] >= 0) {
      unsigned EvenElt = (unsigned)M[i];
      if (EvenElt != i)
        Result0Order0 = false;
      if (EvenElt != i + 1)
        Result1Order0 = false;
      if (EvenElt != NumElts + i)
        Result0Order1 = false;
      if (EvenElt != NumElts + i + 1)
        Result1Order1 = false;
    }
    if (M[i + 1] >= 0) {
      unsigned OddElt = (unsigned)M[i + 1];
      if (OddElt != NumElts + i)
        Result0Order0 = false;
      if (OddElt != NumElts + i + 1)
        Result1Order0 = false;
      if (OddElt != i)
        Result0Order1 = false;
      if (OddElt != i + 1)
        Result1Order1 = false;
    }
  }
 
  if (Result0Order0 + Result1Order0 + Result0Order1 + Result1Order1 != 1)
    return false;
 
  WhichResultOut = (Result0Order0 || Result0Order1) ? 0 : 1;
  OperandOrderOut = (Result0Order0 || Result1Order0) ? 0 : 1;
  return true;
}
 
/// isTRN_v_undef_Mask - Special case of isTRNMask for canonical form of
/// "vector_shuffle v, v", i.e., "vector_shuffle v, undef".
/// Mask is e.g., <0, 0, 2, 2> instead of <0, 4, 2, 6>.
inline bool isTRN_v_undef_Mask(ArrayRef<int> M, unsigned NumElts,
                               unsigned &WhichResult) {
  if (NumElts % 2 != 0)
    return false;
  WhichResult = (M[0] == 0 ? 0 : 1);
  for (unsigned i = 0; i < NumElts; i += 2) {
    if ((M[i] >= 0 && (unsigned)M[i] != i + WhichResult) ||
        (M[i + 1] >= 0 && (unsigned)M[i + 1] != i + WhichResult))
      return false;
  }
  return true;
}
 
/// 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.)
inline bool isREVMask(ArrayRef<int> M, unsigned EltSize, unsigned NumElts,
                      unsigned BlockSize) {
  assert((BlockSize == 16 || BlockSize == 32 || BlockSize == 64 ||
          BlockSize == 128) &&
         "Only possible block sizes for REV are: 16, 32, 64, 128");
 
  unsigned BlockElts = M[0] + 1;
  // If the first shuffle index is UNDEF, be optimistic.
  if (M[0] < 0)
    BlockElts = BlockSize / EltSize;
 
  if (BlockSize <= EltSize || BlockSize != BlockElts * EltSize)
    return false;
 
  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))
      return false;
  }
 
  return true;
}
 
/// Check if an EXT instruction can handle the shuffle mask when the
/// vector sources of the shuffle are different.
inline bool isEXTMask(ArrayRef<int> M, unsigned NumElts, bool &ReverseEXT,
                      unsigned &Imm) {
  // Look for the first non-undef element.
  const int *FirstRealElt = find_if(M, [](int Elt) { return Elt >= 0; });
  if (FirstRealElt == M.end())
    return false;
 
  // Benefit from APInt to handle overflow when calculating expected element.
  unsigned MaskBits = APInt(32, NumElts * 2).logBase2();
  APInt ExpectedElt = APInt(MaskBits, *FirstRealElt + 1, /*isSigned=*/false,
                            /*implicitTrunc=*/true);
  // The following shuffle indices must be the successive elements after the
  // first real element.
  if (std::any_of(FirstRealElt + 1, M.end(),
                  [&](int Elt) { return Elt != ExpectedElt++ && Elt >= 0; }))
    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 <2*NumElts-2, 2*NumElts-1, 0, 1, ...>.
  // ExpectedElt is the last mask index plus 1.
  Imm = ExpectedElt.getZExtValue();
 
  // There are two difference cases requiring to reverse input vectors.
  // For example, for vector <4 x i32> we have the following cases,
  // Case 1: shufflevector(<4 x i32>, <4 x i32>, <-1, -1, -1, 0>)
  // Case 2: shufflevector(<4 x i32>, <4 x i32>, <-1, -1, 7, 0>)
  // For both cases, we finally use mask <5, 6, 7, 0>, which requires
  // to reverse two input vectors.
  ReverseEXT = false;
  if (Imm < NumElts)
    ReverseEXT = true;
  else
    Imm -= NumElts;
 
  return true;
}
 
/// check if an EXT instruction can handle the shuffle mask when the
/// vector sources of the shuffle are the same. i.e <2,3,0,1>
inline bool isSingletonEXTMask(ArrayRef<int> M, unsigned NumElts,
                               unsigned &Imm) {
  // Assume that the first shuffle index is not UNDEF.  Fail if it is.
  if (M[0] < 0)
    return false;
 
  Imm = M[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, just follow it
    // back to index zero and keep going.
    ++ExpectedElt;
    if (ExpectedElt == NumElts)
      ExpectedElt = 0;
 
    if (M[I] < 0)
      continue; // Ignore UNDEF indices.
    if (ExpectedElt != static_cast<unsigned>(M[I]))
      return false;
  }
 
  return true;
}
 
/// isDUPQMask - matches a splat of equivalent lanes within segments of a given
///              number of elements.
inline std::optional<unsigned> isDUPQMask(ArrayRef<int> Mask, unsigned Segments,
                                          unsigned SegmentSize) {
  unsigned Lane = unsigned(Mask[0]);
 
  // Make sure there's no size changes.
  if (SegmentSize * Segments != Mask.size())
    return std::nullopt;
 
  // Check the first index corresponds to one of the lanes in the first segment.
  if (Lane >= SegmentSize)
    return std::nullopt;
 
  // Check that all lanes match the first, adjusted for segment.
  // Undef/poison lanes (<0) are also accepted.
  if (all_of(enumerate(Mask), [&](auto P) {
        const unsigned SegmentIndex = P.index() / SegmentSize;
        return P.value() < 0 ||
               unsigned(P.value()) == Lane + SegmentIndex * SegmentSize;
      }))
    return Lane;
 
  return std::nullopt;
}
 
/// isDUPFirstSegmentMask - matches a splat of the first 128b segment.
inline bool isDUPFirstSegmentMask(ArrayRef<int> Mask, unsigned Segments,
                                  unsigned SegmentSize) {
  // Make sure there's no size changes.
  if (SegmentSize * Segments != Mask.size())
    return false;
 
  // Check that all lanes refer to the equivalent lane in the first segment.
  // Undef/poison lanes (<0) are also accepted.
  return all_of(enumerate(Mask), [&](auto P) {
    const unsigned IndexWithinSegment = P.index() % SegmentSize;
    return P.value() < 0 || unsigned(P.value()) == IndexWithinSegment;
  });
}
 
/// ShuffleEntry - Represent a shuffle entry in the decomposion of a vector
/// shuffle. i.e. a vector shuffle LHS, RHS, Mask can be built using a list of
/// ShuffleEntry, by performing Op to either LHS, RHS or one of the previous
/// ShuffleEntrys in the lists.
struct ShuffleEntry {
  // The supported operations. The first 16 match those generated by
  // PerfectShuffle.cpp.
  enum Operation {
    OP_COPY = 0, // Copy, used for things like <u,u,u,3> to say it is <0,1,2,3>
    OP_VREV,
    OP_VDUP0,
    OP_VDUP1,
    OP_VDUP2,
    OP_VDUP3,
    OP_VEXT1,
    OP_VEXT2,
    OP_VEXT3,
    OP_VUZPL,   // VUZP, left result
    OP_VUZPR,   // VUZP, right result
    OP_VZIPL,   // VZIP, left result
    OP_VZIPR,   // VZIP, right result
    OP_VTRNL,   // VTRN, left result
    OP_VTRNR,   // VTRN, right result
    OP_MOVLANE, // Move lane. RHSID is the lane to move into
 
    OP_MOVLANE32,
    OP_MOVLANE64,
  };
 
  /// Special IDs for the LHS and RHS values.
  enum IDs {
    LHS = 0xfe,
    RHS = 0xff,
  };
 
  Operation Op;
  unsigned LHSID;
  unsigned RHSID;
};
 
inline unsigned
generatePerfectShuffleFromTable(unsigned PFTableIndex,
                                SmallVector<ShuffleEntry> &Entries) {
  unsigned PFEntry = PerfectShuffleTable[PFTableIndex];
  ShuffleEntry::Operation OpNum =
      (ShuffleEntry::Operation)((PFEntry >> 26) & 0x0F);
  unsigned LHSID = (PFEntry >> 13) & ((1 << 13) - 1);
  unsigned RHSID = (PFEntry >> 0) & ((1 << 13) - 1);
 
  if (OpNum == ShuffleEntry::OP_COPY) {
    if (LHSID == (1 * 9 + 2) * 9 + 3)
      return ShuffleEntry::LHS;
    assert(LHSID == ((4 * 9 + 5) * 9 + 6) * 9 + 7 && "Illegal OP_COPY!");
    return ShuffleEntry::RHS;
  }
 
  unsigned OpLHS = generatePerfectShuffleFromTable(LHSID, Entries);
 
  if (OpNum == ShuffleEntry::OP_MOVLANE) {
    // Decompose a PerfectShuffle ID to get the Mask for lane Elt
    auto getPFIDLane = [](unsigned ID, int Elt) -> int {
      assert(Elt < 4 && "Expected Perfect Lanes to be less than 4");
      Elt = 3 - Elt;
      while (Elt > 0) {
        ID /= 9;
        Elt--;
      }
      return (ID % 9 == 8) ? -1 : ID % 9;
    };
 
    // OP_MOVLANE are either D movs (if bit 0x4 is set) or S movs. D movs
    // convert into a higher type.
    if (RHSID & 0x4) {
      int MaskElt = getPFIDLane(PFTableIndex, (RHSID & 0x01) << 1) >> 1;
      if (MaskElt == -1)
        MaskElt =
            (getPFIDLane(PFTableIndex, ((RHSID & 0x01) << 1) + 1) - 1) >> 1;
      assert(MaskElt >= 0 && "Didn't expect an undef movlane index!");
      unsigned ExtLane = MaskElt < 2 ? MaskElt : (MaskElt - 2);
      Entries.push_back({ShuffleEntry::OP_MOVLANE64, OpLHS,
                         (MaskElt >= 2) << 16 | ExtLane << 8 | (RHSID & 0x3)});
    } else {
      int MaskElt = getPFIDLane(PFTableIndex, RHSID);
      assert(MaskElt >= 0 && "Didn't expect an undef movlane index!");
      unsigned ExtLane = MaskElt < 4 ? MaskElt : (MaskElt - 4);
      Entries.push_back({ShuffleEntry::OP_MOVLANE32, OpLHS,
                         (MaskElt >= 4) << 16 | ExtLane << 8 | (RHSID & 0x3)});
    }
    return Entries.size() - 1;
  }
 
  unsigned OpRHS = generatePerfectShuffleFromTable(RHSID, Entries);
 
  Entries.push_back({OpNum, OpLHS, OpRHS});
  return Entries.size() - 1;
}
 
/// generatePerfectShuffle - Given a Mask, attempt to generate the optimal
/// sequence of instructions using zip/uzp/trn/dup/etc. Currently uses the
/// perfect shuffle tables.
inline bool generatePerfectShuffle(ArrayRef<int> Mask, unsigned NumElts,
                                   SmallVector<ShuffleEntry> &Entries) {
  assert(NumElts == 4 && "Only 4 element masks supported at the moment");
 
  // Compute the index in the perfect shuffle table.
  unsigned PFIndexes[4];
  for (unsigned i = 0; i != 4; ++i) {
    if (Mask[i] < 0)
      PFIndexes[i] = 8;
    else
      PFIndexes[i] = Mask[i];
  }
 
  unsigned PFTableIndex = PFIndexes[0] * 9 * 9 * 9 + PFIndexes[1] * 9 * 9 +
                          PFIndexes[2] * 9 + PFIndexes[3];
  unsigned Idx = generatePerfectShuffleFromTable(PFTableIndex, Entries);
  return Idx == Entries.size() - 1;
}
 
} // namespace llvm
 
#endif