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/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;
}
 
/// 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;
  });
}
 
} // namespace llvm
 
#endif