//===------- VectorCombine.cpp - Optimize partial vector operations -------===//
//
// 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 pass optimizes scalar/vector interactions using target cost models. The
// transforms implemented here may not fit in traditional loop-based or SLP
// vectorization passes.
//
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/Vectorize/VectorCombine.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AssumptionCache.h"
#include "llvm/Analysis/BasicAliasAnalysis.h"
#include "llvm/Analysis/GlobalsModRef.h"
#include "llvm/Analysis/InstSimplifyFolder.h"
#include "llvm/Analysis/Loads.h"
#include "llvm/Analysis/TargetFolder.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/Analysis/VectorUtils.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/PatternMatch.h"
#include "llvm/IR/ProfDataUtils.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/KnownBits.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/LoopUtils.h"
#include <numeric>
#include <optional>
#include <queue>
#include <set>

#define DEBUG_TYPE "vector-combine"
#include "llvm/Transforms/Utils/InstructionWorklist.h"

using namespace llvm;
using namespace llvm::PatternMatch;

STATISTIC(NumVecLoad, "Number of vector loads formed");
STATISTIC(NumVecCmp, "Number of vector compares formed");
STATISTIC(NumVecBO, "Number of vector binops formed");
STATISTIC(NumVecCmpBO, "Number of vector compare + binop formed");
STATISTIC(NumShufOfBitcast, "Number of shuffles moved after bitcast");
STATISTIC(NumScalarOps, "Number of scalar unary + binary ops formed");
STATISTIC(NumScalarCmp, "Number of scalar compares formed");
STATISTIC(NumScalarIntrinsic, "Number of scalar intrinsic calls formed");

static cl::opt<bool> DisableVectorCombine(
    "disable-vector-combine", cl::init(false), cl::Hidden,
    cl::desc("Disable all vector combine transforms"));

static cl::opt<bool> DisableBinopExtractShuffle(
    "disable-binop-extract-shuffle", cl::init(false), cl::Hidden,
    cl::desc("Disable binop extract to shuffle transforms"));

static cl::opt<unsigned> MaxInstrsToScan(
    "vector-combine-max-scan-instrs", cl::init(30), cl::Hidden,
    cl::desc("Max number of instructions to scan for vector combining."));

static const unsigned InvalidIndex = std::numeric_limits<unsigned>::max();

namespace {
class VectorCombine {
public:
  VectorCombine(Function &F, const TargetTransformInfo &TTI,
                const DominatorTree &DT, AAResults &AA, AssumptionCache &AC,
                const DataLayout *DL, TTI::TargetCostKind CostKind,
                bool TryEarlyFoldsOnly)
      : F(F), Builder(*F.getParent(), InstSimplifyFolder(*DL)), TTI(TTI),
        DT(DT), AA(AA), DL(DL), CostKind(CostKind),
        SQ(*DL, /*TLI=*/nullptr, &DT, &AC),
        TryEarlyFoldsOnly(TryEarlyFoldsOnly) {}

  bool run();

private:
  Function &F;
  IRBuilder<InstSimplifyFolder> Builder;
  const TargetTransformInfo &TTI;
  const DominatorTree &DT;
  AAResults &AA;
  const DataLayout *DL;
  TTI::TargetCostKind CostKind;
  const SimplifyQuery SQ;

  /// If true, only perform beneficial early IR transforms. Do not introduce new
  /// vector operations.
  bool TryEarlyFoldsOnly;

  InstructionWorklist Worklist;

  /// Next instruction to iterate. It will be updated when it is erased by
  /// RecursivelyDeleteTriviallyDeadInstructions.
  Instruction *NextInst;

  // TODO: Direct calls from the top-level "run" loop use a plain "Instruction"
  //       parameter. That should be updated to specific sub-classes because the
  //       run loop was changed to dispatch on opcode.
  bool vectorizeLoadInsert(Instruction &I);
  bool widenSubvectorLoad(Instruction &I);
  ExtractElementInst *getShuffleExtract(ExtractElementInst *Ext0,
                                        ExtractElementInst *Ext1,
                                        unsigned PreferredExtractIndex) const;
  bool isExtractExtractCheap(ExtractElementInst *Ext0, ExtractElementInst *Ext1,
                             const Instruction &I,
                             ExtractElementInst *&ConvertToShuffle,
                             unsigned PreferredExtractIndex);
  Value *foldExtExtCmp(Value *V0, Value *V1, Value *ExtIndex, Instruction &I);
  Value *foldExtExtBinop(Value *V0, Value *V1, Value *ExtIndex, Instruction &I);
  bool foldExtractExtract(Instruction &I);
  bool foldInsExtFNeg(Instruction &I);
  bool foldInsExtBinop(Instruction &I);
  bool foldInsExtVectorToShuffle(Instruction &I);
  bool foldInsertScalarPartsToShuffle(Instruction &I);
  bool foldBitOpOfCastops(Instruction &I);
  bool foldBitOpOfCastConstant(Instruction &I);
  bool foldBitcastShuffle(Instruction &I);
  bool scalarizeOpOrCmp(Instruction &I);
  bool foldExtractedCmps(Instruction &I);
  bool foldSelectsFromBitcast(Instruction &I);
  bool foldBinopOfReductions(Instruction &I);
  bool foldInsertElementsToStores(Instruction &I);
  bool scalarizeLoad(Instruction &I);
  bool scalarizeLoadExtract(LoadInst *LI, VectorType *VecTy, Value *Ptr);
  bool scalarizeLoadBitcast(LoadInst *LI, VectorType *VecTy, Value *Ptr);
  bool scalarizeExtExtract(Instruction &I);
  bool foldConcatOfBoolMasks(Instruction &I);
  bool foldPermuteOfBinops(Instruction &I);
  bool foldShuffleOfBinops(Instruction &I);
  bool foldShuffleOfSelects(Instruction &I);
  bool foldShuffleOfCastops(Instruction &I);
  bool foldShuffleOfShuffles(Instruction &I);
  bool foldPermuteOfIntrinsic(Instruction &I);
  bool foldShufflesOfLengthChangingShuffles(Instruction &I);
  bool foldShuffleOfIntrinsics(Instruction &I);
  bool foldShuffleToIdentity(Instruction &I);
  bool foldShuffleFromReductions(Instruction &I);
  bool foldShuffleChainsToReduce(Instruction &I);
  bool foldCastFromReductions(Instruction &I);
  bool foldSignBitReductionCmp(Instruction &I);
  bool foldReductionZeroTest(Instruction &I);
  bool foldICmpEqZeroVectorReduce(Instruction &I);
  bool foldEquivalentReductionCmp(Instruction &I);
  bool foldReduceAddCmpZero(Instruction &I);
  bool foldSelectShuffle(Instruction &I, bool FromReduction = false);
  bool foldInterleaveIntrinsics(Instruction &I);
  bool foldDeinterleaveIntrinsics(Instruction &I);
  bool foldBitcastOfVPLoad(Instruction &I);
  bool foldBitOrderReverseAndSwap(Instruction &I);
  bool shrinkType(Instruction &I);
  bool shrinkLoadForShuffles(Instruction &I);
  bool shrinkPhiOfShuffles(Instruction &I);
  bool foldInterleaveOfDeinterleaveChains(Instruction &I);

  void replaceValue(Instruction &Old, Value &New, bool Erase = true) {
    LLVM_DEBUG(dbgs() << "VC: Replacing: " << Old << '\n');
    LLVM_DEBUG(dbgs() << "         With: " << New << '\n');
    Old.replaceAllUsesWith(&New);
    if (auto *NewI = dyn_cast<Instruction>(&New)) {
      New.takeName(&Old);
      Worklist.pushUsersToWorkList(*NewI);
      Worklist.pushValue(NewI);
    }
    if (Erase && isInstructionTriviallyDead(&Old)) {
      eraseInstruction(Old);
    } else {
      Worklist.push(&Old);
    }
  }

  void eraseInstruction(Instruction &I) {
    LLVM_DEBUG(dbgs() << "VC: Erasing: " << I << '\n');
    SmallVector<Value *> Ops(I.operands());
    Worklist.remove(&I);
    I.eraseFromParent();

    // Push remaining users of the operands and then the operand itself - allows
    // further folds that were hindered by OneUse limits.
    SmallPtrSet<Value *, 4> Visited;
    for (Value *Op : Ops) {
      if (!Visited.contains(Op)) {
        if (auto *OpI = dyn_cast<Instruction>(Op)) {
          if (RecursivelyDeleteTriviallyDeadInstructions(
                  OpI, nullptr, nullptr, [&](Value *V) {
                    if (auto *I = dyn_cast<Instruction>(V)) {
                      LLVM_DEBUG(dbgs() << "VC: Erased: " << *I << '\n');
                      Worklist.remove(I);
                      if (I == NextInst)
                        NextInst = NextInst->getNextNode();
                      Visited.insert(I);
                    }
                  }))
            continue;
          Worklist.pushUsersToWorkList(*OpI);
          Worklist.pushValue(OpI);
        }
      }
    }
  }
};
} // namespace

/// Return the source operand of a potentially bitcasted value. If there is no
/// bitcast, return the input value itself.
static Value *peekThroughBitcasts(Value *V) {
  while (auto *BitCast = dyn_cast<BitCastInst>(V))
    V = BitCast->getOperand(0);
  return V;
}

/// Helper to peek through bitcasts to the same value.
static bool isEquivBitcast(Value *X, Value *Y) {
  return X->getType() == Y->getType() &&
         peekThroughBitcasts(X) == peekThroughBitcasts(Y);
}

static bool canWidenLoad(LoadInst *Load, const TargetTransformInfo &TTI) {
  // Do not widen load if atomic/volatile or under asan/hwasan/memtag/tsan.
  // The widened load may load data from dirty regions or create data races
  // non-existent in the source.
  if (!Load || !Load->isSimple() || !Load->hasOneUse() ||
      Load->getFunction()->hasFnAttribute(Attribute::SanitizeMemTag) ||
      mustSuppressSpeculation(*Load))
    return false;

  // We are potentially transforming byte-sized (8-bit) memory accesses, so make
  // sure we have all of our type-based constraints in place for this target.
  Type *ScalarTy = Load->getType()->getScalarType();
  uint64_t ScalarSize = ScalarTy->getPrimitiveSizeInBits();
  unsigned MinVectorSize = TTI.getMinVectorRegisterBitWidth();
  if (!ScalarSize || !MinVectorSize || MinVectorSize % ScalarSize != 0 ||
      ScalarSize % 8 != 0)
    return false;

  return true;
}

bool VectorCombine::vectorizeLoadInsert(Instruction &I) {
  // Match insert into fixed vector of scalar value.
  // TODO: Handle non-zero insert index.
  Value *Scalar;
  if (!match(&I,
             m_InsertElt(m_Poison(), m_OneUse(m_Value(Scalar)), m_ZeroInt())))
    return false;

  // Optionally match an extract from another vector.
  Value *X;
  bool HasExtract = match(Scalar, m_ExtractElt(m_Value(X), m_ZeroInt()));
  if (!HasExtract)
    X = Scalar;

  auto *Load = dyn_cast<LoadInst>(X);
  if (!canWidenLoad(Load, TTI))
    return false;

  Type *ScalarTy = Scalar->getType();
  uint64_t ScalarSize = ScalarTy->getPrimitiveSizeInBits();
  unsigned MinVectorSize = TTI.getMinVectorRegisterBitWidth();

  // Check safety of replacing the scalar load with a larger vector load.
  // We use minimal alignment (maximum flexibility) because we only care about
  // the dereferenceable region. When calculating cost and creating a new op,
  // we may use a larger value based on alignment attributes.
  Value *SrcPtr = Load->getPointerOperand()->stripPointerCasts();
  assert(isa<PointerType>(SrcPtr->getType()) && "Expected a pointer type");

  unsigned MinVecNumElts = MinVectorSize / ScalarSize;
  auto *MinVecTy = VectorType::get(ScalarTy, MinVecNumElts, false);
  unsigned OffsetEltIndex = 0;
  Align Alignment = Load->getAlign();
  if (!isSafeToLoadUnconditionally(SrcPtr, MinVecTy, Align(1),
                                   SQ.getWithInstruction(Load))) {
    // It is not safe to load directly from the pointer, but we can still peek
    // through gep offsets and check if it safe to load from a base address with
    // updated alignment. If it is, we can shuffle the element(s) into place
    // after loading.
    unsigned OffsetBitWidth = DL->getIndexTypeSizeInBits(SrcPtr->getType());
    APInt Offset(OffsetBitWidth, 0);
    SrcPtr = SrcPtr->stripAndAccumulateInBoundsConstantOffsets(*DL, Offset);

    // We want to shuffle the result down from a high element of a vector, so
    // the offset must be positive.
    if (Offset.isNegative())
      return false;

    // The offset must be a multiple of the scalar element to shuffle cleanly
    // in the element's size.
    uint64_t ScalarSizeInBytes = ScalarSize / 8;
    if (Offset.urem(ScalarSizeInBytes) != 0)
      return false;

    // If we load MinVecNumElts, will our target element still be loaded?
    APInt OffsetEltIndexAP = Offset.udiv(ScalarSizeInBytes);
    if (OffsetEltIndexAP.uge(MinVecNumElts))
      return false;
    OffsetEltIndex = OffsetEltIndexAP.getZExtValue();

    if (!isSafeToLoadUnconditionally(SrcPtr, MinVecTy, Align(1),
                                     SQ.getWithInstruction(Load)))
      return false;

    // Update alignment with offset value. Note that the offset could be negated
    // to more accurately represent "(new) SrcPtr - Offset = (old) SrcPtr", but
    // negation does not change the result of the alignment calculation.
    Alignment = commonAlignment(Alignment, Offset.getZExtValue());
  }

  // Original pattern: insertelt undef, load [free casts of] PtrOp, 0
  // Use the greater of the alignment on the load or its source pointer.
  Alignment = std::max(SrcPtr->getPointerAlignment(*DL), Alignment);
  Type *LoadTy = Load->getType();
  unsigned AS = Load->getPointerAddressSpace();
  InstructionCost OldCost =
      TTI.getMemoryOpCost(Instruction::Load, LoadTy, Alignment, AS, CostKind);
  APInt DemandedElts = APInt::getOneBitSet(MinVecNumElts, 0);
  OldCost +=
      TTI.getScalarizationOverhead(MinVecTy, DemandedElts,
                                   /* Insert */ true, HasExtract, CostKind);

  // New pattern: load VecPtr
  InstructionCost NewCost =
      TTI.getMemoryOpCost(Instruction::Load, MinVecTy, Alignment, AS, CostKind);
  // Optionally, we are shuffling the loaded vector element(s) into place.
  // For the mask set everything but element 0 to undef to prevent poison from
  // propagating from the extra loaded memory. This will also optionally
  // shrink/grow the vector from the loaded size to the output size.
  // We assume this operation has no cost in codegen if there was no offset.
  // Note that we could use freeze to avoid poison problems, but then we might
  // still need a shuffle to change the vector size.
  auto *Ty = cast<FixedVectorType>(I.getType());
  unsigned OutputNumElts = Ty->getNumElements();
  SmallVector<int, 16> Mask(OutputNumElts, PoisonMaskElem);
  assert(OffsetEltIndex < MinVecNumElts && "Address offset too big");
  Mask[0] = OffsetEltIndex;
  if (OffsetEltIndex)
    NewCost += TTI.getShuffleCost(TTI::SK_PermuteSingleSrc, Ty, MinVecTy,
                                  CostKind, Mask);

  // We can aggressively convert to the vector form because the backend can
  // invert this transform if it does not result in a performance win.
  if (OldCost < NewCost || !NewCost.isValid())
    return false;

  // It is safe and potentially profitable to load a vector directly:
  // inselt undef, load Scalar, 0 --> load VecPtr
  IRBuilder<> Builder(Load);
  Value *CastedPtr =
      Builder.CreatePointerBitCastOrAddrSpaceCast(SrcPtr, Builder.getPtrTy(AS));
  Value *VecLd = Builder.CreateAlignedLoad(MinVecTy, CastedPtr, Alignment);
  VecLd = Builder.CreateShuffleVector(VecLd, Mask);

  replaceValue(I, *VecLd);
  ++NumVecLoad;
  return true;
}

/// If we are loading a vector and then inserting it into a larger vector with
/// undefined elements, try to load the larger vector and eliminate the insert.
/// This removes a shuffle in IR and may allow combining of other loaded values.
bool VectorCombine::widenSubvectorLoad(Instruction &I) {
  // Match subvector insert of fixed vector.
  auto *Shuf = cast<ShuffleVectorInst>(&I);
  if (!Shuf->isIdentityWithPadding())
    return false;

  // Allow a non-canonical shuffle mask that is choosing elements from op1.
  unsigned NumOpElts =
      cast<FixedVectorType>(Shuf->getOperand(0)->getType())->getNumElements();
  unsigned OpIndex = any_of(Shuf->getShuffleMask(), [&NumOpElts](int M) {
    return M >= (int)(NumOpElts);
  });

  auto *Load = dyn_cast<LoadInst>(Shuf->getOperand(OpIndex));
  if (!canWidenLoad(Load, TTI))
    return false;

  // We use minimal alignment (maximum flexibility) because we only care about
  // the dereferenceable region. When calculating cost and creating a new op,
  // we may use a larger value based on alignment attributes.
  auto *Ty = cast<FixedVectorType>(I.getType());
  Value *SrcPtr = Load->getPointerOperand()->stripPointerCasts();
  assert(isa<PointerType>(SrcPtr->getType()) && "Expected a pointer type");
  Align Alignment = Load->getAlign();
  if (!isSafeToLoadUnconditionally(SrcPtr, Ty, Align(1),
                                   SQ.getWithInstruction(Load)))
    return false;

  Alignment = std::max(SrcPtr->getPointerAlignment(*DL), Alignment);
  Type *LoadTy = Load->getType();
  unsigned AS = Load->getPointerAddressSpace();

  // Original pattern: insert_subvector (load PtrOp)
  // This conservatively assumes that the cost of a subvector insert into an
  // undef value is 0. We could add that cost if the cost model accurately
  // reflects the real cost of that operation.
  InstructionCost OldCost =
      TTI.getMemoryOpCost(Instruction::Load, LoadTy, Alignment, AS, CostKind);

  // New pattern: load PtrOp
  InstructionCost NewCost =
      TTI.getMemoryOpCost(Instruction::Load, Ty, Alignment, AS, CostKind);

  // We can aggressively convert to the vector form because the backend can
  // invert this transform if it does not result in a performance win.
  if (OldCost < NewCost || !NewCost.isValid())
    return false;

  IRBuilder<> Builder(Load);
  Value *CastedPtr =
      Builder.CreatePointerBitCastOrAddrSpaceCast(SrcPtr, Builder.getPtrTy(AS));
  Value *VecLd = Builder.CreateAlignedLoad(Ty, CastedPtr, Alignment);
  replaceValue(I, *VecLd);
  ++NumVecLoad;
  return true;
}

/// Determine which, if any, of the inputs should be replaced by a shuffle
/// followed by extract from a different index.
ExtractElementInst *VectorCombine::getShuffleExtract(
    ExtractElementInst *Ext0, ExtractElementInst *Ext1,
    unsigned PreferredExtractIndex = InvalidIndex) const {
  auto *Index0C = dyn_cast<ConstantInt>(Ext0->getIndexOperand());
  auto *Index1C = dyn_cast<ConstantInt>(Ext1->getIndexOperand());
  assert(Index0C && Index1C && "Expected constant extract indexes");

  unsigned Index0 = Index0C->getZExtValue();
  unsigned Index1 = Index1C->getZExtValue();

  // If the extract indexes are identical, no shuffle is needed.
  if (Index0 == Index1)
    return nullptr;

  Type *VecTy = Ext0->getVectorOperand()->getType();
  assert(VecTy == Ext1->getVectorOperand()->getType() && "Need matching types");
  InstructionCost Cost0 =
      TTI.getVectorInstrCost(*Ext0, VecTy, CostKind, Index0);
  InstructionCost Cost1 =
      TTI.getVectorInstrCost(*Ext1, VecTy, CostKind, Index1);

  // If both costs are invalid no shuffle is needed
  if (!Cost0.isValid() && !Cost1.isValid())
    return nullptr;

  // We are extracting from 2 different indexes, so one operand must be shuffled
  // before performing a vector operation and/or extract. The more expensive
  // extract will be replaced by a shuffle.
  if (Cost0 > Cost1)
    return Ext0;
  if (Cost1 > Cost0)
    return Ext1;

  // If the costs are equal and there is a preferred extract index, shuffle the
  // opposite operand.
  if (PreferredExtractIndex == Index0)
    return Ext1;
  if (PreferredExtractIndex == Index1)
    return Ext0;

  // Otherwise, replace the extract with the higher index.
  return Index0 > Index1 ? Ext0 : Ext1;
}

/// Compare the relative costs of 2 extracts followed by scalar operation vs.
/// vector operation(s) followed by extract. Return true if the existing
/// instructions are cheaper than a vector alternative. Otherwise, return false
/// and if one of the extracts should be transformed to a shufflevector, set
/// \p ConvertToShuffle to that extract instruction.
bool VectorCombine::isExtractExtractCheap(ExtractElementInst *Ext0,
                                          ExtractElementInst *Ext1,
                                          const Instruction &I,
                                          ExtractElementInst *&ConvertToShuffle,
                                          unsigned PreferredExtractIndex) {
  auto *Ext0IndexC = dyn_cast<ConstantInt>(Ext0->getIndexOperand());
  auto *Ext1IndexC = dyn_cast<ConstantInt>(Ext1->getIndexOperand());
  assert(Ext0IndexC && Ext1IndexC && "Expected constant extract indexes");

  unsigned Opcode = I.getOpcode();
  Value *Ext0Src = Ext0->getVectorOperand();
  Value *Ext1Src = Ext1->getVectorOperand();
  Type *ScalarTy = Ext0->getType();
  auto *VecTy = cast<VectorType>(Ext0Src->getType());
  InstructionCost ScalarOpCost, VectorOpCost;

  // Get cost estimates for scalar and vector versions of the operation.
  bool IsBinOp = Instruction::isBinaryOp(Opcode);
  if (IsBinOp) {
    ScalarOpCost = TTI.getArithmeticInstrCost(Opcode, ScalarTy, CostKind);
    VectorOpCost = TTI.getArithmeticInstrCost(Opcode, VecTy, CostKind);
  } else {
    assert((Opcode == Instruction::ICmp || Opcode == Instruction::FCmp) &&
           "Expected a compare");
    CmpInst::Predicate Pred = cast<CmpInst>(I).getPredicate();
    ScalarOpCost = TTI.getCmpSelInstrCost(
        Opcode, ScalarTy, CmpInst::makeCmpResultType(ScalarTy), Pred, CostKind);
    VectorOpCost = TTI.getCmpSelInstrCost(
        Opcode, VecTy, CmpInst::makeCmpResultType(VecTy), Pred, CostKind);
  }

  // Get cost estimates for the extract elements. These costs will factor into
  // both sequences.
  unsigned Ext0Index = Ext0IndexC->getZExtValue();
  unsigned Ext1Index = Ext1IndexC->getZExtValue();

  InstructionCost Extract0Cost =
      TTI.getVectorInstrCost(*Ext0, VecTy, CostKind, Ext0Index);
  InstructionCost Extract1Cost =
      TTI.getVectorInstrCost(*Ext1, VecTy, CostKind, Ext1Index);

  // A more expensive extract will always be replaced by a splat shuffle.
  // For example, if Ext0 is more expensive:
  // opcode (extelt V0, Ext0), (ext V1, Ext1) -->
  // extelt (opcode (splat V0, Ext0), V1), Ext1
  // TODO: Evaluate whether that always results in lowest cost. Alternatively,
  //       check the cost of creating a broadcast shuffle and shuffling both
  //       operands to element 0.
  unsigned BestExtIndex = Extract0Cost > Extract1Cost ? Ext0Index : Ext1Index;
  unsigned BestInsIndex = Extract0Cost > Extract1Cost ? Ext1Index : Ext0Index;
  InstructionCost CheapExtractCost = std::min(Extract0Cost, Extract1Cost);

  // Extra uses of the extracts mean that we include those costs in the
  // vector total because those instructions will not be eliminated.
  InstructionCost OldCost, NewCost;
  if (Ext0Src == Ext1Src && Ext0Index == Ext1Index) {
    // Handle a special case. If the 2 extracts are identical, adjust the
    // formulas to account for that. The extra use charge allows for either the
    // CSE'd pattern or an unoptimized form with identical values:
    // opcode (extelt V, C), (extelt V, C) --> extelt (opcode V, V), C
    bool HasUseTax = Ext0 == Ext1 ? !Ext0->hasNUses(2)
                                  : !Ext0->hasOneUse() || !Ext1->hasOneUse();
    OldCost = CheapExtractCost + ScalarOpCost;
    NewCost = VectorOpCost + CheapExtractCost + HasUseTax * CheapExtractCost;
  } else {
    // Handle the general case. Each extract is actually a different value:
    // opcode (extelt V0, C0), (extelt V1, C1) --> extelt (opcode V0, V1), C
    OldCost = Extract0Cost + Extract1Cost + ScalarOpCost;
    NewCost = VectorOpCost + CheapExtractCost +
              !Ext0->hasOneUse() * Extract0Cost +
              !Ext1->hasOneUse() * Extract1Cost;
  }

  ConvertToShuffle = getShuffleExtract(Ext0, Ext1, PreferredExtractIndex);
  if (ConvertToShuffle) {
    if (IsBinOp && DisableBinopExtractShuffle)
      return true;

    // If we are extracting from 2 different indexes, then one operand must be
    // shuffled before performing the vector operation. The shuffle mask is
    // poison except for 1 lane that is being translated to the remaining
    // extraction lane. Therefore, it is a splat shuffle. Ex:
    // ShufMask = { poison, poison, 0, poison }
    // TODO: The cost model has an option for a "broadcast" shuffle
    //       (splat-from-element-0), but no option for a more general splat.
    if (auto *FixedVecTy = dyn_cast<FixedVectorType>(VecTy)) {
      SmallVector<int> ShuffleMask(FixedVecTy->getNumElements(),
                                   PoisonMaskElem);
      ShuffleMask[BestInsIndex] = BestExtIndex;
      NewCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
                                    VecTy, VecTy, CostKind, ShuffleMask, 0,
                                    nullptr, {ConvertToShuffle});
    } else {
      NewCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
                                    VecTy, VecTy, CostKind, {}, 0, nullptr,
                                    {ConvertToShuffle});
    }
  }

  LLVM_DEBUG(dbgs() << "Found a binop of extractions: " << I << "\n  OldCost: "
                    << OldCost << " vs NewCost: " << NewCost << "\n");

  // Aggressively form a vector op if the cost is equal because the transform
  // may enable further optimization.
  // Codegen can reverse this transform (scalarize) if it was not profitable.
  return OldCost < NewCost;
}

/// Create a shuffle that translates (shifts) 1 element from the input vector
/// to a new element location.
static Value *createShiftShuffle(Value *Vec, unsigned OldIndex,
                                 unsigned NewIndex, IRBuilderBase &Builder) {
  // The shuffle mask is poison except for 1 lane that is being translated
  // to the new element index. Example for OldIndex == 2 and NewIndex == 0:
  // ShufMask = { 2, poison, poison, poison }
  auto *VecTy = cast<FixedVectorType>(Vec->getType());
  SmallVector<int, 32> ShufMask(VecTy->getNumElements(), PoisonMaskElem);
  ShufMask[NewIndex] = OldIndex;
  return Builder.CreateShuffleVector(Vec, ShufMask, "shift");
}

/// Given an extract element instruction with constant index operand, shuffle
/// the source vector (shift the scalar element) to a NewIndex for extraction.
/// Return null if the input can be constant folded, so that we are not creating
/// unnecessary instructions.
static Value *translateExtract(ExtractElementInst *ExtElt, unsigned NewIndex,
                               IRBuilderBase &Builder) {
  // Shufflevectors can only be created for fixed-width vectors.
  Value *X = ExtElt->getVectorOperand();
  if (!isa<FixedVectorType>(X->getType()))
    return nullptr;

  // If the extract can be constant-folded, this code is unsimplified. Defer
  // to other passes to handle that.
  Value *C = ExtElt->getIndexOperand();
  assert(isa<ConstantInt>(C) && "Expected a constant index operand");
  if (isa<Constant>(X))
    return nullptr;

  Value *Shuf = createShiftShuffle(X, cast<ConstantInt>(C)->getZExtValue(),
                                   NewIndex, Builder);
  return Shuf;
}

/// Try to reduce extract element costs by converting scalar compares to vector
/// compares followed by extract.
/// cmp (ext0 V0, ExtIndex), (ext1 V1, ExtIndex)
Value *VectorCombine::foldExtExtCmp(Value *V0, Value *V1, Value *ExtIndex,
                                    Instruction &I) {
  assert(isa<CmpInst>(&I) && "Expected a compare");

  // cmp Pred (extelt V0, ExtIndex), (extelt V1, ExtIndex)
  //   --> extelt (cmp Pred V0, V1), ExtIndex
  ++NumVecCmp;
  CmpInst::Predicate Pred = cast<CmpInst>(&I)->getPredicate();
  Value *VecCmp = Builder.CreateCmp(Pred, V0, V1);
  return Builder.CreateExtractElement(VecCmp, ExtIndex, "foldExtExtCmp");
}

/// Try to reduce extract element costs by converting scalar binops to vector
/// binops followed by extract.
/// bo (ext0 V0, ExtIndex), (ext1 V1, ExtIndex)
Value *VectorCombine::foldExtExtBinop(Value *V0, Value *V1, Value *ExtIndex,
                                      Instruction &I) {
  assert(isa<BinaryOperator>(&I) && "Expected a binary operator");

  // bo (extelt V0, ExtIndex), (extelt V1, ExtIndex)
  //   --> extelt (bo V0, V1), ExtIndex
  ++NumVecBO;
  Value *VecBO = Builder.CreateBinOp(cast<BinaryOperator>(&I)->getOpcode(), V0,
                                     V1, "foldExtExtBinop");

  // All IR flags are safe to back-propagate because any potential poison
  // created in unused vector elements is discarded by the extract.
  if (auto *VecBOInst = dyn_cast<Instruction>(VecBO))
    VecBOInst->copyIRFlags(&I);

  return Builder.CreateExtractElement(VecBO, ExtIndex, "foldExtExtBinop");
}

/// Match an instruction with extracted vector operands.
bool VectorCombine::foldExtractExtract(Instruction &I) {
  // It is not safe to transform things like div, urem, etc. because we may
  // create undefined behavior when executing those on unknown vector elements.
  if (!isSafeToSpeculativelyExecute(&I))
    return false;

  Instruction *I0, *I1;
  CmpPredicate Pred = CmpInst::BAD_ICMP_PREDICATE;
  if (!match(&I, m_Cmp(Pred, m_Instruction(I0), m_Instruction(I1))) &&
      !match(&I, m_BinOp(m_Instruction(I0), m_Instruction(I1))))
    return false;

  Value *V0, *V1;
  uint64_t C0, C1;
  if (!match(I0, m_ExtractElt(m_Value(V0), m_ConstantInt(C0))) ||
      !match(I1, m_ExtractElt(m_Value(V1), m_ConstantInt(C1))) ||
      V0->getType() != V1->getType())
    return false;

  // For fixed-width vectors, reject out-of-bounds extract indexes
  if (auto *FixedVecTy = dyn_cast<FixedVectorType>(V0->getType())) {
    unsigned NumElts = FixedVecTy->getNumElements();
    if (C0 >= NumElts || C1 >= NumElts)
      return false;
  }

  // If the scalar value 'I' is going to be re-inserted into a vector, then try
  // to create an extract to that same element. The extract/insert can be
  // reduced to a "select shuffle".
  // TODO: If we add a larger pattern match that starts from an insert, this
  //       probably becomes unnecessary.
  auto *Ext0 = cast<ExtractElementInst>(I0);
  auto *Ext1 = cast<ExtractElementInst>(I1);
  uint64_t InsertIndex = InvalidIndex;
  if (I.hasOneUse())
    match(I.user_back(),
          m_InsertElt(m_Value(), m_Value(), m_ConstantInt(InsertIndex)));

  ExtractElementInst *ExtractToChange;
  if (isExtractExtractCheap(Ext0, Ext1, I, ExtractToChange, InsertIndex))
    return false;

  Value *ExtOp0 = Ext0->getVectorOperand();
  Value *ExtOp1 = Ext1->getVectorOperand();

  if (ExtractToChange) {
    unsigned CheapExtractIdx = ExtractToChange == Ext0 ? C1 : C0;
    Value *NewExtOp =
        translateExtract(ExtractToChange, CheapExtractIdx, Builder);
    if (!NewExtOp)
      return false;
    if (ExtractToChange == Ext0)
      ExtOp0 = NewExtOp;
    else
      ExtOp1 = NewExtOp;
  }

  Value *ExtIndex = ExtractToChange == Ext0 ? Ext1->getIndexOperand()
                                            : Ext0->getIndexOperand();
  Value *NewExt = Pred != CmpInst::BAD_ICMP_PREDICATE
                      ? foldExtExtCmp(ExtOp0, ExtOp1, ExtIndex, I)
                      : foldExtExtBinop(ExtOp0, ExtOp1, ExtIndex, I);
  Worklist.push(Ext0);
  Worklist.push(Ext1);
  replaceValue(I, *NewExt);
  return true;
}

/// Try to replace an extract + scalar fneg + insert with a vector fneg +
/// shuffle.
bool VectorCombine::foldInsExtFNeg(Instruction &I) {
  // Match an insert (op (extract)) pattern.
  Value *DstVec;
  uint64_t ExtIdx, InsIdx;
  Instruction *FNeg;
  if (!match(&I, m_InsertElt(m_Value(DstVec), m_OneUse(m_Instruction(FNeg)),
                             m_ConstantInt(InsIdx))))
    return false;

  // Note: This handles the canonical fneg instruction and "fsub -0.0, X".
  Value *SrcVec;
  Instruction *Extract;
  if (!match(FNeg, m_FNeg(m_CombineAnd(
                       m_Instruction(Extract),
                       m_ExtractElt(m_Value(SrcVec), m_ConstantInt(ExtIdx))))))
    return false;

  auto *DstVecTy = cast<FixedVectorType>(DstVec->getType());
  auto *DstVecScalarTy = DstVecTy->getScalarType();
  auto *SrcVecTy = dyn_cast<FixedVectorType>(SrcVec->getType());
  if (!SrcVecTy || DstVecScalarTy != SrcVecTy->getScalarType())
    return false;

  // Ignore if insert/extract index is out of bounds or destination vector has
  // one element
  unsigned NumDstElts = DstVecTy->getNumElements();
  unsigned NumSrcElts = SrcVecTy->getNumElements();
  if (ExtIdx > NumSrcElts || InsIdx >= NumDstElts || NumDstElts == 1)
    return false;

  // We are inserting the negated element into the same lane that we extracted
  // from. This is equivalent to a select-shuffle that chooses all but the
  // negated element from the destination vector.
  SmallVector<int> Mask(NumDstElts);
  std::iota(Mask.begin(), Mask.end(), 0);
  Mask[InsIdx] = (ExtIdx % NumDstElts) + NumDstElts;
  InstructionCost OldCost =
      TTI.getArithmeticInstrCost(Instruction::FNeg, DstVecScalarTy, CostKind) +
      TTI.getVectorInstrCost(I, DstVecTy, CostKind, InsIdx);

  // If the extract has one use, it will be eliminated, so count it in the
  // original cost. If it has more than one use, ignore the cost because it will
  // be the same before/after.
  if (Extract->hasOneUse())
    OldCost += TTI.getVectorInstrCost(*Extract, SrcVecTy, CostKind, ExtIdx);

  InstructionCost NewCost =
      TTI.getArithmeticInstrCost(Instruction::FNeg, SrcVecTy, CostKind) +
      TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, DstVecTy,
                         DstVecTy, CostKind, Mask);

  bool NeedLenChg = SrcVecTy->getNumElements() != NumDstElts;
  // If the lengths of the two vectors are not equal,
  // we need to add a length-change vector. Add this cost.
  SmallVector<int> SrcMask;
  if (NeedLenChg) {
    SrcMask.assign(NumDstElts, PoisonMaskElem);
    SrcMask[ExtIdx % NumDstElts] = ExtIdx;
    NewCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
                                  DstVecTy, SrcVecTy, CostKind, SrcMask);
  }

  LLVM_DEBUG(dbgs() << "Found an insertion of (extract)fneg : " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");
  if (NewCost > OldCost)
    return false;

  Value *NewShuf, *LenChgShuf = nullptr;
  // insertelt DstVec, (fneg (extractelt SrcVec, Index)), Index
  Value *VecFNeg = Builder.CreateFNegFMF(SrcVec, FNeg);
  if (NeedLenChg) {
    // shuffle DstVec, (shuffle (fneg SrcVec), poison, SrcMask), Mask
    LenChgShuf = Builder.CreateShuffleVector(VecFNeg, SrcMask);
    NewShuf = Builder.CreateShuffleVector(DstVec, LenChgShuf, Mask);
    Worklist.pushValue(LenChgShuf);
  } else {
    // shuffle DstVec, (fneg SrcVec), Mask
    NewShuf = Builder.CreateShuffleVector(DstVec, VecFNeg, Mask);
  }

  Worklist.pushValue(VecFNeg);
  replaceValue(I, *NewShuf);
  return true;
}

/// Try to fold insert(binop(x,y),binop(a,b),idx)
///         --> binop(insert(x,a,idx),insert(y,b,idx))
bool VectorCombine::foldInsExtBinop(Instruction &I) {
  BinaryOperator *VecBinOp, *SclBinOp;
  uint64_t Index;
  if (!match(&I,
             m_InsertElt(m_OneUse(m_BinOp(VecBinOp)),
                         m_OneUse(m_BinOp(SclBinOp)), m_ConstantInt(Index))))
    return false;

  // TODO: Add support for addlike etc.
  Instruction::BinaryOps BinOpcode = VecBinOp->getOpcode();
  if (BinOpcode != SclBinOp->getOpcode())
    return false;

  auto *ResultTy = dyn_cast<FixedVectorType>(I.getType());
  if (!ResultTy)
    return false;

  // TODO: Attempt to detect m_ExtractElt for scalar operands and convert to
  // shuffle?

  InstructionCost OldCost = TTI.getInstructionCost(&I, CostKind) +
                            TTI.getInstructionCost(VecBinOp, CostKind) +
                            TTI.getInstructionCost(SclBinOp, CostKind);
  InstructionCost NewCost =
      TTI.getArithmeticInstrCost(BinOpcode, ResultTy, CostKind) +
      TTI.getVectorInstrCost(Instruction::InsertElement, ResultTy, CostKind,
                             Index, VecBinOp->getOperand(0),
                             SclBinOp->getOperand(0)) +
      TTI.getVectorInstrCost(Instruction::InsertElement, ResultTy, CostKind,
                             Index, VecBinOp->getOperand(1),
                             SclBinOp->getOperand(1));

  LLVM_DEBUG(dbgs() << "Found an insertion of two binops: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");
  if (NewCost > OldCost)
    return false;

  Value *NewIns0 = Builder.CreateInsertElement(VecBinOp->getOperand(0),
                                               SclBinOp->getOperand(0), Index);
  Value *NewIns1 = Builder.CreateInsertElement(VecBinOp->getOperand(1),
                                               SclBinOp->getOperand(1), Index);
  Value *NewBO = Builder.CreateBinOp(BinOpcode, NewIns0, NewIns1);

  // Intersect flags from the old binops.
  if (auto *NewInst = dyn_cast<Instruction>(NewBO)) {
    NewInst->copyIRFlags(VecBinOp);
    NewInst->andIRFlags(SclBinOp);
  }

  Worklist.pushValue(NewIns0);
  Worklist.pushValue(NewIns1);
  replaceValue(I, *NewBO);
  return true;
}

/// Match: bitop(castop(x), castop(y)) -> castop(bitop(x, y))
/// Supports: bitcast, trunc, sext, zext
bool VectorCombine::foldBitOpOfCastops(Instruction &I) {
  // Check if this is a bitwise logic operation
  auto *BinOp = dyn_cast<BinaryOperator>(&I);
  if (!BinOp || !BinOp->isBitwiseLogicOp())
    return false;

  // Get the cast instructions
  auto *LHSCast = dyn_cast<CastInst>(BinOp->getOperand(0));
  auto *RHSCast = dyn_cast<CastInst>(BinOp->getOperand(1));
  if (!LHSCast || !RHSCast) {
    LLVM_DEBUG(dbgs() << "  One or both operands are not cast instructions\n");
    return false;
  }

  // Both casts must be the same type
  Instruction::CastOps CastOpcode = LHSCast->getOpcode();
  if (CastOpcode != RHSCast->getOpcode())
    return false;

  // Only handle supported cast operations
  switch (CastOpcode) {
  case Instruction::BitCast:
  case Instruction::Trunc:
  case Instruction::SExt:
  case Instruction::ZExt:
    break;
  default:
    return false;
  }

  Value *LHSSrc = LHSCast->getOperand(0);
  Value *RHSSrc = RHSCast->getOperand(0);

  // Source types must match
  if (LHSSrc->getType() != RHSSrc->getType())
    return false;

  auto *SrcTy = LHSSrc->getType();
  auto *DstTy = I.getType();
  // Bitcasts can handle scalar/vector mixes, such as i16 -> <16 x i1>.
  // Other casts only handle vector types with integer elements.
  if (CastOpcode != Instruction::BitCast &&
      (!isa<FixedVectorType>(SrcTy) || !isa<FixedVectorType>(DstTy)))
    return false;

  // Only integer scalar/vector values are legal for bitwise logic operations.
  if (!SrcTy->getScalarType()->isIntegerTy() ||
      !DstTy->getScalarType()->isIntegerTy())
    return false;

  // Cost Check :
  // OldCost = bitlogic + 2*casts
  // NewCost = bitlogic + cast

  // Calculate specific costs for each cast with instruction context
  InstructionCost LHSCastCost = TTI.getCastInstrCost(
      CastOpcode, DstTy, SrcTy, TTI::CastContextHint::None, CostKind, LHSCast);
  InstructionCost RHSCastCost = TTI.getCastInstrCost(
      CastOpcode, DstTy, SrcTy, TTI::CastContextHint::None, CostKind, RHSCast);

  InstructionCost OldCost =
      TTI.getArithmeticInstrCost(BinOp->getOpcode(), DstTy, CostKind) +
      LHSCastCost + RHSCastCost;

  // For new cost, we can't provide an instruction (it doesn't exist yet)
  InstructionCost GenericCastCost = TTI.getCastInstrCost(
      CastOpcode, DstTy, SrcTy, TTI::CastContextHint::None, CostKind);

  InstructionCost NewCost =
      TTI.getArithmeticInstrCost(BinOp->getOpcode(), SrcTy, CostKind) +
      GenericCastCost;

  // Account for multi-use casts using specific costs
  if (!LHSCast->hasOneUse())
    NewCost += LHSCastCost;
  if (!RHSCast->hasOneUse())
    NewCost += RHSCastCost;

  LLVM_DEBUG(dbgs() << "foldBitOpOfCastops: OldCost=" << OldCost
                    << " NewCost=" << NewCost << "\n");

  if (NewCost > OldCost)
    return false;

  // Create the operation on the source type
  Value *NewOp = Builder.CreateBinOp(BinOp->getOpcode(), LHSSrc, RHSSrc,
                                     BinOp->getName() + ".inner");
  if (auto *NewBinOp = dyn_cast<BinaryOperator>(NewOp))
    NewBinOp->copyIRFlags(BinOp);

  Worklist.pushValue(NewOp);

  // Create the cast operation directly to ensure we get a new instruction
  Instruction *NewCast = CastInst::Create(CastOpcode, NewOp, I.getType());

  // Preserve cast instruction flags
  NewCast->copyIRFlags(LHSCast);
  NewCast->andIRFlags(RHSCast);

  // Insert the new instruction
  Value *Result = Builder.Insert(NewCast);

  replaceValue(I, *Result);
  return true;
}

/// Match:
// bitop(castop(x), C) ->
// bitop(castop(x), castop(InvC)) ->
// castop(bitop(x, InvC))
// Supports: bitcast
bool VectorCombine::foldBitOpOfCastConstant(Instruction &I) {
  Instruction *LHS;
  Constant *C;

  // Check if this is a bitwise logic operation
  if (!match(&I, m_c_BitwiseLogic(m_Instruction(LHS), m_Constant(C))))
    return false;

  // Get the cast instructions
  auto *LHSCast = dyn_cast<CastInst>(LHS);
  if (!LHSCast)
    return false;

  Instruction::CastOps CastOpcode = LHSCast->getOpcode();

  // Only handle supported cast operations
  switch (CastOpcode) {
  case Instruction::BitCast:
  case Instruction::ZExt:
  case Instruction::SExt:
  case Instruction::Trunc:
    break;
  default:
    return false;
  }

  Value *LHSSrc = LHSCast->getOperand(0);

  auto *SrcTy = LHSSrc->getType();
  auto *DstTy = I.getType();
  // Bitcasts can handle scalar/vector mixes, such as i16 -> <16 x i1>.
  // Other casts only handle vector types with integer elements.
  if (CastOpcode != Instruction::BitCast &&
      (!isa<FixedVectorType>(SrcTy) || !isa<FixedVectorType>(DstTy)))
    return false;

  // Only integer scalar/vector values are legal for bitwise logic operations.
  if (!SrcTy->getScalarType()->isIntegerTy() ||
      !DstTy->getScalarType()->isIntegerTy())
    return false;

  // Find the constant InvC, such that castop(InvC) equals to C.
  PreservedCastFlags RHSFlags;
  Constant *InvC = getLosslessInvCast(C, SrcTy, CastOpcode, *DL, &RHSFlags);
  if (!InvC)
    return false;

  // Cost Check :
  // OldCost = bitlogic + cast
  // NewCost = bitlogic + cast

  // Calculate specific costs for each cast with instruction context
  InstructionCost LHSCastCost = TTI.getCastInstrCost(
      CastOpcode, DstTy, SrcTy, TTI::CastContextHint::None, CostKind, LHSCast);

  InstructionCost OldCost =
      TTI.getArithmeticInstrCost(I.getOpcode(), DstTy, CostKind) + LHSCastCost;

  // For new cost, we can't provide an instruction (it doesn't exist yet)
  InstructionCost GenericCastCost = TTI.getCastInstrCost(
      CastOpcode, DstTy, SrcTy, TTI::CastContextHint::None, CostKind);

  InstructionCost NewCost =
      TTI.getArithmeticInstrCost(I.getOpcode(), SrcTy, CostKind) +
      GenericCastCost;

  // Account for multi-use casts using specific costs
  if (!LHSCast->hasOneUse())
    NewCost += LHSCastCost;

  LLVM_DEBUG(dbgs() << "foldBitOpOfCastConstant: OldCost=" << OldCost
                    << " NewCost=" << NewCost << "\n");

  if (NewCost > OldCost)
    return false;

  // Create the operation on the source type
  Value *NewOp = Builder.CreateBinOp((Instruction::BinaryOps)I.getOpcode(),
                                     LHSSrc, InvC, I.getName() + ".inner");
  if (auto *NewBinOp = dyn_cast<BinaryOperator>(NewOp))
    NewBinOp->copyIRFlags(&I);

  Worklist.pushValue(NewOp);

  // Create the cast operation directly to ensure we get a new instruction
  Instruction *NewCast = CastInst::Create(CastOpcode, NewOp, I.getType());

  // Preserve cast instruction flags
  if (RHSFlags.NNeg)
    NewCast->setNonNeg();
  if (RHSFlags.NUW)
    NewCast->setHasNoUnsignedWrap();
  if (RHSFlags.NSW)
    NewCast->setHasNoSignedWrap();

  NewCast->andIRFlags(LHSCast);

  // Insert the new instruction
  Value *Result = Builder.Insert(NewCast);

  replaceValue(I, *Result);
  return true;
}

/// If this is a bitcast of a shuffle, try to bitcast the source vector to the
/// destination type followed by shuffle. This can enable further transforms by
/// moving bitcasts or shuffles together.
bool VectorCombine::foldBitcastShuffle(Instruction &I) {
  Value *V0, *V1;
  ArrayRef<int> Mask;
  if (!match(&I, m_BitCast(m_OneUse(
                     m_Shuffle(m_Value(V0), m_Value(V1), m_Mask(Mask))))))
    return false;

  // 1) Do not fold bitcast shuffle for scalable type. First, shuffle cost for
  // scalable type is unknown; Second, we cannot reason if the narrowed shuffle
  // mask for scalable type is a splat or not.
  // 2) Disallow non-vector casts.
  // TODO: We could allow any shuffle.
  auto *DestTy = dyn_cast<FixedVectorType>(I.getType());
  auto *SrcTy = dyn_cast<FixedVectorType>(V0->getType());
  if (!DestTy || !SrcTy)
    return false;

  unsigned DestEltSize = DestTy->getScalarSizeInBits();
  unsigned SrcEltSize = SrcTy->getScalarSizeInBits();
  if (SrcTy->getPrimitiveSizeInBits() % DestEltSize != 0)
    return false;

  bool IsUnary = isa<UndefValue>(V1);

  // For binary shuffles, only fold bitcast(shuffle(X,Y))
  // if it won't increase the number of bitcasts.
  if (!IsUnary) {
    auto *BCTy0 = dyn_cast<FixedVectorType>(peekThroughBitcasts(V0)->getType());
    auto *BCTy1 = dyn_cast<FixedVectorType>(peekThroughBitcasts(V1)->getType());
    if (!(BCTy0 && BCTy0->getElementType() == DestTy->getElementType()) &&
        !(BCTy1 && BCTy1->getElementType() == DestTy->getElementType()))
      return false;
  }

  SmallVector<int, 16> NewMask;
  if (DestEltSize <= SrcEltSize) {
    // The bitcast is from wide to narrow/equal elements. The shuffle mask can
    // always be expanded to the equivalent form choosing narrower elements.
    if (SrcEltSize % DestEltSize != 0)
      return false;
    unsigned ScaleFactor = SrcEltSize / DestEltSize;
    narrowShuffleMaskElts(ScaleFactor, Mask, NewMask);
  } else {
    // The bitcast is from narrow elements to wide elements. The shuffle mask
    // must choose consecutive elements to allow casting first.
    if (DestEltSize % SrcEltSize != 0)
      return false;
    unsigned ScaleFactor = DestEltSize / SrcEltSize;
    if (!widenShuffleMaskElts(ScaleFactor, Mask, NewMask))
      return false;
  }

  // Bitcast the shuffle src - keep its original width but using the destination
  // scalar type.
  unsigned NumSrcElts = SrcTy->getPrimitiveSizeInBits() / DestEltSize;
  auto *NewShuffleTy =
      FixedVectorType::get(DestTy->getScalarType(), NumSrcElts);
  auto *OldShuffleTy =
      FixedVectorType::get(SrcTy->getScalarType(), Mask.size());
  unsigned NumOps = IsUnary ? 1 : 2;

  // The new shuffle must not cost more than the old shuffle.
  TargetTransformInfo::ShuffleKind SK =
      IsUnary ? TargetTransformInfo::SK_PermuteSingleSrc
              : TargetTransformInfo::SK_PermuteTwoSrc;

  InstructionCost NewCost =
      TTI.getShuffleCost(SK, DestTy, NewShuffleTy, CostKind, NewMask) +
      (NumOps * TTI.getCastInstrCost(Instruction::BitCast, NewShuffleTy, SrcTy,
                                     TargetTransformInfo::CastContextHint::None,
                                     CostKind));
  InstructionCost OldCost =
      TTI.getShuffleCost(SK, OldShuffleTy, SrcTy, CostKind, Mask) +
      TTI.getCastInstrCost(Instruction::BitCast, DestTy, OldShuffleTy,
                           TargetTransformInfo::CastContextHint::None,
                           CostKind);

  LLVM_DEBUG(dbgs() << "Found a bitcasted shuffle: " << I << "\n  OldCost: "
                    << OldCost << " vs NewCost: " << NewCost << "\n");

  if (NewCost > OldCost || !NewCost.isValid())
    return false;

  // bitcast (shuf V0, V1, MaskC) --> shuf (bitcast V0), (bitcast V1), MaskC'
  ++NumShufOfBitcast;
  Value *CastV0 = Builder.CreateBitCast(peekThroughBitcasts(V0), NewShuffleTy);
  Value *CastV1 = Builder.CreateBitCast(peekThroughBitcasts(V1), NewShuffleTy);
  Value *Shuf = Builder.CreateShuffleVector(CastV0, CastV1, NewMask);
  replaceValue(I, *Shuf);
  return true;
}

/// Match a vector op/compare/intrinsic with at least one
/// inserted scalar operand and convert to scalar op/cmp/intrinsic followed
/// by insertelement.
bool VectorCombine::scalarizeOpOrCmp(Instruction &I) {
  auto *UO = dyn_cast<UnaryOperator>(&I);
  auto *BO = dyn_cast<BinaryOperator>(&I);
  auto *CI = dyn_cast<CmpInst>(&I);
  auto *II = dyn_cast<IntrinsicInst>(&I);
  if (!UO && !BO && !CI && !II)
    return false;

  // TODO: Allow intrinsics with different argument types
  if (II) {
    if (!isTriviallyVectorizable(II->getIntrinsicID()))
      return false;
    for (auto [Idx, Arg] : enumerate(II->args()))
      if (Arg->getType() != II->getType() &&
          !isVectorIntrinsicWithScalarOpAtArg(II->getIntrinsicID(), Idx, &TTI))
        return false;
  }

  // Do not convert the vector condition of a vector select into a scalar
  // condition. That may cause problems for codegen because of differences in
  // boolean formats and register-file transfers.
  // TODO: Can we account for that in the cost model?
  if (CI)
    for (User *U : I.users())
      if (match(U, m_Select(m_Specific(&I), m_Value(), m_Value())))
        return false;

  // Match constant vectors or scalars being inserted into constant vectors:
  // vec_op [VecC0 | (inselt VecC0, V0, Index)], ...
  SmallVector<Value *> VecCs, ScalarOps;
  std::optional<uint64_t> Index;

  auto Ops = II ? II->args() : I.operands();
  for (auto [OpNum, Op] : enumerate(Ops)) {
    Constant *VecC;
    Value *V;
    uint64_t InsIdx = 0;
    if (match(Op.get(), m_InsertElt(m_Constant(VecC), m_Value(V),
                                    m_ConstantInt(InsIdx)))) {
      // Bail if any inserts are out of bounds.
      VectorType *OpTy = cast<VectorType>(Op->getType());
      if (OpTy->getElementCount().getKnownMinValue() <= InsIdx)
        return false;
      // All inserts must have the same index.
      // TODO: Deal with mismatched index constants and variable indexes?
      if (!Index)
        Index = InsIdx;
      else if (InsIdx != *Index)
        return false;
      VecCs.push_back(VecC);
      ScalarOps.push_back(V);
    } else if (II && isVectorIntrinsicWithScalarOpAtArg(II->getIntrinsicID(),
                                                        OpNum, &TTI)) {
      VecCs.push_back(Op.get());
      ScalarOps.push_back(Op.get());
    } else if (match(Op.get(), m_Constant(VecC))) {
      VecCs.push_back(VecC);
      ScalarOps.push_back(nullptr);
    } else {
      return false;
    }
  }

  // Bail if all operands are constant.
  if (!Index.has_value())
    return false;

  VectorType *VecTy = cast<VectorType>(I.getType());
  Type *ScalarTy = VecTy->getScalarType();
  assert(VecTy->isVectorTy() &&
         (ScalarTy->isIntegerTy() || ScalarTy->isFloatingPointTy() ||
          ScalarTy->isPointerTy()) &&
         "Unexpected types for insert element into binop or cmp");

  unsigned Opcode = I.getOpcode();
  InstructionCost ScalarOpCost, VectorOpCost;
  if (CI) {
    CmpInst::Predicate Pred = CI->getPredicate();
    ScalarOpCost = TTI.getCmpSelInstrCost(
        Opcode, ScalarTy, CmpInst::makeCmpResultType(ScalarTy), Pred, CostKind);
    VectorOpCost = TTI.getCmpSelInstrCost(
        Opcode, VecTy, CmpInst::makeCmpResultType(VecTy), Pred, CostKind);
  } else if (UO || BO) {
    ScalarOpCost = TTI.getArithmeticInstrCost(Opcode, ScalarTy, CostKind);
    VectorOpCost = TTI.getArithmeticInstrCost(Opcode, VecTy, CostKind);
  } else {
    IntrinsicCostAttributes ScalarICA(
        II->getIntrinsicID(), ScalarTy,
        SmallVector<Type *>(II->arg_size(), ScalarTy));
    ScalarOpCost = TTI.getIntrinsicInstrCost(ScalarICA, CostKind);
    IntrinsicCostAttributes VectorICA(
        II->getIntrinsicID(), VecTy,
        SmallVector<Type *>(II->arg_size(), VecTy));
    VectorOpCost = TTI.getIntrinsicInstrCost(VectorICA, CostKind);
  }

  // Fold the vector constants in the original vectors into a new base vector to
  // get more accurate cost modelling.
  Value *NewVecC = nullptr;
  if (CI)
    NewVecC = simplifyCmpInst(CI->getPredicate(), VecCs[0], VecCs[1], SQ);
  else if (UO)
    NewVecC =
        simplifyUnOp(UO->getOpcode(), VecCs[0], UO->getFastMathFlags(), SQ);
  else if (BO)
    NewVecC = simplifyBinOp(BO->getOpcode(), VecCs[0], VecCs[1], SQ);
  else if (II)
    NewVecC = simplifyCall(II, II->getCalledOperand(), VecCs, SQ);

  if (!NewVecC)
    return false;

  // Get cost estimate for the insert element. This cost will factor into
  // both sequences.
  InstructionCost OldCost = VectorOpCost;
  InstructionCost NewCost =
      ScalarOpCost + TTI.getVectorInstrCost(Instruction::InsertElement, VecTy,
                                            CostKind, *Index, NewVecC);

  for (auto [Idx, Op, VecC, Scalar] : enumerate(Ops, VecCs, ScalarOps)) {
    if (!Scalar || (II && isVectorIntrinsicWithScalarOpAtArg(
                              II->getIntrinsicID(), Idx, &TTI)))
      continue;
    InstructionCost InsertCost = TTI.getVectorInstrCost(
        Instruction::InsertElement, VecTy, CostKind, *Index, VecC, Scalar);
    OldCost += InsertCost;
    NewCost += !Op->hasOneUse() * InsertCost;
  }

  // We want to scalarize unless the vector variant actually has lower cost.
  if (OldCost < NewCost || !NewCost.isValid())
    return false;

  // vec_op (inselt VecC0, V0, Index), (inselt VecC1, V1, Index) -->
  // inselt NewVecC, (scalar_op V0, V1), Index
  if (CI)
    ++NumScalarCmp;
  else if (UO || BO)
    ++NumScalarOps;
  else
    ++NumScalarIntrinsic;

  // For constant cases, extract the scalar element, this should constant fold.
  for (auto [OpIdx, Scalar, VecC] : enumerate(ScalarOps, VecCs))
    if (!Scalar)
      ScalarOps[OpIdx] = ConstantExpr::getExtractElement(
          cast<Constant>(VecC), Builder.getInt64(*Index));

  Value *Scalar;
  // We need to pass the flags during the creation of instrucitons. Constant
  // folding might remove the instructions, so post setting the flags might
  // pollute the later instructions.
  if (CI) {
    if (FPMathOperator *FPMO = dyn_cast<FPMathOperator>(&I)) {
      Scalar = Builder.CreateFCmpFMF(CI->getPredicate(), ScalarOps[0],
                                     ScalarOps[1], FPMO->getFastMathFlags(),
                                     CI->getName() + ".scalar");
    } else {
      Scalar = Builder.CreateICmp(CI->getPredicate(), ScalarOps[0],
                                  ScalarOps[1], CI->getName() + ".scalar");
    }
  } else if (UO) {
    Scalar = Builder.CreateUnOpFMF(UO->getOpcode(), ScalarOps[0], UO,
                                   UO->getName() + ".scalar");
  } else if (BO) {
    if (OverflowingBinaryOperator *OBO =
            dyn_cast<OverflowingBinaryOperator>(&I)) {
      Scalar = Builder.CreateNoWrapBinOp(
          BO->getOpcode(), ScalarOps[0], ScalarOps[1], OBO->hasNoUnsignedWrap(),
          OBO->hasNoSignedWrap(), BO->getName() + ".scalar");
    } else if (PossiblyDisjointInst *PDI = dyn_cast<PossiblyDisjointInst>(&I)) {
      Scalar = Builder.CreateOr(ScalarOps[0], ScalarOps[1],
                                BO->getName() + ".scalar", PDI->isDisjoint());
    } else if (PossiblyExactOperator *PEO =
                   dyn_cast<PossiblyExactOperator>(&I)) {
      Scalar =
          Builder.CreateExactBinOp(BO->getOpcode(), ScalarOps[0], ScalarOps[1],
                                   PEO->isExact(), BO->getName() + ".scalar");
    } else if (FPMathOperator *FPMO = dyn_cast<FPMathOperator>(&I)) {
      Scalar = Builder.CreateBinOpFMF(BO->getOpcode(), ScalarOps[0],
                                      ScalarOps[1], FPMO->getFastMathFlags(),
                                      BO->getName() + ".scalar");
    } else {
      Scalar = Builder.CreateBinOp(BO->getOpcode(), ScalarOps[0], ScalarOps[1],
                                   BO->getName() + ".scalar");
    }
  } else {
    FastMathFlags FMF;
    if (auto *FPMO = dyn_cast<FPMathOperator>(&I))
      FMF = FPMO->getFastMathFlags();
    Scalar = Builder.CreateIntrinsic(ScalarTy, II->getIntrinsicID(), ScalarOps,
                                     FMF, II->getName() + ".scalar");
  }

  Value *Insert = Builder.CreateInsertElement(NewVecC, Scalar, *Index);
  replaceValue(I, *Insert);
  return true;
}

/// Try to combine a scalar binop + 2 scalar compares of extracted elements of
/// a vector into vector operations followed by extract. Note: The SLP pass
/// may miss this pattern because of implementation problems.
bool VectorCombine::foldExtractedCmps(Instruction &I) {
  auto *BI = dyn_cast<BinaryOperator>(&I);

  // We are looking for a scalar binop of booleans.
  // binop i1 (cmp Pred I0, C0), (cmp Pred I1, C1)
  if (!BI || !I.getType()->isIntegerTy(1))
    return false;

  // The compare predicates should match, and each compare should have a
  // constant operand.
  Value *B0 = I.getOperand(0), *B1 = I.getOperand(1);
  Instruction *I0, *I1;
  Constant *C0, *C1;
  CmpPredicate P0, P1;
  if (!match(B0, m_Cmp(P0, m_Instruction(I0), m_Constant(C0))) ||
      !match(B1, m_Cmp(P1, m_Instruction(I1), m_Constant(C1))))
    return false;

  auto MatchingPred = CmpPredicate::getMatching(P0, P1);
  if (!MatchingPred)
    return false;

  // The compare operands must be extracts of the same vector with constant
  // extract indexes.
  Value *X;
  uint64_t Index0, Index1;
  if (!match(I0, m_ExtractElt(m_Value(X), m_ConstantInt(Index0))) ||
      !match(I1, m_ExtractElt(m_Specific(X), m_ConstantInt(Index1))))
    return false;

  auto *Ext0 = cast<ExtractElementInst>(I0);
  auto *Ext1 = cast<ExtractElementInst>(I1);
  ExtractElementInst *ConvertToShuf = getShuffleExtract(Ext0, Ext1, CostKind);
  if (!ConvertToShuf)
    return false;
  assert((ConvertToShuf == Ext0 || ConvertToShuf == Ext1) &&
         "Unknown ExtractElementInst");

  // The original scalar pattern is:
  // binop i1 (cmp Pred (ext X, Index0), C0), (cmp Pred (ext X, Index1), C1)
  CmpInst::Predicate Pred = *MatchingPred;
  unsigned CmpOpcode =
      CmpInst::isFPPredicate(Pred) ? Instruction::FCmp : Instruction::ICmp;
  auto *VecTy = dyn_cast<FixedVectorType>(X->getType());
  if (!VecTy)
    return false;

  if (Index0 >= VecTy->getNumElements() || Index1 >= VecTy->getNumElements())
    return false;

  InstructionCost Ext0Cost =
      TTI.getVectorInstrCost(*Ext0, VecTy, CostKind, Index0);
  InstructionCost Ext1Cost =
      TTI.getVectorInstrCost(*Ext1, VecTy, CostKind, Index1);
  InstructionCost CmpCost = TTI.getCmpSelInstrCost(
      CmpOpcode, I0->getType(), CmpInst::makeCmpResultType(I0->getType()), Pred,
      CostKind);

  InstructionCost OldCost =
      Ext0Cost + Ext1Cost + CmpCost * 2 +
      TTI.getArithmeticInstrCost(I.getOpcode(), I.getType(), CostKind);

  // The proposed vector pattern is:
  // vcmp = cmp Pred X, VecC
  // ext (binop vNi1 vcmp, (shuffle vcmp, Index1)), Index0
  int CheapIndex = ConvertToShuf == Ext0 ? Index1 : Index0;
  int ExpensiveIndex = ConvertToShuf == Ext0 ? Index0 : Index1;
  auto *CmpTy = cast<FixedVectorType>(CmpInst::makeCmpResultType(VecTy));
  InstructionCost NewCost = TTI.getCmpSelInstrCost(
      CmpOpcode, VecTy, CmpInst::makeCmpResultType(VecTy), Pred, CostKind);
  SmallVector<int, 32> ShufMask(VecTy->getNumElements(), PoisonMaskElem);
  ShufMask[CheapIndex] = ExpensiveIndex;
  NewCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc, CmpTy,
                                CmpTy, CostKind, ShufMask);
  NewCost += TTI.getArithmeticInstrCost(I.getOpcode(), CmpTy, CostKind);
  NewCost += TTI.getVectorInstrCost(*Ext0, CmpTy, CostKind, CheapIndex);
  NewCost += Ext0->hasOneUse() ? 0 : Ext0Cost;
  NewCost += Ext1->hasOneUse() ? 0 : Ext1Cost;

  // Aggressively form vector ops if the cost is equal because the transform
  // may enable further optimization.
  // Codegen can reverse this transform (scalarize) if it was not profitable.
  if (OldCost < NewCost || !NewCost.isValid())
    return false;

  // Create a vector constant from the 2 scalar constants.
  SmallVector<Constant *, 32> CmpC(VecTy->getNumElements(),
                                   PoisonValue::get(VecTy->getElementType()));
  CmpC[Index0] = C0;
  CmpC[Index1] = C1;
  Value *VCmp = Builder.CreateCmp(Pred, X, ConstantVector::get(CmpC));
  Value *Shuf = createShiftShuffle(VCmp, ExpensiveIndex, CheapIndex, Builder);
  Value *LHS = ConvertToShuf == Ext0 ? Shuf : VCmp;
  Value *RHS = ConvertToShuf == Ext0 ? VCmp : Shuf;
  Value *VecLogic = Builder.CreateBinOp(BI->getOpcode(), LHS, RHS);
  Value *NewExt = Builder.CreateExtractElement(VecLogic, CheapIndex);
  replaceValue(I, *NewExt);
  ++NumVecCmpBO;
  return true;
}

/// Try to fold scalar selects that select between extracted elements and zero
/// into extracting from a vector select. This is rooted at the bitcast.
///
/// This pattern arises when a vector is bitcast to a smaller element type,
/// elements are extracted, and then conditionally selected with zero:
///
///   %bc = bitcast <4 x i32> %src to <16 x i8>
///   %e0 = extractelement <16 x i8> %bc, i32 0
///   %s0 = select i1 %cond, i8 %e0, i8 0
///   %e1 = extractelement <16 x i8> %bc, i32 1
///   %s1 = select i1 %cond, i8 %e1, i8 0
///   ...
///
/// Transforms to:
///   %sel = select i1 %cond, <4 x i32> %src, <4 x i32> zeroinitializer
///   %bc = bitcast <4 x i32> %sel to <16 x i8>
///   %e0 = extractelement <16 x i8> %bc, i32 0
///   %e1 = extractelement <16 x i8> %bc, i32 1
///   ...
///
/// This is profitable because vector select on wider types produces fewer
/// select/cndmask instructions than scalar selects on each element.
bool VectorCombine::foldSelectsFromBitcast(Instruction &I) {
  auto *BC = dyn_cast<BitCastInst>(&I);
  if (!BC)
    return false;

  FixedVectorType *SrcVecTy = dyn_cast<FixedVectorType>(BC->getSrcTy());
  FixedVectorType *DstVecTy = dyn_cast<FixedVectorType>(BC->getDestTy());
  if (!SrcVecTy || !DstVecTy)
    return false;

  // Source must be 32-bit or 64-bit elements, destination must be smaller
  // integer elements. Zero in all these types is all-bits-zero.
  Type *SrcEltTy = SrcVecTy->getElementType();
  Type *DstEltTy = DstVecTy->getElementType();
  unsigned SrcEltBits = SrcEltTy->getPrimitiveSizeInBits();
  unsigned DstEltBits = DstEltTy->getPrimitiveSizeInBits();

  if (SrcEltBits != 32 && SrcEltBits != 64)
    return false;

  if (!DstEltTy->isIntegerTy() || DstEltBits >= SrcEltBits)
    return false;

  // Check profitability using TTI before collecting users.
  Type *CondTy = CmpInst::makeCmpResultType(DstEltTy);
  Type *VecCondTy = CmpInst::makeCmpResultType(SrcVecTy);

  InstructionCost ScalarSelCost =
      TTI.getCmpSelInstrCost(Instruction::Select, DstEltTy, CondTy,
                             CmpInst::BAD_ICMP_PREDICATE, CostKind);
  InstructionCost VecSelCost =
      TTI.getCmpSelInstrCost(Instruction::Select, SrcVecTy, VecCondTy,
                             CmpInst::BAD_ICMP_PREDICATE, CostKind);

  // We need at least this many selects for vectorization to be profitable.
  // VecSelCost < ScalarSelCost * NumSelects => NumSelects > VecSelCost /
  // ScalarSelCost
  if (!ScalarSelCost.isValid() || ScalarSelCost == 0)
    return false;

  unsigned MinSelects = (VecSelCost.getValue() / ScalarSelCost.getValue()) + 1;

  // Quick check: if bitcast doesn't have enough users, bail early.
  if (!BC->hasNUsesOrMore(MinSelects))
    return false;

  // Collect all select users that match the pattern, grouped by condition.
  // Pattern: select i1 %cond, (extractelement %bc, idx), 0
  DenseMap<Value *, SmallVector<SelectInst *, 8>> CondToSelects;

  for (User *U : BC->users()) {
    auto *Ext = dyn_cast<ExtractElementInst>(U);
    if (!Ext)
      continue;

    for (User *ExtUser : Ext->users()) {
      Value *Cond;
      // Match: select i1 %cond, %ext, 0
      if (match(ExtUser, m_Select(m_Value(Cond), m_Specific(Ext), m_Zero())) &&
          Cond->getType()->isIntegerTy(1))
        CondToSelects[Cond].push_back(cast<SelectInst>(ExtUser));
    }
  }

  if (CondToSelects.empty())
    return false;

  bool MadeChange = false;
  Value *SrcVec = BC->getOperand(0);

  // Process each group of selects with the same condition.
  for (auto [Cond, Selects] : CondToSelects) {
    // Only profitable if vector select cost < total scalar select cost.
    if (Selects.size() < MinSelects) {
      LLVM_DEBUG(dbgs() << "VectorCombine: foldSelectsFromBitcast not "
                        << "profitable (VecCost=" << VecSelCost
                        << ", ScalarCost=" << ScalarSelCost
                        << ", NumSelects=" << Selects.size() << ")\n");
      continue;
    }

    // Create the vector select and bitcast once for this condition.
    auto InsertPt = std::next(BC->getIterator());

    if (auto *CondInst = dyn_cast<Instruction>(Cond))
      if (DT.dominates(BC, CondInst))
        InsertPt = std::next(CondInst->getIterator());

    Builder.SetInsertPoint(InsertPt);
    Value *VecSel =
        Builder.CreateSelect(Cond, SrcVec, Constant::getNullValue(SrcVecTy));
    Value *NewBC = Builder.CreateBitCast(VecSel, DstVecTy);

    // Replace each scalar select with an extract from the new bitcast.
    for (SelectInst *Sel : Selects) {
      auto *Ext = cast<ExtractElementInst>(Sel->getTrueValue());
      Value *Idx = Ext->getIndexOperand();

      Builder.SetInsertPoint(Sel);
      Value *NewExt = Builder.CreateExtractElement(NewBC, Idx);
      replaceValue(*Sel, *NewExt);
      MadeChange = true;
    }

    LLVM_DEBUG(dbgs() << "VectorCombine: folded " << Selects.size()
                      << " selects into vector select\n");
  }

  return MadeChange;
}

static void analyzeCostOfVecReduction(const IntrinsicInst &II,
                                      TTI::TargetCostKind CostKind,
                                      const TargetTransformInfo &TTI,
                                      InstructionCost &CostBeforeReduction,
                                      InstructionCost &CostAfterReduction) {
  Instruction *Op0, *Op1;
  auto *RedOp = dyn_cast<Instruction>(II.getOperand(0));
  auto *VecRedTy = cast<VectorType>(II.getOperand(0)->getType());
  unsigned ReductionOpc =
      getArithmeticReductionInstruction(II.getIntrinsicID());
  if (RedOp && match(RedOp, m_ZExtOrSExt(m_Value()))) {
    bool IsUnsigned = isa<ZExtInst>(RedOp);
    auto *ExtType = cast<VectorType>(RedOp->getOperand(0)->getType());

    CostBeforeReduction =
        TTI.getCastInstrCost(RedOp->getOpcode(), VecRedTy, ExtType,
                             TTI::CastContextHint::None, CostKind, RedOp);
    CostAfterReduction =
        TTI.getExtendedReductionCost(ReductionOpc, IsUnsigned, II.getType(),
                                     ExtType, FastMathFlags(), CostKind);
    return;
  }
  if (RedOp && II.getIntrinsicID() == Intrinsic::vector_reduce_add &&
      match(RedOp,
            m_ZExtOrSExt(m_Mul(m_Instruction(Op0), m_Instruction(Op1)))) &&
      match(Op0, m_ZExtOrSExt(m_Value())) &&
      Op0->getOpcode() == Op1->getOpcode() &&
      Op0->getOperand(0)->getType() == Op1->getOperand(0)->getType() &&
      (Op0->getOpcode() == RedOp->getOpcode() || Op0 == Op1)) {
    // Matched reduce.add(ext(mul(ext(A), ext(B)))
    bool IsUnsigned = isa<ZExtInst>(Op0);
    auto *ExtType = cast<VectorType>(Op0->getOperand(0)->getType());
    VectorType *MulType = VectorType::get(Op0->getType(), VecRedTy);

    InstructionCost ExtCost =
        TTI.getCastInstrCost(Op0->getOpcode(), MulType, ExtType,
                             TTI::CastContextHint::None, CostKind, Op0);
    InstructionCost MulCost =
        TTI.getArithmeticInstrCost(Instruction::Mul, MulType, CostKind);
    InstructionCost Ext2Cost =
        TTI.getCastInstrCost(RedOp->getOpcode(), VecRedTy, MulType,
                             TTI::CastContextHint::None, CostKind, RedOp);

    CostBeforeReduction = ExtCost * 2 + MulCost + Ext2Cost;
    CostAfterReduction = TTI.getMulAccReductionCost(
        IsUnsigned, ReductionOpc, II.getType(), ExtType, CostKind);
    return;
  }
  CostAfterReduction = TTI.getArithmeticReductionCost(ReductionOpc, VecRedTy,
                                                      std::nullopt, CostKind);
}

bool VectorCombine::foldBinopOfReductions(Instruction &I) {
  Instruction::BinaryOps BinOpOpc = cast<BinaryOperator>(&I)->getOpcode();
  Intrinsic::ID ReductionIID = getReductionForBinop(BinOpOpc);
  if (BinOpOpc == Instruction::Sub)
    ReductionIID = Intrinsic::vector_reduce_add;
  if (ReductionIID == Intrinsic::not_intrinsic)
    return false;
  // FP reductions have a start-value operand that this fold doesn't handle.
  if (ReductionIID == Intrinsic::vector_reduce_fadd ||
      ReductionIID == Intrinsic::vector_reduce_fmul)
    return false;

  auto checkIntrinsicAndGetItsArgument = [](Value *V,
                                            Intrinsic::ID IID) -> Value * {
    auto *II = dyn_cast<IntrinsicInst>(V);
    if (!II)
      return nullptr;
    if (II->getIntrinsicID() == IID && II->hasOneUse())
      return II->getArgOperand(0);
    return nullptr;
  };

  Value *V0 = checkIntrinsicAndGetItsArgument(I.getOperand(0), ReductionIID);
  if (!V0)
    return false;
  Value *V1 = checkIntrinsicAndGetItsArgument(I.getOperand(1), ReductionIID);
  if (!V1)
    return false;

  auto *VTy = cast<VectorType>(V0->getType());
  if (V1->getType() != VTy)
    return false;
  const auto &II0 = *cast<IntrinsicInst>(I.getOperand(0));
  const auto &II1 = *cast<IntrinsicInst>(I.getOperand(1));
  unsigned ReductionOpc =
      getArithmeticReductionInstruction(II0.getIntrinsicID());

  InstructionCost OldCost = 0;
  InstructionCost NewCost = 0;
  InstructionCost CostOfRedOperand0 = 0;
  InstructionCost CostOfRed0 = 0;
  InstructionCost CostOfRedOperand1 = 0;
  InstructionCost CostOfRed1 = 0;
  analyzeCostOfVecReduction(II0, CostKind, TTI, CostOfRedOperand0, CostOfRed0);
  analyzeCostOfVecReduction(II1, CostKind, TTI, CostOfRedOperand1, CostOfRed1);
  OldCost = CostOfRed0 + CostOfRed1 + TTI.getInstructionCost(&I, CostKind);
  NewCost =
      CostOfRedOperand0 + CostOfRedOperand1 +
      TTI.getArithmeticInstrCost(BinOpOpc, VTy, CostKind) +
      TTI.getArithmeticReductionCost(ReductionOpc, VTy, std::nullopt, CostKind);
  if (NewCost >= OldCost || !NewCost.isValid())
    return false;

  LLVM_DEBUG(dbgs() << "Found two mergeable reductions: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");
  Value *VectorBO;
  if (BinOpOpc == Instruction::Or)
    VectorBO = Builder.CreateOr(V0, V1, "",
                                cast<PossiblyDisjointInst>(I).isDisjoint());
  else
    VectorBO = Builder.CreateBinOp(BinOpOpc, V0, V1);

  Value *Rdx = Builder.CreateIntrinsic(ReductionIID, {VTy}, {VectorBO});
  replaceValue(I, *Rdx);
  return true;
}

// Check if memory is modified, freed, or synchronized between two instrs in
// the same BB.
static bool isMemModifiedBetween(BasicBlock::iterator Begin,
                                 BasicBlock::iterator End,
                                 const MemoryLocation &Loc, AAResults &AA) {
  unsigned NumScanned = 0;
  if (std::any_of(Begin, End, [&](const Instruction &Instr) {
        return isModSet(AA.getModRefInfo(&Instr, Loc)) ||
               ++NumScanned > MaxInstrsToScan;
      }))
    return true;

  // willNotFreeBetween expects instructions rather than iterators. An empty
  // range cannot free or synchronize, so avoid dereferencing its end.
  return Begin != End && !willNotFreeBetween(&*Begin, &*End);
}

namespace {
/// Helper class to indicate whether a vector index can be safely scalarized and
/// if a freeze needs to be inserted.
class ScalarizationResult {
  enum class StatusTy { Unsafe, Safe, SafeWithFreeze };

  StatusTy Status;
  Value *ToFreeze;

  ScalarizationResult(StatusTy Status, Value *ToFreeze = nullptr)
      : Status(Status), ToFreeze(ToFreeze) {}

public:
  ScalarizationResult(const ScalarizationResult &Other) = default;
  ~ScalarizationResult() {
    assert(!ToFreeze && "freeze() not called with ToFreeze being set");
  }

  static ScalarizationResult unsafe() { return {StatusTy::Unsafe}; }
  static ScalarizationResult safe() { return {StatusTy::Safe}; }
  static ScalarizationResult safeWithFreeze(Value *ToFreeze) {
    return {StatusTy::SafeWithFreeze, ToFreeze};
  }

  /// Returns true if the index can be scalarize without requiring a freeze.
  bool isSafe() const { return Status == StatusTy::Safe; }
  /// Returns true if the index cannot be scalarized.
  bool isUnsafe() const { return Status == StatusTy::Unsafe; }
  /// Returns true if the index can be scalarize, but requires inserting a
  /// freeze.
  bool isSafeWithFreeze() const { return Status == StatusTy::SafeWithFreeze; }

  /// Reset the state of Unsafe and clear ToFreze if set.
  void discard() {
    ToFreeze = nullptr;
    Status = StatusTy::Unsafe;
  }

  /// Freeze the ToFreeze and update the use in \p User to use it.
  void freeze(IRBuilderBase &Builder, Instruction &UserI) {
    assert(isSafeWithFreeze() &&
           "should only be used when freezing is required");
    assert(is_contained(ToFreeze->users(), &UserI) &&
           "UserI must be a user of ToFreeze");
    IRBuilder<>::InsertPointGuard Guard(Builder);
    Builder.SetInsertPoint(cast<Instruction>(&UserI));
    Value *Frozen =
        Builder.CreateFreeze(ToFreeze, ToFreeze->getName() + ".frozen");
    for (Use &U : make_early_inc_range((UserI.operands())))
      if (U.get() == ToFreeze)
        U.set(Frozen);

    ToFreeze = nullptr;
  }
};
} // namespace

/// Check if it is legal to scalarize a memory access to \p VecTy at index \p
/// Idx. \p Idx must access a valid vector element.
static ScalarizationResult canScalarizeAccess(VectorType *VecTy, Value *Idx,
                                              const SimplifyQuery &SQ) {
  // We do checks for both fixed vector types and scalable vector types.
  // This is the number of elements of fixed vector types,
  // or the minimum number of elements of scalable vector types.
  uint64_t NumElements = VecTy->getElementCount().getKnownMinValue();
  unsigned IntWidth = Idx->getType()->getScalarSizeInBits();

  if (auto *C = dyn_cast<ConstantInt>(Idx)) {
    if (C->getValue().ult(NumElements))
      return ScalarizationResult::safe();
    return ScalarizationResult::unsafe();
  }

  // Always unsafe if the index type can't handle all inbound values.
  if (!llvm::isUIntN(IntWidth, NumElements))
    return ScalarizationResult::unsafe();

  APInt Zero(IntWidth, 0);
  APInt MaxElts(IntWidth, NumElements);
  ConstantRange ValidIndices(Zero, MaxElts);
  ConstantRange IdxRange(IntWidth, true);

  if (isGuaranteedNotToBePoison(Idx, SQ.AC, SQ.CtxI, SQ.DT)) {
    if (ValidIndices.contains(
            computeConstantRange(Idx, /*ForSigned=*/false, SQ)))
      return ScalarizationResult::safe();
    return ScalarizationResult::unsafe();
  }

  // If the index may be poison, check if we can insert a freeze before the
  // range of the index is restricted.
  Value *IdxBase;
  ConstantInt *CI;
  if (match(Idx, m_And(m_Value(IdxBase), m_ConstantInt(CI)))) {
    IdxRange = IdxRange.binaryAnd(CI->getValue());
  } else if (match(Idx, m_URem(m_Value(IdxBase), m_ConstantInt(CI)))) {
    IdxRange = IdxRange.urem(CI->getValue());
  }

  if (ValidIndices.contains(IdxRange))
    return ScalarizationResult::safeWithFreeze(IdxBase);
  return ScalarizationResult::unsafe();
}

/// Return the GEP index type if the unsigned vector index \p Idx can be
/// represented by an inbounds GEP. A null result means that the maximum byte
/// offset cannot be represented by the pointer's signed GEP index type.
///
///   unsigned lane range
///                |
///                v
///   MaxByteOffset = MaxLane * element store size
///                |
///                +-- unavailable or outside signed GEP range --> reject
///                |
///                v
///   valid range --> use the pointer's GEP index type
static IntegerType *getScalarizedGEPIndexInfo(VectorType *VecTy, Value *Idx,
                                              Type *PtrTy,
                                              const DataLayout &DL) {
  auto *GEPIndexTy = cast<IntegerType>(DL.getIndexType(PtrTy));
  unsigned GEPBits = GEPIndexTy->getBitWidth();
  uint64_t NumElements = VecTy->getElementCount().getKnownMinValue();

  uint64_t MaxLane = NumElements - 1;
  if (auto *C = dyn_cast<ConstantInt>(Idx)) {
    if (C->getValue().uge(NumElements))
      return nullptr;
    MaxLane = C->getZExtValue();
  }

  Type *ElemTy = VecTy->getElementType();
  if (!DL.typeSizeEqualsStoreSize(ElemTy))
    return nullptr;

  TypeSize ElemStride = DL.getTypeStoreSize(ElemTy);
  if (ElemStride.isScalable())
    return nullptr;

  // Compare both values in a common width:
  //
  //   MaxLane (uint64_t) * ElemStride (uint64_t)   signed_max(GEPBits)
  //                         |                               |
  //                         v                               v
  //             ByteOffset (up to 128 bits)        sext to WideBits
  //                         \                               /
  //                          +------------ ugt ------------+
  //                                       |
  //                                greater -> reject
  //
  // WideBits = max(GEPBits, 128) prevents the multiplication from wrapping
  // and preserves the GEP limit during the comparison.
  unsigned WideBits = std::max(GEPBits, 128u);
  APInt MaxLaneValue(WideBits, MaxLane);
  APInt ByteOffset = MaxLaneValue;
  ByteOffset *= APInt(WideBits, ElemStride.getFixedValue());
  APInt MaxGEPOffset = APInt::getSignedMaxValue(GEPBits).sext(WideBits);
  // Reject offsets outside the GEP's positive signed range. Compare as
  // unsigned because the full 128-bit product may set its sign bit.
  if (ByteOffset.ugt(MaxGEPOffset))
    return nullptr;

  return GEPIndexTy;
}

/// Materialize an index for a scalarized GEP after profitability is known.
/// Vector element indices are unsigned, but GEP sign-extends narrow integer
/// indices. Widen a narrow index explicitly so its unsigned value is retained.
static Value *materializeScalarizedGEPIndex(Value *Idx, IntegerType *GEPIndexTy,
                                            IRBuilderBase &Builder) {
  unsigned SrcBits = Idx->getType()->getIntegerBitWidth();
  unsigned DstBits = GEPIndexTy->getBitWidth();
  if (SrcBits >= DstBits)
    return Idx;

  return Builder.CreateZExt(Idx, GEPIndexTy, Idx->getName() + ".gepidx");
}

/// The memory operation on a vector of \p ScalarType had alignment of
/// \p VectorAlignment. Compute the maximal, but conservatively correct,
/// alignment that will be valid for the memory operation on a single scalar
/// element of the same type with index \p Idx.
static Align computeAlignmentAfterScalarization(Align VectorAlignment,
                                                Type *ScalarType, Value *Idx,
                                                const DataLayout &DL) {
  if (auto *C = dyn_cast<ConstantInt>(Idx))
    return commonAlignment(VectorAlignment,
                           C->getZExtValue() * DL.getTypeStoreSize(ScalarType));
  return commonAlignment(VectorAlignment, DL.getTypeStoreSize(ScalarType));
}

/// Fold a vector store fed by a single-use insertelement chain into scalar
/// stores.
///
/// Before:
///
///   %p --> vector load --> insert %x, lane 1 --> insert %y, lane 3
///                                                      |
///                                                      v
///                                              vector store to %p
///
///   Vector lanes:        [ 0 ] [ 1 ] [ 2 ] [ 3 ]
///   Stored value:        [ old |  x  | old |  y  ]  (one vector store)
///
/// After:
///
///                  +--> GEP(%p, lane 1) --> store %x
///   %p -------------+
///                  +--> GEP(%p, lane 3) --> store %y
///
///   Vector lanes:        [ 0 ] [ 1 ] [ 2 ] [ 3 ]
///   Scalar stores:              x             y
///                            store@1       store@3
///
///   Step 1. Gate:
///        target supports vector-element GEP addressing
///
///   Step 2. Trace:
///        vector store <-- insertelement <-- ... <-- insertelement <-- load
///
///   Steps 3-5. Validate:
///        reject unprofitable full overwrites; require simple accesses, a
///        common address/block, no memory write in between, and scalarizable
///        indices.
bool VectorCombine::foldInsertElementsToStores(Instruction &I) {
  // Step 1: The target must support addressing a vector element with a GEP.
  if (!TTI.allowVectorElementIndexingUsingGEP())
    return false;

  auto *SI = cast<StoreInst>(&I);
  if (!SI->isSimple() || !isa<VectorType>(SI->getValueOperand()->getType()))
    return false;

  // Step 2: Collect a single-use insertelement chain, starting at the vector
  // store and walking back to the candidate load.
  Value *Source = SI->getValueOperand();
  SmallVector<std::pair<Value *, Value *>, 4> InsertElements;
  Value *Base = Source;
  while (auto *Insert = dyn_cast<InsertElementInst>(Base)) {
    if (!Insert->hasOneUse())
      break;
    Value *InsertVal = Insert->getOperand(1);
    Value *Idx = Insert->getOperand(2);
    InsertElements.push_back({InsertVal, Idx});
    Base = Insert->getOperand(0);
  }

  if (InsertElements.empty())
    return false;

  // The backwards walk collected the inserts in reverse program order. Restore
  // it now so later scalar stores preserve writes to duplicate/equal indices.
  std::reverse(InsertElements.begin(), InsertElements.end());
  auto *Load = dyn_cast<LoadInst>(Base);
  if (!Load)
    return false;
  auto *VecTy = cast<VectorType>(SI->getValueOperand()->getType());

  // Step 3: Avoid replacing a complete overwrite with scalar stores when every
  // lane receives the same value; keeping the vector operation is preferable.
  if (auto *FVT = dyn_cast<FixedVectorType>(VecTy)) {
    if (InsertElements.size() == FVT->getNumElements()) {
      Value *FirstVal = InsertElements.front().first;
      if (all_of(InsertElements,
                 [FirstVal](const auto &Elt) { return Elt.first == FirstVal; }))
        return false;
    }
  }
  Value *SrcAddr = Load->getPointerOperand()->stripPointerCasts();
  // Step 4: Establish the load/store update is legal: both accesses are simple,
  // have the same base address and block, have scalar elements whose type size
  // equals their store size, and no intervening operation modifies the updated
  // memory.
  if (!Load->isSimple() || Load->getParent() != SI->getParent() ||
      !DL->typeSizeEqualsStoreSize(Load->getType()->getScalarType()) ||
      SrcAddr != SI->getPointerOperand()->stripPointerCasts())
    return false;

  if (isMemModifiedBetween(Load->getIterator(), SI->getIterator(),
                           MemoryLocation::get(SI), AA))
    return false;

  // Step 5: Validate every index before changing IR. A safe-with-freeze result
  // is recorded by ScalarizationResult, so discard it until profitability is
  // known; otherwise a rejected candidate could leave a freeze behind.
  for (auto [InsertVal, Idx] : InsertElements) {
    auto ScalarizableIdx =
        canScalarizeAccess(VecTy, Idx, SQ.getWithInstruction(&I));
    if (ScalarizableIdx.isUnsafe())
      return false;

    auto GEPIndex =
        getScalarizedGEPIndexInfo(VecTy, Idx, SI->getPointerOperandType(), *DL);
    if (!GEPIndex) {
      ScalarizableIdx.discard();
      return false;
    }

    // We are only checking legality here. Do not mutate IR before the
    // profitability check, but also do not leave a pending ToFreeze behind.
    ScalarizableIdx.discard();
  }

  InstructionCost OldCost = TTI.getMemoryOpCost(
      Instruction::Store, SI->getValueOperand()->getType(), SI->getAlign(),
      SI->getPointerAddressSpace(), CostKind);

  if (Load->hasOneUse())
    OldCost += TTI.getMemoryOpCost(Instruction::Load, Load->getType(),
                                   Load->getAlign(),
                                   Load->getPointerAddressSpace(), CostKind);

  for (auto [InsertVal, Idx] : InsertElements) {
    int Index = -1;
    if (auto *CIdx = dyn_cast<ConstantInt>(Idx))
      Index = CIdx->getZExtValue();

    OldCost += TTI.getVectorInstrCost(Instruction::InsertElement, VecTy,
                                      CostKind, Index);
  }

  InstructionCost NewCost = 0;
  // This transform replaces insertelement operations on a single vector with
  // GEPs and scalar stores, so assume constant-index GEP offsets stay within
  // addressing-mode ranges that getGEPCost considers TCC_Free. Cost only GEPs
  // with dynamic indices.
  for (auto [InsertVal, Idx] : InsertElements) {
    if (isa<ConstantInt>(Idx))
      continue;
    const Value *GEPIndices[] = {ConstantInt::get(Idx->getType(), 0), Idx};
    NewCost += TTI.getGEPCost(VecTy, SI->getPointerOperand(), GEPIndices,
                              CostKind, InsertVal->getType());
  }

  for (auto [InsertVal, Idx] : InsertElements) {
    Align ScalarOpAlignment = computeAlignmentAfterScalarization(
        std::max(SI->getAlign(), Load->getAlign()), InsertVal->getType(), Idx,
        *DL);

    NewCost += TTI.getMemoryOpCost(Instruction::Store, InsertVal->getType(),
                                   ScalarOpAlignment,
                                   SI->getPointerAddressSpace(), CostKind);
  }

  LLVM_DEBUG(dbgs() << "Found an insert-elements vector store scalarization "
                       "candidate: "
                    << I << "\n"
                    << "  NumInserts: " << InsertElements.size() << "\n"
                    << "  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");

  if (OldCost <= NewCost)
    return false;

  for (auto [InsertVal, Idx] : InsertElements) {
    auto ScalarizableIdx =
        canScalarizeAccess(VecTy, Idx, SQ.getWithInstruction(&I));
    assert(!ScalarizableIdx.isUnsafe() && "already checked above");

    if (ScalarizableIdx.isSafeWithFreeze())
      ScalarizableIdx.freeze(Builder, *cast<Instruction>(Idx));
  }

  Worklist.push(Load);
  StoreInst *LastStore = nullptr;
  for (auto [InsertVal, Idx] : InsertElements) {
    auto ScalarizableIdx =
        canScalarizeAccess(VecTy, Idx, SQ.getWithInstruction(&I));
    if (ScalarizableIdx.isUnsafe())
      return false;

    IntegerType *GEPIndexTy =
        getScalarizedGEPIndexInfo(VecTy, Idx, SI->getPointerOperandType(), *DL);

    Value *GEPIdx = materializeScalarizedGEPIndex(Idx, GEPIndexTy, Builder);
    Value *GEP = Builder.CreateInBoundsGEP(
        SI->getValueOperand()->getType(), SI->getPointerOperand(),
        {ConstantInt::get(GEPIdx->getType(), 0), GEPIdx});

    LastStore = Builder.CreateStore(InsertVal, GEP);
    LastStore->copyMetadata(*SI);

    // The new GEP may change the pointer operand, so !invariant.group cannot
    // be transferred to the scalar store.
    LastStore->setMetadata(LLVMContext::MD_invariant_group, nullptr);
    Align ScalarOpAlignment = computeAlignmentAfterScalarization(
        std::max(SI->getAlign(), Load->getAlign()), InsertVal->getType(), Idx,
        *DL);
    LastStore->setAlignment(ScalarOpAlignment);
  }

  replaceValue(I, *LastStore);
  eraseInstruction(I);
  return true;
}

/// Try to scalarize vector loads feeding extractelement or bitcast
/// instructions.
bool VectorCombine::scalarizeLoad(Instruction &I) {
  Value *Ptr;
  if (!match(&I, m_Load(m_Value(Ptr))))
    return false;

  auto *LI = cast<LoadInst>(&I);
  auto *VecTy = cast<VectorType>(LI->getType());

  // The isSimple() check could be isUnordered(), but for now we cowardly
  // refuse to handle even unordered atomics.
  if (!LI->isSimple() || !DL->typeSizeEqualsStoreSize(VecTy->getScalarType()))
    return false;

  bool AllExtracts = true;
  bool AllBitcasts = true;
  Instruction *LastCheckedInst = LI;
  unsigned NumInstChecked = 0;

  // Check what type of users we have (must either all be extracts or
  // bitcasts) and ensure no memory modifications between the load and
  // its users.
  for (User *U : LI->users()) {
    auto *UI = dyn_cast<Instruction>(U);
    if (!UI || UI->getParent() != LI->getParent())
      return false;

    // If any user is waiting to be erased, then bail out as this will
    // distort the cost calculation and possibly lead to infinite loops.
    if (UI->use_empty())
      return false;

    if (!isa<ExtractElementInst>(UI))
      AllExtracts = false;
    if (!isa<BitCastInst>(UI))
      AllBitcasts = false;

    // Check if any instruction between the load and the user may modify memory.
    if (LastCheckedInst->comesBefore(UI)) {
      for (Instruction &I :
           make_range(std::next(LI->getIterator()), UI->getIterator())) {
        // Bail out if we reached the check limit or the instruction may write
        // to memory.
        if (NumInstChecked == MaxInstrsToScan || I.mayWriteToMemory())
          return false;
        NumInstChecked++;
      }
      LastCheckedInst = UI;
    }
  }

  if (AllExtracts)
    return scalarizeLoadExtract(LI, VecTy, Ptr);
  if (AllBitcasts)
    return scalarizeLoadBitcast(LI, VecTy, Ptr);
  return false;
}

/// Try to scalarize vector loads feeding extractelement instructions.
bool VectorCombine::scalarizeLoadExtract(LoadInst *LI, VectorType *VecTy,
                                         Value *Ptr) {
  if (!TTI.allowVectorElementIndexingUsingGEP())
    return false;

  DenseMap<ExtractElementInst *, ScalarizationResult> NeedFreeze;
  DenseMap<ExtractElementInst *, IntegerType *> GEPIndexInfos;
  llvm::scope_exit FailureGuard([&]() {
    // If the transform is aborted, discard the ScalarizationResults.
    for (auto &Pair : NeedFreeze)
      Pair.second.discard();
  });

  InstructionCost OriginalCost =
      TTI.getMemoryOpCost(Instruction::Load, VecTy, LI->getAlign(),
                          LI->getPointerAddressSpace(), CostKind);
  InstructionCost ScalarizedCost = 0;

  for (User *U : LI->users()) {
    auto *UI = cast<ExtractElementInst>(U);

    auto ScalarIdx = canScalarizeAccess(VecTy, UI->getIndexOperand(),
                                        SQ.getWithInstruction(LI));
    if (ScalarIdx.isUnsafe())
      return false;

    IntegerType *GEPIndex = getScalarizedGEPIndexInfo(
        VecTy, UI->getIndexOperand(), LI->getPointerOperandType(), *DL);
    if (!GEPIndex) {
      ScalarIdx.discard();
      return false;
    }

    GEPIndexInfos.try_emplace(UI, GEPIndex);

    if (ScalarIdx.isSafeWithFreeze()) {
      NeedFreeze.try_emplace(UI, ScalarIdx);
      ScalarIdx.discard();
    }

    auto *Index = dyn_cast<ConstantInt>(UI->getIndexOperand());
    OriginalCost +=
        TTI.getVectorInstrCost(Instruction::ExtractElement, VecTy, CostKind,
                               Index ? Index->getZExtValue() : -1);
    ScalarizedCost +=
        TTI.getMemoryOpCost(Instruction::Load, VecTy->getElementType(),
                            Align(1), LI->getPointerAddressSpace(), CostKind);
    ScalarizedCost += TTI.getAddressComputationCost(LI->getPointerOperandType(),
                                                    nullptr, nullptr, CostKind);
    if (!Index && UI->getIndexOperand()->getType()->getIntegerBitWidth() <
                      GEPIndex->getBitWidth())
      ScalarizedCost += TTI.getCastInstrCost(
          Instruction::ZExt, GEPIndex, UI->getIndexOperand()->getType(),
          TTI::CastContextHint::None, CostKind);
  }

  LLVM_DEBUG(dbgs() << "Found all extractions of a vector load: " << *LI
                    << "\n  LoadExtractCost: " << OriginalCost
                    << " vs ScalarizedCost: " << ScalarizedCost << "\n");

  if (ScalarizedCost > OriginalCost)
    return false;
  if (ScalarizedCost == OriginalCost && !LI->hasOneUse())
    return false;

  // Ensure we add the load back to the worklist BEFORE its users so they can
  // erased in the correct order.
  Worklist.push(LI);

  Type *ElemType = VecTy->getElementType();

  // Replace extracts with narrow scalar loads.
  for (User *U : LI->users()) {
    auto *EI = cast<ExtractElementInst>(U);
    Value *Idx = EI->getIndexOperand();

    // Insert 'freeze' for poison indexes.
    if (auto It = NeedFreeze.find(EI); It != NeedFreeze.end())
      It->second.freeze(Builder, *cast<Instruction>(Idx));

    Builder.SetInsertPoint(EI);
    auto It = GEPIndexInfos.find(EI);
    assert(It != GEPIndexInfos.end() &&
           "Missing scalarized GEP index information");
    Value *GEPIdx = materializeScalarizedGEPIndex(Idx, It->second, Builder);
    Value *GEP = Builder.CreateInBoundsGEP(
        VecTy, Ptr, {ConstantInt::get(GEPIdx->getType(), 0), GEPIdx});
    auto *NewLoad = cast<LoadInst>(
        Builder.CreateLoad(ElemType, GEP, EI->getName() + ".scalar"));

    Align ScalarOpAlignment =
        computeAlignmentAfterScalarization(LI->getAlign(), ElemType, Idx, *DL);
    NewLoad->setAlignment(ScalarOpAlignment);

    if (auto *ConstIdx = dyn_cast<ConstantInt>(Idx)) {
      size_t Offset = ConstIdx->getZExtValue() * DL->getTypeStoreSize(ElemType);
      AAMDNodes OldAAMD = LI->getAAMetadata();
      NewLoad->setAAMetadata(OldAAMD.adjustForAccess(Offset, ElemType, *DL));
    }

    replaceValue(*EI, *NewLoad, false);
  }

  FailureGuard.release();
  return true;
}

/// Try to scalarize vector loads feeding bitcast instructions.
bool VectorCombine::scalarizeLoadBitcast(LoadInst *LI, VectorType *VecTy,
                                         Value *Ptr) {
  InstructionCost OriginalCost =
      TTI.getMemoryOpCost(Instruction::Load, VecTy, LI->getAlign(),
                          LI->getPointerAddressSpace(), CostKind);

  if (!isa<FixedVectorType>(VecTy))
    return false;

  Type *TargetScalarType = nullptr;
  unsigned VecBitWidth = DL->getTypeSizeInBits(VecTy);

  for (User *U : LI->users()) {
    auto *BC = cast<BitCastInst>(U);

    Type *DestTy = BC->getDestTy();
    if (!DestTy->isIntegerTy() && !DestTy->isFloatingPointTy())
      return false;

    unsigned DestBitWidth = DL->getTypeSizeInBits(DestTy);
    if (DestBitWidth != VecBitWidth)
      return false;

    // All bitcasts must target the same scalar type.
    if (!TargetScalarType)
      TargetScalarType = DestTy;
    else if (TargetScalarType != DestTy)
      return false;

    OriginalCost +=
        TTI.getCastInstrCost(Instruction::BitCast, TargetScalarType, VecTy,
                             TTI.getCastContextHint(BC), CostKind, BC);
  }

  if (!TargetScalarType)
    return false;

  assert(!LI->user_empty() && "Unexpected load without bitcast users");
  InstructionCost ScalarizedCost =
      TTI.getMemoryOpCost(Instruction::Load, TargetScalarType, LI->getAlign(),
                          LI->getPointerAddressSpace(), CostKind);

  LLVM_DEBUG(dbgs() << "Found vector load feeding only bitcasts: " << *LI
                    << "\n  OriginalCost: " << OriginalCost
                    << " vs ScalarizedCost: " << ScalarizedCost << "\n");

  if (ScalarizedCost >= OriginalCost)
    return false;

  // Ensure we add the load back to the worklist BEFORE its users so they can
  // erased in the correct order.
  Worklist.push(LI);

  Builder.SetInsertPoint(LI);
  auto *ScalarLoad =
      Builder.CreateLoad(TargetScalarType, Ptr, LI->getName() + ".scalar");
  ScalarLoad->setAlignment(LI->getAlign());
  ScalarLoad->copyMetadata(*LI);

  // Replace all bitcast users with the scalar load.
  for (User *U : LI->users()) {
    auto *BC = cast<BitCastInst>(U);
    replaceValue(*BC, *ScalarLoad, false);
  }

  return true;
}

bool VectorCombine::scalarizeExtExtract(Instruction &I) {
  if (!TTI.allowVectorElementIndexingUsingGEP())
    return false;
  auto *Ext = dyn_cast<ZExtInst>(&I);
  if (!Ext)
    return false;

  // Try to convert a vector zext feeding only extracts to a set of scalar
  //   (Src << ExtIdx *Size) & (Size -1)
  // if profitable   .
  auto *SrcTy = dyn_cast<FixedVectorType>(Ext->getOperand(0)->getType());
  if (!SrcTy)
    return false;
  auto *DstTy = cast<FixedVectorType>(Ext->getType());

  Type *ScalarDstTy = DstTy->getElementType();
  if (DL->getTypeSizeInBits(SrcTy) != DL->getTypeSizeInBits(ScalarDstTy))
    return false;

  InstructionCost VectorCost =
      TTI.getCastInstrCost(Instruction::ZExt, DstTy, SrcTy,
                           TTI::CastContextHint::None, CostKind, Ext);
  unsigned ExtCnt = 0;
  bool ExtLane0 = false;
  for (User *U : Ext->users()) {
    uint64_t Idx;
    if (!match(U, m_ExtractElt(m_Value(), m_ConstantInt(Idx))))
      return false;
    // An out-of-bounds extractelement produces poison; bail out rather
    // than computing a shift amount that overflows the packed type.
    if (Idx >= SrcTy->getNumElements())
      return false;
    if (cast<Instruction>(U)->use_empty())
      continue;
    ExtCnt += 1;
    ExtLane0 |= !Idx;
    VectorCost += TTI.getVectorInstrCost(Instruction::ExtractElement, DstTy,
                                         CostKind, Idx, U);
  }

  InstructionCost ScalarCost =
      ExtCnt * TTI.getArithmeticInstrCost(
                   Instruction::And, ScalarDstTy, CostKind,
                   {TTI::OK_AnyValue, TTI::OP_None},
                   {TTI::OK_NonUniformConstantValue, TTI::OP_None}) +
      (ExtCnt - ExtLane0) *
          TTI.getArithmeticInstrCost(
              Instruction::LShr, ScalarDstTy, CostKind,
              {TTI::OK_AnyValue, TTI::OP_None},
              {TTI::OK_NonUniformConstantValue, TTI::OP_None});
  if (ScalarCost > VectorCost)
    return false;

  Value *ScalarV = Ext->getOperand(0);
  if (!isGuaranteedNotToBePoison(ScalarV, SQ.AC, dyn_cast<Instruction>(ScalarV),
                                 SQ.DT)) {
    // Check wether all lanes are extracted, all extracts trigger UB
    // on poison, and the last extract (and hence all previous ones)
    // are guaranteed to execute if Ext executes.  If so, we do not
    // need to insert a freeze.
    SmallDenseSet<ConstantInt *, 8> ExtractedLanes;
    bool AllExtractsTriggerUB = true;
    ExtractElementInst *LastExtract = nullptr;
    BasicBlock *ExtBB = Ext->getParent();
    for (User *U : Ext->users()) {
      auto *Extract = cast<ExtractElementInst>(U);
      if (Extract->getParent() != ExtBB || !programUndefinedIfPoison(Extract)) {
        AllExtractsTriggerUB = false;
        break;
      }
      ExtractedLanes.insert(cast<ConstantInt>(Extract->getIndexOperand()));
      if (!LastExtract || LastExtract->comesBefore(Extract))
        LastExtract = Extract;
    }
    if (ExtractedLanes.size() != DstTy->getNumElements() ||
        !AllExtractsTriggerUB ||
        !isGuaranteedToTransferExecutionToSuccessor(Ext->getIterator(),
                                                    LastExtract->getIterator()))
      ScalarV = Builder.CreateFreeze(ScalarV);
  }
  ScalarV = Builder.CreateBitCast(
      ScalarV,
      IntegerType::get(SrcTy->getContext(), DL->getTypeSizeInBits(SrcTy)));
  uint64_t SrcEltSizeInBits = DL->getTypeSizeInBits(SrcTy->getElementType());
  uint64_t TotalBits = DL->getTypeSizeInBits(SrcTy);
  APInt EltBitMask = APInt::getLowBitsSet(TotalBits, SrcEltSizeInBits);
  Type *PackedTy = IntegerType::get(SrcTy->getContext(), TotalBits);
  Value *Mask = ConstantInt::get(PackedTy, EltBitMask);
  for (User *U : Ext->users()) {
    auto *Extract = cast<ExtractElementInst>(U);
    uint64_t Idx =
        cast<ConstantInt>(Extract->getIndexOperand())->getZExtValue();
    uint64_t ShiftAmt =
        DL->isBigEndian()
            ? (TotalBits - SrcEltSizeInBits - Idx * SrcEltSizeInBits)
            : (Idx * SrcEltSizeInBits);
    Value *LShr = Builder.CreateLShr(ScalarV, ShiftAmt);
    Value *And = Builder.CreateAnd(LShr, Mask);
    U->replaceAllUsesWith(And);
  }
  return true;
}

/// Try to fold "(or (zext (bitcast X)), (shl (zext (bitcast Y)), C))"
/// to "(bitcast (concat X, Y))"
/// where X/Y are bitcasted from i1 mask vectors.
bool VectorCombine::foldConcatOfBoolMasks(Instruction &I) {
  Type *Ty = I.getType();
  if (!Ty->isIntegerTy())
    return false;

  // TODO: Add big endian test coverage
  if (DL->isBigEndian())
    return false;

  // Restrict to disjoint cases so the mask vectors aren't overlapping.
  Instruction *X, *Y;
  if (!match(&I, m_DisjointOr(m_Instruction(X), m_Instruction(Y))))
    return false;

  // Allow both sources to contain shl, to handle more generic pattern:
  // "(or (shl (zext (bitcast X)), C1), (shl (zext (bitcast Y)), C2))"
  Value *SrcX;
  uint64_t ShAmtX = 0;
  if (!match(X, m_OneUse(m_ZExt(m_OneUse(m_BitCast(m_Value(SrcX)))))) &&
      !match(X, m_OneUse(
                    m_Shl(m_OneUse(m_ZExt(m_OneUse(m_BitCast(m_Value(SrcX))))),
                          m_ConstantInt(ShAmtX)))))
    return false;

  Value *SrcY;
  uint64_t ShAmtY = 0;
  if (!match(Y, m_OneUse(m_ZExt(m_OneUse(m_BitCast(m_Value(SrcY)))))) &&
      !match(Y, m_OneUse(
                    m_Shl(m_OneUse(m_ZExt(m_OneUse(m_BitCast(m_Value(SrcY))))),
                          m_ConstantInt(ShAmtY)))))
    return false;

  // Canonicalize larger shift to the RHS.
  if (ShAmtX > ShAmtY) {
    std::swap(X, Y);
    std::swap(SrcX, SrcY);
    std::swap(ShAmtX, ShAmtY);
  }

  // Ensure both sources are matching vXi1 bool mask types, and that the shift
  // difference is the mask width so they can be easily concatenated together.
  uint64_t ShAmtDiff = ShAmtY - ShAmtX;
  unsigned NumSHL = (ShAmtX > 0) + (ShAmtY > 0);
  unsigned BitWidth = Ty->getPrimitiveSizeInBits();
  auto *MaskTy = dyn_cast<FixedVectorType>(SrcX->getType());
  if (!MaskTy || SrcX->getType() != SrcY->getType() ||
      !MaskTy->getElementType()->isIntegerTy(1) ||
      MaskTy->getNumElements() != ShAmtDiff ||
      MaskTy->getNumElements() > (BitWidth / 2))
    return false;

  auto *ConcatTy = FixedVectorType::getDoubleElementsVectorType(MaskTy);
  auto *ConcatIntTy =
      Type::getIntNTy(Ty->getContext(), ConcatTy->getNumElements());
  auto *MaskIntTy = Type::getIntNTy(Ty->getContext(), ShAmtDiff);

  SmallVector<int, 32> ConcatMask(ConcatTy->getNumElements());
  std::iota(ConcatMask.begin(), ConcatMask.end(), 0);

  // TODO: Is it worth supporting multi use cases?
  InstructionCost OldCost = 0;
  OldCost += TTI.getArithmeticInstrCost(Instruction::Or, Ty, CostKind);
  OldCost +=
      NumSHL * TTI.getArithmeticInstrCost(Instruction::Shl, Ty, CostKind);
  OldCost += 2 * TTI.getCastInstrCost(Instruction::ZExt, Ty, MaskIntTy,
                                      TTI::CastContextHint::None, CostKind);
  OldCost += 2 * TTI.getCastInstrCost(Instruction::BitCast, MaskIntTy, MaskTy,
                                      TTI::CastContextHint::None, CostKind);

  InstructionCost NewCost = 0;
  NewCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, ConcatTy,
                                MaskTy, CostKind, ConcatMask);
  NewCost += TTI.getCastInstrCost(Instruction::BitCast, ConcatIntTy, ConcatTy,
                                  TTI::CastContextHint::None, CostKind);
  if (Ty != ConcatIntTy)
    NewCost += TTI.getCastInstrCost(Instruction::ZExt, Ty, ConcatIntTy,
                                    TTI::CastContextHint::None, CostKind);
  if (ShAmtX > 0)
    NewCost += TTI.getArithmeticInstrCost(Instruction::Shl, Ty, CostKind);

  LLVM_DEBUG(dbgs() << "Found a concatenation of bitcasted bool masks: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");

  if (NewCost > OldCost)
    return false;

  // Build bool mask concatenation, bitcast back to scalar integer, and perform
  // any residual zero-extension or shifting.
  Value *Concat = Builder.CreateShuffleVector(SrcX, SrcY, ConcatMask);
  Worklist.pushValue(Concat);

  Value *Result = Builder.CreateBitCast(Concat, ConcatIntTy);

  if (Ty != ConcatIntTy) {
    Worklist.pushValue(Result);
    Result = Builder.CreateZExt(Result, Ty);
  }

  if (ShAmtX > 0) {
    Worklist.pushValue(Result);
    Result = Builder.CreateShl(Result, ShAmtX);
  }

  replaceValue(I, *Result);
  return true;
}

/// Try to convert "shuffle (binop (shuffle, shuffle)), undef"
///           -->  "binop (shuffle), (shuffle)".
bool VectorCombine::foldPermuteOfBinops(Instruction &I) {
  BinaryOperator *BinOp;
  ArrayRef<int> OuterMask;
  if (!match(&I, m_Shuffle(m_BinOp(BinOp), m_Undef(), m_Mask(OuterMask))))
    return false;

  // Don't introduce poison into div/rem.
  if (BinOp->isIntDivRem() && llvm::is_contained(OuterMask, PoisonMaskElem))
    return false;

  Value *Op00, *Op01, *Op10, *Op11;
  ArrayRef<int> Mask0, Mask1;
  bool Match0 = match(BinOp->getOperand(0),
                      m_Shuffle(m_Value(Op00), m_Value(Op01), m_Mask(Mask0)));
  bool Match1 = match(BinOp->getOperand(1),
                      m_Shuffle(m_Value(Op10), m_Value(Op11), m_Mask(Mask1)));
  if (!Match0 && !Match1)
    return false;

  Op00 = Match0 ? Op00 : BinOp->getOperand(0);
  Op01 = Match0 ? Op01 : BinOp->getOperand(0);
  Op10 = Match1 ? Op10 : BinOp->getOperand(1);
  Op11 = Match1 ? Op11 : BinOp->getOperand(1);

  Instruction::BinaryOps Opcode = BinOp->getOpcode();
  auto *ShuffleDstTy = dyn_cast<FixedVectorType>(I.getType());
  auto *BinOpTy = dyn_cast<FixedVectorType>(BinOp->getType());
  auto *Op0Ty = dyn_cast<FixedVectorType>(Op00->getType());
  auto *Op1Ty = dyn_cast<FixedVectorType>(Op10->getType());
  if (!ShuffleDstTy || !BinOpTy || !Op0Ty || !Op1Ty)
    return false;

  unsigned NumSrcElts = BinOpTy->getNumElements();

  // Don't accept shuffles that reference the second operand in
  // div/rem or if its an undef arg.
  if ((BinOp->isIntDivRem() || !isa<PoisonValue>(I.getOperand(1))) &&
      any_of(OuterMask, [NumSrcElts](int M) { return M >= (int)NumSrcElts; }))
    return false;

  // Merge outer / inner (or identity if no match) shuffles.
  SmallVector<int> NewMask0, NewMask1;
  for (int M : OuterMask) {
    if (M < 0 || M >= (int)NumSrcElts) {
      NewMask0.push_back(PoisonMaskElem);
      NewMask1.push_back(PoisonMaskElem);
    } else {
      NewMask0.push_back(Match0 ? Mask0[M] : M);
      NewMask1.push_back(Match1 ? Mask1[M] : M);
    }
  }

  unsigned NumOpElts = Op0Ty->getNumElements();
  bool IsIdentity0 = ShuffleDstTy == Op0Ty &&
      all_of(NewMask0, [NumOpElts](int M) { return M < (int)NumOpElts; }) &&
      ShuffleVectorInst::isIdentityMask(NewMask0, NumOpElts);
  bool IsIdentity1 = ShuffleDstTy == Op1Ty &&
      all_of(NewMask1, [NumOpElts](int M) { return M < (int)NumOpElts; }) &&
      ShuffleVectorInst::isIdentityMask(NewMask1, NumOpElts);

  InstructionCost NewCost = 0;
  // Try to merge shuffles across the binop if the new shuffles are not costly.
  InstructionCost BinOpCost =
      TTI.getArithmeticInstrCost(Opcode, BinOpTy, CostKind);
  InstructionCost OldCost =
      BinOpCost + TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
                                     ShuffleDstTy, BinOpTy, CostKind, OuterMask,
                                     0, nullptr, {BinOp}, &I);
  if (!BinOp->hasOneUse())
    NewCost += BinOpCost;

  if (Match0) {
    InstructionCost Shuf0Cost = TTI.getShuffleCost(
        TargetTransformInfo::SK_PermuteTwoSrc, BinOpTy, Op0Ty, CostKind, Mask0,
        0, nullptr, {Op00, Op01}, cast<Instruction>(BinOp->getOperand(0)));
    OldCost += Shuf0Cost;
    if (!BinOp->hasOneUse() || !BinOp->getOperand(0)->hasOneUse())
      NewCost += Shuf0Cost;
  }
  if (Match1) {
    InstructionCost Shuf1Cost = TTI.getShuffleCost(
        TargetTransformInfo::SK_PermuteTwoSrc, BinOpTy, Op1Ty, CostKind, Mask1,
        0, nullptr, {Op10, Op11}, cast<Instruction>(BinOp->getOperand(1)));
    OldCost += Shuf1Cost;
    if (!BinOp->hasOneUse() || !BinOp->getOperand(1)->hasOneUse())
      NewCost += Shuf1Cost;
  }

  NewCost += TTI.getArithmeticInstrCost(Opcode, ShuffleDstTy, CostKind);

  if (!IsIdentity0)
    NewCost +=
        TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, ShuffleDstTy,
                           Op0Ty, CostKind, NewMask0, 0, nullptr, {Op00, Op01});
  if (!IsIdentity1)
    NewCost +=
        TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, ShuffleDstTy,
                           Op1Ty, CostKind, NewMask1, 0, nullptr, {Op10, Op11});

  LLVM_DEBUG(dbgs() << "Found a shuffle feeding a shuffled binop: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");

  // If costs are equal, still fold as we reduce instruction count.
  if (NewCost > OldCost)
    return false;

  Value *LHS =
      IsIdentity0 ? Op00 : Builder.CreateShuffleVector(Op00, Op01, NewMask0);
  Value *RHS =
      IsIdentity1 ? Op10 : Builder.CreateShuffleVector(Op10, Op11, NewMask1);
  Value *NewBO = Builder.CreateBinOp(Opcode, LHS, RHS);

  // Intersect flags from the old binops.
  if (auto *NewInst = dyn_cast<Instruction>(NewBO))
    NewInst->copyIRFlags(BinOp);

  Worklist.pushValue(LHS);
  Worklist.pushValue(RHS);
  replaceValue(I, *NewBO);
  return true;
}

/// Try to convert "shuffle (binop), (binop)" into "binop (shuffle), (shuffle)".
/// Try to convert "shuffle (cmpop), (cmpop)" into "cmpop (shuffle), (shuffle)".
bool VectorCombine::foldShuffleOfBinops(Instruction &I) {
  ArrayRef<int> OldMask;
  Instruction *LHS, *RHS;
  if (!match(&I, m_Shuffle(m_Instruction(LHS), m_Instruction(RHS),
                           m_Mask(OldMask))))
    return false;

  // TODO: Add support for addlike etc.
  if (LHS->getOpcode() != RHS->getOpcode())
    return false;

  Value *X, *Y, *Z, *W;
  bool IsCommutative = false;
  CmpPredicate PredLHS = CmpInst::BAD_ICMP_PREDICATE;
  CmpPredicate PredRHS = CmpInst::BAD_ICMP_PREDICATE;
  if (match(LHS, m_BinOp(m_Value(X), m_Value(Y))) &&
      match(RHS, m_BinOp(m_Value(Z), m_Value(W)))) {
    auto *BO = cast<BinaryOperator>(LHS);
    // Don't introduce poison into div/rem.
    if (llvm::is_contained(OldMask, PoisonMaskElem) && BO->isIntDivRem())
      return false;
    IsCommutative = BinaryOperator::isCommutative(BO->getOpcode());
  } else if (match(LHS, m_Cmp(PredLHS, m_Value(X), m_Value(Y))) &&
             match(RHS, m_Cmp(PredRHS, m_Value(Z), m_Value(W))) &&
             (CmpInst::Predicate)PredLHS == (CmpInst::Predicate)PredRHS) {
    IsCommutative = cast<CmpInst>(LHS)->isCommutative();
  } else
    return false;

  auto *ShuffleDstTy = dyn_cast<FixedVectorType>(I.getType());
  auto *BinResTy = dyn_cast<FixedVectorType>(LHS->getType());
  auto *BinOpTy = dyn_cast<FixedVectorType>(X->getType());
  if (!ShuffleDstTy || !BinResTy || !BinOpTy || X->getType() != Z->getType())
    return false;

  bool SameBinOp = LHS == RHS;
  unsigned NumSrcElts = BinOpTy->getNumElements();

  // If we have something like "add X, Y" and "add Z, X", swap ops to match.
  if (IsCommutative && X != Z && Y != W && (X == W || Y == Z))
    std::swap(X, Y);

  auto ConvertToUnary = [NumSrcElts](int &M) {
    if (M >= (int)NumSrcElts)
      M -= NumSrcElts;
  };

  SmallVector<int> NewMask0(OldMask);
  TargetTransformInfo::ShuffleKind SK0 = TargetTransformInfo::SK_PermuteTwoSrc;
  TTI::OperandValueInfo Op0Info = TTI.commonOperandInfo(X, Z);
  if (X == Z) {
    llvm::for_each(NewMask0, ConvertToUnary);
    SK0 = TargetTransformInfo::SK_PermuteSingleSrc;
    Z = PoisonValue::get(BinOpTy);
  }

  SmallVector<int> NewMask1(OldMask);
  TargetTransformInfo::ShuffleKind SK1 = TargetTransformInfo::SK_PermuteTwoSrc;
  TTI::OperandValueInfo Op1Info = TTI.commonOperandInfo(Y, W);
  if (Y == W) {
    llvm::for_each(NewMask1, ConvertToUnary);
    SK1 = TargetTransformInfo::SK_PermuteSingleSrc;
    W = PoisonValue::get(BinOpTy);
  }

  // Try to replace a binop with a shuffle if the shuffle is not costly.
  // When SameBinOp, only count the binop cost once.
  InstructionCost LHSCost = TTI.getInstructionCost(LHS, CostKind);
  InstructionCost RHSCost = TTI.getInstructionCost(RHS, CostKind);

  InstructionCost OldCost = LHSCost;
  if (!SameBinOp) {
    OldCost += RHSCost;
  }
  OldCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc,
                                ShuffleDstTy, BinResTy, CostKind, OldMask, 0,
                                nullptr, {LHS, RHS}, &I);

  // Handle shuffle(binop(shuffle(x),y),binop(z,shuffle(w))) style patterns
  // where one use shuffles have gotten split across the binop/cmp. These
  // often allow a major reduction in total cost that wouldn't happen as
  // individual folds.
  auto MergeInner = [&](Value *&Op, int Offset, MutableArrayRef<int> Mask,
                        TTI::TargetCostKind CostKind) -> bool {
    Value *InnerOp;
    ArrayRef<int> InnerMask;
    if (match(Op, m_OneUse(m_Shuffle(m_Value(InnerOp), m_Undef(),
                                     m_Mask(InnerMask)))) &&
        InnerOp->getType() == Op->getType() &&
        all_of(InnerMask,
               [NumSrcElts](int M) { return M < (int)NumSrcElts; })) {
      for (int &M : Mask)
        if (Offset <= M && M < (int)(Offset + NumSrcElts)) {
          M = InnerMask[M - Offset];
          M = 0 <= M ? M + Offset : M;
        }
      OldCost += TTI.getInstructionCost(cast<Instruction>(Op), CostKind);
      Op = InnerOp;
      return true;
    }
    return false;
  };
  bool ReducedInstCount = false;
  ReducedInstCount |= MergeInner(X, 0, NewMask0, CostKind);
  ReducedInstCount |= MergeInner(Y, 0, NewMask1, CostKind);
  ReducedInstCount |= MergeInner(Z, NumSrcElts, NewMask0, CostKind);
  ReducedInstCount |= MergeInner(W, NumSrcElts, NewMask1, CostKind);
  bool SingleSrcBinOp = (X == Y) && (Z == W) && (NewMask0 == NewMask1);
  // SingleSrcBinOp only reduces instruction count if we also eliminate the
  // original binop(s). If binops have multiple uses, they won't be eliminated.
  ReducedInstCount |= SingleSrcBinOp && LHS->hasOneUser() && RHS->hasOneUser();

  // For concat shuffles of i1 vectors where both binops are one-use, the
  // transform keeps the same instruction count but canonicalises to a single
  // wider binop, enabling downstream folds (e.g. NOT(XOR(concat(a,b),
  // concat(c,d))) -> XNOR(concat(a,b),concat(c,d)) on AVX-512 mask regs).
  // Restrict to BinaryOperator (not CmpInst) since narrow comparisons may
  // be cheaper than wide ones on some targets (e.g. AVX-512 vpcmpeq).
  ReducedInstCount |= cast<ShuffleVectorInst>(&I)->isConcat() &&
                      I.getType()->getScalarType()->isIntegerTy(1) &&
                      isa<BinaryOperator>(LHS) && LHS->hasOneUser() &&
                      RHS->hasOneUser();

  auto *ShuffleCmpTy =
      FixedVectorType::get(BinOpTy->getElementType(), ShuffleDstTy);
  InstructionCost NewCost = TTI.getShuffleCost(
      SK0, ShuffleCmpTy, BinOpTy, CostKind, NewMask0, 0, nullptr, {X, Z});
  if (!SingleSrcBinOp)
    NewCost += TTI.getShuffleCost(SK1, ShuffleCmpTy, BinOpTy, CostKind,
                                  NewMask1, 0, nullptr, {Y, W});

  if (PredLHS == CmpInst::BAD_ICMP_PREDICATE) {
    NewCost += TTI.getArithmeticInstrCost(LHS->getOpcode(), ShuffleDstTy,
                                          CostKind, Op0Info, Op1Info);
  } else {
    NewCost +=
        TTI.getCmpSelInstrCost(LHS->getOpcode(), ShuffleCmpTy, ShuffleDstTy,
                               PredLHS, CostKind, Op0Info, Op1Info);
  }
  // If LHS/RHS have other uses, we need to account for the cost of keeping
  // the original instructions. When SameBinOp, only add the cost once.
  if (!LHS->hasOneUser())
    NewCost += LHSCost;
  if (!SameBinOp && !RHS->hasOneUser())
    NewCost += RHSCost;

  LLVM_DEBUG(dbgs() << "Found a shuffle feeding two binops: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");

  // If either shuffle will constant fold away, then fold for the same cost as
  // we will reduce the instruction count.
  ReducedInstCount |= (isa<Constant>(X) && isa<Constant>(Z)) ||
                      (isa<Constant>(Y) && isa<Constant>(W));
  if (ReducedInstCount ? (NewCost > OldCost) : (NewCost >= OldCost))
    return false;

  Value *Shuf0 = Builder.CreateShuffleVector(X, Z, NewMask0);
  Value *Shuf1 =
      SingleSrcBinOp ? Shuf0 : Builder.CreateShuffleVector(Y, W, NewMask1);
  Value *NewBO = PredLHS == CmpInst::BAD_ICMP_PREDICATE
                     ? Builder.CreateBinOp(
                           cast<BinaryOperator>(LHS)->getOpcode(), Shuf0, Shuf1)
                     : Builder.CreateCmp(PredLHS, Shuf0, Shuf1);

  // Intersect flags from the old binops.
  if (auto *NewInst = dyn_cast<Instruction>(NewBO)) {
    NewInst->copyIRFlags(LHS);
    NewInst->andIRFlags(RHS);
  }

  Worklist.pushValue(Shuf0);
  Worklist.pushValue(Shuf1);
  replaceValue(I, *NewBO);
  return true;
}

/// Try to convert,
/// (shuffle(select(c1,t1,f1)), (select(c2,t2,f2)), m) into
/// (select (shuffle c1,c2,m), (shuffle t1,t2,m), (shuffle f1,f2,m))
bool VectorCombine::foldShuffleOfSelects(Instruction &I) {
  ArrayRef<int> Mask;
  Value *C1, *T1, *F1, *C2, *T2, *F2;
  if (!match(&I, m_Shuffle(m_Select(m_Value(C1), m_Value(T1), m_Value(F1)),
                           m_Select(m_Value(C2), m_Value(T2), m_Value(F2)),
                           m_Mask(Mask))))
    return false;

  auto *Sel1 = cast<Instruction>(I.getOperand(0));
  auto *Sel2 = cast<Instruction>(I.getOperand(1));

  auto *C1VecTy = dyn_cast<FixedVectorType>(C1->getType());
  auto *C2VecTy = dyn_cast<FixedVectorType>(C2->getType());
  if (!C1VecTy || !C2VecTy || C1VecTy != C2VecTy)
    return false;

  auto *SI0FOp = dyn_cast<FPMathOperator>(I.getOperand(0));
  auto *SI1FOp = dyn_cast<FPMathOperator>(I.getOperand(1));
  // SelectInsts must have the same FMF.
  if (((SI0FOp == nullptr) != (SI1FOp == nullptr)) ||
      ((SI0FOp != nullptr) &&
       (SI0FOp->getFastMathFlags() != SI1FOp->getFastMathFlags())))
    return false;

  auto *SrcVecTy = cast<FixedVectorType>(T1->getType());
  auto *DstVecTy = cast<FixedVectorType>(I.getType());
  auto SK = TargetTransformInfo::SK_PermuteTwoSrc;
  auto SelOp = Instruction::Select;

  InstructionCost CostSel1 = TTI.getCmpSelInstrCost(
      SelOp, SrcVecTy, C1VecTy, CmpInst::BAD_ICMP_PREDICATE, CostKind);
  InstructionCost CostSel2 = TTI.getCmpSelInstrCost(
      SelOp, SrcVecTy, C2VecTy, CmpInst::BAD_ICMP_PREDICATE, CostKind);

  InstructionCost OldCost =
      CostSel1 + CostSel2 +
      TTI.getShuffleCost(SK, DstVecTy, SrcVecTy, CostKind, Mask, 0, nullptr,
                         {I.getOperand(0), I.getOperand(1)}, &I);

  InstructionCost NewCost = TTI.getShuffleCost(
      SK, FixedVectorType::get(C1VecTy->getScalarType(), Mask.size()), C1VecTy,
      CostKind, Mask, 0, nullptr, {C1, C2});
  NewCost += TTI.getShuffleCost(SK, DstVecTy, SrcVecTy, CostKind, Mask, 0,
                                nullptr, {T1, T2});
  NewCost += TTI.getShuffleCost(SK, DstVecTy, SrcVecTy, CostKind, Mask, 0,
                                nullptr, {F1, F2});
  auto *C1C2ShuffledVecTy = FixedVectorType::get(
      Type::getInt1Ty(I.getContext()), DstVecTy->getNumElements());
  NewCost += TTI.getCmpSelInstrCost(SelOp, DstVecTy, C1C2ShuffledVecTy,
                                    CmpInst::BAD_ICMP_PREDICATE, CostKind);

  if (!Sel1->hasOneUse())
    NewCost += CostSel1;
  if (!Sel2->hasOneUse())
    NewCost += CostSel2;

  LLVM_DEBUG(dbgs() << "Found a shuffle feeding two selects: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");
  if (NewCost > OldCost)
    return false;

  Value *ShuffleCmp = Builder.CreateShuffleVector(C1, C2, Mask);
  Value *ShuffleTrue = Builder.CreateShuffleVector(T1, T2, Mask);
  Value *ShuffleFalse = Builder.CreateShuffleVector(F1, F2, Mask);
  Value *NewSel;
  // We presuppose that the SelectInsts have the same FMF.
  if (SI0FOp)
    NewSel = Builder.CreateSelectFMF(ShuffleCmp, ShuffleTrue, ShuffleFalse,
                                     SI0FOp->getFastMathFlags());
  else
    NewSel = Builder.CreateSelect(ShuffleCmp, ShuffleTrue, ShuffleFalse);

  Worklist.pushValue(ShuffleCmp);
  Worklist.pushValue(ShuffleTrue);
  Worklist.pushValue(ShuffleFalse);
  replaceValue(I, *NewSel);
  return true;
}

/// Try to convert "shuffle (castop), (castop)" with a shared castop operand
/// into "castop (shuffle)".
bool VectorCombine::foldShuffleOfCastops(Instruction &I) {
  Value *V0, *V1;
  ArrayRef<int> OldMask;
  if (!match(&I, m_Shuffle(m_Value(V0), m_Value(V1), m_Mask(OldMask))))
    return false;

  // Check whether this is a binary shuffle.
  bool IsBinaryShuffle = !isa<UndefValue>(V1);

  auto *C0 = dyn_cast<CastInst>(V0);
  auto *C1 = dyn_cast<CastInst>(V1);
  if (!C0 || (IsBinaryShuffle && !C1))
    return false;

  Instruction::CastOps Opcode = C0->getOpcode();

  // If this is allowed, foldShuffleOfCastops can get stuck in a loop
  // with foldBitcastOfShuffle. Reject in favor of foldBitcastOfShuffle.
  if (!IsBinaryShuffle && Opcode == Instruction::BitCast)
    return false;

  if (IsBinaryShuffle) {
    if (C0->getSrcTy() != C1->getSrcTy())
      return false;
    // Handle shuffle(zext_nneg(x), sext(y)) -> sext(shuffle(x,y)) folds.
    if (Opcode != C1->getOpcode()) {
      if (match(C0, m_SExtLike(m_Value())) && match(C1, m_SExtLike(m_Value())))
        Opcode = Instruction::SExt;
      else
        return false;
    }
  }

  auto *ShuffleDstTy = dyn_cast<FixedVectorType>(I.getType());
  auto *CastDstTy = dyn_cast<FixedVectorType>(C0->getDestTy());
  auto *CastSrcTy = dyn_cast<FixedVectorType>(C0->getSrcTy());
  if (!ShuffleDstTy || !CastDstTy || !CastSrcTy)
    return false;

  unsigned NumSrcElts = CastSrcTy->getNumElements();
  unsigned NumDstElts = CastDstTy->getNumElements();
  assert((NumDstElts == NumSrcElts || Opcode == Instruction::BitCast) &&
         "Only bitcasts expected to alter src/dst element counts");

  // Check for bitcasting of unscalable vector types.
  // e.g. <32 x i40> -> <40 x i32>
  if (NumDstElts != NumSrcElts && (NumSrcElts % NumDstElts) != 0 &&
      (NumDstElts % NumSrcElts) != 0)
    return false;

  SmallVector<int, 16> NewMask;
  if (NumSrcElts >= NumDstElts) {
    // The bitcast is from wide to narrow/equal elements. The shuffle mask can
    // always be expanded to the equivalent form choosing narrower elements.
    assert(NumSrcElts % NumDstElts == 0 && "Unexpected shuffle mask");
    unsigned ScaleFactor = NumSrcElts / NumDstElts;
    narrowShuffleMaskElts(ScaleFactor, OldMask, NewMask);
  } else {
    // The bitcast is from narrow elements to wide elements. The shuffle mask
    // must choose consecutive elements to allow casting first.
    assert(NumDstElts % NumSrcElts == 0 && "Unexpected shuffle mask");
    unsigned ScaleFactor = NumDstElts / NumSrcElts;
    if (!widenShuffleMaskElts(ScaleFactor, OldMask, NewMask))
      return false;
  }

  auto *NewShuffleDstTy =
      FixedVectorType::get(CastSrcTy->getScalarType(), NewMask.size());

  // Try to replace a castop with a shuffle if the shuffle is not costly.
  InstructionCost CostC0 =
      TTI.getCastInstrCost(C0->getOpcode(), CastDstTy, CastSrcTy,
                           TTI::CastContextHint::None, CostKind, C0);

  TargetTransformInfo::ShuffleKind ShuffleKind;
  if (IsBinaryShuffle)
    ShuffleKind = TargetTransformInfo::SK_PermuteTwoSrc;
  else
    ShuffleKind = TargetTransformInfo::SK_PermuteSingleSrc;

  InstructionCost OldCost = CostC0;
  OldCost += TTI.getShuffleCost(ShuffleKind, ShuffleDstTy, CastDstTy, CostKind,
                                OldMask, 0, nullptr, {}, &I);

  InstructionCost NewCost = TTI.getShuffleCost(ShuffleKind, NewShuffleDstTy,
                                               CastSrcTy, CostKind, NewMask);
  NewCost += TTI.getCastInstrCost(Opcode, ShuffleDstTy, NewShuffleDstTy,
                                  TTI::CastContextHint::None, CostKind);
  if (!C0->hasOneUse())
    NewCost += CostC0;
  if (IsBinaryShuffle) {
    InstructionCost CostC1 =
        TTI.getCastInstrCost(C1->getOpcode(), CastDstTy, CastSrcTy,
                             TTI::CastContextHint::None, CostKind, C1);
    OldCost += CostC1;
    if (!C1->hasOneUse())
      NewCost += CostC1;
  }

  LLVM_DEBUG(dbgs() << "Found a shuffle feeding two casts: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");
  if (NewCost > OldCost)
    return false;

  Value *Shuf;
  if (IsBinaryShuffle)
    Shuf = Builder.CreateShuffleVector(C0->getOperand(0), C1->getOperand(0),
                                       NewMask);
  else
    Shuf = Builder.CreateShuffleVector(C0->getOperand(0), NewMask);

  Value *Cast = Builder.CreateCast(Opcode, Shuf, ShuffleDstTy);

  // Intersect flags from the old casts.
  if (auto *NewInst = dyn_cast<Instruction>(Cast)) {
    NewInst->copyIRFlags(C0);
    if (IsBinaryShuffle)
      NewInst->andIRFlags(C1);
  }

  Worklist.pushValue(Shuf);
  replaceValue(I, *Cast);
  return true;
}

/// Try to convert any of:
/// "shuffle (shuffle x, y), (shuffle y, x)"
/// "shuffle (shuffle x, undef), (shuffle y, undef)"
/// "shuffle (shuffle x, undef), y"
/// "shuffle x, (shuffle y, undef)"
/// into "shuffle x, y".
bool VectorCombine::foldShuffleOfShuffles(Instruction &I) {
  ArrayRef<int> OuterMask;
  Value *OuterV0, *OuterV1;
  if (!match(&I,
             m_Shuffle(m_Value(OuterV0), m_Value(OuterV1), m_Mask(OuterMask))))
    return false;

  ArrayRef<int> InnerMask0, InnerMask1;
  Value *X0, *X1, *Y0, *Y1;
  bool Match0 =
      match(OuterV0, m_Shuffle(m_Value(X0), m_Value(Y0), m_Mask(InnerMask0)));
  bool Match1 =
      match(OuterV1, m_Shuffle(m_Value(X1), m_Value(Y1), m_Mask(InnerMask1)));
  if (!Match0 && !Match1)
    return false;

  // If the outer shuffle is a permute, then create a fake inner all-poison
  // shuffle. This is easier than accounting for length-changing shuffles below.
  SmallVector<int, 16> PoisonMask1;
  if (!Match1 && isa<PoisonValue>(OuterV1)) {
    X1 = X0;
    Y1 = Y0;
    PoisonMask1.append(InnerMask0.size(), PoisonMaskElem);
    InnerMask1 = PoisonMask1;
    Match1 = true; // fake match
  }

  X0 = Match0 ? X0 : OuterV0;
  Y0 = Match0 ? Y0 : OuterV0;
  X1 = Match1 ? X1 : OuterV1;
  Y1 = Match1 ? Y1 : OuterV1;
  auto *ShuffleDstTy = dyn_cast<FixedVectorType>(I.getType());
  auto *ShuffleSrcTy = dyn_cast<FixedVectorType>(X0->getType());
  auto *ShuffleImmTy = dyn_cast<FixedVectorType>(OuterV0->getType());
  if (!ShuffleDstTy || !ShuffleSrcTy || !ShuffleImmTy ||
      X0->getType() != X1->getType())
    return false;

  unsigned NumSrcElts = ShuffleSrcTy->getNumElements();
  unsigned NumImmElts = ShuffleImmTy->getNumElements();

  // Attempt to merge shuffles, matching upto 2 source operands.
  // Replace index to a poison arg with PoisonMaskElem.
  // Bail if either inner masks reference an undef arg.
  SmallVector<int, 16> NewMask(OuterMask);
  Value *NewX = nullptr, *NewY = nullptr;
  for (int &M : NewMask) {
    Value *Src = nullptr;
    if (0 <= M && M < (int)NumImmElts) {
      Src = OuterV0;
      if (Match0) {
        M = InnerMask0[M];
        Src = M >= (int)NumSrcElts ? Y0 : X0;
        M = M >= (int)NumSrcElts ? (M - NumSrcElts) : M;
      }
    } else if (M >= (int)NumImmElts) {
      Src = OuterV1;
      M -= NumImmElts;
      if (Match1) {
        M = InnerMask1[M];
        Src = M >= (int)NumSrcElts ? Y1 : X1;
        M = M >= (int)NumSrcElts ? (M - NumSrcElts) : M;
      }
    }
    if (Src && M != PoisonMaskElem) {
      assert(0 <= M && M < (int)NumSrcElts && "Unexpected shuffle mask index");
      if (isa<UndefValue>(Src)) {
        // We've referenced an undef element - if its poison, update the shuffle
        // mask, else bail.
        if (!isa<PoisonValue>(Src))
          return false;
        M = PoisonMaskElem;
        continue;
      }
      if (!NewX || NewX == Src) {
        NewX = Src;
        continue;
      }
      if (!NewY || NewY == Src) {
        M += NumSrcElts;
        NewY = Src;
        continue;
      }
      return false;
    }
  }

  if (!NewX) {
    replaceValue(I, *PoisonValue::get(ShuffleDstTy));
    return true;
  }

  if (!NewY)
    NewY = PoisonValue::get(ShuffleSrcTy);

  // Have we folded to an Identity shuffle?
  if (ShuffleVectorInst::isIdentityMask(NewMask, NumSrcElts)) {
    replaceValue(I, *NewX);
    return true;
  }

  // Try to merge the shuffles if the new shuffle is not costly.
  InstructionCost InnerCost0 = 0;
  if (Match0)
    InnerCost0 = TTI.getInstructionCost(cast<User>(OuterV0), CostKind);

  InstructionCost InnerCost1 = 0;
  if (Match1)
    InnerCost1 = TTI.getInstructionCost(cast<User>(OuterV1), CostKind);

  InstructionCost OuterCost = TTI.getInstructionCost(&I, CostKind);

  InstructionCost OldCost = InnerCost0 + InnerCost1 + OuterCost;

  bool IsUnary = all_of(NewMask, [&](int M) { return M < (int)NumSrcElts; });
  TargetTransformInfo::ShuffleKind SK =
      IsUnary ? TargetTransformInfo::SK_PermuteSingleSrc
              : TargetTransformInfo::SK_PermuteTwoSrc;
  InstructionCost NewCost =
      TTI.getShuffleCost(SK, ShuffleDstTy, ShuffleSrcTy, CostKind, NewMask, 0,
                         nullptr, {NewX, NewY});
  if (!OuterV0->hasOneUse())
    NewCost += InnerCost0;
  if (!OuterV1->hasOneUse())
    NewCost += InnerCost1;

  LLVM_DEBUG(dbgs() << "Found a shuffle feeding two shuffles: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");
  if (NewCost > OldCost)
    return false;

  Value *Shuf = Builder.CreateShuffleVector(NewX, NewY, NewMask);
  replaceValue(I, *Shuf);
  return true;
}

/// Try to convert a chain of length-preserving shuffles that are fed by
/// length-changing shuffles from the same source, e.g. a chain of length 3:
///
///   "shuffle (shuffle (shuffle x, (shuffle y, undef)),
///                                 (shuffle y, undef)),
//                                  (shuffle y, undef)"
///
/// into a single shuffle fed by a length-changing shuffle:
///
///   "shuffle x, (shuffle y, undef)"
///
/// Such chains arise e.g. from folding extract/insert sequences.
bool VectorCombine::foldShufflesOfLengthChangingShuffles(Instruction &I) {
  FixedVectorType *TrunkType = dyn_cast<FixedVectorType>(I.getType());
  if (!TrunkType)
    return false;

  unsigned ChainLength = 0;
  SmallVector<int> Mask;
  SmallVector<int> YMask;
  InstructionCost OldCost = 0;
  InstructionCost NewCost = 0;
  Value *Trunk = &I;
  unsigned NumTrunkElts = TrunkType->getNumElements();
  Value *Y = nullptr;

  for (;;) {
    // Match the current trunk against (commutations of) the pattern
    // "shuffle trunk', (shuffle y, undef)"
    ArrayRef<int> OuterMask;
    Value *OuterV0, *OuterV1;
    if (ChainLength != 0 && !Trunk->hasOneUse())
      break;
    if (!match(Trunk, m_Shuffle(m_Value(OuterV0), m_Value(OuterV1),
                                m_Mask(OuterMask))))
      break;
    if (OuterV0->getType() != TrunkType) {
      // This shuffle is not length-preserving, so it cannot be part of the
      // chain.
      break;
    }

    ArrayRef<int> InnerMask0, InnerMask1;
    Value *A0, *A1, *B0, *B1;
    bool Match0 =
        match(OuterV0, m_Shuffle(m_Value(A0), m_Value(B0), m_Mask(InnerMask0)));
    bool Match1 =
        match(OuterV1, m_Shuffle(m_Value(A1), m_Value(B1), m_Mask(InnerMask1)));
    bool Match0Leaf = Match0 && A0->getType() != I.getType();
    bool Match1Leaf = Match1 && A1->getType() != I.getType();
    if (Match0Leaf == Match1Leaf) {
      // Only handle the case of exactly one leaf in each step. The "two leaves"
      // case is handled by foldShuffleOfShuffles.
      break;
    }

    SmallVector<int> CommutedOuterMask;
    if (Match0Leaf) {
      std::swap(OuterV0, OuterV1);
      std::swap(InnerMask0, InnerMask1);
      std::swap(A0, A1);
      std::swap(B0, B1);
      llvm::append_range(CommutedOuterMask, OuterMask);
      for (int &M : CommutedOuterMask) {
        if (M == PoisonMaskElem)
          continue;
        if (M < (int)NumTrunkElts)
          M += NumTrunkElts;
        else
          M -= NumTrunkElts;
      }
      OuterMask = CommutedOuterMask;
    }
    if (!OuterV1->hasOneUse())
      break;

    if (!isa<UndefValue>(A1)) {
      if (!Y)
        Y = A1;
      else if (Y != A1)
        break;
    }
    if (!isa<UndefValue>(B1)) {
      if (!Y)
        Y = B1;
      else if (Y != B1)
        break;
    }

    auto *YType = cast<FixedVectorType>(A1->getType());
    int NumLeafElts = YType->getNumElements();
    SmallVector<int> LocalYMask(InnerMask1);
    for (int &M : LocalYMask) {
      if (M >= NumLeafElts)
        M -= NumLeafElts;
    }

    InstructionCost LocalOldCost =
        TTI.getInstructionCost(cast<User>(Trunk), CostKind) +
        TTI.getInstructionCost(cast<User>(OuterV1), CostKind);

    // Handle the initial (start of chain) case.
    if (!ChainLength) {
      Mask.assign(OuterMask);
      YMask.assign(LocalYMask);
      OldCost = NewCost = LocalOldCost;
      Trunk = OuterV0;
      ChainLength++;
      continue;
    }

    // For the non-root case, first attempt to combine masks.
    SmallVector<int> NewYMask(YMask);
    bool Valid = true;
    for (auto [CombinedM, LeafM] : llvm::zip(NewYMask, LocalYMask)) {
      if (LeafM == -1 || CombinedM == LeafM)
        continue;
      if (CombinedM == -1) {
        CombinedM = LeafM;
      } else {
        Valid = false;
        break;
      }
    }
    if (!Valid)
      break;

    SmallVector<int> NewMask;
    NewMask.reserve(NumTrunkElts);
    for (int M : Mask) {
      if (M < 0 || M >= static_cast<int>(NumTrunkElts))
        NewMask.push_back(M);
      else
        NewMask.push_back(OuterMask[M]);
    }

    // Break the chain if adding this new step complicates the shuffles such
    // that it would increase the new cost by more than the old cost of this
    // step.
    InstructionCost LocalNewCost =
        TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc, TrunkType,
                           YType, CostKind, NewYMask) +
        TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, TrunkType,
                           TrunkType, CostKind, NewMask);

    if (LocalNewCost >= NewCost && LocalOldCost < LocalNewCost - NewCost)
      break;

    LLVM_DEBUG({
      if (ChainLength == 1) {
        dbgs() << "Found chain of shuffles fed by length-changing shuffles: "
               << I << '\n';
      }
      dbgs() << "  next chain link: " << *Trunk << '\n'
             << "  old cost: " << (OldCost + LocalOldCost)
             << " new cost: " << LocalNewCost << '\n';
    });

    Mask = NewMask;
    YMask = NewYMask;
    OldCost += LocalOldCost;
    NewCost = LocalNewCost;
    Trunk = OuterV0;
    ChainLength++;
  }
  if (ChainLength <= 1)
    return false;

  // Bail out if all leaves were poison.
  if (!Y)
    return false;

  if (llvm::all_of(Mask, [&](int M) {
        return M < 0 || M >= static_cast<int>(NumTrunkElts);
      })) {
    // Produce a canonical simplified form if all elements are sourced from Y.
    for (int &M : Mask) {
      if (M >= static_cast<int>(NumTrunkElts))
        M = YMask[M - NumTrunkElts];
    }
    Value *Root =
        Builder.CreateShuffleVector(Y, PoisonValue::get(Y->getType()), Mask);
    replaceValue(I, *Root);
    return true;
  }

  Value *Leaf =
      Builder.CreateShuffleVector(Y, PoisonValue::get(Y->getType()), YMask);
  Value *Root = Builder.CreateShuffleVector(Trunk, Leaf, Mask);
  replaceValue(I, *Root);
  return true;
}

/// Try to convert
/// "shuffle (intrinsic), (intrinsic)" into "intrinsic (shuffle), (shuffle)".
bool VectorCombine::foldShuffleOfIntrinsics(Instruction &I) {
  Value *V0, *V1;
  ArrayRef<int> OldMask;
  if (!match(&I, m_Shuffle(m_Value(V0), m_Value(V1), m_Mask(OldMask))))
    return false;

  auto *II0 = dyn_cast<IntrinsicInst>(V0);
  auto *II1 = dyn_cast<IntrinsicInst>(V1);
  if (!II0 || !II1)
    return false;

  Intrinsic::ID IID = II0->getIntrinsicID();
  if (IID != II1->getIntrinsicID())
    return false;
  InstructionCost CostII0 =
      TTI.getIntrinsicInstrCost(IntrinsicCostAttributes(IID, *II0), CostKind);
  InstructionCost CostII1 =
      TTI.getIntrinsicInstrCost(IntrinsicCostAttributes(IID, *II1), CostKind);

  auto *ShuffleDstTy = dyn_cast<FixedVectorType>(I.getType());
  auto *II0Ty = dyn_cast<FixedVectorType>(II0->getType());
  if (!ShuffleDstTy || !II0Ty)
    return false;

  if (!isTriviallyVectorizable(IID))
    return false;

  for (unsigned Idx = 0, E = II0->arg_size(); Idx != E; ++Idx) {
    Value *Arg0 = II0->getArgOperand(Idx);
    Value *Arg1 = II1->getArgOperand(Idx);
    if (isVectorIntrinsicWithScalarOpAtArg(IID, Idx, &TTI)) {
      // Scalar operands must be identical.
      if (Arg0 != Arg1)
        return false;
    } else if (Arg0->getType() != Arg1->getType()) {
      // The corresponding vector operands are shuffled together, so they must
      // share the same type. For intrinsics overloaded on their operand type
      // (e.g. llvm.fptosi.sat), two calls can produce the same result type
      // from different operand types; shuffling those would be invalid.
      return false;
    }
  }

  InstructionCost OldCost =
      CostII0 + CostII1 +
      TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, ShuffleDstTy,
                         II0Ty, CostKind, OldMask, 0, nullptr, {II0, II1}, &I);

  SmallVector<Type *> NewArgsTy;
  InstructionCost NewCost = 0;
  SmallDenseSet<std::pair<Value *, Value *>> SeenOperandPairs;
  for (unsigned Idx = 0, E = II0->arg_size(); Idx != E; ++Idx) {
    if (isVectorIntrinsicWithScalarOpAtArg(IID, Idx, &TTI)) {
      NewArgsTy.push_back(II0->getArgOperand(Idx)->getType());
    } else {
      auto *VecTy = cast<FixedVectorType>(II0->getArgOperand(Idx)->getType());
      auto *ArgTy = FixedVectorType::get(VecTy->getElementType(),
                                         ShuffleDstTy->getNumElements());
      NewArgsTy.push_back(ArgTy);
      std::pair<Value *, Value *> OperandPair =
          std::make_pair(II0->getArgOperand(Idx), II1->getArgOperand(Idx));
      if (!SeenOperandPairs.insert(OperandPair).second) {
        // We've already computed the cost for this operand pair.
        continue;
      }
      NewCost += TTI.getShuffleCost(
          TargetTransformInfo::SK_PermuteTwoSrc, ArgTy, VecTy, CostKind,
          OldMask, 0, nullptr,
          {II0->getArgOperand(Idx), II1->getArgOperand(Idx)});
    }
  }
  IntrinsicCostAttributes NewAttr(IID, ShuffleDstTy, NewArgsTy);

  NewCost += TTI.getIntrinsicInstrCost(NewAttr, CostKind);
  if (!II0->hasOneUse())
    NewCost += CostII0;
  if (II1 != II0 && !II1->hasOneUse())
    NewCost += CostII1;

  LLVM_DEBUG(dbgs() << "Found a shuffle feeding two intrinsics: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");

  if (NewCost > OldCost)
    return false;

  SmallVector<Value *> NewArgs;
  SmallDenseMap<std::pair<Value *, Value *>, Value *> ShuffleCache;
  for (unsigned Idx = 0, E = II0->arg_size(); Idx != E; ++Idx) {
    if (isVectorIntrinsicWithScalarOpAtArg(IID, Idx, &TTI)) {
      NewArgs.push_back(II0->getArgOperand(Idx));
    } else {
      std::pair<Value *, Value *> OperandPair =
          std::make_pair(II0->getArgOperand(Idx), II1->getArgOperand(Idx));
      auto It = ShuffleCache.find(OperandPair);
      if (It != ShuffleCache.end()) {
        // Reuse previously created shuffle for this operand pair.
        NewArgs.push_back(It->second);
        continue;
      }
      Value *Shuf = Builder.CreateShuffleVector(
          II0->getArgOperand(Idx), II1->getArgOperand(Idx), OldMask);
      ShuffleCache[OperandPair] = Shuf;
      NewArgs.push_back(Shuf);
      Worklist.pushValue(Shuf);
    }
  }
  Value *NewIntrinsic = Builder.CreateIntrinsic(ShuffleDstTy, IID, NewArgs);

  // Intersect flags from the old intrinsics.
  if (auto *NewInst = dyn_cast<Instruction>(NewIntrinsic)) {
    NewInst->copyIRFlags(II0);
    NewInst->andIRFlags(II1);
  }

  replaceValue(I, *NewIntrinsic);
  return true;
}

/// Try to convert
/// "shuffle (intrinsic), (poison/undef)" into "intrinsic (shuffle)".
bool VectorCombine::foldPermuteOfIntrinsic(Instruction &I) {
  Value *V0;
  ArrayRef<int> Mask;
  if (!match(&I, m_Shuffle(m_Value(V0), m_Undef(), m_Mask(Mask))))
    return false;

  auto *II0 = dyn_cast<IntrinsicInst>(V0);
  if (!II0)
    return false;

  auto *ShuffleDstTy = dyn_cast<FixedVectorType>(I.getType());
  auto *IntrinsicSrcTy = dyn_cast<FixedVectorType>(II0->getType());
  if (!ShuffleDstTy || !IntrinsicSrcTy)
    return false;

  // Validate it's a pure permute, mask should only reference the first vector
  unsigned NumSrcElts = IntrinsicSrcTy->getNumElements();
  if (any_of(Mask, [NumSrcElts](int M) { return M >= (int)NumSrcElts; }))
    return false;

  Intrinsic::ID IID = II0->getIntrinsicID();
  if (!isTriviallyVectorizable(IID))
    return false;

  // Cost analysis
  InstructionCost IntrinsicCost =
      TTI.getIntrinsicInstrCost(IntrinsicCostAttributes(IID, *II0), CostKind);
  InstructionCost OldCost =
      IntrinsicCost +
      TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc, ShuffleDstTy,
                         IntrinsicSrcTy, CostKind, Mask, 0, nullptr, {V0}, &I);

  SmallVector<Type *> NewArgsTy;
  InstructionCost NewCost = 0;
  for (unsigned I = 0, E = II0->arg_size(); I != E; ++I) {
    if (isVectorIntrinsicWithScalarOpAtArg(IID, I, &TTI)) {
      NewArgsTy.push_back(II0->getArgOperand(I)->getType());
    } else {
      auto *VecTy = cast<FixedVectorType>(II0->getArgOperand(I)->getType());
      auto *ArgTy = FixedVectorType::get(VecTy->getElementType(),
                                         ShuffleDstTy->getNumElements());
      NewArgsTy.push_back(ArgTy);
      NewCost += TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
                                    ArgTy, VecTy, CostKind, Mask, 0, nullptr,
                                    {II0->getArgOperand(I)});
    }
  }
  IntrinsicCostAttributes NewAttr(IID, ShuffleDstTy, NewArgsTy);
  NewCost += TTI.getIntrinsicInstrCost(NewAttr, CostKind);

  // If the intrinsic has multiple uses, we need to account for the cost of
  // keeping the original intrinsic around.
  if (!II0->hasOneUse())
    NewCost += IntrinsicCost;

  LLVM_DEBUG(dbgs() << "Found a permute of intrinsic: " << I << "\n  OldCost: "
                    << OldCost << " vs NewCost: " << NewCost << "\n");

  if (NewCost > OldCost)
    return false;

  // Transform
  SmallVector<Value *> NewArgs;
  for (unsigned I = 0, E = II0->arg_size(); I != E; ++I) {
    if (isVectorIntrinsicWithScalarOpAtArg(IID, I, &TTI)) {
      NewArgs.push_back(II0->getArgOperand(I));
    } else {
      Value *Shuf = Builder.CreateShuffleVector(II0->getArgOperand(I), Mask);
      NewArgs.push_back(Shuf);
      Worklist.pushValue(Shuf);
    }
  }

  Value *NewIntrinsic = Builder.CreateIntrinsic(ShuffleDstTy, IID, NewArgs);

  if (auto *NewInst = dyn_cast<Instruction>(NewIntrinsic))
    NewInst->copyIRFlags(II0);

  replaceValue(I, *NewIntrinsic);
  return true;
}

using InstLane = std::pair<Value *, int>;

static InstLane lookThroughShuffles(Value *V, int Lane) {
  while (auto *SV = dyn_cast<ShuffleVectorInst>(V)) {
    unsigned NumElts =
        cast<FixedVectorType>(SV->getOperand(0)->getType())->getNumElements();
    int M = SV->getMaskValue(Lane);
    if (M < 0)
      return {nullptr, PoisonMaskElem};
    if (static_cast<unsigned>(M) < NumElts) {
      V = SV->getOperand(0);
      Lane = M;
    } else {
      V = SV->getOperand(1);
      Lane = M - NumElts;
    }
  }
  return InstLane{V, Lane};
}

static SmallVector<InstLane>
generateInstLaneVectorFromOperand(ArrayRef<InstLane> Item, int Op) {
  SmallVector<InstLane> NItem;
  for (InstLane IL : Item) {
    auto [U, Lane] = IL;
    InstLane OpLane =
        U ? lookThroughShuffles(cast<Instruction>(U)->getOperand(Op), Lane)
          : InstLane{nullptr, PoisonMaskElem};
    NItem.emplace_back(OpLane);
  }
  return NItem;
}

/// Detect concat of multiple values into a vector
static bool isFreeConcat(ArrayRef<InstLane> Item, TTI::TargetCostKind CostKind,
                         const TargetTransformInfo &TTI) {
  auto *Ty = cast<FixedVectorType>(Item.front().first->getType());
  unsigned NumElts = Ty->getNumElements();
  if (Item.size() == NumElts || NumElts == 1 || Item.size() % NumElts != 0)
    return false;

  // Check that the concat is free, usually meaning that the type will be split
  // during legalization.
  SmallVector<int, 16> ConcatMask(NumElts * 2);
  std::iota(ConcatMask.begin(), ConcatMask.end(), 0);
  if (TTI.getShuffleCost(TTI::SK_PermuteTwoSrc,
                         FixedVectorType::get(Ty->getScalarType(), NumElts * 2),
                         Ty, CostKind, ConcatMask) != 0)
    return false;

  unsigned NumSlices = Item.size() / NumElts;
  // Currently we generate a tree of shuffles for the concats, which limits us
  // to a power2.
  if (!isPowerOf2_32(NumSlices))
    return false;
  for (unsigned Slice = 0; Slice < NumSlices; ++Slice) {
    Value *SliceV = Item[Slice * NumElts].first;
    if (!SliceV || SliceV->getType() != Ty)
      return false;
    for (unsigned Elt = 0; Elt < NumElts; ++Elt) {
      auto [V, Lane] = Item[Slice * NumElts + Elt];
      if (Lane != static_cast<int>(Elt) || SliceV != V)
        return false;
    }
  }
  return true;
}

static Value *
generateNewInstTree(ArrayRef<InstLane> Item, Use *From,
                    const DenseSet<std::pair<Value *, Use *>> &IdentityLeafs,
                    const DenseSet<std::pair<Value *, Use *>> &SplatLeafs,
                    const DenseSet<std::pair<Value *, Use *>> &ConcatLeafs,
                    IRBuilderBase &Builder, InstructionWorklist &WorkList,
                    const TargetTransformInfo *TTI) {
  auto [FrontV, FrontLane] = Item.front();

  if (IdentityLeafs.contains(std::make_pair(FrontV, From))) {
    return FrontV;
  }
  if (SplatLeafs.contains(std::make_pair(FrontV, From))) {
    SmallVector<int, 16> Mask(Item.size(), FrontLane);
    return Builder.CreateShuffleVector(FrontV, Mask);
  }
  if (ConcatLeafs.contains(std::make_pair(FrontV, From))) {
    unsigned NumElts =
        cast<FixedVectorType>(FrontV->getType())->getNumElements();
    SmallVector<Value *> Values(Item.size() / NumElts, nullptr);
    for (unsigned S = 0; S < Values.size(); ++S)
      Values[S] = Item[S * NumElts].first;

    while (Values.size() > 1) {
      NumElts *= 2;
      SmallVector<int, 16> Mask(NumElts, 0);
      std::iota(Mask.begin(), Mask.end(), 0);
      SmallVector<Value *> NewValues(Values.size() / 2, nullptr);
      for (unsigned S = 0; S < NewValues.size(); ++S)
        NewValues[S] =
            Builder.CreateShuffleVector(Values[S * 2], Values[S * 2 + 1], Mask);
      Values = NewValues;
    }
    return Values[0];
  }

  auto *I = cast<Instruction>(FrontV);

  // Handle vector bitcasts that change element count. We cannot use
  // generateInstLaneVectorFromOperand for these because the lane indices
  // don't map 1:1 through the bitcast.
  if (auto *BitCast = dyn_cast<BitCastInst>(I)) {
    auto *BCDstTy = dyn_cast<FixedVectorType>(BitCast->getDestTy());
    auto *BCSrcTy = dyn_cast<FixedVectorType>(BitCast->getSrcTy());
    if (BCDstTy && BCSrcTy &&
        BCDstTy->getElementCount() != BCSrcTy->getElementCount()) {
      unsigned DstElts = BCDstTy->getNumElements();
      unsigned SrcElts = BCSrcTy->getNumElements();
      SmallVector<InstLane> NewItem;
      if (DstElts > SrcElts) {
        // Widening: compress operand Item.
        unsigned R = DstElts / SrcElts;
        if (Item.size() % R != 0)
          return nullptr;
        for (unsigned Idx = 0, E = Item.size(); Idx < E; Idx += R) {
          auto [V, Lane] = Item[Idx];
          if (!V) {
            NewItem.push_back({nullptr, PoisonMaskElem});
            continue;
          }
          NewItem.push_back(
              lookThroughShuffles(cast<Operator>(V)->getOperand(0), Lane / R));
        }
      } else {
        // Narrowing: expand operand Item.
        unsigned R = SrcElts / DstElts;
        for (auto [V, Lane] : Item) {
          if (!V) {
            NewItem.append(R, {nullptr, PoisonMaskElem});
            continue;
          }
          Value *Op = cast<Operator>(V)->getOperand(0);
          for (unsigned J = 0; J < R; ++J)
            NewItem.push_back(lookThroughShuffles(Op, Lane * R + J));
        }
      }
      Value *Op = generateNewInstTree(NewItem, &BitCast->getOperandUse(0),
                                      IdentityLeafs, SplatLeafs, ConcatLeafs,
                                      Builder, WorkList, TTI);
      WorkList.pushValue(Op);
      return Builder.CreateBitCast(
          Op, FixedVectorType::get(BCDstTy->getScalarType(), Item.size()));
    }
  }
  auto *II = dyn_cast<IntrinsicInst>(I);
  unsigned NumOps = I->getNumOperands() - (II ? 1 : 0);
  SmallVector<Value *> Ops(NumOps);
  for (unsigned Idx = 0; Idx < NumOps; Idx++) {
    if (II &&
        isVectorIntrinsicWithScalarOpAtArg(II->getIntrinsicID(), Idx, TTI)) {
      Ops[Idx] = II->getOperand(Idx);
      continue;
    }
    Ops[Idx] = generateNewInstTree(
        generateInstLaneVectorFromOperand(Item, Idx), &I->getOperandUse(Idx),
        IdentityLeafs, SplatLeafs, ConcatLeafs, Builder, WorkList, TTI);
    // Don't re-queue the operand of a bitcast we just regenerated. Doing so
    // lets foldBitcastShuffle sink the bitcast back into a shuffle(bitcast),
    // which foldShuffleToIdentity then re-matches as the same superfluous
    // identity - an infinite loop between the two folds.
    if (!isa<BitCastInst>(I))
      WorkList.pushValue(Ops[Idx]);
  }

  SmallVector<Value *, 8> ValueList;
  for (const auto &Lane : Item)
    if (Lane.first)
      ValueList.push_back(Lane.first);

  Type *DstTy =
      FixedVectorType::get(I->getType()->getScalarType(), Item.size());
  if (auto *BI = dyn_cast<BinaryOperator>(I)) {
    auto *Value = Builder.CreateBinOp((Instruction::BinaryOps)BI->getOpcode(),
                                      Ops[0], Ops[1]);
    propagateIRFlags(Value, ValueList);
    return Value;
  }
  if (auto *CI = dyn_cast<CmpInst>(I)) {
    auto *Value = Builder.CreateCmp(CI->getPredicate(), Ops[0], Ops[1]);
    propagateIRFlags(Value, ValueList);
    return Value;
  }
  if (auto *SI = dyn_cast<SelectInst>(I)) {
    auto *Value = Builder.CreateSelect(Ops[0], Ops[1], Ops[2], "", SI);
    propagateIRFlags(Value, ValueList);
    return Value;
  }
  if (auto *CI = dyn_cast<CastInst>(I)) {
    auto *Value = Builder.CreateCast(CI->getOpcode(), Ops[0], DstTy);
    propagateIRFlags(Value, ValueList);
    return Value;
  }
  if (II) {
    auto *Value = Builder.CreateIntrinsic(DstTy, II->getIntrinsicID(), Ops);
    propagateIRFlags(Value, ValueList);
    return Value;
  }
  assert(isa<UnaryInstruction>(I) && "Unexpected instruction type in Generate");
  auto *Value =
      Builder.CreateUnOp((Instruction::UnaryOps)I->getOpcode(), Ops[0]);
  propagateIRFlags(Value, ValueList);
  return Value;
}

// Starting from a shuffle, look up through operands tracking the shuffled index
// of each lane. If we can simplify away the shuffles to identities then
// do so.
bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
  auto *Ty = dyn_cast<FixedVectorType>(I.getType());
  if (!Ty || I.use_empty())
    return false;

  SmallVector<InstLane> Start(Ty->getNumElements());
  for (unsigned M = 0, E = Ty->getNumElements(); M < E; ++M)
    Start[M] = lookThroughShuffles(&I, M);

  SmallVector<std::pair<SmallVector<InstLane>, Use *>> Candidates;
  Candidates.push_back(std::make_pair(Start, &*I.use_begin()));
  DenseSet<std::pair<Value *, Use *>> IdentityLeafs, SplatLeafs, ConcatLeafs;
  unsigned NumVisited = 0;
  bool TraversedElCountChangingBitcast = false;

  while (!Candidates.empty()) {
    if (++NumVisited > MaxInstrsToScan)
      return false;

    auto ItemFrom = Candidates.pop_back_val();
    auto Item = ItemFrom.first;
    auto From = ItemFrom.second;
    auto [FrontV, FrontLane] = Item.front();

    // If we found an undef first lane then bail out to keep things simple.
    if (!FrontV)
      return false;

    // Look for an identity value.
    if (FrontLane == 0 &&
        cast<FixedVectorType>(FrontV->getType())->getNumElements() ==
            Item.size() &&
        all_of(drop_begin(enumerate(Item)), [Item](const auto &E) {
          Value *FrontV = Item.front().first;
          return !E.value().first || (isEquivBitcast(E.value().first, FrontV) &&
                                      E.value().second == (int)E.index());
        })) {
      IdentityLeafs.insert(std::make_pair(FrontV, From));
      continue;
    }
    // Look for constants, for the moment only supporting constant splats.
    if (auto *C = dyn_cast<Constant>(FrontV);
        C && C->getSplatValue() &&
        all_of(drop_begin(Item), [Item](InstLane &IL) {
          Value *FrontV = Item.front().first;
          Value *V = IL.first;
          return !V || (isa<Constant>(V) &&
                        cast<Constant>(V)->getSplatValue() ==
                            cast<Constant>(FrontV)->getSplatValue());
        })) {
      SplatLeafs.insert(std::make_pair(FrontV, From));
      continue;
    }
    // Look for a splat value.
    if (all_of(drop_begin(Item), [Item](InstLane &IL) {
          auto [FrontV, FrontLane] = Item.front();
          auto [V, Lane] = IL;
          return !V || (V == FrontV && Lane == FrontLane);
        })) {
      SplatLeafs.insert(std::make_pair(FrontV, From));
      continue;
    }

    // We need each element to be the same type of value, and check that each
    // element has a single use.
    auto CheckLaneIsEquivalentToFirst = [Item](InstLane IL) {
      Value *FrontV = Item.front().first;
      if (!IL.first)
        return true;
      Value *V = IL.first;
      if (auto *I = dyn_cast<Instruction>(V); I && !I->hasOneUser())
        return false;
      if (V->getValueID() != FrontV->getValueID())
        return false;
      if (auto *CI = dyn_cast<CmpInst>(V))
        if (CI->getPredicate() != cast<CmpInst>(FrontV)->getPredicate())
          return false;
      if (auto *CI = dyn_cast<CastInst>(V))
        if (CI->getSrcTy()->getScalarType() !=
            cast<CastInst>(FrontV)->getSrcTy()->getScalarType())
          return false;
      if (auto *SI = dyn_cast<SelectInst>(V))
        if (!isa<VectorType>(SI->getOperand(0)->getType()) ||
            SI->getOperand(0)->getType() !=
                cast<SelectInst>(FrontV)->getOperand(0)->getType())
          return false;
      if (isa<CallInst>(V) && !isa<IntrinsicInst>(V))
        return false;
      auto *II = dyn_cast<IntrinsicInst>(V);
      return !II || (isa<IntrinsicInst>(FrontV) &&
                     II->getIntrinsicID() ==
                         cast<IntrinsicInst>(FrontV)->getIntrinsicID() &&
                     !II->hasOperandBundles());
    };
    if (all_of(drop_begin(Item), CheckLaneIsEquivalentToFirst)) {
      // Check the operator is one that we support.
      if (isa<BinaryOperator, CmpInst>(FrontV)) {
        //  We exclude div/rem in case they hit UB from poison lanes.
        if (auto *BO = dyn_cast<BinaryOperator>(FrontV);
            BO && BO->isIntDivRem())
          return false;
        Candidates.emplace_back(generateInstLaneVectorFromOperand(Item, 0),
                                &cast<Instruction>(FrontV)->getOperandUse(0));
        Candidates.emplace_back(generateInstLaneVectorFromOperand(Item, 1),
                                &cast<Instruction>(FrontV)->getOperandUse(1));
        continue;
      } else if (isa<UnaryOperator, TruncInst, ZExtInst, SExtInst, FPToSIInst,
                     FPToUIInst, SIToFPInst, UIToFPInst>(FrontV)) {
        Candidates.emplace_back(generateInstLaneVectorFromOperand(Item, 0),
                                &cast<Instruction>(FrontV)->getOperandUse(0));
        continue;
      } else if (auto *BitCast = dyn_cast<BitCastInst>(FrontV)) {
        auto *BCDstTy = dyn_cast<FixedVectorType>(BitCast->getDestTy());
        auto *BCSrcTy = dyn_cast<FixedVectorType>(BitCast->getSrcTy());
        if (BCDstTy && BCSrcTy) {
          ElementCount DstEC = BCDstTy->getElementCount();
          ElementCount SrcEC = BCSrcTy->getElementCount();
          if (DstEC == SrcEC) {
            // Same element count - simple pass-through.
            Candidates.emplace_back(generateInstLaneVectorFromOperand(Item, 0),
                                    &BitCast->getOperandUse(0));
            continue;
          }
          unsigned DstElts = DstEC.getFixedValue();
          unsigned SrcElts = SrcEC.getFixedValue();
          if (DstElts > SrcElts && DstElts % SrcElts == 0) {
            // Widening bitcast (e.g. <2 x i32> -> <4 x i16>). Compress
            // consecutive groups of R destination lanes into one source
            // lane.
            unsigned R = DstElts / SrcElts;
            SmallVector<InstLane> NItem;
            bool Valid = Item.size() % R == 0;
            for (unsigned Idx = 0, E = Item.size(); Valid && Idx < E;
                 Idx += R) {
              auto [V0, L0] = Item[Idx];
              if (!V0) {
                if (any_of(ArrayRef(Item).slice(Idx + 1, R - 1),
                           [](InstLane IL) { return IL.first != nullptr; })) {
                  Valid = false;
                  break;
                }
                NItem.push_back({nullptr, PoisonMaskElem});
                continue;
              }
              if (L0 % R != 0) {
                Valid = false;
                break;
              }
              for (unsigned J = 1; J < R; ++J) {
                auto [VJ, LJ] = Item[Idx + J];
                if (!VJ || VJ != V0 || LJ != L0 + (int)J) {
                  Valid = false;
                  break;
                }
              }
              if (!Valid)
                break;
              NItem.push_back(lookThroughShuffles(
                  cast<Operator>(V0)->getOperand(0), L0 / R));
            }
            if (Valid) {
              TraversedElCountChangingBitcast = true;
              Candidates.emplace_back(NItem, &BitCast->getOperandUse(0));
              continue;
            }
          } else if (SrcElts > DstElts && SrcElts % DstElts == 0) {
            // Narrowing bitcast (e.g. <4 x i16> -> <2 x i32>). Expand
            // each destination lane into R source lanes.
            unsigned R = SrcElts / DstElts;
            SmallVector<InstLane> NItem;
            for (auto [V, Lane] : Item) {
              if (!V) {
                NItem.append(R, {nullptr, PoisonMaskElem});
                continue;
              }
              Value *Op = cast<Operator>(V)->getOperand(0);
              for (unsigned J = 0; J < R; ++J)
                NItem.push_back(lookThroughShuffles(Op, Lane * R + J));
            }
            TraversedElCountChangingBitcast = true;
            Candidates.emplace_back(NItem, &BitCast->getOperandUse(0));
            continue;
          }
        }
      } else if (auto *Sel = dyn_cast<SelectInst>(FrontV)) {
        Candidates.emplace_back(generateInstLaneVectorFromOperand(Item, 0),
                                &Sel->getOperandUse(0));
        Candidates.emplace_back(generateInstLaneVectorFromOperand(Item, 1),
                                &Sel->getOperandUse(1));
        Candidates.emplace_back(generateInstLaneVectorFromOperand(Item, 2),
                                &Sel->getOperandUse(2));
        continue;
      } else if (auto *II = dyn_cast<IntrinsicInst>(FrontV);
                 II && isTriviallyVectorizable(II->getIntrinsicID()) &&
                 !II->hasOperandBundles()) {
        for (unsigned Op = 0, E = II->getNumOperands() - 1; Op < E; Op++) {
          if (isVectorIntrinsicWithScalarOpAtArg(II->getIntrinsicID(), Op,
                                                 &TTI)) {
            if (!all_of(drop_begin(Item), [Item, Op](InstLane &IL) {
                  Value *FrontV = Item.front().first;
                  Value *V = IL.first;
                  return !V || (cast<Instruction>(V)->getOperand(Op) ==
                                cast<Instruction>(FrontV)->getOperand(Op));
                }))
              return false;
            continue;
          }
          Candidates.emplace_back(
              generateInstLaneVectorFromOperand(Item, Op),
              &cast<Instruction>(FrontV)->getOperandUse(Op));
        }
        continue;
      }
    }

    if (isFreeConcat(Item, CostKind, TTI)) {
      ConcatLeafs.insert(std::make_pair(FrontV, From));
      continue;
    }

    return false;
  }

  if (NumVisited <= 1)
    return false;

  // If the only non-leaf node traversed was a single bitcast that changes
  // element count, the fold would just commute the bitcast and shuffle.
  // foldBitcastShuffle does the reverse transform, causing an infinite loop.
  if (NumVisited == 2 && TraversedElCountChangingBitcast)
    return false;

  LLVM_DEBUG(dbgs() << "Found a superfluous identity shuffle: " << I << "\n");

  // If we got this far, we know the shuffles are superfluous and can be
  // removed. Scan through again and generate the new tree of instructions.
  Builder.SetInsertPoint(&I);
  Value *V =
      generateNewInstTree(Start, &*I.use_begin(), IdentityLeafs, SplatLeafs,
                          ConcatLeafs, Builder, Worklist, &TTI);
  replaceValue(I, *V);
  return true;
}

/// Given a commutative reduction, the order of the input lanes does not alter
/// the results. We can use this to remove certain shuffles feeding the
/// reduction, removing the need to shuffle at all.
bool VectorCombine::foldShuffleFromReductions(Instruction &I) {
  auto *II = dyn_cast<IntrinsicInst>(&I);
  if (!II)
    return false;
  switch (II->getIntrinsicID()) {
  case Intrinsic::vector_reduce_add:
  case Intrinsic::vector_reduce_mul:
  case Intrinsic::vector_reduce_and:
  case Intrinsic::vector_reduce_or:
  case Intrinsic::vector_reduce_xor:
  case Intrinsic::vector_reduce_smin:
  case Intrinsic::vector_reduce_smax:
  case Intrinsic::vector_reduce_umin:
  case Intrinsic::vector_reduce_umax:
    break;
  default:
    return false;
  }

  // Find all the inputs when looking through operations that do not alter the
  // lane order (binops, for example). Currently we look for a single shuffle,
  // and can ignore splat values.
  std::queue<Value *> Worklist;
  SmallPtrSet<Value *, 4> Visited;
  ShuffleVectorInst *Shuffle = nullptr;
  if (auto *Op = dyn_cast<Instruction>(I.getOperand(0)))
    Worklist.push(Op);

  while (!Worklist.empty()) {
    Value *CV = Worklist.front();
    Worklist.pop();
    if (Visited.contains(CV))
      continue;

    // Splats don't change the order, so can be safely ignored.
    if (isSplatValue(CV))
      continue;

    Visited.insert(CV);

    if (auto *CI = dyn_cast<Instruction>(CV)) {
      if (CI->isBinaryOp()) {
        for (auto *Op : CI->operand_values())
          Worklist.push(Op);
        continue;
      } else if (auto *SV = dyn_cast<ShuffleVectorInst>(CI)) {
        if (Shuffle && Shuffle != SV)
          return false;
        Shuffle = SV;
        continue;
      }
    }

    // Anything else is currently an unknown node.
    return false;
  }

  if (!Shuffle)
    return false;

  // Check all uses of the binary ops and shuffles are also included in the
  // lane-invariant operations (Visited should be the list of lanewise
  // instructions, including the shuffle that we found).
  for (auto *V : Visited)
    for (auto *U : V->users())
      if (!Visited.contains(U) && U != &I)
        return false;

  FixedVectorType *VecType =
      dyn_cast<FixedVectorType>(II->getOperand(0)->getType());
  if (!VecType)
    return false;
  FixedVectorType *ShuffleInputType =
      dyn_cast<FixedVectorType>(Shuffle->getOperand(0)->getType());
  if (!ShuffleInputType)
    return false;
  unsigned NumInputElts = ShuffleInputType->getNumElements();

  // Find the mask from sorting the lanes into order. This is most likely to
  // become a identity or concat mask. Undef elements are pushed to the end.
  SmallVector<int> ConcatMask;
  Shuffle->getShuffleMask(ConcatMask);
  sort(ConcatMask, [](int X, int Y) { return (unsigned)X < (unsigned)Y; });
  bool UsesSecondVec =
      any_of(ConcatMask, [&](int M) { return M >= (int)NumInputElts; });

  InstructionCost OldCost = TTI.getShuffleCost(
      UsesSecondVec ? TTI::SK_PermuteTwoSrc : TTI::SK_PermuteSingleSrc, VecType,
      ShuffleInputType, CostKind, Shuffle->getShuffleMask());
  InstructionCost NewCost = TTI.getShuffleCost(
      UsesSecondVec ? TTI::SK_PermuteTwoSrc : TTI::SK_PermuteSingleSrc, VecType,
      ShuffleInputType, CostKind, ConcatMask);

  LLVM_DEBUG(dbgs() << "Found a reduction feeding from a shuffle: " << *Shuffle
                    << "\n");
  LLVM_DEBUG(dbgs() << "  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");
  bool MadeChanges = false;
  if (NewCost < OldCost) {
    Builder.SetInsertPoint(Shuffle);
    Value *NewShuffle = Builder.CreateShuffleVector(
        Shuffle->getOperand(0), Shuffle->getOperand(1), ConcatMask);
    LLVM_DEBUG(dbgs() << "Created new shuffle: " << *NewShuffle << "\n");
    replaceValue(*Shuffle, *NewShuffle);
    return true;
  }

  // See if we can re-use foldSelectShuffle, getting it to reduce the size of
  // the shuffle into a nicer order, as it can ignore the order of the shuffles.
  MadeChanges |= foldSelectShuffle(*Shuffle, true);
  return MadeChanges;
}

/// Try to fold a chain of shuffles and ops feeding extractelement(..., 0)
/// into llvm.vector.reduce.*, by tracking which lanes contribute to the
/// extracted lane and reducing the widest vector whose lanes each contribute
/// once.
///
/// For example:
///
///   %lo = shufflevector <4 x i32> %a, poison, <2 x i32> <i32 0, i32 1>
///   %hi = shufflevector <4 x i32> %a, poison, <2 x i32> <i32 2, i32 3>
///   %s  = add <2 x i32> %lo, %hi
///   %sh = shufflevector <2 x i32> %s, poison, <2 x i32> <i32 1, i32 poison>
///   %r  = add <2 x i32> %s, %sh
///   %e  = extractelement <2 x i32> %r, i64 0
///
/// transforms to:
///
///   %e = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %a)
bool VectorCombine::foldShuffleChainsToReduce(Instruction &I) {
  Value *VecOpEE;
  if (!match(&I, m_ExtractElt(m_Value(VecOpEE), m_Zero())))
    return false;

  auto *FVT = dyn_cast<FixedVectorType>(VecOpEE->getType());
  if (!FVT)
    return false;

  if (FVT->getNumElements() < 2)
    return false;

  std::optional<Instruction::BinaryOps> CommonBinOp;
  std::optional<Intrinsic::ID> CommonCallOp;

  if (auto *BO = dyn_cast<BinaryOperator>(VecOpEE)) {
    if (!getReductionForBinop(BO->getOpcode()))
      return false;
    CommonBinOp = BO->getOpcode();
  } else if (auto *MMI = dyn_cast<MinMaxIntrinsic>(VecOpEE)) {
    CommonCallOp = MMI->getIntrinsicID();
  } else {
    return false;
  }

  // For floating-point reductions, track FMF intersection across all binops.
  FastMathFlags CommonFMF;
  bool IsFloatReduction = false;

  // A chain node is one we walk through, either a matching-opcode binop/min-max
  // or a single-source shuffle. Anything else is a leaf source.
  auto IsChainNode = [&](Value *V) {
    if (auto *BO = dyn_cast<BinaryOperator>(V))
      return CommonBinOp && BO->getOpcode() == *CommonBinOp;
    if (auto *MMI = dyn_cast<MinMaxIntrinsic>(V))
      return CommonCallOp && MMI->getIntrinsicID() == *CommonCallOp;
    if (auto *SVI = dyn_cast<ShuffleVectorInst>(V))
      return isa<PoisonValue>(SVI->getOperand(1));
    return false;
  };

  // Collect the chain, building Nodes in postorder. Bail if the chain is empty
  // or exceeds MaxChainNodes.
  constexpr unsigned MaxChainNodes = 32;
  SmallSetVector<Value *, 16> Nodes;
  SmallSetVector<Value *, 4> Sources;
  unsigned NumVisited = 0;
  auto AddSource = [&](Value *V) {
    if (!isa<FixedVectorType>(V->getType()))
      return false;
    Sources.insert(V);
    return true;
  };
  auto Walk = [&](Value *V, auto &&Walk) -> bool {
    if (Nodes.contains(V) || Sources.contains(V))
      return true;
    if (++NumVisited > MaxChainNodes)
      return false;
    if (!IsChainNode(V))
      return AddSource(V);
    // Chain shuffles always have poison as op1, so only op0 matters.
    auto *U = cast<Instruction>(V);
    unsigned NumOps = isa<ShuffleVectorInst>(U) ? 1 : 2;
    for (unsigned I = 0; I != NumOps; ++I)
      if (!Walk(U->getOperand(I), Walk))
        return false;
    if (isa<ShuffleVectorInst>(U) || Nodes.contains(U->getOperand(0)) ||
        Nodes.contains(U->getOperand(1))) {
      Nodes.insert(V);
      return true;
    }
    // Both operands are leaves so treat this binop as a source rather than
    // walking into it.
    return AddSource(V);
  };
  if (!Walk(VecOpEE, Walk) || Nodes.empty())
    return false;

  bool IsIdempotent =
      CommonCallOp || (CommonBinOp && Instruction::isIdempotent(*CommonBinOp));

  // For FP reductions, require reassoc on every binop and collect FMF.
  for (Value *V : Nodes) {
    auto *BinOp = dyn_cast<BinaryOperator>(V);
    if (!BinOp || !BinOp->getType()->isFPOrFPVectorTy())
      continue;
    if (!BinOp->hasAllowReassoc())
      return false;
    if (!IsFloatReduction) {
      CommonFMF = BinOp->getFastMathFlags();
      IsFloatReduction = true;
    } else {
      CommonFMF &= BinOp->getFastMathFlags();
    }
  }

  // Top-down demanded elements. For each chain value, track which lanes feed
  // the extracted lane 0 and which feed it more than once. Reverse postorder
  // visits every use before its value. A binop forwards its demand to both
  // operands and a shuffle follows its mask back to the source lane.
  struct Demand {
    APInt Lanes;
    APInt Duplicates;
  };
  DenseMap<Value *, Demand> Demands;
  auto DemandOf = [&](Value *V) -> Demand & {
    unsigned N = cast<FixedVectorType>(V->getType())->getNumElements();
    Demand &D = Demands[V];
    if (D.Lanes.getBitWidth() != N)
      D.Lanes = D.Duplicates = APInt::getZero(N);
    return D;
  };
  DemandOf(VecOpEE).Lanes.setBit(0);
  for (Value *V : reverse(Nodes)) {
    Demand DV = Demands.lookup(V);
    if (DV.Lanes.isZero())
      continue;
    if (auto *SVI = dyn_cast<ShuffleVectorInst>(V)) {
      ArrayRef<int> Mask = SVI->getShuffleMask();
      Demand &DS = DemandOf(SVI->getOperand(0));
      for (unsigned I = 0, E = Mask.size(); I != E; ++I) {
        // Skip lanes that are undemanded or map to poison.
        if (!DV.Lanes[I] || Mask[I] < 0 ||
            (unsigned)Mask[I] >= DS.Lanes.getBitWidth())
          continue;
        if (DS.Lanes[Mask[I]] || DV.Duplicates[I])
          DS.Duplicates.setBit(Mask[I]);
        DS.Lanes.setBit(Mask[I]);
      }
    } else {
      auto *U = cast<User>(V);
      for (Value *Op : {U->getOperand(0), U->getOperand(1)}) {
        Demand &DOp = DemandOf(Op);
        // Lanes demanded through more than one path accumulate in Duplicates.
        DOp.Duplicates |= DV.Duplicates | (DOp.Lanes & DV.Lanes);
        DOp.Lanes |= DV.Lanes;
      }
    }
  }

  // Reducing V replaces the entire chain, so every contribution to the result
  // must flow through V. Reject if anything above V reads outside the chain.
  auto CoversChain = [&](Value *V) {
    SmallVector<Value *, 8> Worklist(1, VecOpEE);
    SmallPtrSet<Value *, 8> Seen;
    Seen.insert(VecOpEE);
    while (!Worklist.empty()) {
      auto *U = cast<Instruction>(Worklist.pop_back_val());
      unsigned NumOps = isa<ShuffleVectorInst>(U) ? 1 : 2;
      for (unsigned I = 0; I != NumOps; ++I) {
        Value *Op = U->getOperand(I);
        if (Op == V || !Seen.insert(Op).second)
          continue;
        if (!Nodes.contains(Op))
          return false;
        Worklist.push_back(Op);
      }
    }
    return true;
  };

  // Reduce a single cleanly demanded source if there is one, otherwise the
  // deepest intermediate that covers the chain.
  struct ReductionCut {
    Value *Src;
    APInt Elts;
  };
  std::optional<ReductionCut> Cut;
  for (Value *S : Sources) {
    auto It = Demands.find(S);
    if (It == Demands.end() || It->second.Lanes.isZero())
      continue;
    if (!IsIdempotent && !It->second.Duplicates.isZero()) {
      Cut.reset();
      break;
    }
    if (!Cut) {
      Cut = ReductionCut{S, It->second.Lanes};
      continue;
    }
    if (!isEquivBitcast(Cut->Src, S)) {
      Cut.reset();
      break;
    }
    if (!IsIdempotent && !(Cut->Elts & It->second.Lanes).isZero()) {
      Cut.reset();
      break;
    }
    Cut->Elts |= It->second.Lanes;
  }
  if (!Cut) {
    for (Value *V : Nodes) {
      if (!isa<BinaryOperator>(V) && !isa<MinMaxIntrinsic>(V))
        continue;
      auto It = Demands.find(V);
      if (It == Demands.end() || !It->second.Lanes.isAllOnes())
        continue;
      if (!IsIdempotent && !It->second.Duplicates.isZero())
        continue;
      if (!CoversChain(V))
        continue;
      Cut = ReductionCut{V, It->second.Lanes};
      break;
    }
  }
  // Reducing one lane is just an extract and can refold forever.
  if (!Cut || Cut->Elts.popcount() < 2)
    return false;

  Intrinsic::ID ReducedOp =
      (CommonCallOp ? getMinMaxReductionIntrinsicID(*CommonCallOp)
                    : getReductionForBinop(*CommonBinOp));
  if (!ReducedOp)
    return false;

  InstructionCost OrigCost = 0;
  for (Value *V : Nodes)
    OrigCost += TTI.getInstructionCost(cast<Instruction>(V), CostKind);

  auto *SrcVT = cast<FixedVectorType>(Cut->Src->getType());
  bool IsPartialReduction = !Cut->Elts.isAllOnes();
  FixedVectorType *ReduceVecTy =
      IsPartialReduction
          ? FixedVectorType::get(FVT->getElementType(), Cut->Elts.popcount())
          : SrcVT;

  SmallVector<int> ExtractMask;
  InstructionCost NewCost = 0;
  if (IsPartialReduction) {
    for (unsigned I = 0, E = Cut->Elts.getBitWidth(); I != E; ++I)
      if (Cut->Elts[I])
        ExtractMask.push_back(I);
    unsigned SubIdx = 0, SubLen;
    auto SK = Cut->Elts.isShiftedMask(SubIdx, SubLen)
                  ? TargetTransformInfo::SK_ExtractSubvector
                  : TargetTransformInfo::SK_PermuteSingleSrc;
    NewCost += TTI.getShuffleCost(SK, ReduceVecTy, SrcVT, CostKind, ExtractMask,
                                  SubIdx, ReduceVecTy);
  }

  IntrinsicCostAttributes ICA(
      ReducedOp, ReduceVecTy->getElementType(),
      IsFloatReduction
          ? SmallVector<Type *, 2>{ReduceVecTy->getElementType(), ReduceVecTy}
          : SmallVector<Type *, 2>{ReduceVecTy},
      IsFloatReduction ? CommonFMF : FastMathFlags());
  NewCost += TTI.getIntrinsicInstrCost(ICA, CostKind);

  LLVM_DEBUG(dbgs() << "Found reduction shuffle chain: " << I << "\n OldCost : "
                    << OrigCost << " vs NewCost: " << NewCost << "\n");

  if (!OrigCost.isValid() || !NewCost.isValid())
    return false;

  if (VecOpEE->hasOneUse() ? (NewCost > OrigCost) : (NewCost >= OrigCost))
    return false;

  Value *ReduceInput = Cut->Src;
  if (IsPartialReduction)
    ReduceInput = Builder.CreateShuffleVector(Cut->Src, ExtractMask);

  Value *ReducedResult;
  if (IsFloatReduction) {
    Value *Identity = ConstantExpr::getBinOpIdentity(
        *CommonBinOp, ReduceVecTy->getElementType(), /*AllowRHSConstant=*/false,
        CommonFMF.noSignedZeros());
    ReducedResult = Builder.CreateIntrinsic(ReducedOp, {ReduceVecTy},
                                            {Identity, ReduceInput}, CommonFMF);
  } else {
    ReducedResult =
        Builder.CreateIntrinsic(ReducedOp, {ReduceVecTy}, {ReduceInput});
  }
  replaceValue(I, *ReducedResult);

  return true;
}

/// Determine if its more efficient to fold:
///   reduce(trunc(x)) -> trunc(reduce(x)).
///   reduce(sext(x))  -> sext(reduce(x)).
///   reduce(zext(x))  -> zext(reduce(x)).
bool VectorCombine::foldCastFromReductions(Instruction &I) {
  auto *II = dyn_cast<IntrinsicInst>(&I);
  if (!II)
    return false;

  bool TruncOnly = false;
  Intrinsic::ID IID = II->getIntrinsicID();
  switch (IID) {
  case Intrinsic::vector_reduce_add:
  case Intrinsic::vector_reduce_mul:
    TruncOnly = true;
    break;
  case Intrinsic::vector_reduce_and:
  case Intrinsic::vector_reduce_or:
  case Intrinsic::vector_reduce_xor:
    break;
  default:
    return false;
  }

  unsigned ReductionOpc = getArithmeticReductionInstruction(IID);
  Value *ReductionSrc = I.getOperand(0);

  Value *Src;
  if (!match(ReductionSrc, m_OneUse(m_Trunc(m_Value(Src)))) &&
      (TruncOnly || !match(ReductionSrc, m_OneUse(m_ZExtOrSExt(m_Value(Src))))))
    return false;

  auto CastOpc =
      (Instruction::CastOps)cast<Instruction>(ReductionSrc)->getOpcode();

  auto *SrcTy = cast<VectorType>(Src->getType());
  auto *ReductionSrcTy = cast<VectorType>(ReductionSrc->getType());
  Type *ResultTy = I.getType();

  InstructionCost OldCost = TTI.getArithmeticReductionCost(
      ReductionOpc, ReductionSrcTy, std::nullopt, CostKind);
  OldCost += TTI.getCastInstrCost(CastOpc, ReductionSrcTy, SrcTy,
                                  TTI::CastContextHint::None, CostKind,
                                  cast<CastInst>(ReductionSrc));
  InstructionCost NewCost =
      TTI.getArithmeticReductionCost(ReductionOpc, SrcTy, std::nullopt,
                                     CostKind) +
      TTI.getCastInstrCost(CastOpc, ResultTy, ReductionSrcTy->getScalarType(),
                           TTI::CastContextHint::None, CostKind);

  if (OldCost <= NewCost || !NewCost.isValid())
    return false;

  Value *NewReduction = Builder.CreateIntrinsic(SrcTy->getScalarType(),
                                                II->getIntrinsicID(), {Src});
  Value *NewCast = Builder.CreateCast(CastOpc, NewReduction, ResultTy);
  replaceValue(I, *NewCast);
  return true;
}

/// Fold:
///   icmp pred (reduce.{add,or,and,umax,umin}(signbit_extract(x))), C
/// into:
///   icmp sgt/slt (reduce.{or,umax,and,umin}(x)), -1/0
///
/// Sign-bit reductions produce values with known semantics:
///   - reduce.{or,umax}: 0 if no element is negative, 1 if any is
///   - reduce.{and,umin}: 1 if all elements are negative, 0 if any isn't
///   - reduce.add: count of negative elements (0 to NumElts)
///
/// Both lshr and ashr are supported:
///   - lshr produces 0 or 1, so reduce.add range is [0, N]
///   - ashr produces 0 or -1, so reduce.add range is [-N, 0]
///
/// The fold generalizes to multiple source vectors combined with the same
/// operation as the reduction. For example:
///   reduce.or(or(shr A, shr B)) conceptually extends the vector
/// For reduce.add, this changes the count to M*N where M is the number of
/// source vectors.
///
/// We transform to a direct sign check on the original vector using
/// reduce.{or,umax} or reduce.{and,umin}.
///
/// In spirit, it's similar to foldSignBitCheck in InstCombine.
bool VectorCombine::foldSignBitReductionCmp(Instruction &I) {
  CmpPredicate Pred;
  IntrinsicInst *ReduceOp;
  const APInt *CmpVal;
  if (!match(&I,
             m_ICmp(Pred, m_OneUse(m_AnyIntrinsic(ReduceOp)), m_APInt(CmpVal))))
    return false;

  Intrinsic::ID OrigIID = ReduceOp->getIntrinsicID();
  switch (OrigIID) {
  case Intrinsic::vector_reduce_or:
  case Intrinsic::vector_reduce_umax:
  case Intrinsic::vector_reduce_and:
  case Intrinsic::vector_reduce_umin:
  case Intrinsic::vector_reduce_add:
    break;
  default:
    return false;
  }

  Value *ReductionSrc = ReduceOp->getArgOperand(0);
  auto *VecTy = dyn_cast<FixedVectorType>(ReductionSrc->getType());
  if (!VecTy)
    return false;

  unsigned BitWidth = VecTy->getScalarSizeInBits();
  if (BitWidth == 1)
    return false;

  unsigned NumElts = VecTy->getNumElements();

  // Determine the expected tree opcode for multi-vector patterns.
  // The tree opcode must match the reduction's underlying operation.
  //
  // TODO: for pairs of equivalent operators, we should match both,
  //       not only the most common.
  Instruction::BinaryOps TreeOpcode;
  switch (OrigIID) {
  case Intrinsic::vector_reduce_or:
  case Intrinsic::vector_reduce_umax:
    TreeOpcode = Instruction::Or;
    break;
  case Intrinsic::vector_reduce_and:
  case Intrinsic::vector_reduce_umin:
    TreeOpcode = Instruction::And;
    break;
  case Intrinsic::vector_reduce_add:
    TreeOpcode = Instruction::Add;
    break;
  default:
    llvm_unreachable("Unexpected intrinsic");
  }

  // Collect sign-bit extraction leaves from an associative tree of TreeOpcode.
  // The tree conceptually extends the vector being reduced.
  SmallVector<Value *, 8> Worklist;
  SmallVector<Value *, 8> Sources; // Original vectors (X in shr X, BW-1)
  Worklist.push_back(ReductionSrc);
  std::optional<bool> IsAShr;
  constexpr unsigned MaxSources = 8;

  // Calculate old cost: all shifts + tree ops + reduction
  InstructionCost OldCost = TTI.getInstructionCost(ReduceOp, CostKind);

  while (!Worklist.empty() && Worklist.size() <= MaxSources &&
         Sources.size() <= MaxSources) {
    Value *V = Worklist.pop_back_val();

    // Try to match sign-bit extraction: shr X, (bitwidth-1)
    Value *X;
    if (match(V, m_OneUse(m_Shr(m_Value(X), m_SpecificInt(BitWidth - 1))))) {
      auto *Shr = cast<Instruction>(V);

      // All shifts must be the same type (all lshr or all ashr)
      bool ThisIsAShr = Shr->getOpcode() == Instruction::AShr;
      if (!IsAShr)
        IsAShr = ThisIsAShr;
      else if (*IsAShr != ThisIsAShr)
        return false;

      Sources.push_back(X);

      // As part of the fold, we remove all of the shifts, so we need to keep
      // track of their costs.
      OldCost += TTI.getInstructionCost(Shr, CostKind);

      continue;
    }

    // Try to extend through a tree node of the expected opcode
    Value *A, *B;
    if (!match(V, m_OneUse(m_BinOp(TreeOpcode, m_Value(A), m_Value(B)))))
      return false;

    // We are potentially replacing these operations as well, so we add them
    // to the costs.
    OldCost += TTI.getInstructionCost(cast<Instruction>(V), CostKind);

    Worklist.push_back(A);
    Worklist.push_back(B);
  }

  // Must have at least one source and not exceed limit
  if (Sources.empty() || Sources.size() > MaxSources ||
      Worklist.size() > MaxSources || !IsAShr)
    return false;

  unsigned NumSources = Sources.size();

  // For reduce.add, the total count must fit as a signed integer.
  // Range is [0, M*N] for lshr or [-M*N, 0] for ashr.
  if (OrigIID == Intrinsic::vector_reduce_add &&
      !isIntN(BitWidth, NumSources * NumElts))
    return false;

  // Compute the boundary value when all elements are negative:
  // - Per-element contribution: 1 for lshr, -1 for ashr
  // - For add: M*N (total elements across all sources); for others: just 1
  unsigned Count =
      (OrigIID == Intrinsic::vector_reduce_add) ? NumSources * NumElts : 1;
  APInt NegativeVal(CmpVal->getBitWidth(), Count);
  if (*IsAShr)
    NegativeVal.negate();

  // Range is [min(0, AllNegVal), max(0, AllNegVal)]
  APInt Zero = APInt::getZero(CmpVal->getBitWidth());
  APInt RangeLow = APIntOps::smin(Zero, NegativeVal);
  APInt RangeHigh = APIntOps::smax(Zero, NegativeVal);

  // Determine comparison semantics:
  // - IsEq: true for equality test, false for inequality
  // - TestsNegative: true if testing against AllNegVal, false for zero
  //
  // In addition to EQ/NE against 0 or AllNegVal, we support inequalities
  // that fold to boundary tests given the narrow value range:
  //   < RangeHigh  -> != RangeHigh
  //   > RangeHigh-1 -> == RangeHigh
  //   > RangeLow   -> != RangeLow
  //   < RangeLow+1 -> == RangeLow
  //
  // For inequalities, we work with signed predicates only. Unsigned predicates
  // are canonicalized to signed when the range is non-negative (where they are
  // equivalent). When the range includes negative values, unsigned predicates
  // would have different semantics due to wrap-around, so we reject them.
  if (!ICmpInst::isEquality(Pred) && !ICmpInst::isSigned(Pred)) {
    if (RangeLow.isNegative())
      return false;
    Pred = ICmpInst::getSignedPredicate(Pred);
  }

  bool IsEq;
  bool TestsNegative;
  if (ICmpInst::isEquality(Pred)) {
    if (CmpVal->isZero()) {
      TestsNegative = false;
    } else if (*CmpVal == NegativeVal) {
      TestsNegative = true;
    } else {
      return false;
    }
    IsEq = Pred == ICmpInst::ICMP_EQ;
  } else if (Pred == ICmpInst::ICMP_SLT && *CmpVal == RangeHigh) {
    IsEq = false;
    TestsNegative = (RangeHigh == NegativeVal);
  } else if (Pred == ICmpInst::ICMP_SGT && *CmpVal == RangeHigh - 1) {
    IsEq = true;
    TestsNegative = (RangeHigh == NegativeVal);
  } else if (Pred == ICmpInst::ICMP_SGT && *CmpVal == RangeLow) {
    IsEq = false;
    TestsNegative = (RangeLow == NegativeVal);
  } else if (Pred == ICmpInst::ICMP_SLT && *CmpVal == RangeLow + 1) {
    IsEq = true;
    TestsNegative = (RangeLow == NegativeVal);
  } else {
    return false;
  }

  // For this fold we support four types of checks:
  //
  //   1. All lanes are negative - AllNeg
  //   2. All lanes are non-negative - AllNonNeg
  //   3. At least one negative lane - AnyNeg
  //   4. At least one non-negative lane - AnyNonNeg
  //
  // For each case, we can generate the following code:
  //
  //   1. AllNeg    - reduce.and/umin(X) < 0
  //   2. AllNonNeg -  reduce.or/umax(X) > -1
  //   3. AnyNeg    -  reduce.or/umax(X) < 0
  //   4. AnyNonNeg - reduce.and/umin(X) > -1
  //
  // The table below shows the aggregation of all supported cases
  // using these four cases.
  //
  //   Reduction   | == 0      | != 0      | == MAX    | != MAX
  //   ------------+-----------+-----------+-----------+-----------
  //   or/umax     | AllNonNeg | AnyNeg    | AnyNeg    | AllNonNeg
  //   and/umin    | AnyNonNeg | AllNeg    | AllNeg    | AnyNonNeg
  //   add         | AllNonNeg | AnyNeg    | AllNeg    | AnyNonNeg
  //
  // NOTE: MAX = 1 for or/and/umax/umin, and the vector size N for add
  //
  // For easier codegen and check inversion, we use the following encoding:
  //
  //   1. Bit-3 === requires or/umax (1) or and/umin (0) check
  //   2. Bit-2 === checks < 0 (1) or > -1 (0)
  //   3. Bit-1 === universal (1) or existential (0) check
  //
  //   AnyNeg    = 0b110: uses or/umax,  checks negative, any-check
  //   AllNonNeg = 0b101: uses or/umax,  checks non-neg,  all-check
  //   AnyNonNeg = 0b000: uses and/umin, checks non-neg,  any-check
  //   AllNeg    = 0b011: uses and/umin, checks negative, all-check
  //
  // XOR with 0b011 inverts the check (swaps all/any and neg/non-neg).
  //
  enum CheckKind : unsigned {
    AnyNonNeg = 0b000,
    AllNeg = 0b011,
    AllNonNeg = 0b101,
    AnyNeg = 0b110,
  };
  // Return true if we fold this check into or/umax and false for and/umin
  auto RequiresOr = [](CheckKind C) -> bool { return C & 0b100; };
  // Return true if we should check if result is negative and false otherwise
  auto IsNegativeCheck = [](CheckKind C) -> bool { return C & 0b010; };
  // Logically invert the check
  auto Invert = [](CheckKind C) { return CheckKind(C ^ 0b011); };

  CheckKind Base;
  switch (OrigIID) {
  case Intrinsic::vector_reduce_or:
  case Intrinsic::vector_reduce_umax:
    Base = TestsNegative ? AnyNeg : AllNonNeg;
    break;
  case Intrinsic::vector_reduce_and:
  case Intrinsic::vector_reduce_umin:
    Base = TestsNegative ? AllNeg : AnyNonNeg;
    break;
  case Intrinsic::vector_reduce_add:
    Base = TestsNegative ? AllNeg : AllNonNeg;
    break;
  default:
    llvm_unreachable("Unexpected intrinsic");
  }

  CheckKind Check = IsEq ? Base : Invert(Base);

  auto PickCheaper = [&](Intrinsic::ID Arith, Intrinsic::ID MinMax) {
    InstructionCost ArithCost =
        TTI.getArithmeticReductionCost(getArithmeticReductionInstruction(Arith),
                                       VecTy, std::nullopt, CostKind);
    InstructionCost MinMaxCost =
        TTI.getMinMaxReductionCost(getMinMaxReductionIntrinsicOp(MinMax), VecTy,
                                   FastMathFlags(), CostKind);
    return ArithCost <= MinMaxCost ? std::make_pair(Arith, ArithCost)
                                   : std::make_pair(MinMax, MinMaxCost);
  };

  // Choose output reduction based on encoding's MSB
  auto [NewIID, NewCost] = RequiresOr(Check)
                               ? PickCheaper(Intrinsic::vector_reduce_or,
                                             Intrinsic::vector_reduce_umax)
                               : PickCheaper(Intrinsic::vector_reduce_and,
                                             Intrinsic::vector_reduce_umin);

  // Add cost of combining multiple sources with or/and
  if (NumSources > 1) {
    unsigned CombineOpc =
        RequiresOr(Check) ? Instruction::Or : Instruction::And;
    NewCost += TTI.getArithmeticInstrCost(CombineOpc, VecTy, CostKind) *
               (NumSources - 1);
  }

  LLVM_DEBUG(dbgs() << "Found sign-bit reduction cmp: " << I << "\n  OldCost: "
                    << OldCost << " vs NewCost: " << NewCost << "\n");

  if (NewCost > OldCost)
    return false;

  // Generate the combined input and reduction
  Builder.SetInsertPoint(&I);
  Type *ScalarTy = VecTy->getScalarType();

  Value *Input;
  if (NumSources == 1) {
    Input = Sources[0];
  } else {
    // Combine sources with or/and based on check type
    Input = RequiresOr(Check) ? Builder.CreateOr(Sources)
                              : Builder.CreateAnd(Sources);
  }

  Value *NewReduce = Builder.CreateIntrinsic(ScalarTy, NewIID, {Input});
  Value *NewCmp = IsNegativeCheck(Check) ? Builder.CreateIsNeg(NewReduce)
                                         : Builder.CreateIsNotNeg(NewReduce);
  replaceValue(I, *NewCmp);
  return true;
}

/// Fold a zero test of reduce.or or reduce.umax into a boolean reduction.
///
/// Vectorization may produce IR that compares the result of a scalar reduction
/// with zero. Depending on the target, lowering a reduction and a scalar
/// comparison separately can cost more than reducing lane-wise comparison
/// results. This fold creates the latter form only when it is not costlier.
///
/// Before:
///   %r = call iT @llvm.vector.reduce.or.vNiT(<N x iT> %x)
///   %cmp = icmp ne iT %r, 0
///
/// After:
///   %lane.cmp = icmp ne <N x iT> %x, zeroinitializer
///   %cmp = call i1 @llvm.vector.reduce.or.vNi1(<N x i1> %lane.cmp)
///
/// `reduce.or` and `reduce.umax` are non-zero when at least one lane is
/// non-zero. Therefore, `icmp ne` uses the existential `reduce.or` test.
/// Conversely, `icmp eq` must check that every lane is zero, so it uses the
/// universal `reduce.and` test.
///
/// Before:
///   %r = call iT @llvm.vector.reduce.umax.vNiT(<N x iT> %x)
///   %cmp = icmp eq iT %r, 0
///
/// After:
///   %lane.cmp = icmp eq <N x iT> %x, zeroinitializer
///   %cmp = call i1 @llvm.vector.reduce.and.vNi1(<N x i1> %lane.cmp)
bool VectorCombine::foldReductionZeroTest(Instruction &I) {
  CmpPredicate Pred;
  Value *Op;

  if (!match(&I, m_c_ICmp(Pred, m_Value(Op), m_Zero())) ||
      !ICmpInst::isEquality(Pred))
    return false;

  auto *II = dyn_cast<IntrinsicInst>(Op);
  if (!II || !II->hasOneUse())
    return false;

  auto ReduceID = II->getIntrinsicID();
  if (ReduceID != Intrinsic::vector_reduce_or &&
      ReduceID != Intrinsic::vector_reduce_umax)
    return false;

  Value *Vec = II->getArgOperand(0);
  auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType());
  if (!VecTy || !VecTy->getElementType()->isIntegerTy())
    return false;

  // Map the scalar zero test to an any-lane or all-lane boolean reduction.
  Intrinsic::ID NewIID = (Pred == ICmpInst::ICMP_NE)
                             ? Intrinsic::vector_reduce_or
                             : Intrinsic::vector_reduce_and;

  // This is not an unconditional canonicalization: compare the cost of the
  // original scalar reduction and compare with the vector compare and i1
  // reduction replacement for both reduce.or and reduce.umax.
  InstructionCost OldCost = TTI.getInstructionCost(II, CostKind) +
                            TTI.getInstructionCost(&I, CostKind);

  auto *CmpTy = cast<VectorType>(CmpInst::makeCmpResultType(VecTy));
  InstructionCost NewCost =
      TTI.getCmpSelInstrCost(Instruction::ICmp, VecTy, CmpTy, Pred, CostKind);
  NewCost += TTI.getArithmeticReductionCost(
      getArithmeticReductionInstruction(NewIID), CmpTy, std::nullopt, CostKind);

  LLVM_DEBUG(dbgs() << "Found a reduction zero test: " << I << "\n  OldCost: "
                    << OldCost << " vs NewCost: " << NewCost << "\n");

  if (!OldCost.isValid() || !NewCost.isValid() || NewCost > OldCost)
    return false;

  Builder.SetInsertPoint(&I);
  Value *NewCmp = Builder.CreateICmp(Pred, Vec, Constant::getNullValue(VecTy));
  Value *NewReduce = Builder.CreateIntrinsic(NewIID, {CmpTy}, {NewCmp});
  replaceValue(I, *NewReduce);
  return true;
}

/// vector.reduce.OP f(X_i) == 0 -> vector.reduce.OP X_i == 0
///
/// We can prove it for cases when:
///
///   1.  OP X_i == 0 <=> \forall i \in [1, N] X_i == 0
///   1'. OP X_i == 0 <=> \exists j \in [1, N] X_j == 0
///   2.  f(x) == 0 <=> x == 0
///
/// From 1 and 2 (or 1' and 2), we can infer that
///
///   OP f(X_i) == 0 <=> OP X_i == 0.
///
///                  (1)
///   OP f(X_i) == 0 <=> \forall i \in [1, N] f(X_i) == 0
///                  (2)
///                  <=> \forall i \in [1, N] X_i == 0
///                  (1)
///                  <=> OP(X_i) == 0
///
/// For some of the OP's and f's, we need to have domain constraints on X
/// to ensure properties 1 (or 1') and 2.
bool VectorCombine::foldICmpEqZeroVectorReduce(Instruction &I) {
  CmpPredicate Pred;
  Value *Op;
  if (!match(&I, m_ICmp(Pred, m_Value(Op), m_Zero())) ||
      !ICmpInst::isEquality(Pred))
    return false;

  auto *II = dyn_cast<IntrinsicInst>(Op);
  if (!II)
    return false;

  switch (II->getIntrinsicID()) {
  case Intrinsic::vector_reduce_add:
  case Intrinsic::vector_reduce_or:
  case Intrinsic::vector_reduce_umin:
  case Intrinsic::vector_reduce_umax:
  case Intrinsic::vector_reduce_smin:
  case Intrinsic::vector_reduce_smax:
    break;
  default:
    return false;
  }

  Value *InnerOp = II->getArgOperand(0);

  // TODO: fixed vector type might be too restrictive
  if (!II->hasOneUse() || !isa<FixedVectorType>(InnerOp->getType()))
    return false;

  Value *X = nullptr;

  // Check for zero-preserving operations where f(x) = 0 <=> x = 0
  //
  //   1. f(x) = shl nuw x, y for arbitrary y
  //   2. f(x) = mul nuw x, c for defined c != 0
  //   3. f(x) = zext x
  //   4. f(x) = sext x
  //   5. f(x) = neg x
  //
  if (!(match(InnerOp, m_NUWShl(m_Value(X), m_Value())) ||      // Case 1
        match(InnerOp, m_NUWMul(m_Value(X), m_NonZeroInt())) || // Case 2
        match(InnerOp, m_ZExt(m_Value(X))) ||                   // Case 3
        match(InnerOp, m_SExt(m_Value(X))) ||                   // Case 4
        match(InnerOp, m_Neg(m_Value(X)))                       // Case 5
        ))
    return false;

  SimplifyQuery S = SQ.getWithInstruction(&I);
  auto *XTy = cast<FixedVectorType>(X->getType());

  // Check for domain constraints for all supported reductions.
  //
  //  a. OR X_i   - has property 1  for every X
  //  b. UMAX X_i - has property 1  for every X
  //  c. UMIN X_i - has property 1' for every X
  //  d. SMAX X_i - has property 1  for X >= 0
  //  e. SMIN X_i - has property 1' for X >= 0
  //  f. ADD X_i  - has property 1  for X >= 0 && ADD X_i doesn't sign wrap
  //
  // In order for the proof to work, we need 1 (or 1') to be true for both
  // OP f(X_i) and OP X_i and that's why below we check constraints twice.
  //
  // NOTE: ADD X_i holds property 1 for a mirror case as well, i.e. when
  //       X <= 0 && ADD X_i doesn't sign wrap. However, due to the nature
  //       of known bits, we can't reasonably hold knowledge of "either 0
  //       or negative".
  switch (II->getIntrinsicID()) {
  case Intrinsic::vector_reduce_add: {
    // We need to check that both X_i and f(X_i) have enough leading
    // zeros to not overflow.
    KnownBits KnownX = computeKnownBits(X, S);
    KnownBits KnownFX = computeKnownBits(InnerOp, S);
    unsigned NumElems = XTy->getNumElements();
    // Adding N elements loses at most ceil(log2(N)) leading bits.
    unsigned LostBits = Log2_32_Ceil(NumElems);
    unsigned LeadingZerosX = KnownX.countMinLeadingZeros();
    unsigned LeadingZerosFX = KnownFX.countMinLeadingZeros();
    // Need at least one leading zero left after summation to ensure no overflow
    if (LeadingZerosX <= LostBits || LeadingZerosFX <= LostBits)
      return false;

    // We are not checking whether X or f(X) are positive explicitly because
    // we implicitly checked for it when we checked if both cases have enough
    // leading zeros to not wrap addition.
    break;
  }
  case Intrinsic::vector_reduce_smin:
  case Intrinsic::vector_reduce_smax:
    // Check whether X >= 0 and f(X) >= 0
    if (!isKnownNonNegative(InnerOp, S) || !isKnownNonNegative(X, S))
      return false;

    break;
  default:
    break;
  };

  LLVM_DEBUG(dbgs() << "Found a reduction to 0 comparison with removable op: "
                    << *II << "\n");

  // For zext/sext, check if the transform is profitable using cost model.
  // For other operations (shl, mul, neg), we're removing an instruction
  // while keeping the same reduction type, so it's always profitable.
  if (isa<ZExtInst>(InnerOp) || isa<SExtInst>(InnerOp)) {
    auto *FXTy = cast<FixedVectorType>(InnerOp->getType());
    Intrinsic::ID IID = II->getIntrinsicID();

    InstructionCost ExtCost = TTI.getCastInstrCost(
        cast<CastInst>(InnerOp)->getOpcode(), FXTy, XTy,
        TTI::CastContextHint::None, CostKind, cast<CastInst>(InnerOp));

    InstructionCost OldReduceCost, NewReduceCost;
    switch (IID) {
    case Intrinsic::vector_reduce_add:
    case Intrinsic::vector_reduce_or:
      OldReduceCost = TTI.getArithmeticReductionCost(
          getArithmeticReductionInstruction(IID), FXTy, std::nullopt, CostKind);
      NewReduceCost = TTI.getArithmeticReductionCost(
          getArithmeticReductionInstruction(IID), XTy, std::nullopt, CostKind);
      break;
    case Intrinsic::vector_reduce_umin:
    case Intrinsic::vector_reduce_umax:
    case Intrinsic::vector_reduce_smin:
    case Intrinsic::vector_reduce_smax:
      OldReduceCost = TTI.getMinMaxReductionCost(
          getMinMaxReductionIntrinsicOp(IID), FXTy, FastMathFlags(), CostKind);
      NewReduceCost = TTI.getMinMaxReductionCost(
          getMinMaxReductionIntrinsicOp(IID), XTy, FastMathFlags(), CostKind);
      break;
    default:
      llvm_unreachable("Unexpected reduction");
    }

    InstructionCost OldCost = OldReduceCost + ExtCost;
    InstructionCost NewCost =
        NewReduceCost + (InnerOp->hasOneUse() ? 0 : ExtCost);

    LLVM_DEBUG(dbgs() << "Found a removable extension before reduction: "
                      << *InnerOp << "\n  OldCost: " << OldCost
                      << " vs NewCost: " << NewCost << "\n");

    // We consider transformation to still be potentially beneficial even
    // when the costs are the same because we might remove a use from f(X)
    // and unlock other optimizations. Equal costs would just mean that we
    // didn't make it worse in the worst case.
    if (NewCost > OldCost)
      return false;
  }

  // Since we support zext and sext as f, we might change the scalar type
  // of the intrinsic.
  Type *Ty = XTy->getScalarType();
  Value *NewReduce = Builder.CreateIntrinsic(Ty, II->getIntrinsicID(), {X});
  Value *NewCmp =
      Builder.CreateICmp(Pred, NewReduce, ConstantInt::getNullValue(Ty));
  replaceValue(I, *NewCmp);
  return true;
}

/// Fold comparisons of reduce.or/reduce.and with reduce.umax/reduce.umin
/// based on cost, preserving the comparison semantics.
///
/// We use two fundamental properties for each pair:
///
///   1. or(X) == 0 <=> umax(X) == 0
///   2. or(X) == 1 <=> umax(X) == 1
///   3. sign(or(X)) == sign(umax(X))
///
///   1. and(X) == -1 <=> umin(X) == -1
///   2. and(X) == -2 <=> umin(X) == -2
///   3. sign(and(X)) == sign(umin(X))
///
/// From these we can infer the following transformations:
///   a. or(X) ==/!= 0   <-> umax(X) ==/!= 0
///   b. or(X) s< 0      <-> umax(X) s< 0
///   c. or(X) s> -1     <-> umax(X) s> -1
///   d. or(X) s< 1      <-> umax(X) s< 1
///   e. or(X) ==/!= 1   <-> umax(X) ==/!= 1
///   f. or(X) s< 2      <-> umax(X) s< 2
///   g. and(X) ==/!= -1 <-> umin(X) ==/!= -1
///   h. and(X) s< 0     <-> umin(X) s< 0
///   i. and(X) s> -1    <-> umin(X) s> -1
///   j. and(X) s> -2    <-> umin(X) s> -2
///   k. and(X) ==/!= -2 <-> umin(X) ==/!= -2
///   l. and(X) s> -3    <-> umin(X) s> -3
///
bool VectorCombine::foldEquivalentReductionCmp(Instruction &I) {
  CmpPredicate Pred;
  Value *ReduceOp;
  const APInt *CmpVal;
  if (!match(&I, m_ICmp(Pred, m_Value(ReduceOp), m_APInt(CmpVal))))
    return false;

  auto *II = dyn_cast<IntrinsicInst>(ReduceOp);
  if (!II || !II->hasOneUse())
    return false;

  const auto IsValidOrUmaxCmp = [&]() {
    // or === umax for i1
    if (CmpVal->getBitWidth() == 1)
      return true;

    // Cases a and e
    bool IsEquality =
        (CmpVal->isZero() || CmpVal->isOne()) && ICmpInst::isEquality(Pred);
    // Case c
    bool IsPositive = CmpVal->isAllOnes() && Pred == ICmpInst::ICMP_SGT;
    // Cases b, d, and f
    bool IsNegative = (CmpVal->isZero() || CmpVal->isOne() || *CmpVal == 2) &&
                      Pred == ICmpInst::ICMP_SLT;
    return IsEquality || IsPositive || IsNegative;
  };

  const auto IsValidAndUminCmp = [&]() {
    // and === umin for i1
    if (CmpVal->getBitWidth() == 1)
      return true;

    const auto LeadingOnes = CmpVal->countl_one();

    // Cases g and k
    bool IsEquality =
        (CmpVal->isAllOnes() || LeadingOnes + 1 == CmpVal->getBitWidth()) &&
        ICmpInst::isEquality(Pred);
    // Case h
    bool IsNegative = CmpVal->isZero() && Pred == ICmpInst::ICMP_SLT;
    // Cases i, j, and l
    bool IsPositive =
        // if the number has at least N - 2 leading ones
        // and the two LSBs are:
        //   - 1 x 1 -> -1
        //   - 1 x 0 -> -2
        //   - 0 x 1 -> -3
        LeadingOnes + 2 >= CmpVal->getBitWidth() &&
        ((*CmpVal)[0] || (*CmpVal)[1]) && Pred == ICmpInst::ICMP_SGT;
    return IsEquality || IsNegative || IsPositive;
  };

  Intrinsic::ID OriginalIID = II->getIntrinsicID();
  Intrinsic::ID AlternativeIID;

  // Check if this is a valid comparison pattern and determine the alternate
  // reduction intrinsic.
  switch (OriginalIID) {
  case Intrinsic::vector_reduce_or:
    if (!IsValidOrUmaxCmp())
      return false;
    AlternativeIID = Intrinsic::vector_reduce_umax;
    break;
  case Intrinsic::vector_reduce_umax:
    if (!IsValidOrUmaxCmp())
      return false;
    AlternativeIID = Intrinsic::vector_reduce_or;
    break;
  case Intrinsic::vector_reduce_and:
    if (!IsValidAndUminCmp())
      return false;
    AlternativeIID = Intrinsic::vector_reduce_umin;
    break;
  case Intrinsic::vector_reduce_umin:
    if (!IsValidAndUminCmp())
      return false;
    AlternativeIID = Intrinsic::vector_reduce_and;
    break;
  default:
    return false;
  }

  Value *X = II->getArgOperand(0);
  auto *VecTy = dyn_cast<FixedVectorType>(X->getType());
  if (!VecTy)
    return false;

  const auto GetReductionCost = [&](Intrinsic::ID IID) -> InstructionCost {
    unsigned ReductionOpc = getArithmeticReductionInstruction(IID);
    if (ReductionOpc != Instruction::ICmp)
      return TTI.getArithmeticReductionCost(ReductionOpc, VecTy, std::nullopt,
                                            CostKind);
    return TTI.getMinMaxReductionCost(getMinMaxReductionIntrinsicOp(IID), VecTy,
                                      FastMathFlags(), CostKind);
  };

  InstructionCost OrigCost = GetReductionCost(OriginalIID);
  InstructionCost AltCost = GetReductionCost(AlternativeIID);

  LLVM_DEBUG(dbgs() << "Found equivalent reduction cmp: " << I
                    << "\n  OrigCost: " << OrigCost
                    << " vs AltCost: " << AltCost << "\n");

  if (AltCost >= OrigCost)
    return false;

  Builder.SetInsertPoint(&I);
  Type *ScalarTy = VecTy->getScalarType();
  Value *NewReduce = Builder.CreateIntrinsic(ScalarTy, AlternativeIID, {X});
  Value *NewCmp =
      Builder.CreateICmp(Pred, NewReduce, ConstantInt::get(ScalarTy, *CmpVal));

  replaceValue(I, *NewCmp);
  return true;
}

/// Used by foldReduceAddCmpZero to check if we can prove that a value is
/// non-positive.
/// KnownBits cannot see sext <? x i1> as non-positive: each top bit equals a
/// single unknown input bit, which a per-bit lattice cannot track. The fold's
/// target shape is popcount-style sums of <N x i1> valid/invalid masks (e.g.
/// ray-intersection hits) tested for any-hit.
/// Previous attempts to approximate the known bits of such expressions were
/// using a fully recursive value tracking approach to infer a constant range
/// but ultimately turned to be too expensive in compile time.
static bool isKnownNonPositive(const Value *V, const SimplifyQuery &SQ,
                               unsigned Depth = 0) {
  constexpr unsigned MaxLocalDepth = 2;
  if (Depth > MaxLocalDepth)
    return false;

  auto NumSignBits = [&](const Value *X) {
    return ComputeNumSignBits(X, SQ.DL, SQ.AC, SQ.CtxI, SQ.DT);
  };
  if (NumSignBits(V) == V->getType()->getScalarSizeInBits())
    return true;

  Value *A, *B;
  if (match(V, m_Add(m_Value(A), m_Value(B))))
    return NumSignBits(A) >= 2 && NumSignBits(B) >= 2 &&
           isKnownNonPositive(A, SQ, Depth + 1) &&
           isKnownNonPositive(B, SQ, Depth + 1);

  return computeKnownBits(V, SQ).isNonPositive();
}

/// Fold (icmp pred (reduce.add X), 0) to (icmp pred' (reduce.or X), 0) when X
/// has lanes known to all be non-negative or all non-positive, so that
/// sum == 0 iff every lane is 0. Falls back to reduce.umax if reduce.or is
/// more expensive on the target.
bool VectorCombine::foldReduceAddCmpZero(Instruction &I) {
  CmpPredicate Pred;
  Value *Vec;
  if (!match(&I, m_ICmp(Pred,
                        m_OneUse(m_Intrinsic<Intrinsic::vector_reduce_add>(
                            m_Value(Vec))),
                        m_Zero())))
    return false;

  auto *VecTy = dyn_cast<FixedVectorType>(Vec->getType());
  if (!VecTy || VecTy->getNumElements() < 2)
    return false;

  SimplifyQuery Q = SQ.getWithInstruction(&I);
  bool IsNonNegative = isKnownNonNegative(Vec, Q);
  bool IsNonPositive = !IsNonNegative && isKnownNonPositive(Vec, Q);
  if (!IsNonNegative && !IsNonPositive)
    return false;

  // Summing NumElts lanes can consume up to log2(NumElts) sign bits. Require
  // strictly more headroom than that so the sum cannot wrap to zero.
  unsigned NumElts = VecTy->getNumElements();
  unsigned NumSignBits = ComputeNumSignBits(Vec, *DL, SQ.AC, &I, &DT);
  if (Log2_32(NumElts) >= NumSignBits)
    return false;

  ICmpInst::Predicate NewPred;
  switch (Pred) {
  case ICmpInst::ICMP_EQ:
  case ICmpInst::ICMP_ULE:
  case ICmpInst::ICMP_SLE:
  case ICmpInst::ICMP_SGE:
    NewPred = ICmpInst::ICMP_EQ;
    break;
  case ICmpInst::ICMP_NE:
  case ICmpInst::ICMP_UGT:
  case ICmpInst::ICMP_SGT:
  case ICmpInst::ICMP_SLT:
    NewPred = ICmpInst::ICMP_NE;
    break;
  default:
    return false;
  }

  // SGT and SLE on a non-positive tree, and SLT and SGE on a non-negative
  // tree, are tautologies (always true or always false). Leave those to
  // InstCombine rather than mapping them here. Remaining signed inequalities
  // also need one extra sign bit so the sum cannot flip sign.
  if (!IsNonNegative &&
      (Pred == ICmpInst::ICMP_SGT || Pred == ICmpInst::ICMP_SLE))
    return false;
  if (!IsNonPositive &&
      (Pred == ICmpInst::ICMP_SLT || Pred == ICmpInst::ICMP_SGE))
    return false;
  if ((Pred == ICmpInst::ICMP_SGT || Pred == ICmpInst::ICMP_SLE ||
       Pred == ICmpInst::ICMP_SLT || Pred == ICmpInst::ICMP_SGE) &&
      Log2_32(NumElts) >= NumSignBits - 1)
    return false;

  InstructionCost OrigCost = TTI.getArithmeticReductionCost(
      Instruction::Add, VecTy, std::nullopt, CostKind);
  InstructionCost OrCost = TTI.getArithmeticReductionCost(
      Instruction::Or, VecTy, std::nullopt, CostKind);
  InstructionCost UmaxCost = TTI.getMinMaxReductionCost(
      Intrinsic::umax, VecTy, FastMathFlags(), CostKind);
  if (!OrCost.isValid() && !UmaxCost.isValid())
    return false;
  bool UseOr = OrCost.isValid() && (!UmaxCost.isValid() || OrCost <= UmaxCost);
  InstructionCost AltCost = UseOr ? OrCost : UmaxCost;
  if (AltCost > OrigCost)
    return false;

  Builder.SetInsertPoint(&I);
  Value *NewReduce = UseOr ? Builder.CreateOrReduce(Vec)
                           : Builder.CreateIntrinsic(
                                 Intrinsic::vector_reduce_umax, {VecTy}, {Vec});
  Worklist.pushValue(NewReduce);
  Value *NewCmp = Builder.CreateICmp(
      NewPred, NewReduce, ConstantInt::getNullValue(VecTy->getScalarType()));
  replaceValue(I, *NewCmp);
  return true;
}

/// Returns true if this ShuffleVectorInst eventually feeds into a
/// vector reduction intrinsic (e.g., vector_reduce_add) by only following
/// chains of shuffles and binary operators (in any combination/order).
/// The search does not go deeper than the given Depth.
static bool feedsIntoVectorReduction(ShuffleVectorInst *SVI) {
  constexpr unsigned MaxVisited = 32;
  SmallPtrSet<Instruction *, 8> Visited;
  SmallVector<Instruction *, 4> WorkList;
  bool FoundReduction = false;

  WorkList.push_back(SVI);
  while (!WorkList.empty()) {
    Instruction *I = WorkList.pop_back_val();
    for (User *U : I->users()) {
      auto *UI = cast<Instruction>(U);
      if (!UI || !Visited.insert(UI).second)
        continue;
      if (Visited.size() > MaxVisited)
        return false;
      if (auto *II = dyn_cast<IntrinsicInst>(UI)) {
        // More than one reduction reached
        if (FoundReduction)
          return false;
        switch (II->getIntrinsicID()) {
        case Intrinsic::vector_reduce_add:
        case Intrinsic::vector_reduce_mul:
        case Intrinsic::vector_reduce_and:
        case Intrinsic::vector_reduce_or:
        case Intrinsic::vector_reduce_xor:
        case Intrinsic::vector_reduce_smin:
        case Intrinsic::vector_reduce_smax:
        case Intrinsic::vector_reduce_umin:
        case Intrinsic::vector_reduce_umax:
          FoundReduction = true;
          continue;
        default:
          return false;
        }
      }

      if (!isa<BinaryOperator>(UI) && !isa<ShuffleVectorInst>(UI))
        return false;

      WorkList.emplace_back(UI);
    }
  }
  return FoundReduction;
}

/// This method looks for groups of shuffles acting on binops, of the form:
///  %x = shuffle ...
///  %y = shuffle ...
///  %a = binop %x, %y
///  %b = binop %x, %y
///  shuffle %a, %b, selectmask
/// We may, especially if the shuffle is wider than legal, be able to convert
/// the shuffle to a form where only parts of a and b need to be computed. On
/// architectures with no obvious "select" shuffle, this can reduce the total
/// number of operations if the target reports them as cheaper.
bool VectorCombine::foldSelectShuffle(Instruction &I, bool FromReduction) {
  auto *SVI = cast<ShuffleVectorInst>(&I);
  auto *VT = cast<FixedVectorType>(I.getType());
  auto *Op0 = dyn_cast<Instruction>(SVI->getOperand(0));
  auto *Op1 = dyn_cast<Instruction>(SVI->getOperand(1));
  if (!Op0 || !Op1 || Op0 == Op1 || !Op0->isBinaryOp() || !Op1->isBinaryOp() ||
      VT != Op0->getType())
    return false;

  auto *SVI0A = dyn_cast<Instruction>(Op0->getOperand(0));
  auto *SVI0B = dyn_cast<Instruction>(Op0->getOperand(1));
  auto *SVI1A = dyn_cast<Instruction>(Op1->getOperand(0));
  auto *SVI1B = dyn_cast<Instruction>(Op1->getOperand(1));
  SmallPtrSet<Instruction *, 4> InputShuffles({SVI0A, SVI0B, SVI1A, SVI1B});
  auto checkSVNonOpUses = [&](Instruction *I) {
    if (!I || I->getOperand(0)->getType() != VT)
      return true;
    return any_of(I->users(), [&](User *U) {
      return U != Op0 && U != Op1 &&
             !(isa<ShuffleVectorInst>(U) &&
               (InputShuffles.contains(cast<Instruction>(U)) ||
                isInstructionTriviallyDead(cast<Instruction>(U))));
    });
  };
  if (checkSVNonOpUses(SVI0A) || checkSVNonOpUses(SVI0B) ||
      checkSVNonOpUses(SVI1A) || checkSVNonOpUses(SVI1B))
    return false;

  // Collect all the uses that are shuffles that we can transform together. We
  // may not have a single shuffle, but a group that can all be transformed
  // together profitably.
  SmallVector<ShuffleVectorInst *> Shuffles;
  auto collectShuffles = [&](Instruction *I) {
    for (auto *U : I->users()) {
      auto *SV = dyn_cast<ShuffleVectorInst>(U);
      if (!SV || SV->getType() != VT)
        return false;
      if ((SV->getOperand(0) != Op0 && SV->getOperand(0) != Op1) ||
          (SV->getOperand(1) != Op0 && SV->getOperand(1) != Op1))
        return false;
      if (!llvm::is_contained(Shuffles, SV))
        Shuffles.push_back(SV);
    }
    return true;
  };
  if (!collectShuffles(Op0) || !collectShuffles(Op1))
    return false;
  // From a reduction, we need to be processing a single shuffle, otherwise the
  // other uses will not be lane-invariant.
  if (FromReduction && Shuffles.size() > 1)
    return false;

  // Add any shuffle uses for the shuffles we have found, to include them in our
  // cost calculations.
  if (!FromReduction) {
    for (size_t Idx = 0, E = Shuffles.size(); Idx != E; ++Idx) {
      for (auto *U : Shuffles[Idx]->users()) {
        ShuffleVectorInst *SSV = dyn_cast<ShuffleVectorInst>(U);
        if (SSV && isa<UndefValue>(SSV->getOperand(1)) && SSV->getType() == VT)
          Shuffles.push_back(SSV);
      }
    }
  }

  // For each of the output shuffles, we try to sort all the first vector
  // elements to the beginning, followed by the second array elements at the
  // end. If the binops are legalized to smaller vectors, this may reduce total
  // number of binops. We compute the ReconstructMask mask needed to convert
  // back to the original lane order.
  SmallVector<std::pair<int, int>> V1, V2;
  SmallVector<SmallVector<int>> OrigReconstructMasks;
  int MaxV1Elt = 0, MaxV2Elt = 0;
  unsigned NumElts = VT->getNumElements();
  for (ShuffleVectorInst *SVN : Shuffles) {
    SmallVector<int> Mask;
    SVN->getShuffleMask(Mask);

    // Check the operands are the same as the original, or reversed (in which
    // case we need to commute the mask).
    Value *SVOp0 = SVN->getOperand(0);
    Value *SVOp1 = SVN->getOperand(1);
    if (isa<UndefValue>(SVOp1)) {
      auto *SSV = cast<ShuffleVectorInst>(SVOp0);
      SVOp0 = SSV->getOperand(0);
      SVOp1 = SSV->getOperand(1);
      for (int &Elem : Mask) {
        if (Elem >= static_cast<int>(SSV->getShuffleMask().size()))
          return false;
        Elem = Elem < 0 ? Elem : SSV->getMaskValue(Elem);
      }
    }
    if (SVOp0 == Op1 && SVOp1 == Op0) {
      std::swap(SVOp0, SVOp1);
      ShuffleVectorInst::commuteShuffleMask(Mask, NumElts);
    }
    if (SVOp0 != Op0 || SVOp1 != Op1)
      return false;

    // Calculate the reconstruction mask for this shuffle, as the mask needed to
    // take the packed values from Op0/Op1 and reconstructing to the original
    // order.
    SmallVector<int> ReconstructMask;
    for (unsigned I = 0; I < Mask.size(); I++) {
      if (Mask[I] < 0) {
        ReconstructMask.push_back(-1);
      } else if (Mask[I] < static_cast<int>(NumElts)) {
        MaxV1Elt = std::max(MaxV1Elt, Mask[I]);
        auto It = find_if(V1, [&](const std::pair<int, int> &A) {
          return Mask[I] == A.first;
        });
        if (It != V1.end())
          ReconstructMask.push_back(It - V1.begin());
        else {
          ReconstructMask.push_back(V1.size());
          V1.emplace_back(Mask[I], V1.size());
        }
      } else {
        MaxV2Elt = std::max<int>(MaxV2Elt, Mask[I] - NumElts);
        auto It = find_if(V2, [&](const std::pair<int, int> &A) {
          return Mask[I] - static_cast<int>(NumElts) == A.first;
        });
        if (It != V2.end())
          ReconstructMask.push_back(NumElts + It - V2.begin());
        else {
          ReconstructMask.push_back(NumElts + V2.size());
          V2.emplace_back(Mask[I] - NumElts, NumElts + V2.size());
        }
      }
    }

    // For reductions, we know that the lane ordering out doesn't alter the
    // result. In-order can help simplify the shuffle away.
    if (FromReduction)
      sort(ReconstructMask);
    OrigReconstructMasks.push_back(std::move(ReconstructMask));
  }

  // If the Maximum element used from V1 and V2 are not larger than the new
  // vectors, the vectors are already packes and performing the optimization
  // again will likely not help any further. This also prevents us from getting
  // stuck in a cycle in case the costs do not also rule it out.
  if (V1.empty() || V2.empty() ||
      (MaxV1Elt == static_cast<int>(V1.size()) - 1 &&
       MaxV2Elt == static_cast<int>(V2.size()) - 1))
    return false;

  // GetBaseMaskValue takes one of the inputs, which may either be a shuffle, a
  // shuffle of another shuffle, or not a shuffle (that is treated like a
  // identity shuffle).
  auto GetBaseMaskValue = [&](Instruction *I, int M) {
    auto *SV = dyn_cast<ShuffleVectorInst>(I);
    if (!SV)
      return M;
    if (isa<UndefValue>(SV->getOperand(1)))
      if (auto *SSV = dyn_cast<ShuffleVectorInst>(SV->getOperand(0)))
        if (InputShuffles.contains(SSV))
          return SSV->getMaskValue(SV->getMaskValue(M));
    return SV->getMaskValue(M);
  };

  // Attempt to sort the inputs my ascending mask values to make simpler input
  // shuffles and push complex shuffles down to the uses. We sort on the first
  // of the two input shuffle orders, to try and get at least one input into a
  // nice order.
  auto SortBase = [&](Instruction *A, std::pair<int, int> X,
                      std::pair<int, int> Y) {
    int MXA = GetBaseMaskValue(A, X.first);
    int MYA = GetBaseMaskValue(A, Y.first);
    return MXA < MYA;
  };
  stable_sort(V1, [&](std::pair<int, int> A, std::pair<int, int> B) {
    return SortBase(SVI0A, A, B);
  });
  stable_sort(V2, [&](std::pair<int, int> A, std::pair<int, int> B) {
    return SortBase(SVI1A, A, B);
  });
  // Calculate our ReconstructMasks from the OrigReconstructMasks and the
  // modified order of the input shuffles.
  SmallVector<SmallVector<int>> ReconstructMasks;
  for (const auto &Mask : OrigReconstructMasks) {
    SmallVector<int> ReconstructMask;
    for (int M : Mask) {
      auto FindIndex = [](const SmallVector<std::pair<int, int>> &V, int M) {
        auto It = find_if(V, [M](auto A) { return A.second == M; });
        assert(It != V.end() && "Expected all entries in Mask");
        return std::distance(V.begin(), It);
      };
      if (M < 0)
        ReconstructMask.push_back(-1);
      else if (M < static_cast<int>(NumElts)) {
        ReconstructMask.push_back(FindIndex(V1, M));
      } else {
        ReconstructMask.push_back(NumElts + FindIndex(V2, M));
      }
    }
    ReconstructMasks.push_back(std::move(ReconstructMask));
  }

  // Calculate the masks needed for the new input shuffles, which get padded
  // with undef
  SmallVector<int> V1A, V1B, V2A, V2B;
  for (unsigned I = 0; I < V1.size(); I++) {
    V1A.push_back(GetBaseMaskValue(SVI0A, V1[I].first));
    V1B.push_back(GetBaseMaskValue(SVI0B, V1[I].first));
  }
  for (unsigned I = 0; I < V2.size(); I++) {
    V2A.push_back(GetBaseMaskValue(SVI1A, V2[I].first));
    V2B.push_back(GetBaseMaskValue(SVI1B, V2[I].first));
  }
  while (V1A.size() < NumElts) {
    V1A.push_back(PoisonMaskElem);
    V1B.push_back(PoisonMaskElem);
  }
  while (V2A.size() < NumElts) {
    V2A.push_back(PoisonMaskElem);
    V2B.push_back(PoisonMaskElem);
  }

  auto AddShuffleCost = [&](InstructionCost C, Instruction *I) {
    auto *SV = dyn_cast<ShuffleVectorInst>(I);
    if (!SV)
      return C;
    return C + TTI.getShuffleCost(isa<UndefValue>(SV->getOperand(1))
                                      ? TTI::SK_PermuteSingleSrc
                                      : TTI::SK_PermuteTwoSrc,
                                  VT, VT, CostKind, SV->getShuffleMask());
  };
  auto AddShuffleMaskCost = [&](InstructionCost C, ArrayRef<int> Mask) {
    return C +
           TTI.getShuffleCost(TTI::SK_PermuteTwoSrc, VT, VT, CostKind, Mask);
  };

  unsigned ElementSize = VT->getElementType()->getPrimitiveSizeInBits();
  unsigned MaxVectorSize =
      TTI.getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector);
  unsigned MaxElementsInVector = MaxVectorSize / ElementSize;
  if (MaxElementsInVector == 0)
    return false;
  // When there are multiple shufflevector operations on the same input,
  // especially when the vector length is larger than the register size,
  // identical shuffle patterns may occur across different groups of elements.
  // To avoid overestimating the cost by counting these repeated shuffles more
  // than once, we only account for unique shuffle patterns. This adjustment
  // prevents inflated costs in the cost model for wide vectors split into
  // several register-sized groups.
  std::set<SmallVector<int, 4>> UniqueShuffles;
  auto AddShuffleMaskAdjustedCost = [&](InstructionCost C, ArrayRef<int> Mask) {
    // Compute the cost for performing the shuffle over the full vector.
    auto ShuffleCost =
        TTI.getShuffleCost(TTI::SK_PermuteTwoSrc, VT, VT, CostKind, Mask);
    unsigned NumFullVectors = Mask.size() / MaxElementsInVector;
    if (NumFullVectors < 2)
      return C + ShuffleCost;
    SmallVector<int, 4> SubShuffle(MaxElementsInVector);
    unsigned NumUniqueGroups = 0;
    unsigned NumGroups = Mask.size() / MaxElementsInVector;
    // For each group of MaxElementsInVector contiguous elements,
    // collect their shuffle pattern and insert into the set of unique patterns.
    for (unsigned I = 0; I < NumFullVectors; ++I) {
      for (unsigned J = 0; J < MaxElementsInVector; ++J)
        SubShuffle[J] = Mask[MaxElementsInVector * I + J];
      if (UniqueShuffles.insert(SubShuffle).second)
        NumUniqueGroups += 1;
    }
    return C + ShuffleCost * NumUniqueGroups / NumGroups;
  };
  auto AddShuffleAdjustedCost = [&](InstructionCost C, Instruction *I) {
    auto *SV = dyn_cast<ShuffleVectorInst>(I);
    if (!SV)
      return C;
    SmallVector<int, 16> Mask;
    SV->getShuffleMask(Mask);
    return AddShuffleMaskAdjustedCost(C, Mask);
  };
  // Check that input consists of ShuffleVectors applied to the same input
  auto AllShufflesHaveSameOperands =
      [](SmallPtrSetImpl<Instruction *> &InputShuffles) {
        if (InputShuffles.size() < 2)
          return false;
        ShuffleVectorInst *FirstSV =
            dyn_cast<ShuffleVectorInst>(*InputShuffles.begin());
        if (!FirstSV)
          return false;

        Value *In0 = FirstSV->getOperand(0), *In1 = FirstSV->getOperand(1);
        return std::all_of(
            std::next(InputShuffles.begin()), InputShuffles.end(),
            [&](Instruction *I) {
              ShuffleVectorInst *SV = dyn_cast<ShuffleVectorInst>(I);
              return SV && SV->getOperand(0) == In0 && SV->getOperand(1) == In1;
            });
      };

  // Get the costs of the shuffles + binops before and after with the new
  // shuffle masks.
  InstructionCost CostBefore =
      TTI.getArithmeticInstrCost(Op0->getOpcode(), VT, CostKind) +
      TTI.getArithmeticInstrCost(Op1->getOpcode(), VT, CostKind);
  CostBefore += std::accumulate(Shuffles.begin(), Shuffles.end(),
                                InstructionCost(0), AddShuffleCost);
  if (AllShufflesHaveSameOperands(InputShuffles)) {
    UniqueShuffles.clear();
    CostBefore += std::accumulate(InputShuffles.begin(), InputShuffles.end(),
                                  InstructionCost(0), AddShuffleAdjustedCost);
  } else {
    CostBefore += std::accumulate(InputShuffles.begin(), InputShuffles.end(),
                                  InstructionCost(0), AddShuffleCost);
  }

  // The new binops will be unused for lanes past the used shuffle lengths.
  // These types attempt to get the correct cost for that from the target.
  FixedVectorType *Op0SmallVT =
      FixedVectorType::get(VT->getScalarType(), V1.size());
  FixedVectorType *Op1SmallVT =
      FixedVectorType::get(VT->getScalarType(), V2.size());
  InstructionCost CostAfter =
      TTI.getArithmeticInstrCost(Op0->getOpcode(), Op0SmallVT, CostKind) +
      TTI.getArithmeticInstrCost(Op1->getOpcode(), Op1SmallVT, CostKind);
  UniqueShuffles.clear();
  CostAfter += std::accumulate(ReconstructMasks.begin(), ReconstructMasks.end(),
                               InstructionCost(0), AddShuffleMaskAdjustedCost);
  std::set<SmallVector<int>> OutputShuffleMasks({V1A, V1B, V2A, V2B});
  CostAfter +=
      std::accumulate(OutputShuffleMasks.begin(), OutputShuffleMasks.end(),
                      InstructionCost(0), AddShuffleMaskCost);

  LLVM_DEBUG(dbgs() << "Found a binop select shuffle pattern: " << I << "\n");
  LLVM_DEBUG(dbgs() << "  CostBefore: " << CostBefore
                    << " vs CostAfter: " << CostAfter << "\n");
  if (CostBefore < CostAfter ||
      (CostBefore == CostAfter && !feedsIntoVectorReduction(SVI)))
    return false;

  // The cost model has passed, create the new instructions.
  auto GetShuffleOperand = [&](Instruction *I, unsigned Op) -> Value * {
    auto *SV = dyn_cast<ShuffleVectorInst>(I);
    if (!SV)
      return I;
    if (isa<UndefValue>(SV->getOperand(1)))
      if (auto *SSV = dyn_cast<ShuffleVectorInst>(SV->getOperand(0)))
        if (InputShuffles.contains(SSV))
          return SSV->getOperand(Op);
    return SV->getOperand(Op);
  };
  Builder.SetInsertPoint(*SVI0A->getInsertionPointAfterDef());
  Value *NSV0A = Builder.CreateShuffleVector(GetShuffleOperand(SVI0A, 0),
                                             GetShuffleOperand(SVI0A, 1), V1A);
  Builder.SetInsertPoint(*SVI0B->getInsertionPointAfterDef());
  Value *NSV0B = Builder.CreateShuffleVector(GetShuffleOperand(SVI0B, 0),
                                             GetShuffleOperand(SVI0B, 1), V1B);
  Builder.SetInsertPoint(*SVI1A->getInsertionPointAfterDef());
  Value *NSV1A = Builder.CreateShuffleVector(GetShuffleOperand(SVI1A, 0),
                                             GetShuffleOperand(SVI1A, 1), V2A);
  Builder.SetInsertPoint(*SVI1B->getInsertionPointAfterDef());
  Value *NSV1B = Builder.CreateShuffleVector(GetShuffleOperand(SVI1B, 0),
                                             GetShuffleOperand(SVI1B, 1), V2B);
  Builder.SetInsertPoint(Op0);
  Value *NOp0 = Builder.CreateBinOp((Instruction::BinaryOps)Op0->getOpcode(),
                                    NSV0A, NSV0B);
  if (auto *I = dyn_cast<Instruction>(NOp0))
    I->copyIRFlags(Op0, true);
  Builder.SetInsertPoint(Op1);
  Value *NOp1 = Builder.CreateBinOp((Instruction::BinaryOps)Op1->getOpcode(),
                                    NSV1A, NSV1B);
  if (auto *I = dyn_cast<Instruction>(NOp1))
    I->copyIRFlags(Op1, true);

  for (int S = 0, E = ReconstructMasks.size(); S != E; S++) {
    Builder.SetInsertPoint(Shuffles[S]);
    Value *NSV = Builder.CreateShuffleVector(NOp0, NOp1, ReconstructMasks[S]);
    replaceValue(*Shuffles[S], *NSV, false);
  }

  Worklist.pushValue(NSV0A);
  Worklist.pushValue(NSV0B);
  Worklist.pushValue(NSV1A);
  Worklist.pushValue(NSV1B);
  return true;
}

/// Check if instruction depends on ZExt and this ZExt can be moved after the
/// instruction. Move ZExt if it is profitable. For example:
///     logic(zext(x),y) -> zext(logic(x,trunc(y)))
///     lshr((zext(x),y) -> zext(lshr(x,trunc(y)))
/// Cost model calculations takes into account if zext(x) has other users and
/// whether it can be propagated through them too.
bool VectorCombine::shrinkType(Instruction &I) {
  Value *ZExted, *OtherOperand;
  if (!match(&I, m_c_BitwiseLogic(m_ZExt(m_Value(ZExted)),
                                  m_Value(OtherOperand))) &&
      !match(&I, m_LShr(m_ZExt(m_Value(ZExted)), m_Value(OtherOperand))))
    return false;

  Value *ZExtOperand = I.getOperand(I.getOperand(0) == OtherOperand ? 1 : 0);

  auto *BigTy = cast<FixedVectorType>(I.getType());
  auto *SmallTy = cast<FixedVectorType>(ZExted->getType());
  unsigned BW = SmallTy->getElementType()->getPrimitiveSizeInBits();

  if (I.getOpcode() == Instruction::LShr) {
    // Check that the shift amount is less than the number of bits in the
    // smaller type. Otherwise, the smaller lshr will return a poison value.
    KnownBits ShAmtKB = computeKnownBits(I.getOperand(1), *DL);
    if (ShAmtKB.getMaxValue().uge(BW))
      return false;
  } else {
    // Check that the expression overall uses at most the same number of bits as
    // ZExted
    KnownBits KB = computeKnownBits(&I, *DL);
    if (KB.countMaxActiveBits() > BW)
      return false;
  }

  // Calculate costs of leaving current IR as it is and moving ZExt operation
  // later, along with adding truncates if needed
  InstructionCost ZExtCost = TTI.getCastInstrCost(
      Instruction::ZExt, BigTy, SmallTy,
      TargetTransformInfo::CastContextHint::None, CostKind);
  InstructionCost CurrentCost = ZExtCost;
  InstructionCost ShrinkCost = 0;

  // Calculate total cost and check that we can propagate through all ZExt users
  for (User *U : ZExtOperand->users()) {
    auto *UI = cast<Instruction>(U);
    if (UI == &I) {
      CurrentCost +=
          TTI.getArithmeticInstrCost(UI->getOpcode(), BigTy, CostKind);
      ShrinkCost +=
          TTI.getArithmeticInstrCost(UI->getOpcode(), SmallTy, CostKind);
      ShrinkCost += ZExtCost;
      continue;
    }

    if (!Instruction::isBinaryOp(UI->getOpcode()))
      return false;

    // Check if we can propagate ZExt through its other users
    KnownBits KB = computeKnownBits(UI, *DL);
    if (KB.countMaxActiveBits() > BW)
      return false;

    CurrentCost += TTI.getArithmeticInstrCost(UI->getOpcode(), BigTy, CostKind);
    ShrinkCost +=
        TTI.getArithmeticInstrCost(UI->getOpcode(), SmallTy, CostKind);
    ShrinkCost += ZExtCost;
  }

  // If the other instruction operand is not a constant, we'll need to
  // generate a truncate instruction. So we have to adjust cost
  if (!isa<Constant>(OtherOperand))
    ShrinkCost += TTI.getCastInstrCost(
        Instruction::Trunc, SmallTy, BigTy,
        TargetTransformInfo::CastContextHint::None, CostKind);

  // If the cost of shrinking types and leaving the IR is the same, we'll lean
  // towards modifying the IR because shrinking opens opportunities for other
  // shrinking optimisations.
  if (ShrinkCost > CurrentCost)
    return false;

  Builder.SetInsertPoint(&I);
  Value *Op0 = ZExted;
  Value *Op1 = Builder.CreateTrunc(OtherOperand, SmallTy);
  // Keep the order of operands the same
  if (I.getOperand(0) == OtherOperand)
    std::swap(Op0, Op1);
  Value *NewBinOp =
      Builder.CreateBinOp((Instruction::BinaryOps)I.getOpcode(), Op0, Op1);
  if (auto *NewBinOpI = dyn_cast<Instruction>(NewBinOp)) {
    NewBinOpI->copyIRFlags(&I);
    NewBinOpI->copyMetadata(I);
  }
  Value *NewZExtr = Builder.CreateZExt(NewBinOp, BigTy);
  replaceValue(I, *NewZExtr);
  return true;
}

/// insert (DstVec, (extract SrcVec, ExtIdx), InsIdx) -->
/// shuffle (DstVec, SrcVec, Mask)
bool VectorCombine::foldInsExtVectorToShuffle(Instruction &I) {
  Value *DstVec, *SrcVec;
  uint64_t ExtIdx, InsIdx;
  if (!match(&I,
             m_InsertElt(m_Value(DstVec),
                         m_ExtractElt(m_Value(SrcVec), m_ConstantInt(ExtIdx)),
                         m_ConstantInt(InsIdx))))
    return false;

  auto *DstVecTy = dyn_cast<FixedVectorType>(I.getType());
  auto *SrcVecTy = dyn_cast<FixedVectorType>(SrcVec->getType());
  // We can try combining vectors with different element sizes.
  if (!DstVecTy || !SrcVecTy ||
      SrcVecTy->getElementType() != DstVecTy->getElementType())
    return false;

  unsigned NumDstElts = DstVecTy->getNumElements();
  unsigned NumSrcElts = SrcVecTy->getNumElements();
  if (InsIdx >= NumDstElts || ExtIdx >= NumSrcElts || NumDstElts == 1)
    return false;

  // Insertion into poison is a cheaper single operand shuffle.
  TargetTransformInfo::ShuffleKind SK;
  SmallVector<int> Mask(NumDstElts, PoisonMaskElem);

  bool NeedExpOrNarrow = NumSrcElts != NumDstElts;
  bool NeedDstSrcSwap = isa<PoisonValue>(DstVec) && !isa<UndefValue>(SrcVec);
  if (NeedDstSrcSwap) {
    SK = TargetTransformInfo::SK_PermuteSingleSrc;
    Mask[InsIdx] = ExtIdx % NumDstElts;
    std::swap(DstVec, SrcVec);
  } else {
    SK = TargetTransformInfo::SK_PermuteTwoSrc;
    std::iota(Mask.begin(), Mask.end(), 0);
    Mask[InsIdx] = (ExtIdx % NumDstElts) + NumDstElts;
  }

  // Cost
  auto *Ins = cast<InsertElementInst>(&I);
  auto *Ext = cast<ExtractElementInst>(I.getOperand(1));
  InstructionCost InsCost =
      TTI.getVectorInstrCost(*Ins, DstVecTy, CostKind, InsIdx);
  InstructionCost ExtCost =
      TTI.getVectorInstrCost(*Ext, DstVecTy, CostKind, ExtIdx);
  InstructionCost OldCost = ExtCost + InsCost;

  InstructionCost NewCost = 0;
  SmallVector<int> ExtToVecMask;
  if (!NeedExpOrNarrow) {
    // Ignore 'free' identity insertion shuffle.
    // TODO: getShuffleCost should return TCC_Free for Identity shuffles.
    if (!ShuffleVectorInst::isIdentityMask(Mask, NumSrcElts))
      NewCost += TTI.getShuffleCost(SK, DstVecTy, DstVecTy, CostKind, Mask, 0,
                                    nullptr, {DstVec, SrcVec});
  } else {
    // When creating a length-changing-vector, always try to keep the relevant
    // element in an equivalent position, so that bulk shuffles are more likely
    // to be useful.
    ExtToVecMask.assign(NumDstElts, PoisonMaskElem);
    ExtToVecMask[ExtIdx % NumDstElts] = ExtIdx;
    // Add cost for expanding or narrowing
    NewCost = TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
                                 DstVecTy, SrcVecTy, CostKind, ExtToVecMask);
    NewCost += TTI.getShuffleCost(SK, DstVecTy, DstVecTy, CostKind, Mask);
  }

  if (!Ext->hasOneUse())
    NewCost += ExtCost;

  LLVM_DEBUG(dbgs() << "Found a insert/extract shuffle-like pair: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");

  if (OldCost < NewCost)
    return false;

  if (NeedExpOrNarrow) {
    if (!NeedDstSrcSwap)
      SrcVec = Builder.CreateShuffleVector(SrcVec, ExtToVecMask);
    else
      DstVec = Builder.CreateShuffleVector(DstVec, ExtToVecMask);
  }

  // Canonicalize undef param to RHS to help further folds.
  if (isa<UndefValue>(DstVec) && !isa<UndefValue>(SrcVec)) {
    ShuffleVectorInst::commuteShuffleMask(Mask, NumDstElts);
    std::swap(DstVec, SrcVec);
  }

  Value *Shuf = Builder.CreateShuffleVector(DstVec, SrcVec, Mask);
  replaceValue(I, *Shuf);

  return true;
}

/// Try to replace a chain of insertelements of parts of the same scalar with a
/// bitcast and a shuffle (little endian):
///   insert (insert poison, (trunc (lshr X, 32)), 0), (trunc X), 1 -->
///   shuffle (bitcast X to <2 x i32>), poison, <1, 0>
bool VectorCombine::foldInsertScalarPartsToShuffle(Instruction &I) {
  auto *VecTy = dyn_cast<FixedVectorType>(I.getType());
  if (!VecTy)
    return false;

  // Start from the last insertelement of the chain.
  if (I.hasOneUse() && isa<InsertElementInst>(I.user_back()))
    return false;

  Type *EltTy = VecTy->getElementType();
  if ((!EltTy->isIntegerTy() && !EltTy->isIEEELikeFPTy()) ||
      !DL->typeSizeEqualsStoreSize(EltTy))
    return false;
  unsigned EltBits = EltTy->getPrimitiveSizeInBits();
  unsigned NumElts = VecTy->getNumElements();

  Value *Src = nullptr;
  unsigned NumSrcElts = 0;
  SmallVector<int> Mask(NumElts, PoisonMaskElem);
  APInt DemandedElts = APInt::getZero(NumElts);
  InstructionCost OldCost = 0;
  Value *Vec = &I;
  while (auto *Ins = dyn_cast<InsertElementInst>(Vec)) {
    if (Ins != &I && !Ins->hasOneUse())
      return false;
    uint64_t Idx;
    if (!match(Ins->getOperand(2), m_ConstantInt(Idx)) || Idx >= NumElts)
      return false;
    Vec = Ins->getOperand(0);
    // A later insert to the same element overrides this one.
    if (DemandedElts[Idx])
      continue;
    DemandedElts.setBit(Idx);

    // Match (bitcast (trunc (lshr X, ShAmt))), the bitcast and shift being
    // optional.
    Value *Elt = Ins->getOperand(1);
    Value *Trunc = Elt;
    match(Trunc, m_BitCast(m_Value(Trunc)));
    Value *X;
    if (!match(Trunc, m_Trunc(m_Value(X))) || !X->getType()->isIntegerTy() ||
        Trunc->getType()->getPrimitiveSizeInBits() != EltBits)
      return false;
    Value *Shift = nullptr;
    uint64_t ShAmt = 0;
    if (match(X, m_LShr(m_Value(), m_ConstantInt(ShAmt)))) {
      Shift = X;
      X = cast<Instruction>(Shift)->getOperand(0);
    }

    if (!Src) {
      unsigned SrcBits = X->getType()->getIntegerBitWidth();
      if (SrcBits % EltBits)
        return false;
      Src = X;
      NumSrcElts = SrcBits / EltBits;
    } else if (X != Src) {
      return false;
    }
    uint64_t Part = ShAmt / EltBits;
    if (ShAmt % EltBits || Part >= NumSrcElts)
      return false;
    Mask[Idx] = DL->isBigEndian() ? NumSrcElts - 1 - Part : Part;

    // The scalar ops die with the chain if it is their only user.
    for (Value *V : {Elt == Trunc ? nullptr : Elt, Trunc, Shift}) {
      if (!V)
        continue;
      if (!V->hasOneUse())
        break;
      OldCost += TTI.getInstructionCost(cast<Instruction>(V), CostKind);
    }
  }
  // Elements that are not inserted become poison, so the base must be poison
  // unless every element is inserted.
  if (!Src || (!isa<PoisonValue>(Vec) && !DemandedElts.isAllOnes()))
    return false;
  // Inserting a single part is a scalar insert or a splat, whose canonical
  // insertelement (+ splat shuffle) form is better left alone.
  if (all_equal(
          make_filter_range(Mask, [](int M) { return M != PoisonMaskElem; })))
    return false;

  OldCost += TTI.getScalarizationOverhead(VecTy, DemandedElts, /*Insert=*/true,
                                          /*Extract=*/false, CostKind);

  auto *SrcVecTy = FixedVectorType::get(EltTy, NumSrcElts);
  InstructionCost NewCost =
      TTI.getCastInstrCost(Instruction::BitCast, SrcVecTy, Src->getType(),
                           TTI::CastContextHint::None, CostKind);
  bool IsIdentity = NumSrcElts == NumElts &&
                    ShuffleVectorInst::isIdentityMask(Mask, NumSrcElts);
  if (!IsIdentity)
    NewCost += TTI.getShuffleCost(TTI::SK_PermuteSingleSrc, VecTy, SrcVecTy,
                                  CostKind, Mask);

  LLVM_DEBUG(dbgs() << "Found an insertelement chain of scalar parts: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");
  if (!OldCost.isValid() || !NewCost.isValid() || NewCost >= OldCost)
    return false;

  Value *Cast = Builder.CreateBitCast(Src, SrcVecTy);
  Value *Shuf = IsIdentity ? Cast : Builder.CreateShuffleVector(Cast, Mask);
  replaceValue(I, *Shuf);
  return true;
}

/// Return the number of data operands of \p Inst.
static unsigned getNumDataOperands(const Instruction *Inst) {
  if (auto *CB = dyn_cast<CallBase>(Inst))
    return CB->arg_size(); // Exclude callee operand and bundles.
  return Inst->getNumOperands();
}

/// Return true if \p Inst is an elementwise operation that can be rebuilt at a
/// wider element count.
static bool isSupportedElementwise(Instruction *Inst) {
  auto *ResultTy = dyn_cast<VectorType>(Inst->getType());
  if (!ResultTy || !isSafeToSpeculativelyExecute(Inst))
    return false;

  if (auto *II = dyn_cast<IntrinsicInst>(Inst)) {
    if (II->hasOperandBundles() ||
        !isTriviallyVectorizable(II->getIntrinsicID()))
      return false;
  } else if (!isa<BinaryOperator, UnaryOperator, CastInst, CmpInst, SelectInst,
                  FreezeInst>(Inst)) {
    return false;
  }

  // Reject operations that change the element-count.
  // E.g., bitcast <vscale x 4 x i16> %v to <vscale x 8 x i8>
  for (unsigned Op = 0, E = getNumDataOperands(Inst); Op != E; ++Op) {
    auto *OperandTy = dyn_cast<VectorType>(Inst->getOperand(Op)->getType());
    if (OperandTy &&
        OperandTy->getElementCount() != ResultTy->getElementCount())
      return false;
  }

  return true;
}

/// Return the common splat value of \p Values.
static Value *getCommonSplatValue(ArrayRef<Value *> Values) {
  auto GetSplatOrScalar = [](Value *V) {
    return isa<VectorType>(V->getType()) ? getSplatValue(V) : V;
  };

  Value *CommonValue = GetSplatOrScalar(Values.front());
  if (!CommonValue)
    return nullptr;
  for (Value *V : Values.drop_front())
    if (GetSplatOrScalar(V) != CommonValue)
      return nullptr;
  return CommonValue;
}

/// Return the common deinterleave intrinsic if \p Members are its extracts in
/// field order.
static IntrinsicInst *getCommonDeinterleavedSource(ArrayRef<Value *> Members) {
  IntrinsicInst *Deinterleave = nullptr;
  for (const auto &[Index, Member] : enumerate(Members)) {
    auto *Extract = dyn_cast<ExtractValueInst>(Member);
    if (!Extract || Extract->getNumIndices() != 1 ||
        *Extract->idx_begin() != Index)
      return nullptr;

    auto *Current = dyn_cast<IntrinsicInst>(Extract->getAggregateOperand());
    if (!Current || (Deinterleave && Current != Deinterleave))
      return nullptr;
    Deinterleave = Current;
  }
  unsigned Factor = Members.size();
  if (Deinterleave->hasOperandBundles() ||
      getDeinterleaveIntrinsicFactor(Deinterleave->getIntrinsicID()) !=
          Factor ||
      !Deinterleave->hasNUndroppableUses(Factor))
    return nullptr;
  return Deinterleave;
}

/// Return the operand at \p OperandIndex of each \p Members.
static SmallVector<Value *, 8> getInstrOperandsAtIdx(ArrayRef<Value *> Members,
                                                     unsigned OperandIndex) {
  SmallVector<Value *, 8> Operands;
  for (Value *Member : Members)
    Operands.push_back(cast<Instruction>(Member)->getOperand(OperandIndex));
  return Operands;
}

/// Check whether the tree of elementwise operations each feeding \p Members
/// can be rebuilt at the interleaved width.
static bool canWidenDeinterleavedOperations(ArrayRef<Value *> Members,
                                            unsigned &NumScanned) {
  unsigned Factor = Members.size();
  if (getCommonDeinterleavedSource(Members))
    return true;
  if (NumScanned + Factor > MaxInstrsToScan)
    return false;
  NumScanned += Factor;

  auto *FirstInst = dyn_cast<Instruction>(Members.front());
  if (!FirstInst || !isSupportedElementwise(FirstInst) ||
      !FirstInst->getSingleUndroppableUse())
    return false;

  for (Value *Member : Members.drop_front()) {
    auto *Inst = dyn_cast<Instruction>(Member);
    if (!Inst || !Inst->getSingleUndroppableUse() ||
        !FirstInst->isSameOperationAs(Inst, Instruction::CompareCallTargets))
      return false;
  }

  // Scalars operands should be equal among all members.
  // Vector operands should be a common splat value or can be widened.
  for (unsigned Op = 0, E = getNumDataOperands(FirstInst); Op != E; ++Op) {
    SmallVector<Value *, 8> Operands = getInstrOperandsAtIdx(Members, Op);
    if (!getCommonSplatValue(Operands) &&
        !canWidenDeinterleavedOperations(Operands, NumScanned))
      return false;
  }
  return true;
}

static Value *createWideInstruction(Instruction *NarrowInst,
                                    ArrayRef<Value *> NewOperands,
                                    VectorType *WideResultTy,
                                    IRBuilder<InstSimplifyFolder> &Builder) {
  if (isa<BinaryOperator, UnaryOperator>(NarrowInst))
    return Builder.CreateNAryOp(NarrowInst->getOpcode(), NewOperands);
  if (auto *Cast = dyn_cast<CastInst>(NarrowInst))
    return Builder.CreateCast(Cast->getOpcode(), NewOperands[0], WideResultTy);
  if (auto *Cmp = dyn_cast<CmpInst>(NarrowInst))
    return Builder.CreateCmp(Cmp->getPredicate(), NewOperands[0],
                             NewOperands[1]);
  if (isa<SelectInst>(NarrowInst))
    return Builder.CreateSelect(
        NewOperands[0], NewOperands[1], NewOperands[2], /*Name=*/"",
        ProfcheckDisableMetadataFixes ? nullptr : NarrowInst);
  if (isa<FreezeInst>(NarrowInst))
    return Builder.CreateFreeze(NewOperands[0]);
  if (auto *II = dyn_cast<IntrinsicInst>(NarrowInst))
    return Builder.CreateIntrinsic(WideResultTy, II->getIntrinsicID(),
                                   NewOperands);
  llvm_unreachable("Unsupported instruction");
}

static Value *
widenDeinterleavedOperations(ArrayRef<Value *> Members, ElementCount WideEC,
                             IRBuilder<InstSimplifyFolder> &Builder) {
  if (auto *Deinterleave = getCommonDeinterleavedSource(Members)) {
    Value *Source = Deinterleave->getArgOperand(0);
    assert(cast<VectorType>(Source->getType())->getElementCount() == WideEC &&
           "deinterleave source must have the interleaved element count");
    return Source;
  }

  auto *NarrowInst = cast<Instruction>(Members.front());
  unsigned NumOperands = getNumDataOperands(NarrowInst);
  SmallVector<Value *, 4> NewOperands;
  NewOperands.reserve(NumOperands);
  for (unsigned Op = 0; Op != NumOperands; ++Op) {
    SmallVector<Value *, 8> Operands = getInstrOperandsAtIdx(Members, Op);
    Value *NewOperand = nullptr;
    if ((NewOperand = getCommonSplatValue(Operands))) {
      if (isa<VectorType>(Operands.front()->getType())) {
        Builder.SetCurrentDebugLocation(NarrowInst->getDebugLoc());
        NewOperand = Builder.CreateVectorSplat(WideEC, NewOperand);
      } else {
        assert(all_equal(Operands) && "expected all operands to be equal");
      }
    } else {
      NewOperand = widenDeinterleavedOperations(Operands, WideEC, Builder);
    }
    NewOperands.push_back(NewOperand);
  }

  Builder.SetCurrentDebugLocation(NarrowInst->getDebugLoc());
  auto *WideResultTy =
      VectorType::get(NarrowInst->getType()->getScalarType(), WideEC);
  Value *NewValue =
      createWideInstruction(NarrowInst, NewOperands, WideResultTy, Builder);
  if (auto *NewInst = dyn_cast<Instruction>(NewValue)) {
    propagateIRFlags(NewInst, Members);
    propagateMetadata(NewInst, Members);
  }
  return NewValue;
}

/// Fold away vector.deinterleave/interleave intrinsics with matching trees of
/// elementwise operations between them.
///
/// For example:
///  ```
///  %d = call { <2 x i16>, <2 x i16> } @deinterleave2.v4i16(<4 x i16> %v)
///  %f0 = extractvalue { <2 x i16>, <2 x i16> } %d, 0
///  %f1 = extractvalue { <2 x i16>, <2 x i16> } %d, 1
///
///  %u0 = add <2 x i16> %f0, splat (i16 3)
///  %u1 = add <2 x i16> %f1, splat (i16 3)
///
///  %r = call <4 x i16> @interleave2.v4i16(<2 x i16> %u0, <2 x i16> %u1)
///  ```
/// Folds to:
///  ```
///  %r = add <4 x i16> %v, splat (i16 3)
///  ```
/// And with two sources:
///  ```
///  %da = call { <2 x i16>, <2 x i16> } @deinterleave2.v4i16(<4 x i16> %a)
///  %a0 = extractvalue { <2 x i16>, <2 x i16> } %da, 0
///  %a1 = extractvalue { <2 x i16>, <2 x i16> } %da, 1
///  %db = call { <2 x i16>, <2 x i16> } @deinterleave2.v4i16(<4 x i16> %b)
///  %b0 = extractvalue { <2 x i16>, <2 x i16> } %db, 0
///  %b1 = extractvalue { <2 x i16>, <2 x i16> } %db, 1
///
///  %m0 = mul <2 x i16> %a0, %b0
///  %m1 = mul <2 x i16> %a1, %b1
///
///  %r = call <4 x i16> @interleave2.v4i16(<2 x i16> %m0, <2 x i16> %m1)
///  ```
/// Folds to:
///  ```
///  %r = mul <4 x i16> %a, %b
///  ```
bool VectorCombine::foldInterleaveOfDeinterleaveChains(Instruction &I) {
  auto *Interleave = dyn_cast<IntrinsicInst>(&I);
  if (!Interleave)
    return false;

  unsigned Factor = getInterleaveIntrinsicFactor(Interleave->getIntrinsicID());
  if (!Factor || Interleave->hasOperandBundles())
    return false;

  SmallVector<Value *, 8> RootMembers(Interleave->args());
  unsigned NumScanned = 0;
  if (!canWidenDeinterleavedOperations(RootMembers, NumScanned))
    return false;

  ElementCount WideEC = cast<VectorType>(I.getType())->getElementCount();
  Builder.SetInsertPoint(Interleave);
  Value *WideValue = widenDeinterleavedOperations(RootMembers, WideEC, Builder);
  assert(WideValue->getType() == Interleave->getType());
  replaceValue(*Interleave, *WideValue);
  return true;
}

/// If we're interleaving 2 constant splats, for instance `<vscale x 8 x i32>
/// <splat of 666>` and `<vscale x 8 x i32> <splat of 777>`, we can create a
/// larger splat `<vscale x 8 x i64> <splat of ((777 << 32) | 666)>` first
/// before casting it back into `<vscale x 16 x i32>`.
bool VectorCombine::foldInterleaveIntrinsics(Instruction &I) {
  if (foldInterleaveOfDeinterleaveChains(I))
    return true;

  const APInt *SplatVal0, *SplatVal1;
  if (!match(&I, m_Intrinsic<Intrinsic::vector_interleave2>(
                     m_APInt(SplatVal0), m_APInt(SplatVal1))))
    return false;

  LLVM_DEBUG(dbgs() << "VC: Folding interleave2 with two splats: " << I
                    << "\n");

  auto *VTy =
      cast<VectorType>(cast<IntrinsicInst>(I).getArgOperand(0)->getType());
  auto *ExtVTy = VectorType::getExtendedElementVectorType(VTy);
  unsigned Width = VTy->getElementType()->getIntegerBitWidth();

  // Just in case the cost of interleave2 intrinsic and bitcast are both
  // invalid, in which case we want to bail out, we use <= rather
  // than < here. Even they both have valid and equal costs, it's probably
  // not a good idea to emit a high-cost constant splat.
  if (TTI.getInstructionCost(&I, CostKind) <=
      TTI.getCastInstrCost(Instruction::BitCast, I.getType(), ExtVTy,
                           TTI::CastContextHint::None, CostKind)) {
    LLVM_DEBUG(dbgs() << "VC: The cost to cast from " << *ExtVTy << " to "
                      << *I.getType() << " is too high.\n");
    return false;
  }

  APInt NewSplatVal = SplatVal1->zext(Width * 2);
  NewSplatVal <<= Width;
  NewSplatVal |= SplatVal0->zext(Width * 2);
  auto *NewSplat = ConstantVector::getSplat(
      ExtVTy->getElementCount(), ConstantInt::get(F.getContext(), NewSplatVal));

  IRBuilder<> Builder(&I);
  replaceValue(I, *Builder.CreateBitCast(NewSplat, I.getType()));
  return true;
}

/// Given this sequence:
/// ```
/// %d = llvm.vector.deinterleave2 <vscale x 16 x i32> %v
/// %f0 = extractvalue { <vscale x 8 x i32>, <vscale x 8 x i32> } %d, 0
/// %f1 = extractvalue { <vscale x 8 x i32>, <vscale x 8 x i32> } %d, 1
///
/// %low0 = and <vscale x 8 x i32> %f0, splat (i32 65535)
/// %low1 = shl <vscale x 8 x i32> %f1, splat (i32 16)
/// %merge0 = or disjoint <vscale x 8 x i32> %low0, %low1
///
/// %high0 = and <vscale x 8 x i32> %f1, splat (i32 -65536)
/// %high1 = lshr <vscale x 8 x i32> %f0, splat (i32 16)
/// %merge1 = or disjoint <vscale x 8 x i32> %high0, %high1
/// ```
/// It is actually just de-interleaving a 16-bit vector with double the
/// vector length. More generally speaking, it's de-interleaving on a vector
/// with half the element width as the original vector.
///
/// Therefore, we can turn it into:
/// ```
/// %narrow.v = bitcast <vscale x 16 x i32> %v to <vscale x 32 x i16>
/// %d = llvm.vector.deinterleave2 <vscale x 32 x i16> %narrow.v
/// %f0 = extractvalue { <vscale x 16 x i16>, <vscale x 16 x i16> } %d, 0
/// %f1 = extractvalue { <vscale x 16 x i16>, <vscale x 16 x i16> } %d, 1
///
/// %merge0 = bitcast <vscale x 16 x i16> %f0 to <vscale x 8 x i32>
/// %merge1 = bitcast <vscale x 16 x i16> %f1 to <vscale x 8 x i32>
/// ```
bool VectorCombine::foldDeinterleaveIntrinsics(Instruction &I) {
  // This pattern involves bitcast that is not compatible with big endian.
  if (DL->isBigEndian())
    return false;

  using namespace PatternMatch;
  Value *DeinterleavedVal;
  if (!match(&I, m_Deinterleave2(m_Value(DeinterleavedVal))))
    return false;

  VectorType *VecTy = cast<VectorType>(DeinterleavedVal->getType());
  IntegerType *ElementTy = dyn_cast<IntegerType>(VecTy->getElementType());
  if (!ElementTy)
    return false;
  unsigned ElementWidth = ElementTy->getBitWidth();
  if (ElementWidth < 2 || !isPowerOf2_32(ElementWidth))
    return false;
  unsigned HalfElementWidth = ElementWidth / 2;

  if (!I.hasNUses(2))
    return false;
  std::array<ExtractValueInst *, 2> OrigFields{};
  for (User *Usr : I.users()) {
    auto *E = dyn_cast<ExtractValueInst>(Usr);
    // The deinterleave result can only be used by extractions.
    if (!E || E->getNumIndices() != 1)
      return false;
    unsigned Idx = *E->idx_begin();
    // A single field cannot be extracted more than once.
    if (Idx >= 2 || OrigFields[Idx] || !E->hasNUses(2))
      return false;
    OrigFields[Idx] = E;
  }

  // Find the merge instruction (i.e. OR) first.
  SmallVector<Instruction *, 2> MergeInsts;
  for (auto *FieldUsr : OrigFields[0]->users()) {
    if (!FieldUsr->hasOneUse() || !isa<Instruction>(FieldUsr->user_back()))
      return false;
    MergeInsts.push_back(cast<Instruction>(FieldUsr->user_back()));
  }
  assert(MergeInsts.size() == 2);

  // Pattern match bottom-up from the merge instructions.
  auto MatchMerge = [&](void) -> bool {
    APInt LoMask = APInt::getLowBitsSet(ElementWidth, HalfElementWidth);
    APInt HiMask = APInt::getHighBitsSet(ElementWidth, HalfElementWidth);
    return match(MergeInsts[0],
                 m_c_Or(m_And(m_Specific(OrigFields[0]), m_SpecificInt(LoMask)),
                        m_Shl(m_Specific(OrigFields[1]),
                              m_SpecificInt(HalfElementWidth)))) &&
           match(MergeInsts[1],
                 m_c_Or(m_And(m_Specific(OrigFields[1]), m_SpecificInt(HiMask)),
                        m_LShr(m_Specific(OrigFields[0]),
                               m_SpecificInt(HalfElementWidth))));
  };
  if (!MatchMerge()) {
    std::swap(MergeInsts[0], MergeInsts[1]);
    if (!MatchMerge())
      return false;
  }

  // Profitability check.
  InstructionCost OldCost =
      TTI.getInstructionCost(MergeInsts[0], CostKind) +
      TTI.getInstructionCost(cast<Instruction>(MergeInsts[0]->getOperand(0)),
                             CostKind) +
      TTI.getInstructionCost(cast<Instruction>(MergeInsts[0]->getOperand(1)),
                             CostKind);
  // There are two fields (assuming SHL has the same cost as LSHR).
  OldCost *= 2;

  auto *NewFieldTy = VecTy->getWithNewBitWidth(HalfElementWidth);
  auto *NewVecTy =
      VectorType::getDoubleElementsVectorType(cast<VectorType>(NewFieldTy));
  InstructionCost NewCost =
      TTI.getCastInstrCost(Instruction::BitCast, VecTy, NewVecTy,
                           TTI::CastContextHint::None, CostKind) +
      TTI.getCastInstrCost(Instruction::BitCast, NewFieldTy,
                           MergeInsts[0]->getType(), TTI::CastContextHint::None,
                           CostKind) *
          2;
  if (OldCost <= NewCost || !NewCost.isValid()) {
    LLVM_DEBUG(
        dbgs() << "VC: New deinterleave2 sequence cost (" << NewCost << ")"
               << " is higher than that of the old one (" << OldCost << ")\n");
    return false;
  }

  // Do the replacement.
  IRBuilder<> Builder(&I);
  Value *NewVecCast = Builder.CreateBitCast(DeinterleavedVal, NewVecTy);
  Value *NewDeinterleave = Builder.CreateIntrinsic(
      Intrinsic::vector_deinterleave2, {NewVecTy}, {NewVecCast});
  for (auto [Idx, MergeInst] : enumerate(MergeInsts)) {
    Value *NewField = Builder.CreateExtractValue(NewDeinterleave, Idx);
    NewField = Builder.CreateBitCast(NewField, MergeInst->getType());
    replaceValue(*MergeInst, *NewField);
  }

  return true;
}

bool VectorCombine::foldBitcastOfVPLoad(Instruction &I) {
  const DataLayout &DL = I.getDataLayout();
  auto *Cast = dyn_cast<CastInst>(&I);
  if (!Cast || !Cast->isNoopCast(DL) || !isa<VectorType>(Cast->getDestTy()))
    return false;

  // Fold away bit casts of the loaded value by loading the desired type,
  // if the mask is all-ones.
  Value *EVL;
  auto *II = dyn_cast<VPIntrinsic>(I.getOperand(0));
  if (!II || !match(II, m_OneUse(m_Intrinsic<Intrinsic::vp_load>(
                            m_Value(), m_AllOnes(), m_Value(EVL)))))
    return false;

  VectorType *OrigVecTy = cast<VectorType>(II->getType());
  Align OrigAlign =
      DL.getValueOrABITypeAlignment(II->getPointerAlignment(), OrigVecTy);
  ElementCount OrigVecCnt = OrigVecTy->getElementCount();
  VectorType *NewVecTy = cast<VectorType>(Cast->getDestTy());
  ElementCount NewVecCnt = NewVecTy->getElementCount();

  // Right now we only support cases where the NewVec is longer, because for
  // cases where it's shorter, we have to be sure that EVL can be exactly
  // divided, otherwise it might yield incorrect results or even page faults
  // (if we round-up during the division).
  if (!(OrigVecCnt.isScalable() == NewVecCnt.isScalable() &&
        NewVecCnt.hasKnownScalarFactor(OrigVecCnt)))
    return false;

  InstructionCost OldCost =
      TTI.getMemIntrinsicInstrCost({Intrinsic::vp_load, OrigVecTy,
                                    II->getMemoryPointerParam(), false,
                                    OrigAlign},
                                   CostKind) +
      TTI.getCastInstrCost(Instruction::BitCast, Cast->getType(), OrigVecTy,
                           TTI::CastContextHint::None, CostKind);
  InstructionCost NewCost = TTI.getMemIntrinsicInstrCost(
      {Intrinsic::vp_load, NewVecTy, II->getMemoryPointerParam(), false,
       OrigAlign},
      CostKind);
  LLVM_DEBUG(dbgs() << "foldBitcastOfVPLoad: OldCost=" << OldCost
                    << " NewCost=" << NewCost << "\n");
  if (NewCost > OldCost || !NewCost.isValid())
    return false;

  Builder.SetInsertPoint(II);
  unsigned Factor = NewVecCnt.getKnownScalarFactor(OrigVecCnt);
  Value *NewEVL = Builder.CreateNUWMul(EVL, Builder.getInt32(Factor));
  Value *NewMask = Builder.CreateVectorSplat(NewVecCnt, Builder.getTrue());
  CallInst *NewVP = Builder.CreateIntrinsicWithoutFolding(
      NewVecTy, Intrinsic::vp_load,
      {II->getMemoryPointerParam(), NewMask, NewEVL});
  // Preserve the original alignment.
  NewVP->addParamAttrs(
      0, AttrBuilder(II->getContext()).addAlignmentAttr(OrigAlign));
  replaceValue(*Cast, *NewVP);
  return true;
}
/// Fold the following cases into a single byte-level bit-reverse operation
/// and accepts bswap and bitreverse intrinsics:
///   bswap(bitreverse(x))  --> bitcast(bitreverse(bitcast(x)))
///   bitreverse(bswap(x)) <--> bitcast(bitreverse(bitcast(x)))
/// The direction of the fold is cost-model driven.
/// Also supports:
///   bitcast(bitreverse(bitcast(x))) --> bitreverse(fshl(x))
bool VectorCombine::foldBitOrderReverseAndSwap(Instruction &I) {
  Value *X;

  if (match(&I, m_BitCast(m_BitReverse(m_BitCast(m_Value(X)))))) {
    Type *Ty = X->getType();
    Type *VecTy = I.getOperand(0)->getType();
    // Detect the case when bitreversing every octet in X individually. Then we
    // can use bswap to reorder the octets before doing a single bitreverse.
    bool CanUseBswap =
        Ty->isIntegerTy() && Ty == I.getType() && isa<FixedVectorType>(VecTy) &&
        cast<FixedVectorType>(VecTy)->getElementType()->isIntegerTy(8) &&
        Ty->getIntegerBitWidth() % 16 == 0;
    // Detect the case when bitreversing upper and lower half of X
    // individually. Then we can use fshl as a rotate operation, to swap the
    // halves before doing a single bitreverse.
    bool CanUseFshl =
        Ty->isIntegerTy() && Ty == I.getType() && isa<FixedVectorType>(VecTy) &&
        cast<FixedVectorType>(VecTy)->getElementType()->isIntegerTy() &&
        cast<FixedVectorType>(VecTy)->getNumElements() == 2;
    if (CanUseBswap || CanUseFshl) {
      auto *InnerCall = dyn_cast<Instruction>(I.getOperand(0));
      if (!InnerCall)
        return false;
      auto *InnerBitCast = dyn_cast<BitCastInst>(InnerCall->getOperand(0));
      if (!InnerBitCast)
        return false;
      Constant *HalfBW = ConstantInt::get(Ty, Ty->getIntegerBitWidth() / 2);
      InstructionCost OldCost = TTI.getInstructionCost(InnerBitCast, CostKind) +
                                TTI.getInstructionCost(InnerCall, CostKind) +
                                TTI.getInstructionCost(&I, CostKind);
      IntrinsicCostAttributes ICABSwap(Intrinsic::bswap, Ty, {Ty});
      IntrinsicCostAttributes ICABFshl(Intrinsic::fshl, Ty, {X, X, HalfBW},
                                       {Ty, Ty, Ty});
      IntrinsicCostAttributes ICABRev(Intrinsic::bitreverse, Ty, {Ty});
      InstructionCost NewCost =
          TTI.getIntrinsicInstrCost(CanUseBswap ? ICABSwap : ICABFshl,
                                    CostKind) +
          TTI.getIntrinsicInstrCost(ICABRev, CostKind);
      if (!InnerCall->hasOneUse())
        NewCost += TTI.getInstructionCost(InnerCall, CostKind) +
                   TTI.getInstructionCost(InnerBitCast, CostKind);
      else if (!InnerBitCast->hasOneUse())
        NewCost += TTI.getInstructionCost(InnerBitCast, CostKind);
      LLVM_DEBUG(dbgs() << "Found bitreverse vector roundtrip: " << I
                        << "\n  OldCost: " << OldCost
                        << " vs NewCost: " << NewCost << "\n");
      if (NewCost.isValid() && NewCost < OldCost) {
        Builder.SetInsertPoint(&I);
        Value *Swap =
            CanUseBswap
                ? Builder.CreateUnaryIntrinsic(Intrinsic::bswap, X)
                : Builder.CreateIntrinsic(Ty, Intrinsic::fshl, {X, X, HalfBW});
        Worklist.pushValue(Swap);
        Value *BRev = Builder.CreateUnaryIntrinsic(Intrinsic::bitreverse, Swap);
        replaceValue(I, *BRev);
        return true;
      }
    }
  }

  if (!match(&I, m_BitReverse(m_BSwap(m_Value(X)))) &&
      !match(&I, m_BSwap(m_BitReverse(m_Value(X)))))
    return false;
  Type *Ty = I.getType();
  Type *I8Ty = Builder.getInt8Ty();
  TypeSize ElementSize = DL->getTypeStoreSize(Ty);
  ElementCount NewVecCnt = ElementCount::get(ElementSize.getKnownMinValue(),
                                             ElementSize.isScalable());
  Type *NewVecTy = VectorType::get(I8Ty, NewVecCnt);
  auto *II = cast<IntrinsicInst>(&I);
  auto *InnerII = cast<IntrinsicInst>(II->getArgOperand(0));
  // OldCost = cost of bitreverse/bswap + cost of bswap/bitreverse
  InstructionCost OldCost = TTI.getInstructionCost(II, CostKind) +
                            TTI.getInstructionCost(InnerII, CostKind);
  // NewCost = cost of bitcast to byte vector +
  //           cost of bitreverse/bswap on byte vector +
  //           cost of bitcast back to original type
  InstructionCost CastToVecCost = TTI.getCastInstrCost(
      Instruction::BitCast, NewVecTy, Ty, TTI::CastContextHint::None, CostKind);
  InstructionCost CastToOrigCost = TTI.getCastInstrCost(
      Instruction::BitCast, Ty, NewVecTy, TTI::CastContextHint::None, CostKind);
  IntrinsicCostAttributes ICANew(Intrinsic::bitreverse, NewVecTy, {NewVecTy});
  InstructionCost NewIntrinsicCost =
      TTI.getIntrinsicInstrCost(ICANew, CostKind);
  InstructionCost NewCost = CastToVecCost + NewIntrinsicCost + CastToOrigCost;
  if (!InnerII->hasOneUse())
    NewCost += TTI.getInstructionCost(InnerII, CostKind);
  LLVM_DEBUG(dbgs() << "Found bitorder reverse and swap: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");
  if (!NewCost.isValid() || NewCost >= OldCost)
    return false;
  // Perform transform: bitcast(arg, <N x i8>), bitreverse, bitcast back
  Builder.SetInsertPoint(II);
  Value *CastToVec = Builder.CreateBitCast(X, NewVecTy);
  Value *NewCall =
      Builder.CreateUnaryIntrinsic(Intrinsic::bitreverse, CastToVec);
  Value *CastToOrig = Builder.CreateBitCast(NewCall, Ty);
  replaceValue(I, *CastToOrig);
  return true;
}

/// Given the maximum shuffle index and load vector type, compute the number of
/// elements for the shrunk load, rounding up to the next full vector register
/// boundary to avoid scalar remainders that legalize poorly.
static unsigned getAlignedNumElements(unsigned MaxIdx, FixedVectorType *LoadTy,
                                      const TargetTransformInfo &TTI,
                                      const DataLayout &DL) {
  unsigned RawNumElements = MaxIdx + 1u;
  Type *ElemTy = LoadTy->getElementType();
  // Skip alignment for illegal element types.
  if (!TTI.isTypeLegal(ElemTy))
    return RawNumElements;

  TypeSize ElemSize = DL.getTypeSizeInBits(ElemTy);
  if (ElemSize.isScalable() || ElemSize.isZero())
    return RawNumElements;

  TypeSize RegSize =
      TTI.getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector);
  if (RegSize.isScalable() || RegSize.isZero())
    return RawNumElements;

  unsigned ElemsPerReg = RegSize.getFixedValue() / ElemSize.getFixedValue();
  // If the load already fits in a register, keep the exact size.
  // Otherwise round up to the next full register boundary.
  if (ElemsPerReg == 0 || RawNumElements <= ElemsPerReg)
    return RawNumElements;

  return alignTo(RawNumElements, ElemsPerReg);
}

// Attempt to shrink loads that are only used by shufflevector instructions.
bool VectorCombine::shrinkLoadForShuffles(Instruction &I) {
  auto *OldLoad = dyn_cast<LoadInst>(&I);
  if (!OldLoad || !OldLoad->isSimple())
    return false;

  auto *OldLoadTy = dyn_cast<FixedVectorType>(OldLoad->getType());
  if (!OldLoadTy)
    return false;

  unsigned const OldNumElements = OldLoadTy->getNumElements();

  // Search all uses of load. If all uses are shufflevector instructions, and
  // the second operands are all poison values, find the minimum and maximum
  // indices of the vector elements referenced by all shuffle masks.
  // Otherwise return `std::nullopt`.
  using IndexRange = std::pair<int, int>;
  auto GetIndexRangeInShuffles = [&]() -> std::optional<IndexRange> {
    IndexRange OutputRange = IndexRange(OldNumElements, -1);
    for (llvm::Use &Use : I.uses()) {
      // Ensure all uses match the required pattern.
      User *Shuffle = Use.getUser();
      ArrayRef<int> Mask;

      if (!match(Shuffle,
                 m_Shuffle(m_Specific(OldLoad), m_Undef(), m_Mask(Mask))))
        return std::nullopt;

      // Ignore shufflevector instructions that have no uses.
      if (Shuffle->use_empty())
        continue;

      // Find the min and max indices used by the shufflevector instruction.
      for (int Index : Mask) {
        if (Index >= 0 && Index < static_cast<int>(OldNumElements)) {
          OutputRange.first = std::min(Index, OutputRange.first);
          OutputRange.second = std::max(Index, OutputRange.second);
        }
      }
    }

    if (OutputRange.second < OutputRange.first)
      return std::nullopt;

    return OutputRange;
  };

  // Get the range of vector elements used by shufflevector instructions.
  if (std::optional<IndexRange> Indices = GetIndexRangeInShuffles()) {
    unsigned const NewNumElements =
        getAlignedNumElements(Indices->second, OldLoadTy, TTI, *DL);

    // If the range of vector elements is smaller than the full load, attempt
    // to create a smaller load.
    if (NewNumElements < OldNumElements) {
      IRBuilder Builder(&I);
      Builder.SetCurrentDebugLocation(I.getDebugLoc());

      // Calculate costs of old and new ops.
      Type *ElemTy = OldLoadTy->getElementType();
      FixedVectorType *NewLoadTy = FixedVectorType::get(ElemTy, NewNumElements);
      Value *PtrOp = OldLoad->getPointerOperand();

      InstructionCost OldCost = TTI.getMemoryOpCost(
          Instruction::Load, OldLoad->getType(), OldLoad->getAlign(),
          OldLoad->getPointerAddressSpace(), CostKind);
      InstructionCost NewCost =
          TTI.getMemoryOpCost(Instruction::Load, NewLoadTy, OldLoad->getAlign(),
                              OldLoad->getPointerAddressSpace(), CostKind);

      using UseEntry = std::pair<ShuffleVectorInst *, std::vector<int>>;
      SmallVector<UseEntry, 4u> NewUses;
      unsigned const MaxIndex = NewNumElements * 2u;

      for (llvm::Use &Use : I.uses()) {
        auto *Shuffle = cast<ShuffleVectorInst>(Use.getUser());

        // Ignore shufflevector instructions that have no uses.
        if (Shuffle->use_empty())
          continue;

        ArrayRef<int> OldMask = Shuffle->getShuffleMask();

        // Create entry for new use.
        NewUses.push_back({Shuffle, OldMask});

        // Validate mask indices.
        for (int Index : OldMask) {
          if (Index >= static_cast<int>(MaxIndex))
            return false;
        }

        // Update costs.
        OldCost +=
            TTI.getShuffleCost(TTI::SK_PermuteSingleSrc, Shuffle->getType(),
                               OldLoadTy, CostKind, OldMask);
        NewCost +=
            TTI.getShuffleCost(TTI::SK_PermuteSingleSrc, Shuffle->getType(),
                               NewLoadTy, CostKind, OldMask);
      }

      LLVM_DEBUG(
          dbgs() << "Found a load used only by shufflevector instructions: "
                 << I << "\n  OldCost: " << OldCost
                 << " vs NewCost: " << NewCost << "\n");

      if (OldCost < NewCost || !NewCost.isValid())
        return false;

      // Create new load of smaller vector.
      auto *NewLoad = cast<LoadInst>(
          Builder.CreateAlignedLoad(NewLoadTy, PtrOp, OldLoad->getAlign()));
      NewLoad->copyMetadata(I);

      // Replace all uses.
      for (UseEntry &Use : NewUses) {
        ShuffleVectorInst *Shuffle = Use.first;
        std::vector<int> &NewMask = Use.second;

        Builder.SetInsertPoint(Shuffle);
        Builder.SetCurrentDebugLocation(Shuffle->getDebugLoc());
        Value *NewShuffle = Builder.CreateShuffleVector(
            NewLoad, PoisonValue::get(NewLoadTy), NewMask);

        replaceValue(*Shuffle, *NewShuffle, false);
      }

      return true;
    }
  }
  return false;
}

// Attempt to narrow a phi of shufflevector instructions where the two incoming
// values have the same operands but different masks. If the two shuffle masks
// are offsets of one another we can use one branch to rotate the incoming
// vector and perform one larger shuffle after the phi.
bool VectorCombine::shrinkPhiOfShuffles(Instruction &I) {
  auto *Phi = dyn_cast<PHINode>(&I);
  if (!Phi || Phi->getNumIncomingValues() != 2u)
    return false;

  Value *Op = nullptr;
  ArrayRef<int> Mask0;
  ArrayRef<int> Mask1;

  if (!match(Phi->getOperand(0u),
             m_OneUse(m_Shuffle(m_Value(Op), m_Poison(), m_Mask(Mask0)))) ||
      !match(Phi->getOperand(1u),
             m_OneUse(m_Shuffle(m_Specific(Op), m_Poison(), m_Mask(Mask1)))))
    return false;

  auto *Shuf = cast<ShuffleVectorInst>(Phi->getOperand(0u));

  // Ensure result vectors are wider than the argument vector.
  auto *InputVT = cast<FixedVectorType>(Op->getType());
  auto *ResultVT = cast<FixedVectorType>(Shuf->getType());
  auto const InputNumElements = InputVT->getNumElements();

  if (InputNumElements >= ResultVT->getNumElements())
    return false;

  // Take the difference of the two shuffle masks at each index. Ignore poison
  // values at the same index in both masks.
  SmallVector<int, 16> NewMask;
  NewMask.reserve(Mask0.size());

  for (auto [M0, M1] : zip(Mask0, Mask1)) {
    if (M0 >= 0 && M1 >= 0)
      NewMask.push_back(M0 - M1);
    else if (M0 == -1 && M1 == -1)
      continue;
    else
      return false;
  }

  // Ensure all elements of the new mask are equal. If the difference between
  // the incoming mask elements is the same, the two must be constant offsets
  // of one another.
  if (NewMask.empty() || !all_equal(NewMask))
    return false;

  // Create new mask using difference of the two incoming masks.
  int MaskOffset = NewMask[0u];
  unsigned Index = (InputNumElements + MaskOffset) % InputNumElements;
  NewMask.clear();

  for (unsigned I = 0u; I < InputNumElements; ++I) {
    NewMask.push_back(Index);
    Index = (Index + 1u) % InputNumElements;
  }

  // Calculate costs for worst cases and compare.
  auto const Kind = TTI::SK_PermuteSingleSrc;
  auto OldCost =
      std::max(TTI.getShuffleCost(Kind, ResultVT, InputVT, CostKind, Mask0),
               TTI.getShuffleCost(Kind, ResultVT, InputVT, CostKind, Mask1));
  auto NewCost = TTI.getShuffleCost(Kind, InputVT, InputVT, CostKind, NewMask) +
                 TTI.getShuffleCost(Kind, ResultVT, InputVT, CostKind, Mask1);

  LLVM_DEBUG(dbgs() << "Found a phi of mergeable shuffles: " << I
                    << "\n  OldCost: " << OldCost << " vs NewCost: " << NewCost
                    << "\n");

  if (NewCost > OldCost)
    return false;

  // Create new shuffles and narrowed phi.
  auto Builder = IRBuilder(Shuf);
  Builder.SetCurrentDebugLocation(Shuf->getDebugLoc());
  auto *PoisonVal = PoisonValue::get(InputVT);
  auto *NewShuf0 = Builder.CreateShuffleVector(Op, PoisonVal, NewMask);
  Worklist.push(cast<Instruction>(NewShuf0));

  Builder.SetInsertPoint(Phi);
  Builder.SetCurrentDebugLocation(Phi->getDebugLoc());
  auto *NewPhi = Builder.CreatePHI(NewShuf0->getType(), 2u);
  NewPhi->addIncoming(NewShuf0, Phi->getIncomingBlock(0u));
  NewPhi->addIncoming(Op, Phi->getIncomingBlock(1u));

  Builder.SetInsertPoint(*NewPhi->getInsertionPointAfterDef());
  PoisonVal = PoisonValue::get(NewPhi->getType());
  auto *NewShuf1 = Builder.CreateShuffleVector(NewPhi, PoisonVal, Mask1);

  replaceValue(*Phi, *NewShuf1);
  return true;
}

/// This is the entry point for all transforms. Pass manager differences are
/// handled in the callers of this function.
bool VectorCombine::run() {
  if (DisableVectorCombine)
    return false;

  // Don't attempt vectorization if the target does not support vectors.
  if (!TTI.getNumberOfRegisters(TTI.getRegisterClassForType(/*Vector*/ true)))
    return false;

  LLVM_DEBUG(dbgs() << "\n\nVECTORCOMBINE on " << F.getName() << "\n");

  auto FoldInst = [this](Instruction &I) {
    Builder.SetInsertPoint(&I);
    bool IsVectorType = isa<VectorType>(I.getType());
    bool IsFixedVectorType = isa<FixedVectorType>(I.getType());
    auto Opcode = I.getOpcode();

    LLVM_DEBUG(dbgs() << "VC: Visiting: " << I << '\n');

    // These folds should be beneficial regardless of when this pass is run
    // in the optimization pipeline.
    // The type checking is for run-time efficiency. We can avoid wasting time
    // dispatching to folding functions if there's no chance of matching.
    if (IsFixedVectorType) {
      switch (Opcode) {
      case Instruction::InsertElement:
        if (vectorizeLoadInsert(I))
          return true;
        break;
      case Instruction::ShuffleVector:
        if (widenSubvectorLoad(I))
          return true;
        break;
      default:
        break;
      }
    }

    // This transform works with scalable and fixed vectors
    // TODO: Identify and allow other scalable transforms
    if (IsVectorType) {
      if (scalarizeOpOrCmp(I))
        return true;
      if (scalarizeLoad(I))
        return true;
      if (scalarizeExtExtract(I))
        return true;
      if (foldInterleaveIntrinsics(I))
        return true;
      if (foldBitcastOfVPLoad(I))
        return true;
    }

    if (foldDeinterleaveIntrinsics(I))
      return true;

    if (Opcode == Instruction::Store)
      if (foldInsertElementsToStores(I))
        return true;

    // If this is an early pipeline invocation of this pass, we are done.
    if (TryEarlyFoldsOnly)
      return false;

    if (Opcode == Instruction::Call)
      if (foldBitOrderReverseAndSwap(I))
        return true;
    if (Opcode == Instruction::BitCast)
      if (foldBitOrderReverseAndSwap(I))
        return true;

    // Otherwise, try folds that improve codegen but may interfere with
    // early IR canonicalizations.
    // The type checking is for run-time efficiency. We can avoid wasting time
    // dispatching to folding functions if there's no chance of matching.
    if (IsFixedVectorType) {
      switch (Opcode) {
      case Instruction::InsertElement:
        if (foldInsExtFNeg(I))
          return true;
        if (foldInsExtBinop(I))
          return true;
        if (foldInsExtVectorToShuffle(I))
          return true;
        if (foldInsertScalarPartsToShuffle(I))
          return true;
        break;
      case Instruction::ShuffleVector:
        if (foldPermuteOfBinops(I))
          return true;
        if (foldShuffleOfBinops(I))
          return true;
        if (foldShuffleOfSelects(I))
          return true;
        if (foldShuffleOfCastops(I))
          return true;
        if (foldShuffleOfShuffles(I))
          return true;
        if (foldPermuteOfIntrinsic(I))
          return true;
        if (foldShufflesOfLengthChangingShuffles(I))
          return true;
        if (foldShuffleOfIntrinsics(I))
          return true;
        if (foldSelectShuffle(I))
          return true;
        if (foldShuffleToIdentity(I))
          return true;
        break;
      case Instruction::Load:
        if (shrinkLoadForShuffles(I))
          return true;
        break;
      case Instruction::BitCast:
        if (foldBitcastShuffle(I))
          return true;
        if (foldSelectsFromBitcast(I))
          return true;
        break;
      case Instruction::And:
      case Instruction::Or:
      case Instruction::Xor:
        if (foldBitOpOfCastops(I))
          return true;
        if (foldBitOpOfCastConstant(I))
          return true;
        break;
      case Instruction::PHI:
        if (shrinkPhiOfShuffles(I))
          return true;
        break;
      default:
        if (shrinkType(I))
          return true;
        break;
      }
    } else {
      switch (Opcode) {
      case Instruction::Call:
        if (foldShuffleFromReductions(I))
          return true;
        if (foldCastFromReductions(I))
          return true;
        break;
      case Instruction::ExtractElement:
        if (foldShuffleChainsToReduce(I))
          return true;
        break;
      case Instruction::ICmp:
        if (foldSignBitReductionCmp(I))
          return true;
        if (foldICmpEqZeroVectorReduce(I))
          return true;
        if (foldReductionZeroTest(I))
          return true;
        if (foldEquivalentReductionCmp(I))
          return true;
        if (foldReduceAddCmpZero(I))
          return true;
        [[fallthrough]];
      case Instruction::FCmp:
        if (foldExtractExtract(I))
          return true;
        break;
      case Instruction::Or:
        if (foldConcatOfBoolMasks(I))
          return true;
        [[fallthrough]];
      default:
        if (Instruction::isBinaryOp(Opcode)) {
          if (foldExtractExtract(I))
            return true;
          if (foldExtractedCmps(I))
            return true;
          if (foldBinopOfReductions(I))
            return true;
        }
        break;
      }
    }
    return false;
  };

  bool MadeChange = false;
  for (BasicBlock &BB : F) {
    // Ignore unreachable basic blocks.
    if (!DT.isReachableFromEntry(&BB))
      continue;
    // Use early increment range so that we can erase instructions in loop.
    // make_early_inc_range is not applicable here, as the next iterator may
    // be invalidated by RecursivelyDeleteTriviallyDeadInstructions.
    // We manually maintain the next instruction and update it when it is about
    // to be deleted.
    Instruction *I = &BB.front();
    while (I) {
      NextInst = I->getNextNode();
      if (!I->isDebugOrPseudoInst())
        MadeChange |= FoldInst(*I);
      I = NextInst;
    }
  }

  NextInst = nullptr;

  while (!Worklist.isEmpty()) {
    Instruction *I = Worklist.removeOne();
    if (!I)
      continue;

    if (isInstructionTriviallyDead(I)) {
      eraseInstruction(*I);
      continue;
    }

    MadeChange |= FoldInst(*I);
  }

  return MadeChange;
}

PreservedAnalyses VectorCombinePass::run(Function &F,
                                         FunctionAnalysisManager &FAM) {
  auto &AC = FAM.getResult<AssumptionAnalysis>(F);
  TargetTransformInfo &TTI = FAM.getResult<TargetIRAnalysis>(F);
  DominatorTree &DT = FAM.getResult<DominatorTreeAnalysis>(F);
  AAResults &AA = FAM.getResult<AAManager>(F);
  const DataLayout *DL = &F.getDataLayout();
  TTI::TargetCostKind CostKind =
      F.hasOptSize() ? TTI::TCK_CodeSize : TTI::TCK_RecipThroughput;
  VectorCombine Combiner(F, TTI, DT, AA, AC, DL, CostKind, TryEarlyFoldsOnly);
  if (!Combiner.run())
    return PreservedAnalyses::all();
  PreservedAnalyses PA;
  PA.preserveSet<CFGAnalyses>();
  return PA;
}