diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp index 5b556058cc762..d13770a35c108 100644 --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -7303,14 +7303,34 @@ LoopVectorizationPlanner::precomputeCosts(VPlan &Plan, ElementCount VF, // The legacy cost model has special logic to compute the cost of in-loop // reductions, which may be smaller than the sum of all instructions involved - // in the reduction. + // in the reduction. For AnyOf reductions, VPlan codegen may remove the select + // which the legacy cost model uses to assign cost. Pre-compute their costs + // for now. // TODO: Switch to costing based on VPlan once the logic has been ported. for (const auto &[RedPhi, RdxDesc] : Legal->getReductionVars()) { if (ForceTargetInstructionCost.getNumOccurrences()) continue; - if (!CM.isInLoopReduction(RedPhi)) + if (!CM.isInLoopReduction(RedPhi) && + !RecurrenceDescriptor::isAnyOfRecurrenceKind( + RdxDesc.getRecurrenceKind())) + continue; + + // AnyOf reduction codegen may remove the select. To match the legacy cost + // model, pre-compute the cost for AnyOf reductions here. + if (RecurrenceDescriptor::isAnyOfRecurrenceKind( + RdxDesc.getRecurrenceKind())) { + auto *Select = cast(*find_if( + RedPhi->users(), [](User *U) { return isa(U); })); + assert(!CostCtx.SkipCostComputation.contains(Select) && + "reduction op visited multiple times"); + CostCtx.SkipCostComputation.insert(Select); + auto ReductionCost = CostCtx.getLegacyCost(Select, VF); + LLVM_DEBUG(dbgs() << "Cost of " << ReductionCost << " for VF " << VF + << ":\n any-of reduction " << *Select << "\n"); + Cost += ReductionCost; continue; + } const auto &ChainOps = RdxDesc.getReductionOpChain(RedPhi, OrigLoop); SetVector ChainOpsAndOperands(ChainOps.begin(),