[SCEV] Refactor isHighCostExpansionHelper

[samparker]

Summary

As discussed in D76434, to enable the cost of constants, the helper function has been reorganised:

• A struct has been introduced to hold SCEV operand information so that we know the user of the operand, as well as the operand index. The Worklist now uses instead instead of a bare SCEV.
• The costing of each SCEV, and collection of its operands, is now performed in a helper function.
• Some SCEVExprs have been modified to provide methods that allow us to iterate through their operands.

[samparker]

Updated isHighCostExpansion to use a SCEVOperand instead of a SCEV.

[samparker]

Moved AddRecExpr code into costAndCollectOperands.

[lebedev.ri]

Please separate SCEV changes into preparatory review.

llvm/include/llvm/Analysis/ScalarEvolutionExpressions.h

92 ↗

(On Diff #285952)

Doesn't this return empty range?

276 ↗

(On Diff #285952)

static_assert(sizeof(SmallVector<const SCEV*, 2>) == 4*sizeof(SCEV*), ""); holds, so this is a major bloat.
This should be std::array<>.

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

44

The variables could use doxygen comments

236

Why std::unique_ptr?

[samparker]

• Extracted the SCEV changes and rebased.
• Removed the use of unique_pointer.
• Moved the cost adjustment for NAry operands into the helper. These now also use the extra operand checks that were performed for AddRecs.

[lebedev.ri]

Please can you explain why the calculated costs changed for those X86 tests?

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

47

/// LLVM IR instruction opcode

[samparker]

Please can you explain why the calculated costs changed for those X86 tests?

These were introduced once the Add and Mul expressions were costed in the same manner as when we're calculating for an AddRec. I'll remove the change for this patch and add it back in a follow-up.

[samparker]

• Updated comment
• Reverted back to the original costing for Add and Mul.

[lebedev.ri]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2056–2064

I guess i'm starting getting confused again and thus stalling, but.

In D86050#2232843, @samparker wrote:

Please can you explain why the calculated costs changed for those X86 tests?

These were introduced once the Add and Mul expressions were costed in the same manner as when we're calculating for an AddRec.

Right. But why was that done in the first place? Was that intentional?
That doesn't look correct to me.

Plain Add and Mul are *not* identical to AddRec.

I'll remove the change for this patch and add it back in a follow-up.

[samparker]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2056–2064

It was intentional... Why wouldn't we want to check for constant zero for plain Adds and constant one for plain Muls?

[lebedev.ri]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2056–2064

Why wouldn't we want to check for constant zero for plain Adds and constant one for plain Muls?

Because they shouldn't be there, they should be canonicalized away during SCEV creation.

But my main question is, why did that result in *higher* cost, as it can be seen from the blocked transforms?

[samparker]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2056–2064

I don't know, it didn't make sense to me. Anyway, it's not in the patch anymore.

[lebedev.ri]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2056–2064

I don't know, it didn't make sense to me.

Honestly, recalling how many of these costmodelling patches were there recently,
i'm now somewhat worried that i didn't pay attention to them..

[samparker]

@lebedev.ri Are you okay with this now?

[lebedev.ri]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2189–2200

Move this back into case scAddRecExpr: please

2223–2225

I'm sorry, but i don't like this.
The old code would loudly fail should anyone introduce a new SCEV type and forget to update this,
new code will silently treat such a case as free.

[samparker]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2223–2225

Ah, good point.

[samparker]

Rebased and addressed comments.

[lebedev.ri]

Thanks, almost there i think.

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

239

I'd maybe suggest (-1, -1, Expr) to really differentiate this case.

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2189–2200

Please mark resolved comments as such.

2228

As a fixme, i suppose you might want to cost 1<<SC->getAPInt() constant?

2290–2291

for (auto I = enumerate(S->operands()))
  Worklist.emplace_back(Opc, I.index(), I.value());

2326–2339

This is very much not identical to the old code.
I think you want to do

// 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.

int Cost =
  costAndCollectOperands<SCEVUDivExpr>(WorkItem, TTI, CostKind, Worklist);

[nikic]

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

52

Nit: = nullptr here seems unnecessary, as the constructor always initializes it.

[samparker]

Thanks, addressed comments.

[samparker]

Removed nullptr assignment.

[lebedev.ri]

Alright, i think this is NFC now..
Thanks

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

241–243

potentially const SCEVOperand WorkItem =

411

Here and elsewhere: const SCEVOperand &WorkItem - it shouldn't be ever mutated.

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2187

Opcodes needs a good comment, even i'm currently having trouble understanding it.

[samparker]

Thanks! Added a comment and const before committing.

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

 bool isSafeToExpand(const SCEV *S, ScalarEvolution &SE);
 
 /// Return true if the given expression is safe to expand in the sense that
 /// all materialized values are defined and safe to speculate at the specified
 /// location and their operands are defined at this location.
 bool isSafeToExpandAt(const SCEV *S, const Instruction *InsertionPoint,
                       ScalarEvolution &SE);
 
+/// struct for holding enough information to help calculate the cost of the
+/// given SCEV when expanded into IR.
+struct SCEVOperand {
+  explicit SCEVOperand(unsigned Opc, int Idx, const SCEV *S) :
+    ParentOpcode(Opc), OperandIdx(Idx), S(S) { }
+  /// LLVM instruction opcode that uses the operand.
+  unsigned ParentOpcode;
+  /// The use index of an expanded instruction.
+  int OperandIdx;
+  /// The SCEV operand to be costed.
+  const SCEV* S = nullptr;
+};
+
 /// This class uses information about analyze scalars to rewrite expressions
 /// in canonical form.
 ///
 /// Clients should create an instance of this class when rewriting is needed,
 /// and destroy it when finished to allow the release of the associated
 /// memory.
 class SCEVExpander : public SCEVVisitor<SCEVExpander, Value *> {
   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;
+    SmallVector<SCEVOperand, 8> Worklist;
     SmallPtrSet<const SCEV *, 8> Processed;
     int BudgetRemaining = Budget * TargetTransformInfo::TCC_Basic;
-    Worklist.emplace_back(Expr);
+    Worklist.emplace_back(-1, -1, Expr);
     while (!Worklist.empty()) {
-      const SCEV *S = Worklist.pop_back_val();
-      if (isHighCostExpansionHelper(S, L, *At, BudgetRemaining, *TTI, Processed,
-                                    Worklist))
+      SCEVOperand WorkItem = Worklist.pop_back_val();
+      if (isHighCostExpansionHelper(WorkItem, 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
private:
 
   /// Insert code to directly compute the specified SCEV expression into the
   /// program. The code is inserted into the specified block. If \p
   /// Root is true, this indicates that \p SH is the top-level expression to
   /// expand passed from an external client call.
   Value *expandCodeForImpl(const SCEV *SH, Type *Ty, Instruction *I, bool Root);
 
   /// Recursive helper function for isHighCostExpansion.
-  bool isHighCostExpansionHelper(const SCEV *S, Loop *L, const Instruction &At,
-                                 int &BudgetRemaining,
-                                 const TargetTransformInfo &TTI,
-                                 SmallPtrSetImpl<const SCEV *> &Processed,
-                                 SmallVectorImpl<const SCEV *> &Worklist);
+  bool isHighCostExpansionHelper(
+    SCEVOperand &WorkItem, Loop *L, const Instruction &At, int &BudgetRemaining,
+    SmallVectorImpl<SCEVOperand> &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/Transforms/Utils/ScalarEvolutionExpander.cpp

if (!OrigPhiRef) {
       if (Phi->getType()->isIntegerTy() && TTI &&
           TTI->isTruncateFree(Phi->getType(), Phis.back()->getType())) {
         // This phi can be freely truncated to the narrowest phi type. Map the
         // truncated expression to it so it will be reused for narrow types.
         const SCEV *TruncExpr =
           SE.getTruncateExpr(SE.getSCEV(Phi), Phis.back()->getType());
         ExprToIVMap[TruncExpr] = Phi;
       }
       continue;
     }
 
     // Replacing a pointer phi with an integer phi or vice-versa doesn't make
     // sense.
     if (OrigPhiRef->getType()->isPointerTy() != Phi->getType()->isPointerTy())
       continue;
 
     if (BasicBlock *LatchBlock = L->getLoopLatch()) {
       Instruction *OrigInc = dyn_cast<Instruction>(
           OrigPhiRef->getIncomingValueForBlock(LatchBlock));
       Instruction *IsomorphicInc =
           dyn_cast<Instruction>(Phi->getIncomingValueForBlock(LatchBlock));
 
       if (OrigInc && IsomorphicInc) {
         // If this phi has the same width but is more canonical, replace the
         // original with it. As part of the "more canonical" determination,
SCEVExpander::getRelatedExistingExpansion(const SCEV *S, const Instruction *At,
 
   // There is potential to make this significantly smarter, but this simple
   // 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,
-    SmallVectorImpl<const SCEV *> &Worklist) {
-  if (BudgetRemaining < 0)
-    return true; // Already run out of budget, give up.
+template<typename T> static int costAndCollectOperands(
+  SCEVOperand &WorkItem, const TargetTransformInfo &TTI,
+  TargetTransformInfo::TargetCostKind CostKind,
+  SmallVectorImpl<SCEVOperand> &Worklist) {
+
+  const T *S = cast<T>(WorkItem.S);
+  int Cost = 0;
+  SmallVector<unsigned, 2> Opcodes;
+
+  auto CastCost = [&](unsigned Opcode) {
+    Opcodes.push_back(Opcode);
+    return TTI.getCastInstrCost(Opcode, S->getType(),
+                                S->getOperand(0)->getType(),
+                                TTI::CastContextHint::None, CostKind);
+  };
 
-  // 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.
+  auto ArithCost = [&](unsigned Opcode, unsigned NumRequired) {
+    Opcodes.push_back(Opcode);
+    return NumRequired *
+      TTI.getArithmeticInstrCost(Opcode, S->getType(), CostKind);
+  };
 
-  // If we can find an existing value for this scev available at the point "At"
-  // then consider the expression cheap.
-  if (getRelatedExistingExpansion(S, &At, L))
-    return false; // Consider the expression to be free.
+  auto CmpSelCost = [&](unsigned Opcode, unsigned NumRequired) {
+    Opcodes.push_back(Opcode);
+    Type *OpType = S->getOperand(0)->getType();
+    return NumRequired *
+      TTI.getCmpSelInstrCost(Opcode, OpType,
+                             CmpInst::makeCmpResultType(OpType), CostKind);
+  };
 
   switch (S->getSCEVType()) {
+  default:
+    llvm_unreachable("No other scev expressions possible.");
   case scUnknown:
   case scConstant:
-    return false; // Assume to be zero-cost.
-  }
-
-  TargetTransformInfo::TargetCostKind CostKind =
-    TargetTransformInfo::TCK_RecipThroughput;
-
-  if (auto *CastExpr = dyn_cast<SCEVCastExpr>(S)) {
-    unsigned Opcode;
-    switch (S->getSCEVType()) {
-    case scTruncate:
-      Opcode = Instruction::Trunc;
-      break;
-    case scZeroExtend:
-      Opcode = Instruction::ZExt;
-      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, /*Dst=*/S->getType(),
-        /*Src=*/Op->getType(), TTI::CastContextHint::None, CostKind);
-    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(),
-                                       CostKind);
-        // 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().
-        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(),
-                                   CostKind);
-    Worklist.insert(Worklist.end(), {UDivExpr->getLHS(), UDivExpr->getRHS()});
-    return false; // Will answer upon next entry into this function.
+    return 0;
+    Cost = CastCost(Instruction::Trunc);
+    Cost = CastCost(Instruction::ZExt);
+    Cost = CastCost(Instruction::SExt);
+    break;
+  case scUDivExpr: {
+    unsigned Opcode = Instruction::UDiv;
+    if (auto *SC = dyn_cast<SCEVConstant>(S->getOperand(1)))
+      if (SC->getAPInt().isPowerOf2())
+        Opcode = Instruction::LShr;
+    Cost = ArithCost(Opcode, 1);
+  case scAddExpr:
+    Cost = ArithCost(Instruction::Add, S->getNumOperands() - 1);
+    break;
+  case scMulExpr:
+    // TODO: this is a very pessimistic cost modelling for Mul,
+    // because of Bin Pow algorithm actually used by the expander,
+    // see SCEVExpander::visitMulExpr(), ExpandOpBinPowN().
+    Cost = ArithCost(Instruction::Mul, S->getNumOperands() - 1);
+    break;
+  case scSMaxExpr:
+  case scUMaxExpr:
+  case scSMinExpr:
+  case scUMinExpr: {
+    Cost += CmpSelCost(Instruction::ICmp, S->getNumOperands() - 1);
+    Cost += CmpSelCost(Instruction::Select, S->getNumOperands() - 1);
+    break;
   }
-
-  if (const auto *NAry = dyn_cast<SCEVAddRecExpr>(S)) {
-    Type *OpType = NAry->getType();
-
-    assert(NAry->getNumOperands() >= 2 &&
-           "Polynomial should be at least linear");
-
-    int AddCost =
-      TTI.getArithmeticInstrCost(Instruction::Add, OpType, CostKind);
-    int MulCost =
-      TTI.getArithmeticInstrCost(Instruction::Mul, OpType, CostKind);
-
+  case scAddRecExpr: {
     // In this polynominal, we may have some zero operands, and we shouldn't
     // really charge for those. So how many non-zero coeffients are there?
-    int NumTerms = llvm::count_if(NAry->operands(),
-                                  [](const SCEV *S) { return !S->isZero(); });
+    int NumTerms = llvm::count_if(S->operands(), [](const SCEV *Op) {
+                                    return !Op->isZero();
+                                  });
+
     assert(NumTerms >= 1 && "Polynominal should have at least one term.");
-    assert(!(*std::prev(NAry->operands().end()))->isZero() &&
+    assert(!(*std::prev(S->operands().end()))->isZero() &&
            "Last operand should not be zero");
 
-    // Much like with normal add expr, the polynominal will require
-    // one less addition than the number of it's terms.
-    BudgetRemaining -= AddCost * (NumTerms - 1);
-    if (BudgetRemaining < 0)
-      return true;
-
     // Ignoring constant term (operand 0), how many of the coeffients are u> 1?
     int NumNonZeroDegreeNonOneTerms =
-        llvm::count_if(make_range(std::next(NAry->op_begin()), NAry->op_end()),
-                       [](const SCEV *S) {
-                         auto *SConst = dyn_cast<SCEVConstant>(S);
-                         return !SConst || SConst->getAPInt().ugt(1);
-                       });
+      llvm::count_if(S->operands(), [](const SCEV *Op) {
+                      auto *SConst = dyn_cast<SCEVConstant>(Op);
+
+    // Much like with normal add expr, the polynominal will require
+    // one less addition than the number of it's terms.
+    int AddCost = ArithCost(Instruction::Add, NumTerms - 1);
     // Here, *each* one of those will require a multiplication.
-    BudgetRemaining -= MulCost * NumNonZeroDegreeNonOneTerms;
-    if (BudgetRemaining < 0)
-      return true;
+    int MulCost = ArithCost(Instruction::Mul, NumNonZeroDegreeNonOneTerms);
+    Cost = AddCost + MulCost;
 
     // What is the degree of this polynominal?
-    int PolyDegree = NAry->getNumOperands() - 1;
+    int PolyDegree = S->getNumOperands() - 1;
     assert(PolyDegree >= 1 && "Should be at least affine.");
 
     // The final term will be:
     //   Op_{PolyDegree} * x ^ {PolyDegree}
     // Where  x ^ {PolyDegree}  will again require PolyDegree-1 mul operations.
     // 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;
+    Cost += MulCost * (PolyDegree - 1);
+  }
+  }
 
-    // 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?
-  }
+  for (unsigned Opc : Opcodes)
+    for (auto I : enumerate(S->operands()))
+      Worklist.emplace_back(Opc, I.index(), I.value());
+  return Cost;
 
-  if (const SCEVNAryExpr *NAry = dyn_cast<SCEVNAryExpr>(S)) {
-    Type *OpType = NAry->getType();
+bool SCEVExpander::isHighCostExpansionHelper(
+    SCEVOperand &WorkItem, Loop *L, const Instruction &At, int &BudgetRemaining,
+    const TargetTransformInfo &TTI, SmallPtrSetImpl<const SCEV *> &Processed,
+    SmallVectorImpl<SCEVOperand> &Worklist) {
+  if (BudgetRemaining < 0)
+    return true; // Already run out of budget, give up.
+
+  const SCEV *S = WorkItem.S;
+  // 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"
+  // then consider the expression cheap.
+  if (getRelatedExistingExpansion(S, &At, L))
+    return false; // Consider the expression to be free.
 
-    int PairCost;
-    switch (S->getSCEVType()) {
-    case scAddExpr:
-      PairCost =
-        TTI.getArithmeticInstrCost(Instruction::Add, OpType, CostKind);
-      break;
-    case scMulExpr:
-      // TODO: this is a very pessimistic cost modelling for Mul,
-      // because of Bin Pow algorithm actually used by the expander,
-      // see SCEVExpander::visitMulExpr(), ExpandOpBinPowN().
-      PairCost =
-        TTI.getArithmeticInstrCost(Instruction::Mul, OpType, CostKind);
-      break;
-    case scSMaxExpr:
-    case scUMaxExpr:
-    case scSMinExpr:
-    case scUMinExpr:
-      PairCost = TTI.getCmpSelInstrCost(Instruction::ICmp, OpType,
-                                        CmpInst::makeCmpResultType(OpType),
-                                        CostKind) +
-                 TTI.getCmpSelInstrCost(Instruction::Select, OpType,
-                                        CmpInst::makeCmpResultType(OpType),
-                                        CostKind);
-      break;
-    default:
-      llvm_unreachable("There are no other variants here.");
-    }
+  case scUnknown:
+  case scConstant:
+    return false; // Assume to be zero-cost.
 
+  TargetTransformInfo::TargetCostKind CostKind =
+    TargetTransformInfo::TCK_RecipThroughput;
+
+  if (isa<SCEVCastExpr>(S)) {
+    int Cost =
+      costAndCollectOperands<SCEVCastExpr>(WorkItem, TTI, CostKind, Worklist);
+    BudgetRemaining -= Cost;
+    return false; // Will answer upon next entry into this function.
+  } else if (isa<SCEVUDivExpr>(S)) {
+    // 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.
+
+    int Cost =
+      costAndCollectOperands<SCEVUDivExpr>(WorkItem, TTI, CostKind, Worklist);
+    // Need to count the cost of this UDiv.
+    BudgetRemaining -= Cost;
+    return false; // Will answer upon next entry into this function.
+  } else if (const SCEVNAryExpr *NAry = dyn_cast<SCEVNAryExpr>(S)) {
     assert(NAry->getNumOperands() > 1 &&
            "Nary expr should have more than 1 operand.");
     // 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;
-
-    // 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?
-  }
-
+    int Cost =
+      costAndCollectOperands<SCEVNAryExpr>(WorkItem, TTI, CostKind, Worklist);
+    BudgetRemaining -= Cost;
+    return BudgetRemaining < 0;
+  } else if (const auto *NAry = dyn_cast<SCEVAddRecExpr>(S)) {
+    assert(NAry->getNumOperands() >= 2 &&
+           "Polynomial should be at least linear");
+    BudgetRemaining -= costAndCollectOperands<SCEVAddRecExpr>(
+      WorkItem, TTI, CostKind, Worklist);
+    return BudgetRemaining < 0;
+  } else
     llvm_unreachable("No other scev expressions possible.");
 }
 
 Value *SCEVExpander::expandCodeForPredicate(const SCEVPredicate *Pred,
                                             Instruction *IP) {
   assert(IP);
   switch (Pred->getKind()) {
   case SCEVPredicate::P_Union:
     return expandUnionPredicate(cast<SCEVUnionPredicate>(Pred), IP);