diff --git a/llvm/include/llvm/Analysis/ScalarEvolutionExpander.h b/llvm/include/llvm/Analysis/ScalarEvolutionExpander.h --- a/llvm/include/llvm/Analysis/ScalarEvolutionExpander.h +++ b/llvm/include/llvm/Analysis/ScalarEvolutionExpander.h @@ -16,6 +16,7 @@ #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/Optional.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/Analysis/ScalarEvolutionExpressions.h" #include "llvm/Analysis/ScalarEvolutionNormalization.h" #include "llvm/Analysis/TargetFolder.h" @@ -186,10 +187,18 @@ assert(At && "This function requires At instruction to be provided."); if (!TTI) // In assert-less builds, avoid crashing return true; // by always claiming to be high-cost. + SmallVector Worklist; SmallPtrSet Processed; int BudgetRemaining = Budget * TargetTransformInfo::TCC_Basic; - return isHighCostExpansionHelper(Expr, L, *At, BudgetRemaining, *TTI, - Processed); + Worklist.emplace_back(Expr); + while (!Worklist.empty()) { + const SCEV *S = Worklist.pop_back_val(); + if (isHighCostExpansionHelper(S, L, *At, BudgetRemaining, *TTI, + Processed, Worklist)) + return true; + } + assert(BudgetRemaining >= 0 && "Should have returned from inner loop."); + return false; } /// This method returns the canonical induction variable of the specified @@ -334,7 +343,8 @@ bool isHighCostExpansionHelper(const SCEV *S, Loop *L, const Instruction &At, int &BudgetRemaining, const TargetTransformInfo &TTI, - SmallPtrSetImpl &Processed); + SmallPtrSetImpl &Processed, + SmallVectorImpl &Worklist); /// Insert the specified binary operator, doing a small amount of work to /// avoid inserting an obviously redundant operation, and hoisting to an diff --git a/llvm/lib/Analysis/ScalarEvolutionExpander.cpp b/llvm/lib/Analysis/ScalarEvolutionExpander.cpp --- a/llvm/lib/Analysis/ScalarEvolutionExpander.cpp +++ b/llvm/lib/Analysis/ScalarEvolutionExpander.cpp @@ -2137,7 +2137,8 @@ bool SCEVExpander::isHighCostExpansionHelper( const SCEV *S, Loop *L, const Instruction &At, int &BudgetRemaining, - const TargetTransformInfo &TTI, SmallPtrSetImpl &Processed) { + const TargetTransformInfo &TTI, SmallPtrSetImpl &Processed, + SmallVectorImpl &Worklist) { if (BudgetRemaining < 0) return true; // Already run out of budget, give up. @@ -2172,13 +2173,12 @@ llvm_unreachable("There are no other cast types."); } const SCEV *Op = CastExpr->getOperand(); - BudgetRemaining -= - TTI.getCastInstrCost(Opcode, S->getType(), Op->getType()); - return isHighCostExpansionHelper(Op, L, At, BudgetRemaining, TTI, - Processed); + BudgetRemaining -= TTI.getCastInstrCost(Opcode, /*Dst=*/S->getType(), + /*Src=*/Op->getType()); + Worklist.emplace_back(Op); + return false; // Will answer upon next entry into this function. } - if (auto *UDivExpr = dyn_cast(S)) { // If the divisor is a power of two count this as a logical right-shift. if (auto *SC = dyn_cast(UDivExpr->getRHS())) { @@ -2188,8 +2188,8 @@ // Note that we don't count the cost of RHS, because it is a constant, // and we consider those to be free. But if that changes, we would need // to log2() it first before calling isHighCostExpansionHelper(). - return isHighCostExpansionHelper(UDivExpr->getLHS(), L, At, - BudgetRemaining, TTI, Processed); + Worklist.emplace_back(UDivExpr->getLHS()); + return false; // Will answer upon next entry into this function. } } @@ -2208,10 +2208,8 @@ // Need to count the cost of this UDiv. BudgetRemaining -= TTI.getArithmeticInstrCost(Instruction::UDiv, S->getType()); - return isHighCostExpansionHelper(UDivExpr->getLHS(), L, At, BudgetRemaining, - TTI, Processed) || - isHighCostExpansionHelper(UDivExpr->getRHS(), L, At, BudgetRemaining, - TTI, Processed); + Worklist.insert(Worklist.end(), {UDivExpr->getLHS(), UDivExpr->getRHS()}); + return false; // Will answer upon next entry into this function. } if (const auto *NAry = dyn_cast(S)) { @@ -2264,12 +2262,9 @@ return true; // And finally, the operands themselves should fit within the budget. - for (const SCEV *Op : NAry->operands()) { - if (isHighCostExpansionHelper(Op, L, At, BudgetRemaining, TTI, Processed)) - return true; - } - - return BudgetRemaining < 0; + Worklist.insert(Worklist.end(), NAry->operands().begin(), + NAry->operands().end()); + return false; // So far so good, though ops may be too costly? } if (const SCEVNAryExpr *NAry = dyn_cast(S)) { @@ -2301,15 +2296,16 @@ assert(NAry->getNumOperands() > 1 && "Nary expr should have more than 1 operand."); - for (const SCEV *Op : NAry->operands()) { - if (isHighCostExpansionHelper(Op, L, At, BudgetRemaining, TTI, Processed)) - return true; - if (Op == *NAry->op_begin()) - continue; - BudgetRemaining -= PairCost; - } + // The simple nary expr will require one less op (or pair of ops) + // than the number of it's terms. + BudgetRemaining -= PairCost * (NAry->getNumOperands() - 1); + if (BudgetRemaining < 0) + return true; - return BudgetRemaining < 0; + // And finally, the operands themselves should fit within the budget. + Worklist.insert(Worklist.end(), NAry->operands().begin(), + NAry->operands().end()); + return false; // So far so good, though ops may be too costly? } llvm_unreachable("No other scev expressions possible.");