diff --git a/llvm/lib/Transforms/Vectorize/CMakeLists.txt b/llvm/lib/Transforms/Vectorize/CMakeLists.txt index 220e71491811e..87d75390a1d69 100644 --- a/llvm/lib/Transforms/Vectorize/CMakeLists.txt +++ b/llvm/lib/Transforms/Vectorize/CMakeLists.txt @@ -41,6 +41,7 @@ add_llvm_component_library(LLVMVectorize VPlanUnroll.cpp VPlanVerifier.cpp VPlanUtils.cpp + VPlanWideningDecisions.cpp ADDITIONAL_HEADER_DIRS ${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp index 9990fd142df02..6b836ff6c2010 100644 --- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp @@ -43,136 +43,6 @@ using namespace llvm; using namespace VPlanPatternMatch; using namespace SCEVPatternMatch; -/// If the pointer operand \p Addr of a memory access is an affine AddRec -/// w.r.t. \p L with a constant stride, return the stride in units of -/// \p AccessTy. Otherwise return std::nullopt. -static std::optional getConstantStride(VPValue *Addr, Type *AccessTy, - PredicatedScalarEvolution &PSE, - const Loop *L) { - const SCEV *AddrSCEV = vputils::getSCEVExprForVPValue(Addr, PSE, L); - auto *AddRec = dyn_cast(AddrSCEV); - if (!AddRec) - return {}; - - return getStrideFromAddRec(AddRec, L, AccessTy, /*Ptr=*/nullptr, PSE); -} - -bool VPlanTransforms::tryToConvertVPInstructionsToVPRecipes( - VPlan &Plan, const TargetLibraryInfo &TLI, PredicatedScalarEvolution &PSE, - Loop *OuterLoop) { - - ReversePostOrderTraversal> RPOT( - Plan.getVectorLoopRegion()); - for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly(RPOT)) { - // Skip blocks outside region - if (!VPBB->getParent()) - break; - VPRecipeBase *Term = VPBB->getTerminator(); - auto EndIter = Term ? Term->getIterator() : VPBB->end(); - // Introduce each ingredient into VPlan. - for (VPRecipeBase &Ingredient : - make_early_inc_range(make_range(VPBB->begin(), EndIter))) { - - VPValue *VPV = Ingredient.getVPSingleValue(); - if (!VPV->getUnderlyingValue()) - continue; - - Instruction *Inst = cast(VPV->getUnderlyingValue()); - - // Atomic accesses and fences have ordering/atomicity semantics that - // cannot be preserved by lane-wise widening. - if (isa(Inst)) - return false; - - VPRecipeBase *NewRecipe = nullptr; - if (auto *PhiR = dyn_cast(&Ingredient)) { - auto *Phi = cast(PhiR->getUnderlyingValue()); - NewRecipe = new VPWidenPHIRecipe(PhiR->operands(), PhiR->getDebugLoc(), - Phi->getName()); - } else if (auto *VPI = dyn_cast(&Ingredient)) { - assert(!isa(Inst) && "phis should be handled above"); - // Create VPWidenMemoryRecipe for loads and stores. - if (LoadInst *Load = dyn_cast(Inst)) { - bool IsConsecutive = - getConstantStride(VPI->getOperand(0), VPI->getScalarType(), PSE, - OuterLoop) == 1; - NewRecipe = new VPWidenLoadRecipe(*Load, Ingredient.getOperand(0), - nullptr /*Mask*/, IsConsecutive, - *VPI, Ingredient.getDebugLoc()); - } else if (StoreInst *Store = dyn_cast(Inst)) { - bool IsConsecutive = - getConstantStride(VPI->getOperand(1), - VPI->getOperand(0)->getScalarType(), PSE, - OuterLoop) == 1; - NewRecipe = new VPWidenStoreRecipe( - *Store, Ingredient.getOperand(1), Ingredient.getOperand(0), - nullptr /*Mask*/, IsConsecutive, *VPI, Ingredient.getDebugLoc()); - } else if (GetElementPtrInst *GEP = dyn_cast(Inst)) { - NewRecipe = new VPWidenGEPRecipe(GEP->getSourceElementType(), - Ingredient.operands(), *VPI, - Ingredient.getDebugLoc(), GEP); - } else if (CallInst *CI = dyn_cast(Inst)) { - Intrinsic::ID VectorID = getVectorIntrinsicIDForCall(CI, &TLI); - if (VectorID == Intrinsic::not_intrinsic) - return false; - - // The noalias.scope.decl intrinsic declares a noalias scope that - // is valid for a single iteration. Emitting it as a single-scalar - // replicate would incorrectly extend the scope across multiple - // original iterations packed into one vector iteration. - // FIXME: If we want to vectorize this loop, then we have to drop - // all the associated !alias.scope and !noalias. - if (VectorID == Intrinsic::experimental_noalias_scope_decl) - return false; - - // These intrinsics are recognized by getVectorIntrinsicIDForCall - // but are not widenable. Emit them as replicate instead of widening. - if (VectorID == Intrinsic::assume || - VectorID == Intrinsic::lifetime_end || - VectorID == Intrinsic::lifetime_start || - VectorID == Intrinsic::sideeffect || - VectorID == Intrinsic::pseudoprobe) { - // If the operand of llvm.assume holds before vectorization, it will - // also hold per lane. - // llvm.pseudoprobe requires to be duplicated per lane for accurate - // sample count. - const bool IsSingleScalar = VectorID != Intrinsic::assume && - VectorID != Intrinsic::pseudoprobe; - NewRecipe = new VPReplicateRecipe(CI, Ingredient.operands(), - /*IsSingleScalar=*/IsSingleScalar, - /*Mask=*/nullptr, *VPI, *VPI, - Ingredient.getDebugLoc()); - } else { - NewRecipe = new VPWidenIntrinsicRecipe( - *CI, VectorID, drop_end(Ingredient.operands()), CI->getType(), - VPIRFlags(*CI), *VPI, CI->getDebugLoc()); - } - } else if (auto *CI = dyn_cast(Inst)) { - NewRecipe = new VPWidenCastRecipe( - CI->getOpcode(), Ingredient.getOperand(0), CI->getType(), CI, - VPIRFlags(*CI), VPIRMetadata(*CI)); - } else { - NewRecipe = new VPWidenRecipe(*Inst, Ingredient.operands(), *VPI, - *VPI, Ingredient.getDebugLoc()); - } - } else { - assert(isa(&Ingredient) && - "inductions must be created earlier"); - continue; - } - - NewRecipe->insertBefore(&Ingredient); - if (NewRecipe->getNumDefinedValues() == 1) - VPV->replaceAllUsesWith(NewRecipe->getVPSingleValue()); - else - assert(NewRecipe->getNumDefinedValues() == 0 && - "Only recpies with zero or one defined values expected"); - Ingredient.eraseFromParent(); - } - } - return true; -} - /// Helper for extra no-alias checks via known-safe recipe and SCEV. class SinkStoreInfo { SmallPtrSet ExcludeRecipes; @@ -710,109 +580,6 @@ void VPlanTransforms::removeDeadRecipes(VPlan &Plan) { } } -/// Legalize VPWidenPointerInductionRecipe, by replacing it with a PtrAdd -/// (IndStart, ScalarIVSteps (0, Step)) if only its scalar values are used, as -/// VPWidenPointerInductionRecipe will generate vectors only. If some users -/// require vectors while other require scalars, the scalar uses need to extract -/// the scalars from the generated vectors (Note that this is different to how -/// int/fp inductions are handled). Legalize extract-from-ends using uniform -/// VPReplicateRecipe of wide inductions to use regular VPReplicateRecipe, so -/// the correct end value is available. Also optimize -/// VPWidenIntOrFpInductionRecipe, if any of its users needs scalar values, by -/// providing them scalar steps built on the canonical scalar IV and update the -/// original IV's users. This is an optional optimization to reduce the needs of -/// vector extracts. -static void legalizeAndOptimizeInductions(VPlan &Plan) { - VPBasicBlock *HeaderVPBB = Plan.getVectorLoopRegion()->getEntryBasicBlock(); - bool HasOnlyVectorVFs = !Plan.hasScalarVFOnly(); - VPBuilder Builder(HeaderVPBB, HeaderVPBB->getFirstNonPhi()); - for (VPRecipeBase &Phi : HeaderVPBB->phis()) { - auto *PhiR = dyn_cast(&Phi); - if (!PhiR) - continue; - - // Try to narrow wide and replicating recipes to uniform recipes, based on - // VPlan analysis. - // TODO: Apply to all recipes in the future, to replace legacy uniformity - // analysis. - auto Users = vputils::collectUsersRecursively(PhiR); - for (VPUser *U : reverse(Users)) { - auto *Def = dyn_cast(U); - auto *RepR = dyn_cast(U); - // Skip recipes that shouldn't be narrowed. - if (!Def || !isa(Def) || - Def->user_empty() || !Def->getUnderlyingValue() || - (RepR && (RepR->isSingleScalar() || RepR->isPredicated()))) - continue; - - // Skip recipes that may have other lanes than their first used. - if (!vputils::isSingleScalar(Def) && !vputils::onlyFirstLaneUsed(Def)) - continue; - - // TODO: Support scalarizing ExtractValue. - if (match(Def, - m_Binary(m_VPValue(), m_VPValue()))) - continue; - - auto *Clone = VPBuilder::createSingleScalarOp( - Def->getUnderlyingInstr()->getOpcode(), Def->operands(), - /*Mask=*/nullptr, *Def, {}, DebugLoc::getUnknown(), - Def->getUnderlyingInstr()); - Clone->insertAfter(Def); - Def->replaceAllUsesWith(Clone); - } - - // Replace wide pointer inductions which have only their scalars used by - // PtrAdd(IndStart, ScalarIVSteps (0, Step)). - if (auto *PtrIV = dyn_cast(&Phi)) { - if (!Plan.hasScalarVFOnly() && - !PtrIV->onlyScalarsGenerated(Plan.hasScalableVF())) - continue; - - VPValue *PtrAdd = - vputils::scalarizeVPWidenPointerInduction(PtrIV, Plan, Builder); - PtrIV->replaceAllUsesWith(PtrAdd); - continue; - } - - // Replace widened induction with scalar steps for users that only use - // scalars. - auto *WideIV = cast(&Phi); - if (HasOnlyVectorVFs && none_of(WideIV->users(), [WideIV](VPUser *U) { - return U->usesScalars(WideIV); - })) - continue; - - const InductionDescriptor &ID = WideIV->getInductionDescriptor(); - VPIRFlags::WrapFlagsTy WrapFlags; - // We can preserve nuw when the step is non-negative. - const APInt *Step; - if (match(WideIV->getStepValue(), m_APInt(Step)) && Step->isNonNegative()) - WrapFlags = {static_cast(WideIV->getNoWrapFlagsOrNone().HasNUW), - false}; - VPScalarIVStepsRecipe *Steps = vputils::createScalarIVSteps( - Plan, ID.getKind(), ID.getInductionOpcode(), - dyn_cast_or_null(ID.getInductionBinOp()), - WideIV->getTruncInst(), WideIV->getStartValue(), WideIV->getStepValue(), - WideIV->getDebugLoc(), Builder, WrapFlags); - - // Update scalar users of IV to use Step instead. - if (!HasOnlyVectorVFs) { - assert(!Plan.hasScalableVF() && - "plans containing a scalar VF cannot also include scalable VFs"); - WideIV->replaceAllUsesWith(Steps); - } else { - bool HasScalableVF = Plan.hasScalableVF(); - WideIV->replaceUsesWithIf(Steps, - [WideIV, HasScalableVF](VPUser &U, unsigned) { - if (HasScalableVF) - return U.usesFirstLaneOnly(WideIV); - return U.usesScalars(WideIV); - }); - } - } -} - /// Check if \p VPV is an untruncated wide induction, either before or after the /// increment. If so return the header IV (before the increment), otherwise /// return null. @@ -1710,120 +1477,6 @@ static void reassociateHeaderMask(VPlan &Plan) { } } -static std::optional -getUnmaskedDivRemOpcode(Intrinsic::ID ID) { - switch (ID) { - case Intrinsic::masked_udiv: - return Instruction::UDiv; - case Intrinsic::masked_sdiv: - return Instruction::SDiv; - case Intrinsic::masked_urem: - return Instruction::URem; - case Intrinsic::masked_srem: - return Instruction::SRem; - default: - return {}; - } -} - -static void narrowToSingleScalarRecipes(VPlan &Plan) { - if (Plan.hasScalarVFOnly()) - return; - - for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly( - vp_depth_first_deep(Plan.getEntry()))) { - for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) { - if (!isa(&R)) - continue; - auto *RepR = dyn_cast(&R); - if (RepR && (RepR->isSingleScalar() || RepR->isPredicated())) - continue; - - auto *RepOrWidenR = cast(&R); - if (RepR && RepR->getOpcode() == Instruction::Store && - vputils::isSingleScalar(RepR->getOperand(1))) { - auto *Clone = new VPReplicateRecipe( - RepOrWidenR->getUnderlyingInstr(), RepOrWidenR->operands(), - true /*IsSingleScalar*/, nullptr /*Mask*/, *RepR /*Flags*/, - *RepR /*Metadata*/, RepR->getDebugLoc()); - Clone->insertBefore(RepOrWidenR); - VPBuilder Builder(Clone); - VPValue *ExtractOp = Clone->getOperand(0); - if (vputils::isUniformAcrossVFsAndUFs(RepR->getOperand(1))) - ExtractOp = - Builder.createNaryOp(VPInstruction::ExtractLastPart, ExtractOp); - ExtractOp = - Builder.createNaryOp(VPInstruction::ExtractLastLane, ExtractOp); - Clone->setOperand(0, ExtractOp); - RepR->eraseFromParent(); - continue; - } - - // Narrow llvm.masked.{u,s}{div,rem} intrinsics with a safe divisor. - if (auto *IntrR = dyn_cast(RepOrWidenR)) { - if (!vputils::onlyFirstLaneUsed(IntrR)) - continue; - auto Opc = getUnmaskedDivRemOpcode(IntrR->getVectorIntrinsicID()); - if (!Opc) - continue; - VPBuilder Builder(IntrR); - VPValue *SafeDivisor = Builder.createSelect( - IntrR->getOperand(2), IntrR->getOperand(1), - Plan.getConstantInt(IntrR->getScalarType(), 1)); - VPValue *Clone = Builder.createNaryOp( - *Opc, {IntrR->getOperand(0), SafeDivisor}, - VPIRFlags::getDefaultFlags(*Opc), IntrR->getDebugLoc()); - IntrR->replaceAllUsesWith(Clone); - IntrR->eraseFromParent(); - continue; - } - - // Skip recipes that aren't single scalars. - if (!vputils::isSingleScalar(RepOrWidenR)) - continue; - - // Predicate to check if a user of Op introduces extra broadcasts. - auto IntroducesBCastOf = [](const VPValue *Op) { - return [Op](const VPUser *U) { - if (auto *VPI = dyn_cast(U)) { - if (is_contained({VPInstruction::ExtractLastLane, - VPInstruction::ExtractLastPart, - VPInstruction::ExtractPenultimateElement}, - VPI->getOpcode())) - return false; - } - return !U->usesScalars(Op); - }; - }; - - if (any_of(RepOrWidenR->users(), IntroducesBCastOf(RepOrWidenR)) && - none_of(RepOrWidenR->operands(), [&](VPValue *Op) { - if (any_of( - make_filter_range(Op->users(), not_equal_to(RepOrWidenR)), - IntroducesBCastOf(Op))) - return false; - // Non-constant live-ins require broadcasts, while constants do not - // need explicit broadcasts. - bool LiveInNeedsBroadcast = - isa(Op) && !isa(Op); - auto *OpR = dyn_cast(Op); - return LiveInNeedsBroadcast || (OpR && OpR->isSingleScalar()); - })) - continue; - - auto *Clone = VPBuilder::createSingleScalarOp( - vputils::getOpcode(RepOrWidenR), RepOrWidenR->operands(), - /*Mask=*/nullptr, *RepOrWidenR, {}, DebugLoc::getUnknown(), - RepOrWidenR->getUnderlyingInstr()); - Clone->insertBefore(RepOrWidenR); - RepOrWidenR->replaceAllUsesWith(Clone); - if (vputils::isDeadRecipe(*RepOrWidenR)) - RepOrWidenR->eraseFromParent(); - } - } -} - /// Try to see if all of \p Blend's masks share a common value logically and'ed /// and remove it from the masks. static void removeCommonBlendMask(VPBlendRecipe *Blend) { @@ -2807,135 +2460,6 @@ void VPlanTransforms::dropPoisonGeneratingRecipes(VPlan &Plan) { } } -void VPlanTransforms::createInterleaveGroups( - VPlan &Plan, - const SmallPtrSetImpl *> - &InterleaveGroups, - const bool &EpilogueAllowed) { - if (InterleaveGroups.empty()) - return; - - DenseMap IRMemberToRecipe; - for (VPBasicBlock *VPBB : - VPBlockUtils::blocksOnly(vp_depth_first_shallow( - Plan.getVectorLoopRegion()->getEntryBasicBlock()))) - for (VPRecipeBase &R : make_filter_range(*VPBB, [](VPRecipeBase &R) { - return isa(&R); - })) { - auto *MemR = cast(&R); - IRMemberToRecipe[&MemR->getIngredient()] = MemR; - } - - // Interleave memory: for each Interleave Group we marked earlier as relevant - // for this VPlan, replace the Recipes widening its memory instructions with a - // single VPInterleaveRecipe at its insertion point. - VPDominatorTree VPDT(Plan); - for (const auto *IG : InterleaveGroups) { - VPWidenMemoryRecipe *Start = nullptr; - Instruction *StartMember = nullptr; - for (auto *Member : IG->members()) - if (VPWidenMemoryRecipe *R = IRMemberToRecipe.lookup(Member)) { - StartMember = Member; - Start = R; - break; - } - if (!StartMember) // All member recipes are dead, so the group is dead. - continue; - VPIRMetadata InterleaveMD(*Start); - SmallVector StoredValues; - for (unsigned I = 0; I < IG->getFactor(); ++I) { - Instruction *MemberI = IG->getMember(I); - if (!MemberI) - continue; - if (VPWidenMemoryRecipe *MemoryR = IRMemberToRecipe.lookup(MemberI)) { - if (auto *StoreR = dyn_cast(MemoryR->getAsRecipe())) - StoredValues.push_back(StoreR->getStoredValue()); - InterleaveMD.intersect(*MemoryR); - } else { - InterleaveMD.intersect(VPIRMetadata(*MemberI)); - } - } - - bool NeedsMaskForGaps = - (IG->requiresScalarEpilogue() && !EpilogueAllowed) || - (!StoredValues.empty() && !IG->isFull()); - - Instruction *IRInsertPos = IG->getInsertPos(); - auto *InsertPos = IRMemberToRecipe.lookup(IRInsertPos); - if (!InsertPos) { - // InsertPos member is dead: find a new member that is alive. - assert(isa(Start->getAsRecipe()) && - "Dead member in non-load group?"); - InsertPos = Start; - for (Instruction *Member : IG->members()) - if (VPWidenMemoryRecipe *MemberR = IRMemberToRecipe.lookup(Member)) - if (VPDT.properlyDominates(MemberR->getAsRecipe(), - InsertPos->getAsRecipe())) - InsertPos = MemberR; - IRInsertPos = &InsertPos->getIngredient(); - } - VPRecipeBase *InsertPosR = InsertPos->getAsRecipe(); - - GEPNoWrapFlags NW = GEPNoWrapFlags::none(); - if (auto *Gep = dyn_cast( - getLoadStorePointerOperand(IRInsertPos)->stripPointerCasts())) - NW = Gep->getNoWrapFlags().withoutNoUnsignedWrap(); - - // Get or create the start address for the interleave group. - VPValue *Addr = Start->getAddr(); - VPRecipeBase *AddrDef = Addr->getDefiningRecipe(); - if (IG->getIndex(StartMember) != 0 || - (AddrDef && !VPDT.properlyDominates(AddrDef, InsertPosR))) { - // Either member zero's recipe is dead, or we cannot re-use the address of - // member zero because it does not dominate the insert position. Instead, - // use the address of the insert position and create a PtrAdd adjusting it - // to the address of member zero. - // TODO: Hoist Addr's defining recipe (and any operands as needed) to - // InsertPos or sink loads above zero members to join it. - assert(IG->getIndex(IRInsertPos) != 0 && - "index of insert position shouldn't be zero"); - auto &DL = IRInsertPos->getDataLayout(); - APInt Offset(32, - DL.getTypeAllocSize(getLoadStoreType(IRInsertPos)) * - IG->getIndex(IRInsertPos), - /*IsSigned=*/true); - VPValue *OffsetVPV = Plan.getConstantInt(-Offset); - VPBuilder B(InsertPosR); - Addr = B.createNoWrapPtrAdd(InsertPos->getAddr(), OffsetVPV, NW); - } - // If the group is reverse, adjust the index to refer to the last vector - // lane instead of the first. We adjust the index from the first vector - // lane, rather than directly getting the pointer for lane VF - 1, because - // the pointer operand of the interleaved access is supposed to be uniform. - if (IG->isReverse()) { - auto *ReversePtr = new VPVectorEndPointerRecipe( - Addr, &Plan.getVF(), getLoadStoreType(IRInsertPos), - -(int64_t)IG->getFactor(), NW, InsertPosR->getDebugLoc()); - ReversePtr->insertBefore(InsertPosR); - Addr = ReversePtr; - } - auto *VPIG = new VPInterleaveRecipe( - IG, Addr, StoredValues, InsertPos->getMask(), NeedsMaskForGaps, - InterleaveMD, InsertPosR->getDebugLoc()); - VPIG->insertBefore(InsertPosR); - - unsigned J = 0; - for (unsigned i = 0; i < IG->getFactor(); ++i) - if (Instruction *Member = IG->getMember(i)) { - VPWidenMemoryRecipe *MemberR = IRMemberToRecipe.lookup(Member); - if (!Member->getType()->isVoidTy()) { - if (MemberR) { - VPValue *OriginalV = MemberR->getAsRecipe()->getVPSingleValue(); - OriginalV->replaceAllUsesWith(VPIG->getVPValue(J)); - } - J++; - } - if (MemberR) - MemberR->getAsRecipe()->eraseFromParent(); - } - } -} - /// Returns the VPValue representing the uncountable exit comparison used by /// AnyOf if the recipes it depends on can be traced back to live-ins and /// the addresses (in GEP/PtrAdd form) of any (non-masked) load used in @@ -3464,264 +2988,6 @@ bool VPlanTransforms::handleUncountableEarlyExits( return true; } -/// This function tries convert extended in-loop reductions to -/// VPExpressionRecipe and clamp the \p Range if it is beneficial and -/// valid. The created recipe must be decomposed to its constituent -/// recipes before execution. -static VPExpressionRecipe * -tryToMatchAndCreateExtendedReduction(VPReductionRecipe *Red, VPCostContext &Ctx, - VFRange &Range) { - Type *RedTy = Red->getScalarType(); - VPValue *VecOp = Red->getVecOp(); - - assert(!Red->isPartialReduction() && - "This path does not support partial reductions"); - - // Clamp the range if using extended-reduction is profitable. - auto IsExtendedRedValidAndClampRange = - [&](unsigned Opcode, Instruction::CastOps ExtOpc, Type *SrcTy) -> bool { - return LoopVectorizationPlanner::getDecisionAndClampRange( - [&](ElementCount VF) { - auto *SrcVecTy = cast(toVectorTy(SrcTy, VF)); - TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput; - - InstructionCost ExtRedCost = InstructionCost::getInvalid(); - InstructionCost ExtCost = - cast(VecOp)->computeCost(VF, Ctx); - InstructionCost RedCost = Red->computeCost(VF, Ctx); - - assert(!RedTy->isFloatingPointTy() && - "getExtendedReductionCost only supports integer types"); - ExtRedCost = Ctx.TTI.getExtendedReductionCost( - Opcode, ExtOpc == Instruction::CastOps::ZExt, RedTy, SrcVecTy, - Red->getFastMathFlagsOrNone(), CostKind); - return ExtRedCost.isValid() && ExtRedCost < ExtCost + RedCost; - }, - Range); - }; - - VPValue *A; - // Match reduce(ext)). - if (match(VecOp, m_Isa(m_ZExtOrSExt(m_VPValue(A)))) && - IsExtendedRedValidAndClampRange( - RecurrenceDescriptor::getOpcode(Red->getRecurrenceKind()), - cast(VecOp)->getOpcode(), A->getScalarType())) - return new VPExpressionRecipe(cast(VecOp), Red); - - return nullptr; -} - -/// This function tries convert extended in-loop reductions to -/// VPExpressionRecipe and clamp the \p Range if it is beneficial -/// and valid. The created VPExpressionRecipe must be decomposed to its -/// constituent recipes before execution. Patterns of the -/// VPExpressionRecipe: -/// reduce.add(mul(...)), -/// reduce.add(mul(ext(A), ext(B))), -/// reduce.add(ext(mul(ext(A), ext(B)))). -/// reduce.fadd(fmul(ext(A), ext(B))) -static VPExpressionRecipe * -tryToMatchAndCreateMulAccumulateReduction(VPReductionRecipe *Red, - VPCostContext &Ctx, VFRange &Range) { - unsigned Opcode = RecurrenceDescriptor::getOpcode(Red->getRecurrenceKind()); - if (Opcode != Instruction::Add && Opcode != Instruction::Sub && - Opcode != Instruction::FAdd) - return nullptr; - - assert(!Red->isPartialReduction() && - "This path does not support partial reductions"); - Type *RedTy = Red->getScalarType(); - - // Clamp the range if using multiply-accumulate-reduction is profitable. - auto IsMulAccValidAndClampRange = - [&](VPWidenRecipe *Mul, VPWidenCastRecipe *Ext0, VPWidenCastRecipe *Ext1, - VPWidenCastRecipe *OuterExt) -> bool { - return LoopVectorizationPlanner::getDecisionAndClampRange( - [&](ElementCount VF) { - TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput; - Type *SrcTy = Ext0 ? Ext0->getOperand(0)->getScalarType() : RedTy; - InstructionCost MulAccCost; - - // getMulAccReductionCost for in-loop reductions does not support - // mixed or floating-point extends. - if (Ext0 && Ext1 && - (Ext0->getOpcode() != Ext1->getOpcode() || - Ext0->getOpcode() == Instruction::CastOps::FPExt)) - return false; - - bool IsZExt = - !Ext0 || Ext0->getOpcode() == Instruction::CastOps::ZExt; - auto *SrcVecTy = cast(toVectorTy(SrcTy, VF)); - MulAccCost = Ctx.TTI.getMulAccReductionCost(IsZExt, Opcode, RedTy, - SrcVecTy, CostKind); - - InstructionCost MulCost = Mul->computeCost(VF, Ctx); - InstructionCost RedCost = Red->computeCost(VF, Ctx); - InstructionCost ExtCost = 0; - if (Ext0) - ExtCost += Ext0->computeCost(VF, Ctx); - if (Ext1) - ExtCost += Ext1->computeCost(VF, Ctx); - if (OuterExt) - ExtCost += OuterExt->computeCost(VF, Ctx); - - return MulAccCost.isValid() && - MulAccCost < ExtCost + MulCost + RedCost; - }, - Range); - }; - - VPValue *VecOp = Red->getVecOp(); - VPRecipeBase *Sub = nullptr; - VPValue *A, *B; - VPValue *Tmp = nullptr; - - if (RedTy->isFloatingPointTy()) - return nullptr; - - // Sub reductions could have a sub between the add reduction and vec op. - if (match(VecOp, m_Sub(m_ZeroInt(), m_VPValue(Tmp)))) { - Sub = VecOp->getDefiningRecipe(); - VecOp = Tmp; - } - - // If ValB is a constant and can be safely extended, truncate it to the same - // type as ExtA's operand, then extend it to the same type as ExtA. This - // creates two uniform extends that can more easily be matched by the rest of - // the bundling code. The ExtB reference, ValB and operand 1 of Mul are all - // replaced with the new extend of the constant. - auto ExtendAndReplaceConstantOp = [](VPWidenCastRecipe *ExtA, - VPWidenCastRecipe *&ExtB, VPValue *&ValB, - VPWidenRecipe *Mul) { - if (!ExtA || ExtB || !isa(ValB)) - return; - Type *NarrowTy = ExtA->getOperand(0)->getScalarType(); - Instruction::CastOps ExtOpc = ExtA->getOpcode(); - const APInt *Const; - if (!match(ValB, m_APInt(Const)) || - !llvm::canConstantBeExtended( - Const, NarrowTy, TTI::getPartialReductionExtendKind(ExtOpc))) - return; - // The truncate ensures that the type of each extended operand is the - // same, and it's been proven that the constant can be extended from - // NarrowTy safely. Necessary since ExtA's extended operand would be - // e.g. an i8, while the const will likely be an i32. This will be - // elided by later optimisations. - VPBuilder Builder(Mul); - auto *Trunc = - Builder.createWidenCast(Instruction::CastOps::Trunc, ValB, NarrowTy); - Type *WideTy = ExtA->getScalarType(); - ValB = ExtB = Builder.createWidenCast(ExtOpc, Trunc, WideTy); - Mul->setOperand(1, ExtB); - }; - - // Try to match reduce.add(mul(...)). - if (match(VecOp, m_Mul(m_VPValue(A), m_VPValue(B)))) { - auto *RecipeA = dyn_cast(A); - auto *RecipeB = dyn_cast(B); - auto *Mul = cast(VecOp); - - // Convert reduce.add(mul(ext, const)) to reduce.add(mul(ext, ext(const))) - ExtendAndReplaceConstantOp(RecipeA, RecipeB, B, Mul); - - // Match reduce.add/sub(mul(ext, ext)). - if (RecipeA && RecipeB && match(RecipeA, m_ZExtOrSExt(m_VPValue())) && - match(RecipeB, m_ZExtOrSExt(m_VPValue())) && - IsMulAccValidAndClampRange(Mul, RecipeA, RecipeB, nullptr)) { - if (Sub) - return new VPExpressionRecipe(RecipeA, RecipeB, Mul, - cast(Sub), Red); - return new VPExpressionRecipe(RecipeA, RecipeB, Mul, Red); - } - // TODO: Add an expression type for this variant with a negated mul - if (!Sub && IsMulAccValidAndClampRange(Mul, nullptr, nullptr, nullptr)) - return new VPExpressionRecipe(Mul, Red); - } - // TODO: Add an expression type for negated versions of other expression - // variants. - if (Sub) - return nullptr; - - // Match reduce.add(ext(mul(A, B))). - if (match(VecOp, m_ZExtOrSExt(m_Mul(m_VPValue(A), m_VPValue(B))))) { - auto *Ext = cast(VecOp); - auto *Mul = cast(Ext->getOperand(0)); - auto *Ext0 = dyn_cast(A); - auto *Ext1 = dyn_cast(B); - - // reduce.add(ext(mul(ext, const))) - // -> reduce.add(ext(mul(ext, ext(const)))) - ExtendAndReplaceConstantOp(Ext0, Ext1, B, Mul); - - // reduce.add(ext(mul(ext(A), ext(B)))) - // -> reduce.add(mul(wider_ext(A), wider_ext(B))) - // The inner extends must either have the same opcode as the outer extend or - // be the same, in which case the multiply can never result in a negative - // value and the outer extend can be folded away by doing wider - // extends for the operands of the mul. - if (Ext0 && Ext1 && - (Ext->getOpcode() == Ext0->getOpcode() || Ext0 == Ext1) && - Ext0->getOpcode() == Ext1->getOpcode() && - IsMulAccValidAndClampRange(Mul, Ext0, Ext1, Ext) && Mul->hasOneUse()) { - auto *NewExt0 = new VPWidenCastRecipe( - Ext0->getOpcode(), Ext0->getOperand(0), Ext->getScalarType(), nullptr, - *Ext0, *Ext0, Ext0->getDebugLoc()); - NewExt0->insertBefore(Ext0); - - VPWidenCastRecipe *NewExt1 = NewExt0; - if (Ext0 != Ext1) { - NewExt1 = new VPWidenCastRecipe(Ext1->getOpcode(), Ext1->getOperand(0), - Ext->getScalarType(), nullptr, *Ext1, - *Ext1, Ext1->getDebugLoc()); - NewExt1->insertBefore(Ext1); - } - auto *NewMul = Mul->cloneWithOperands({NewExt0, NewExt1}); - NewMul->insertBefore(Mul); - Ext->replaceAllUsesWith(NewMul); - Ext->eraseFromParent(); - Mul->eraseFromParent(); - return new VPExpressionRecipe(NewExt0, NewExt1, NewMul, Red); - } - } - return nullptr; -} - -/// This function tries to create abstract recipes from the reduction recipe for -/// following optimizations and cost estimation. -static void tryToCreateAbstractReductionRecipe(VPReductionRecipe *Red, - VPCostContext &Ctx, - VFRange &Range) { - // Creation of VPExpressions for partial reductions is entirely handled in - // transformToPartialReduction. - assert(!Red->isPartialReduction() && - "This path does not support partial reductions"); - - VPExpressionRecipe *AbstractR = nullptr; - auto IP = std::next(Red->getIterator()); - auto *VPBB = Red->getParent(); - if (auto *MulAcc = tryToMatchAndCreateMulAccumulateReduction(Red, Ctx, Range)) - AbstractR = MulAcc; - else if (auto *ExtRed = tryToMatchAndCreateExtendedReduction(Red, Ctx, Range)) - AbstractR = ExtRed; - // Cannot create abstract inloop reduction recipes. - if (!AbstractR) - return; - - AbstractR->insertBefore(*VPBB, IP); - Red->replaceAllUsesWith(AbstractR); -} - -void VPlanTransforms::convertToAbstractRecipes(VPlan &Plan, VPCostContext &Ctx, - VFRange &Range) { - for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly( - vp_depth_first_deep(Plan.getVectorLoopRegion()))) { - for (VPRecipeBase &R : make_early_inc_range(*VPBB)) { - if (auto *Red = dyn_cast(&R)) - tryToCreateAbstractReductionRecipe(Red, Ctx, Range); - } - } -} - // Collect common metadata from a group of replicate recipes by intersecting // metadata from all recipes in the group. static VPIRMetadata getCommonMetadata(ArrayRef Recipes) { @@ -3912,377 +3178,6 @@ void VPlanTransforms::sinkPredicatedStores(VPlan &Plan, } } -/// Returns true if \p V is VPWidenLoadRecipe or VPInterleaveRecipe that can be -/// converted to a narrower recipe. \p V is used by a wide recipe that feeds a -/// store interleave group at index \p Idx, \p WideMember0 is the recipe feeding -/// the same interleave group at index 0. A VPWidenLoadRecipe can be narrowed to -/// an index-independent load if it feeds all wide ops at all indices (\p OpV -/// must be the operand at index \p OpIdx for both the recipe at lane 0, \p -/// WideMember0). A VPInterleaveRecipe can be narrowed to a wide load, if \p V -/// is defined at \p Idx of a load interleave group. -/// A live-in or recipe defined outside the loop region can be converted, if it -/// is the same across all lanes, or we can create a BuildVector for it. -static bool canNarrowLoad(VPSingleDefRecipe *WideMember0, unsigned OpIdx, - VPValue *OpV, unsigned Idx, bool IsScalable) { - VPValue *Member0Op = WideMember0->getOperand(OpIdx); - if (Member0Op->isDefinedOutsideLoopRegions()) { - // Operand matches Member0, broadcast across all fields for both live-ins - // and recipes. - if (Member0Op == OpV) - return true; - // Otherwise distinct per-field VPValues are assembled into a BuildVector. - return !IsScalable && OpV->isDefinedOutsideLoopRegions() && - OpV->getScalarType() == Member0Op->getScalarType(); - } - VPRecipeBase *Member0OpR = Member0Op->getDefiningRecipe(); - if (auto *W = dyn_cast(Member0OpR)) - // For scalable VFs, the narrowed plan processes vscale iterations at once, - // so a shared wide load cannot be narrowed to a uniform scalar; bail out. - return !IsScalable && !W->getMask() && W->isConsecutive() && - Member0Op == OpV; - if (auto *IR = dyn_cast(Member0OpR)) - return IR->getInterleaveGroup()->isFull() && IR->getVPValue(Idx) == OpV; - return false; -} - -static bool canNarrowOps(ArrayRef Ops, bool IsScalable) { - SmallVector Ops0; - auto *WideMember0 = dyn_cast(Ops[0]); - if (!WideMember0) - return false; - for (VPValue *V : Ops) { - if (!isa(V)) - return false; - auto *R = cast(V); - if (vputils::getOpcode(R) != vputils::getOpcode(WideMember0)) - return false; - if (R->getScalarType() != WideMember0->getScalarType()) - return false; - if (R->hasPredicate() && R->getPredicate() != WideMember0->getPredicate()) - return false; - } - - for (unsigned Idx = 0; Idx != WideMember0->getNumOperands(); ++Idx) { - SmallVector OpsI; - for (VPValue *Op : Ops) - OpsI.push_back(Op->getDefiningRecipe()->getOperand(Idx)); - - if (canNarrowOps(OpsI, IsScalable)) - continue; - - if (any_of(enumerate(OpsI), [WideMember0, Idx, IsScalable](const auto &P) { - const auto &[OpIdx, OpV] = P; - return !canNarrowLoad(WideMember0, Idx, OpV, OpIdx, IsScalable); - })) - return false; - } - - return true; -} - -/// Returns VF from \p VFs if \p IR is a full interleave group with factor and -/// number of members both equal to VF. The interleave group must also access -/// the full vector width. -static std::optional -isConsecutiveInterleaveGroup(VPInterleaveRecipe *InterleaveR, - ArrayRef VFs, - const TargetTransformInfo &TTI) { - if (!InterleaveR || InterleaveR->getMask()) - return std::nullopt; - - Type *GroupElementTy = nullptr; - if (InterleaveR->getStoredValues().empty()) { - GroupElementTy = InterleaveR->getVPValue(0)->getScalarType(); - if (!all_of(InterleaveR->definedValues(), [GroupElementTy](VPValue *Op) { - return Op->getScalarType() == GroupElementTy; - })) - return std::nullopt; - } else { - GroupElementTy = InterleaveR->getStoredValues()[0]->getScalarType(); - if (!all_of(InterleaveR->getStoredValues(), [GroupElementTy](VPValue *Op) { - return Op->getScalarType() == GroupElementTy; - })) - return std::nullopt; - } - - auto IG = InterleaveR->getInterleaveGroup(); - if (IG->getFactor() != IG->getNumMembers()) - return std::nullopt; - - auto GetVectorBitWidthForVF = [&TTI](ElementCount VF) { - TypeSize Size = TTI.getRegisterBitWidth( - VF.isFixed() ? TargetTransformInfo::RGK_FixedWidthVector - : TargetTransformInfo::RGK_ScalableVector); - assert(Size.isScalable() == VF.isScalable() && - "if Size is scalable, VF must be scalable and vice versa"); - return Size.getKnownMinValue(); - }; - - for (ElementCount VF : VFs) { - unsigned MinVal = VF.getKnownMinValue(); - unsigned GroupSize = GroupElementTy->getScalarSizeInBits() * MinVal; - if (IG->getFactor() == MinVal && GroupSize == GetVectorBitWidthForVF(VF)) - return {VF}; - } - return std::nullopt; -} - -/// Returns true if \p VPValue is a narrow VPValue. -static bool isAlreadyNarrow(VPValue *VPV) { - if (isa(VPV)) - return true; - auto *RepR = dyn_cast(VPV); - return RepR && RepR->isSingleScalar(); -} - -// Convert the wide recipes defining the VPValues in \p Members feeding an -// interleave group to a single narrow variant. The first member is reused as -// the narrowed recipe. BuildVectors for live-in operands are inserted into \p -// Preheader. -static VPValue *narrowInterleaveGroupOp(ArrayRef Members, - SmallPtrSetImpl &NarrowedOps, - VPBasicBlock *Preheader) { - VPValue *V = Members.front(); - if (NarrowedOps.contains(V)) - return V; - - if (V->isDefinedOutsideLoopRegions()) { - assert(all_of(Members, - [V](VPValue *M) { - return M->isDefinedOutsideLoopRegions() && - M->getScalarType() == V->getScalarType(); - }) && - "expected distinct loop-invariant values of matching scalar type"); - auto *BV = new VPInstruction(VPInstruction::BuildVector, Members); - Preheader->appendRecipe(BV); - NarrowedOps.insert(BV); - return BV; - } - - if (isAlreadyNarrow(V)) - return V; - - VPRecipeBase *R = V->getDefiningRecipe(); - if (isa(R)) { - auto *WideMember0 = cast(R); - for (VPValue *Member : Members.drop_front()) - WideMember0->intersectFlags(*cast(Member)); - for (unsigned Idx = 0, E = WideMember0->getNumOperands(); Idx != E; ++Idx) { - SmallVector OpsI; - for (VPValue *Member : Members) - OpsI.push_back(Member->getDefiningRecipe()->getOperand(Idx)); - WideMember0->setOperand( - Idx, narrowInterleaveGroupOp(OpsI, NarrowedOps, Preheader)); - } - return V; - } - - if (auto *LoadGroup = dyn_cast(R)) { - // Narrow interleave group to wide load, as transformed VPlan will only - // process one original iteration. - auto *LI = cast(LoadGroup->getInterleaveGroup()->getInsertPos()); - auto *L = VPBuilder(LoadGroup).createWidenLoad( - *LI, LoadGroup->getAddr(), LoadGroup->getMask(), /*Consecutive=*/true, - *LoadGroup, LoadGroup->getDebugLoc()); - NarrowedOps.insert(L); - return L; - } - - if (auto *RepR = dyn_cast(R)) { - assert(RepR->isSingleScalar() && RepR->getOpcode() == Instruction::Load && - "must be a single scalar load"); - NarrowedOps.insert(RepR); - return RepR; - } - - auto *WideLoad = cast(R); - VPValue *PtrOp = WideLoad->getAddr(); - if (auto *VecPtr = dyn_cast(PtrOp)) - PtrOp = VecPtr->getOperand(0); - // Narrow wide load to uniform scalar load, as transformed VPlan will only - // process one original iteration. - auto *N = new VPReplicateRecipe(&WideLoad->getIngredient(), {PtrOp}, - /*IsUniform*/ true, - /*Mask*/ nullptr, {}, *WideLoad); - N->insertBefore(WideLoad); - NarrowedOps.insert(N); - return N; -} - -std::unique_ptr -VPlanTransforms::narrowInterleaveGroups(VPlan &Plan, - const TargetTransformInfo &TTI) { - VPRegionBlock *VectorLoop = Plan.getVectorLoopRegion(); - - if (!VectorLoop) - return nullptr; - - // Only handle single-block loops for now. - if (VectorLoop->getEntryBasicBlock() != VectorLoop->getExitingBasicBlock()) - return nullptr; - - // Skip plans when we may not be able to properly narrow. - VPBasicBlock *Exiting = VectorLoop->getExitingBasicBlock(); - if (!match(&Exiting->back(), m_BranchOnCount())) - return nullptr; - - assert(match(&Exiting->back(), - m_BranchOnCount(m_Add(m_VPValue(), m_Specific(&Plan.getVFxUF())), - m_Specific(&Plan.getVectorTripCount()))) && - "unexpected branch-on-count"); - - SmallVector StoreGroups; - std::optional VFToOptimize; - for (auto &R : *VectorLoop->getEntryBasicBlock()) { - if (isa(&R) && - vputils::onlyFirstLaneUsed(cast(&R))) - continue; - - // Bail out on recipes not supported at the moment: - // * phi recipes other than the canonical induction - // * recipes writing to memory except interleave groups - // Only support plans with a canonical induction phi. - if (R.isPhi()) - return nullptr; - - auto *InterleaveR = dyn_cast(&R); - if (R.mayWriteToMemory() && !InterleaveR) - return nullptr; - - // Bail out if any recipe defines a vector value used outside the - // vector loop region. - if (any_of(R.definedValues(), [&](VPValue *V) { - return any_of(V->users(), [&](VPUser *U) { - auto *UR = cast(U); - return UR->getParent()->getParent() != VectorLoop; - }); - })) - return nullptr; - - // All other ops are allowed, but we reject uses that cannot be converted - // when checking all allowed consumers (store interleave groups) below. - if (!InterleaveR) - continue; - - // Try to find a single VF, where all interleave groups are consecutive and - // saturate the full vector width. If we already have a candidate VF, check - // if it is applicable for the current InterleaveR, otherwise look for a - // suitable VF across the Plan's VFs. - SmallVector VFs = - VFToOptimize ? SmallVector({*VFToOptimize}) - : to_vector(Plan.vectorFactors()); - std::optional NarrowedVF = - isConsecutiveInterleaveGroup(InterleaveR, VFs, TTI); - if (!NarrowedVF || (VFToOptimize && NarrowedVF != VFToOptimize)) - return nullptr; - VFToOptimize = NarrowedVF; - - // Skip read interleave groups. - if (InterleaveR->getStoredValues().empty()) - continue; - - // Narrow interleave groups, if all operands are already matching narrow - // ops. - auto *Member0 = InterleaveR->getStoredValues()[0]; - if (isAlreadyNarrow(Member0) && - all_of(InterleaveR->getStoredValues(), equal_to(Member0))) { - StoreGroups.push_back(InterleaveR); - continue; - } - - // For now, we only support full interleave groups storing load interleave - // groups. - if (all_of(enumerate(InterleaveR->getStoredValues()), [](auto Op) { - VPRecipeBase *DefR = Op.value()->getDefiningRecipe(); - if (!DefR) - return false; - auto *IR = dyn_cast(DefR); - return IR && IR->getInterleaveGroup()->isFull() && - IR->getVPValue(Op.index()) == Op.value(); - })) { - StoreGroups.push_back(InterleaveR); - continue; - } - - // Check if all values feeding InterleaveR are matching wide recipes, which - // operands that can be narrowed. - if (!canNarrowOps(InterleaveR->getStoredValues(), - VFToOptimize->isScalable())) - return nullptr; - StoreGroups.push_back(InterleaveR); - } - - if (StoreGroups.empty()) - return nullptr; - - VPBasicBlock *MiddleVPBB = Plan.getMiddleBlock(); - bool RequiresScalarEpilogue = - MiddleVPBB->getNumSuccessors() == 1 && - MiddleVPBB->getSingleSuccessor() == Plan.getScalarPreheader(); - // Bail out for tail-folding (middle block with a single successor to exit). - if (MiddleVPBB->getNumSuccessors() != 2 && !RequiresScalarEpilogue) - return nullptr; - - // All interleave groups in Plan can be narrowed for VFToOptimize. Split the - // original Plan into 2: a) a new clone which contains all VFs of Plan, except - // VFToOptimize, and b) the original Plan with VFToOptimize as single VF. - // TODO: Handle cases where only some interleave groups can be narrowed. - std::unique_ptr NewPlan; - if (size(Plan.vectorFactors()) != 1) { - NewPlan = std::unique_ptr(Plan.duplicate()); - Plan.setVF(*VFToOptimize); - NewPlan->removeVF(*VFToOptimize); - } - - // Convert InterleaveGroup \p R to a single VPWidenLoadRecipe. - SmallPtrSet NarrowedOps; - VPBasicBlock *Preheader = Plan.getVectorPreheader(); - // Narrow operation tree rooted at store groups. - for (auto *StoreGroup : StoreGroups) { - VPValue *Res = narrowInterleaveGroupOp(StoreGroup->getStoredValues(), - NarrowedOps, Preheader); - auto *SI = - cast(StoreGroup->getInterleaveGroup()->getInsertPos()); - VPBuilder(StoreGroup) - .createWidenStore(*SI, StoreGroup->getAddr(), Res, nullptr, - /*Consecutive=*/true, *StoreGroup, - StoreGroup->getDebugLoc()); - StoreGroup->eraseFromParent(); - } - - // Adjust induction to reflect that the transformed plan only processes one - // original iteration. - VPInstruction *CanIVInc = vputils::findCanonicalIVIncrement(Plan); - Type *CanIVTy = VectorLoop->getCanonicalIVType(); - VPBasicBlock *VectorPH = Plan.getVectorPreheader(); - VPBuilder PHBuilder(VectorPH, VectorPH->begin()); - - VPValue *UF = &Plan.getUF(); - VPValue *Step; - if (VFToOptimize->isScalable()) { - VPValue *VScale = - PHBuilder.createElementCount(CanIVTy, ElementCount::getScalable(1)); - Step = PHBuilder.createOverflowingOp(Instruction::Mul, {VScale, UF}, - {true, false}); - Plan.getVF().replaceAllUsesWith(VScale); - } else { - Step = UF; - Plan.getVF().replaceAllUsesWith(Plan.getConstantInt(CanIVTy, 1)); - } - // Materialize vector trip count with the narrowed step. - materializeVectorTripCount(Plan, VectorPH, /*TailByMasking=*/false, - RequiresScalarEpilogue, Step); - - CanIVInc->setOperand(1, Step); - Plan.getVFxUF().replaceAllUsesWith(Step); - - removeDeadRecipes(Plan); - assert(none_of(*VectorLoop->getEntryBasicBlock(), - IsaPred) && - "All VPVectorPointerRecipes should have been removed"); - return NewPlan; -} - void VPlanTransforms::adjustFirstOrderRecurrenceMiddleUsers(VPlan &Plan, VFRange &Range) { VPRegionBlock *VectorRegion = Plan.getVectorLoopRegion(); @@ -4649,1192 +3544,3 @@ void VPlanTransforms::optimizeFindIVReductions(VPlan &Plan, PhiR->eraseFromParent(); } } - -namespace { - -using ExtendKind = TTI::PartialReductionExtendKind; -struct ReductionExtend { - Type *SrcType = nullptr; - ExtendKind Kind = ExtendKind::PR_None; -}; - -/// Describes the extends used to compute the extended reduction operand. -/// ExtendB is optional. If ExtendB is present, ExtendsUser is a binary -/// operation. -struct ExtendedReductionOperand { - /// The recipe that consumes the extends. - VPWidenRecipe *ExtendsUser = nullptr; - /// Extend descriptions (inputs to getPartialReductionCost). - ReductionExtend ExtendA, ExtendB; -}; - -/// A chain of recipes that form a partial reduction. Matches either -/// reduction_bin_op (extended op, accumulator), or -/// reduction_bin_op (accumulator, extended op). -/// The possible forms of the "extended op" are listed in -/// matchExtendedReductionOperand. -struct VPPartialReductionChain { - /// The top-level binary operation that forms the reduction to a scalar - /// after the loop body. - VPWidenRecipe *ReductionBinOp = nullptr; - /// The user of the extends that is then reduced. - ExtendedReductionOperand ExtendedOp; - /// The recurrence kind for the entire partial reduction chain. - /// This allows distinguishing between Sub and AddWithSub recurrences, - /// when the ReductionBinOp is a Instruction::Sub. - RecurKind RK; - /// The index of the accumulator operand of ReductionBinOp. The extended op - /// is `1 - AccumulatorOpIdx`. - unsigned AccumulatorOpIdx; - unsigned ScaleFactor; - /// Optional blend to represent predication for the block that updates the - /// reduction. - VPBlendRecipe *Blend = nullptr; -}; - -// Return the incoming index of the single-use value in the blend, which is -// expected to be the predicated reduction update. -static std::optional -getBlendReductionUpdateValueIdx(VPBlendRecipe *Blend) { - assert(Blend && !Blend->isNormalized() && - Blend->getNumIncomingValues() == 2 && - "Expected a non-normalized blend with two incoming values"); - bool FirstIncomingHasOneUse = Blend->getIncomingValue(0)->hasOneUse(); - - // Only the update value should have one use (the blend). The previous - // value should always have at least two uses, the blend and the reduction. - if (FirstIncomingHasOneUse == Blend->getIncomingValue(1)->hasOneUse()) - return std::nullopt; - return FirstIncomingHasOneUse ? 0 : 1; -} - -static VPSingleDefRecipe * -optimizeExtendsForPartialReduction(VPSingleDefRecipe *Op) { - // reduce.add(mul(ext(A), C)) - // -> reduce.add(mul(ext(A), ext(trunc(C)))) - const APInt *Const; - if (match(Op, m_Mul(m_ZExtOrSExt(m_VPValue()), m_APInt(Const)))) { - auto *ExtA = cast(Op->getOperand(0)); - Instruction::CastOps ExtOpc = ExtA->getOpcode(); - Type *NarrowTy = ExtA->getOperand(0)->getScalarType(); - if (!Op->hasOneUse() || - !llvm::canConstantBeExtended( - Const, NarrowTy, TTI::getPartialReductionExtendKind(ExtOpc))) - return Op; - - VPBuilder Builder(Op); - auto *Trunc = Builder.createWidenCast(Instruction::CastOps::Trunc, - Op->getOperand(1), NarrowTy); - Type *WideTy = ExtA->getScalarType(); - Op->setOperand(1, Builder.createWidenCast(ExtOpc, Trunc, WideTy)); - return Op; - } - - // reduce.add(abs(sub(ext(A), ext(B)))) - // -> reduce.add(ext(absolute-difference(A, B))) - VPValue *X, *Y; - if (match(Op, m_WidenIntrinsic(m_Sub( - m_ZExtOrSExt(m_VPValue(X)), m_ZExtOrSExt(m_VPValue(Y)))))) { - auto *Sub = Op->getOperand(0)->getDefiningRecipe(); - auto *Ext = cast(Sub->getOperand(0)); - assert(Ext->getOpcode() == - cast(Sub->getOperand(1))->getOpcode() && - "Expected both the LHS and RHS extends to be the same"); - bool IsSigned = Ext->getOpcode() == Instruction::SExt; - VPBuilder Builder(Op); - Type *SrcTy = X->getScalarType(); - auto *FreezeX = Builder.insert(new VPWidenRecipe(Instruction::Freeze, {X})); - auto *FreezeY = Builder.insert(new VPWidenRecipe(Instruction::Freeze, {Y})); - auto *Max = Builder.insert( - new VPWidenIntrinsicRecipe(IsSigned ? Intrinsic::smax : Intrinsic::umax, - {FreezeX, FreezeY}, SrcTy)); - auto *Min = Builder.insert( - new VPWidenIntrinsicRecipe(IsSigned ? Intrinsic::smin : Intrinsic::umin, - {FreezeX, FreezeY}, SrcTy)); - auto *AbsDiff = - Builder.insert(new VPWidenRecipe(Instruction::Sub, {Max, Min})); - return Builder.createWidenCast(Instruction::CastOps::ZExt, AbsDiff, - Op->getScalarType()); - } - - // reduce.add(ext(mul(ext(A), ext(B)))) - // -> reduce.add(mul(wider_ext(A), wider_ext(B))) - // TODO: Support this optimization for float types. - if (match(Op, m_ZExtOrSExt(m_Mul(m_ZExtOrSExt(m_VPValue()), - m_ZExtOrSExt(m_VPValue()))))) { - auto *Ext = cast(Op); - auto *Mul = cast(Ext->getOperand(0)); - auto *MulLHS = cast(Mul->getOperand(0)); - auto *MulRHS = cast(Mul->getOperand(1)); - if (!Mul->hasOneUse() || - (Ext->getOpcode() != MulLHS->getOpcode() && MulLHS != MulRHS) || - MulLHS->getOpcode() != MulRHS->getOpcode()) - return Op; - VPBuilder Builder(Mul); - auto *NewLHS = Builder.createWidenCast( - MulLHS->getOpcode(), MulLHS->getOperand(0), Ext->getScalarType()); - auto *NewRHS = MulLHS == MulRHS - ? NewLHS - : Builder.createWidenCast(MulRHS->getOpcode(), - MulRHS->getOperand(0), - Ext->getScalarType()); - auto *NewMul = Mul->cloneWithOperands({NewLHS, NewRHS}); - Builder.insert(NewMul); - Op->replaceAllUsesWith(NewMul); - Op->eraseFromParent(); - Mul->eraseFromParent(); - return NewMul; - } - - return Op; -} - -static VPExpressionRecipe * -createPartialReductionExpression(VPReductionRecipe *Red) { - VPValue *VecOp = Red->getVecOp(); - - // reduce.[f]add(ext(op)) - // -> VPExpressionRecipe(op, red) - if (match(VecOp, m_WidenAnyExtend(m_VPValue()))) - return new VPExpressionRecipe(cast(VecOp), Red); - - // reduce.[f]add(neg(ext(op))) - // -> VPExpressionRecipe(op, sub/neg, red) - if (match(VecOp, m_AnyNeg(m_WidenAnyExtend(m_VPValue())))) { - auto *Neg = cast(VecOp); - auto *Ext = - cast(Neg->getOperand(Neg->getNumOperands() - 1)); - return new VPExpressionRecipe(Ext, Neg, Red); - } - - // reduce.[f]add([f]mul(ext(a), ext(b))) - // -> VPExpressionRecipe(a, b, mul, red) - if (match(VecOp, m_FMul(m_FPExt(m_VPValue()), m_FPExt(m_VPValue()))) || - match(VecOp, - m_Mul(m_ZExtOrSExt(m_VPValue()), m_ZExtOrSExt(m_VPValue())))) { - auto *Mul = cast(VecOp); - auto *ExtA = cast(Mul->getOperand(0)); - auto *ExtB = cast(Mul->getOperand(1)); - return new VPExpressionRecipe(ExtA, ExtB, Mul, Red); - } - - // reduce.fadd(fneg(fmul(fpext(a), fpext(b)))) - // -> VPExpressionRecipe(a, b, fmul, fsub, red) - if (match(VecOp, - m_FNeg(m_FMul(m_FPExt(m_VPValue()), m_FPExt(m_VPValue()))))) { - auto *FNeg = cast(VecOp); - auto *FMul = cast(FNeg->getOperand(0)); - auto *ExtA = cast(FMul->getOperand(0)); - auto *ExtB = cast(FMul->getOperand(1)); - return new VPExpressionRecipe(ExtA, ExtB, FMul, FNeg, Red); - } - - // reduce.add(neg(mul(ext(a), ext(b)))) - // -> VPExpressionRecipe(a, b, mul, sub, red) - if (match(VecOp, m_Sub(m_ZeroInt(), m_Mul(m_ZExtOrSExt(m_VPValue()), - m_ZExtOrSExt(m_VPValue()))))) { - auto *Sub = cast(VecOp); - auto *Mul = cast(Sub->getOperand(1)); - auto *ExtA = cast(Mul->getOperand(0)); - auto *ExtB = cast(Mul->getOperand(1)); - return new VPExpressionRecipe(ExtA, ExtB, Mul, Sub, Red); - } - - llvm_unreachable("Unsupported expression"); -} - -// Helper to transform a partial reduction chain into a partial reduction -// recipe. Assumes profitability has been checked. -static void transformToPartialReduction(const VPPartialReductionChain &Chain, - VPlan &Plan, - VPReductionPHIRecipe *RdxPhi) { - VPWidenRecipe *WidenRecipe = Chain.ReductionBinOp; - assert(WidenRecipe->getNumOperands() == 2 && "Expected binary operation"); - - VPValue *Accumulator = WidenRecipe->getOperand(Chain.AccumulatorOpIdx); - auto *ExtendedOp = cast( - WidenRecipe->getOperand(1 - Chain.AccumulatorOpIdx)); - - // FIXME: Do these transforms before invoking the cost-model. - ExtendedOp = optimizeExtendsForPartialReduction(ExtendedOp); - - // Sub-reductions can be implemented in two ways: - // (1) negate the operand in the vector loop (the default way). - // (2) subtract the reduced value from the init value in the middle block. - // Both ways keep the reduction itself as an 'add' reduction. - // - // The ISD nodes for partial reductions don't support folding the - // sub/negation into its operands because the following is not a valid - // transformation: - // sub(0, mul(ext(a), ext(b))) - // -> mul(ext(a), ext(sub(0, b))) - // - // It's therefore better to choose option (2) such that the partial - // reduction is always positive (starting at '0') and to do a final - // subtract in the middle block. - if ((WidenRecipe->getOpcode() == Instruction::Sub && - Chain.RK != RecurKind::Sub) || - (WidenRecipe->getOpcode() == Instruction::FSub && - Chain.RK != RecurKind::FSub)) { - VPBuilder Builder(WidenRecipe); - Type *ElemTy = ExtendedOp->getScalarType(); - VPWidenRecipe *NegRecipe; - if (WidenRecipe->getOpcode() == Instruction::FSub) { - NegRecipe = - new VPWidenRecipe(Instruction::FNeg, {ExtendedOp}, VPIRFlags(), - VPIRMetadata(), DebugLoc::getUnknown()); - } else { - auto *Zero = Plan.getZero(ElemTy); - NegRecipe = - new VPWidenRecipe(Instruction::Sub, {Zero, ExtendedOp}, VPIRFlags(), - VPIRMetadata(), DebugLoc::getUnknown()); - } - Builder.insert(NegRecipe); - ExtendedOp = NegRecipe; - } - - // Check if WidenRecipe is the final result of the reduction. If so, look - // through the Select recipe introduced by tail-folding, otherwise look - // through any Blend recipe introduced by predication for the block. - VPValue *ExitSearch = - Chain.Blend ? cast(Chain.Blend) : cast(WidenRecipe); - - VPValue *Cond = nullptr; - VPValue *ExitValue = cast_or_null( - findUserOf(ExitSearch, m_Select(m_VPValue(Cond), m_Specific(ExitSearch), - m_Specific(RdxPhi)))); - - if (Chain.Blend) { - std::optional BlendReductionIdx = - getBlendReductionUpdateValueIdx(Chain.Blend); - assert(BlendReductionIdx && - Chain.Blend->getIncomingValue(*BlendReductionIdx) == WidenRecipe && - "Expected blend to contain the reduction update"); - VPValue *BlendCond = Chain.Blend->getMask(*BlendReductionIdx); - Cond = ExitValue ? VPBuilder(WidenRecipe) - .createLogicalAnd(Cond, BlendCond, - WidenRecipe->getDebugLoc()) - : BlendCond; - } - - bool IsLastInChain = RdxPhi->getBackedgeValue() == WidenRecipe || - RdxPhi->getBackedgeValue() == ExitValue || - RdxPhi->getBackedgeValue() == Chain.Blend; - assert((!ExitValue || IsLastInChain) && - "if we found ExitValue, it must match RdxPhi's backedge value"); - - Type *PhiType = RdxPhi->getScalarType(); - RecurKind RdxKind = - PhiType->isFloatingPointTy() ? RecurKind::FAdd : RecurKind::Add; - auto *PartialRed = new VPReductionRecipe( - RdxKind, - RdxKind == RecurKind::FAdd ? WidenRecipe->getFastMathFlagsOrNone() - : FastMathFlags(), - WidenRecipe->getUnderlyingInstr(), Accumulator, ExtendedOp, Cond, - RdxUnordered{/*VFScaleFactor=*/Chain.ScaleFactor}); - PartialRed->insertBefore(WidenRecipe); - - if (ExitValue) - ExitValue->replaceAllUsesWith(PartialRed); - if (Chain.Blend) - Chain.Blend->replaceAllUsesWith(PartialRed); - WidenRecipe->replaceAllUsesWith(PartialRed); - - // For cost-model purposes, fold this into a VPExpression. - VPExpressionRecipe *E = createPartialReductionExpression(PartialRed); - E->insertBefore(WidenRecipe); - PartialRed->replaceAllUsesWith(E); - - // We only need to update the PHI node once, which is when we find the - // last reduction in the chain. - if (!IsLastInChain) - return; - - // Scale the PHI and ReductionStartVector by the VFScaleFactor - assert(RdxPhi->getVFScaleFactor() == 1 && "scale factor must not be set"); - RdxPhi->setVFScaleFactor(Chain.ScaleFactor); - - auto *StartInst = cast(RdxPhi->getStartValue()); - assert(StartInst->getOpcode() == VPInstruction::ReductionStartVector); - auto *NewScaleFactor = Plan.getConstantInt(32, Chain.ScaleFactor); - StartInst->setOperand(2, NewScaleFactor); - - // If this is the last value in a sub-reduction chain, then update the PHI - // node to start at `0` and update the reduction-result to subtract from - // the PHI's start value. - if (Chain.RK != RecurKind::Sub && Chain.RK != RecurKind::FSub) - return; - - VPValue *OldStartValue = StartInst->getOperand(0); - StartInst->setOperand(0, StartInst->getOperand(1)); - - // Replace reduction_result by 'sub (startval, reductionresult)'. - VPInstruction *RdxResult = vputils::findComputeReductionResult(RdxPhi); - assert(RdxResult && "Could not find reduction result"); - - VPBuilder Builder = VPBuilder::getToInsertAfter(RdxResult); - unsigned SubOpc = Chain.RK == RecurKind::FSub ? Instruction::BinaryOps::FSub - : Instruction::BinaryOps::Sub; - VPInstruction *NewResult = Builder.createNaryOp( - SubOpc, {OldStartValue, RdxResult}, VPIRFlags::getDefaultFlags(SubOpc), - RdxPhi->getDebugLoc()); - RdxResult->replaceUsesWithIf( - NewResult, - [&NewResult](VPUser &U, unsigned Idx) { return &U != NewResult; }); -} - -/// Returns the cost of a link in a partial-reduction chain for a given VF. -static InstructionCost -getPartialReductionLinkCost(VPCostContext &CostCtx, - const VPPartialReductionChain &Link, - ElementCount VF) { - Type *RdxType = Link.ReductionBinOp->getScalarType(); - const ExtendedReductionOperand &ExtendedOp = Link.ExtendedOp; - std::optional BinOpc = std::nullopt; - // If ExtendB is not none, then the "ExtendsUser" is the binary operation. - if (ExtendedOp.ExtendB.Kind != ExtendKind::PR_None) - BinOpc = ExtendedOp.ExtendsUser->getOpcode(); - - std::optional Flags; - if (RdxType->isFloatingPointTy()) - Flags = Link.ReductionBinOp->getFastMathFlagsOrNone(); - - auto GetLinkOpcode = [&Link]() -> unsigned { - switch (Link.RK) { - case RecurKind::Sub: - return Instruction::Add; - case RecurKind::FSub: - return Instruction::FAdd; - default: - return Link.ReductionBinOp->getOpcode(); - } - }; - - return CostCtx.TTI.getPartialReductionCost( - GetLinkOpcode(), ExtendedOp.ExtendA.SrcType, ExtendedOp.ExtendB.SrcType, - RdxType, VF, ExtendedOp.ExtendA.Kind, ExtendedOp.ExtendB.Kind, BinOpc, - CostCtx.CostKind, Flags); -} - -static ExtendKind getPartialReductionExtendKind(VPWidenCastRecipe *Cast) { - return TTI::getPartialReductionExtendKind(Cast->getOpcode()); -} - -/// Checks if \p Op (which is an operand of \p UpdateR) is an extended reduction -/// operand. This is an operand where the source of the value (e.g. a load) has -/// been extended (sext, zext, or fpext) before it is used in the reduction. -/// -/// Possible forms matched by this function: -/// - UpdateR(PrevValue, ext(...)) -/// - UpdateR(PrevValue, mul(ext(...), ext(...))) -/// - UpdateR(PrevValue, mul(ext(...), Constant)) -/// - UpdateR(PrevValue, ext(mul(ext(...), ext(...)))) -/// - UpdateR(PrevValue, ext(mul(ext(...), Constant))) -/// - UpdateR(PrevValue, abs(sub(ext(...), ext(...))) -/// -/// Note: The second operand of UpdateR corresponds to \p Op in the examples. -static std::optional -matchExtendedReductionOperand(VPWidenRecipe *UpdateR, VPValue *Op) { - assert(is_contained(UpdateR->operands(), Op) && - "Op should be operand of UpdateR"); - - // Try matching an absolute difference operand of the form - // `abs(sub(ext(A), ext(B)))`. This will be later transformed into - // `ext(absolute-difference(A, B))`. This allows us to perform the absolute - // difference on a wider type and get the extend for "free" from the partial - // reduction. - VPValue *X, *Y; - if (Op->hasOneUse() && - match(Op, m_WidenIntrinsic( - m_OneUse(m_Sub(m_WidenAnyExtend(m_VPValue(X)), - m_WidenAnyExtend(m_VPValue(Y))))))) { - auto *Abs = cast(Op); - auto *Sub = cast(Abs->getOperand(0)); - auto *LHSExt = cast(Sub->getOperand(0)); - auto *RHSExt = cast(Sub->getOperand(1)); - Type *LHSInputType = X->getScalarType(); - Type *RHSInputType = Y->getScalarType(); - if (LHSInputType != RHSInputType || - LHSExt->getOpcode() != RHSExt->getOpcode()) - return std::nullopt; - // Note: This is essentially the same as matching ext(...) as we will - // rewrite this operand to ext(absolute-difference(A, B)). - return ExtendedReductionOperand{ - Sub, - /*ExtendA=*/{LHSInputType, getPartialReductionExtendKind(LHSExt)}, - /*ExtendB=*/{}}; - } - - std::optional OuterExtKind; - if (match(Op, m_WidenAnyExtend(m_VPValue()))) { - auto *CastRecipe = cast(Op); - VPValue *CastSource = CastRecipe->getOperand(0); - OuterExtKind = getPartialReductionExtendKind(CastRecipe); - if (match(CastSource, m_Mul(m_VPValue(), m_VPValue())) || - match(CastSource, m_FMul(m_VPValue(), m_VPValue()))) { - // Match: ext(mul(...)) - // Record the outer extend kind and set `Op` to the mul. We can then match - // this as a binary operation. Note: We can optimize out the outer extend - // by widening the inner extends to match it. See - // optimizeExtendsForPartialReduction. - Op = CastSource; - } else { - return ExtendedReductionOperand{ - UpdateR, - /*ExtendA=*/{CastSource->getScalarType(), *OuterExtKind}, - /*ExtendB=*/{}}; - } - } - - if (!Op->hasOneUse()) - return std::nullopt; - - VPWidenRecipe *MulOp = dyn_cast(Op); - if (!MulOp || - !is_contained({Instruction::Mul, Instruction::FMul}, MulOp->getOpcode())) - return std::nullopt; - - // The rest of the matching assumes `Op` is a (possibly extended) mul - // operation. - - VPValue *LHS = MulOp->getOperand(0); - VPValue *RHS = MulOp->getOperand(1); - - // The LHS of the operation must always be an extend. - if (!match(LHS, m_WidenAnyExtend(m_VPValue()))) - return std::nullopt; - - auto *LHSCast = cast(LHS); - Type *LHSInputType = LHSCast->getOperand(0)->getScalarType(); - ExtendKind LHSExtendKind = getPartialReductionExtendKind(LHSCast); - - // The RHS of the operation can be an extend or a constant integer. - const APInt *RHSConst = nullptr; - VPWidenCastRecipe *RHSCast = nullptr; - if (match(RHS, m_WidenAnyExtend(m_VPValue()))) - RHSCast = cast(RHS); - else if (!match(RHS, m_APInt(RHSConst)) || - !canConstantBeExtended(RHSConst, LHSInputType, LHSExtendKind)) - return std::nullopt; - - // The outer extend kind must match the inner extends for folding. - for (VPWidenCastRecipe *Cast : {LHSCast, RHSCast}) - if (Cast && OuterExtKind && - getPartialReductionExtendKind(Cast) != OuterExtKind) - return std::nullopt; - - Type *RHSInputType = LHSInputType; - ExtendKind RHSExtendKind = LHSExtendKind; - if (RHSCast) { - RHSInputType = RHSCast->getOperand(0)->getScalarType(); - RHSExtendKind = getPartialReductionExtendKind(RHSCast); - } - - return ExtendedReductionOperand{ - MulOp, {LHSInputType, LHSExtendKind}, {RHSInputType, RHSExtendKind}}; -} - -/// Examines each operation in the reduction chain corresponding to \p RedPhiR, -/// and determines if the target can use a cheaper operation with a wider -/// per-iteration input VF and narrower PHI VF. If successful, returns the chain -/// of operations in the reduction. -static std::optional> -getScaledReductions(VPReductionPHIRecipe *RedPhiR) { - // Get the backedge value from the reduction PHI and find the - // ComputeReductionResult that uses it (directly or through a select for - // predicated reductions). - auto *RdxResult = vputils::findComputeReductionResult(RedPhiR); - if (!RdxResult) - return std::nullopt; - VPValue *ExitValue = RdxResult->getOperand(0); - match(ExitValue, m_Select(m_VPValue(), m_VPValue(ExitValue), m_VPValue())); - - SmallVector Chain; - RecurKind RK = RedPhiR->getRecurrenceKind(); - Type *PhiType = RedPhiR->getScalarType(); - TypeSize PHISize = PhiType->getPrimitiveSizeInBits(); - - // Work backwards from the ExitValue examining each reduction operation. - VPValue *CurrentValue = ExitValue; - while (CurrentValue != RedPhiR) { - VPBlendRecipe *Blend = dyn_cast(CurrentValue); - std::optional BlendReductionIdx; - if (Blend) { - assert(!Blend->isNormalized() && "Expect Blend not to be normalized."); - if (Blend->getNumIncomingValues() != 2) - return std::nullopt; - - BlendReductionIdx = getBlendReductionUpdateValueIdx(Blend); - if (!BlendReductionIdx) - return std::nullopt; - - CurrentValue = Blend->getIncomingValue(*BlendReductionIdx); - } - - auto *UpdateR = dyn_cast(CurrentValue); - if (!UpdateR || !Instruction::isBinaryOp(UpdateR->getOpcode())) - return std::nullopt; - - VPValue *Op = UpdateR->getOperand(1); - VPValue *PrevValue = UpdateR->getOperand(0); - - // Find the extended operand. The other operand (PrevValue) is the next link - // in the reduction chain. - std::optional ExtendedOp = - matchExtendedReductionOperand(UpdateR, Op); - if (!ExtendedOp) { - ExtendedOp = matchExtendedReductionOperand(UpdateR, PrevValue); - if (!ExtendedOp) - return std::nullopt; - std::swap(Op, PrevValue); - } - - // Look for VPBlend(reduce(PrevValue, Op), PrevValue), where - // reduce is equal to CurrentValue. This can be lowered as - // a conditional reduction by hoisting the select to the inputs. - if (Blend && Blend->getIncomingValue(1 - *BlendReductionIdx) != PrevValue) - return std::nullopt; - - Type *ExtSrcType = ExtendedOp->ExtendA.SrcType; - TypeSize ExtSrcSize = ExtSrcType->getPrimitiveSizeInBits(); - if (!PHISize.hasKnownScalarFactor(ExtSrcSize)) - return std::nullopt; - - VPPartialReductionChain Link( - {UpdateR, *ExtendedOp, RK, - PrevValue == UpdateR->getOperand(0) ? 0U : 1U, - static_cast(PHISize.getKnownScalarFactor(ExtSrcSize)), - Blend}); - Chain.push_back(Link); - CurrentValue = PrevValue; - } - - // The chain links were collected by traversing backwards from the exit value. - // Reverse the chains so they are in program order. - std::reverse(Chain.begin(), Chain.end()); - return Chain; -} -} // namespace - -void VPlanTransforms::createPartialReductions(VPlan &Plan, - VPCostContext &CostCtx, - VFRange &Range) { - // Find all possible valid partial reductions, grouping chains by their PHI. - // This grouping allows invalidating the whole chain, if any link is not a - // valid partial reduction. - MapVector> - ChainsByPhi; - VPBasicBlock *HeaderVPBB = Plan.getVectorLoopRegion()->getEntryBasicBlock(); - for (VPRecipeBase &R : HeaderVPBB->phis()) { - auto *RedPhiR = dyn_cast(&R); - if (!RedPhiR) - continue; - - if (auto Chains = getScaledReductions(RedPhiR)) - ChainsByPhi.try_emplace(RedPhiR, std::move(*Chains)); - } - - if (ChainsByPhi.empty()) - return; - - // Build set of partial reduction operations and blends for user validation - // and a map of reduction bin ops to their scale factors for scale validation. - SmallPtrSet PartialReductionOps; - SmallPtrSet PartialReductionBlends; - DenseMap ScaledReductionMap; - for (const auto &[_, Chains] : ChainsByPhi) - for (const VPPartialReductionChain &Chain : Chains) { - PartialReductionOps.insert(Chain.ExtendedOp.ExtendsUser); - if (Chain.Blend) - PartialReductionBlends.insert(Chain.Blend); - ScaledReductionMap[Chain.ReductionBinOp] = Chain.ScaleFactor; - } - - // A partial reduction is invalid if any of its extends are used by - // something that isn't another partial reduction. This is because the - // extends are intended to be lowered along with the reduction itself. - auto ExtendUsersValid = [&](VPValue *Ext) { - return !isa(Ext) || all_of(Ext->users(), [&](VPUser *U) { - return PartialReductionOps.contains(cast(U)); - }); - }; - - auto IsProfitablePartialReductionChainForVF = - [&](ArrayRef Chain, ElementCount VF) -> bool { - InstructionCost PartialCost = 0, RegularCost = 0; - - // The chain is a profitable partial reduction chain if the cost of handling - // the entire chain is cheaper when using partial reductions than when - // handling the entire chain using regular reductions. - for (const VPPartialReductionChain &Link : Chain) { - const ExtendedReductionOperand &ExtendedOp = Link.ExtendedOp; - InstructionCost LinkCost = getPartialReductionLinkCost(CostCtx, Link, VF); - if (!LinkCost.isValid()) - return false; - - PartialCost += LinkCost; - RegularCost += Link.ReductionBinOp->computeCost(VF, CostCtx); - // If ExtendB is not none, then the "ExtendsUser" is the binary operation. - if (ExtendedOp.ExtendB.Kind != ExtendKind::PR_None) - RegularCost += ExtendedOp.ExtendsUser->computeCost(VF, CostCtx); - for (VPValue *Op : ExtendedOp.ExtendsUser->operands()) - if (auto *Extend = dyn_cast(Op)) - RegularCost += Extend->computeCost(VF, CostCtx); - } - return PartialCost.isValid() && PartialCost < RegularCost; - }; - - // Validate chains: check that extends are only used by partial reductions, - // and that reduction bin ops are only used by other partial reductions with - // matching scale factors, are outside the loop region or the select - // introduced by tail-folding. Otherwise we would create users of scaled - // reductions where the types of the other operands don't match. - for (auto &[RedPhiR, Chains] : ChainsByPhi) { - for (const VPPartialReductionChain &Chain : Chains) { - if (!all_of(Chain.ExtendedOp.ExtendsUser->operands(), ExtendUsersValid)) { - Chains.clear(); - break; - } - auto UseIsValid = [&, RedPhiR = RedPhiR](VPUser *U) { - if (auto *PhiR = dyn_cast(U)) - return PhiR == RedPhiR; - auto *R = cast(U); - - if (auto *Blend = dyn_cast(R)) - return Blend == Chain.Blend || PartialReductionBlends.contains(Blend); - - return Chain.ScaleFactor == ScaledReductionMap.lookup_or(R, 0) || - match(R, m_ComputeReductionResult( - m_Specific(Chain.ReductionBinOp))) || - match(R, m_Select(m_VPValue(), m_Specific(Chain.ReductionBinOp), - m_Specific(RedPhiR))); - }; - if (!all_of(Chain.ReductionBinOp->users(), UseIsValid)) { - Chains.clear(); - break; - } - - // Check if the compute-reduction-result is used by a sunk store. - // TODO: Also form partial reductions in those cases. - if (auto *RdxResult = vputils::findComputeReductionResult(RedPhiR)) { - if (any_of(RdxResult->users(), [](VPUser *U) { - auto *RepR = dyn_cast(U); - return RepR && RepR->getOpcode() == Instruction::Store; - })) { - Chains.clear(); - break; - } - } - } - - // Clear the chain if it is not profitable. - if (!LoopVectorizationPlanner::getDecisionAndClampRange( - [&, &Chains = Chains](ElementCount VF) { - return IsProfitablePartialReductionChainForVF(Chains, VF); - }, - Range)) - Chains.clear(); - } - - for (auto &[Phi, Chains] : ChainsByPhi) - for (const VPPartialReductionChain &Chain : Chains) - transformToPartialReduction(Chain, Plan, Phi); -} - -void VPlanTransforms::makeMemOpWideningDecisions(VPlan &Plan, VFRange &Range, - VPRecipeBuilder &RecipeBuilder, - VPCostContext &CostCtx) { - // Collect all loads/stores first. We will start with ones having simpler - // decisions followed by more complex ones that are potentially - // guided/dependent on the simpler ones. - SmallVector MemOps; - for (VPBasicBlock *VPBB : - VPBlockUtils::blocksOnly(vp_depth_first_shallow( - Plan.getVectorLoopRegion()->getEntryBasicBlock()))) { - for (VPRecipeBase &R : *VPBB) { - auto *VPI = dyn_cast(&R); - if (VPI && VPI->getUnderlyingValue() && - is_contained({Instruction::Load, Instruction::Store}, - VPI->getOpcode())) - MemOps.push_back(VPI); - } - } - - // Few helpers to process different kinds of memory operations. - - // To be used as argument to `VPlanTransforms::runPass` which explicitly - // specified pass name, hence `VPlan &` parameter. - auto ProcessSubset = [&](VPlan &, auto ProcessVPInst) { - SmallVector RemainingMemOps; - for (VPInstruction *VPI : MemOps) { - if (!ProcessVPInst(VPI)) - RemainingMemOps.push_back(VPI); - } - - MemOps.clear(); - std::swap(MemOps, RemainingMemOps); - }; - - auto ReplaceWith = [&](VPInstruction *VPI, VPRecipeBase *New) { - assert(New->getParent() && "New recipe must have been inserted"); - if (VPI->getOpcode() == Instruction::Load) - VPI->replaceAllUsesWith(New->getVPSingleValue()); - VPI->eraseFromParent(); - - // VPI has been processed. - return true; - }; - - auto Scalarize = [&](VPInstruction *VPI) { - return ReplaceWith(VPI, VPBuilder(VPI).insert( - RecipeBuilder.handleReplication(VPI, Range))); - }; - - VPBasicBlock *MiddleVPBB = Plan.getMiddleBlock(); - VPBuilder FinalRedStoresBuilder(MiddleVPBB, MiddleVPBB->getFirstNonPhi()); - VPlanTransforms::runPass( - "lowerMemoryIdioms", ProcessSubset, Plan, [&](VPInstruction *VPI) { - if (RecipeBuilder.replaceWithFinalIfReductionStore( - VPI, FinalRedStoresBuilder)) - return true; - - // Filter out scalar VPlan for the remaining idioms. - if (LoopVectorizationPlanner::getDecisionAndClampRange( - [](ElementCount VF) { return VF.isScalar(); }, Range)) - return false; - - if (VPHistogramRecipe *Histogram = RecipeBuilder.widenIfHistogram(VPI)) - return ReplaceWith(VPI, VPBuilder(VPI).insert(Histogram)); - - return false; - }); - - // Filter out scalar VPlan for the remaining memory operations. - if (LoopVectorizationPlanner::getDecisionAndClampRange( - [](ElementCount VF) { return VF.isScalar(); }, Range)) - return; - - // If the instruction's allocated size doesn't equal it's type size, it - // requires padding and will be scalarized. - VPlanTransforms::runPass( - "scalarizeMemOpsWithIrregularTypes", ProcessSubset, Plan, - [&](VPInstruction *VPI) { - Instruction *I = VPI->getUnderlyingInstr(); - if (hasIrregularType(getLoadStoreType(I), I->getDataLayout())) - return Scalarize(VPI); - - return false; - }); - - if (!RecipeBuilder.prefersVectorizedAddressing()) { - VPlanTransforms::runPass( - "makeVPlanMemOpDecision", ProcessSubset, Plan, [&](VPInstruction *VPI) { - Instruction *I = VPI->getUnderlyingInstr(); - bool IsLoad = VPI->getOpcode() == Instruction::Load; - if (RecipeBuilder.isPredicatedInst(I) || !IsLoad || - !vputils::isUsedByLoadStoreAddress(VPI)) - return false; - - // Scalarize loads used as addresses, matching the legacy CM. The load - // is single-scalar if the pointer is loop-invariant, otherwise it is - // replicated per-lane. No mask is needed as the load is not - // predicated. - VPValue *Ptr = VPI->getOperand(0); - const SCEV *PtrSCEV = - vputils::getSCEVExprForVPValue(Ptr, CostCtx.PSE, CostCtx.L); - bool IsSingleScalarLoad = - !isa(PtrSCEV) && - CostCtx.PSE.getSE()->isLoopInvariant(PtrSCEV, CostCtx.L); - - ReplaceWith(VPI, - VPBuilder(VPI).insert(new VPReplicateRecipe( - I, Ptr, /*IsSingleScalar=*/IsSingleScalarLoad, - /*Mask=*/nullptr, *VPI, *VPI, VPI->getDebugLoc()))); - return true; - }); - } - - // Widen unit-stride consecutive accesses, matching the legacy CM. Both - // forward (stride +1) and reverse (stride -1) accesses are handled. - VPlanTransforms::runPass( - "widenConsecutiveMemOps", ProcessSubset, Plan, [&](VPInstruction *VPI) { - Instruction *I = VPI->getUnderlyingInstr(); - bool IsLoad = VPI->getOpcode() == Instruction::Load; - VPValue *Ptr = VPI->getOperand(!IsLoad); - Type *ScalarTy = - IsLoad ? VPI->getScalarType() : VPI->getOperand(0)->getScalarType(); - std::optional Stride = - getConstantStride(Ptr, ScalarTy, CostCtx.PSE, CostCtx.L); - if (Stride != 1 && Stride != -1) - return false; - bool Reverse = Stride == -1; - - // A predicated access can only be widened (rather than scalarized) if - // the target supports a masked load/store for it. - // TODO: Determine if a load/store needs predication directly in VPlan. - bool IsPredicated = RecipeBuilder.isPredicatedInst(I); - if (IsPredicated && !CostCtx.Config.isLegalMaskedLoadOrStore( - IsLoad, ScalarTy, getLoadStoreAlignment(I), - getLoadStoreAddressSpace(I))) - return false; - - VPBuilder Builder(VPI); - VPSingleDefRecipe *VectorPtr = Builder.createConsecutiveVectorPointer( - Ptr, ScalarTy, Reverse, VPI->getDebugLoc()); - - VPValue *Mask = IsPredicated ? VPI->getMask() : nullptr; - // Reverse the mask so it matches the reversed access order. - if (Reverse && Mask) - Mask = Builder.createNaryOp(VPInstruction::Reverse, Mask, - VPI->getDebugLoc()); - - if (IsLoad) { - VPSingleDefRecipe *Load = Builder.createWidenLoad( - *cast(I), VectorPtr, Mask, - /*Consecutive=*/true, *VPI, VPI->getDebugLoc()); - // Reverse the loaded values back into program order. - if (Reverse) - Load = Builder.createNaryOp(VPInstruction::Reverse, Load, - VPI->getDebugLoc()); - return ReplaceWith(VPI, Load); - } - - VPValue *StoredVal = VPI->getOperand(0); - if (Reverse) - // Reverse the stored values so they are written in descending order. - StoredVal = Builder.createNaryOp(VPInstruction::Reverse, StoredVal, - VPI->getDebugLoc()); - - auto *StoreR = Builder.createWidenStore( - *cast(I), VectorPtr, StoredVal, Mask, - /*Consecutive=*/true, *VPI, VPI->getDebugLoc()); - return ReplaceWith(VPI, StoreR); - }); - - VPlanTransforms::runPass("delegateMemOpWideningToLegacyCM", ProcessSubset, - Plan, [&](VPInstruction *VPI) { - if (VPRecipeBase *Recipe = - RecipeBuilder.tryToWidenMemory(VPI, Range)) - return ReplaceWith(VPI, Recipe); - - return Scalarize(VPI); - }); -} - -void VPlanTransforms::makeScalarizationDecisions(VPlan &Plan, VFRange &Range) { - if (LoopVectorizationPlanner::getDecisionAndClampRange( - [&](ElementCount VF) { return VF.isScalar(); }, Range)) - return; - - PostOrderTraversal> POT( - Plan.getEntry()); - for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly(POT)) { - for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) { - auto *VPI = dyn_cast(&R); - if (!VPI) - continue; - - auto *I = cast_or_null(VPI->getUnderlyingValue()); - // Wouldn't be able to create a `VPReplicateRecipe` anyway. - if (!I) - continue; - - // If executing other lanes produces side-effects we can't avoid them. - if (VPI->mayHaveSideEffects()) - continue; - - // We want to drop the mask operand, verify we can safely do that. - if (VPI->isMasked() && !VPI->isSafeToSpeculativelyExecute()) - continue; - - // Avoid rewriting IV increment as that interferes with - // `removeRedundantCanonicalIVs`. - if (VPI->getOpcode() == Instruction::Add && - any_of(VPI->operands(), IsaPred)) - continue; - - // Other lanes are needed - can't drop them. - if (!vputils::onlyFirstLaneUsed(VPI)) - continue; - - auto *Recipe = VPBuilder::createSingleScalarOp( - VPI->getOpcode(), VPI->operandsWithoutMask(), /*Mask=*/nullptr, *VPI, - *VPI, VPI->getDebugLoc(), I); - Recipe->insertBefore(VPI); - VPI->replaceAllUsesWith(Recipe); - VPI->eraseFromParent(); - } - } -} - -/// Returns true if \p Info's parameter kinds are compatible with \p Args. -static bool areVFParamsOk(const VFInfo &Info, ArrayRef Args, - PredicatedScalarEvolution &PSE, const Loop *L) { - ScalarEvolution *SE = PSE.getSE(); - return all_of(Info.Shape.Parameters, [&](VFParameter Param) { - switch (Param.ParamKind) { - case VFParamKind::Vector: - case VFParamKind::GlobalPredicate: - return true; - case VFParamKind::OMP_Uniform: - return SE->isSCEVable(Args[Param.ParamPos]->getScalarType()) && - SE->isLoopInvariant( - vputils::getSCEVExprForVPValue(Args[Param.ParamPos], PSE, L), - L); - case VFParamKind::OMP_Linear: - return match(vputils::getSCEVExprForVPValue(Args[Param.ParamPos], PSE, L), - m_scev_AffineAddRec( - m_SCEV(), m_scev_SpecificSInt(Param.LinearStepOrPos), - m_SpecificLoop(L))); - default: - return false; - } - }); -} - -/// Find a vector variant of \p CI for \p VF, respecting \p MaskRequired. -/// Returns the variant function, or nullptr. Masked variants are assumed to -/// take the mask as a trailing parameter. -static Function *findVectorVariant(CallInst *CI, ArrayRef Args, - ElementCount VF, bool MaskRequired, - PredicatedScalarEvolution &PSE, - const Loop *L) { - if (CI->isNoBuiltin()) - return nullptr; - auto Mappings = VFDatabase::getMappings(*CI); - const auto *It = find_if(Mappings, [&](const VFInfo &Info) { - return Info.Shape.VF == VF && (!MaskRequired || Info.isMasked()) && - areVFParamsOk(Info, Args, PSE, L); - }); - if (It == Mappings.end()) - return nullptr; - return CI->getModule()->getFunction(It->VectorName); -} - -namespace { -/// The outcome of choosing how to widen a call at a given VF. -struct CallWideningDecision { - enum class KindTy { Scalarize, Intrinsic, VectorVariant }; - CallWideningDecision(KindTy Kind, Function *Variant = nullptr) - : Kind(Kind), Variant(Variant) {} - KindTy Kind; - - /// Set when Kind == VectorVariant. - Function *Variant; - - bool operator==(const CallWideningDecision &Other) const { - return Kind == Other.Kind && Variant == Other.Variant; - } -}; -} // namespace - -/// Pick the cheapest widening for the call \p VPI at \p VF among scalarization, -/// vector intrinsic, and vector library variant. -static CallWideningDecision decideCallWidening(VPInstruction &VPI, - ArrayRef Ops, - ElementCount VF, - VPCostContext &CostCtx) { - auto *CI = cast(VPI.getUnderlyingInstr()); - - // Scalar VFs and calls forced or known to scalarize always replicate. - if (VF.isScalar() || CostCtx.willBeScalarized(CI, VF)) - return CallWideningDecision::KindTy::Scalarize; - - auto *CalledFn = cast( - VPI.getOperand(VPI.getNumOperandsWithoutMask() - 1)->getLiveInIRValue()); - Type *ResultTy = VPI.getScalarType(); - Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, &CostCtx.TLI); - bool MaskRequired = CostCtx.isMaskRequired(CI); - - // Pseudo intrinsics (assume, lifetime, ...) are always scalarized. - if (ID && VPCostContext::isFreeScalarIntrinsic(ID)) - return CallWideningDecision::KindTy::Scalarize; - - InstructionCost ScalarCost = - VPReplicateRecipe::computeCallCost(CalledFn, ResultTy, Ops, - /*IsSingleScalar=*/false, VF, CostCtx); - - Function *VecFunc = - findVectorVariant(CI, Ops, VF, MaskRequired, CostCtx.PSE, CostCtx.L); - InstructionCost VecCallCost = InstructionCost::getInvalid(); - if (VecFunc) - VecCallCost = VPWidenCallRecipe::computeCallCost(VecFunc, CostCtx); - - // Prefer the intrinsic if it is at least as cheap as scalarizing and any - // available vector variant. - if (ID) { - InstructionCost IntrinsicCost = - VPWidenIntrinsicRecipe::computeCallCost(ID, Ops, VPI, VF, CostCtx); - if (IntrinsicCost.isValid() && ScalarCost >= IntrinsicCost && - (!VecFunc || VecCallCost >= IntrinsicCost)) - return CallWideningDecision::KindTy::Intrinsic; - } - - // Otherwise, use a vector library variant when it beats scalarizing. - if (VecFunc && ScalarCost >= VecCallCost) - return {CallWideningDecision::KindTy::VectorVariant, VecFunc}; - - return CallWideningDecision::KindTy::Scalarize; -} - -void VPlanTransforms::makeCallWideningDecisions(VPlan &Plan, VFRange &Range, - VPRecipeBuilder &RecipeBuilder, - VPCostContext &CostCtx) { - for (VPBasicBlock *VPBB : VPBlockUtils::blocksAs( - vp_depth_first_shallow(Plan.getVectorLoopRegion()->getEntry()))) { - for (VPRecipeBase &R : make_early_inc_range(*VPBB)) { - auto *VPI = dyn_cast(&R); - if (!VPI || !VPI->getUnderlyingValue() || - VPI->getOpcode() != Instruction::Call) - continue; - - auto *CI = cast(VPI->getUnderlyingInstr()); - SmallVector Ops(VPI->op_begin(), - VPI->op_begin() + CI->arg_size()); - - CallWideningDecision Decision = - decideCallWidening(*VPI, Ops, Range.Start, CostCtx); - LoopVectorizationPlanner::getDecisionAndClampRange( - [&](ElementCount VF) { - return Decision == decideCallWidening(*VPI, Ops, VF, CostCtx); - }, - Range); - - VPSingleDefRecipe *Replacement = nullptr; - switch (Decision.Kind) { - case CallWideningDecision::KindTy::Intrinsic: { - Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, &CostCtx.TLI); - Type *ResultTy = VPI->getScalarType(); - Replacement = new VPWidenIntrinsicRecipe(*CI, ID, Ops, ResultTy, *VPI, - *VPI, VPI->getDebugLoc()); - break; - } - case CallWideningDecision::KindTy::VectorVariant: { - // Masked variants take the mask as a trailing parameter, so they have - // one more parameter than the original call's arguments. - if (Decision.Variant->arg_size() > Ops.size()) { - VPValue *Mask = VPI->isMasked() ? VPI->getMask() : Plan.getTrue(); - Ops.push_back(Mask); - } - Ops.push_back(VPI->getOperand(VPI->getNumOperandsWithoutMask() - 1)); - Replacement = new VPWidenCallRecipe(CI, Decision.Variant, Ops, *VPI, - *VPI, VPI->getDebugLoc()); - break; - } - case CallWideningDecision::KindTy::Scalarize: - Replacement = RecipeBuilder.handleReplication(VPI, Range); - break; - } - - Replacement->insertBefore(VPI); - VPI->replaceAllUsesWith(Replacement); - VPI->eraseFromParent(); - } - } -} - -void VPlanTransforms::convertToStridedAccesses(VPlan &Plan, - PredicatedScalarEvolution &PSE, - Loop &L, VPCostContext &Ctx, - VFRange &Range) { - if (Plan.hasScalarVFOnly()) - return; - - VPRegionBlock *VectorLoop = Plan.getVectorLoopRegion(); - VPValue *I32VF = nullptr; - for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly( - vp_depth_first_shallow(VectorLoop->getEntry()))) { - for (VPRecipeBase &R : make_early_inc_range(*VPBB)) { - auto *LoadR = dyn_cast(&R); - // TODO: Support strided store. - // TODO: Transform reverse access into strided access with -1 stride. - // TODO: Transform gather/scatter with uniform address into strided access - // with 0 stride. - // TODO: Transform interleave access into multiple strided accesses. - if (!LoadR || LoadR->isConsecutive()) - continue; - - VPValue *Ptr = LoadR->getAddr(); - // Check if this is a strided access by analyzing the address SCEV for an - // affine addRec. - const SCEV *PtrSCEV = vputils::getSCEVExprForVPValue(Ptr, PSE, &L); - const SCEV *Start; - const SCEVConstant *Step; - // TODO: Support non-constant loop invariant stride. - if (!match(PtrSCEV, - m_scev_AffineAddRec(m_SCEV(Start), m_SCEVConstant(Step), - m_SpecificLoop(&L)))) - continue; - - Type *LoadTy = LoadR->getScalarType(); - Align Alignment = LoadR->getAlign(); - auto IsProfitable = [&](ElementCount VF) { - Type *DataTy = toVectorTy(LoadTy, VF); - if (!Ctx.TTI.isLegalStridedLoadStore(DataTy, Alignment)) - return false; - const InstructionCost CurrentCost = LoadR->computeCost(VF, Ctx); - const InstructionCost StridedLoadStoreCost = - VPWidenMemIntrinsicRecipe::computeMemIntrinsicCost( - Intrinsic::experimental_vp_strided_load, DataTy, - LoadR->isMasked(), Alignment, Ctx); - return StridedLoadStoreCost < CurrentCost; - }; - - if (!LoopVectorizationPlanner::getDecisionAndClampRange(IsProfitable, - Range)) - continue; - - // Invalidate the legacy widening decision so the cost of replaced load is - // not counted during precomputeCosts. - // TODO: Remove once the legacy exit cost computation is retired. - for (ElementCount VF : Range) - Ctx.invalidateWideningDecision(&LoadR->getIngredient(), VF); - - // Get VF as i32 for the vector length operand. - if (!I32VF) { - VPBuilder Builder(Plan.getVectorPreheader()); - I32VF = Builder.createScalarZExtOrTrunc( - &Plan.getVF(), Type::getInt32Ty(Plan.getContext()), - DebugLoc::getUnknown()); - } - - VPBuilder Builder(LoadR); - // Create the base pointer of strided access. - // TODO: reuse VPDerivedIVRecipe for base pointer computation when it - // supports a general VPValue as the start value. - VPValue *StartVPV = - VPSCEVExpander(Builder, *PSE.getSE(), LoadR->getDebugLoc()) - .tryToExpand(Start); - if (!StartVPV) - StartVPV = VPBuilder(Plan.getEntry()).createExpandSCEV(Start); - VPValue *StrideInBytes = Plan.getOrAddLiveIn(Step->getValue()); - Type *IndexTy = Plan.getDataLayout().getIndexType(Ptr->getScalarType()); - assert(IndexTy == StrideInBytes->getScalarType() && - "Stride type from SCEV must match the index type"); - VPValue *CanIV = Builder.createScalarZExtOrTrunc( - VectorLoop->getCanonicalIV(), IndexTy, DebugLoc::getUnknown()); - auto *AddRecPtr = cast(PtrSCEV); - auto *Offset = Builder.createOverflowingOp( - Instruction::Mul, {CanIV, StrideInBytes}, - {AddRecPtr->hasNoUnsignedWrap(), /*HasNSW=*/false}); - GEPNoWrapFlags NWFlags = AddRecPtr->hasNoUnsignedWrap() - ? GEPNoWrapFlags::noUnsignedWrap() - : GEPNoWrapFlags::none(); - VPValue *BasePtr = Builder.createNoWrapPtrAdd(StartVPV, Offset, NWFlags); - - // Create a new vector pointer for strided access. - VPValue *NewPtr = Builder.createVectorPointer( - BasePtr, Type::getInt8Ty(Plan.getContext()), StrideInBytes, NWFlags, - LoadR->getDebugLoc()); - - VPValue *Mask = LoadR->getMask(); - if (!Mask) - Mask = Plan.getTrue(); - auto *StridedLoad = Builder.createWidenMemIntrinsic( - Intrinsic::experimental_vp_strided_load, - {NewPtr, StrideInBytes, Mask, I32VF}, LoadTy, Alignment, *LoadR, - LoadR->getDebugLoc()); - LoadR->replaceAllUsesWith(StridedLoad); - } - } -} diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h index 4b59d37150ff4..5b7178e226bef 100644 --- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h +++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h @@ -398,6 +398,26 @@ struct VPlanTransforms { static void convertToAbstractRecipes(VPlan &Plan, VPCostContext &Ctx, VFRange &Range); + /// Legalize VPWidenPointerInductionRecipe, by replacing it with a PtrAdd + /// (IndStart, ScalarIVSteps (0, Step)) if only its scalar values are used, as + /// VPWidenPointerInductionRecipe will generate vectors only. If some users + /// require vectors while other require scalars, the scalar uses need to + /// extract the scalars from the generated vectors (Note that this is + /// different to how int/fp inductions are handled). Legalize + /// extract-from-ends using uniform VPReplicateRecipe of wide inductions to + /// use regular VPReplicateRecipe, so the correct end value is available. Also + /// optimize VPWidenIntOrFpInductionRecipe, if any of its users needs scalar + /// values, by providing them scalar steps built on the canonical scalar IV + /// and update the original IV's users. This is an optional optimization to + /// reduce the needs of vector extracts. + static void legalizeAndOptimizeInductions(VPlan &Plan); + + /// Try to narrow wide and replicating recipes to single scalar recipes for + /// loops with the same VF, when the values are known to be uniform. Also + /// narrows masked div/rem intrinsics with a safe divisor to unmasked scalar + /// operations. + static void narrowToSingleScalarRecipes(VPlan &Plan); + /// Perform instcombine-like simplifications on recipes in \p Plan. static void simplifyRecipes(VPlan &Plan); diff --git a/llvm/lib/Transforms/Vectorize/VPlanWideningDecisions.cpp b/llvm/lib/Transforms/Vectorize/VPlanWideningDecisions.cpp new file mode 100644 index 0000000000000..5605fd4e9b0f6 --- /dev/null +++ b/llvm/lib/Transforms/Vectorize/VPlanWideningDecisions.cpp @@ -0,0 +1,2333 @@ +//===- VPlanWideningDecisions.cpp - VPlan-based widening decisions --------===// +// +// 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 +// +//===----------------------------------------------------------------------===// +/// +/// \file +/// This file implements VPlan-based widening decisions, which convert +/// initial recipes into widened, scalarized, interleaved or otherwise +/// specialized recipes. +/// +//===----------------------------------------------------------------------===// + +#include "VPRecipeBuilder.h" +#include "VPlan.h" +#include "VPlanAnalysis.h" +#include "VPlanCFG.h" +#include "VPlanDominatorTree.h" +#include "VPlanHelpers.h" +#include "VPlanPatternMatch.h" +#include "VPlanTransforms.h" +#include "VPlanUtils.h" +#include "VPlanVerifier.h" +#include "llvm/ADT/APInt.h" +#include "llvm/ADT/PostOrderIterator.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SetVector.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Analysis/IVDescriptors.h" +#include "llvm/Analysis/InstSimplifyFolder.h" +#include "llvm/Analysis/Loads.h" +#include "llvm/Analysis/LoopAccessAnalysis.h" +#include "llvm/Analysis/LoopInfo.h" +#include "llvm/Analysis/MemoryLocation.h" +#include "llvm/Analysis/ScalarEvolutionPatternMatch.h" +#include "llvm/Analysis/ScopedNoAliasAA.h" +#include "llvm/Analysis/VectorUtils.h" +#include "llvm/IR/Intrinsics.h" +#include "llvm/IR/MDBuilder.h" +#include "llvm/IR/Metadata.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/TypeSize.h" +#include "llvm/Transforms/Utils/LoopUtils.h" +#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h" + +using namespace llvm; +using namespace VPlanPatternMatch; +using namespace SCEVPatternMatch; + +/// If the pointer operand \p Addr of a memory access is an affine AddRec +/// w.r.t. \p L with a constant stride, return the stride in units of +/// \p AccessTy. Otherwise return std::nullopt. +static std::optional getConstantStride(VPValue *Addr, Type *AccessTy, + PredicatedScalarEvolution &PSE, + const Loop *L) { + const SCEV *AddrSCEV = vputils::getSCEVExprForVPValue(Addr, PSE, L); + auto *AddRec = dyn_cast(AddrSCEV); + if (!AddRec) + return {}; + + return getStrideFromAddRec(AddRec, L, AccessTy, /*Ptr=*/nullptr, PSE); +} + +bool VPlanTransforms::tryToConvertVPInstructionsToVPRecipes( + VPlan &Plan, const TargetLibraryInfo &TLI, PredicatedScalarEvolution &PSE, + Loop *OuterLoop) { + + ReversePostOrderTraversal> RPOT( + Plan.getVectorLoopRegion()); + for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly(RPOT)) { + // Skip blocks outside region + if (!VPBB->getParent()) + break; + VPRecipeBase *Term = VPBB->getTerminator(); + auto EndIter = Term ? Term->getIterator() : VPBB->end(); + // Introduce each ingredient into VPlan. + for (VPRecipeBase &Ingredient : + make_early_inc_range(make_range(VPBB->begin(), EndIter))) { + + VPValue *VPV = Ingredient.getVPSingleValue(); + if (!VPV->getUnderlyingValue()) + continue; + + Instruction *Inst = cast(VPV->getUnderlyingValue()); + + // Atomic accesses and fences have ordering/atomicity semantics that + // cannot be preserved by lane-wise widening. + if (isa(Inst)) + return false; + + VPRecipeBase *NewRecipe = nullptr; + if (auto *PhiR = dyn_cast(&Ingredient)) { + auto *Phi = cast(PhiR->getUnderlyingValue()); + NewRecipe = new VPWidenPHIRecipe(PhiR->operands(), PhiR->getDebugLoc(), + Phi->getName()); + } else if (auto *VPI = dyn_cast(&Ingredient)) { + assert(!isa(Inst) && "phis should be handled above"); + // Create VPWidenMemoryRecipe for loads and stores. + if (LoadInst *Load = dyn_cast(Inst)) { + bool IsConsecutive = + getConstantStride(VPI->getOperand(0), VPI->getScalarType(), PSE, + OuterLoop) == 1; + NewRecipe = new VPWidenLoadRecipe(*Load, Ingredient.getOperand(0), + nullptr /*Mask*/, IsConsecutive, + *VPI, Ingredient.getDebugLoc()); + } else if (StoreInst *Store = dyn_cast(Inst)) { + bool IsConsecutive = + getConstantStride(VPI->getOperand(1), + VPI->getOperand(0)->getScalarType(), PSE, + OuterLoop) == 1; + NewRecipe = new VPWidenStoreRecipe( + *Store, Ingredient.getOperand(1), Ingredient.getOperand(0), + nullptr /*Mask*/, IsConsecutive, *VPI, Ingredient.getDebugLoc()); + } else if (GetElementPtrInst *GEP = dyn_cast(Inst)) { + NewRecipe = new VPWidenGEPRecipe(GEP->getSourceElementType(), + Ingredient.operands(), *VPI, + Ingredient.getDebugLoc(), GEP); + } else if (CallInst *CI = dyn_cast(Inst)) { + Intrinsic::ID VectorID = getVectorIntrinsicIDForCall(CI, &TLI); + if (VectorID == Intrinsic::not_intrinsic) + return false; + + // The noalias.scope.decl intrinsic declares a noalias scope that + // is valid for a single iteration. Emitting it as a single-scalar + // replicate would incorrectly extend the scope across multiple + // original iterations packed into one vector iteration. + // FIXME: If we want to vectorize this loop, then we have to drop + // all the associated !alias.scope and !noalias. + if (VectorID == Intrinsic::experimental_noalias_scope_decl) + return false; + + // These intrinsics are recognized by getVectorIntrinsicIDForCall + // but are not widenable. Emit them as replicate instead of widening. + if (VectorID == Intrinsic::assume || + VectorID == Intrinsic::lifetime_end || + VectorID == Intrinsic::lifetime_start || + VectorID == Intrinsic::sideeffect || + VectorID == Intrinsic::pseudoprobe) { + // If the operand of llvm.assume holds before vectorization, it will + // also hold per lane. + // llvm.pseudoprobe requires to be duplicated per lane for accurate + // sample count. + const bool IsSingleScalar = VectorID != Intrinsic::assume && + VectorID != Intrinsic::pseudoprobe; + NewRecipe = new VPReplicateRecipe(CI, Ingredient.operands(), + /*IsSingleScalar=*/IsSingleScalar, + /*Mask=*/nullptr, *VPI, *VPI, + Ingredient.getDebugLoc()); + } else { + NewRecipe = new VPWidenIntrinsicRecipe( + *CI, VectorID, drop_end(Ingredient.operands()), CI->getType(), + VPIRFlags(*CI), *VPI, CI->getDebugLoc()); + } + } else if (auto *CI = dyn_cast(Inst)) { + NewRecipe = new VPWidenCastRecipe( + CI->getOpcode(), Ingredient.getOperand(0), CI->getType(), CI, + VPIRFlags(*CI), VPIRMetadata(*CI)); + } else { + NewRecipe = new VPWidenRecipe(*Inst, Ingredient.operands(), *VPI, + *VPI, Ingredient.getDebugLoc()); + } + } else { + assert(isa(&Ingredient) && + "inductions must be created earlier"); + continue; + } + + NewRecipe->insertBefore(&Ingredient); + if (NewRecipe->getNumDefinedValues() == 1) + VPV->replaceAllUsesWith(NewRecipe->getVPSingleValue()); + else + assert(NewRecipe->getNumDefinedValues() == 0 && + "Only recpies with zero or one defined values expected"); + Ingredient.eraseFromParent(); + } + } + return true; +} + +/// This function tries convert extended in-loop reductions to +/// VPExpressionRecipe and clamp the \p Range if it is beneficial and +/// valid. The created recipe must be decomposed to its constituent +/// recipes before execution. +static VPExpressionRecipe * +tryToMatchAndCreateExtendedReduction(VPReductionRecipe *Red, VPCostContext &Ctx, + VFRange &Range) { + Type *RedTy = Red->getScalarType(); + VPValue *VecOp = Red->getVecOp(); + + assert(!Red->isPartialReduction() && + "This path does not support partial reductions"); + + // Clamp the range if using extended-reduction is profitable. + auto IsExtendedRedValidAndClampRange = + [&](unsigned Opcode, Instruction::CastOps ExtOpc, Type *SrcTy) -> bool { + return LoopVectorizationPlanner::getDecisionAndClampRange( + [&](ElementCount VF) { + auto *SrcVecTy = cast(toVectorTy(SrcTy, VF)); + TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput; + + InstructionCost ExtRedCost = InstructionCost::getInvalid(); + InstructionCost ExtCost = + cast(VecOp)->computeCost(VF, Ctx); + InstructionCost RedCost = Red->computeCost(VF, Ctx); + + assert(!RedTy->isFloatingPointTy() && + "getExtendedReductionCost only supports integer types"); + ExtRedCost = Ctx.TTI.getExtendedReductionCost( + Opcode, ExtOpc == Instruction::CastOps::ZExt, RedTy, SrcVecTy, + Red->getFastMathFlagsOrNone(), CostKind); + return ExtRedCost.isValid() && ExtRedCost < ExtCost + RedCost; + }, + Range); + }; + + VPValue *A; + // Match reduce(ext)). + if (match(VecOp, m_Isa(m_ZExtOrSExt(m_VPValue(A)))) && + IsExtendedRedValidAndClampRange( + RecurrenceDescriptor::getOpcode(Red->getRecurrenceKind()), + cast(VecOp)->getOpcode(), A->getScalarType())) + return new VPExpressionRecipe(cast(VecOp), Red); + + return nullptr; +} + +/// This function tries convert extended in-loop reductions to +/// VPExpressionRecipe and clamp the \p Range if it is beneficial +/// and valid. The created VPExpressionRecipe must be decomposed to its +/// constituent recipes before execution. Patterns of the +/// VPExpressionRecipe: +/// reduce.add(mul(...)), +/// reduce.add(mul(ext(A), ext(B))), +/// reduce.add(ext(mul(ext(A), ext(B)))). +/// reduce.fadd(fmul(ext(A), ext(B))) +static VPExpressionRecipe * +tryToMatchAndCreateMulAccumulateReduction(VPReductionRecipe *Red, + VPCostContext &Ctx, VFRange &Range) { + unsigned Opcode = RecurrenceDescriptor::getOpcode(Red->getRecurrenceKind()); + if (Opcode != Instruction::Add && Opcode != Instruction::Sub && + Opcode != Instruction::FAdd) + return nullptr; + + assert(!Red->isPartialReduction() && + "This path does not support partial reductions"); + Type *RedTy = Red->getScalarType(); + + // Clamp the range if using multiply-accumulate-reduction is profitable. + auto IsMulAccValidAndClampRange = + [&](VPWidenRecipe *Mul, VPWidenCastRecipe *Ext0, VPWidenCastRecipe *Ext1, + VPWidenCastRecipe *OuterExt) -> bool { + return LoopVectorizationPlanner::getDecisionAndClampRange( + [&](ElementCount VF) { + TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput; + Type *SrcTy = Ext0 ? Ext0->getOperand(0)->getScalarType() : RedTy; + InstructionCost MulAccCost; + + // getMulAccReductionCost for in-loop reductions does not support + // mixed or floating-point extends. + if (Ext0 && Ext1 && + (Ext0->getOpcode() != Ext1->getOpcode() || + Ext0->getOpcode() == Instruction::CastOps::FPExt)) + return false; + + bool IsZExt = + !Ext0 || Ext0->getOpcode() == Instruction::CastOps::ZExt; + auto *SrcVecTy = cast(toVectorTy(SrcTy, VF)); + MulAccCost = Ctx.TTI.getMulAccReductionCost(IsZExt, Opcode, RedTy, + SrcVecTy, CostKind); + + InstructionCost MulCost = Mul->computeCost(VF, Ctx); + InstructionCost RedCost = Red->computeCost(VF, Ctx); + InstructionCost ExtCost = 0; + if (Ext0) + ExtCost += Ext0->computeCost(VF, Ctx); + if (Ext1) + ExtCost += Ext1->computeCost(VF, Ctx); + if (OuterExt) + ExtCost += OuterExt->computeCost(VF, Ctx); + + return MulAccCost.isValid() && + MulAccCost < ExtCost + MulCost + RedCost; + }, + Range); + }; + + VPValue *VecOp = Red->getVecOp(); + VPRecipeBase *Sub = nullptr; + VPValue *A, *B; + VPValue *Tmp = nullptr; + + if (RedTy->isFloatingPointTy()) + return nullptr; + + // Sub reductions could have a sub between the add reduction and vec op. + if (match(VecOp, m_Sub(m_ZeroInt(), m_VPValue(Tmp)))) { + Sub = VecOp->getDefiningRecipe(); + VecOp = Tmp; + } + + // If ValB is a constant and can be safely extended, truncate it to the same + // type as ExtA's operand, then extend it to the same type as ExtA. This + // creates two uniform extends that can more easily be matched by the rest of + // the bundling code. The ExtB reference, ValB and operand 1 of Mul are all + // replaced with the new extend of the constant. + auto ExtendAndReplaceConstantOp = [](VPWidenCastRecipe *ExtA, + VPWidenCastRecipe *&ExtB, VPValue *&ValB, + VPWidenRecipe *Mul) { + if (!ExtA || ExtB || !isa(ValB)) + return; + Type *NarrowTy = ExtA->getOperand(0)->getScalarType(); + Instruction::CastOps ExtOpc = ExtA->getOpcode(); + const APInt *Const; + if (!match(ValB, m_APInt(Const)) || + !llvm::canConstantBeExtended( + Const, NarrowTy, TTI::getPartialReductionExtendKind(ExtOpc))) + return; + // The truncate ensures that the type of each extended operand is the + // same, and it's been proven that the constant can be extended from + // NarrowTy safely. Necessary since ExtA's extended operand would be + // e.g. an i8, while the const will likely be an i32. This will be + // elided by later optimisations. + VPBuilder Builder(Mul); + auto *Trunc = + Builder.createWidenCast(Instruction::CastOps::Trunc, ValB, NarrowTy); + Type *WideTy = ExtA->getScalarType(); + ValB = ExtB = Builder.createWidenCast(ExtOpc, Trunc, WideTy); + Mul->setOperand(1, ExtB); + }; + + // Try to match reduce.add(mul(...)). + if (match(VecOp, m_Mul(m_VPValue(A), m_VPValue(B)))) { + auto *RecipeA = dyn_cast(A); + auto *RecipeB = dyn_cast(B); + auto *Mul = cast(VecOp); + + // Convert reduce.add(mul(ext, const)) to reduce.add(mul(ext, ext(const))) + ExtendAndReplaceConstantOp(RecipeA, RecipeB, B, Mul); + + // Match reduce.add/sub(mul(ext, ext)). + if (RecipeA && RecipeB && match(RecipeA, m_ZExtOrSExt(m_VPValue())) && + match(RecipeB, m_ZExtOrSExt(m_VPValue())) && + IsMulAccValidAndClampRange(Mul, RecipeA, RecipeB, nullptr)) { + if (Sub) + return new VPExpressionRecipe(RecipeA, RecipeB, Mul, + cast(Sub), Red); + return new VPExpressionRecipe(RecipeA, RecipeB, Mul, Red); + } + // TODO: Add an expression type for this variant with a negated mul + if (!Sub && IsMulAccValidAndClampRange(Mul, nullptr, nullptr, nullptr)) + return new VPExpressionRecipe(Mul, Red); + } + // TODO: Add an expression type for negated versions of other expression + // variants. + if (Sub) + return nullptr; + + // Match reduce.add(ext(mul(A, B))). + if (match(VecOp, m_ZExtOrSExt(m_Mul(m_VPValue(A), m_VPValue(B))))) { + auto *Ext = cast(VecOp); + auto *Mul = cast(Ext->getOperand(0)); + auto *Ext0 = dyn_cast(A); + auto *Ext1 = dyn_cast(B); + + // reduce.add(ext(mul(ext, const))) + // -> reduce.add(ext(mul(ext, ext(const)))) + ExtendAndReplaceConstantOp(Ext0, Ext1, B, Mul); + + // reduce.add(ext(mul(ext(A), ext(B)))) + // -> reduce.add(mul(wider_ext(A), wider_ext(B))) + // The inner extends must either have the same opcode as the outer extend or + // be the same, in which case the multiply can never result in a negative + // value and the outer extend can be folded away by doing wider + // extends for the operands of the mul. + if (Ext0 && Ext1 && + (Ext->getOpcode() == Ext0->getOpcode() || Ext0 == Ext1) && + Ext0->getOpcode() == Ext1->getOpcode() && + IsMulAccValidAndClampRange(Mul, Ext0, Ext1, Ext) && Mul->hasOneUse()) { + auto *NewExt0 = new VPWidenCastRecipe( + Ext0->getOpcode(), Ext0->getOperand(0), Ext->getScalarType(), nullptr, + *Ext0, *Ext0, Ext0->getDebugLoc()); + NewExt0->insertBefore(Ext0); + + VPWidenCastRecipe *NewExt1 = NewExt0; + if (Ext0 != Ext1) { + NewExt1 = new VPWidenCastRecipe(Ext1->getOpcode(), Ext1->getOperand(0), + Ext->getScalarType(), nullptr, *Ext1, + *Ext1, Ext1->getDebugLoc()); + NewExt1->insertBefore(Ext1); + } + auto *NewMul = Mul->cloneWithOperands({NewExt0, NewExt1}); + NewMul->insertBefore(Mul); + Ext->replaceAllUsesWith(NewMul); + Ext->eraseFromParent(); + Mul->eraseFromParent(); + return new VPExpressionRecipe(NewExt0, NewExt1, NewMul, Red); + } + } + return nullptr; +} + +/// This function tries to create abstract recipes from the reduction recipe for +/// following optimizations and cost estimation. +static void tryToCreateAbstractReductionRecipe(VPReductionRecipe *Red, + VPCostContext &Ctx, + VFRange &Range) { + // Creation of VPExpressions for partial reductions is entirely handled in + // transformToPartialReduction. + assert(!Red->isPartialReduction() && + "This path does not support partial reductions"); + + VPExpressionRecipe *AbstractR = nullptr; + auto IP = std::next(Red->getIterator()); + auto *VPBB = Red->getParent(); + if (auto *MulAcc = tryToMatchAndCreateMulAccumulateReduction(Red, Ctx, Range)) + AbstractR = MulAcc; + else if (auto *ExtRed = tryToMatchAndCreateExtendedReduction(Red, Ctx, Range)) + AbstractR = ExtRed; + // Cannot create abstract inloop reduction recipes. + if (!AbstractR) + return; + + AbstractR->insertBefore(*VPBB, IP); + Red->replaceAllUsesWith(AbstractR); +} + +void VPlanTransforms::convertToAbstractRecipes(VPlan &Plan, VPCostContext &Ctx, + VFRange &Range) { + for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly( + vp_depth_first_deep(Plan.getVectorLoopRegion()))) { + for (VPRecipeBase &R : make_early_inc_range(*VPBB)) { + if (auto *Red = dyn_cast(&R)) + tryToCreateAbstractReductionRecipe(Red, Ctx, Range); + } + } +} + +namespace { + +using ExtendKind = TTI::PartialReductionExtendKind; +struct ReductionExtend { + Type *SrcType = nullptr; + ExtendKind Kind = ExtendKind::PR_None; +}; + +/// Describes the extends used to compute the extended reduction operand. +/// ExtendB is optional. If ExtendB is present, ExtendsUser is a binary +/// operation. +struct ExtendedReductionOperand { + /// The recipe that consumes the extends. + VPWidenRecipe *ExtendsUser = nullptr; + /// Extend descriptions (inputs to getPartialReductionCost). + ReductionExtend ExtendA, ExtendB; +}; + +/// A chain of recipes that form a partial reduction. Matches either +/// reduction_bin_op (extended op, accumulator), or +/// reduction_bin_op (accumulator, extended op). +/// The possible forms of the "extended op" are listed in +/// matchExtendedReductionOperand. +struct VPPartialReductionChain { + /// The top-level binary operation that forms the reduction to a scalar + /// after the loop body. + VPWidenRecipe *ReductionBinOp = nullptr; + /// The user of the extends that is then reduced. + ExtendedReductionOperand ExtendedOp; + /// The recurrence kind for the entire partial reduction chain. + /// This allows distinguishing between Sub and AddWithSub recurrences, + /// when the ReductionBinOp is a Instruction::Sub. + RecurKind RK; + /// The index of the accumulator operand of ReductionBinOp. The extended op + /// is `1 - AccumulatorOpIdx`. + unsigned AccumulatorOpIdx; + unsigned ScaleFactor; + /// Optional blend to represent predication for the block that updates the + /// reduction. + VPBlendRecipe *Blend = nullptr; +}; + +// Return the incoming index of the single-use value in the blend, which is +// expected to be the predicated reduction update. +static std::optional +getBlendReductionUpdateValueIdx(VPBlendRecipe *Blend) { + assert(Blend && !Blend->isNormalized() && + Blend->getNumIncomingValues() == 2 && + "Expected a non-normalized blend with two incoming values"); + bool FirstIncomingHasOneUse = Blend->getIncomingValue(0)->hasOneUse(); + + // Only the update value should have one use (the blend). The previous + // value should always have at least two uses, the blend and the reduction. + if (FirstIncomingHasOneUse == Blend->getIncomingValue(1)->hasOneUse()) + return std::nullopt; + return FirstIncomingHasOneUse ? 0 : 1; +} + +static VPSingleDefRecipe * +optimizeExtendsForPartialReduction(VPSingleDefRecipe *Op) { + // reduce.add(mul(ext(A), C)) + // -> reduce.add(mul(ext(A), ext(trunc(C)))) + const APInt *Const; + if (match(Op, m_Mul(m_ZExtOrSExt(m_VPValue()), m_APInt(Const)))) { + auto *ExtA = cast(Op->getOperand(0)); + Instruction::CastOps ExtOpc = ExtA->getOpcode(); + Type *NarrowTy = ExtA->getOperand(0)->getScalarType(); + if (!Op->hasOneUse() || + !llvm::canConstantBeExtended( + Const, NarrowTy, TTI::getPartialReductionExtendKind(ExtOpc))) + return Op; + + VPBuilder Builder(Op); + auto *Trunc = Builder.createWidenCast(Instruction::CastOps::Trunc, + Op->getOperand(1), NarrowTy); + Type *WideTy = ExtA->getScalarType(); + Op->setOperand(1, Builder.createWidenCast(ExtOpc, Trunc, WideTy)); + return Op; + } + + // reduce.add(abs(sub(ext(A), ext(B)))) + // -> reduce.add(ext(absolute-difference(A, B))) + VPValue *X, *Y; + if (match(Op, m_WidenIntrinsic(m_Sub( + m_ZExtOrSExt(m_VPValue(X)), m_ZExtOrSExt(m_VPValue(Y)))))) { + auto *Sub = Op->getOperand(0)->getDefiningRecipe(); + auto *Ext = cast(Sub->getOperand(0)); + assert(Ext->getOpcode() == + cast(Sub->getOperand(1))->getOpcode() && + "Expected both the LHS and RHS extends to be the same"); + bool IsSigned = Ext->getOpcode() == Instruction::SExt; + VPBuilder Builder(Op); + Type *SrcTy = X->getScalarType(); + auto *FreezeX = Builder.insert(new VPWidenRecipe(Instruction::Freeze, {X})); + auto *FreezeY = Builder.insert(new VPWidenRecipe(Instruction::Freeze, {Y})); + auto *Max = Builder.insert( + new VPWidenIntrinsicRecipe(IsSigned ? Intrinsic::smax : Intrinsic::umax, + {FreezeX, FreezeY}, SrcTy)); + auto *Min = Builder.insert( + new VPWidenIntrinsicRecipe(IsSigned ? Intrinsic::smin : Intrinsic::umin, + {FreezeX, FreezeY}, SrcTy)); + auto *AbsDiff = + Builder.insert(new VPWidenRecipe(Instruction::Sub, {Max, Min})); + return Builder.createWidenCast(Instruction::CastOps::ZExt, AbsDiff, + Op->getScalarType()); + } + + // reduce.add(ext(mul(ext(A), ext(B)))) + // -> reduce.add(mul(wider_ext(A), wider_ext(B))) + // TODO: Support this optimization for float types. + if (match(Op, m_ZExtOrSExt(m_Mul(m_ZExtOrSExt(m_VPValue()), + m_ZExtOrSExt(m_VPValue()))))) { + auto *Ext = cast(Op); + auto *Mul = cast(Ext->getOperand(0)); + auto *MulLHS = cast(Mul->getOperand(0)); + auto *MulRHS = cast(Mul->getOperand(1)); + if (!Mul->hasOneUse() || + (Ext->getOpcode() != MulLHS->getOpcode() && MulLHS != MulRHS) || + MulLHS->getOpcode() != MulRHS->getOpcode()) + return Op; + VPBuilder Builder(Mul); + auto *NewLHS = Builder.createWidenCast( + MulLHS->getOpcode(), MulLHS->getOperand(0), Ext->getScalarType()); + auto *NewRHS = MulLHS == MulRHS + ? NewLHS + : Builder.createWidenCast(MulRHS->getOpcode(), + MulRHS->getOperand(0), + Ext->getScalarType()); + auto *NewMul = Mul->cloneWithOperands({NewLHS, NewRHS}); + Builder.insert(NewMul); + Op->replaceAllUsesWith(NewMul); + Op->eraseFromParent(); + Mul->eraseFromParent(); + return NewMul; + } + + return Op; +} + +static VPExpressionRecipe * +createPartialReductionExpression(VPReductionRecipe *Red) { + VPValue *VecOp = Red->getVecOp(); + + // reduce.[f]add(ext(op)) + // -> VPExpressionRecipe(op, red) + if (match(VecOp, m_WidenAnyExtend(m_VPValue()))) + return new VPExpressionRecipe(cast(VecOp), Red); + + // reduce.[f]add(neg(ext(op))) + // -> VPExpressionRecipe(op, sub/neg, red) + if (match(VecOp, m_AnyNeg(m_WidenAnyExtend(m_VPValue())))) { + auto *Neg = cast(VecOp); + auto *Ext = + cast(Neg->getOperand(Neg->getNumOperands() - 1)); + return new VPExpressionRecipe(Ext, Neg, Red); + } + + // reduce.[f]add([f]mul(ext(a), ext(b))) + // -> VPExpressionRecipe(a, b, mul, red) + if (match(VecOp, m_FMul(m_FPExt(m_VPValue()), m_FPExt(m_VPValue()))) || + match(VecOp, + m_Mul(m_ZExtOrSExt(m_VPValue()), m_ZExtOrSExt(m_VPValue())))) { + auto *Mul = cast(VecOp); + auto *ExtA = cast(Mul->getOperand(0)); + auto *ExtB = cast(Mul->getOperand(1)); + return new VPExpressionRecipe(ExtA, ExtB, Mul, Red); + } + + // reduce.fadd(fneg(fmul(fpext(a), fpext(b)))) + // -> VPExpressionRecipe(a, b, fmul, fsub, red) + if (match(VecOp, + m_FNeg(m_FMul(m_FPExt(m_VPValue()), m_FPExt(m_VPValue()))))) { + auto *FNeg = cast(VecOp); + auto *FMul = cast(FNeg->getOperand(0)); + auto *ExtA = cast(FMul->getOperand(0)); + auto *ExtB = cast(FMul->getOperand(1)); + return new VPExpressionRecipe(ExtA, ExtB, FMul, FNeg, Red); + } + + // reduce.add(neg(mul(ext(a), ext(b)))) + // -> VPExpressionRecipe(a, b, mul, sub, red) + if (match(VecOp, m_Sub(m_ZeroInt(), m_Mul(m_ZExtOrSExt(m_VPValue()), + m_ZExtOrSExt(m_VPValue()))))) { + auto *Sub = cast(VecOp); + auto *Mul = cast(Sub->getOperand(1)); + auto *ExtA = cast(Mul->getOperand(0)); + auto *ExtB = cast(Mul->getOperand(1)); + return new VPExpressionRecipe(ExtA, ExtB, Mul, Sub, Red); + } + + llvm_unreachable("Unsupported expression"); +} + +// Helper to transform a partial reduction chain into a partial reduction +// recipe. Assumes profitability has been checked. +static void transformToPartialReduction(const VPPartialReductionChain &Chain, + VPlan &Plan, + VPReductionPHIRecipe *RdxPhi) { + VPWidenRecipe *WidenRecipe = Chain.ReductionBinOp; + assert(WidenRecipe->getNumOperands() == 2 && "Expected binary operation"); + + VPValue *Accumulator = WidenRecipe->getOperand(Chain.AccumulatorOpIdx); + auto *ExtendedOp = cast( + WidenRecipe->getOperand(1 - Chain.AccumulatorOpIdx)); + + // FIXME: Do these transforms before invoking the cost-model. + ExtendedOp = optimizeExtendsForPartialReduction(ExtendedOp); + + // Sub-reductions can be implemented in two ways: + // (1) negate the operand in the vector loop (the default way). + // (2) subtract the reduced value from the init value in the middle block. + // Both ways keep the reduction itself as an 'add' reduction. + // + // The ISD nodes for partial reductions don't support folding the + // sub/negation into its operands because the following is not a valid + // transformation: + // sub(0, mul(ext(a), ext(b))) + // -> mul(ext(a), ext(sub(0, b))) + // + // It's therefore better to choose option (2) such that the partial + // reduction is always positive (starting at '0') and to do a final + // subtract in the middle block. + if ((WidenRecipe->getOpcode() == Instruction::Sub && + Chain.RK != RecurKind::Sub) || + (WidenRecipe->getOpcode() == Instruction::FSub && + Chain.RK != RecurKind::FSub)) { + VPBuilder Builder(WidenRecipe); + Type *ElemTy = ExtendedOp->getScalarType(); + VPWidenRecipe *NegRecipe; + if (WidenRecipe->getOpcode() == Instruction::FSub) { + NegRecipe = + new VPWidenRecipe(Instruction::FNeg, {ExtendedOp}, VPIRFlags(), + VPIRMetadata(), DebugLoc::getUnknown()); + } else { + auto *Zero = Plan.getZero(ElemTy); + NegRecipe = + new VPWidenRecipe(Instruction::Sub, {Zero, ExtendedOp}, VPIRFlags(), + VPIRMetadata(), DebugLoc::getUnknown()); + } + Builder.insert(NegRecipe); + ExtendedOp = NegRecipe; + } + + // Check if WidenRecipe is the final result of the reduction. If so, look + // through the Select recipe introduced by tail-folding, otherwise look + // through any Blend recipe introduced by predication for the block. + VPValue *ExitSearch = + Chain.Blend ? cast(Chain.Blend) : cast(WidenRecipe); + + VPValue *Cond = nullptr; + VPValue *ExitValue = cast_or_null( + findUserOf(ExitSearch, m_Select(m_VPValue(Cond), m_Specific(ExitSearch), + m_Specific(RdxPhi)))); + + if (Chain.Blend) { + std::optional BlendReductionIdx = + getBlendReductionUpdateValueIdx(Chain.Blend); + assert(BlendReductionIdx && + Chain.Blend->getIncomingValue(*BlendReductionIdx) == WidenRecipe && + "Expected blend to contain the reduction update"); + VPValue *BlendCond = Chain.Blend->getMask(*BlendReductionIdx); + Cond = ExitValue ? VPBuilder(WidenRecipe) + .createLogicalAnd(Cond, BlendCond, + WidenRecipe->getDebugLoc()) + : BlendCond; + } + + bool IsLastInChain = RdxPhi->getBackedgeValue() == WidenRecipe || + RdxPhi->getBackedgeValue() == ExitValue || + RdxPhi->getBackedgeValue() == Chain.Blend; + assert((!ExitValue || IsLastInChain) && + "if we found ExitValue, it must match RdxPhi's backedge value"); + + Type *PhiType = RdxPhi->getScalarType(); + RecurKind RdxKind = + PhiType->isFloatingPointTy() ? RecurKind::FAdd : RecurKind::Add; + auto *PartialRed = new VPReductionRecipe( + RdxKind, + RdxKind == RecurKind::FAdd ? WidenRecipe->getFastMathFlagsOrNone() + : FastMathFlags(), + WidenRecipe->getUnderlyingInstr(), Accumulator, ExtendedOp, Cond, + RdxUnordered{/*VFScaleFactor=*/Chain.ScaleFactor}); + PartialRed->insertBefore(WidenRecipe); + + if (ExitValue) + ExitValue->replaceAllUsesWith(PartialRed); + if (Chain.Blend) + Chain.Blend->replaceAllUsesWith(PartialRed); + WidenRecipe->replaceAllUsesWith(PartialRed); + + // For cost-model purposes, fold this into a VPExpression. + VPExpressionRecipe *E = createPartialReductionExpression(PartialRed); + E->insertBefore(WidenRecipe); + PartialRed->replaceAllUsesWith(E); + + // We only need to update the PHI node once, which is when we find the + // last reduction in the chain. + if (!IsLastInChain) + return; + + // Scale the PHI and ReductionStartVector by the VFScaleFactor + assert(RdxPhi->getVFScaleFactor() == 1 && "scale factor must not be set"); + RdxPhi->setVFScaleFactor(Chain.ScaleFactor); + + auto *StartInst = cast(RdxPhi->getStartValue()); + assert(StartInst->getOpcode() == VPInstruction::ReductionStartVector); + auto *NewScaleFactor = Plan.getConstantInt(32, Chain.ScaleFactor); + StartInst->setOperand(2, NewScaleFactor); + + // If this is the last value in a sub-reduction chain, then update the PHI + // node to start at `0` and update the reduction-result to subtract from + // the PHI's start value. + if (Chain.RK != RecurKind::Sub && Chain.RK != RecurKind::FSub) + return; + + VPValue *OldStartValue = StartInst->getOperand(0); + StartInst->setOperand(0, StartInst->getOperand(1)); + + // Replace reduction_result by 'sub (startval, reductionresult)'. + VPInstruction *RdxResult = vputils::findComputeReductionResult(RdxPhi); + assert(RdxResult && "Could not find reduction result"); + + VPBuilder Builder = VPBuilder::getToInsertAfter(RdxResult); + unsigned SubOpc = Chain.RK == RecurKind::FSub ? Instruction::BinaryOps::FSub + : Instruction::BinaryOps::Sub; + VPInstruction *NewResult = Builder.createNaryOp( + SubOpc, {OldStartValue, RdxResult}, VPIRFlags::getDefaultFlags(SubOpc), + RdxPhi->getDebugLoc()); + RdxResult->replaceUsesWithIf( + NewResult, + [&NewResult](VPUser &U, unsigned Idx) { return &U != NewResult; }); +} + +/// Returns the cost of a link in a partial-reduction chain for a given VF. +static InstructionCost +getPartialReductionLinkCost(VPCostContext &CostCtx, + const VPPartialReductionChain &Link, + ElementCount VF) { + Type *RdxType = Link.ReductionBinOp->getScalarType(); + const ExtendedReductionOperand &ExtendedOp = Link.ExtendedOp; + std::optional BinOpc = std::nullopt; + // If ExtendB is not none, then the "ExtendsUser" is the binary operation. + if (ExtendedOp.ExtendB.Kind != ExtendKind::PR_None) + BinOpc = ExtendedOp.ExtendsUser->getOpcode(); + + std::optional Flags; + if (RdxType->isFloatingPointTy()) + Flags = Link.ReductionBinOp->getFastMathFlagsOrNone(); + + auto GetLinkOpcode = [&Link]() -> unsigned { + switch (Link.RK) { + case RecurKind::Sub: + return Instruction::Add; + case RecurKind::FSub: + return Instruction::FAdd; + default: + return Link.ReductionBinOp->getOpcode(); + } + }; + + return CostCtx.TTI.getPartialReductionCost( + GetLinkOpcode(), ExtendedOp.ExtendA.SrcType, ExtendedOp.ExtendB.SrcType, + RdxType, VF, ExtendedOp.ExtendA.Kind, ExtendedOp.ExtendB.Kind, BinOpc, + CostCtx.CostKind, Flags); +} + +static ExtendKind getPartialReductionExtendKind(VPWidenCastRecipe *Cast) { + return TTI::getPartialReductionExtendKind(Cast->getOpcode()); +} + +/// Checks if \p Op (which is an operand of \p UpdateR) is an extended reduction +/// operand. This is an operand where the source of the value (e.g. a load) has +/// been extended (sext, zext, or fpext) before it is used in the reduction. +/// +/// Possible forms matched by this function: +/// - UpdateR(PrevValue, ext(...)) +/// - UpdateR(PrevValue, mul(ext(...), ext(...))) +/// - UpdateR(PrevValue, mul(ext(...), Constant)) +/// - UpdateR(PrevValue, ext(mul(ext(...), ext(...)))) +/// - UpdateR(PrevValue, ext(mul(ext(...), Constant))) +/// - UpdateR(PrevValue, abs(sub(ext(...), ext(...))) +/// +/// Note: The second operand of UpdateR corresponds to \p Op in the examples. +static std::optional +matchExtendedReductionOperand(VPWidenRecipe *UpdateR, VPValue *Op) { + assert(is_contained(UpdateR->operands(), Op) && + "Op should be operand of UpdateR"); + + // Try matching an absolute difference operand of the form + // `abs(sub(ext(A), ext(B)))`. This will be later transformed into + // `ext(absolute-difference(A, B))`. This allows us to perform the absolute + // difference on a wider type and get the extend for "free" from the partial + // reduction. + VPValue *X, *Y; + if (Op->hasOneUse() && + match(Op, m_WidenIntrinsic( + m_OneUse(m_Sub(m_WidenAnyExtend(m_VPValue(X)), + m_WidenAnyExtend(m_VPValue(Y))))))) { + auto *Abs = cast(Op); + auto *Sub = cast(Abs->getOperand(0)); + auto *LHSExt = cast(Sub->getOperand(0)); + auto *RHSExt = cast(Sub->getOperand(1)); + Type *LHSInputType = X->getScalarType(); + Type *RHSInputType = Y->getScalarType(); + if (LHSInputType != RHSInputType || + LHSExt->getOpcode() != RHSExt->getOpcode()) + return std::nullopt; + // Note: This is essentially the same as matching ext(...) as we will + // rewrite this operand to ext(absolute-difference(A, B)). + return ExtendedReductionOperand{ + Sub, + /*ExtendA=*/{LHSInputType, getPartialReductionExtendKind(LHSExt)}, + /*ExtendB=*/{}}; + } + + std::optional OuterExtKind; + if (match(Op, m_WidenAnyExtend(m_VPValue()))) { + auto *CastRecipe = cast(Op); + VPValue *CastSource = CastRecipe->getOperand(0); + OuterExtKind = getPartialReductionExtendKind(CastRecipe); + if (match(CastSource, m_Mul(m_VPValue(), m_VPValue())) || + match(CastSource, m_FMul(m_VPValue(), m_VPValue()))) { + // Match: ext(mul(...)) + // Record the outer extend kind and set `Op` to the mul. We can then match + // this as a binary operation. Note: We can optimize out the outer extend + // by widening the inner extends to match it. See + // optimizeExtendsForPartialReduction. + Op = CastSource; + } else { + return ExtendedReductionOperand{ + UpdateR, + /*ExtendA=*/{CastSource->getScalarType(), *OuterExtKind}, + /*ExtendB=*/{}}; + } + } + + if (!Op->hasOneUse()) + return std::nullopt; + + VPWidenRecipe *MulOp = dyn_cast(Op); + if (!MulOp || + !is_contained({Instruction::Mul, Instruction::FMul}, MulOp->getOpcode())) + return std::nullopt; + + // The rest of the matching assumes `Op` is a (possibly extended) mul + // operation. + + VPValue *LHS = MulOp->getOperand(0); + VPValue *RHS = MulOp->getOperand(1); + + // The LHS of the operation must always be an extend. + if (!match(LHS, m_WidenAnyExtend(m_VPValue()))) + return std::nullopt; + + auto *LHSCast = cast(LHS); + Type *LHSInputType = LHSCast->getOperand(0)->getScalarType(); + ExtendKind LHSExtendKind = getPartialReductionExtendKind(LHSCast); + + // The RHS of the operation can be an extend or a constant integer. + const APInt *RHSConst = nullptr; + VPWidenCastRecipe *RHSCast = nullptr; + if (match(RHS, m_WidenAnyExtend(m_VPValue()))) + RHSCast = cast(RHS); + else if (!match(RHS, m_APInt(RHSConst)) || + !canConstantBeExtended(RHSConst, LHSInputType, LHSExtendKind)) + return std::nullopt; + + // The outer extend kind must match the inner extends for folding. + for (VPWidenCastRecipe *Cast : {LHSCast, RHSCast}) + if (Cast && OuterExtKind && + getPartialReductionExtendKind(Cast) != OuterExtKind) + return std::nullopt; + + Type *RHSInputType = LHSInputType; + ExtendKind RHSExtendKind = LHSExtendKind; + if (RHSCast) { + RHSInputType = RHSCast->getOperand(0)->getScalarType(); + RHSExtendKind = getPartialReductionExtendKind(RHSCast); + } + + return ExtendedReductionOperand{ + MulOp, {LHSInputType, LHSExtendKind}, {RHSInputType, RHSExtendKind}}; +} + +/// Examines each operation in the reduction chain corresponding to \p RedPhiR, +/// and determines if the target can use a cheaper operation with a wider +/// per-iteration input VF and narrower PHI VF. If successful, returns the chain +/// of operations in the reduction. +static std::optional> +getScaledReductions(VPReductionPHIRecipe *RedPhiR) { + // Get the backedge value from the reduction PHI and find the + // ComputeReductionResult that uses it (directly or through a select for + // predicated reductions). + auto *RdxResult = vputils::findComputeReductionResult(RedPhiR); + if (!RdxResult) + return std::nullopt; + VPValue *ExitValue = RdxResult->getOperand(0); + match(ExitValue, m_Select(m_VPValue(), m_VPValue(ExitValue), m_VPValue())); + + SmallVector Chain; + RecurKind RK = RedPhiR->getRecurrenceKind(); + Type *PhiType = RedPhiR->getScalarType(); + TypeSize PHISize = PhiType->getPrimitiveSizeInBits(); + + // Work backwards from the ExitValue examining each reduction operation. + VPValue *CurrentValue = ExitValue; + while (CurrentValue != RedPhiR) { + VPBlendRecipe *Blend = dyn_cast(CurrentValue); + std::optional BlendReductionIdx; + if (Blend) { + assert(!Blend->isNormalized() && "Expect Blend not to be normalized."); + if (Blend->getNumIncomingValues() != 2) + return std::nullopt; + + BlendReductionIdx = getBlendReductionUpdateValueIdx(Blend); + if (!BlendReductionIdx) + return std::nullopt; + + CurrentValue = Blend->getIncomingValue(*BlendReductionIdx); + } + + auto *UpdateR = dyn_cast(CurrentValue); + if (!UpdateR || !Instruction::isBinaryOp(UpdateR->getOpcode())) + return std::nullopt; + + VPValue *Op = UpdateR->getOperand(1); + VPValue *PrevValue = UpdateR->getOperand(0); + + // Find the extended operand. The other operand (PrevValue) is the next link + // in the reduction chain. + std::optional ExtendedOp = + matchExtendedReductionOperand(UpdateR, Op); + if (!ExtendedOp) { + ExtendedOp = matchExtendedReductionOperand(UpdateR, PrevValue); + if (!ExtendedOp) + return std::nullopt; + std::swap(Op, PrevValue); + } + + // Look for VPBlend(reduce(PrevValue, Op), PrevValue), where + // reduce is equal to CurrentValue. This can be lowered as + // a conditional reduction by hoisting the select to the inputs. + if (Blend && Blend->getIncomingValue(1 - *BlendReductionIdx) != PrevValue) + return std::nullopt; + + Type *ExtSrcType = ExtendedOp->ExtendA.SrcType; + TypeSize ExtSrcSize = ExtSrcType->getPrimitiveSizeInBits(); + if (!PHISize.hasKnownScalarFactor(ExtSrcSize)) + return std::nullopt; + + VPPartialReductionChain Link( + {UpdateR, *ExtendedOp, RK, + PrevValue == UpdateR->getOperand(0) ? 0U : 1U, + static_cast(PHISize.getKnownScalarFactor(ExtSrcSize)), + Blend}); + Chain.push_back(Link); + CurrentValue = PrevValue; + } + + // The chain links were collected by traversing backwards from the exit value. + // Reverse the chains so they are in program order. + std::reverse(Chain.begin(), Chain.end()); + return Chain; +} +} // namespace + +void VPlanTransforms::createPartialReductions(VPlan &Plan, + VPCostContext &CostCtx, + VFRange &Range) { + // Find all possible valid partial reductions, grouping chains by their PHI. + // This grouping allows invalidating the whole chain, if any link is not a + // valid partial reduction. + MapVector> + ChainsByPhi; + VPBasicBlock *HeaderVPBB = Plan.getVectorLoopRegion()->getEntryBasicBlock(); + for (VPRecipeBase &R : HeaderVPBB->phis()) { + auto *RedPhiR = dyn_cast(&R); + if (!RedPhiR) + continue; + + if (auto Chains = getScaledReductions(RedPhiR)) + ChainsByPhi.try_emplace(RedPhiR, std::move(*Chains)); + } + + if (ChainsByPhi.empty()) + return; + + // Build set of partial reduction operations and blends for user validation + // and a map of reduction bin ops to their scale factors for scale validation. + SmallPtrSet PartialReductionOps; + SmallPtrSet PartialReductionBlends; + DenseMap ScaledReductionMap; + for (const auto &[_, Chains] : ChainsByPhi) + for (const VPPartialReductionChain &Chain : Chains) { + PartialReductionOps.insert(Chain.ExtendedOp.ExtendsUser); + if (Chain.Blend) + PartialReductionBlends.insert(Chain.Blend); + ScaledReductionMap[Chain.ReductionBinOp] = Chain.ScaleFactor; + } + + // A partial reduction is invalid if any of its extends are used by + // something that isn't another partial reduction. This is because the + // extends are intended to be lowered along with the reduction itself. + auto ExtendUsersValid = [&](VPValue *Ext) { + return !isa(Ext) || all_of(Ext->users(), [&](VPUser *U) { + return PartialReductionOps.contains(cast(U)); + }); + }; + + auto IsProfitablePartialReductionChainForVF = + [&](ArrayRef Chain, ElementCount VF) -> bool { + InstructionCost PartialCost = 0, RegularCost = 0; + + // The chain is a profitable partial reduction chain if the cost of handling + // the entire chain is cheaper when using partial reductions than when + // handling the entire chain using regular reductions. + for (const VPPartialReductionChain &Link : Chain) { + const ExtendedReductionOperand &ExtendedOp = Link.ExtendedOp; + InstructionCost LinkCost = getPartialReductionLinkCost(CostCtx, Link, VF); + if (!LinkCost.isValid()) + return false; + + PartialCost += LinkCost; + RegularCost += Link.ReductionBinOp->computeCost(VF, CostCtx); + // If ExtendB is not none, then the "ExtendsUser" is the binary operation. + if (ExtendedOp.ExtendB.Kind != ExtendKind::PR_None) + RegularCost += ExtendedOp.ExtendsUser->computeCost(VF, CostCtx); + for (VPValue *Op : ExtendedOp.ExtendsUser->operands()) + if (auto *Extend = dyn_cast(Op)) + RegularCost += Extend->computeCost(VF, CostCtx); + } + return PartialCost.isValid() && PartialCost < RegularCost; + }; + + // Validate chains: check that extends are only used by partial reductions, + // and that reduction bin ops are only used by other partial reductions with + // matching scale factors, are outside the loop region or the select + // introduced by tail-folding. Otherwise we would create users of scaled + // reductions where the types of the other operands don't match. + for (auto &[RedPhiR, Chains] : ChainsByPhi) { + for (const VPPartialReductionChain &Chain : Chains) { + if (!all_of(Chain.ExtendedOp.ExtendsUser->operands(), ExtendUsersValid)) { + Chains.clear(); + break; + } + auto UseIsValid = [&, RedPhiR = RedPhiR](VPUser *U) { + if (auto *PhiR = dyn_cast(U)) + return PhiR == RedPhiR; + auto *R = cast(U); + + if (auto *Blend = dyn_cast(R)) + return Blend == Chain.Blend || PartialReductionBlends.contains(Blend); + + return Chain.ScaleFactor == ScaledReductionMap.lookup_or(R, 0) || + match(R, m_ComputeReductionResult( + m_Specific(Chain.ReductionBinOp))) || + match(R, m_Select(m_VPValue(), m_Specific(Chain.ReductionBinOp), + m_Specific(RedPhiR))); + }; + if (!all_of(Chain.ReductionBinOp->users(), UseIsValid)) { + Chains.clear(); + break; + } + + // Check if the compute-reduction-result is used by a sunk store. + // TODO: Also form partial reductions in those cases. + if (auto *RdxResult = vputils::findComputeReductionResult(RedPhiR)) { + if (any_of(RdxResult->users(), [](VPUser *U) { + auto *RepR = dyn_cast(U); + return RepR && RepR->getOpcode() == Instruction::Store; + })) { + Chains.clear(); + break; + } + } + } + + // Clear the chain if it is not profitable. + if (!LoopVectorizationPlanner::getDecisionAndClampRange( + [&, &Chains = Chains](ElementCount VF) { + return IsProfitablePartialReductionChainForVF(Chains, VF); + }, + Range)) + Chains.clear(); + } + + for (auto &[Phi, Chains] : ChainsByPhi) + for (const VPPartialReductionChain &Chain : Chains) + transformToPartialReduction(Chain, Plan, Phi); +} + +void VPlanTransforms::makeMemOpWideningDecisions(VPlan &Plan, VFRange &Range, + VPRecipeBuilder &RecipeBuilder, + VPCostContext &CostCtx) { + // Collect all loads/stores first. We will start with ones having simpler + // decisions followed by more complex ones that are potentially + // guided/dependent on the simpler ones. + SmallVector MemOps; + for (VPBasicBlock *VPBB : + VPBlockUtils::blocksOnly(vp_depth_first_shallow( + Plan.getVectorLoopRegion()->getEntryBasicBlock()))) { + for (VPRecipeBase &R : *VPBB) { + auto *VPI = dyn_cast(&R); + if (VPI && VPI->getUnderlyingValue() && + is_contained({Instruction::Load, Instruction::Store}, + VPI->getOpcode())) + MemOps.push_back(VPI); + } + } + + // Few helpers to process different kinds of memory operations. + + // To be used as argument to `VPlanTransforms::runPass` which explicitly + // specified pass name, hence `VPlan &` parameter. + auto ProcessSubset = [&](VPlan &, auto ProcessVPInst) { + SmallVector RemainingMemOps; + for (VPInstruction *VPI : MemOps) { + if (!ProcessVPInst(VPI)) + RemainingMemOps.push_back(VPI); + } + + MemOps.clear(); + std::swap(MemOps, RemainingMemOps); + }; + + auto ReplaceWith = [&](VPInstruction *VPI, VPRecipeBase *New) { + assert(New->getParent() && "New recipe must have been inserted"); + if (VPI->getOpcode() == Instruction::Load) + VPI->replaceAllUsesWith(New->getVPSingleValue()); + VPI->eraseFromParent(); + + // VPI has been processed. + return true; + }; + + auto Scalarize = [&](VPInstruction *VPI) { + return ReplaceWith(VPI, VPBuilder(VPI).insert( + RecipeBuilder.handleReplication(VPI, Range))); + }; + + VPBasicBlock *MiddleVPBB = Plan.getMiddleBlock(); + VPBuilder FinalRedStoresBuilder(MiddleVPBB, MiddleVPBB->getFirstNonPhi()); + VPlanTransforms::runPass( + "lowerMemoryIdioms", ProcessSubset, Plan, [&](VPInstruction *VPI) { + if (RecipeBuilder.replaceWithFinalIfReductionStore( + VPI, FinalRedStoresBuilder)) + return true; + + // Filter out scalar VPlan for the remaining idioms. + if (LoopVectorizationPlanner::getDecisionAndClampRange( + [](ElementCount VF) { return VF.isScalar(); }, Range)) + return false; + + if (VPHistogramRecipe *Histogram = RecipeBuilder.widenIfHistogram(VPI)) + return ReplaceWith(VPI, VPBuilder(VPI).insert(Histogram)); + + return false; + }); + + // Filter out scalar VPlan for the remaining memory operations. + if (LoopVectorizationPlanner::getDecisionAndClampRange( + [](ElementCount VF) { return VF.isScalar(); }, Range)) + return; + + // If the instruction's allocated size doesn't equal it's type size, it + // requires padding and will be scalarized. + VPlanTransforms::runPass( + "scalarizeMemOpsWithIrregularTypes", ProcessSubset, Plan, + [&](VPInstruction *VPI) { + Instruction *I = VPI->getUnderlyingInstr(); + if (hasIrregularType(getLoadStoreType(I), I->getDataLayout())) + return Scalarize(VPI); + + return false; + }); + + if (!RecipeBuilder.prefersVectorizedAddressing()) { + VPlanTransforms::runPass( + "makeVPlanMemOpDecision", ProcessSubset, Plan, [&](VPInstruction *VPI) { + Instruction *I = VPI->getUnderlyingInstr(); + bool IsLoad = VPI->getOpcode() == Instruction::Load; + if (RecipeBuilder.isPredicatedInst(I) || !IsLoad || + !vputils::isUsedByLoadStoreAddress(VPI)) + return false; + + // Scalarize loads used as addresses, matching the legacy CM. The load + // is single-scalar if the pointer is loop-invariant, otherwise it is + // replicated per-lane. No mask is needed as the load is not + // predicated. + VPValue *Ptr = VPI->getOperand(0); + const SCEV *PtrSCEV = + vputils::getSCEVExprForVPValue(Ptr, CostCtx.PSE, CostCtx.L); + bool IsSingleScalarLoad = + !isa(PtrSCEV) && + CostCtx.PSE.getSE()->isLoopInvariant(PtrSCEV, CostCtx.L); + + ReplaceWith(VPI, + VPBuilder(VPI).insert(new VPReplicateRecipe( + I, Ptr, /*IsSingleScalar=*/IsSingleScalarLoad, + /*Mask=*/nullptr, *VPI, *VPI, VPI->getDebugLoc()))); + return true; + }); + } + + // Widen unit-stride consecutive accesses, matching the legacy CM. Both + // forward (stride +1) and reverse (stride -1) accesses are handled. + VPlanTransforms::runPass( + "widenConsecutiveMemOps", ProcessSubset, Plan, [&](VPInstruction *VPI) { + Instruction *I = VPI->getUnderlyingInstr(); + bool IsLoad = VPI->getOpcode() == Instruction::Load; + VPValue *Ptr = VPI->getOperand(!IsLoad); + Type *ScalarTy = + IsLoad ? VPI->getScalarType() : VPI->getOperand(0)->getScalarType(); + std::optional Stride = + getConstantStride(Ptr, ScalarTy, CostCtx.PSE, CostCtx.L); + if (Stride != 1 && Stride != -1) + return false; + bool Reverse = Stride == -1; + + // A predicated access can only be widened (rather than scalarized) if + // the target supports a masked load/store for it. + // TODO: Determine if a load/store needs predication directly in VPlan. + bool IsPredicated = RecipeBuilder.isPredicatedInst(I); + if (IsPredicated && !CostCtx.Config.isLegalMaskedLoadOrStore( + IsLoad, ScalarTy, getLoadStoreAlignment(I), + getLoadStoreAddressSpace(I))) + return false; + + VPBuilder Builder(VPI); + VPSingleDefRecipe *VectorPtr = Builder.createConsecutiveVectorPointer( + Ptr, ScalarTy, Reverse, VPI->getDebugLoc()); + + VPValue *Mask = IsPredicated ? VPI->getMask() : nullptr; + // Reverse the mask so it matches the reversed access order. + if (Reverse && Mask) + Mask = Builder.createNaryOp(VPInstruction::Reverse, Mask, + VPI->getDebugLoc()); + + if (IsLoad) { + VPSingleDefRecipe *Load = Builder.createWidenLoad( + *cast(I), VectorPtr, Mask, + /*Consecutive=*/true, *VPI, VPI->getDebugLoc()); + // Reverse the loaded values back into program order. + if (Reverse) + Load = Builder.createNaryOp(VPInstruction::Reverse, Load, + VPI->getDebugLoc()); + return ReplaceWith(VPI, Load); + } + + VPValue *StoredVal = VPI->getOperand(0); + if (Reverse) + // Reverse the stored values so they are written in descending order. + StoredVal = Builder.createNaryOp(VPInstruction::Reverse, StoredVal, + VPI->getDebugLoc()); + + auto *StoreR = Builder.createWidenStore( + *cast(I), VectorPtr, StoredVal, Mask, + /*Consecutive=*/true, *VPI, VPI->getDebugLoc()); + return ReplaceWith(VPI, StoreR); + }); + + VPlanTransforms::runPass("delegateMemOpWideningToLegacyCM", ProcessSubset, + Plan, [&](VPInstruction *VPI) { + if (VPRecipeBase *Recipe = + RecipeBuilder.tryToWidenMemory(VPI, Range)) + return ReplaceWith(VPI, Recipe); + + return Scalarize(VPI); + }); +} + +void VPlanTransforms::makeScalarizationDecisions(VPlan &Plan, VFRange &Range) { + if (LoopVectorizationPlanner::getDecisionAndClampRange( + [&](ElementCount VF) { return VF.isScalar(); }, Range)) + return; + + PostOrderTraversal> POT( + Plan.getEntry()); + for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly(POT)) { + for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) { + auto *VPI = dyn_cast(&R); + if (!VPI) + continue; + + auto *I = cast_or_null(VPI->getUnderlyingValue()); + // Wouldn't be able to create a `VPReplicateRecipe` anyway. + if (!I) + continue; + + // If executing other lanes produces side-effects we can't avoid them. + if (VPI->mayHaveSideEffects()) + continue; + + // We want to drop the mask operand, verify we can safely do that. + if (VPI->isMasked() && !VPI->isSafeToSpeculativelyExecute()) + continue; + + // Avoid rewriting IV increment as that interferes with + // `removeRedundantCanonicalIVs`. + if (VPI->getOpcode() == Instruction::Add && + any_of(VPI->operands(), IsaPred)) + continue; + + // Other lanes are needed - can't drop them. + if (!vputils::onlyFirstLaneUsed(VPI)) + continue; + + auto *Recipe = VPBuilder::createSingleScalarOp( + VPI->getOpcode(), VPI->operandsWithoutMask(), /*Mask=*/nullptr, *VPI, + *VPI, VPI->getDebugLoc(), I); + Recipe->insertBefore(VPI); + VPI->replaceAllUsesWith(Recipe); + VPI->eraseFromParent(); + } + } +} + +/// Returns true if \p Info's parameter kinds are compatible with \p Args. +static bool areVFParamsOk(const VFInfo &Info, ArrayRef Args, + PredicatedScalarEvolution &PSE, const Loop *L) { + ScalarEvolution *SE = PSE.getSE(); + return all_of(Info.Shape.Parameters, [&](VFParameter Param) { + switch (Param.ParamKind) { + case VFParamKind::Vector: + case VFParamKind::GlobalPredicate: + return true; + case VFParamKind::OMP_Uniform: + return SE->isSCEVable(Args[Param.ParamPos]->getScalarType()) && + SE->isLoopInvariant( + vputils::getSCEVExprForVPValue(Args[Param.ParamPos], PSE, L), + L); + case VFParamKind::OMP_Linear: + return match(vputils::getSCEVExprForVPValue(Args[Param.ParamPos], PSE, L), + m_scev_AffineAddRec( + m_SCEV(), m_scev_SpecificSInt(Param.LinearStepOrPos), + m_SpecificLoop(L))); + default: + return false; + } + }); +} + +/// Find a vector variant of \p CI for \p VF, respecting \p MaskRequired. +/// Returns the variant function, or nullptr. Masked variants are assumed to +/// take the mask as a trailing parameter. +static Function *findVectorVariant(CallInst *CI, ArrayRef Args, + ElementCount VF, bool MaskRequired, + PredicatedScalarEvolution &PSE, + const Loop *L) { + if (CI->isNoBuiltin()) + return nullptr; + auto Mappings = VFDatabase::getMappings(*CI); + const auto *It = find_if(Mappings, [&](const VFInfo &Info) { + return Info.Shape.VF == VF && (!MaskRequired || Info.isMasked()) && + areVFParamsOk(Info, Args, PSE, L); + }); + if (It == Mappings.end()) + return nullptr; + return CI->getModule()->getFunction(It->VectorName); +} + +namespace { +/// The outcome of choosing how to widen a call at a given VF. +struct CallWideningDecision { + enum class KindTy { Scalarize, Intrinsic, VectorVariant }; + CallWideningDecision(KindTy Kind, Function *Variant = nullptr) + : Kind(Kind), Variant(Variant) {} + KindTy Kind; + + /// Set when Kind == VectorVariant. + Function *Variant; + + bool operator==(const CallWideningDecision &Other) const { + return Kind == Other.Kind && Variant == Other.Variant; + } +}; +} // namespace + +/// Pick the cheapest widening for the call \p VPI at \p VF among scalarization, +/// vector intrinsic, and vector library variant. +static CallWideningDecision decideCallWidening(VPInstruction &VPI, + ArrayRef Ops, + ElementCount VF, + VPCostContext &CostCtx) { + auto *CI = cast(VPI.getUnderlyingInstr()); + + // Scalar VFs and calls forced or known to scalarize always replicate. + if (VF.isScalar() || CostCtx.willBeScalarized(CI, VF)) + return CallWideningDecision::KindTy::Scalarize; + + auto *CalledFn = cast( + VPI.getOperand(VPI.getNumOperandsWithoutMask() - 1)->getLiveInIRValue()); + Type *ResultTy = VPI.getScalarType(); + Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, &CostCtx.TLI); + bool MaskRequired = CostCtx.isMaskRequired(CI); + + // Pseudo intrinsics (assume, lifetime, ...) are always scalarized. + if (ID && VPCostContext::isFreeScalarIntrinsic(ID)) + return CallWideningDecision::KindTy::Scalarize; + + InstructionCost ScalarCost = + VPReplicateRecipe::computeCallCost(CalledFn, ResultTy, Ops, + /*IsSingleScalar=*/false, VF, CostCtx); + + Function *VecFunc = + findVectorVariant(CI, Ops, VF, MaskRequired, CostCtx.PSE, CostCtx.L); + InstructionCost VecCallCost = InstructionCost::getInvalid(); + if (VecFunc) + VecCallCost = VPWidenCallRecipe::computeCallCost(VecFunc, CostCtx); + + // Prefer the intrinsic if it is at least as cheap as scalarizing and any + // available vector variant. + if (ID) { + InstructionCost IntrinsicCost = + VPWidenIntrinsicRecipe::computeCallCost(ID, Ops, VPI, VF, CostCtx); + if (IntrinsicCost.isValid() && ScalarCost >= IntrinsicCost && + (!VecFunc || VecCallCost >= IntrinsicCost)) + return CallWideningDecision::KindTy::Intrinsic; + } + + // Otherwise, use a vector library variant when it beats scalarizing. + if (VecFunc && ScalarCost >= VecCallCost) + return {CallWideningDecision::KindTy::VectorVariant, VecFunc}; + + return CallWideningDecision::KindTy::Scalarize; +} + +void VPlanTransforms::makeCallWideningDecisions(VPlan &Plan, VFRange &Range, + VPRecipeBuilder &RecipeBuilder, + VPCostContext &CostCtx) { + for (VPBasicBlock *VPBB : VPBlockUtils::blocksAs( + vp_depth_first_shallow(Plan.getVectorLoopRegion()->getEntry()))) { + for (VPRecipeBase &R : make_early_inc_range(*VPBB)) { + auto *VPI = dyn_cast(&R); + if (!VPI || !VPI->getUnderlyingValue() || + VPI->getOpcode() != Instruction::Call) + continue; + + auto *CI = cast(VPI->getUnderlyingInstr()); + SmallVector Ops(VPI->op_begin(), + VPI->op_begin() + CI->arg_size()); + + CallWideningDecision Decision = + decideCallWidening(*VPI, Ops, Range.Start, CostCtx); + LoopVectorizationPlanner::getDecisionAndClampRange( + [&](ElementCount VF) { + return Decision == decideCallWidening(*VPI, Ops, VF, CostCtx); + }, + Range); + + VPSingleDefRecipe *Replacement = nullptr; + switch (Decision.Kind) { + case CallWideningDecision::KindTy::Intrinsic: { + Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, &CostCtx.TLI); + Type *ResultTy = VPI->getScalarType(); + Replacement = new VPWidenIntrinsicRecipe(*CI, ID, Ops, ResultTy, *VPI, + *VPI, VPI->getDebugLoc()); + break; + } + case CallWideningDecision::KindTy::VectorVariant: { + // Masked variants take the mask as a trailing parameter, so they have + // one more parameter than the original call's arguments. + if (Decision.Variant->arg_size() > Ops.size()) { + VPValue *Mask = VPI->isMasked() ? VPI->getMask() : Plan.getTrue(); + Ops.push_back(Mask); + } + Ops.push_back(VPI->getOperand(VPI->getNumOperandsWithoutMask() - 1)); + Replacement = new VPWidenCallRecipe(CI, Decision.Variant, Ops, *VPI, + *VPI, VPI->getDebugLoc()); + break; + } + case CallWideningDecision::KindTy::Scalarize: + Replacement = RecipeBuilder.handleReplication(VPI, Range); + break; + } + + Replacement->insertBefore(VPI); + VPI->replaceAllUsesWith(Replacement); + VPI->eraseFromParent(); + } + } +} + +void VPlanTransforms::createInterleaveGroups( + VPlan &Plan, + const SmallPtrSetImpl *> + &InterleaveGroups, + const bool &EpilogueAllowed) { + if (InterleaveGroups.empty()) + return; + + DenseMap IRMemberToRecipe; + for (VPBasicBlock *VPBB : + VPBlockUtils::blocksOnly(vp_depth_first_shallow( + Plan.getVectorLoopRegion()->getEntryBasicBlock()))) + for (VPRecipeBase &R : make_filter_range(*VPBB, [](VPRecipeBase &R) { + return isa(&R); + })) { + auto *MemR = cast(&R); + IRMemberToRecipe[&MemR->getIngredient()] = MemR; + } + + // Interleave memory: for each Interleave Group we marked earlier as relevant + // for this VPlan, replace the Recipes widening its memory instructions with a + // single VPInterleaveRecipe at its insertion point. + VPDominatorTree VPDT(Plan); + for (const auto *IG : InterleaveGroups) { + VPWidenMemoryRecipe *Start = nullptr; + Instruction *StartMember = nullptr; + for (auto *Member : IG->members()) + if (VPWidenMemoryRecipe *R = IRMemberToRecipe.lookup(Member)) { + StartMember = Member; + Start = R; + break; + } + if (!StartMember) // All member recipes are dead, so the group is dead. + continue; + VPIRMetadata InterleaveMD(*Start); + SmallVector StoredValues; + for (unsigned I = 0; I < IG->getFactor(); ++I) { + Instruction *MemberI = IG->getMember(I); + if (!MemberI) + continue; + if (VPWidenMemoryRecipe *MemoryR = IRMemberToRecipe.lookup(MemberI)) { + if (auto *StoreR = dyn_cast(MemoryR->getAsRecipe())) + StoredValues.push_back(StoreR->getStoredValue()); + InterleaveMD.intersect(*MemoryR); + } else { + InterleaveMD.intersect(VPIRMetadata(*MemberI)); + } + } + + bool NeedsMaskForGaps = + (IG->requiresScalarEpilogue() && !EpilogueAllowed) || + (!StoredValues.empty() && !IG->isFull()); + + Instruction *IRInsertPos = IG->getInsertPos(); + auto *InsertPos = IRMemberToRecipe.lookup(IRInsertPos); + if (!InsertPos) { + // InsertPos member is dead: find a new member that is alive. + assert(isa(Start->getAsRecipe()) && + "Dead member in non-load group?"); + InsertPos = Start; + for (Instruction *Member : IG->members()) + if (VPWidenMemoryRecipe *MemberR = IRMemberToRecipe.lookup(Member)) + if (VPDT.properlyDominates(MemberR->getAsRecipe(), + InsertPos->getAsRecipe())) + InsertPos = MemberR; + IRInsertPos = &InsertPos->getIngredient(); + } + VPRecipeBase *InsertPosR = InsertPos->getAsRecipe(); + + GEPNoWrapFlags NW = GEPNoWrapFlags::none(); + if (auto *Gep = dyn_cast( + getLoadStorePointerOperand(IRInsertPos)->stripPointerCasts())) + NW = Gep->getNoWrapFlags().withoutNoUnsignedWrap(); + + // Get or create the start address for the interleave group. + VPValue *Addr = Start->getAddr(); + VPRecipeBase *AddrDef = Addr->getDefiningRecipe(); + if (IG->getIndex(StartMember) != 0 || + (AddrDef && !VPDT.properlyDominates(AddrDef, InsertPosR))) { + // Either member zero's recipe is dead, or we cannot re-use the address of + // member zero because it does not dominate the insert position. Instead, + // use the address of the insert position and create a PtrAdd adjusting it + // to the address of member zero. + // TODO: Hoist Addr's defining recipe (and any operands as needed) to + // InsertPos or sink loads above zero members to join it. + assert(IG->getIndex(IRInsertPos) != 0 && + "index of insert position shouldn't be zero"); + auto &DL = IRInsertPos->getDataLayout(); + APInt Offset(32, + DL.getTypeAllocSize(getLoadStoreType(IRInsertPos)) * + IG->getIndex(IRInsertPos), + /*IsSigned=*/true); + VPValue *OffsetVPV = Plan.getConstantInt(-Offset); + VPBuilder B(InsertPosR); + Addr = B.createNoWrapPtrAdd(InsertPos->getAddr(), OffsetVPV, NW); + } + // If the group is reverse, adjust the index to refer to the last vector + // lane instead of the first. We adjust the index from the first vector + // lane, rather than directly getting the pointer for lane VF - 1, because + // the pointer operand of the interleaved access is supposed to be uniform. + if (IG->isReverse()) { + auto *ReversePtr = new VPVectorEndPointerRecipe( + Addr, &Plan.getVF(), getLoadStoreType(IRInsertPos), + -(int64_t)IG->getFactor(), NW, InsertPosR->getDebugLoc()); + ReversePtr->insertBefore(InsertPosR); + Addr = ReversePtr; + } + auto *VPIG = new VPInterleaveRecipe( + IG, Addr, StoredValues, InsertPos->getMask(), NeedsMaskForGaps, + InterleaveMD, InsertPosR->getDebugLoc()); + VPIG->insertBefore(InsertPosR); + + unsigned J = 0; + for (unsigned i = 0; i < IG->getFactor(); ++i) + if (Instruction *Member = IG->getMember(i)) { + VPWidenMemoryRecipe *MemberR = IRMemberToRecipe.lookup(Member); + if (!Member->getType()->isVoidTy()) { + if (MemberR) { + VPValue *OriginalV = MemberR->getAsRecipe()->getVPSingleValue(); + OriginalV->replaceAllUsesWith(VPIG->getVPValue(J)); + } + J++; + } + if (MemberR) + MemberR->getAsRecipe()->eraseFromParent(); + } + } +} + +void VPlanTransforms::convertToStridedAccesses(VPlan &Plan, + PredicatedScalarEvolution &PSE, + Loop &L, VPCostContext &Ctx, + VFRange &Range) { + if (Plan.hasScalarVFOnly()) + return; + + VPRegionBlock *VectorLoop = Plan.getVectorLoopRegion(); + VPValue *I32VF = nullptr; + for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly( + vp_depth_first_shallow(VectorLoop->getEntry()))) { + for (VPRecipeBase &R : make_early_inc_range(*VPBB)) { + auto *LoadR = dyn_cast(&R); + // TODO: Support strided store. + // TODO: Transform reverse access into strided access with -1 stride. + // TODO: Transform gather/scatter with uniform address into strided access + // with 0 stride. + // TODO: Transform interleave access into multiple strided accesses. + if (!LoadR || LoadR->isConsecutive()) + continue; + + VPValue *Ptr = LoadR->getAddr(); + // Check if this is a strided access by analyzing the address SCEV for an + // affine addRec. + const SCEV *PtrSCEV = vputils::getSCEVExprForVPValue(Ptr, PSE, &L); + const SCEV *Start; + const SCEVConstant *Step; + // TODO: Support non-constant loop invariant stride. + if (!match(PtrSCEV, + m_scev_AffineAddRec(m_SCEV(Start), m_SCEVConstant(Step), + m_SpecificLoop(&L)))) + continue; + + Type *LoadTy = LoadR->getScalarType(); + Align Alignment = LoadR->getAlign(); + auto IsProfitable = [&](ElementCount VF) { + Type *DataTy = toVectorTy(LoadTy, VF); + if (!Ctx.TTI.isLegalStridedLoadStore(DataTy, Alignment)) + return false; + const InstructionCost CurrentCost = LoadR->computeCost(VF, Ctx); + const InstructionCost StridedLoadStoreCost = + VPWidenMemIntrinsicRecipe::computeMemIntrinsicCost( + Intrinsic::experimental_vp_strided_load, DataTy, + LoadR->isMasked(), Alignment, Ctx); + return StridedLoadStoreCost < CurrentCost; + }; + + if (!LoopVectorizationPlanner::getDecisionAndClampRange(IsProfitable, + Range)) + continue; + + // Invalidate the legacy widening decision so the cost of replaced load is + // not counted during precomputeCosts. + // TODO: Remove once the legacy exit cost computation is retired. + for (ElementCount VF : Range) + Ctx.invalidateWideningDecision(&LoadR->getIngredient(), VF); + + // Get VF as i32 for the vector length operand. + if (!I32VF) { + VPBuilder Builder(Plan.getVectorPreheader()); + I32VF = Builder.createScalarZExtOrTrunc( + &Plan.getVF(), Type::getInt32Ty(Plan.getContext()), + DebugLoc::getUnknown()); + } + + VPBuilder Builder(LoadR); + // Create the base pointer of strided access. + // TODO: reuse VPDerivedIVRecipe for base pointer computation when it + // supports a general VPValue as the start value. + VPValue *StartVPV = + VPSCEVExpander(Builder, *PSE.getSE(), LoadR->getDebugLoc()) + .tryToExpand(Start); + if (!StartVPV) + StartVPV = VPBuilder(Plan.getEntry()).createExpandSCEV(Start); + VPValue *StrideInBytes = Plan.getOrAddLiveIn(Step->getValue()); + Type *IndexTy = Plan.getDataLayout().getIndexType(Ptr->getScalarType()); + assert(IndexTy == StrideInBytes->getScalarType() && + "Stride type from SCEV must match the index type"); + VPValue *CanIV = Builder.createScalarZExtOrTrunc( + VectorLoop->getCanonicalIV(), IndexTy, DebugLoc::getUnknown()); + auto *AddRecPtr = cast(PtrSCEV); + auto *Offset = Builder.createOverflowingOp( + Instruction::Mul, {CanIV, StrideInBytes}, + {AddRecPtr->hasNoUnsignedWrap(), /*HasNSW=*/false}); + GEPNoWrapFlags NWFlags = AddRecPtr->hasNoUnsignedWrap() + ? GEPNoWrapFlags::noUnsignedWrap() + : GEPNoWrapFlags::none(); + VPValue *BasePtr = Builder.createNoWrapPtrAdd(StartVPV, Offset, NWFlags); + + // Create a new vector pointer for strided access. + VPValue *NewPtr = Builder.createVectorPointer( + BasePtr, Type::getInt8Ty(Plan.getContext()), StrideInBytes, NWFlags, + LoadR->getDebugLoc()); + + VPValue *Mask = LoadR->getMask(); + if (!Mask) + Mask = Plan.getTrue(); + auto *StridedLoad = Builder.createWidenMemIntrinsic( + Intrinsic::experimental_vp_strided_load, + {NewPtr, StrideInBytes, Mask, I32VF}, LoadTy, Alignment, *LoadR, + LoadR->getDebugLoc()); + LoadR->replaceAllUsesWith(StridedLoad); + } + } +} + +void VPlanTransforms::legalizeAndOptimizeInductions(VPlan &Plan) { + VPBasicBlock *HeaderVPBB = Plan.getVectorLoopRegion()->getEntryBasicBlock(); + bool HasOnlyVectorVFs = !Plan.hasScalarVFOnly(); + VPBuilder Builder(HeaderVPBB, HeaderVPBB->getFirstNonPhi()); + for (VPRecipeBase &Phi : HeaderVPBB->phis()) { + auto *PhiR = dyn_cast(&Phi); + if (!PhiR) + continue; + + // Try to narrow wide and replicating recipes to uniform recipes, based on + // VPlan analysis. + // TODO: Apply to all recipes in the future, to replace legacy uniformity + // analysis. + auto Users = vputils::collectUsersRecursively(PhiR); + for (VPUser *U : reverse(Users)) { + auto *Def = dyn_cast(U); + auto *RepR = dyn_cast(U); + // Skip recipes that shouldn't be narrowed. + if (!Def || !isa(Def) || + Def->user_empty() || !Def->getUnderlyingValue() || + (RepR && (RepR->isSingleScalar() || RepR->isPredicated()))) + continue; + + // Skip recipes that may have other lanes than their first used. + if (!vputils::isSingleScalar(Def) && !vputils::onlyFirstLaneUsed(Def)) + continue; + + // TODO: Support scalarizing ExtractValue. + if (match(Def, + m_Binary(m_VPValue(), m_VPValue()))) + continue; + + auto *Clone = VPBuilder::createSingleScalarOp( + Def->getUnderlyingInstr()->getOpcode(), Def->operands(), + /*Mask=*/nullptr, *Def, {}, DebugLoc::getUnknown(), + Def->getUnderlyingInstr()); + Clone->insertAfter(Def); + Def->replaceAllUsesWith(Clone); + } + + // Replace wide pointer inductions which have only their scalars used by + // PtrAdd(IndStart, ScalarIVSteps (0, Step)). + if (auto *PtrIV = dyn_cast(&Phi)) { + if (!Plan.hasScalarVFOnly() && + !PtrIV->onlyScalarsGenerated(Plan.hasScalableVF())) + continue; + + VPValue *PtrAdd = + vputils::scalarizeVPWidenPointerInduction(PtrIV, Plan, Builder); + PtrIV->replaceAllUsesWith(PtrAdd); + continue; + } + + // Replace widened induction with scalar steps for users that only use + // scalars. + auto *WideIV = cast(&Phi); + if (HasOnlyVectorVFs && none_of(WideIV->users(), [WideIV](VPUser *U) { + return U->usesScalars(WideIV); + })) + continue; + + const InductionDescriptor &ID = WideIV->getInductionDescriptor(); + VPIRFlags::WrapFlagsTy WrapFlags; + // We can preserve nuw when the step is non-negative. + const APInt *Step; + if (match(WideIV->getStepValue(), m_APInt(Step)) && Step->isNonNegative()) + WrapFlags = {static_cast(WideIV->getNoWrapFlagsOrNone().HasNUW), + false}; + VPScalarIVStepsRecipe *Steps = vputils::createScalarIVSteps( + Plan, ID.getKind(), ID.getInductionOpcode(), + dyn_cast_or_null(ID.getInductionBinOp()), + WideIV->getTruncInst(), WideIV->getStartValue(), WideIV->getStepValue(), + WideIV->getDebugLoc(), Builder, WrapFlags); + + // Update scalar users of IV to use Step instead. + if (!HasOnlyVectorVFs) { + assert(!Plan.hasScalableVF() && + "plans containing a scalar VF cannot also include scalable VFs"); + WideIV->replaceAllUsesWith(Steps); + } else { + bool HasScalableVF = Plan.hasScalableVF(); + WideIV->replaceUsesWithIf(Steps, + [WideIV, HasScalableVF](VPUser &U, unsigned) { + if (HasScalableVF) + return U.usesFirstLaneOnly(WideIV); + return U.usesScalars(WideIV); + }); + } + } +} + +static std::optional +getUnmaskedDivRemOpcode(Intrinsic::ID ID) { + switch (ID) { + case Intrinsic::masked_udiv: + return Instruction::UDiv; + case Intrinsic::masked_sdiv: + return Instruction::SDiv; + case Intrinsic::masked_urem: + return Instruction::URem; + case Intrinsic::masked_srem: + return Instruction::SRem; + default: + return {}; + } +} + +void VPlanTransforms::narrowToSingleScalarRecipes(VPlan &Plan) { + if (Plan.hasScalarVFOnly()) + return; + + for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly( + vp_depth_first_deep(Plan.getEntry()))) { + for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) { + if (!isa(&R)) + continue; + auto *RepR = dyn_cast(&R); + if (RepR && (RepR->isSingleScalar() || RepR->isPredicated())) + continue; + + auto *RepOrWidenR = cast(&R); + if (RepR && RepR->getOpcode() == Instruction::Store && + vputils::isSingleScalar(RepR->getOperand(1))) { + auto *Clone = new VPReplicateRecipe( + RepOrWidenR->getUnderlyingInstr(), RepOrWidenR->operands(), + true /*IsSingleScalar*/, nullptr /*Mask*/, *RepR /*Flags*/, + *RepR /*Metadata*/, RepR->getDebugLoc()); + Clone->insertBefore(RepOrWidenR); + VPBuilder Builder(Clone); + VPValue *ExtractOp = Clone->getOperand(0); + if (vputils::isUniformAcrossVFsAndUFs(RepR->getOperand(1))) + ExtractOp = + Builder.createNaryOp(VPInstruction::ExtractLastPart, ExtractOp); + ExtractOp = + Builder.createNaryOp(VPInstruction::ExtractLastLane, ExtractOp); + Clone->setOperand(0, ExtractOp); + RepR->eraseFromParent(); + continue; + } + + // Narrow llvm.masked.{u,s}{div,rem} intrinsics with a safe divisor. + if (auto *IntrR = dyn_cast(RepOrWidenR)) { + if (!vputils::onlyFirstLaneUsed(IntrR)) + continue; + auto Opc = getUnmaskedDivRemOpcode(IntrR->getVectorIntrinsicID()); + if (!Opc) + continue; + VPBuilder Builder(IntrR); + VPValue *SafeDivisor = Builder.createSelect( + IntrR->getOperand(2), IntrR->getOperand(1), + Plan.getConstantInt(IntrR->getScalarType(), 1)); + VPValue *Clone = Builder.createNaryOp( + *Opc, {IntrR->getOperand(0), SafeDivisor}, + VPIRFlags::getDefaultFlags(*Opc), IntrR->getDebugLoc()); + IntrR->replaceAllUsesWith(Clone); + IntrR->eraseFromParent(); + continue; + } + + // Skip recipes that aren't single scalars. + if (!vputils::isSingleScalar(RepOrWidenR)) + continue; + + // Predicate to check if a user of Op introduces extra broadcasts. + auto IntroducesBCastOf = [](const VPValue *Op) { + return [Op](const VPUser *U) { + if (auto *VPI = dyn_cast(U)) { + if (is_contained({VPInstruction::ExtractLastLane, + VPInstruction::ExtractLastPart, + VPInstruction::ExtractPenultimateElement}, + VPI->getOpcode())) + return false; + } + return !U->usesScalars(Op); + }; + }; + + if (any_of(RepOrWidenR->users(), IntroducesBCastOf(RepOrWidenR)) && + none_of(RepOrWidenR->operands(), [&](VPValue *Op) { + if (any_of( + make_filter_range(Op->users(), not_equal_to(RepOrWidenR)), + IntroducesBCastOf(Op))) + return false; + // Non-constant live-ins require broadcasts, while constants do not + // need explicit broadcasts. + bool LiveInNeedsBroadcast = + isa(Op) && !isa(Op); + auto *OpR = dyn_cast(Op); + return LiveInNeedsBroadcast || (OpR && OpR->isSingleScalar()); + })) + continue; + + auto *Clone = VPBuilder::createSingleScalarOp( + vputils::getOpcode(RepOrWidenR), RepOrWidenR->operands(), + /*Mask=*/nullptr, *RepOrWidenR, {}, DebugLoc::getUnknown(), + RepOrWidenR->getUnderlyingInstr()); + Clone->insertBefore(RepOrWidenR); + RepOrWidenR->replaceAllUsesWith(Clone); + if (vputils::isDeadRecipe(*RepOrWidenR)) + RepOrWidenR->eraseFromParent(); + } + } +} + +/// Returns true if \p V is VPWidenLoadRecipe or VPInterleaveRecipe that can be +/// converted to a narrower recipe. \p V is used by a wide recipe that feeds a +/// store interleave group at index \p Idx, \p WideMember0 is the recipe feeding +/// the same interleave group at index 0. A VPWidenLoadRecipe can be narrowed to +/// an index-independent load if it feeds all wide ops at all indices (\p OpV +/// must be the operand at index \p OpIdx for both the recipe at lane 0, \p +/// WideMember0). A VPInterleaveRecipe can be narrowed to a wide load, if \p V +/// is defined at \p Idx of a load interleave group. +/// A live-in or recipe defined outside the loop region can be converted, if it +/// is the same across all lanes, or we can create a BuildVector for it. +static bool canNarrowLoad(VPSingleDefRecipe *WideMember0, unsigned OpIdx, + VPValue *OpV, unsigned Idx, bool IsScalable) { + VPValue *Member0Op = WideMember0->getOperand(OpIdx); + if (Member0Op->isDefinedOutsideLoopRegions()) { + // Operand matches Member0, broadcast across all fields for both live-ins + // and recipes. + if (Member0Op == OpV) + return true; + // Otherwise distinct per-field VPValues are assembled into a BuildVector. + return !IsScalable && OpV->isDefinedOutsideLoopRegions() && + OpV->getScalarType() == Member0Op->getScalarType(); + } + VPRecipeBase *Member0OpR = Member0Op->getDefiningRecipe(); + if (auto *W = dyn_cast(Member0OpR)) + // For scalable VFs, the narrowed plan processes vscale iterations at once, + // so a shared wide load cannot be narrowed to a uniform scalar; bail out. + return !IsScalable && !W->getMask() && W->isConsecutive() && + Member0Op == OpV; + if (auto *IR = dyn_cast(Member0OpR)) + return IR->getInterleaveGroup()->isFull() && IR->getVPValue(Idx) == OpV; + return false; +} + +static bool canNarrowOps(ArrayRef Ops, bool IsScalable) { + SmallVector Ops0; + auto *WideMember0 = dyn_cast(Ops[0]); + if (!WideMember0) + return false; + for (VPValue *V : Ops) { + if (!isa(V)) + return false; + auto *R = cast(V); + if (vputils::getOpcode(R) != vputils::getOpcode(WideMember0)) + return false; + if (R->getScalarType() != WideMember0->getScalarType()) + return false; + if (R->hasPredicate() && R->getPredicate() != WideMember0->getPredicate()) + return false; + } + + for (unsigned Idx = 0; Idx != WideMember0->getNumOperands(); ++Idx) { + SmallVector OpsI; + for (VPValue *Op : Ops) + OpsI.push_back(Op->getDefiningRecipe()->getOperand(Idx)); + + if (canNarrowOps(OpsI, IsScalable)) + continue; + + if (any_of(enumerate(OpsI), [WideMember0, Idx, IsScalable](const auto &P) { + const auto &[OpIdx, OpV] = P; + return !canNarrowLoad(WideMember0, Idx, OpV, OpIdx, IsScalable); + })) + return false; + } + + return true; +} + +/// Returns VF from \p VFs if \p IR is a full interleave group with factor and +/// number of members both equal to VF. The interleave group must also access +/// the full vector width. +static std::optional +isConsecutiveInterleaveGroup(VPInterleaveRecipe *InterleaveR, + ArrayRef VFs, + const TargetTransformInfo &TTI) { + if (!InterleaveR || InterleaveR->getMask()) + return std::nullopt; + + Type *GroupElementTy = nullptr; + if (InterleaveR->getStoredValues().empty()) { + GroupElementTy = InterleaveR->getVPValue(0)->getScalarType(); + if (!all_of(InterleaveR->definedValues(), [GroupElementTy](VPValue *Op) { + return Op->getScalarType() == GroupElementTy; + })) + return std::nullopt; + } else { + GroupElementTy = InterleaveR->getStoredValues()[0]->getScalarType(); + if (!all_of(InterleaveR->getStoredValues(), [GroupElementTy](VPValue *Op) { + return Op->getScalarType() == GroupElementTy; + })) + return std::nullopt; + } + + auto IG = InterleaveR->getInterleaveGroup(); + if (IG->getFactor() != IG->getNumMembers()) + return std::nullopt; + + auto GetVectorBitWidthForVF = [&TTI](ElementCount VF) { + TypeSize Size = TTI.getRegisterBitWidth( + VF.isFixed() ? TargetTransformInfo::RGK_FixedWidthVector + : TargetTransformInfo::RGK_ScalableVector); + assert(Size.isScalable() == VF.isScalable() && + "if Size is scalable, VF must be scalable and vice versa"); + return Size.getKnownMinValue(); + }; + + for (ElementCount VF : VFs) { + unsigned MinVal = VF.getKnownMinValue(); + unsigned GroupSize = GroupElementTy->getScalarSizeInBits() * MinVal; + if (IG->getFactor() == MinVal && GroupSize == GetVectorBitWidthForVF(VF)) + return {VF}; + } + return std::nullopt; +} + +/// Returns true if \p VPValue is a narrow VPValue. +static bool isAlreadyNarrow(VPValue *VPV) { + if (isa(VPV)) + return true; + auto *RepR = dyn_cast(VPV); + return RepR && RepR->isSingleScalar(); +} + +// Convert the wide recipes defining the VPValues in \p Members feeding an +// interleave group to a single narrow variant. The first member is reused as +// the narrowed recipe. BuildVectors for live-in operands are inserted into \p +// Preheader. +static VPValue *narrowInterleaveGroupOp(ArrayRef Members, + SmallPtrSetImpl &NarrowedOps, + VPBasicBlock *Preheader) { + VPValue *V = Members.front(); + if (NarrowedOps.contains(V)) + return V; + + if (V->isDefinedOutsideLoopRegions()) { + assert(all_of(Members, + [V](VPValue *M) { + return M->isDefinedOutsideLoopRegions() && + M->getScalarType() == V->getScalarType(); + }) && + "expected distinct loop-invariant values of matching scalar type"); + auto *BV = new VPInstruction(VPInstruction::BuildVector, Members); + Preheader->appendRecipe(BV); + NarrowedOps.insert(BV); + return BV; + } + + if (isAlreadyNarrow(V)) + return V; + + VPRecipeBase *R = V->getDefiningRecipe(); + if (isa(R)) { + auto *WideMember0 = cast(R); + for (VPValue *Member : Members.drop_front()) + WideMember0->intersectFlags(*cast(Member)); + for (unsigned Idx = 0, E = WideMember0->getNumOperands(); Idx != E; ++Idx) { + SmallVector OpsI; + for (VPValue *Member : Members) + OpsI.push_back(Member->getDefiningRecipe()->getOperand(Idx)); + WideMember0->setOperand( + Idx, narrowInterleaveGroupOp(OpsI, NarrowedOps, Preheader)); + } + return V; + } + + if (auto *LoadGroup = dyn_cast(R)) { + // Narrow interleave group to wide load, as transformed VPlan will only + // process one original iteration. + auto *LI = cast(LoadGroup->getInterleaveGroup()->getInsertPos()); + auto *L = VPBuilder(LoadGroup).createWidenLoad( + *LI, LoadGroup->getAddr(), LoadGroup->getMask(), /*Consecutive=*/true, + *LoadGroup, LoadGroup->getDebugLoc()); + NarrowedOps.insert(L); + return L; + } + + if (auto *RepR = dyn_cast(R)) { + assert(RepR->isSingleScalar() && RepR->getOpcode() == Instruction::Load && + "must be a single scalar load"); + NarrowedOps.insert(RepR); + return RepR; + } + + auto *WideLoad = cast(R); + VPValue *PtrOp = WideLoad->getAddr(); + if (auto *VecPtr = dyn_cast(PtrOp)) + PtrOp = VecPtr->getOperand(0); + // Narrow wide load to uniform scalar load, as transformed VPlan will only + // process one original iteration. + auto *N = new VPReplicateRecipe(&WideLoad->getIngredient(), {PtrOp}, + /*IsUniform*/ true, + /*Mask*/ nullptr, {}, *WideLoad); + N->insertBefore(WideLoad); + NarrowedOps.insert(N); + return N; +} + +std::unique_ptr +VPlanTransforms::narrowInterleaveGroups(VPlan &Plan, + const TargetTransformInfo &TTI) { + VPRegionBlock *VectorLoop = Plan.getVectorLoopRegion(); + + if (!VectorLoop) + return nullptr; + + // Only handle single-block loops for now. + if (VectorLoop->getEntryBasicBlock() != VectorLoop->getExitingBasicBlock()) + return nullptr; + + // Skip plans when we may not be able to properly narrow. + VPBasicBlock *Exiting = VectorLoop->getExitingBasicBlock(); + if (!match(&Exiting->back(), m_BranchOnCount())) + return nullptr; + + assert(match(&Exiting->back(), + m_BranchOnCount(m_Add(m_VPValue(), m_Specific(&Plan.getVFxUF())), + m_Specific(&Plan.getVectorTripCount()))) && + "unexpected branch-on-count"); + + SmallVector StoreGroups; + std::optional VFToOptimize; + for (auto &R : *VectorLoop->getEntryBasicBlock()) { + if (isa(&R) && + vputils::onlyFirstLaneUsed(cast(&R))) + continue; + + // Bail out on recipes not supported at the moment: + // * phi recipes other than the canonical induction + // * recipes writing to memory except interleave groups + // Only support plans with a canonical induction phi. + if (R.isPhi()) + return nullptr; + + auto *InterleaveR = dyn_cast(&R); + if (R.mayWriteToMemory() && !InterleaveR) + return nullptr; + + // Bail out if any recipe defines a vector value used outside the + // vector loop region. + if (any_of(R.definedValues(), [&](VPValue *V) { + return any_of(V->users(), [&](VPUser *U) { + auto *UR = cast(U); + return UR->getParent()->getParent() != VectorLoop; + }); + })) + return nullptr; + + // All other ops are allowed, but we reject uses that cannot be converted + // when checking all allowed consumers (store interleave groups) below. + if (!InterleaveR) + continue; + + // Try to find a single VF, where all interleave groups are consecutive and + // saturate the full vector width. If we already have a candidate VF, check + // if it is applicable for the current InterleaveR, otherwise look for a + // suitable VF across the Plan's VFs. + SmallVector VFs = + VFToOptimize ? SmallVector({*VFToOptimize}) + : to_vector(Plan.vectorFactors()); + std::optional NarrowedVF = + isConsecutiveInterleaveGroup(InterleaveR, VFs, TTI); + if (!NarrowedVF || (VFToOptimize && NarrowedVF != VFToOptimize)) + return nullptr; + VFToOptimize = NarrowedVF; + + // Skip read interleave groups. + if (InterleaveR->getStoredValues().empty()) + continue; + + // Narrow interleave groups, if all operands are already matching narrow + // ops. + auto *Member0 = InterleaveR->getStoredValues()[0]; + if (isAlreadyNarrow(Member0) && + all_of(InterleaveR->getStoredValues(), equal_to(Member0))) { + StoreGroups.push_back(InterleaveR); + continue; + } + + // For now, we only support full interleave groups storing load interleave + // groups. + if (all_of(enumerate(InterleaveR->getStoredValues()), [](auto Op) { + VPRecipeBase *DefR = Op.value()->getDefiningRecipe(); + if (!DefR) + return false; + auto *IR = dyn_cast(DefR); + return IR && IR->getInterleaveGroup()->isFull() && + IR->getVPValue(Op.index()) == Op.value(); + })) { + StoreGroups.push_back(InterleaveR); + continue; + } + + // Check if all values feeding InterleaveR are matching wide recipes, which + // operands that can be narrowed. + if (!canNarrowOps(InterleaveR->getStoredValues(), + VFToOptimize->isScalable())) + return nullptr; + StoreGroups.push_back(InterleaveR); + } + + if (StoreGroups.empty()) + return nullptr; + + VPBasicBlock *MiddleVPBB = Plan.getMiddleBlock(); + bool RequiresScalarEpilogue = + MiddleVPBB->getNumSuccessors() == 1 && + MiddleVPBB->getSingleSuccessor() == Plan.getScalarPreheader(); + // Bail out for tail-folding (middle block with a single successor to exit). + if (MiddleVPBB->getNumSuccessors() != 2 && !RequiresScalarEpilogue) + return nullptr; + + // All interleave groups in Plan can be narrowed for VFToOptimize. Split the + // original Plan into 2: a) a new clone which contains all VFs of Plan, except + // VFToOptimize, and b) the original Plan with VFToOptimize as single VF. + // TODO: Handle cases where only some interleave groups can be narrowed. + std::unique_ptr NewPlan; + if (size(Plan.vectorFactors()) != 1) { + NewPlan = std::unique_ptr(Plan.duplicate()); + Plan.setVF(*VFToOptimize); + NewPlan->removeVF(*VFToOptimize); + } + + // Convert InterleaveGroup \p R to a single VPWidenLoadRecipe. + SmallPtrSet NarrowedOps; + VPBasicBlock *Preheader = Plan.getVectorPreheader(); + // Narrow operation tree rooted at store groups. + for (auto *StoreGroup : StoreGroups) { + VPValue *Res = narrowInterleaveGroupOp(StoreGroup->getStoredValues(), + NarrowedOps, Preheader); + auto *SI = + cast(StoreGroup->getInterleaveGroup()->getInsertPos()); + VPBuilder(StoreGroup) + .createWidenStore(*SI, StoreGroup->getAddr(), Res, nullptr, + /*Consecutive=*/true, *StoreGroup, + StoreGroup->getDebugLoc()); + StoreGroup->eraseFromParent(); + } + + // Adjust induction to reflect that the transformed plan only processes one + // original iteration. + VPInstruction *CanIVInc = vputils::findCanonicalIVIncrement(Plan); + Type *CanIVTy = VectorLoop->getCanonicalIVType(); + VPBasicBlock *VectorPH = Plan.getVectorPreheader(); + VPBuilder PHBuilder(VectorPH, VectorPH->begin()); + + VPValue *UF = &Plan.getUF(); + VPValue *Step; + if (VFToOptimize->isScalable()) { + VPValue *VScale = + PHBuilder.createElementCount(CanIVTy, ElementCount::getScalable(1)); + Step = PHBuilder.createOverflowingOp(Instruction::Mul, {VScale, UF}, + {true, false}); + Plan.getVF().replaceAllUsesWith(VScale); + } else { + Step = UF; + Plan.getVF().replaceAllUsesWith(Plan.getConstantInt(CanIVTy, 1)); + } + // Materialize vector trip count with the narrowed step. + materializeVectorTripCount(Plan, VectorPH, /*TailByMasking=*/false, + RequiresScalarEpilogue, Step); + + CanIVInc->setOperand(1, Step); + Plan.getVFxUF().replaceAllUsesWith(Step); + + removeDeadRecipes(Plan); + assert(none_of(*VectorLoop->getEntryBasicBlock(), + IsaPred) && + "All VPVectorPointerRecipes should have been removed"); + return NewPlan; +}