llvm/lib/Transforms/Vectorize/SLPVectorizer/SLPMemoryUtils.cpp

//===- SLPMemoryUtils.cpp - SLP pointer/stride helpers --------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "SLPMemoryUtils.h"
#include "SLPCompatibilityAnalysis.h"
#include "SLPCostAnalysis.h"
#include "SLPTypeUtils.h"
#include "SLPUtils.h"

#include "llvm/ADT/APInt.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Analysis/Loads.h"
#include "llvm/Analysis/LoopAccessAnalysis.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/Support/InstructionCost.h"

#include <algorithm>
#include <limits>
#include <optional>
#include <set>
#include <tuple>
#include <utility>

using namespace llvm;

namespace llvm::slpvectorizer {

Value *createWidenedStridedCast(IRBuilderBase &Builder, Value *V, Type *DstTy,
                                const DataLayout &DL) {
  bool ToPtr = cast<VectorType>(DstTy)->getElementType()->isPointerTy();
  if (ToPtr == cast<VectorType>(V->getType())->getElementType()->isPointerTy())
    return Builder.CreateBitOrPointerCast(V, DstTy);
  if (ToPtr)
    return Builder.CreateIntToPtr(
        Builder.CreateBitCast(V, DL.getIntPtrType(DstTy)), DstTy);
  return Builder.CreateBitCast(
      Builder.CreatePtrToInt(V, DL.getIntPtrType(V->getType())), DstTy);
}

ConstantInt *getStrideBytesIfConstant(Value *Stride, Type *ScalarTy,
                                      const DataLayout &DL, bool IsReverse) {
  auto *CI = dyn_cast_or_null<ConstantInt>(Stride);
  if (!CI)
    return nullptr;

  uint64_t ElementSize = DL.getTypeAllocSize(ScalarTy).getFixedValue();
  APInt Bytes = CI->getValue() * ElementSize;
  return ConstantInt::get(CI->getContext(), IsReverse ? -Bytes : Bytes);
}

bool arePointersCompatible(Value *Ptr1, Value *Ptr2,
                           const TargetLibraryInfo &TLI, unsigned MaxDepth,
                           bool CompareOpcodes) {
  if (getUnderlyingObject(Ptr1, MaxDepth) !=
      getUnderlyingObject(Ptr2, MaxDepth))
    return false;
  auto *GEP1 = dyn_cast<GetElementPtrInst>(Ptr1);
  auto *GEP2 = dyn_cast<GetElementPtrInst>(Ptr2);
  return (!GEP1 || GEP1->getNumOperands() == 2) &&
         (!GEP2 || GEP2->getNumOperands() == 2) &&
         (((!GEP1 || isConstant(GEP1->getOperand(1))) &&
           (!GEP2 || isConstant(GEP2->getOperand(1)))) ||
          !CompareOpcodes ||
          (GEP1 && GEP2 &&
           getSameOpcode({GEP1->getOperand(1), GEP2->getOperand(1)}, TLI)));
}

/// Calculates minimal alignment as a common alignment.
template <typename T> Align computeCommonAlignment(ArrayRef<Value *> VL) {
  Align CommonAlignment = cast<T>(VL.consume_front())->getAlign();
  for (Value *V : VL)
    CommonAlignment = std::min(CommonAlignment, cast<T>(V)->getAlign());
  return CommonAlignment;
}

template Align computeCommonAlignment<LoadInst>(ArrayRef<Value *>);
template Align computeCommonAlignment<StoreInst>(ArrayRef<Value *>);

const SCEV *calculateRtStride(ArrayRef<Value *> PointerOps, Type *ElemTy,
                              const DataLayout &DL, ScalarEvolution &SE,
                              SmallVectorImpl<unsigned> &SortedIndices) {
  SmallVector<const SCEV *> SCEVs;
  const SCEV *PtrSCEVLowest = nullptr;
  const SCEV *PtrSCEVHighest = nullptr;
  // Find lower/upper pointers from the PointerOps (i.e. with lowest and highest
  // addresses).
  for (Value *Ptr : PointerOps) {
    const SCEV *PtrSCEV = SE.getSCEV(Ptr);
    if (!PtrSCEV)
      return nullptr;
    SCEVs.push_back(PtrSCEV);
    if (!PtrSCEVLowest && !PtrSCEVHighest) {
      PtrSCEVLowest = PtrSCEVHighest = PtrSCEV;
      continue;
    }
    const SCEV *Diff = SE.getMinusSCEV(PtrSCEV, PtrSCEVLowest);
    if (isa<SCEVCouldNotCompute>(Diff))
      return nullptr;
    if (Diff->isNonConstantNegative()) {
      PtrSCEVLowest = PtrSCEV;
      continue;
    }
    const SCEV *Diff1 = SE.getMinusSCEV(PtrSCEVHighest, PtrSCEV);
    if (isa<SCEVCouldNotCompute>(Diff1))
      return nullptr;
    if (Diff1->isNonConstantNegative()) {
      PtrSCEVHighest = PtrSCEV;
      continue;
    }
  }
  // Dist = PtrSCEVHighest - PtrSCEVLowest;
  const SCEV *Dist = SE.getMinusSCEV(PtrSCEVHighest, PtrSCEVLowest);
  if (isa<SCEVCouldNotCompute>(Dist))
    return nullptr;
  int Size = DL.getTypeStoreSize(ElemTy);
  auto TryGetStride = [&](const SCEV *Dist,
                          const SCEV *Multiplier) -> const SCEV * {
    if (const auto *M = dyn_cast<SCEVMulExpr>(Dist)) {
      if (M->getOperand(0) == Multiplier)
        return M->getOperand(1);
      if (M->getOperand(1) == Multiplier)
        return M->getOperand(0);
      return nullptr;
    }
    if (Multiplier == Dist)
      return SE.getConstant(Dist->getType(), 1);
    return SE.getUDivExactExpr(Dist, Multiplier);
  };
  // Stride_in_elements = Dist / element_size * (num_elems - 1).
  const SCEV *Stride = nullptr;
  if (Size != 1 || SCEVs.size() > 1) {
    const SCEV *Sz = SE.getConstant(Dist->getType(), Size * (SCEVs.size() - 1));
    Stride = TryGetStride(Dist, Sz);
    if (!Stride)
      return nullptr;
  }
  if (!Stride || isa<SCEVConstant>(Stride))
    return nullptr;
  // Iterate through all pointers and check if all distances are
  // unique multiple of Stride.
  using DistOrdPair = std::pair<int64_t, int>;
  auto Compare = llvm::less_first();
  std::set<DistOrdPair, decltype(Compare)> Offsets(Compare);
  bool IsConsecutive = true;
  for (const auto [Idx, PtrSCEV] : enumerate(SCEVs)) {
    unsigned Dist = 0;
    if (PtrSCEV != PtrSCEVLowest) {
      const SCEV *Diff = SE.getMinusSCEV(PtrSCEV, PtrSCEVLowest);
      const SCEV *Coeff = TryGetStride(Diff, Stride);
      if (!Coeff)
        return nullptr;
      const auto *SC = dyn_cast<SCEVConstant>(Coeff);
      if (!SC || isa<SCEVCouldNotCompute>(SC))
        return nullptr;
      if (!SE.getMinusSCEV(PtrSCEV, SE.getAddExpr(PtrSCEVLowest,
                                                  SE.getMulExpr(Stride, SC)))
               ->isZero())
        return nullptr;
      Dist = SC->getAPInt().getZExtValue();
    }
    // If the strides are not the same or repeated, we can't vectorize.
    if ((Dist / Size) * Size != Dist || (Dist / Size) >= SCEVs.size())
      return nullptr;
    auto Res = Offsets.emplace(Dist, Idx);
    if (!Res.second)
      return nullptr;
    // Consecutive order if the inserted element is the last one.
    IsConsecutive = IsConsecutive && std::next(Res.first) == Offsets.end();
  }
  SortedIndices.clear();
  if (!IsConsecutive) {
    // Fill SortedIndices array only if it is non-consecutive.
    SortedIndices.resize(PointerOps.size());
    for (const auto [Idx, Pair] : enumerate(Offsets))
      SortedIndices[Idx] = Pair.second;
  }
  return Stride;
}

std::optional<int64_t>
SLPAnalysisCache::getPointersDiff(Type *ElemTyA, Value *PtrA, Type *ElemTyB,
                                  Value *PtrB, const DataLayout &DL,
                                  ScalarEvolution &SE, bool StrictCheck,
                                  bool CheckType) {
  if (PtrA == PtrB)
    return 0;
  unsigned Flags = unsigned(StrictCheck) | (unsigned(CheckType) << 1);
  PtrsDiffKeyT Key(ElemTyA, PtrA, ElemTyB, PtrB, Flags);
  auto [It, Inserted] = PtrsDiffCache.try_emplace(Key);
  if (!Inserted)
    return It->second;
  std::optional<int64_t> Diff = llvm::getPointersDiff(
      ElemTyA, PtrA, ElemTyB, PtrB, DL, SE, StrictCheck, CheckType);
  // Set before the reverse insertion below, which may rehash.
  It->second = Diff;
  // The distance is counted in elements of ElemTyA, so querying the pair the
  // other way round gives the negation, provided both elements have the same
  // size. A failed query is recorded for this order only.
  if (Diff && *Diff != std::numeric_limits<int64_t>::min() &&
      DL.getTypeStoreSize(ElemTyA) == DL.getTypeStoreSize(ElemTyB))
    PtrsDiffCache.try_emplace(PtrsDiffKeyT(ElemTyB, PtrB, ElemTyA, PtrA, Flags),
                              -*Diff);
  return Diff;
}

bool SLPAnalysisCache::isGuaranteedNotToBePoison(const Value *V,
                                                 AssumptionCache *AC) {
  auto [It, Inserted] = NotPoisonCache.try_emplace(V, false);
  if (Inserted)
    It->second = llvm::isGuaranteedNotToBePoison(V, AC);
  return It->second;
}

/// Builds compress-like mask for shuffles for the given \p PointerOps, ordered
/// with \p Order.
/// \return true if the mask represents strided access, false - otherwise.
static bool buildCompressMask(ArrayRef<Value *> PointerOps,
                              ArrayRef<unsigned> Order, Type *ScalarTy,
                              const DataLayout &DL, ScalarEvolution &SE,
                              SmallVectorImpl<int> &CompressMask,
                              SLPAnalysisCache &Cache) {
  const unsigned Sz = PointerOps.size();
  CompressMask.assign(Sz, PoisonMaskElem);
  // The first element always set.
  CompressMask[0] = 0;
  // Check if the mask represents strided access.
  std::optional<unsigned> Stride = 0;
  Value *Ptr0 = Order.empty() ? PointerOps.front() : PointerOps[Order.front()];
  for (unsigned I : seq<unsigned>(1, Sz)) {
    Value *Ptr = Order.empty() ? PointerOps[I] : PointerOps[Order[I]];
    std::optional<int64_t> OptPos =
        Cache.getPointersDiff(ScalarTy, Ptr0, ScalarTy, Ptr, DL, SE);
    if (!OptPos || OptPos > std::numeric_limits<unsigned>::max())
      return false;
    unsigned Pos = static_cast<unsigned>(*OptPos);
    CompressMask[I] = Pos;
    if (!Stride)
      continue;
    if (*Stride == 0) {
      *Stride = Pos;
      continue;
    }
    if (Pos != *Stride * I)
      Stride.reset();
  }
  return Stride.has_value();
}

/// Checks if the \p VL can be transformed to a (masked)load + compress or
/// (masked) interleaved load.
bool isMaskedLoadCompress(
    ArrayRef<Value *> VL, ArrayRef<Value *> PointerOps,
    ArrayRef<unsigned> Order, const TargetTransformInfo &TTI,
    const DataLayout &DL, ScalarEvolution &SE, AssumptionCache &AC,
    const DominatorTree &DT, const TargetLibraryInfo &TLI,
    const TargetTransformInfo::TargetCostKind CostKind,
    const function_ref<bool(Value *)> AreAllUsersVectorized, bool ReVec,
    bool &IsMasked, unsigned &InterleaveFactor,
    SmallVectorImpl<int> &CompressMask, VectorType *&LoadVecTy,
    SLPAnalysisCache &Cache) {
  InterleaveFactor = 0;
  Type *ScalarTy = VL.front()->getType();
  const size_t Sz = VL.size();
  auto *VecTy = cast<VectorType>(getWidenedType(ScalarTy, Sz));
  SmallVector<int> Mask;
  if (!Order.empty())
    inversePermutation(Order, Mask);
  // Check external uses.
  for (const auto [I, V] : enumerate(VL)) {
    if (AreAllUsersVectorized(V))
      continue;
    InstructionCost ExtractCost =
        TTI.getVectorInstrCost(Instruction::ExtractElement, VecTy, CostKind,
                               Mask.empty() ? I : Mask[I]);
    InstructionCost ScalarCost =
        TTI.getInstructionCost(cast<Instruction>(V), CostKind);
    if (ExtractCost <= ScalarCost)
      return false;
  }
  Value *Ptr0;
  Value *PtrN;
  if (Order.empty()) {
    Ptr0 = PointerOps.front();
    PtrN = PointerOps.back();
  } else {
    Ptr0 = PointerOps[Order.front()];
    PtrN = PointerOps[Order.back()];
  }
  std::optional<int64_t> Diff =
      Cache.getPointersDiff(ScalarTy, Ptr0, ScalarTy, PtrN, DL, SE);
  if (!Diff)
    return false;
  const size_t MaxRegSize =
      TTI.getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector)
          .getFixedValue();
  // Check for very large distances between elements.
  if (*Diff / Sz >= MaxRegSize / 8)
    return false;
  LoadVecTy = cast<FixedVectorType>(getWidenedType(ScalarTy, *Diff + 1));
  auto *LI = cast<LoadInst>(Order.empty() ? VL.front() : VL[Order.front()]);
  Align CommonAlignment = LI->getAlign();
  SimplifyQuery SQ(
      DL, &TLI, &DT, &AC,
      cast<LoadInst>(Order.empty() ? VL.back() : VL[Order.back()]));
  IsMasked = !isSafeToLoadUnconditionally(Ptr0, LoadVecTy, CommonAlignment, SQ);
  if (IsMasked && !TTI.isLegalMaskedLoad(LoadVecTy, CommonAlignment,
                                         LI->getPointerAddressSpace()))
    return false;
  // TODO: perform the analysis of each scalar load for better
  // safe-load-unconditionally analysis.
  bool IsStrided = buildCompressMask(PointerOps, Order, ScalarTy, DL, SE,
                                     CompressMask, Cache);
  assert(CompressMask.size() >= 2 && "At least two elements are required");
  SmallVector<Value *> OrderedPointerOps(PointerOps);
  if (!Order.empty())
    reorderScalars(OrderedPointerOps, Mask);
  auto [ScalarGEPCost, VectorGEPCost] =
      getGEPCosts(TTI, OrderedPointerOps, OrderedPointerOps.front(),
                  Instruction::Load, CostKind, ScalarTy, LoadVecTy);
  // The cost of scalar loads.
  InstructionCost ScalarLoadsCost =
      accumulate(VL, InstructionCost(),
                 [&](InstructionCost C, Value *V) {
                   return C + TTI.getInstructionCost(cast<Instruction>(V),
                                                     CostKind);
                 }) +
      ScalarGEPCost;
  APInt DemandedElts = APInt::getAllOnes(Sz);
  InstructionCost GatherCost =
      getScalarizationOverhead(TTI, ReVec, ScalarTy, VecTy, DemandedElts,
                               /*Insert=*/true,
                               /*Extract=*/false, CostKind) +
      ScalarLoadsCost;
  InstructionCost LoadCost = 0;
  if (IsMasked) {
    LoadCost = TTI.getMemIntrinsicInstrCost(
        MemIntrinsicCostAttributes(Intrinsic::masked_load, LoadVecTy,
                                   CommonAlignment,
                                   LI->getPointerAddressSpace()),
        CostKind);
  } else {
    LoadCost =
        TTI.getMemoryOpCost(Instruction::Load, LoadVecTy, CommonAlignment,
                            LI->getPointerAddressSpace(), CostKind,
                            TTI::getOperandInfo(LI->getPointerOperand()));
  }
  if (IsStrided && !IsMasked && Order.empty()) {
    // Check for potential segmented(interleaved) loads.
    VectorType *AlignedLoadVecTy = cast<VectorType>(getWidenedType(
        ScalarTy,
        getFullVectorNumberOfElements(TTI, ScalarTy, *Diff + 1, ReVec)));
    SimplifyQuery SQ(DL, &TLI, &DT, &AC, cast<LoadInst>(VL.back()));
    if (!isSafeToLoadUnconditionally(Ptr0, AlignedLoadVecTy, CommonAlignment,
                                     SQ))
      AlignedLoadVecTy = LoadVecTy;
    if (TTI.isLegalInterleavedAccessType(AlignedLoadVecTy, CompressMask[1],
                                         CommonAlignment,
                                         LI->getPointerAddressSpace())) {
      InstructionCost InterleavedCost =
          VectorGEPCost + TTI.getInterleavedMemoryOpCost(
                              Instruction::Load, AlignedLoadVecTy,
                              CompressMask[1], {}, CommonAlignment,
                              LI->getPointerAddressSpace(), CostKind, IsMasked);
      if (InterleavedCost < GatherCost) {
        InterleaveFactor = CompressMask[1];
        LoadVecTy = AlignedLoadVecTy;
        return true;
      }
    }
  }
  // Estimating the compression shuffle cost below can be extremely expensive
  // for a very wide LoadVecTy, which is split into a large number of vector
  // registers (see processShuffleMasks). The shuffle cost is always
  // non-negative, so if the load cost alone already reaches the gather cost the
  // masked-load-compress cannot be profitable. Bail out before the costly
  // shuffle cost estimation in that case.
  if (VectorGEPCost + LoadCost >= GatherCost)
    return false;
  // For REVEC the scalars are vectors, and the emitted shuffle uses the
  // compress mask expanded to the individual elements. Cost that mask.
  SmallVector<int> CostMask(CompressMask.begin(), CompressMask.end());
  transformScalarShuffleIndiciesToVector(getNumElements(ScalarTy), CostMask);
  InstructionCost CompressCost = getShuffleCost(TTI, TTI::SK_PermuteSingleSrc,
                                                LoadVecTy, CostKind, CostMask);
  if (!Order.empty()) {
    SmallVector<int> NewMask(Sz, PoisonMaskElem);
    for (unsigned I : seq<unsigned>(Sz)) {
      NewMask[I] = CompressMask[Mask[I]];
    }
    CompressMask.swap(NewMask);
  }
  InstructionCost TotalVecCost = VectorGEPCost + LoadCost + CompressCost;
  return TotalVecCost < GatherCost;
}

/// Checks if the \p VL can be transformed to a (masked)load + compress or
/// (masked) interleaved load.
bool isMaskedLoadCompress(
    ArrayRef<Value *> VL, ArrayRef<Value *> PointerOps,
    ArrayRef<unsigned> Order, const TargetTransformInfo &TTI,
    const DataLayout &DL, ScalarEvolution &SE, AssumptionCache &AC,
    const DominatorTree &DT, const TargetLibraryInfo &TLI,
    const TargetTransformInfo::TargetCostKind CostKind,
    const function_ref<bool(Value *)> AreAllUsersVectorized, bool ReVec,
    SLPAnalysisCache &Cache) {
  bool IsMasked;
  unsigned InterleaveFactor;
  SmallVector<int> CompressMask;
  VectorType *LoadVecTy;
  return isMaskedLoadCompress(VL, PointerOps, Order, TTI, DL, SE, AC, DT, TLI,
                              CostKind, AreAllUsersVectorized, ReVec, IsMasked,
                              InterleaveFactor, CompressMask, LoadVecTy, Cache);
}

bool canWidenLoadsOverUndefLanes(ArrayRef<Value *> VL, const DataLayout &DL,
                                 ScalarEvolution &SE, AssumptionCache &AC,
                                 const DominatorTree &DT,
                                 const TargetLibraryInfo &TLI) {
  if (count_if(VL, IsaPred<LoadInst>) < 2)
    return false;
  const auto *FirstIt = find_if(VL, IsaPred<LoadInst>);
  const unsigned FirstLane = std::distance(VL.begin(), FirstIt);
  auto *FirstLI = cast<LoadInst>(*FirstIt);
  // The wide load may read the memory with a different tag.
  if (FirstLI->getFunction()->hasFnAttribute(Attribute::SanitizeMemTag))
    return false;
  Type *ScalarTy = FirstLI->getType();
  Value *Ptr0 = FirstLI->getPointerOperand();
  LoadInst *LastLI = FirstLI;
  for (auto [Lane, V] : enumerate(VL)) {
    if (isa<UndefValue>(V))
      continue;
    auto *LI = dyn_cast<LoadInst>(V);
    if (!LI || !LI->isSimple() || mustSuppressSpeculation(*LI))
      return false;
    std::optional<int64_t> Dist =
        getPointersDiff(ScalarTy, Ptr0, ScalarTy, LI->getPointerOperand(), DL,
                        SE, /*StrictCheck=*/true);
    if (!Dist || *Dist != static_cast<int64_t>(Lane) - FirstLane)
      return false;
    if (LastLI->comesBefore(LI))
      LastLI = LI;
  }
  // The wide load starts before the first load, so check the range as the one
  // starting at the underlying object.
  const uint64_t EltSize = DL.getTypeAllocSize(ScalarTy).getFixedValue();
  const unsigned BitWidth = DL.getIndexTypeSizeInBits(Ptr0->getType());
  APInt Offset(BitWidth, 0);
  const Value *Base = Ptr0->stripAndAccumulateConstantOffsets(
      DL, Offset, /*AllowNonInbounds=*/true);
  bool Overflow = false;
  APInt Start = Offset.ssub_ov(APInt(BitWidth, FirstLane * EltSize), Overflow);
  if (Overflow || Start.isNegative())
    return false;
  APInt End = Start.sadd_ov(APInt(BitWidth, VL.size() * EltSize), Overflow);
  return !Overflow &&
         isDereferenceableAndAlignedPointer(
             Base, Align(1), End, SimplifyQuery(DL, &TLI, &DT, &AC, LastLI));
}

uint64_t getVectorLoadShift(ArrayRef<Value *> Scalars) {
  const auto *It = find_if(Scalars, IsaPred<LoadInst>);
  auto *LI = cast<LoadInst>(*It);
  return std::distance(Scalars.begin(), It) *
         LI->getDataLayout().getTypeAllocSize(LI->getType()).getFixedValue();
}

/// Checks if the stores \p VL with pointers \p PointerOps can be lowered as a
/// single masked store. On success \p StoreVecTy is the widened store type and
/// \p ReuseShuffleIndices is the expand mask that places each stored value at
/// its element offset from the base (poison in the gaps).
bool isMaskedStoreCompress(ArrayRef<Value *> VL, ArrayRef<Value *> PointerOps,
                           ArrayRef<unsigned> Order,
                           const TargetTransformInfo &TTI, const DataLayout &DL,
                           ScalarEvolution &SE, Align CommonAlignment,
                           SmallVectorImpl<int> &ReuseShuffleIndices,
                           FixedVectorType *&StoreVecTy,
                           SLPAnalysisCache &Cache) {
  Type *ScalarTy = cast<StoreInst>(VL.front())->getValueOperand()->getType();
  const size_t Sz = VL.size();
  // Only simple scalar element types are supported.
  if (Sz < 2 || (!ScalarTy->isIntOrPtrTy() && !ScalarTy->isFloatingPointTy()))
    return false;
  Value *Ptr0 = Order.empty() ? PointerOps.front() : PointerOps[Order.front()];
  Value *PtrN = Order.empty() ? PointerOps.back() : PointerOps[Order.back()];
  std::optional<int64_t> Diff =
      Cache.getPointersDiff(ScalarTy, Ptr0, ScalarTy, PtrN, DL, SE);
  if (!Diff || *Diff <= 0)
    return false;
  // Avoid widened vectors with very large gaps between the stored elements.
  const unsigned MaxRegSize =
      TTI.getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector)
          .getFixedValue();
  const unsigned ScalarBits = DL.getTypeSizeInBits(ScalarTy).getFixedValue();
  if (ScalarBits == 0 ||
      static_cast<uint64_t>(*Diff) / Sz >= MaxRegSize / ScalarBits)
    return false;
  StoreVecTy = cast<FixedVectorType>(getWidenedType(ScalarTy, *Diff + 1));
  unsigned AS = cast<StoreInst>(VL.front())->getPointerAddressSpace();
  if (!TTI.isLegalMaskedStore(StoreVecTy, CommonAlignment, AS,
                              TTI::ConstantMask))
    return false;
  // Build the expand mask: store I (in address-sorted order) is placed at its
  // element offset from the base, other widened lanes are poison.
  ReuseShuffleIndices.assign(*Diff + 1, PoisonMaskElem);
  int64_t Prev = -1;
  for (unsigned I : seq<unsigned>(Sz)) {
    Value *Ptr = Order.empty() ? PointerOps[I] : PointerOps[Order[I]];
    std::optional<int64_t> Off =
        Cache.getPointersDiff(ScalarTy, Ptr0, ScalarTy, Ptr, DL, SE);
    if (!Off || *Off <= Prev || *Off > *Diff)
      return false;
    ReuseShuffleIndices[*Off] = static_cast<int>(I);
    Prev = *Off;
  }
  return true;
}

bool clusterSortPtrAccesses(ArrayRef<Value *> VL, ArrayRef<BasicBlock *> BBs,
                            Type *ElemTy, const DataLayout &DL,
                            ScalarEvolution &SE, unsigned MaxDepth,
                            SmallVectorImpl<unsigned> &SortedIndices,
                            SLPAnalysisCache &Cache) {
  assert(
      all_of(VL, [](const Value *V) { return V->getType()->isPointerTy(); }) &&
      "Expected list of pointer operands.");
  // Map from bases to a vector of (Ptr, Offset, OrigIdx), which we insert each
  // Ptr into, sort and return the sorted indices with values next to one
  // another.
  SmallMapVector<
      std::pair<BasicBlock *, Value *>,
      SmallVector<SmallVector<std::tuple<Value *, int64_t, unsigned>>>, 8>
      Bases;
  Bases
      .try_emplace(std::make_pair(BBs.front(),
                                  getUnderlyingObject(VL.front(), MaxDepth)))
      .first->second.emplace_back()
      .emplace_back(VL.front(), 0U, 0U);

  SortedIndices.clear();
  for (auto [Cnt, Ptr] : enumerate(VL.drop_front())) {
    auto Key = std::make_pair(BBs[Cnt + 1], getUnderlyingObject(Ptr, MaxDepth));
    bool Found = any_of(Bases.try_emplace(Key).first->second,
                        [&, &Cnt = Cnt, &Ptr = Ptr](auto &Base) {
                          std::optional<int64_t> Diff = Cache.getPointersDiff(
                              ElemTy, std::get<0>(Base.front()), ElemTy, Ptr,
                              DL, SE, /*StrictCheck=*/true);
                          if (!Diff)
                            return false;

                          Base.emplace_back(Ptr, *Diff, Cnt + 1);
                          return true;
                        });

    if (!Found) {
      // If we haven't found enough to usefully cluster, return early.
      if (Bases.size() > VL.size() / 2 - 1)
        return false;

      // Not found already - add a new Base
      Bases.find(Key)->second.emplace_back().emplace_back(Ptr, 0, Cnt + 1);
    }
  }

  if (Bases.size() == VL.size())
    return false;

  if (Bases.size() == 1 && (Bases.front().second.size() == 1 ||
                            Bases.front().second.size() == VL.size()))
    return false;

  // For each of the bases sort the pointers by Offset and check if any of the
  // base become consecutively allocated.
  auto ComparePointers = [MaxDepth](Value *Ptr1, Value *Ptr2) {
    SmallPtrSet<Value *, 13> FirstPointers;
    SmallPtrSet<Value *, 13> SecondPointers;
    Value *P1 = Ptr1;
    Value *P2 = Ptr2;
    unsigned Depth = 0;
    while (!FirstPointers.contains(P2) && !SecondPointers.contains(P1)) {
      if (P1 == P2 || Depth > MaxDepth)
        return false;
      FirstPointers.insert(P1);
      SecondPointers.insert(P2);
      P1 = getUnderlyingObject(P1, /*MaxLookup=*/1);
      P2 = getUnderlyingObject(P2, /*MaxLookup=*/1);
      ++Depth;
    }
    assert((FirstPointers.contains(P2) || SecondPointers.contains(P1)) &&
           "Unable to find matching root.");
    return FirstPointers.contains(P2) && !SecondPointers.contains(P1);
  };
  for (auto &Base : Bases) {
    for (auto &Vec : Base.second) {
      if (Vec.size() > 1) {
        stable_sort(Vec, llvm::less_second());
        int64_t InitialOffset = std::get<1>(Vec[0]);
        bool AnyConsecutive =
            all_of(enumerate(Vec), [InitialOffset](const auto &P) {
              return std::get<1>(P.value()) ==
                     int64_t(P.index()) + InitialOffset;
            });
        // Fill SortedIndices array only if it looks worth-while to sort the
        // ptrs.
        if (!AnyConsecutive)
          return false;
      }
    }
    stable_sort(Base.second, [&](const auto &V1, const auto &V2) {
      return ComparePointers(std::get<0>(V1.front()), std::get<0>(V2.front()));
    });
  }

  for (auto &T : Bases)
    for (const auto &Vec : T.second)
      for (const auto &P : Vec)
        SortedIndices.push_back(std::get<2>(P));

  assert(SortedIndices.size() == VL.size() &&
         "Expected SortedIndices to be the size of VL");
  return true;
}

} // namespace llvm::slpvectorizer