//===- VPlanUtils.cpp - VPlan-related utilities ---------------------------===//
//
// 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 "VPlanUtils.h"
#include "LoopVectorizationPlanner.h"
#include "VPlanAnalysis.h"
#include "VPlanCFG.h"
#include "VPlanDominatorTree.h"
#include "VPlanPatternMatch.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Analysis/BlockFrequencyInfoImpl.h"
#include "llvm/Analysis/BranchProbabilityInfo.h"
#include "llvm/Analysis/InstSimplifyFolder.h"
#include "llvm/Analysis/LoopAccessAnalysis.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/MemoryLocation.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Analysis/ScalarEvolutionPatternMatch.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/ProfDataUtils.h"
#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"

using namespace llvm;
using namespace llvm::VPlanPatternMatch;
using namespace llvm::SCEVPatternMatch;

bool vputils::onlyFirstLaneUsed(const VPValue *Def) {
  return all_of(Def->users(),
                [Def](const VPUser *U) { return U->usesFirstLaneOnly(Def); });
}

bool vputils::onlyFirstPartUsed(const VPValue *Def) {
  return all_of(Def->users(),
                [Def](const VPUser *U) { return U->usesFirstPartOnly(Def); });
}

bool vputils::onlyScalarValuesUsed(const VPValue *Def) {
  return all_of(Def->users(),
                [Def](const VPUser *U) { return U->usesScalars(Def); });
}

VPValue *vputils::getOrCreateVPValueForSCEVExpr(VPlan &Plan, const SCEV *Expr) {
  if (auto *E = dyn_cast<SCEVConstant>(Expr))
    return Plan.getOrAddLiveIn(E->getValue());
  // Skip SCEV expansion if Expr is a SCEVUnknown wrapping a non-instruction
  // value. Otherwise the value may be defined in a loop and using it directly
  // will break LCSSA form. The SCEV expansion takes care of preserving LCSSA
  // form.
  auto *U = dyn_cast<SCEVUnknown>(Expr);
  if (U && !isa<Instruction>(U->getValue()))
    return Plan.getOrAddLiveIn(U->getValue());
  auto *Expanded = new VPExpandSCEVRecipe(Expr);
  VPBasicBlock *EntryVPBB = Plan.getEntry();
  auto Iter = EntryVPBB->getFirstNonPhi();
  while (Iter != EntryVPBB->end() && isa<VPIRInstruction>(*Iter))
    ++Iter;
  EntryVPBB->insert(Expanded, Iter);
  return Expanded;
}

/// Returns true if \p V being poison is guaranteed to trigger UB because it
/// propagates to the address of a memory recipe.
static bool poisonGuaranteesUB(const VPValue *V) {
  SmallPtrSet<const VPValue *, 8> Visited;
  SmallVector<const VPValue *, 16> Worklist;

  auto PropagatesPoisonFromRecipeOp = [](const VPRecipeBase *R) {
    if (!isa<VPSingleDefRecipe>(R))
      return false;
    unsigned Opcode = vputils::getOpcode(R->getVPSingleValue());
    return Instruction::isCast(Opcode) || Opcode == Instruction::GetElementPtr;
  };

  Worklist.push_back(V);

  while (!Worklist.empty()) {
    const VPValue *Current = Worklist.pop_back_val();
    if (!Visited.insert(Current).second)
      continue;

    for (VPUser *U : Current->users()) {
      // Check if Current is used as an address operand for load/store.
      auto *R = cast<VPRecipeBase>(U);
      if (auto *MemR = dyn_cast<VPWidenMemoryRecipe>(R)) {
        if (MemR->getAddr() == Current)
          return true;
        continue;
      }
      if (auto *Rep = dyn_cast<VPReplicateRecipe>(U)) {
        unsigned Opcode = Rep->getOpcode();
        if ((Opcode == Instruction::Load && Rep->getOperand(0) == Current) ||
            (Opcode == Instruction::Store && Rep->getOperand(1) == Current))
          return true;
      }

      // Check if poison propagates through this recipe to any of its users.
      for (const VPValue *Op : R->operands()) {
        if (Op == Current && PropagatesPoisonFromRecipeOp(R)) {
          Worklist.push_back(R->getVPSingleValue());
          break;
        }
      }
    }
  }

  return false;
}

GEPNoWrapFlags vputils::getGEPFlagsForPtr(VPValue *Ptr) {
  // Like IR stripPointerCasts, look through GEPs with all-zero indices and
  // casts to find a root GEP VPInstruction.
  while (auto *PtrVPI = dyn_cast<VPInstruction>(Ptr)) {
    unsigned Opcode = PtrVPI->getOpcode();
    if (Opcode == Instruction::GetElementPtr) {
      if (!all_of(drop_begin(PtrVPI->operands()), match_fn(m_ZeroInt())))
        return PtrVPI->getGEPNoWrapFlags();
      Ptr = PtrVPI->getOperand(0);
      continue;
    }
    if (Opcode != Instruction::BitCast && Opcode != Instruction::AddrSpaceCast)
      break;
    Ptr = PtrVPI->getOperand(0);
  }
  return GEPNoWrapFlags::none();
}

const SCEV *vputils::getSCEVExprForVPValue(const VPValue *V,
                                           PredicatedScalarEvolution &PSE,
                                           const Loop *L) {
  ScalarEvolution &SE = *PSE.getSE();
  if (auto *RV = dyn_cast<VPRegionValue>(V)) {
    assert(RV == RV->getDefiningRegion()->getCanonicalIV() &&
           "RegionValue must be canonical IV");
    if (!L)
      return SE.getCouldNotCompute();
    return SE.getAddRecExpr(SE.getZero(RV->getType()), SE.getOne(RV->getType()),
                            L, SCEV::FlagNone);
  }

  if (isa<VPIRValue, VPSymbolicValue>(V)) {
    Value *LiveIn = V->getUnderlyingValue();
    if (LiveIn && SE.isSCEVable(LiveIn->getType()))
      return SE.getSCEV(LiveIn);
    return SE.getCouldNotCompute();
  }

  // Helper to create SCEVs for binary and unary operations.
  auto CreateSCEV = [&](ArrayRef<VPValue *> Ops,
                        function_ref<const SCEV *(ArrayRef<SCEVUse>)> CreateFn)
      -> const SCEV * {
    SmallVector<SCEVUse, 2> SCEVOps;
    for (VPValue *Op : Ops) {
      const SCEV *S = getSCEVExprForVPValue(Op, PSE, L);
      if (isa<SCEVCouldNotCompute>(S))
        return SE.getCouldNotCompute();
      SCEVOps.push_back(S);
    }
    return PSE.getPredicatedSCEV(CreateFn(SCEVOps));
  };

  VPValue *LHSVal, *RHSVal;
  if (match(V, m_Add(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getAddExpr(Ops[0], Ops[1], SCEV::FlagNone, 0);
    });
  if (match(V, m_BinaryOr(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    if (cast<VPRecipeWithIRFlags>(V->getDefiningRecipe())->isDisjoint())
      return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
        return SE.getAddExpr(Ops[0], Ops[1], SCEV::FlagNone, 0);
      });
  if (match(V, m_Sub(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getMinusSCEV(Ops[0], Ops[1], SCEV::FlagNone, 0);
    });
  if (match(V, m_Not(m_VPValue(LHSVal)))) {
    // not X = xor X, -1 = -1 - X
    return CreateSCEV({LHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getMinusSCEV(SE.getMinusOne(Ops[0]->getType()), Ops[0]);
    });
  }
  if (match(V, m_Mul(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getMulExpr(Ops[0], Ops[1], SCEV::FlagNone, 0);
    });
  // Handle shl by constant: x << c is equivalent to x * (1 << c). A shift
  // amount >= the bit width produces poison; do not rewrite it, as
  // getPowerOfTwo requires the power to be in range.
  uint64_t ShiftAmt;
  if (match(V, m_Shl(m_VPValue(LHSVal), m_ConstantInt(ShiftAmt))) &&
      ShiftAmt < LHSVal->getScalarType()->getScalarSizeInBits())
    return CreateSCEV(LHSVal, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getMulExpr(Ops[0],
                           SE.getPowerOfTwo(Ops[0]->getType(), ShiftAmt));
    });
  if (match(V, m_LShr(m_VPValue(LHSVal), m_ConstantInt(ShiftAmt)))) {
    Type *Ty = V->getScalarType();
    if (ShiftAmt < SE.getTypeSizeInBits(Ty))
      return CreateSCEV(LHSVal, [&](ArrayRef<SCEVUse> Ops) {
        return SE.getUDivExpr(Ops[0], SE.getPowerOfTwo(Ty, ShiftAmt));
      });
  }
  if (match(V, m_UDiv(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getUDivExpr(Ops[0], Ops[1]);
    });
  if (match(V, m_URem(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getURemExpr(Ops[0], Ops[1]);
    });
  // A SDiv with non-negative operands is equivalent to an UDiv.
  if (match(V, m_SDiv(m_VPValue(LHSVal), m_VPValue(RHSVal)))) {
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      if (!SE.isKnownNonNegative(Ops[0]) || !SE.isKnownNonNegative(Ops[1]))
        return SE.getCouldNotCompute();
      return SE.getUDivExpr(Ops[0], Ops[1]);
    });
  }
  // A SRem with non-negative operands is equivalent to an URem.
  if (match(V, m_SRem(m_VPValue(LHSVal), m_VPValue(RHSVal)))) {
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      if (!SE.isKnownNonNegative(Ops[0]) || !SE.isKnownNonNegative(Ops[1]))
        return SE.getCouldNotCompute();
      return SE.getURemExpr(Ops[0], Ops[1]);
    });
  }
  // Handle AND with constant mask: x & (2^n - 1) can be represented as x % 2^n.
  const APInt *Mask;
  if (match(V, m_c_BinaryAnd(m_VPValue(LHSVal), m_APInt(Mask))) &&
      (*Mask + 1).isPowerOf2())
    return CreateSCEV({LHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getURemExpr(Ops[0], SE.getConstant(*Mask + 1));
    });
  // SCEV models ptrtoaddr, but not ptrtoint, mirroring createSCEV.
  if (match(V, m_PtrToAddr(m_VPValue(LHSVal))))
    return CreateSCEV({LHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getPtrToAddrExpr(Ops[0]);
    });
  if (match(V, m_Trunc(m_VPValue(LHSVal)))) {
    Type *DestTy = V->getScalarType();
    return CreateSCEV({LHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getTruncateExpr(Ops[0], DestTy);
    });
  }
  if (match(V, m_ZExt(m_VPValue(LHSVal)))) {
    Type *DestTy = V->getScalarType();
    return CreateSCEV({LHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getZeroExtendExpr(Ops[0], DestTy);
    });
  }
  if (match(V, m_SExt(m_VPValue(LHSVal)))) {
    Type *DestTy = V->getScalarType();

    // Mirror SCEV's createSCEV handling for sext(sub nsw): push sign extension
    // onto the operands before computing the subtraction.
    VPValue *SubLHS, *SubRHS;
    auto *SubR = dyn_cast<VPRecipeWithIRFlags>(LHSVal);
    if (match(LHSVal, m_Sub(m_VPValue(SubLHS), m_VPValue(SubRHS))) && SubR &&
        SubR->hasNoSignedWrap() && poisonGuaranteesUB(LHSVal)) {
      const SCEV *V1 = getSCEVExprForVPValue(SubLHS, PSE, L);
      const SCEV *V2 = getSCEVExprForVPValue(SubRHS, PSE, L);
      if (!isa<SCEVCouldNotCompute>(V1) && !isa<SCEVCouldNotCompute>(V2))
        return SE.getMinusSCEV(SE.getSignExtendExpr(V1, DestTy),
                               SE.getSignExtendExpr(V2, DestTy), SCEV::FlagNSW);
    }

    return CreateSCEV({LHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getSignExtendExpr(Ops[0], DestTy);
    });
  }
  if (match(V,
            m_Intrinsic<Intrinsic::umax>(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getUMaxExpr(Ops[0], Ops[1]);
    });
  if (match(V,
            m_Intrinsic<Intrinsic::smax>(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getSMaxExpr(Ops[0], Ops[1]);
    });
  if (match(V,
            m_Intrinsic<Intrinsic::umin>(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getUMinExpr(Ops[0], Ops[1]);
    });
  if (match(V,
            m_Intrinsic<Intrinsic::smin>(m_VPValue(LHSVal), m_VPValue(RHSVal))))
    return CreateSCEV({LHSVal, RHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getSMinExpr(Ops[0], Ops[1]);
    });
  if (match(V, m_Intrinsic<Intrinsic::abs>(m_VPValue(LHSVal), m_VPValue())))
    return CreateSCEV({LHSVal}, [&](ArrayRef<SCEVUse> Ops) {
      // is_int_min_poison is local to this intrinsic: poison on INT_MIN is
      // not proof that the input is never INT_MIN, nor that poison reaches
      // UB. Do not translate it to SCEV's global IsNSW flag.
      return SE.getAbsExpr(Ops[0], /*IsNSW=*/false);
    });

  ArrayRef<VPValue *> Ops;
  Type *SourceElementType;
  if (match(V, m_GetElementPtr(SourceElementType, Ops))) {
    return CreateSCEV(Ops, [&](ArrayRef<SCEVUse> Ops) {
      return SE.getGEPExpr(Ops.front(), Ops.drop_front(), SourceElementType);
    });
  }

  // TODO: Support constructing SCEVs for more recipes as needed.
  const VPRecipeBase *DefR = V->getDefiningRecipe();
  const SCEV *Expr =
      TypeSwitch<const VPRecipeBase *, const SCEV *>(DefR)
          .Case([](const VPExpandSCEVRecipe *R) { return R->getSCEV(); })
          .Case([&SE, &PSE, L](const VPWidenIntOrFpInductionRecipe *R) {
            const SCEV *Step = getSCEVExprForVPValue(R->getStepValue(), PSE, L);
            if (!L || isa<SCEVCouldNotCompute>(Step))
              return SE.getCouldNotCompute();
            const SCEV *Start =
                getSCEVExprForVPValue(R->getStartValue(), PSE, L);
            const SCEV *AddRec =
                SE.getAddRecExpr(Start, Step, L, SCEV::FlagNone);
            if (R->getTruncInst())
              return SE.getTruncateExpr(AddRec, R->getScalarType());
            return AddRec;
          })
          .Case([&SE, &PSE,
                 L](const VPWidenPointerInductionRecipe *R) -> const SCEV * {
            const SCEV *Start =
                getSCEVExprForVPValue(R->getStartValue(), PSE, L);
            if (!L || isa<SCEVCouldNotCompute>(Start))
              return SE.getCouldNotCompute();
            const SCEV *Step = getSCEVExprForVPValue(R->getStepValue(), PSE, L);
            if (isa<SCEVCouldNotCompute>(Step))
              return SE.getCouldNotCompute();
            return SE.getAddRecExpr(Start, Step, L, SCEV::FlagNone);
          })
          .Case([&SE, &PSE, L](const VPDerivedIVRecipe *R) -> const SCEV * {
            const SCEV *Start = getSCEVExprForVPValue(R->getOperand(0), PSE, L);
            const SCEV *IV = getSCEVExprForVPValue(R->getOperand(1), PSE, L);
            const SCEV *Scale = getSCEVExprForVPValue(R->getOperand(2), PSE, L);
            if (any_of(ArrayRef({Start, IV, Scale}),
                       IsaPred<SCEVCouldNotCompute>))
              return SE.getCouldNotCompute();

            return SE.getAddExpr(
                SE.getTruncateOrSignExtend(Start, IV->getType()),
                SE.getMulExpr(
                    IV, SE.getTruncateOrSignExtend(Scale, IV->getType())));
          })
          .Case([&SE, &PSE, L](const VPScalarIVStepsRecipe *R) {
            const SCEV *IV = getSCEVExprForVPValue(R->getOperand(0), PSE, L);
            const SCEV *Step = getSCEVExprForVPValue(R->getOperand(1), PSE, L);
            if (isa<SCEVCouldNotCompute>(IV) || !isa<SCEVConstant>(Step))
              return SE.getCouldNotCompute();
            return SE.getTruncateOrSignExtend(IV, Step->getType());
          })
          .Default(
              [&SE](const VPRecipeBase *) { return SE.getCouldNotCompute(); });

  return PSE.getPredicatedSCEV(Expr);
}

std::optional<int64_t>
vputils::getConstantStride(VPValue *Addr, Type *AccessTy,
                           PredicatedScalarEvolution &PSE, const Loop *L) {
  assert(!hasIrregularType(AccessTy, L->getHeader()->getDataLayout()) &&
         "should not try to widen irregular types");
  const SCEV *AddrSCEV = getSCEVExprForVPValue(Addr, PSE, L);
  auto *AddRec = dyn_cast<SCEVAddRecExpr>(AddrSCEV);
  if (!AddRec)
    return {};

  return getStrideFromAddRec(AddRec, L, AccessTy, /*Ptr=*/nullptr, PSE);
}

bool vputils::isAddressSCEVForCost(const SCEV *Addr, ScalarEvolution &SE,
                                   const Loop *L) {
  // If address is an SCEVAddExpr, we require that all operands must be either
  // be invariant or a (possibly sign-extend) affine AddRec.
  if (auto *PtrAdd = dyn_cast<SCEVAddExpr>(Addr)) {
    return all_of(PtrAdd->operands(), [&SE, L](const SCEV *Op) {
      return SE.isLoopInvariant(Op, L) ||
             match(Op, m_scev_SExt(m_scev_AffineAddRec(m_SCEV(), m_SCEV()))) ||
             match(Op, m_scev_AffineAddRec(m_SCEV(), m_SCEV()));
    });
  }

  // Otherwise, check if address is loop invariant or an affine add recurrence.
  return SE.isLoopInvariant(Addr, L) ||
         match(Addr, m_scev_AffineAddRec(m_SCEV(), m_SCEV()));
}

unsigned vputils::getOpcode(const VPValue *V) {
  return TypeSwitch<const VPValue *, unsigned>(V)
      .Case<VPInstruction, VPWidenRecipe, VPWidenCastRecipe, VPWidenGEPRecipe,
            VPReplicateRecipe, VPWidenPHIRecipe, VPWidenLoadRecipe,
            VPWidenLoadEVLRecipe>([](auto *I) { return I->getOpcode(); })
      .Case<VPVectorPointerRecipe, VPPredInstPHIRecipe, VPScalarIVStepsRecipe>(
          [](auto *I) {
            // For recipes that do not directly map to LLVM IR instructions,
            // assign opcodes after the last VPInstruction opcode (which is also
            // after the last IR Instruction opcode), based on the VPRecipeID.
            return VPInstruction::OpsEnd + 1 + I->getVPRecipeID();
          })
      .Default([](auto *) { return 0; });
}

std::optional<std::pair<bool, unsigned>>
vputils::getOpcodeOrIntrinsicID(const VPValue *V) {
  if (Intrinsic::ID IID = vputils::getIntrinsicID(V))
    return std::make_pair(true, IID);
  if (unsigned Opcode = vputils::getOpcode(V))
    return std::make_pair(false, Opcode);
  return {};
}

/// Returns true if \p Opcode preserves uniformity, i.e., if all operands are
/// uniform, the result will also be uniform.
static bool preservesUniformity(unsigned Opcode) {
  if (Instruction::isBinaryOp(Opcode) || Instruction::isCast(Opcode))
    return true;
  switch (Opcode) {
  case Instruction::Freeze:
  case Instruction::GetElementPtr:
  case Instruction::ICmp:
  case Instruction::FCmp:
  case Instruction::Select:
  case VPInstruction::Not:
  case VPInstruction::Broadcast:
  case VPInstruction::MaskedCond:
  case VPInstruction::PtrAdd:
    return true;
  default:
    return false;
  }
}

bool vputils::isElementwise(const VPValue *V) {
  // TODO: Handle more opcodes and recipes.
  if (!isa<VPInstruction, VPWidenRecipe>(V))
    return false;
  unsigned Opcode = getOpcode(V);
  return Instruction::isUnaryOp(Opcode) || Instruction::isBinaryOp(Opcode);
}

bool vputils::isSingleScalar(const VPValue *VPV) {
  // Live-in, symbolic and canonical-IV region values are single-scalar.
  if (auto *RV = dyn_cast<VPRegionValue>(VPV))
    return RV == RV->getDefiningRegion()->getCanonicalIV();
  if (isa<VPIRValue, VPSymbolicValue>(VPV))
    return true;

  if (auto *Rep = dyn_cast<VPReplicateRecipe>(VPV)) {
    const VPRegionBlock *RegionOfR = Rep->getRegion();
    // Don't consider recipes in replicate regions as uniform yet; their first
    // lane cannot be accessed when executing the replicate region for other
    // lanes.
    if (RegionOfR && RegionOfR->isReplicator())
      return false;
    return Rep->isSingleScalar() || (preservesUniformity(Rep->getOpcode()) &&
                                     all_of(Rep->operands(), isSingleScalar));
  }
  if (isa<VPWidenGEPRecipe, VPBlendRecipe>(VPV))
    return all_of(VPV->getDefiningRecipe()->operands(), isSingleScalar);
  if (auto *WidenR = dyn_cast<VPWidenRecipe>(VPV)) {
    return preservesUniformity(WidenR->getOpcode()) &&
           all_of(WidenR->operands(), isSingleScalar);
  }
  if (auto *VPI = dyn_cast<VPInstruction>(VPV))
    return VPI->isSingleScalar() || VPI->isVectorToScalar() ||
           (preservesUniformity(VPI->getOpcode()) &&
            all_of(VPI->operands(), isSingleScalar));
  if (auto *RR = dyn_cast<VPReductionRecipe>(VPV))
    return !RR->isPartialReduction();
  if (isa<VPVectorPointerRecipe, VPVectorEndPointerRecipe, VPDerivedIVRecipe>(
          VPV))
    return true;
  if (auto *Expr = dyn_cast<VPExpressionRecipe>(VPV))
    return Expr->isVectorToScalar();

  // VPExpandSCEVRecipes must be placed in the entry and are always uniform.
  return isa<VPExpandSCEVRecipe>(VPV);
}

bool vputils::isUniformAcrossVFsAndUFs(const VPValue *V) {
  // Live-ins, symbolic and canonical-IV region values are uniform.
  if (auto *RV = dyn_cast<VPRegionValue>(V))
    return RV == RV->getDefiningRegion()->getCanonicalIV();
  if (isa<VPIRValue, VPSymbolicValue>(V))
    return true;

  const VPRecipeBase *R = V->getDefiningRecipe();
  const VPBasicBlock *VPBB = R ? R->getParent() : nullptr;
  const VPlan *Plan = VPBB ? VPBB->getPlan() : nullptr;
  if (VPBB &&
      (VPBB == Plan->getVectorPreheader() || VPBB == Plan->getEntry())) {
    if (match(R,
              m_VPInstruction<VPInstruction::CanonicalIVIncrementForPart>()) ||
        match(R, m_ExtractVectorForPart(m_VPValue(), m_VPValue())))
      return false;
    return all_of(R->operands(), isUniformAcrossVFsAndUFs);
  }

  return TypeSwitch<const VPRecipeBase *, bool>(R)
      .Case([](const VPDerivedIVRecipe *R) { return true; })
      .Case([](const VPReplicateRecipe *R) {
        // Be conservative about side-effects, except for the
        // known-side-effecting assumes and stores, which we know will be
        // uniform.
        return R->isSingleScalar() &&
               (!R->mayHaveSideEffects() ||
                isa<AssumeInst, StoreInst>(R->getUnderlyingInstr())) &&
               all_of(R->operands(), isUniformAcrossVFsAndUFs);
      })
      .Case([](const VPWidenRecipe *R) {
        return preservesUniformity(R->getOpcode()) &&
               all_of(R->operands(), isUniformAcrossVFsAndUFs);
      })
      .Case([](const VPPhi *) {
        // Bail out on VPPhi, as we can end up in infinite cycles.
        return false;
      })
      .Case([](const VPInstruction *VPI) {
        return (VPI->isSingleScalar() || VPI->isVectorToScalar() ||
                preservesUniformity(VPI->getOpcode())) &&
               all_of(VPI->operands(), isUniformAcrossVFsAndUFs);
      })
      .Case([](const VPWidenCastRecipe *R) {
        // A cast is uniform according to its operand.
        return isUniformAcrossVFsAndUFs(R->getOperand(0));
      })
      .Default([](const VPRecipeBase *) { // A value is considered non-uniform
                                          // unless proven otherwise.
        return false;
      });
}

bool vputils::doesGeneratePerAllLanes(const VPRecipeBase *R) {
  if (auto *RepR = dyn_cast<VPReplicateRecipe>(R))
    return RepR->doesGeneratePerAllLanes();
  if (auto *VPI = dyn_cast<VPInstruction>(R))
    return VPI->doesGeneratePerAllLanes();
  if (auto *SIVSteps = dyn_cast<VPScalarIVStepsRecipe>(R))
    return SIVSteps->doesGeneratePerAllLanes();
  return false;
}

VPBasicBlock *vputils::getFirstLoopHeader(VPlan &Plan, VPDominatorTree &VPDT) {
  auto DepthFirst = vp_depth_first_shallow(Plan.getEntry());
  auto I = find_if(DepthFirst, [&VPDT](VPBlockBase *VPB) {
    return VPBlockUtils::isHeader(VPB, VPDT);
  });
  return I == DepthFirst.end() ? nullptr : cast<VPBasicBlock>(*I);
}

unsigned vputils::getVFScaleFactor(VPRecipeBase *R) {
  if (!R)
    return 1;
  if (auto *RR = dyn_cast<VPReductionPHIRecipe>(R))
    return RR->getVFScaleFactor();
  if (auto *RR = dyn_cast<VPReductionRecipe>(R))
    return RR->getVFScaleFactor();
  if (auto *ER = dyn_cast<VPExpressionRecipe>(R))
    return ER->getVFScaleFactor();
  assert(
      (!isa<VPInstruction>(R) || cast<VPInstruction>(R)->getOpcode() !=
                                     VPInstruction::ReductionStartVector) &&
      "getting scaling factor of reduction-start-vector not implemented yet");
  return 1;
}

bool vputils::cannotHoistOrSinkRecipe(const VPRecipeBase &R, bool Sinking) {
  // Assumes don't alias anything or throw; as long as they're guaranteed to
  // execute, they're safe to hoist. They should however not be sunk, as it
  // would destroy information.
  if (match(&R, m_Intrinsic<Intrinsic::assume>()))
    return Sinking;
  if (R.mayHaveSideEffects() || R.mayReadFromMemory() || R.isPhi())
    return true;
  // Allocas cannot be hoisted.
  auto *RepR = dyn_cast<VPReplicateRecipe>(&R);
  return RepR && RepR->getOpcode() == Instruction::Alloca;
}

SmallVector<VPBasicBlock *>
VPBlockUtils::blocksInSingleSuccessorChainBetween(VPBasicBlock *FirstBB,
                                                  VPBasicBlock *LastBB) {
  assert(FirstBB->getParent() == LastBB->getParent() &&
         "FirstBB and LastBB from different regions");
#ifndef NDEBUG
  bool InSingleSuccChain = false;
  for (VPBlockBase *Succ = FirstBB; Succ; Succ = Succ->getSingleSuccessor())
    InSingleSuccChain |= (Succ == LastBB);
  assert(InSingleSuccChain &&
         "LastBB unreachable from FirstBB in single-successor chain");
#endif
  auto Blocks = to_vector(
      VPBlockUtils::blocksOnly<VPBasicBlock>(vp_depth_first_deep(FirstBB)));
  auto *LastIt = find(Blocks, LastBB);
  assert(LastIt != Blocks.end() &&
         "LastBB unreachable from FirstBB in depth-first traversal");
  Blocks.erase(std::next(LastIt), Blocks.end());
  return Blocks;
}

VPValue *vputils::findIncomingAliasMask(const VPlan &Plan) {
  for (VPRecipeBase &R : *Plan.getVectorPreheader())
    if (match(&R, m_VPInstruction<VPInstruction::IncomingAliasMask>()))
      return cast<VPInstruction>(&R);
  return nullptr;
}

SmallVector<std::pair<VPBasicBlock *, VPIRBasicBlock *>>
vputils::getEarlyExits(const VPlan &Plan, const VPBlockBase *MiddleVPBB) {
  SmallVector<std::pair<VPBasicBlock *, VPIRBasicBlock *>> Exits;
  for (VPIRBasicBlock *ExitVPBB : Plan.getExitBlocks())
    for (VPBlockBase *Pred : ExitVPBB->getPredecessors())
      if (Pred != MiddleVPBB)
        Exits.emplace_back(cast<VPBasicBlock>(Pred), ExitVPBB);
  return Exits;
}

VPScalarIVStepsRecipe *vputils::createScalarIVSteps(
    VPlan &Plan, InductionDescriptor::InductionKind Kind,
    Instruction::BinaryOps InductionOpcode, FPMathOperator *FPBinOp,
    Instruction *TruncI, VPValue *StartV, VPValue *Step, DebugLoc DL,
    VPBuilder &Builder, const VPIRFlags::WrapFlagsTy &Flags) {
  VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
  VPBasicBlock *HeaderVPBB = LoopRegion->getEntryBasicBlock();
  VPValue *CanonicalIV = LoopRegion->getCanonicalIV();
  VPSingleDefRecipe *BaseIV =
      Builder.createDerivedIV(Kind, FPBinOp, StartV, CanonicalIV, Step, Flags);

  // Truncate base induction if needed.
  Type *ResultTy = BaseIV->getScalarType();
  if (TruncI) {
    Type *TruncTy = TruncI->getType();
    assert(ResultTy->getScalarSizeInBits() > TruncTy->getScalarSizeInBits() &&
           "Not truncating.");
    assert(ResultTy->isIntegerTy() && "Truncation requires an integer type");
    BaseIV = Builder.createScalarCast(Instruction::Trunc, BaseIV, TruncTy, DL);
    ResultTy = TruncTy;
  }

  // Truncate step if needed.
  Type *StepTy = Step->getScalarType();
  if (ResultTy != StepTy) {
    assert(StepTy->getScalarSizeInBits() > ResultTy->getScalarSizeInBits() &&
           "Not truncating.");
    assert(StepTy->isIntegerTy() && "Truncation requires an integer type");
    auto *VecPreheader =
        cast<VPBasicBlock>(HeaderVPBB->getSingleHierarchicalPredecessor());
    VPBuilder::InsertPointGuard Guard(Builder);
    Builder.setInsertPoint(VecPreheader);
    Step = Builder.createScalarCast(Instruction::Trunc, Step, ResultTy, DL);
  }
  return Builder.createScalarIVSteps(InductionOpcode, FPBinOp, BaseIV, Step,
                                     &Plan.getVF(), DL);
}

VPValue *
vputils::scalarizeVPWidenPointerInduction(VPWidenPointerInductionRecipe *PtrIV,
                                          VPlan &Plan, VPBuilder &Builder) {
  const InductionDescriptor &ID = PtrIV->getInductionDescriptor();
  VPIRValue *StartV = Plan.getZero(ID.getStep()->getType());
  VPValue *StepV = PtrIV->getOperand(1);
  VPScalarIVStepsRecipe *Steps = createScalarIVSteps(
      Plan, InductionDescriptor::IK_IntInduction, Instruction::Add, nullptr,
      nullptr, StartV, StepV, PtrIV->getDebugLoc(), Builder);

  return Builder.createPtrAdd(PtrIV->getStartValue(), Steps,
                              PtrIV->getDebugLoc(), "next.gep");
}

bool VPBlockUtils::isHeader(const VPBlockBase *VPB,
                            const VPDominatorTree &VPDT) {
  auto *VPBB = dyn_cast<VPBasicBlock>(VPB);
  if (!VPBB)
    return false;

  // If VPBB is in a region R, VPBB is a loop header if R is a loop region with
  // VPBB as its entry, i.e., free of predecessors.
  if (auto *R = VPBB->getParent())
    return !R->isReplicator() && !VPBB->hasPredecessors();

  // A header dominates its second predecessor (the latch), with the other
  // predecessor being the preheader
  return VPB->getPredecessors().size() == 2 &&
         VPDT.dominates(VPB, VPB->getPredecessors()[1]);
}

bool VPBlockUtils::isLatch(const VPBlockBase *VPB,
                           const VPDominatorTree &VPDT) {
  // A latch has a header as its last successor, with its other successors
  // leaving the loop. A preheader OTOH has a header as its first (and only)
  // successor.
  return VPB->getNumSuccessors() >= 2 &&
         VPBlockUtils::isHeader(VPB->getSuccessors().back(), VPDT);
}

std::pair<VPBasicBlock *, VPBasicBlock *>
VPBlockUtils::getPlainCFGHeaderAndLatch(const VPlan &Plan) {
  VPBasicBlock *Header = cast<VPBasicBlock>(
      Plan.getEntry()->getNumSuccessors() == 1
          ? Plan.getEntry()->getSingleSuccessor()
          : Plan.getEntry()->getSuccessors()[1]->getSingleSuccessor());
  assert(Header->getNumPredecessors() == 2 &&
         "Header must have exactly 2  predecessors");
  auto *Latch = cast<VPBasicBlock>(Header->getPredecessors()[1]);
  return {Header, Latch};
}

VPBasicBlock *VPBlockUtils::getPlainCFGMiddleBlock(const VPlan &Plan) {
  return cast<VPBasicBlock>(Plan.getScalarPreheader()->getPredecessors()[0]);
}

VPIRFlags vputils::getFlagsForInduction(const InductionDescriptor &ID,
                                        const VPPhi *PhiR) {
  if (ID.getKind() == InductionDescriptor::IK_FpInduction)
    return ID.getInductionBinOp()->getFastMathFlags();

  // The flags only bound the induction values if the increment directly
  // updates PhiR.
  VPValue *Inc = PhiR->getOperand(1);
  if (match(Inc, m_c_Add(m_Specific(PhiR), m_VPValue())))
    return cast<VPInstruction>(Inc)->getNoWrapFlagsOrNone();

  if (match(Inc, m_Sub(m_Specific(PhiR), m_VPValue()))) {
    // The step of a sub induction is negated, so NUW cannot be preserved. NSW
    // can, if the step is not the signed minimum.
    ConstantInt *Step = ID.getConstIntStepValue();
    bool NSW = cast<VPInstruction>(Inc)->getNoWrapFlagsOrNone().HasNSW &&
               Step && !Step->isMinValue(/*IsSigned=*/true);
    return VPIRFlags::WrapFlagsTy(/*NUW*/ false, NSW);
  }

  return VPIRFlags::WrapFlagsTy(false, false);
}

std::optional<MemoryLocation>
vputils::getMemoryLocation(const VPRecipeBase &R) {
  auto *M = dyn_cast<VPIRMetadata>(&R);
  if (!M)
    return std::nullopt;
  MemoryLocation Loc;
  // Populate noalias metadata from VPIRMetadata.
  if (MDNode *NoAliasMD = M->getMetadata(LLVMContext::MD_noalias))
    Loc.AATags.NoAlias = NoAliasMD;
  if (MDNode *AliasScopeMD = M->getMetadata(LLVMContext::MD_alias_scope))
    Loc.AATags.Scope = AliasScopeMD;
  return Loc;
}

VPInstruction *vputils::findCanonicalIVIncrement(VPlan &Plan) {
  VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
  VPRegionValue *CanIV = LoopRegion->getCanonicalIV();
  assert(CanIV && "Expected loop region to have a canonical IV");

  VPSymbolicValue &VFxUF = Plan.getVFxUF();

  // Check if \p Step matches the expected increment step, accounting for
  // materialization of VFxUF and UF.
  auto IsIncrementStep = [&](VPValue *Step) -> bool {
    if (!VFxUF.isMaterialized())
      return Step == &VFxUF;

    VPSymbolicValue &UF = Plan.getUF();
    if (!UF.isMaterialized())
      return Step == &UF ||
             match(Step, m_c_Mul(m_Specific(&Plan.getUF()), m_VScale()));

    // Alias masking: step is number of active lanes of a dependence mask.
    if (match(Step, m_ZExtOrTruncOrSelf(
                        m_VPInstruction<VPInstruction::NumActiveLanes>())))
      return true;

    unsigned ConcreteUF = Plan.getConcreteUF();
    // Fixed VF: step is just the concrete UF.
    if (match(Step, m_SpecificInt(ConcreteUF)))
      return true;

    // Scalable VF: step involves VScale.
    if (ConcreteUF == 1)
      return match(Step, m_VScale());
    if (match(Step, m_c_Mul(m_SpecificInt(ConcreteUF), m_VScale())))
      return true;
    // mul(VScale, ConcreteUF) may have been simplified to
    // shl(VScale, log2(ConcreteUF)) when ConcreteUF is a power of 2.
    return isPowerOf2_32(ConcreteUF) &&
           match(Step, m_Shl(m_VScale(), m_SpecificInt(Log2_32(ConcreteUF))));
  };

  VPInstruction *Increment = nullptr;
  for (VPUser *U : CanIV->users()) {
    VPValue *Step;
    if (isa<VPInstruction>(U) &&
        match(U, m_c_Add(m_Specific(CanIV), m_VPValue(Step))) &&
        IsIncrementStep(Step)) {
      assert(!Increment && "There must be a unique increment");
      Increment = cast<VPInstruction>(U);
    }
  }

  assert((!VFxUF.isMaterialized() || Increment) &&
         "After materializing VFxUF, an increment must exist");
  assert((!Increment ||
          LoopRegion->hasCanonicalIVNUW() == Increment->hasNoUnsignedWrap()) &&
         "NUW flag in region and increment must match");
  return Increment;
}

/// Find the ComputeReductionResult recipe for \p PhiR, looking through selects
/// inserted for predicated reductions or tail folding.
VPInstruction *vputils::findComputeReductionResult(VPReductionPHIRecipe *PhiR) {
  VPValue *BackedgeVal = PhiR->getBackedgeValue();
  if (auto *Res =
          findUserOf<VPInstruction::ComputeReductionResult>(BackedgeVal))
    return Res;

  // Look through selects inserted for tail folding or predicated reductions.
  VPRecipeBase *SelR =
      findUserOf(BackedgeVal, m_Select(m_VPValue(), m_VPValue(), m_VPValue()));
  if (!SelR)
    return nullptr;
  return findUserOf<VPInstruction::ComputeReductionResult>(
      cast<VPSingleDefRecipe>(SelR));
}

bool vputils::isUsedByLoadStoreAddress(const VPValue *V) {
  SmallPtrSet<const VPValue *, 4> Seen;
  SmallVector<const VPValue *> WorkList = {V};

  while (!WorkList.empty()) {
    const VPValue *Cur = WorkList.pop_back_val();
    if (!Seen.insert(Cur).second)
      continue;

    auto *Blend = dyn_cast<VPBlendRecipe>(Cur);
    // Skip blends that use V only through a compare by checking if any incoming
    // value was already visited.
    if (Blend && none_of(seq<unsigned>(0, Blend->getNumIncomingValues()),
                         [&](unsigned I) {
                           return Seen.contains(Blend->getIncomingValue(I));
                         }))
      continue;

    for (VPUser *U : Cur->users()) {
      if (auto *InterleaveR = dyn_cast<VPInterleaveBase>(U))
        if (InterleaveR->getAddr() == Cur)
          return true;
      // Cur is used as the pointer of a (possibly masked) load (operand 0) or
      // store (operand 1).
      if (match(U, m_CombineOr(m_Unary<Instruction::Load>(m_Specific(Cur)),
                               m_Binary<Instruction::Store>(m_VPValue(),
                                                            m_Specific(Cur)))))
        return true;
      if (auto *MemR = dyn_cast<VPWidenMemoryRecipe>(cast<VPRecipeBase>(U))) {
        if (MemR->getAddr() == Cur && MemR->isConsecutive())
          return true;
      }
    }

    // The legacy cost model only supports scalarization loads/stores with phi
    // addresses, if the phi is directly used as load/store address. Don't
    // traverse further for Blends.
    if (Blend)
      continue;

    // Only traverse further through users that also define a value (and can
    // thus have their own users walked). Skip when Cur is only used as mask ,
    // as well as loads: a loaded value does not depend on the load's operand.
    for (VPUser *U : Cur->users()) {
      auto *VPI = dyn_cast<VPInstruction>(U);
      if (VPI && VPI->getMask() == Cur &&
          none_of(VPI->operandsWithoutMask(), equal_to(Cur)))
        continue;
      if (match(U, m_VPInstruction<Instruction::Load>()))
        continue;
      if (auto *SDR = dyn_cast<VPSingleDefRecipe>(U))
        WorkList.push_back(SDR);
    }
  }
  return false;
}

/// Try to find a loop-invariant IR value for \p S in the plan's entry block
/// that can be reused. Returns the corresponding live-in VPValue, or nullptr
/// if no reusable IR value is found.
VPValue *VPSCEVExpander::tryToReuseIRValue(const SCEV *S) {
  if (isa<SCEVConstant, SCEVUnknown>(S))
    return nullptr;
  VPlan &Plan = Builder.getPlan();
  BasicBlock *PH = cast<VPIRBasicBlock>(Plan.getEntry())->getIRBasicBlock();
  for (Value *V : SE.getSCEVValues(S)) {
    // Only reuse instructions in the plan's entry block, or, when a
    // DominatorTree is available, any instruction that dominates it.
    // Instructions in sibling branches may not dominate the entry block.
    auto *I = dyn_cast<Instruction>(V);
    if (!I)
      return Plan.getOrAddLiveIn(V);
    if (!SE.DT.dominates(I->getParent(), PH))
      continue;
    SmallVector<Instruction *> DropPoisonGeneratingInsts;
    if (!SE.canReuseInstruction(S, I, DropPoisonGeneratingInsts))
      continue;
    for (Instruction *DropI : DropPoisonGeneratingInsts)
      SCEVExpander::dropPoisonGeneratingAnnotationsAndReinfer(SE, DropI);
    return Plan.getOrAddLiveIn(V);
  }
  return nullptr;
}

VPValue *VPSCEVExpander::expand(const SCEV *S) {
  if (VPValue *V = tryToReuseIRValue(S))
    return V;

  switch (S->getSCEVType()) {
  case scConstant:
    return Builder.getPlan().getOrAddLiveIn(cast<SCEVConstant>(S)->getValue());
  case scUnknown:
    return Builder.getPlan().getOrAddLiveIn(cast<SCEVUnknown>(S)->getValue());
  case scVScale:
    return Builder.createVScale(S->getType(), DL);
  case scAddExpr: {
    auto *AddE = cast<SCEVAddExpr>(S);
    VPIRFlags::WrapFlagsTy WrapFlags(AddE->hasNoUnsignedWrap(),
                                     AddE->hasNoSignedWrap());

    // Expand pointer SCEVAddExpr as a ptradd of the pointer base and the
    // integer offset, matching SCEVExpander.
    if (S->getType()->isPointerTy()) {
      VPValue *Base = expand(SE.getPointerBase(S));
      VPValue *Offset = expand(SE.removePointerBase(S));
      GEPNoWrapFlags GEPFlags = WrapFlags.HasNUW
                                    ? GEPNoWrapFlags::noUnsignedWrap()
                                    : GEPNoWrapFlags::none();
      return Builder.createNoWrapPtrAdd(Base, Offset, GEPFlags, DL);
    }

    // Non-constant-negative add operands are expanded negated and subtracted
    // from the running result below, instead of being negated and added.
    auto UseSubtract = [](const SCEV *Op) {
      return Op->isNonConstantNegative();
    };
    // Iterate in reverse so that constants are emitted last, and move the
    // subtracted operands last, matching SCEVExpander's LoopCompare, so that
    // they don't start the running result.
    SmallVector<const SCEV *, 2> SCEVOps(reverse(AddE->operands()));
    stable_sort(SCEVOps, [&](const SCEV *L, const SCEV *R) {
      return !UseSubtract(L) && UseSubtract(R);
    });
    SmallVector<VPValue *, 2> Ops;
    for (const SCEV *Op : SCEVOps) {
      // The first operand starts the result, so it is never subtracted.
      bool Negate = !Ops.empty() && UseSubtract(Op);
      Ops.push_back(expand(Negate ? SE.getNegativeSCEV(Op) : Op));
    }
    VPValue *Result = Ops.front();
    for (auto [Op, OpV] : drop_begin(zip_equal(SCEVOps, Ops))) {
      if (UseSubtract(Op)) {
        // Result + (-Op) == Result - Op, which saves the multiply for the
        // negation. NSW only transfers if negating Op cannot overflow, see
        // ScalarEvolution::getMinusSCEV.
        bool HasNSW =
            WrapFlags.HasNSW && !SE.getSignedRangeMin(Op).isMinSignedValue();
        Result = Builder.createOverflowingOp(Instruction::Sub, {Result, OpV},
                                             {/*HasNUW=*/false, HasNSW}, DL);
        continue;
      }
      Result = Builder.createOverflowingOp(Instruction::Add, {Result, OpV},
                                           WrapFlags, DL);
    }
    return Result;
  }
  case scMulExpr: {
    auto *MulE = cast<SCEVMulExpr>(S);
    VPIRFlags::WrapFlagsTy WrapFlags(MulE->hasNoUnsignedWrap(),
                                     MulE->hasNoSignedWrap());
    SmallVector<VPValue *, 2> Ops;
    for (const SCEV *Op : reverse(MulE->operands()))
      Ops.push_back(expand(Op));
    VPValue *Result = Ops.front();
    for (VPValue *OpV : drop_begin(Ops)) {
      Result = Builder.createOverflowingOp(Instruction::Mul, {Result, OpV},
                                           WrapFlags, DL);
    }
    return Result;
  }
  case scUDivExpr: {
    auto *UDiv = cast<SCEVUDivExpr>(S);
    VPValue *LHS = expand(UDiv->getLHS());
    const SCEV *RHSExpr = UDiv->getRHS();
    VPValue *RHS = expand(RHSExpr);
    if (SafeUDivMode) {
      // Make sure the UDiv's divisor is guaranteed to not be zero/poison, to
      // avoid UB.
      Type *Ty = UDiv->getType();
      bool GuaranteedNotPoison =
          ScalarEvolution::isGuaranteedNotToBePoison(RHSExpr);
      if (!GuaranteedNotPoison)
        RHS = Builder.createFreeze(RHS, DL);
      if (!SE.isKnownNonZero(RHSExpr) || !GuaranteedNotPoison)
        RHS = Builder.createScalarIntrinsic(
            Intrinsic::umax, {RHS, Builder.getPlan().getConstantInt(Ty, 1)}, Ty,
            DL);
    }
    return Builder.createNaryOp(Instruction::UDiv, {LHS, RHS},
                                VPIRFlags::getDefaultFlags(Instruction::UDiv),
                                DL);
  }
  case scTruncate:
  case scZeroExtend:
  case scSignExtend:
  case scPtrToAddr: {
    auto *Cast = cast<SCEVCastExpr>(S);
    VPValue *Op = expand(Cast->getOperand());
    Instruction::CastOps Opcode;
    switch (S->getSCEVType()) {
    case scTruncate:
      Opcode = Instruction::Trunc;
      break;
    case scZeroExtend:
      Opcode = Instruction::ZExt;
      break;
    case scSignExtend:
      Opcode = Instruction::SExt;
      break;
    case scPtrToAddr:
      Opcode = Instruction::PtrToAddr;
      break;
    default:
      llvm_unreachable("Unhandled cast SCEV");
    }

    // When expanding ptrtoaddr, first check if there's an existing ptrtoint we
    // can reuse.
    if (Opcode == Instruction::PtrToAddr) {
      VPlan &Plan = Builder.getPlan();
      BasicBlock *PH = cast<VPIRBasicBlock>(Plan.getEntry())->getIRBasicBlock();
      if (auto *IRV = dyn_cast<VPIRValue>(Op)) {
        if (CastInst *CI = SCEVExpander::findReusableCastForPtrToAddr(
                IRV->getValue(), S->getType(), PH->getDataLayout(),
                [&](const CastInst *CI) {
                  return SE.DT.dominates(CI->getParent(), PH);
                }))
          return Plan.getOrAddLiveIn(CI);
      }
    }

    std::optional<VPIRFlags> Flags;
    if (Opcode == Instruction::ZExt)
      Flags =
          VPIRFlags::NonNegFlagsTy(SE.isKnownNonNegative(Cast->getOperand()));

    return Builder.createScalarCast(Opcode, Op, S->getType(), DL, Flags);
  }
  case scUMaxExpr:
  case scSMaxExpr:
  case scUMinExpr:
  case scSMinExpr:
  case scSequentialUMinExpr: {
    auto *MinMax = cast<SCEVNAryExpr>(S);
    Intrinsic::ID IntrinsicID;
    switch (S->getSCEVType()) {
    case scUMaxExpr:
      IntrinsicID = Intrinsic::umax;
      break;
    case scSMaxExpr:
      IntrinsicID = Intrinsic::smax;
      break;
    case scUMinExpr:
    case scSequentialUMinExpr:
      IntrinsicID = Intrinsic::umin;
      break;
    case scSMinExpr:
      IntrinsicID = Intrinsic::smin;
      break;
    default:
      llvm_unreachable("Unexpected min/max SCEV type");
    }
    // Chain operands in reverse order matching SCEVExpander's expansion of
    // min/max expressions. In SafeUDivMode freeze expansion results of operands
    // other than the first for sequential UMins, to avoid short-circuiting
    // divide-by-0/poison.
    bool IsSequential = S->getSCEVType() == scSequentialUMinExpr;
    Type *ResultTy = MinMax->getType();
    bool PrevSafeMode = SafeUDivMode;
    SmallVector<VPValue *, 2> Ops;
    for (const SCEV *SCEVOp : reverse(MinMax->operands())) {
      bool MayShortCircuit =
          IsSequential && Ops.size() != MinMax->getNumOperands() - 1;
      SafeUDivMode = MayShortCircuit || PrevSafeMode;
      VPValue *OpV = expand(SCEVOp);
      SafeUDivMode = PrevSafeMode;
      if (MayShortCircuit)
        OpV = Builder.createFreeze(OpV, DL);
      Ops.push_back(OpV);
    }
    VPValue *Result = Ops.front();
    for (VPValue *Op : drop_begin(Ops)) {
      if (!ResultTy->isPointerTy()) {
        Result = Builder.createScalarIntrinsic(IntrinsicID, {Result, Op},
                                               ResultTy, DL);
        continue;
      }
      // The min/max intrinsics don't support pointer operands, so expand
      // pointer-typed min/max as cmp + select, matching SCEVExpander.
      VPValue *Cmp = Builder.createICmp(
          MinMaxIntrinsic::getPredicate(IntrinsicID), Result, Op, DL);
      Result = Builder.createSelect(Cmp, Result, Op, DL);
      Function &F = *Builder.getPlan().getIRFunction();
      if (MDNode *MD =
              getExplicitlyUnknownBranchWeightsIfProfiled(F, "scev-expander"))
        cast<VPInstruction>(Result)->setMetadata(LLVMContext::MD_prof, MD);
    }
    return Result;
  }
  case scAddRecExpr: {
    auto *AR = cast<SCEVAddRecExpr>(S);
    VPlan &Plan = Builder.getPlan();
    [[maybe_unused]] BasicBlock *PH =
        cast<VPIRBasicBlock>(Plan.getEntry())->getIRBasicBlock();
    assert(SE.DT.dominates(AR->getLoop()->getHeader(), PH) &&
           "can only expand AddRecs for loops outside VPlan's scope");

    // Try to expand AR by re-using an existing canonical IV in the Plan's
    // entry. A canonical IV must be affine and integer typed.
    if (!AR->isAffine() || !AR->getType()->isIntegerTy())
      return vputils::getOrCreateVPValueForSCEVExpr(Plan, AR);
    auto FoundCanIV =
        find_if(Plan.getEntry()->phis(), [&](const VPRecipeBase &R) {
          if (!SE.isSCEVable(cast<VPIRPhi>(R).getIRPhi().getType()))
            return false;
          const SCEV *Candidate = SE.getSCEV(&cast<VPIRPhi>(R).getIRPhi());
          return match(Candidate,
                       m_scev_AffineAddRec(m_scev_Zero(), m_scev_One(),
                                           m_SpecificLoop(AR->getLoop()))) &&
                 Candidate->getType() == AR->getType();
        });
    if (FoundCanIV == Plan.getEntry()->phis().end())
      return vputils::getOrCreateVPValueForSCEVExpr(Plan, AR);

    // {Start, +, Step} --> Start + IV * Step, since the AddRec is affine.
    // Compute Offset = IV * Step.
    VPValue *Start = expand(AR->getStart());
    Value *CanonicalIV = &cast<VPIRPhi>(FoundCanIV)->getIRPhi();
    VPValue *Offset = expand(
        SE.getMulExpr(SE.getUnknown(CanonicalIV), AR->getStepRecurrence(SE)));

    // Compute Start + Offset with nuw from the AddRec.
    return Builder.createAdd(Start, Offset, DL, "",
                             {AR->hasNoUnsignedWrap(), false});
  }
  case scCouldNotCompute:
    llvm_unreachable("Attempt to expand a SCEVCouldNotCompute");
  }
  llvm_unreachable("Unknown SCEV kind!");
}

bool vputils::isDeadRecipe(VPRecipeBase &R) {
  // Do remove conditional assume instructions as their conditions may be
  // flattened.
  auto *RepR = dyn_cast<VPReplicateRecipe>(&R);
  bool IsConditionalAssume = RepR && RepR->isPredicated() &&
                             match(RepR, m_Intrinsic<Intrinsic::assume>());
  if (IsConditionalAssume)
    return true;

  if (R.mayHaveSideEffects())
    return false;

  // Forbid removing trip-count expressions.
  if (isa<VPExpandSCEVRecipe>(R) &&
      R.getVPSingleValue() == R.getParent()->getPlan()->getTripCount())
    return false;

  // Recipe is dead if no user keeps the recipe alive.
  return all_of(R.definedValues(), [](VPValue *V) { return V->user_empty(); });
}

void vputils::recursivelyDeleteDeadRecipes(VPValue *V) {
  SmallVector<VPValue *> WorkList;
  SmallPtrSet<VPValue *, 8> Seen;
  WorkList.push_back(V);

  while (!WorkList.empty()) {
    VPValue *Cur = WorkList.pop_back_val();
    if (!Seen.insert(Cur).second)
      continue;
    VPRecipeBase *R = Cur->getDefiningRecipe();
    if (!R)
      continue;
    if (!isDeadRecipe(*R))
      continue;
    append_range(WorkList, R->operands());
    R->eraseFromParent();
  }
}

SmallVector<VPUser *> vputils::collectUsersRecursively(VPValue *V) {
  SetVector<VPUser *> Users(llvm::from_range, V->users());
  for (unsigned I = 0; I != Users.size(); ++I) {
    VPRecipeBase *Cur = cast<VPRecipeBase>(Users[I]);
    for (VPValue *V : Cur->definedValues())
      Users.insert_range(V->users());
  }
  return Users.takeVector();
}

/// Returns \p Num / \p Denom as a BranchProbability, clamped so a ratio that is
/// neither zero nor one does not round to zero or one. BlockFrequencyInfo also
/// keeps a zero-weight edge distinguishable from an unreachable one.
static BranchProbability getBranchProbabilityKeepingPartial(uint64_t Num,
                                                            uint64_t Denom) {
  BranchProbability P = BranchProbability::getBranchProbability(Num, Denom);
  if (Num == 0 || Num == Denom)
    return P;
  return BranchProbability::getRaw(std::clamp(
      P.getNumerator(), 1u, BranchProbability::getDenominator() - 1));
}

BranchProbability vputils::getExecutionProbability(BlockFrequency Freq) {
  return getBranchProbabilityKeepingPartial(
      Freq.getFrequency(),
      BlockFrequencyInfoImplBase::BlockMass::getFull().getMass());
}

/// Returns the probability of each successor edge of \p VPBB, computed via
/// BranchProbabilityInfo::getEdgeProbabilitiesFromWeights from the branch
/// weights recorded on its terminator, or std::nullopt if not available.
static std::optional<SmallVector<BranchProbability>>
getSuccessorProbabilities(const VPBasicBlock *VPBB) {
  // With a single successor the edge is always taken and needs no weights.
  if (VPBB->getSingleSuccessor())
    return SmallVector<BranchProbability>{BranchProbability::getOne()};

  SmallVector<uint32_t> Weights;
  auto *Term = dyn_cast_if_present<VPInstruction>(VPBB->getTerminator());
  if (!Term || !extractBranchWeights(Term->getBranchWeights(), Weights) ||
      Weights.size() != VPBB->getNumSuccessors())
    return std::nullopt;
  return BranchProbabilityInfo::getEdgeProbabilitiesFromWeights(Weights);
}

DenseMap<const VPBasicBlock *, std::optional<VPExecutionFrequency>>
vputils::computeExecutionFrequencies(ArrayRef<VPBasicBlock *> Blocks) {
  using BFIBase = BlockFrequencyInfoImplBase;
  assert(!Blocks.empty() && "expected at least the header block");
  // Distribute the header's frequency using BFI. Nodes for blocks are numbered
  // in reverse post-order. Edges leaving Blocks, i.e. a plain CFG's edges to
  // the middle block or to an exit block, exit to a node outside the loop.
  BFIBase BFI;
  BFIBase::BlockNode Header(0), Outside(Blocks.size());
  BFIBase::LoopData &Loop = BFI.Loops.emplace_back(nullptr, Header);
  DenseMap<const VPBlockBase *, BFIBase::BlockNode> Nodes;
  for (auto [Idx, VPBB] : enumerate(Blocks)) {
    Nodes[VPBB] = BFIBase::BlockNode(Idx);
    BFI.Working.emplace_back(BFIBase::BlockNode(Idx)).Loop = &Loop;
  }
  BFI.Working.emplace_back(Outside);
  BFI.Working[Header.Index].getMass() = BFIBase::BlockMass::getFull();

  // Keep track nodes reached via an edge without branch weighs or with
  // estimated ones
  SmallVector<bool> IsUnknown(Blocks.size()), IsEstimated(Blocks.size());
  for (auto [Idx, VPBB] : enumerate(Blocks)) {
    BFIBase::BlockNode Node(Idx);
    auto Probs = getSuccessorProbabilities(VPBB);
    auto *Term = dyn_cast_if_present<VPInstruction>(VPBB->getTerminator());
    bool TermIsEstimated = Term && Term->hasEstimatedBranchWeights();
    BFIBase::Distribution Dist;
    for (auto [SuccIdx, Succ] : enumerate(VPBB->getSuccessors())) {
      BFIBase::BlockNode SuccNode = Nodes.lookup_or(Succ, Outside);
      if (SuccNode != Header && SuccNode != Outside) {
        IsUnknown[SuccNode.Index] |= IsUnknown[Idx] || !Probs;
        IsEstimated[SuccNode.Index] |= IsEstimated[Idx] || TermIsEstimated;
      }
      if (Probs)
        BFI.addToDist(Dist, &Loop, Node, SuccNode,
                      getWeightFromBranchProb((*Probs)[SuccIdx]));
    }
    if (Probs)
      BFI.distributeMass(Node, &Loop, Dist);
  }

  // Round frequencies up to at least 1, so all edges are reached with a
  // non-zero frequency, to distinguish rarely executed blocks from unreachable
  // ones. blocks distinguishable from unreachable ones.
  DenseMap<const VPBasicBlock *, std::optional<VPExecutionFrequency>>
      Frequencies;
  for (auto [Idx, VPBB] : enumerate(Blocks)) {
    std::optional<VPExecutionFrequency> &Freq = Frequencies[VPBB];
    if (IsUnknown[Idx])
      continue;
    uint64_t Mass = BFI.Working[Idx].getMass().getMass();
    Freq.emplace(BlockFrequency(std::max<uint64_t>(Mass, 1)), IsEstimated[Idx]);
  }
  return Frequencies;
}

VPIRValue *vputils::tryToFoldLiveIns(VPSingleDefRecipe &R,
                                     ArrayRef<VPValue *> Operands,
                                     const DataLayout &DL) {
  auto OpcodeOrIID = getOpcodeOrIntrinsicID(&R);
  if (!OpcodeOrIID)
    return nullptr;

  SmallVector<Value *, 4> Ops;
  for (VPValue *Op : Operands) {
    VPValue *Candidate = Op;
    match(Op, m_Broadcast(m_VPValue(Candidate)));
    if (!match(Candidate, m_LiveIn()))
      return nullptr;
    Value *V = Candidate->getUnderlyingValue();
    if (!V)
      return nullptr;
    Ops.push_back(V);
  }

  VPlan &Plan = *R.getParent()->getPlan();
  auto FoldToIRValue = [&]() -> Value * {
    InstSimplifyFolder Folder(DL);
    if (OpcodeOrIID->first) {
      // VPInstructions store the called intrinsic as last operand.
      if (isa<VPInstruction>(R))
        Ops.pop_back();

      auto *RFlags = dyn_cast<VPRecipeWithIRFlags>(&R);
      return Folder.FoldIntrinsic(OpcodeOrIID->second, Ops, R.getScalarType(),
                                  RFlags ? RFlags->getFastMathFlagsOrNone()
                                         : FastMathFlags());
    }
    unsigned Opcode = OpcodeOrIID->second;
    if (Instruction::isBinaryOp(Opcode))
      return Folder.FoldBinOp(static_cast<Instruction::BinaryOps>(Opcode),
                              Ops[0], Ops[1]);
    if (Instruction::isCast(Opcode))
      return Folder.FoldCast(static_cast<Instruction::CastOps>(Opcode), Ops[0],
                             R.getVPSingleValue()->getScalarType());
    switch (Opcode) {
    case VPInstruction::Not:
      return Folder.FoldBinOp(Instruction::BinaryOps::Xor, Ops[0],
                              Constant::getAllOnesValue(Ops[0]->getType()));
    case Instruction::Select:
      return Folder.FoldSelect(Ops[0], Ops[1], Ops[2]);
    case Instruction::ICmp:
    case Instruction::FCmp:
      return Folder.FoldCmp(cast<VPRecipeWithIRFlags>(R).getPredicate(), Ops[0],
                            Ops[1]);
    case Instruction::GetElementPtr: {
      auto &RFlags = cast<VPRecipeWithIRFlags>(R);
      auto *GEP = cast<GetElementPtrInst>(RFlags.getUnderlyingInstr());
      return Folder.FoldGEP(DL, GEP->getSourceElementType(), Ops[0],
                            drop_begin(Ops), RFlags.getGEPNoWrapFlags());
    }
    case VPInstruction::PtrAdd:
    case VPInstruction::WidePtrAdd:
      return Folder.FoldGEP(DL, IntegerType::getInt8Ty(Plan.getContext()),
                            Ops[0], Ops[1],
                            cast<VPRecipeWithIRFlags>(R).getGEPNoWrapFlags());
    // An extract of a live-in is an extract of a broadcast, so return the
    // broadcasted element.
    case Instruction::ExtractElement:
      assert(!Ops[0]->getType()->isVectorTy() && "Live-ins should be scalar");
      return Ops[0];
    case VPInstruction::ActiveLaneMask:
    case VPInstruction::WideActiveLaneMask: {
      uint64_t Multiplier = 1;
      if (Opcode == VPInstruction::WideActiveLaneMask) {
        // Optimizing WideALM can only happen after the Plan is unrolled.
        if (!Plan.isUnrolled())
          return nullptr;
        Multiplier = cast<ConstantInt>(Ops[2])->getZExtValue();
        Ops.pop_back();
      }

      // We rely on the fact that different VPlans are created for the
      // fixed-vector and scalable-vector cases.
      ElementCount MaxVF =
          *max_element(Plan.vectorFactors(), ElementCount::isKnownLT) *
          Multiplier;

      Type *I1Ty = IntegerType::getInt1Ty(Plan.getContext());
      if (auto *C = dyn_cast_if_present<Constant>(Folder.FoldIntrinsic(
              Intrinsic::get_active_lane_mask, Ops,
              VectorType::get(I1Ty, MaxVF), {}, Plan.getIRFunction()))) {
        // We cannot handle vector constants that are not all-true or all-false,
        // because they would not be collapsable to a scalar constant, that
        // would be necessary for live-in simplification.
        if (C->isOneValue())
          return ConstantInt::getTrue(I1Ty);
        if (C->isNullValue())
          return ConstantInt::getFalse(I1Ty);
      }
    }
    }
    return nullptr;
  };

  if (Value *V = FoldToIRValue())
    return Plan.getOrAddLiveIn(V);
  return nullptr;
}

void vputils::detail::pullOutPermutationsImpl(
    VPlan &Plan, function_ref<VPValue *(VPValue *Op)> MatchPerm,
    function_ref<VPSingleDefRecipe *(VPSingleDefRecipe *X)> BuildPerm) {
  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
           vp_depth_first_deep(Plan.getEntry()))) {
    for (VPSingleDefRecipe &Def :
         make_early_inc_range(make_isa_range<VPSingleDefRecipe>(*VPBB))) {
      if (!isElementwise(&Def))
        continue;

      // At least one of the ops must be a permutation.
      if (none_of(Def.operands(), MatchPerm))
        continue;

      // All operands must be a single-use permutation or a live in (splat).
      if (!all_of(Def.operands(), [&MatchPerm](VPValue *Op) {
            return (Op->hasOneUse() && MatchPerm(Op)) || match(Op, m_LiveIn());
          }))
        continue;

      // Remove the inner permutations.
      for (unsigned I = 0, E = Def.getNumOperands(); I != E; ++I)
        if (VPValue *X = MatchPerm(Def.getOperand(I)))
          Def.setOperand(I, X);

      VPSingleDefRecipe *Res = BuildPerm(&Def);
      Res->insertAfter(&Def);
      Def.replaceUsesWithIf(Res, [&Res](VPUser &U) { return &U != Res; });
    }
  }
}

// Implements the algorithm described in "Simple and Efficient Construction of
// Static Single Assignment Form" by Braun et al.
VPValue *vputils::reconstructSSA(VPBasicBlock *VPBB,
                                 DenseMap<VPBasicBlock *, VPValue *> &Defs) {
  assert(!Defs.empty() && "Defs shouldn't be empty");
  assert(
      is_contained(vp_depth_first_shallow(VPBB->getPlan()->getEntry()), VPBB) &&
      "VPBB isn't reachable from entry");
  if (VPValue *Def = Defs.lookup(VPBB))
    return Def;
  // If the entry block is reached and there's still no def, then Defs is
  // missing a definition that covers this path.
  assert(VPBB->getNumPredecessors() && "Not all paths have def");

  if (VPBlockBase *Pred = VPBB->getSinglePredecessor())
    return reconstructSSA(cast<VPBasicBlock>(Pred), Defs);

  // Multiple predecessors, create a join.
  Type *Ty = Defs.begin()->second->getScalarType();
  VPPhi *Phi = VPBuilder(VPBB, VPBB->getFirstNonPhi())
                   .createScalarPhi({}, DebugLoc::getUnknown(), "", {}, Ty);
  Defs[VPBB] = Phi;
  for (auto *Pred : VPBB->predecessors())
    Phi->addIncoming(reconstructSSA(cast<VPBasicBlock>(Pred), Defs));

  // Fold away trivial phis.
  // TODO: Remove phi users which have become trivial too.
  if (all_equal(Phi->incoming_values())) {
    VPValue *Common = Phi->getIncomingValue(0);
    Phi->replaceAllUsesWith(Common);
    for (auto &[_, V] : Defs)
      if (V == Phi)
        V = Common;
    Defs[VPBB] = Common;
    Phi->eraseFromParent();
    return Common;
  }

  return Phi;
}