[NFCI][SCEV] Avoid recursion in SCEVExpander::isHighCostExpansion*()

[lebedev.ri]

Summary

As noted in PR45201,
PR10090 SCEV doesn't always avoid
recursive algorithms, and that causes issues with large expression depths and/or smaller stack sizes.

In SCEVExpander::isHighCostExpansion*() case, the refactoring to avoid recursion
is rather idiomatic. We simply need to place the root expr into a vector,
and iterate over vector elements accounting for the cost of each one,
and add new exprs at the end of the vector, thus achieving FIFO breadth-first traversal.

[mkazantsev]

llvm/include/llvm/Analysis/ScalarEvolutionExpander.h

195

I'm not getting why do you need a worklist. On the 1st iteration of this loop, you put one value to the WorkList and then remove it. isHighCostExpansionHelper always puts one element to the worklist and then returns. Therefore the size of your FIFO is either 0 or 1. Do you really need to use deque for it?

[mkazantsev]

LGTM

llvm/include/llvm/Analysis/ScalarEvolutionExpander.h

195

Sorry, I misread the code. It may put multiple elements into it.

[mkazantsev]

llvm/include/llvm/Analysis/ScalarEvolutionExpander.h

27

Can we go with queue?

[mkazantsev]

As a suggestion: since this helper doesn't create new code, it doesn't seem to be important whether it is FIFO or LIFO. If that is so, I'd suggest using SmallVector and LIFO algorithm instead. It should have better memory footprint. OK if it does as follow-up.

[lebedev.ri]

In D76273#1928590, @mkazantsev wrote:

LGTM

Thank you for the review!

In D76273#1928615, @mkazantsev wrote:

As a suggestion: since this helper doesn't create new code, it doesn't seem to be important whether it is FIFO or LIFO. If that is so, I'd suggest using SmallVector and LIFO algorithm instead. It should have better memory footprint. OK if it does as follow-up.

I agree the order doesn't matter for correctness.
Currently we seem to be doing depth-first traversal (since we recurse before looking at other operands).
I don't have any evidence which approach is better in the sense that it would exhaust the budget faster,
so i'll go back to simple smallvector with back-insertion/extraction.

llvm/include/llvm/Analysis/ScalarEvolutionExpander.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_SCALAREVOLUTIONEXPANDER_H
 #define LLVM_ANALYSIS_SCALAREVOLUTIONEXPANDER_H
 
 #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"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Support/CommandLine.h"
 
 namespace llvm {
   extern cl::opt<unsigned> SCEVCheapExpansionBudget;
 
   /// Return true if the given expression is safe to expand in the sense that
   /// all materialized values are safe to speculate anywhere their operands are
   /// defined.
   bool isSafeToExpand(const SCEV *S, ScalarEvolution &SE);
 
#endif
     /// a less pessimistic cost estimation.
     bool isHighCostExpansion(const SCEV *Expr, Loop *L, unsigned Budget,
                              const TargetTransformInfo *TTI,
                              const Instruction *At) {
       assert(TTI && "This function requires TTI to be provided.");
       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<const SCEV *, 8> Worklist;
       SmallPtrSet<const SCEV *, 8> 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
     /// type for the specified loop (inserting one if there is none).  A
     /// canonical induction variable starts at zero and steps by one on each
     /// iteration.
     PHINode *getOrInsertCanonicalInductionVariable(const Loop *L, Type *Ty);
 
#endif
 
   private:
     LLVMContext &getContext() const { return SE.getContext(); }
 
     /// Recursive helper function for isHighCostExpansion.
     bool isHighCostExpansionHelper(const SCEV *S, Loop *L,
                                    const Instruction &At, int &BudgetRemaining,
                                    const TargetTransformInfo &TTI,
-                                   SmallPtrSetImpl<const SCEV *> &Processed);
+                                   SmallPtrSetImpl<const SCEV *> &Processed,
+                                   SmallVectorImpl<const SCEV *> &Worklist);
 
     /// Insert the specified binary operator, doing a small amount of work to
     /// avoid inserting an obviously redundant operation, and hoisting to an
     /// outer loop when the opportunity is there and it is safe.
     Value *InsertBinop(Instruction::BinaryOps Opcode, Value *LHS, Value *RHS,
                        SCEV::NoWrapFlags Flags, bool IsSafeToHoist);
 
     /// Arrange for there to be a cast of V to Ty at IP, reusing an existing

llvm/lib/Analysis/ScalarEvolutionExpander.cpp

SCEVExpander::getRelatedExistingExpansion(const SCEV *S, const Instruction *At,
   // heuristic already gets some interesting cases.
 
   // Can not find suitable value.
   return None;
 }
 
 bool SCEVExpander::isHighCostExpansionHelper(
     const SCEV *S, Loop *L, const Instruction &At, int &BudgetRemaining,
-    const TargetTransformInfo &TTI, SmallPtrSetImpl<const SCEV *> &Processed) {
+    const TargetTransformInfo &TTI, SmallPtrSetImpl<const SCEV *> &Processed,
+    SmallVectorImpl<const SCEV *> &Worklist) {
   if (BudgetRemaining < 0)
     return true; // Already run out of budget, give up.
 
   // Was the cost of expansion of this expression already accounted for?
   if (!Processed.insert(S).second)
     return false; // We have already accounted for this expression.
 
   // If we can find an existing value for this scev available at the point "At"
case scZeroExtend:
       break;
     case scSignExtend:
       Opcode = Instruction::SExt;
       break;
     default:
       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<SCEVUDivExpr>(S)) {
     // If the divisor is a power of two count this as a logical right-shift.
     if (auto *SC = dyn_cast<SCEVConstant>(UDivExpr->getRHS())) {
       if (SC->getAPInt().isPowerOf2()) {
         BudgetRemaining -=
             TTI.getArithmeticInstrCost(Instruction::LShr, S->getType());
         // 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.
       }
     }
 
     // UDivExpr is very likely a UDiv that ScalarEvolution's HowFarToZero or
     // HowManyLessThans produced to compute a precise expression, rather than a
     // UDiv from the user's code. If we can't find a UDiv in the code with some
     // simple searching, we need to account for it's cost.
 
     // At the beginning of this function we already tried to find existing
     // value for plain 'S'. Now try to lookup 'S + 1' since it is common
     // pattern involving division. This is just a simple search heuristic.
     if (getRelatedExistingExpansion(
             SE.getAddExpr(S, SE.getConstant(S->getType(), 1)), &At, L))
       return false; // Consider it to be free.
 
     // 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<SCEVAddRecExpr>(S)) {
     Type *OpType = NAry->getType();
 
     assert(NAry->getNumOperands() >= 2 &&
            "Polynomial should be at least linear");
 
if (const auto *NAry = dyn_cast<SCEVAddRecExpr>(S)) {
     // Note that  x ^ {PolyDegree} = x * x ^ {PolyDegree-1}  so charging for
     // x ^ {PolyDegree}  will give us  x ^ {2} .. x ^ {PolyDegree-1}  for free.
     // FIXME: this is conservatively correct, but might be overly pessimistic.
     BudgetRemaining -= MulCost * (PolyDegree - 1);
     if (BudgetRemaining < 0)
       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<SCEVNAryExpr>(S)) {
     Type *OpType = NAry->getType();
 
     int PairCost;
     switch (S->getSCEVType()) {
     case scAddExpr:
case scUMinExpr:
                                         CmpInst::makeCmpResultType(OpType));
       break;
     default:
       llvm_unreachable("There are no other variants here.");
     }
 
     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 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.");
 }
 
 Value *SCEVExpander::expandCodeForPredicate(const SCEVPredicate *Pred,
                                             Instruction *IP) {
   assert(IP);