SLPShuffleAnalysis.h

Go to the documentation of this file.

//===- SLPShuffleAnalysis.h - SLP shuffle analysis base ---------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Internal header used by SLPVectorizer.cpp. It defines the base class for
// shuffle cost estimation and shuffle instruction emission. It does not depend
// on BoUpSLP or any other SLP-private type.
//
//===----------------------------------------------------------------------===//
 
#ifndef LLVM_LIB_TRANSFORMS_VECTORIZE_SLPVECTORIZER_SLPSHUFFLEANALYSIS_H
#define LLVM_LIB_TRANSFORMS_VECTORIZE_SLPVECTORIZER_SLPSHUFFLEANALYSIS_H
 
#include "SLPUtils.h"
 
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Instructions.h"
#include "llvm/Support/Casting.h"
 
#include <algorithm>
#include <cassert>
#include <type_traits>
#include <utility>
 
namespace llvm::slpvectorizer {
 
/// The base class for shuffle instruction emission and shuffle cost estimation.
class BaseShuffleAnalysis {
protected:
  Type *ScalarTy = nullptr;
 
  BaseShuffleAnalysis(Type *ScalarTy) : ScalarTy(ScalarTy) {}
 
  /// V is expected to be a vectorized value.
  /// When REVEC is disabled, there is no difference between VF and
  /// VNumElements.
  /// When REVEC is enabled, VF is VNumElements / ScalarTyNumElements.
  /// e.g., if ScalarTy is <4 x Ty> and V1 is <8 x Ty>, 2 is returned instead
  /// of 8.
  unsigned getVF(Value *V) const {
    assert(V && "V cannot be nullptr");
    assert(isa<FixedVectorType>(V->getType()) &&
           "V does not have FixedVectorType");
    assert(ScalarTy && "ScalarTy cannot be nullptr");
    unsigned ScalarTyNumElements = getNumElements(ScalarTy);
    unsigned VNumElements =
        cast<FixedVectorType>(V->getType())->getNumElements();
    assert(VNumElements > ScalarTyNumElements &&
           "the number of elements of V is not large enough");
    assert(VNumElements % ScalarTyNumElements == 0 &&
           "the number of elements of V is not a vectorized value");
    return VNumElements / ScalarTyNumElements;
  }
 
  /// Checks if the mask is an identity mask.
  /// \param IsStrict if is true the function returns false if mask size does
  /// not match vector size.
  static bool isIdentityMask(ArrayRef<int> Mask, const FixedVectorType *VecTy,
                             bool IsStrict) {
    int Limit = Mask.size();
    int VF = VecTy->getNumElements();
    int Index = -1;
    if (VF == Limit && ShuffleVectorInst::isIdentityMask(Mask, Limit))
      return true;
    if (!IsStrict) {
      // Consider extract subvector starting from index 0.
      if (ShuffleVectorInst::isExtractSubvectorMask(Mask, VF, Index) &&
          Index == 0)
        return true;
      // All VF-size submasks are identity (e.g.
      // <poison,poison,poison,poison,0,1,2,poison,poison,1,2,3> etc. for VF 4).
      if (Limit % VF == 0 && all_of(seq<int>(0, Limit / VF), [=](int Idx) {
            ArrayRef<int> Slice = Mask.slice(Idx * VF, VF);
            return all_of(Slice, equal_to(PoisonMaskElem)) ||
                   ShuffleVectorInst::isIdentityMask(Slice, VF);
          }))
        return true;
    }
    return false;
  }
 
  /// Tries to combine 2 different masks into single one.
  /// \param LocalVF Vector length of the permuted input vector. \p Mask may
  /// change the size of the vector, \p LocalVF is the original size of the
  /// shuffled vector.
  static void combineMasks(unsigned LocalVF, SmallVectorImpl<int> &Mask,
                           ArrayRef<int> ExtMask) {
    unsigned VF = Mask.size();
    SmallVector<int> NewMask(ExtMask.size(), PoisonMaskElem);
    for (int I = 0, Sz = ExtMask.size(); I < Sz; ++I) {
      if (ExtMask[I] == PoisonMaskElem)
        continue;
      int MaskedIdx = Mask[ExtMask[I] % VF];
      NewMask[I] =
          MaskedIdx == PoisonMaskElem ? PoisonMaskElem : MaskedIdx % LocalVF;
    }
    Mask.swap(NewMask);
  }
 
  /// Looks through shuffles trying to reduce final number of shuffles in the
  /// code. The function looks through the previously emitted shuffle
  /// instructions and properly mark indices in mask as undef.
  /// For example, given the code
  /// \code
  /// %s1 = shufflevector <2 x ty> %0, poison, <1, 0>
  /// %s2 = shufflevector <2 x ty> %1, poison, <1, 0>
  /// \endcode
  /// and if need to emit shuffle of %s1 and %s2 with mask <1, 0, 3, 2>, it will
  /// look through %s1 and %s2 and select vectors %0 and %1 with mask
  /// <0, 1, 2, 3> for the shuffle.
  /// If 2 operands are of different size, the smallest one will be resized and
  /// the mask recalculated properly.
  /// For example, given the code
  /// \code
  /// %s1 = shufflevector <2 x ty> %0, poison, <1, 0, 1, 0>
  /// %s2 = shufflevector <2 x ty> %1, poison, <1, 0, 1, 0>
  /// \endcode
  /// and if need to emit shuffle of %s1 and %s2 with mask <1, 0, 5, 4>, it will
  /// look through %s1 and %s2 and select vectors %0 and %1 with mask
  /// <0, 1, 2, 3> for the shuffle.
  /// So, it tries to transform permutations to simple vector merge, if
  /// possible.
  /// \param V The input vector which must be shuffled using the given \p Mask.
  /// If the better candidate is found, \p V is set to this best candidate
  /// vector.
  /// \param Mask The input mask for the shuffle. If the best candidate is found
  /// during looking-through-shuffles attempt, it is updated accordingly.
  /// \param SinglePermute true if the shuffle operation is originally a
  /// single-value-permutation. In this case the look-through-shuffles procedure
  /// may look for resizing shuffles as the best candidates.
  /// \return true if the shuffle results in the non-resizing identity shuffle
  /// (and thus can be ignored), false - otherwise.
  static bool peekThroughShuffles(Value *&V, SmallVectorImpl<int> &Mask,
                                  bool SinglePermute) {
    Value *Op = V;
    ShuffleVectorInst *IdentityOp = nullptr;
    SmallVector<int> IdentityMask;
    while (auto *SV = dyn_cast<ShuffleVectorInst>(Op)) {
      // Exit if not a fixed vector type or changing size shuffle.
      auto *SVTy = dyn_cast<FixedVectorType>(SV->getType());
      if (!SVTy)
        break;
      // Remember the identity or broadcast mask, if it is not a resizing
      // shuffle. If no better candidates are found, this Op and Mask will be
      // used in the final shuffle.
      if (isIdentityMask(Mask, SVTy, /*IsStrict=*/false)) {
        if (!IdentityOp || !SinglePermute ||
            (isIdentityMask(Mask, SVTy, /*IsStrict=*/true) &&
             !ShuffleVectorInst::isZeroEltSplatMask(IdentityMask,
                                                    IdentityMask.size()))) {
          IdentityOp = SV;
          // Store current mask in the IdentityMask so later we did not lost
          // this info if IdentityOp is selected as the best candidate for the
          // permutation.
          IdentityMask.assign(Mask);
        }
      }
      // Remember the broadcast mask. If no better candidates are found, this Op
      // and Mask will be used in the final shuffle.
      // Zero splat can be used as identity too, since it might be used with
      // mask <0, 1, 2, ...>, i.e. identity mask without extra reshuffling.
      // E.g. if need to shuffle the vector with the mask <3, 1, 2, 0>, which is
      // expensive, the analysis founds out, that the source vector is just a
      // broadcast, this original mask can be transformed to identity mask <0,
      // 1, 2, 3>.
      // \code
      // %0 = shuffle %v, poison, zeroinitalizer
      // %res = shuffle %0, poison, <3, 1, 2, 0>
      // \endcode
      // may be transformed to
      // \code
      // %0 = shuffle %v, poison, zeroinitalizer
      // %res = shuffle %0, poison, <0, 1, 2, 3>
      // \endcode
      if (SV->isZeroEltSplat()) {
        IdentityOp = SV;
        IdentityMask.assign(Mask);
      }
      int LocalVF = Mask.size();
      if (auto *SVOpTy =
              dyn_cast<FixedVectorType>(SV->getOperand(0)->getType()))
        LocalVF = SVOpTy->getNumElements();
      SmallVector<int> ExtMask(Mask.size(), PoisonMaskElem);
      for (auto [Idx, I] :
           make_filter_range(enumerate(Mask), [&](const auto &P) {
             return P.value() != PoisonMaskElem &&
                    static_cast<unsigned>(P.value()) <
                        SV->getShuffleMask().size();
           }))
        ExtMask[Idx] = SV->getMaskValue(I);
      bool IsOp1Undef = isUndefVector</*isPoisonOnly=*/true>(
                            SV->getOperand(0),
                            buildUseMask(LocalVF, ExtMask, UseMask::FirstArg))
                            .all();
      bool IsOp2Undef = isUndefVector</*isPoisonOnly=*/true>(
                            SV->getOperand(1),
                            buildUseMask(LocalVF, ExtMask, UseMask::SecondArg))
                            .all();
      if (!IsOp1Undef && !IsOp2Undef) {
        // Update mask and mark undef elems.
        for (int &I : make_filter_range(
                 Mask, [](int I) { return I != PoisonMaskElem; })) {
          if (SV->getMaskValue(I % SV->getShuffleMask().size()) ==
              PoisonMaskElem)
            I = PoisonMaskElem;
        }
        break;
      }
      SmallVector<int> ShuffleMask(SV->getShuffleMask());
      combineMasks(LocalVF, ShuffleMask, Mask);
      Mask.swap(ShuffleMask);
      if (IsOp2Undef)
        Op = SV->getOperand(0);
      else
        Op = SV->getOperand(1);
    }
    if (auto *OpTy = dyn_cast<FixedVectorType>(Op->getType());
        !OpTy || !isIdentityMask(Mask, OpTy, SinglePermute) ||
        ShuffleVectorInst::isZeroEltSplatMask(Mask, Mask.size())) {
      if (IdentityOp) {
        V = IdentityOp;
        assert(Mask.size() == IdentityMask.size() &&
               "Expected masks of same sizes.");
        // Clear known poison elements.
        for (auto [I, Idx] : enumerate(Mask))
          if (Idx == PoisonMaskElem)
            IdentityMask[I] = PoisonMaskElem;
        Mask.swap(IdentityMask);
        auto *Shuffle = dyn_cast<ShuffleVectorInst>(V);
        return SinglePermute &&
               (isIdentityMask(Mask, cast<FixedVectorType>(V->getType()),
                               /*IsStrict=*/true) ||
                (Shuffle && Mask.size() == Shuffle->getShuffleMask().size() &&
                 Shuffle->isZeroEltSplat() &&
                 ShuffleVectorInst::isZeroEltSplatMask(Mask, Mask.size()) &&
                 all_of(enumerate(Mask), [&](const auto &P) {
                   return P.value() == PoisonMaskElem ||
                          Shuffle->getShuffleMask()[P.index()] == 0;
                 })));
      }
      V = Op;
      return false;
    }
    V = Op;
    return true;
  }
 
  /// Smart shuffle instruction emission, walks through shuffles trees and
  /// tries to find the best matching vector for the actual shuffle
  /// instruction.
  template <typename T, typename ShuffleBuilderTy, typename... Args>
  static T createShuffle(Value *V1, Value *V2, ArrayRef<int> Mask,
                         ShuffleBuilderTy &Builder, Type *ScalarTy,
                         [[maybe_unused]] bool ReVec, Args... Arguments) {
    assert(V1 && "Expected at least one vector value.");
    unsigned ScalarTyNumElements = getNumElements(ScalarTy);
    SmallVector<int> NewMask(Mask);
    if (ScalarTyNumElements != 1) {
      assert(ReVec && "FixedVectorType is not expected.");
      transformScalarShuffleIndiciesToVector(ScalarTyNumElements, NewMask);
      Mask = NewMask;
    }
    if (V2)
      Builder.resizeToMatch(V1, V2);
    int VF = Mask.size();
    if (auto *FTy = dyn_cast<FixedVectorType>(V1->getType()))
      VF = FTy->getNumElements();
    if (V2 && !isUndefVector</*IsPoisonOnly=*/true>(
                   V2, buildUseMask(VF, Mask, UseMask::SecondArg))
                   .all()) {
      // Peek through shuffles.
      Value *Op1 = V1;
      Value *Op2 = V2;
      int VF =
          cast<VectorType>(V1->getType())->getElementCount().getKnownMinValue();
      SmallVector<int> CombinedMask1(Mask.size(), PoisonMaskElem);
      SmallVector<int> CombinedMask2(Mask.size(), PoisonMaskElem);
      for (int I = 0, E = Mask.size(); I < E; ++I) {
        if (Mask[I] < VF)
          CombinedMask1[I] = Mask[I];
        else
          CombinedMask2[I] = Mask[I] - VF;
      }
      Value *PrevOp1;
      Value *PrevOp2;
      do {
        PrevOp1 = Op1;
        PrevOp2 = Op2;
        (void)peekThroughShuffles(Op1, CombinedMask1, /*SinglePermute=*/false);
        (void)peekThroughShuffles(Op2, CombinedMask2, /*SinglePermute=*/false);
        // Check if we have 2 resizing shuffles - need to peek through operands
        // again.
        if (auto *SV1 = dyn_cast<ShuffleVectorInst>(Op1))
          if (auto *SV2 = dyn_cast<ShuffleVectorInst>(Op2)) {
            SmallVector<int> ExtMask1(Mask.size(), PoisonMaskElem);
            for (const auto &P :
                 make_filter_range(enumerate(CombinedMask1), [](const auto &P) {
                   return P.value() != PoisonMaskElem;
                 }))
              ExtMask1[P.index()] = SV1->getMaskValue(P.value());
            SmallBitVector UseMask1 = buildUseMask(
                cast<FixedVectorType>(SV1->getOperand(1)->getType())
                    ->getNumElements(),
                ExtMask1, UseMask::SecondArg);
            SmallVector<int> ExtMask2(CombinedMask2.size(), PoisonMaskElem);
            for (const auto &P :
                 make_filter_range(enumerate(CombinedMask2), [](const auto &P) {
                   return P.value() != PoisonMaskElem;
                 }))
              ExtMask2[P.index()] = SV2->getMaskValue(P.value());
            SmallBitVector UseMask2 = buildUseMask(
                cast<FixedVectorType>(SV2->getOperand(1)->getType())
                    ->getNumElements(),
                ExtMask2, UseMask::SecondArg);
            if (SV1->getOperand(0)->getType() ==
                    SV2->getOperand(0)->getType() &&
                SV1->getOperand(0)->getType() != SV1->getType() &&
                isUndefVector</*IsPoisonOnly=*/true>(SV1->getOperand(1),
                                                     UseMask1)
                    .all() &&
                isUndefVector</*IsPoisonOnly=*/true>(SV2->getOperand(1),
                                                     UseMask2)
                    .all()) {
              Op1 = SV1->getOperand(0);
              Op2 = SV2->getOperand(0);
              SmallVector<int> ShuffleMask1(SV1->getShuffleMask());
              int LocalVF = ShuffleMask1.size();
              if (auto *FTy = dyn_cast<FixedVectorType>(Op1->getType()))
                LocalVF = FTy->getNumElements();
              combineMasks(LocalVF, ShuffleMask1, CombinedMask1);
              CombinedMask1.swap(ShuffleMask1);
              SmallVector<int> ShuffleMask2(SV2->getShuffleMask());
              LocalVF = ShuffleMask2.size();
              if (auto *FTy = dyn_cast<FixedVectorType>(Op2->getType()))
                LocalVF = FTy->getNumElements();
              combineMasks(LocalVF, ShuffleMask2, CombinedMask2);
              CombinedMask2.swap(ShuffleMask2);
            }
          }
      } while (PrevOp1 != Op1 || PrevOp2 != Op2);
      Builder.resizeToMatch(Op1, Op2);
      VF = std::max(cast<VectorType>(Op1->getType())
                        ->getElementCount()
                        .getKnownMinValue(),
                    cast<VectorType>(Op2->getType())
                        ->getElementCount()
                        .getKnownMinValue());
      for (int I = 0, E = Mask.size(); I < E; ++I) {
        if (CombinedMask2[I] != PoisonMaskElem) {
          assert(CombinedMask1[I] == PoisonMaskElem &&
                 "Expected undefined mask element");
          CombinedMask1[I] = CombinedMask2[I] + (Op1 == Op2 ? 0 : VF);
        }
      }
      if (Op1 == Op2 &&
          (ShuffleVectorInst::isIdentityMask(CombinedMask1, VF) ||
           (ShuffleVectorInst::isZeroEltSplatMask(CombinedMask1, VF) &&
            isa<ShuffleVectorInst>(Op1) &&
            cast<ShuffleVectorInst>(Op1)->getShuffleMask() ==
                ArrayRef(CombinedMask1))))
        return Builder.createIdentity(Op1);
      return Builder.createShuffleVector(
          Op1, Op1 == Op2 ? PoisonValue::get(Op1->getType()) : Op2,
          CombinedMask1);
    }
    if (isa<PoisonValue>(V1))
      return Builder.createPoison(
          cast<VectorType>(V1->getType())->getElementType(), Mask.size());
    bool IsIdentity = peekThroughShuffles(V1, NewMask, /*SinglePermute=*/true);
    assert(V1 && "Expected non-null value after looking through shuffles.");
 
    if (!IsIdentity)
      return Builder.createShuffleVector(V1, NewMask, Arguments...);
    return Builder.createIdentity(V1);
  }
 
  /// Transforms mask \p CommonMask per given \p Mask to make proper set after
  /// shuffle emission.
  static void transformMaskAfterShuffle(MutableArrayRef<int> CommonMask,
                                        ArrayRef<int> Mask) {
    for (unsigned I : seq<unsigned>(CommonMask.size()))
      if (Mask[I] != PoisonMaskElem)
        CommonMask[I] = I;
  }
};
 
/// Does the analysis of the provided shuffle masks and performs the requested
/// actions on the vectors with the given shuffle masks. It tries to do it in
/// several steps.
/// 1. If the Base vector is not poison vector, resizing the very first mask to
/// have common VF and perform action for 2 input vectors (including non-poison
/// Base). Other shuffle masks are combined with the resulting after the 1 stage
/// and processed as a shuffle of 2 elements.
/// 2. If the Base is poison vector and have only 1 shuffle mask, perform the
/// action only for 1 vector with the given mask, if it is not the identity
/// mask.
/// 3. If > 2 masks are used, perform the remaining shuffle actions for 2
/// vectors, combing the masks properly between the steps.
template <typename T>
T *performExtractsShuffleAction(
    MutableArrayRef<std::pair<T *, SmallVector<int>>> ShuffleMask, Value *Base,
    function_ref<unsigned(T *)> GetVF,
    function_ref<std::pair<T *, bool>(T *, ArrayRef<int>, bool)> ResizeAction,
    function_ref<T *(ArrayRef<int>, ArrayRef<T *>)> Action) {
  assert(!ShuffleMask.empty() && "Empty list of shuffles for inserts.");
  SmallVector<int> Mask(ShuffleMask.begin()->second);
  auto VMIt = std::next(ShuffleMask.begin());
  T *Prev = nullptr;
  SmallBitVector UseMask =
      buildUseMask(Mask.size(), Mask, UseMask::UndefsAsMask);
  // The undef elements of the base must be preserved, only the poison base can
  // be dropped.
  SmallBitVector IsBasePoison = isUndefVector<true>(Base, UseMask);
  if (!IsBasePoison.all()) {
    // Base is not poison, need to combine it with the next subvectors.
    std::pair<T *, bool> Res =
        ResizeAction(ShuffleMask.begin()->first, Mask, /*ForSingleMask=*/false);
    for (unsigned Idx = 0, VF = Mask.size(); Idx < VF; ++Idx) {
      if (Mask[Idx] == PoisonMaskElem)
        Mask[Idx] = IsBasePoison.test(Idx) ? PoisonMaskElem : Idx;
      else
        Mask[Idx] = (Res.second ? Idx : Mask[Idx]) + VF;
    }
    if constexpr (std::is_same_v<T, Value>)
      assert(GetVF(Base) == Mask.size() &&
             "Expected base vector of VF number of elements.");
    Prev = Action(Mask, {nullptr, Res.first});
  } else if (ShuffleMask.size() == 1) {
    // Base is poison and only 1 vector is shuffled - perform the action only
    // for single vector, if the mask is not the identity mask.
    std::pair<T *, bool> Res = ResizeAction(ShuffleMask.begin()->first, Mask,
                                            /*ForSingleMask=*/true);
    if (Res.second)
      // Identity mask is found.
      Prev = Res.first;
    else
      Prev = Action(Mask, {ShuffleMask.begin()->first});
  } else {
    // Base is poison and at least 2 input vectors shuffled - perform 2 vectors
    // shuffles step by step, combining shuffle between the steps.
    unsigned Vec1VF = GetVF(ShuffleMask.begin()->first);
    unsigned Vec2VF = GetVF(VMIt->first);
    if (Vec1VF == Vec2VF) {
      // No need to resize the input vectors since they are of the same size, we
      // can shuffle them directly.
      ArrayRef<int> SecMask = VMIt->second;
      for (unsigned I = 0, VF = Mask.size(); I < VF; ++I) {
        if (SecMask[I] != PoisonMaskElem) {
          assert(Mask[I] == PoisonMaskElem && "Multiple uses of scalars.");
          Mask[I] = SecMask[I] + Vec1VF;
        }
      }
      Prev = Action(Mask, {ShuffleMask.begin()->first, VMIt->first});
    } else {
      // Vectors of different sizes - resize and reshuffle.
      std::pair<T *, bool> Res1 = ResizeAction(ShuffleMask.begin()->first, Mask,
                                               /*ForSingleMask=*/false);
      std::pair<T *, bool> Res2 =
          ResizeAction(VMIt->first, VMIt->second, /*ForSingleMask=*/false);
      ArrayRef<int> SecMask = VMIt->second;
      for (unsigned I = 0, VF = Mask.size(); I < VF; ++I) {
        if (Mask[I] != PoisonMaskElem) {
          assert(SecMask[I] == PoisonMaskElem && "Multiple uses of scalars.");
          if (Res1.second)
            Mask[I] = I;
        } else if (SecMask[I] != PoisonMaskElem) {
          assert(Mask[I] == PoisonMaskElem && "Multiple uses of scalars.");
          Mask[I] = (Res2.second ? I : SecMask[I]) + VF;
        }
      }
      Prev = Action(Mask, {Res1.first, Res2.first});
    }
    VMIt = std::next(VMIt);
  }
  [[maybe_unused]] bool IsBaseNotPoison = !IsBasePoison.all();
  // Perform requested actions for the remaining masks/vectors.
  for (auto E = ShuffleMask.end(); VMIt != E; ++VMIt) {
    // Shuffle other input vectors, if any.
    std::pair<T *, bool> Res =
        ResizeAction(VMIt->first, VMIt->second, /*ForSingleMask=*/false);
    ArrayRef<int> SecMask = VMIt->second;
    for (unsigned I = 0, VF = Mask.size(); I < VF; ++I) {
      if (SecMask[I] != PoisonMaskElem) {
        assert((Mask[I] == PoisonMaskElem || IsBaseNotPoison) &&
               "Multiple uses of scalars.");
        Mask[I] = (Res.second ? I : SecMask[I]) + VF;
      } else if (Mask[I] != PoisonMaskElem) {
        Mask[I] = I;
      }
    }
    Prev = Action(Mask, {Prev, Res.first});
  }
  return Prev;
}
 
} // namespace llvm::slpvectorizer
 
#endif // LLVM_LIB_TRANSFORMS_VECTORIZE_SLPVECTORIZER_SLPSHUFFLEANALYSIS_H