SLPUtils.cpp

Go to the documentation of this file.

//===- SLPUtils.cpp - SLP Vectorizer free utility helpers -----------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#include "SLPUtils.h"
 
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Analysis/AssumptionCache.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/Analysis/VectorUtils.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/PatternMatch.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Support/raw_ostream.h"
 
#include <algorithm>
#include <numeric>
#include <string>
#include <type_traits>
 
using namespace llvm;
using namespace llvm::PatternMatch;
 
namespace llvm::slpvectorizer {
 
bool isConstant(Value *V) {
  return isa<Constant>(V) && !isa<ConstantExpr, GlobalValue>(V);
}
 
bool isBinOpIdentityConstant(const Value *V, unsigned Opcode) {
  const auto *CI = dyn_cast<ConstantInt>(V);
  return CI && ConstantExpr::getBinOpIdentity(Opcode, CI->getType()) == CI;
}
 
unsigned getReassocCombineOpcode(unsigned Opcode) {
  switch (Opcode) {
  case Instruction::Sub:
    return Instruction::Add;
  case Instruction::FSub:
    return Instruction::FAdd;
  default:
    return Opcode;
  }
}
 
bool isReassocChainLink(const Instruction *I) {
  if (I->getOpcode() == Instruction::Sub)
    return true;
  if (I->getOpcode() == Instruction::FSub)
    return I->hasAllowReassoc();
  return I->isAssociative();
}
 
bool isVectorLikeInstWithConstOps(Value *V) {
  auto *I = dyn_cast<Instruction>(V);
  // Non-instructions are vector-like only if they are undef.
  if (!I)
    return isa<UndefValue>(V);
  switch (I->getOpcode()) {
  case Instruction::ExtractValue:
  case Instruction::InsertValue:
    return true;
  case Instruction::ExtractElement:
    return isa<FixedVectorType>(I->getOperand(0)->getType()) &&
           isConstant(I->getOperand(1));
  case Instruction::InsertElement:
    return isa<FixedVectorType>(I->getOperand(0)->getType()) &&
           isConstant(I->getOperand(2));
  default:
    return false;
  }
}
 
unsigned getNumElements(Type *Ty) {
  assert(!isa<ScalableVectorType>(Ty) &&
         "ScalableVectorType is not supported.");
  if (isVectorizedTy(Ty))
    return getVectorizedTypeVF(Ty).getFixedValue();
  return 1;
}
 
unsigned getPartNumElems(unsigned Size, unsigned NumParts) {
  return std::min<unsigned>(Size, bit_ceil(divideCeil(Size, NumParts)));
}
 
unsigned getNumElems(unsigned Size, unsigned PartNumElems, unsigned Part) {
  return std::min<unsigned>(PartNumElems, Size - Part * PartNumElems);
}
 
#if !defined(NDEBUG)
std::string shortBundleName(ArrayRef<Value *> VL, int Idx) {
  std::string Result;
  raw_string_ostream OS(Result);
  if (Idx >= 0)
    OS << "Idx: " << Idx << ", ";
  OS << "n=" << VL.size() << " [" << *VL.front() << ", ..]";
  return Result;
}
#endif
 
bool allSameBlock(ArrayRef<Value *> VL) {
  auto *It = find_if(VL, IsaPred<Instruction>);
  if (It == VL.end())
    return false;
  Instruction *I0 = cast<Instruction>(*It);
  if (all_of(VL, isVectorLikeInstWithConstOps))
    return true;
 
  BasicBlock *BB = I0->getParent();
  for (Value *V : make_filter_range(iterator_range(It, VL.end()), [](Value *V) {
         return !isa<PoisonValue>(V);
       })) {
    auto *II = dyn_cast<Instruction>(V);
    if (!II)
      return false;
 
    if (BB != II->getParent())
      return false;
  }
  return true;
}
 
bool allConstant(ArrayRef<Value *> VL) {
  // Constant expressions and globals can't be vectorized like normal integer/FP
  // constants.
  return all_of(VL, isConstant);
}
 
bool isSplat(ArrayRef<Value *> VL) {
  Value *FirstNonUndef = nullptr;
  for (Value *V :
       make_filter_range(VL, [](Value *V) { return !isa<UndefValue>(V); })) {
    if (!FirstNonUndef) {
      FirstNonUndef = V;
      continue;
    }
    if (V != FirstNonUndef)
      return false;
  }
  return FirstNonUndef != nullptr;
}
 
Intrinsic::ID isEquivalentIntrinsicID(Intrinsic::ID LHS, Intrinsic::ID RHS) {
  if (LHS == RHS)
    return RHS;
  if ((LHS == Intrinsic::fma || LHS == Intrinsic::fmuladd) &&
      (RHS == Intrinsic::fma || RHS == Intrinsic::fmuladd))
    return Intrinsic::fma;
  return Intrinsic::not_intrinsic;
}
 
bool isCommutative(const Instruction *I, const Value *ValWithUses,
                   bool IsCopyable) {
  if (auto *Cmp = dyn_cast<CmpInst>(I))
    return Cmp->isCommutative();
  if (auto *BO = dyn_cast<BinaryOperator>(I))
    return BO->isCommutative() ||
           (BO->getOpcode() == Instruction::Sub && ValWithUses->hasUseList() &&
            !ValWithUses->hasNUsesOrMore(UsesLimit) &&
            all_of(
                ValWithUses->uses(),
                [&](const Use &U) {
                  // Commutative, if icmp eq/ne sub, 0
                  CmpPredicate Pred;
                  if (match(U.getUser(),
                            m_ICmp(Pred, m_Specific(U.get()), m_Zero())) &&
                      (Pred == ICmpInst::ICMP_EQ || Pred == ICmpInst::ICMP_NE))
                    return true;
                  // Commutative, if abs(sub nsw, true) or abs(sub, false).
                  ConstantInt *Flag;
                  auto *I = dyn_cast<BinaryOperator>(U.get());
                  return match(U.getUser(),
                               m_Intrinsic<Intrinsic::abs>(
                                   m_Specific(U.get()), m_ConstantInt(Flag))) &&
                         ((!IsCopyable && I && !I->hasNoSignedWrap()) ||
                          Flag->isOne());
                })) ||
           (BO->getOpcode() == Instruction::FSub && ValWithUses->hasUseList() &&
            !ValWithUses->hasNUsesOrMore(UsesLimit) &&
            all_of(ValWithUses->uses(), [](const Use &U) {
              return match(U.getUser(),
                           m_Intrinsic<Intrinsic::fabs>(m_Specific(U.get())));
            }));
  return I->isCommutative();
}
 
bool isCommutative(const Instruction *I) { return isCommutative(I, I); }
 
bool isCommutableOperand(const Instruction *I, Value *ValWithUses, unsigned Op,
                         bool IsCopyable) {
  assert(isCommutative(I, ValWithUses, IsCopyable) &&
         "The instruction is not commutative.");
  if (isa<CmpInst>(I))
    return true;
  if (auto *BO = dyn_cast<BinaryOperator>(I)) {
    switch (BO->getOpcode()) {
    case Instruction::Sub:
    case Instruction::FSub:
      return true;
    default:
      break;
    }
  }
  return I->isCommutableOperand(Op);
}
 
unsigned getNumberOfPotentiallyCommutativeOps(Instruction *I) {
  if (isa<IntrinsicInst>(I) && isCommutative(I)) {
    // IntrinsicInst::isCommutative returns true if swapping the first "two"
    // arguments to the intrinsic produces the same result.
    constexpr unsigned IntrinsicNumOperands = 2;
    return IntrinsicNumOperands;
  }
  return I->getNumOperands();
}
 
std::optional<unsigned> getElementIndex(const Value *Inst, unsigned Offset) {
  if (auto Index = getInsertExtractIndex<InsertElementInst>(Inst, Offset))
    return Index;
  if (auto Index = getInsertExtractIndex<ExtractElementInst>(Inst, Offset))
    return Index;
 
  unsigned Index = Offset;
 
  const auto *IV = dyn_cast<InsertValueInst>(Inst);
  if (!IV)
    return std::nullopt;
 
  Type *CurrentType = IV->getType();
  for (unsigned I : IV->indices()) {
    if (const auto *ST = dyn_cast<StructType>(CurrentType)) {
      Index *= ST->getNumElements();
      CurrentType = ST->getElementType(I);
    } else if (const auto *AT = dyn_cast<ArrayType>(CurrentType)) {
      Index *= AT->getNumElements();
      CurrentType = AT->getElementType();
    } else {
      return std::nullopt;
    }
    Index += I;
  }
  return Index;
}
 
bool allSameOpcode(ArrayRef<Value *> VL) {
  auto *It = find_if(VL, IsaPred<Instruction>);
  if (It == VL.end())
    return true;
  Instruction *MainOp = cast<Instruction>(*It);
  unsigned Opcode = MainOp->getOpcode();
  bool IsCmpOp = isa<CmpInst>(MainOp);
  CmpInst::Predicate BasePred = IsCmpOp ? cast<CmpInst>(MainOp)->getPredicate()
                                        : CmpInst::BAD_ICMP_PREDICATE;
  return all_of(make_range(It, VL.end()), [&](Value *V) {
    if (auto *CI = dyn_cast<CmpInst>(V))
      return BasePred == CI->getPredicate();
    if (auto *I = dyn_cast<Instruction>(V))
      return I->getOpcode() == Opcode;
    return isa<PoisonValue>(V);
  });
}
 
std::optional<unsigned> getExtractIndex(const Instruction *E) {
  unsigned Opcode = E->getOpcode();
  assert((Opcode == Instruction::ExtractElement ||
          Opcode == Instruction::ExtractValue) &&
         "Expected extractelement or extractvalue instruction.");
  if (Opcode == Instruction::ExtractElement) {
    auto *CI = dyn_cast<ConstantInt>(E->getOperand(1));
    if (!CI)
      return std::nullopt;
    // Check if the index is out of bound. We can get the source vector from
    // operand 0.
    unsigned Idx = CI->getZExtValue();
    auto *EE = cast<ExtractElementInst>(E);
    const unsigned VF = getNumElements(EE->getVectorOperandType());
    if (Idx >= VF)
      return std::nullopt;
    return Idx;
  }
  auto *EI = cast<ExtractValueInst>(E);
  if (EI->getNumIndices() != 1)
    return std::nullopt;
  return *EI->idx_begin();
}
 
void inversePermutation(ArrayRef<unsigned> Indices,
                        SmallVectorImpl<int> &Mask) {
  Mask.clear();
  const unsigned E = Indices.size();
  Mask.resize(E, PoisonMaskElem);
  for (unsigned I = 0; I < E; ++I)
    Mask[Indices[I]] = I;
}
 
void reorderScalars(SmallVectorImpl<Value *> &Scalars, ArrayRef<int> Mask) {
  assert(!Mask.empty() && "Expected non-empty mask.");
  SmallVector<Value *> Prev(Scalars.size(),
                            PoisonValue::get(Scalars.front()->getType()));
  Prev.swap(Scalars);
  for (unsigned I = 0, E = Prev.size(); I < E; ++I)
    if (Mask[I] != PoisonMaskElem)
      Scalars[Mask[I]] = Prev[I];
}
 
void reorderReuses(SmallVectorImpl<int> &Reuses, ArrayRef<int> Mask) {
  assert(!Mask.empty() && Reuses.size() == Mask.size() &&
         "Expected non-empty mask.");
  SmallVector<int> Prev(Reuses.begin(), Reuses.end());
  Prev.swap(Reuses);
  for (unsigned I = 0, E = Prev.size(); I < E; ++I)
    if (Mask[I] != PoisonMaskElem)
      Reuses[Mask[I]] = Prev[I];
}
 
void reorderOrder(SmallVectorImpl<unsigned> &Order, ArrayRef<int> Mask,
                  bool BottomOrder) {
  assert(!Mask.empty() && "Expected non-empty mask.");
  unsigned Sz = Mask.size();
  if (BottomOrder) {
    SmallVector<unsigned> PrevOrder;
    if (Order.empty()) {
      PrevOrder.resize(Sz);
      std::iota(PrevOrder.begin(), PrevOrder.end(), 0);
    } else {
      PrevOrder.swap(Order);
    }
    Order.assign(Sz, Sz);
    for (unsigned I = 0; I < Sz; ++I)
      if (Mask[I] != PoisonMaskElem)
        Order[I] = PrevOrder[Mask[I]];
    if (all_of(enumerate(Order), [&](const auto &Data) {
          return Data.value() == Sz || Data.index() == Data.value();
        })) {
      Order.clear();
      return;
    }
    fixupOrderingIndices(Order);
    return;
  }
  SmallVector<int> MaskOrder;
  if (Order.empty()) {
    MaskOrder.resize(Sz);
    std::iota(MaskOrder.begin(), MaskOrder.end(), 0);
  } else {
    inversePermutation(Order, MaskOrder);
  }
  reorderReuses(MaskOrder, Mask);
  if (ShuffleVectorInst::isIdentityMask(MaskOrder, Sz)) {
    Order.clear();
    return;
  }
  Order.assign(Sz, Sz);
  for (unsigned I = 0; I < Sz; ++I)
    if (MaskOrder[I] != PoisonMaskElem)
      Order[MaskOrder[I]] = I;
  fixupOrderingIndices(Order);
}
 
bool isReverseOrder(ArrayRef<unsigned> Order) {
  assert(!Order.empty() &&
         "Order is empty. Please check it before using isReverseOrder.");
  unsigned Sz = Order.size();
  return all_of(enumerate(Order), [&](const auto &Pair) {
    return Pair.value() == Sz || Sz - Pair.index() - 1 == Pair.value();
  });
}
 
bool isRepeatedNonIdentityClusteredMask(ArrayRef<int> Mask, unsigned Sz) {
  ArrayRef<int> FirstCluster = Mask.slice(0, Sz);
  if (ShuffleVectorInst::isIdentityMask(FirstCluster, Sz))
    return false;
  for (unsigned I = Sz, E = Mask.size(); I < E; I += Sz) {
    ArrayRef<int> Cluster = Mask.slice(I, Sz);
    if (Cluster != FirstCluster)
      return false;
  }
  return true;
}
 
void combineOrders(MutableArrayRef<unsigned> Order,
                   ArrayRef<unsigned> SecondaryOrder) {
  assert((SecondaryOrder.empty() || Order.size() == SecondaryOrder.size()) &&
         "Expected same size of orders");
  size_t Sz = Order.size();
  SmallBitVector UsedIndices(Sz);
  for (unsigned Idx : seq<unsigned>(0, Sz)) {
    if (Order[Idx] != Sz)
      UsedIndices.set(Order[Idx]);
  }
  if (SecondaryOrder.empty()) {
    for (unsigned Idx : seq<unsigned>(0, Sz))
      if (Order[Idx] == Sz && !UsedIndices.test(Idx))
        Order[Idx] = Idx;
  } else {
    for (unsigned Idx : seq<unsigned>(0, Sz))
      if (SecondaryOrder[Idx] != Sz && Order[Idx] == Sz &&
          !UsedIndices.test(SecondaryOrder[Idx]))
        Order[Idx] = SecondaryOrder[Idx];
  }
}
 
bool allSameType(ArrayRef<Value *> VL) {
  assert(!VL.empty() && "Expected non-empty list of values.");
  Type *Ty = VL.consume_front()->getType();
  return all_of(VL, [&](Value *V) { return V->getType() == Ty; });
}
 
template <typename T>
std::optional<unsigned> getInsertExtractIndex(const Value *Inst,
                                              unsigned Offset) {
  static_assert(std::is_same_v<T, InsertElementInst> ||
                    std::is_same_v<T, ExtractElementInst>,
                "unsupported T");
  const auto *IE = dyn_cast<T>(Inst);
  if (!IE)
    return std::nullopt;
  // InsertElement: result is the vector, index is op 2.
  // ExtractElement: result is scalar, vector is op 0, index is op 1.
  constexpr bool IsInsert = std::is_same_v<T, InsertElementInst>;
  Type *VecTy = IsInsert ? IE->getType() : IE->getOperand(0)->getType();
  const auto *VT = dyn_cast<FixedVectorType>(VecTy);
  if (!VT)
    return std::nullopt;
  const auto *CI = dyn_cast<ConstantInt>(IE->getOperand(IsInsert ? 2 : 1));
  if (!CI)
    return std::nullopt;
  if (CI->getValue().uge(VT->getNumElements()))
    return std::nullopt;
  unsigned Index = Offset;
  Index *= VT->getNumElements();
  Index += CI->getZExtValue();
  return Index;
}
 
// Only these two specializations are used; instantiate them here so the
// definition can stay out of the header.
template std::optional<unsigned>
getInsertExtractIndex<InsertElementInst>(const Value *, unsigned);
template std::optional<unsigned>
getInsertExtractIndex<ExtractElementInst>(const Value *, unsigned);
 
bool areAllOperandsNonInsts(Value *V) {
  auto *I = dyn_cast<Instruction>(V);
  if (!I)
    return true;
  return !mayHaveNonDefUseDependency(*I) &&
         all_of(make_isa_range<Instruction>(I->operands()),
                [I](Instruction *IO) {
                  return isa<PHINode>(IO) || IO->getParent() != I->getParent();
                });
}
 
bool isUsedOutsideBlock(Value *V) {
  auto *I = dyn_cast<Instruction>(V);
  if (!I)
    return true;
  // Limits the number of uses to save compile time.
  return !I->mayReadOrWriteMemory() && !I->hasNUsesOrMore(UsesLimit) &&
         all_of(I->users(), [I](User *U) {
           auto *IU = dyn_cast<Instruction>(U);
           if (!IU)
             return true;
           return IU->getParent() != I->getParent() || isa<PHINode>(IU);
         });
}
 
bool doesNotNeedToBeScheduled(Value *V) {
  return areAllOperandsNonInsts(V) && isUsedOutsideBlock(V);
}
 
bool doesNotNeedToSchedule(ArrayRef<Value *> VL) {
  return !VL.empty() &&
         (all_of(VL, isUsedOutsideBlock) || all_of(VL, areAllOperandsNonInsts));
}
 
void transformScalarShuffleIndiciesToVector(unsigned VecTyNumElements,
                                            SmallVectorImpl<int> &Mask) {
  // The ShuffleBuilder implementation use shufflevector to splat an "element".
  // But the element have different meaning for SLP (scalar) and REVEC
  // (vector). We need to expand Mask into masks which shufflevector can use
  // directly.
  SmallVector<int> NewMask(Mask.size() * VecTyNumElements);
  for (unsigned I : seq<unsigned>(Mask.size()))
    for (auto [J, MaskV] : enumerate(MutableArrayRef(NewMask).slice(
             I * VecTyNumElements, VecTyNumElements)))
      MaskV = Mask[I] == PoisonMaskElem ? PoisonMaskElem
                                        : Mask[I] * VecTyNumElements + J;
  Mask.swap(NewMask);
}
 
unsigned getShufflevectorNumGroups(ArrayRef<Value *> VL) {
  if (VL.empty())
    return 0;
  if (!all_of(VL, IsaPred<ShuffleVectorInst>))
    return 0;
  auto *SV = cast<ShuffleVectorInst>(VL.front());
  unsigned SVNumElements =
      cast<FixedVectorType>(SV->getOperand(0)->getType())->getNumElements();
  unsigned ShuffleMaskSize = SV->getShuffleMask().size();
  if (SVNumElements % ShuffleMaskSize != 0)
    return 0;
  unsigned GroupSize = SVNumElements / ShuffleMaskSize;
  if (GroupSize == 0 || (VL.size() % GroupSize) != 0)
    return 0;
  unsigned NumGroup = 0;
  for (size_t I = 0, E = VL.size(); I != E; I += GroupSize) {
    auto *SV = cast<ShuffleVectorInst>(VL[I]);
    Value *Src = SV->getOperand(0);
    ArrayRef<Value *> Group = VL.slice(I, GroupSize);
    SmallBitVector ExpectedIndex(GroupSize);
    if (!all_of(Group, [&](Value *V) {
          auto *SV = cast<ShuffleVectorInst>(V);
          // From the same source.
          if (SV->getOperand(0) != Src)
            return false;
          int Index;
          if (!SV->isExtractSubvectorMask(Index))
            return false;
          ExpectedIndex.set(Index / ShuffleMaskSize);
          return true;
        }))
      return 0;
    if (!ExpectedIndex.all())
      return 0;
    ++NumGroup;
  }
  assert(NumGroup == (VL.size() / GroupSize) && "Unexpected number of groups");
  return NumGroup;
}
 
SmallVector<int> calculateShufflevectorMask(ArrayRef<Value *> VL) {
  assert(getShufflevectorNumGroups(VL) && "Not supported shufflevector usage.");
  auto *SV = cast<ShuffleVectorInst>(VL.front());
  unsigned SVNumElements =
      cast<FixedVectorType>(SV->getOperand(0)->getType())->getNumElements();
  SmallVector<int> Mask;
  unsigned AccumulateLength = 0;
  for (Value *V : VL) {
    auto *SV = cast<ShuffleVectorInst>(V);
    for (int M : SV->getShuffleMask())
      Mask.push_back(M == PoisonMaskElem ? PoisonMaskElem
                                         : AccumulateLength + M);
    AccumulateLength += SVNumElements;
  }
  return Mask;
}
 
/// Checks if the vector of instructions can be represented as a shuffle, like:
/// %x0 = extractelement <4 x i8> %x, i32 0
/// %x3 = extractelement <4 x i8> %x, i32 3
/// %y1 = extractelement <4 x i8> %y, i32 1
/// %y2 = extractelement <4 x i8> %y, i32 2
/// %x0x0 = mul i8 %x0, %x0
/// %x3x3 = mul i8 %x3, %x3
/// %y1y1 = mul i8 %y1, %y1
/// %y2y2 = mul i8 %y2, %y2
/// %ins1 = insertelement <4 x i8> poison, i8 %x0x0, i32 0
/// %ins2 = insertelement <4 x i8> %ins1, i8 %x3x3, i32 1
/// %ins3 = insertelement <4 x i8> %ins2, i8 %y1y1, i32 2
/// %ins4 = insertelement <4 x i8> %ins3, i8 %y2y2, i32 3
/// ret <4 x i8> %ins4
/// can be transformed into:
/// %1 = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 3, i32 5,
///                                                         i32 6>
/// %2 = mul <4 x i8> %1, %1
/// ret <4 x i8> %2
/// Mask will return the Shuffle Mask equivalent to the extracted elements.
/// TODO: Can we split off and reuse the shuffle mask detection from
/// ShuffleVectorInst/getShuffleCost?
std::optional<TargetTransformInfo::ShuffleKind>
isFixedVectorShuffle(ArrayRef<Value *> VL, SmallVectorImpl<int> &Mask,
                     AssumptionCache *AC) {
  const auto *It = find_if(VL, IsaPred<ExtractElementInst>);
  if (It == VL.end())
    return std::nullopt;
  unsigned Size = accumulate(VL, 0u, [](unsigned S, Value *V) {
    auto *EI = dyn_cast<ExtractElementInst>(V);
    if (!EI)
      return S;
    auto *VTy = dyn_cast<FixedVectorType>(EI->getVectorOperandType());
    if (!VTy)
      return S;
    return std::max(S, VTy->getNumElements());
  });
 
  Value *Vec1 = nullptr;
  Value *Vec2 = nullptr;
  bool HasNonUndefVec = any_of(make_isa_range<ExtractElementInst>(VL),
                               [&](ExtractElementInst *EE) {
                                 Value *Vec = EE->getVectorOperand();
                                 if (isa<UndefValue>(Vec))
                                   return false;
                                 return isGuaranteedNotToBePoison(Vec, AC);
                               });
  enum ShuffleMode { Unknown, Select, Permute };
  ShuffleMode CommonShuffleMode = Unknown;
  Mask.assign(VL.size(), PoisonMaskElem);
  for (unsigned I = 0, E = VL.size(); I < E; ++I) {
    // Undef, or a copyable lane modeled on an extract main op, can be
    // represented as an undef element in a vector.
    if (isa<UndefValue>(VL[I]))
      continue;
    auto *EI = dyn_cast<ExtractElementInst>(VL[I]);
    if (!EI)
      continue;
    if (isa<ScalableVectorType>(EI->getVectorOperandType()))
      return std::nullopt;
    auto *Vec = EI->getVectorOperand();
    // We can extractelement from undef or poison vector.
    if (isUndefVector</*isPoisonOnly=*/true>(Vec).all())
      continue;
    // All vector operands must have the same number of vector elements.
    if (isa<UndefValue>(Vec)) {
      Mask[I] = I;
    } else {
      if (isa<UndefValue>(EI->getIndexOperand()))
        continue;
      auto *Idx = dyn_cast<ConstantInt>(EI->getIndexOperand());
      if (!Idx)
        return std::nullopt;
      // Undefined behavior if Idx is negative or >= Size.
      if (Idx->getValue().uge(Size))
        continue;
      unsigned IntIdx = Idx->getValue().getZExtValue();
      Mask[I] = IntIdx;
    }
    if (isUndefVector(Vec).all() && HasNonUndefVec)
      continue;
    // For correct shuffling we have to have at most 2 different vector operands
    // in all extractelement instructions.
    if (!Vec1 || Vec1 == Vec) {
      Vec1 = Vec;
    } else if (!Vec2 || Vec2 == Vec) {
      Vec2 = Vec;
      Mask[I] += Size;
    } else {
      return std::nullopt;
    }
    if (CommonShuffleMode == Permute)
      continue;
    // If the extract index is not the same as the operation number, it is a
    // permutation.
    if (Mask[I] % Size != I) {
      CommonShuffleMode = Permute;
      continue;
    }
    CommonShuffleMode = Select;
  }
  // If we're not crossing lanes in different vectors, consider it as blending.
  if (CommonShuffleMode == Select && Vec2)
    return TargetTransformInfo::SK_Select;
  // If Vec2 was never used, we have a permutation of a single vector, otherwise
  // we have permutation of 2 vectors.
  return Vec2 ? TargetTransformInfo::SK_PermuteTwoSrc
              : TargetTransformInfo::SK_PermuteSingleSrc;
}
 
Value *createInsertVector(
    IRBuilderBase &Builder, Value *Vec, Value *V, unsigned Index,
    function_ref<Value *(Value *, Value *, ArrayRef<int>)> Generator) {
  if (isa<PoisonValue>(Vec) && isa<PoisonValue>(V))
    return Vec;
  const unsigned SubVecVF = getNumElements(V->getType());
  // Create shuffle, insertvector requires that index is multiple of
  // the subvector length.
  const unsigned VecVF = getNumElements(Vec->getType());
  SmallVector<int> Mask(VecVF, PoisonMaskElem);
  if (isa<PoisonValue>(Vec)) {
    auto *Begin = std::next(Mask.begin(), Index);
    std::iota(Begin, std::next(Begin, SubVecVF), 0);
    Vec = Builder.CreateShuffleVector(V, Mask);
    return Vec;
  }
  std::iota(Mask.begin(), Mask.end(), 0);
  std::iota(std::next(Mask.begin(), Index),
            std::next(Mask.begin(), Index + SubVecVF), VecVF);
  if (Generator)
    return Generator(Vec, V, Mask);
  // 1. Resize V to the size of Vec.
  SmallVector<int> ResizeMask(VecVF, PoisonMaskElem);
  std::iota(ResizeMask.begin(), std::next(ResizeMask.begin(), SubVecVF), 0);
  V = Builder.CreateShuffleVector(V, ResizeMask);
  // 2. Insert V into Vec.
  return Builder.CreateShuffleVector(Vec, V, Mask);
}
 
Value *createExtractVector(IRBuilderBase &Builder, Value *Vec,
                           unsigned SubVecVF, unsigned Index) {
  SmallVector<int> Mask(SubVecVF, PoisonMaskElem);
  std::iota(Mask.begin(), Mask.end(), Index);
  return Builder.CreateShuffleVector(Vec, Mask);
}
 
SmallBitVector buildUseMask(int VF, ArrayRef<int> Mask, UseMask MaskArg) {
  SmallBitVector UseMask(VF, true);
  for (auto [Idx, Value] : enumerate(Mask)) {
    if (Value == PoisonMaskElem) {
      if (MaskArg == UseMask::UndefsAsMask)
        UseMask.reset(Idx);
      continue;
    }
    if (MaskArg == UseMask::FirstArg && Value < VF)
      UseMask.reset(Value);
    else if (MaskArg == UseMask::SecondArg && Value >= VF)
      UseMask.reset(Value - VF);
  }
  return UseMask;
}
 
template <bool IsPoisonOnly>
SmallBitVector isUndefVector(const Value *V, const SmallBitVector &UseMask) {
  SmallBitVector Res(UseMask.empty() ? 1 : UseMask.size(), true);
  using T = std::conditional_t<IsPoisonOnly, PoisonValue, UndefValue>;
  if (isa<T>(V))
    return Res;
  auto *VecTy = dyn_cast<FixedVectorType>(V->getType());
  if (!VecTy)
    return Res.reset();
  auto *C = dyn_cast<Constant>(V);
  if (!C) {
    if (!UseMask.empty()) {
      const Value *Base = V;
      while (auto *II = dyn_cast<InsertElementInst>(Base)) {
        Base = II->getOperand(0);
        if (isa<T>(II->getOperand(1)))
          continue;
        std::optional<unsigned> Idx = getElementIndex(II);
        if (!Idx) {
          Res.reset();
          return Res;
        }
        if (*Idx < UseMask.size() && !UseMask.test(*Idx))
          Res.reset(*Idx);
      }
      // TODO: Add analysis for shuffles here too.
      if (V == Base) {
        Res.reset();
      } else {
        SmallBitVector SubMask(UseMask.size(), false);
        Res &= isUndefVector<IsPoisonOnly>(Base, SubMask);
      }
    } else {
      Res.reset();
    }
    return Res;
  }
  for (unsigned I = 0, E = VecTy->getNumElements(); I != E; ++I) {
    if (Constant *Elem = C->getAggregateElement(I))
      if (!isa<T>(Elem) &&
          (UseMask.empty() || (I < UseMask.size() && !UseMask.test(I))))
        Res.reset(I);
  }
  return Res;
}
 
template SmallBitVector isUndefVector<false>(const Value *,
                                             const SmallBitVector &);
template SmallBitVector isUndefVector<true>(const Value *,
                                            const SmallBitVector &);
 
bool doesInTreeUserNeedToExtract(Value *Scalar, Instruction *UserInst,
                                 TargetLibraryInfo *TLI,
                                 const TargetTransformInfo *TTI) {
  if (!UserInst)
    return false;
  unsigned Opcode = UserInst->getOpcode();
  switch (Opcode) {
  case Instruction::Load: {
    LoadInst *LI = cast<LoadInst>(UserInst);
    return (LI->getPointerOperand() == Scalar);
  }
  case Instruction::Store: {
    StoreInst *SI = cast<StoreInst>(UserInst);
    return (SI->getPointerOperand() == Scalar);
  }
  case Instruction::Call: {
    CallInst *CI = cast<CallInst>(UserInst);
    Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, TLI);
    return any_of(enumerate(CI->args()), [&](auto &&Arg) {
      return isVectorIntrinsicWithScalarOpAtArg(ID, Arg.index(), TTI) &&
             Arg.value().get() == Scalar;
    });
  }
  default:
    return false;
  }
}
 
MemoryLocation getLocation(Instruction *I) {
  if (StoreInst *SI = dyn_cast<StoreInst>(I))
    return MemoryLocation::get(SI);
  if (LoadInst *LI = dyn_cast<LoadInst>(I))
    return MemoryLocation::get(LI);
  return MemoryLocation();
}
 
bool isSimple(Instruction *I) {
  if (LoadInst *LI = dyn_cast<LoadInst>(I))
    return LI->isSimple();
  if (StoreInst *SI = dyn_cast<StoreInst>(I))
    return SI->isSimple();
  if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(I))
    return !MI->isVolatile();
  return true;
}
 
bool isSelectedBaseLoad(Type *ScalarTy, ArrayRef<Value *> PointerOps,
                        const DataLayout &DL, Value *&TrueBase,
                        Value *&FalseBase,
                        SmallVectorImpl<Value *> &Conditions) {
  TrueBase = nullptr;
  FalseBase = nullptr;
  uint64_t ScalarSize = DL.getTypeStoreSize(ScalarTy);
  Conditions.assign(PointerOps.size(), nullptr);
  for (auto [Idx, P] : enumerate(PointerOps)) {
    Value *Base = P;
    uint64_t Offset = 0;
    if (auto *GEP = dyn_cast<GetElementPtrInst>(P)) {
      APInt OffsetAP(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
      if (!GEP->accumulateConstantOffset(DL, OffsetAP) || OffsetAP.isNegative())
        return false;
      Offset = OffsetAP.getZExtValue();
      Base = GEP->getPointerOperand();
    }
    auto *Sel = dyn_cast<SelectInst>(Base);
    if (!Sel)
      return false;
    Value *T = Sel->getTrueValue();
    Value *F = Sel->getFalseValue();
    if (!TrueBase) {
      if (T == F)
        return false;
      TrueBase = T;
      FalseBase = F;
    } else if (TrueBase != T || FalseBase != F) {
      return false;
    }
    // Lane Idx must be at exactly Base + Idx * sizeof(ScalarTy); codegen reads
    // contiguously from TrueBase/FalseBase starting at lane 0.
    if (Offset != static_cast<uint64_t>(Idx) * ScalarSize)
      return false;
    Conditions[Idx] = Sel->getCondition();
  }
  return TrueBase != nullptr;
}
 
Type *getCommonGEPIndexType(ArrayRef<Value *> VL, Instruction *VL0,
                            function_ref<bool(Value *)> IsGEPLane,
                            const DataLayout &DL) {
  constexpr unsigned IndexIdx = 1;
  Type *VL0Ty = VL0->getOperand(IndexIdx)->getType();
  Type *PtrIdxTy =
      DL.getIndexType(VL0->getOperand(0)->getType()->getScalarType());
  bool AllSameTy = true;
  bool HasNonConstIdx = false;
  bool ConstsFitVL0Ty = true;
  for (Value *V : make_filter_range(VL, IsGEPLane)) {
    Value *Op = cast<GetElementPtrInst>(V)->getOperand(IndexIdx);
    if (Op->getType() != VL0Ty)
      AllSameTy = false;
    auto *CI = dyn_cast<ConstantInt>(Op);
    if (!CI) {
      // Non-constant indices are not cast, they must have the main op type.
      if (Op->getType() != VL0Ty)
        return nullptr;
      HasNonConstIdx = true;
      continue;
    }
    if (!CI->getValue().isSignedIntN(VL0Ty->getIntegerBitWidth()))
      ConstsFitVL0Ty = false;
  }
  if (AllSameTy)
    return VL0Ty;
  if (!HasNonConstIdx || VL0Ty == PtrIdxTy)
    return PtrIdxTy;
  return ConstsFitVL0Ty ? VL0Ty : nullptr;
}
 
bool isCopyableGEPAddressVector(ArrayRef<Value *> PointerOps) {
  SmallPtrSet<Value *, 16> UniquePtrs(llvm::from_range, PointerOps);
  if (UniquePtrs.size() != PointerOps.size())
    return false;
  auto IsConstantOffsetPtr = [](Value *P) {
    auto *GEP = dyn_cast<GetElementPtrInst>(P);
    return !GEP ||
           (GEP->getNumOperands() == 2 && isConstant(GEP->getOperand(1)));
  };
  auto *RefIt = find_if_not(PointerOps, IsConstantOffsetPtr);
  if (RefIt == PointerOps.end())
    return false;
  auto *RefGEP = dyn_cast<GetElementPtrInst>(*RefIt);
  if (!RefGEP || RefGEP->getNumOperands() != 2)
    return false;
  Value *Base = RefGEP->getPointerOperand();
  Type *PtrTy = RefGEP->getType();
  Type *SrcElemTy = RefGEP->getSourceElementType();
  // The stride and the (optional) cast opcode of the runtime indices.
  Value *Stride = nullptr;
  unsigned CastOpcode = 0;
  for (Value *P : PointerOps) {
    if (P->getType() != PtrTy)
      return false;
    if (P == Base)
      continue;
    auto *GEP = dyn_cast<GetElementPtrInst>(P);
    if (!GEP || GEP->getNumOperands() != 2 ||
        GEP->getPointerOperand() != Base ||
        GEP->getSourceElementType() != SrcElemTy)
      return false;
    Value *Idx = GEP->getOperand(1);
    if (isConstant(Idx))
      continue;
    unsigned LaneCastOpcode = 0;
    if (auto *Cast = dyn_cast<CastInst>(Idx)) {
      LaneCastOpcode = Cast->getOpcode();
      Idx = Cast->getOperand(0);
    }
    Value *LaneStride = Idx;
    if (auto *BO = dyn_cast<BinaryOperator>(Idx)) {
      if (isa<Constant>(BO->getOperand(1)))
        LaneStride = BO->getOperand(0);
      else if (isa<Constant>(BO->getOperand(0)))
        LaneStride = BO->getOperand(1);
    }
    if (!Stride) {
      Stride = LaneStride;
      CastOpcode = LaneCastOpcode;
      continue;
    }
    if (LaneStride != Stride || LaneCastOpcode != CastOpcode)
      return false;
  }
  return Stride != nullptr;
}
 
void addMask(SmallVectorImpl<int> &Mask, ArrayRef<int> SubMask,
             bool ExtendingManyInputs) {
  if (SubMask.empty())
    return;
  assert(
      (!ExtendingManyInputs || SubMask.size() > Mask.size() ||
       // Check if input scalars were extended to match the size of other node.
       (SubMask.size() == Mask.size() && Mask.back() == PoisonMaskElem)) &&
      "SubMask with many inputs support must be larger than the mask.");
  if (Mask.empty()) {
    Mask.append(SubMask.begin(), SubMask.end());
    return;
  }
  SmallVector<int> NewMask(SubMask.size(), PoisonMaskElem);
  int TermValue = std::min(Mask.size(), SubMask.size());
  for (int I = 0, E = SubMask.size(); I < E; ++I) {
    if (SubMask[I] == PoisonMaskElem ||
        (!ExtendingManyInputs &&
         (SubMask[I] >= TermValue || Mask[SubMask[I]] >= TermValue)))
      continue;
    NewMask[I] = Mask[SubMask[I]];
  }
  Mask.swap(NewMask);
}
 
void fixupOrderingIndices(MutableArrayRef<unsigned> Order) {
  const size_t Sz = Order.size();
  SmallBitVector UnusedIndices(Sz, /*t=*/true);
  SmallBitVector MaskedIndices(Sz);
  for (unsigned I = 0; I < Sz; ++I) {
    if (Order[I] < Sz)
      UnusedIndices.reset(Order[I]);
    else
      MaskedIndices.set(I);
  }
  if (MaskedIndices.none())
    return;
  assert(UnusedIndices.count() == MaskedIndices.count() &&
         "Non-synced masked/available indices.");
  int Idx = UnusedIndices.find_first();
  int MIdx = MaskedIndices.find_first();
  while (MIdx >= 0) {
    assert(Idx >= 0 && "Indices must be synced.");
    Order[MIdx] = Idx;
    Idx = UnusedIndices.find_next(Idx);
    MIdx = MaskedIndices.find_next(MIdx);
  }
}
 
SmallBitVector getAltInstrMask(ArrayRef<Value *> VL, Type *ScalarTy,
                               unsigned Opcode0, unsigned Opcode1) {
  unsigned ScalarTyNumElements = getNumElements(ScalarTy);
  SmallBitVector OpcodeMask(VL.size() * ScalarTyNumElements, false);
  for (unsigned Lane : seq<unsigned>(VL.size())) {
    if (isa<PoisonValue>(VL[Lane]))
      continue;
    if (cast<Instruction>(VL[Lane])->getOpcode() == Opcode1)
      OpcodeMask.set(Lane * ScalarTyNumElements,
                     Lane * ScalarTyNumElements + ScalarTyNumElements);
  }
  return OpcodeMask;
}
 
SmallVector<Constant *> replicateMask(ArrayRef<Constant *> Val, unsigned VF) {
  assert(none_of(Val, [](Constant *C) { return C->getType()->isVectorTy(); }) &&
         "Expected scalar constants.");
  SmallVector<Constant *> NewVal(Val.size() * VF);
  for (auto [I, V] : enumerate(Val))
    std::fill_n(NewVal.begin() + I * VF, VF, V);
  return NewVal;
}
 
Intrinsic::ID getMaskedDivRemIntrinsic(unsigned Opcode) {
  switch (Opcode) {
  case Instruction::UDiv:
    return Intrinsic::masked_udiv;
  case Instruction::SDiv:
    return Intrinsic::masked_sdiv;
  case Instruction::URem:
    return Intrinsic::masked_urem;
  case Instruction::SRem:
    return Intrinsic::masked_srem;
  default:
    llvm_unreachable("Unexpected opcode");
  }
}
 
/// Returns true if \p I is a part of a single-use chain, computing an address,
/// which does not pay off the vectorization: a constant table is accessed by a
/// gather, while the indices, unrelated between the lanes, require a full
/// buildvector, unlike the ones, shifted by a constant from a common base.
static bool isNonProfitableIndex(const Instruction *I) {
  constexpr unsigned MaxIndexChainLength = 3;
  // A constant shift of a common base is a cheap buildvector, while the loads
  // are vectorized together with the indices, computed from them.
  auto IsProfitableOperand = [](const Value *V) {
    if (isa<Constant>(V))
      return true;
    if (const auto *Cast = dyn_cast<CastInst>(V); Cast && Cast->hasOneUse())
      V = Cast->getOperand(0);
    return isa<LoadInst>(V);
  };
  const User *U = I->user_back();
  for ([[maybe_unused]] unsigned _ : seq<unsigned>(MaxIndexChainLength)) {
    if (const auto *GEP = dyn_cast<GetElementPtrInst>(U))
      return isa<Constant>(GEP->getPointerOperand()) ||
             none_of(I->operand_values(), IsProfitableOperand);
    if (!isa<Instruction>(U) || !U->hasOneUse())
      return false;
    U = U->user_back();
  }
  return false;
}
 
bool isOnceUsedSeed(const Instruction *I) {
  if (!I->hasOneUse() || isNonProfitableIndex(I))
    return false;
  // The operation with the identity or the absorbing constant is folded away
  // before the codegen, the vector node only repacks the lanes.
  if (const auto *BO = dyn_cast<BinaryOperator>(I)) {
    unsigned Opcode = BO->getOpcode();
    Type *Ty = BO->getType();
    for (unsigned Idx : seq<unsigned>(2)) {
      const auto *C = dyn_cast<Constant>(BO->getOperand(Idx));
      if (C && (C == ConstantExpr::getBinOpIdentity(
                         Opcode, Ty, /*AllowRHSConstant=*/Idx == 1) ||
                C == ConstantExpr::getBinOpAbsorber(
                         Opcode, Ty, /*AllowLHSConstant=*/Idx == 0)))
        return false;
    }
  }
  const User *U = I->user_back();
  if (isa<ExtractElementInst, ExtractValueInst>(I))
    return isa<InsertElementInst, InsertValueInst>(U);
  if (isa<CastInst>(I))
    return !isa<FPToSIInst, FPToUIInst>(I) &&
           (!isa<CastInst>(U) || U->hasOneUse());
  return isa<BinaryOperator, UnaryOperator, SelectInst, FreezeInst, CallInst>(
      I);
}
 
Instruction *lookThroughCastRoundTrip(Value *V, bool MustBeElidable) {
  auto *Wide = dyn_cast<FPExtInst>(V);
  if (!Wide || !Wide->hasOneUse())
    return nullptr;
  auto *Narrow = dyn_cast<FPTruncInst>(Wide->getOperand(0));
  if (!Narrow || !Narrow->hasOneUse())
    return nullptr;
  Value *Src = Narrow->getOperand(0);
  if (!isa<Instruction>(Src) || Src->getType() != Wide->getType())
    return nullptr;
  if (MustBeElidable && !(Wide->hasAllowContract() && Wide->hasNoNaNs() &&
                          Wide->hasNoInfs() && Narrow->hasAllowContract()))
    return nullptr;
  return Narrow;
}
 
namespace {
 
/// Shifts and the mask accumulated from the narrow ops on the current path:
/// the shifts above and at the narrow level, the bitwidth of the narrow ops
/// (0 if none) and the mask from the absorbed narrow ands.
struct NarrowedChainState {
  unsigned Shift = 0;
  unsigned NarrowShift = 0;
  unsigned NarrowBW = 0;
  APInt NarrowMask = APInt(1, 0);
 
  /// The mask for the absorbed narrow ops in the leaf type, applied before
  /// widening and shifting; all-ones if nothing was absorbed.
  APInt getMask(unsigned LeafBW) const {
    if (NarrowBW == 0)
      return APInt::getAllOnes(LeafBW);
    return (NarrowMask & (APInt::getAllOnes(NarrowBW) << NarrowShift))
        .lshr(NarrowShift)
        .trunc(LeafBW);
  }
};
 
} // namespace
 
static void
collectNarrowedLeavesImpl(Value *V, unsigned RdxOpcode, unsigned WideBW,
                          NarrowedChainState S, unsigned Depth,
                          unsigned MaxDepth,
                          SmallVectorImpl<NarrowedLeafInfo> &Leaves,
                          SmallVectorImpl<Instruction *> &ChainInsts) {
  if (Depth < MaxDepth) {
    if (auto *Z = dyn_cast<ZExtInst>(V);
        Z && Z->getSrcTy()->isIntegerTy() && !Z->getSrcTy()->isIntegerTy(1)) {
      ChainInsts.push_back(Z);
      return collectNarrowedLeavesImpl(Z->getOperand(0), RdxOpcode, WideBW, S,
                                       Depth + 1, MaxDepth, Leaves, ChainInsts);
    }
    if (auto *BO = dyn_cast<BinaryOperator>(V)) {
      if (BO->getOpcode() == RdxOpcode) {
        ChainInsts.push_back(BO);
        collectNarrowedLeavesImpl(BO->getOperand(0), RdxOpcode, WideBW, S,
                                  Depth + 1, MaxDepth, Leaves, ChainInsts);
        collectNarrowedLeavesImpl(BO->getOperand(1), RdxOpcode, WideBW, S,
                                  Depth + 1, MaxDepth, Leaves, ChainInsts);
        return;
      }
      const APInt *Amt;
      unsigned BW = V->getType()->getScalarSizeInBits();
      auto *Z = dyn_cast<ZExtInst>(BO->getOperand(0));
      if (BO->getOpcode() == Instruction::Shl && Z && S.NarrowBW == 0 &&
          match(BO->getOperand(1), m_APInt(Amt)) && Amt->ult(BW) &&
          Z->getSrcTy()->isIntegerTy() && !Z->getSrcTy()->isIntegerTy(1) &&
          (BW == WideBW ||
           Z->getSrcTy()->getIntegerBitWidth() + Amt->getZExtValue() <= BW) &&
          S.Shift + Amt->getZExtValue() < WideBW) {
        ChainInsts.push_back(BO);
        ChainInsts.push_back(Z);
        S.Shift += Amt->getZExtValue();
        return collectNarrowedLeavesImpl(Z->getOperand(0), RdxOpcode, WideBW, S,
                                         Depth + 1, MaxDepth, Leaves,
                                         ChainInsts);
      }
      // Narrow shls fold into the shift and narrow ands into the mask; the
      // mask clears the bits the shls shift out. Only same-width ops compose
      // on one path, and the combined shift must stay a valid shift amount in
      // both types.
      if (BW < WideBW && (S.NarrowBW == 0 || BW == S.NarrowBW)) {
        if (BO->getOpcode() == Instruction::Shl &&
            match(BO->getOperand(1), m_APInt(Amt)) && Amt->ult(BW) &&
            S.NarrowShift + Amt->getZExtValue() < BW &&
            S.Shift + S.NarrowShift + Amt->getZExtValue() < WideBW) {
          ChainInsts.push_back(BO);
          if (BO->hasNoUnsignedWrap() && S.NarrowBW == 0) {
            S.Shift += Amt->getZExtValue();
            // Lossless shls shift out only known-zero bits; record them as
            // the mask so matching lanes can form a splat.
            S.NarrowBW = BW;
            S.NarrowMask = APInt::getLowBitsSet(BW, BW - Amt->getZExtValue());
          } else {
            if (S.NarrowBW == 0) {
              S.NarrowBW = BW;
              S.NarrowMask = APInt::getAllOnes(BW);
            }
            S.NarrowShift += Amt->getZExtValue();
          }
          return collectNarrowedLeavesImpl(BO->getOperand(0), RdxOpcode, WideBW,
                                           S, Depth + 1, MaxDepth, Leaves,
                                           ChainInsts);
        }
        Value *X;
        if (match(BO, m_c_And(m_Value(X), m_APInt(Amt)))) {
          ChainInsts.push_back(BO);
          if (S.NarrowBW == 0) {
            S.NarrowBW = BW;
            S.NarrowMask = APInt::getAllOnes(BW);
          }
          S.NarrowMask &= *Amt << S.NarrowShift;
          return collectNarrowedLeavesImpl(X, RdxOpcode, WideBW, S, Depth + 1,
                                           MaxDepth, Leaves, ChainInsts);
        }
      }
    }
  }
  Leaves.emplace_back(V, S.Shift + S.NarrowShift,
                      S.getMask(V->getType()->getScalarSizeInBits()));
}
 
void collectNarrowedLeaves(Value *V, unsigned RdxOpcode, unsigned WideBW,
                           unsigned MaxDepth,
                           SmallVectorImpl<NarrowedLeafInfo> &Leaves,
                           SmallVectorImpl<Instruction *> &ChainInsts) {
  collectNarrowedLeavesImpl(V, RdxOpcode, WideBW, NarrowedChainState(),
                            /*Depth=*/0, MaxDepth, Leaves, ChainInsts);
}
 
TargetTransformInfo::TargetCostKind getSLPCostKind(const Function *F) {
  assert(F && "Expected function.");
  return F->hasOptSize() ? TTI::TCK_CodeSize : TTI::TCK_RecipThroughput;
}
 
/// Checks if \p V is a zero-extended sub-field of a wider integer scalar.
/// Returns the source scalar, the field width and the field offset.
static std::optional<std::tuple<Value *, unsigned, unsigned>>
matchExtractedField(Value *V) {
  if (!V->getType()->isIntegerTy())
    return std::nullopt;
  // Field offset for the field-aligned shift amount, if the shifted value of
  // the given bit width keeps at least one full field.
  auto GetFieldOffset = [](const APInt *Amt, unsigned BitWidth,
                           unsigned FieldWidth) -> std::optional<unsigned> {
    uint64_t ShAmt = Amt->getLimitedValue(BitWidth);
    if (ShAmt % FieldWidth != 0 || ShAmt + FieldWidth > BitWidth)
      return std::nullopt;
    return ShAmt / FieldWidth;
  };
  // Checks if the low bits of Val are a sub-field of the given width of a
  // wider integer scalar. Val is a scalar integer, since V is one, and so is
  // the matched source.
  auto MatchLowField =
      [&](Value *Val,
          unsigned FieldWidth) -> std::optional<std::pair<Value *, unsigned>> {
    Value *Src;
    const APInt *Amt;
    // Only the low bits of Val are observed, so lshr and ashr are equivalent.
    if (match(Val, m_Trunc(m_Shr(m_Value(Src), m_APInt(Amt)))) ||
        match(Val, m_Shr(m_Value(Src), m_APInt(Amt)))) {
      if (std::optional<unsigned> Offset = GetFieldOffset(
              Amt, Src->getType()->getIntegerBitWidth(), FieldWidth)) {
        // The truncation of the shifted value keeps the field, look through it.
        match(Src, m_Trunc(m_Value(Src)));
        return std::make_pair(Src, *Offset);
      }
      return std::nullopt;
    }
    if (match(Val, m_Trunc(m_Value(Src))) &&
        Src->getType()->getIntegerBitWidth() >= FieldWidth)
      return std::make_pair(Src, 0u);
    // Val itself is the source of its low field.
    if (Val->getType()->getIntegerBitWidth() > FieldWidth)
      return std::make_pair(Val, 0u);
    return std::nullopt;
  };
  Value *Val;
  const APInt *Mask;
  // and Val, (1 << FieldWidth) - 1 or zext i<FieldWidth> Val - the low bits of
  // Val.
  unsigned FieldWidth = 0;
  if (match(V, m_c_And(m_Value(Val), m_APInt(Mask))) && Mask->isMask())
    FieldWidth = Mask->popcount();
  else if (match(V, m_ZExt(m_Value(Val))))
    FieldWidth = Val->getType()->getIntegerBitWidth();
  if (FieldWidth != 0) {
    if (std::optional<std::pair<Value *, unsigned>> Field =
            MatchLowField(Val, FieldWidth))
      return std::make_tuple(Field->first, FieldWidth, Field->second);
    return std::nullopt;
  }
  unsigned LaneWidth = V->getType()->getIntegerBitWidth();
  Value *Src;
  const APInt *Amt;
  if (match(V, m_Trunc(m_LShr(m_Value(Src), m_APInt(Amt))))) {
    unsigned SrcWidth = Src->getType()->getIntegerBitWidth();
    uint64_t ShAmt = Amt->getLimitedValue(SrcWidth);
    // The field itself, if the lane width is the field width.
    if (std::optional<unsigned> Offset =
            GetFieldOffset(Amt, SrcWidth, LaneWidth))
      return std::make_tuple(Src, LaneWidth, *Offset);
    // The zero-extended top field of the source.
    unsigned FieldWidth = SrcWidth - ShAmt;
    if (FieldWidth > 0 && FieldWidth < LaneWidth && ShAmt % FieldWidth == 0)
      return std::make_tuple(Src, FieldWidth, ShAmt / FieldWidth);
    return std::nullopt;
  }
  if (match(V, m_LShr(m_Value(Src), m_APInt(Amt)))) {
    // The zero-extended top field of the source, if the result keeps exactly
    // one field. Look through a truncation of the shifted value.
    unsigned ShfWidth = Src->getType()->getIntegerBitWidth();
    uint64_t ShAmt = Amt->getLimitedValue(ShfWidth);
    unsigned FieldWidth = ShfWidth - ShAmt;
    if (FieldWidth > 0 && ShAmt % FieldWidth == 0) {
      match(Src, m_Trunc(m_Value(Src)));
      return std::make_tuple(Src, FieldWidth, ShAmt / FieldWidth);
    }
    return std::nullopt;
  }
  if (match(V, m_Trunc(m_Value(Src))))
    return std::make_tuple(Src, LaneWidth, 0u);
  return std::nullopt;
}
 
std::optional<std::tuple<Value *, unsigned, SmallVector<int>>>
matchGatheredExtractedFields(ArrayRef<Value *> VL, const DataLayout &DL) {
  // Splats are emitted as broadcasts, sub-fields of a constant are folded.
  // The bitcast to the field vector maps lane 0 to the least significant
  // field on little-endian targets only.
  if (VL.size() < 2 || !VL.front()->getType()->isIntegerTy() || isSplat(VL) ||
      DL.isBigEndian())
    return std::nullopt;
  Value *Src = nullptr;
  unsigned FieldWidth = 0;
  SmallVector<int> Mask(VL.size(), PoisonMaskElem);
  for (auto [Idx, V] : make_filter_range(enumerate(VL), [](const auto &P) {
         return !isa<UndefValue>(P.value());
       })) {
    if (V->getType() != VL.front()->getType())
      return std::nullopt;
    std::optional<std::tuple<Value *, unsigned, unsigned>> Field =
        matchExtractedField(V);
    if (!Field || (Src && (Src != std::get<0>(*Field) ||
                           FieldWidth != std::get<1>(*Field))))
      return std::nullopt;
    Src = std::get<0>(*Field);
    FieldWidth = std::get<1>(*Field);
    Mask[Idx] = std::get<2>(*Field);
  }
  // The field width is a whole number of bytes and divides the source
  // exactly, same as for the packing layout, so the source bitcasts to the
  // field vector.
  if (!Src || isa<Constant>(Src) || FieldWidth % 8 != 0 ||
      Src->getType()->getIntegerBitWidth() % FieldWidth != 0)
    return std::nullopt;
  // The same field in every lane is a splat, emitted as a broadcast.
  if (all_of(Mask, [First = *find_if(Mask, not_equal_to(PoisonMaskElem))](
                       int MaskElt) {
        return MaskElt == PoisonMaskElem || MaskElt == First;
      }))
    return std::nullopt;
  return std::make_tuple(Src, FieldWidth, std::move(Mask));
}
 
/// Deeper than the standard analysis recursion depth to keep the numeric
/// bound precise through arithmetic carry chains.
constexpr unsigned MaxBitPackAnalysisDepth = MaxAnalysisRecursionDepth + 2;
 
APInt getScalarMaxValue(const Value *V, unsigned Depth) {
  unsigned BitWidth = V->getType()->getScalarSizeInBits();
  const APInt Unknown = APInt::getAllOnes(BitWidth);
  if (Depth > MaxBitPackAnalysisDepth || !V->getType()->isIntegerTy())
    return Unknown;
  const APInt *C, *Amt;
  if (match(V, m_APInt(C)))
    return *C;
  Value *L, *R;
  if (match(V, m_Add(m_Value(L), m_Value(R))) ||
      match(V, m_Or(m_Value(L), m_Value(R))) ||
      match(V, m_Xor(m_Value(L), m_Value(R))))
    return getScalarMaxValue(L, Depth + 1)
        .uadd_sat(getScalarMaxValue(R, Depth + 1));
  if (match(V, m_NUWSub(m_Value(L), m_Value(R))))
    return getScalarMaxValue(L, Depth + 1);
  if (match(V, m_Mul(m_Value(L), m_Value(R))))
    return getScalarMaxValue(L, Depth + 1)
        .umul_sat(getScalarMaxValue(R, Depth + 1));
  if (match(V, m_And(m_Value(L), m_Value(R))))
    return APIntOps::umin(getScalarMaxValue(L, Depth + 1),
                          getScalarMaxValue(R, Depth + 1));
  if (match(V, m_LShr(m_Value(L), m_APInt(Amt))) && Amt->ult(BitWidth))
    return getScalarMaxValue(L, Depth + 1).lshr(*Amt);
  if (match(V, m_Shl(m_Value(L), m_APInt(Amt))) && Amt->ult(BitWidth)) {
    APInt LMax = getScalarMaxValue(L, Depth + 1);
    return LMax.getActiveBits() + Amt->getZExtValue() <= BitWidth
               ? LMax.shl(*Amt)
               : Unknown;
  }
  if (match(V, m_ZExt(m_Value(L))))
    return getScalarMaxValue(L, Depth + 1).zext(BitWidth);
  if (match(V, m_Trunc(m_Value(L)))) {
    APInt Max = getScalarMaxValue(L, Depth + 1);
    return Max.getActiveBits() <= BitWidth ? Max.trunc(BitWidth) : Unknown;
  }
  if (match(V, m_SExt(m_Value(L)))) {
    APInt Max = getScalarMaxValue(L, Depth + 1);
    return Max.isNonNegative() ? Max.zext(BitWidth) : Unknown;
  }
  Value *F;
  if (match(V, m_Select(m_Value(), m_Value(L), m_Value(F))))
    return APIntOps::umax(getScalarMaxValue(L, Depth + 1),
                          getScalarMaxValue(F, Depth + 1));
  return Unknown;
}
 
std::optional<BitPackInfo> computeBitPackInfo(unsigned BitWidth,
                                              ArrayRef<APInt> PossibleBits,
                                              ArrayRef<uint64_t> ShlAmts,
                                              ArrayRef<APInt> Masks) {
  unsigned NumElts = PossibleBits.size();
  BitPackInfo Info;
  Info.LShrAmts.assign(NumElts, 0);
  for (unsigned Idx : seq(NumElts)) {
    APInt Possible = PossibleBits[Idx].shl(ShlAmts[Idx]) & Masks[Idx];
    if (Possible.isZero())
      continue;
    unsigned Lo, W;
    if (!Possible.isShiftedMask(Lo, W))
      return std::nullopt;
    if (Info.FieldWidth == 0) {
      if (W % 8 != 0 || BitWidth % W != 0)
        return std::nullopt;
      Info.FieldWidth = W;
      Info.LaneOfField.assign(BitWidth / W, BitPackInfo::NoLane);
    }
    if (W != Info.FieldWidth || Lo % W != 0)
      return std::nullopt;
    unsigned Field = Lo / W;
    if (Info.LaneOfField[Field] != BitPackInfo::NoLane)
      return std::nullopt;
    Info.LaneOfField[Field] = Idx;
    Info.LShrAmts[Idx] = Lo - ShlAmts[Idx];
  }
  if (Info.FieldWidth == 0)
    return std::nullopt;
  return Info;
}
 
SmallVector<int> getBitPackMask(const BitPackInfo &Info, unsigned NumBytes,
                                unsigned NumElts, unsigned BytesPerLane) {
  unsigned BytesPerField = Info.FieldWidth / 8;
  SmallVector<int> Mask;
  for (unsigned J : seq(NumBytes)) {
    unsigned Lane = Info.LaneOfField[J / BytesPerField];
    Mask.push_back(Lane == BitPackInfo::NoLane
                       ? (int)(NumElts * BytesPerLane)
                       : (int)(Lane * BytesPerLane + J % BytesPerField));
  }
  return Mask;
}
 
Value *buildBitPack(IRBuilderBase &Builder, Value *X, const BitPackInfo &Info,
                    unsigned ShiftWidth, unsigned &NumInsts) {
  NumInsts = 0;
  auto *VecTy = cast<FixedVectorType>(X->getType());
  unsigned BitWidth = VecTy->getScalarSizeInBits();
  assert(BitWidth % 8 == 0 &&
         "The byte-multiple field width divides the result bit width.");
  unsigned NumElts = VecTy->getNumElements();
  Value *Y = X;
  if (ShiftWidth != BitWidth) {
    // Compacting a zext back to its source is free, use it directly.
    if (auto *Z = dyn_cast<ZExtInst>(X);
        Z && Z->getSrcTy()->getScalarSizeInBits() == ShiftWidth)
      Y = Z->getOperand(0);
    else {
      Y = Builder.CreateTrunc(
          Y, FixedVectorType::get(IntegerType::get(X->getContext(), ShiftWidth),
                                  NumElts));
      ++NumInsts;
    }
  }
  if (Info.needsShift()) {
    SmallVector<Constant *> Amts;
    for (uint64_t A : Info.LShrAmts)
      Amts.push_back(
          ConstantInt::get(IntegerType::get(X->getContext(), ShiftWidth), A));
    Y = Builder.CreateLShr(Y, ConstantVector::get(Amts));
    ++NumInsts;
  }
  unsigned InBytes = NumElts * (ShiftWidth / 8);
  auto *ByteTy = FixedVectorType::get(Builder.getInt8Ty(), InBytes);
  SmallVector<int> Mask =
      getBitPackMask(Info, BitWidth / 8, NumElts, ShiftWidth / 8);
  auto *IntTy = IntegerType::get(X->getContext(), BitWidth);
  // A plain byte reversal of the shifted lanes is a bswap.
  if (ShuffleVectorInst::isReverseMask(Mask, InBytes)) {
    NumInsts += 2;
    return Builder.CreateUnaryIntrinsic(Intrinsic::bswap,
                                        Builder.CreateBitCast(Y, IntTy));
  }
  // An identity byte order needs no shuffle.
  if (ShuffleVectorInst::isIdentityMask(Mask, InBytes)) {
    ++NumInsts;
    return Builder.CreateBitCast(Y, IntTy);
  }
  Value *Packed = Builder.CreateShuffleVector(
      Builder.CreateBitCast(Y, ByteTy),
      is_contained(Info.LaneOfField, BitPackInfo::NoLane)
          ? Constant::getNullValue(ByteTy)
          : PoisonValue::get(ByteTy),
      Mask);
  NumInsts += 3;
  return Builder.CreateBitCast(Packed, IntTy);
}
 
} // namespace llvm::slpvectorizer