[SCEVExpander] Migrate costAndCollectOperands to use InstructionCost.

[sdesmalen]

Summary

This patch changes costAndCollectOperands to use InstructionCost for
accumulated cost values.

isHighCostExpansion will return true if the cost is too high or invalid.

[CarolineConcatto]

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

250

Maybe put this assertion before the loop? Is there any reason to check the Cost if BudgetRemaining.isValid() is not valid?
Or maybe you can remove this assert from here, because if we check isHighCostExpansionHelper this test is done in line 2346.

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2188

Why you did not changed Cost to be InstructionCost?

[lebedev.ri]

I think this is conflating use-cases.
I'm not sure the entire InstructionCost should be used as a variable to track full/remaining cost over a several instructions.
Perhaps there should be a[nother] abstraction for that?

[ctetreau]

In D92238#2452418, @lebedev.ri wrote:

I think this is conflating use-cases.
I'm not sure the entire InstructionCost should be used as a variable to track full/remaining cost over a several instructions.
Perhaps there should be a[nother] abstraction for that?

Really, we were already doing this; we just didn't have a special cost type. We could leave the cost as an integer, but then the code would just become way more verbose (unwrapping every InstructionCost after computing them). Using the metaphor of money, a cost is a quantity of money. I think it's reasonable to subtract money from a budget. I suppose we could accumulate a running cost, and compare it with the budget that remains an integer, but that would just introduce another variable for (IMO) no good reason.

[sdesmalen]

• Simplified *BudgetRemaining.getValue() < 0 to BudgetRemaining < 0.
• Changed int Cost; to InstructionCost Cost.
• Rebased patch.

[sdesmalen]

In D92238#2480072, @ctetreau wrote:

In D92238#2452418, @lebedev.ri wrote:

I think this is conflating use-cases.
I'm not sure the entire InstructionCost should be used as a variable to track full/remaining cost over a several instructions.
Perhaps there should be a[nother] abstraction for that?

Really, we were already doing this; we just didn't have a special cost type. We could leave the cost as an integer, but then the code would just become way more verbose (unwrapping every InstructionCost after computing them). Using the metaphor of money, a cost is a quantity of money. I think it's reasonable to subtract money from a budget. I suppose we could accumulate a running cost, and compare it with the budget that remains an integer, but that would just introduce another variable for (IMO) no good reason.

Like @ctetreau I don't see why the algorithm shouldn't be using InstructionCost. The class has enough capabilities that any invalid state now propagates quite naturally without too big changes (InstructionCost behaves mostly like a scalar int that carries extra state).
I've simplified the use of *BudgetRemaining.getValue() < 0 to BudgetRemaining < 0, which hopefully makes the patch more palatable?

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

250

Or maybe you can remove this assert from here, because if we check isHighCostExpansionHelper this test is done in line 2346.

That's true, but the assert here is only to verify the code/values are sane.

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2188

Good catch!

[sdesmalen]

Gentle ping

[RKSimon]

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

20

the headers need sorting - clang-format this

[sdesmalen]

• Reordered includes.
• Changed two more variables to be of type InstructionCost instead of int.

[sdesmalen]

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

20

Fixed, thanks for pointing out!

[CarolineConcatto]

It look to me now.

[lebedev.ri]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2328–2329

What happens without the isValid() check?

[sdesmalen]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2328–2329

Without the isValid() check it defaults to the total ordering for InstructionCost where all valid costs < Invalid.
This means that BudgetRemaining < 0 would evaluate to false, where instead we want to return true from this function to signal this is a high cost expansion.

[lebedev.ri]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2328–2329

Uh oh.

So if every single check now needs to not be forgotten to be prefixed with isValid(),
doesn't that imply that the default is wrong, and it should be Invalid < all valid costs instead?
I'm guessing not, because i guess 0 > BudgetRemaining will then break.

So alternatively, why not at least make it easy to detect this issue,
and instead assert within the InstructionCost that costs must be valid to do such comparison?

[sdesmalen]

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2328–2329

In D91174 @ctetreau argued in favour of having a total ordering over asserting that the costs must be valid (https://reviews.llvm.org/D91174#2422681).

Probably the most prominent use-case where an InstructionCost is used in comparisons is:

if (CostX < CostY)
  // replace with X

Here the total-ordering is desirable, i.e. if CostX is invalid, don't replace with X. Otherwise, replace if CostX is smaller than CostY or CostY is invalid. We want to avoid having to rewrite these cases to:

if ((CostX.isValid() && !CostY.isValid()) ||
     (CostX.isValid() && CostY.isValid() && CostX < Cost))
  // replace with X

// All cases would need to be covered if CostX < CostY would *assert*
// they most both be valid.

I think that in this particular case - working with cost as a budget - there is extra confusion because the code is comparing directly with a constant, which makes it easy to forget that this is not a 'normal' comparison, but is in fact a comparison between InstructionCosts which has total ordering semantics.

What I could do is:

• Mark InstructionCost::operator<(int) operator as delete, so that the above statement no longer compiles. This means only comparisons between InstructionCosts are supported.
• Add a new method bool InstructionCost::isNegativeValue() which returns true if the value is Valid and less than 0.
• Update this patch to use isNegativeValue() instead.

Would that be a suitable solution?

[sdesmalen]

Changed patch to use a new InstructionBudget class.

[sdesmalen]

I have come to the realisation that an abstraction around InstructionCost specific for budgeting is actually quite useful.

@lebedev.ri are you happy with these changes?

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

2328–2329

Last week I've been working to update multiple passes and interfaces to work on InstructionCost, and I found that for most, if not all, other cases where InstructionCost is used the total ordering is sufficient and sufficiently clear. Therefore I don't think my above suggestion makes sense, so please ignore it.

[lebedev.ri]

I guess this is as good as we'll get for the first approximation.
Thanks.

[sdesmalen]

In D92238#2519531, @lebedev.ri wrote:

I guess this is as good as we'll get for the first approximation.
Thanks.

Thanks for your feedback!

[ctetreau]

llvm/include/llvm/Support/InstructionCost.h

238 ↗

(On Diff #318949)

I think this new class is uneccesary. Can't we just write 0 > Cost for any case where we would use Cost.isExceeded()?

[ctetreau]

I'd like you to reconsider adding InstructionBudget. I think this new class is just trying to paper over people's misunderstanding of how InstructionCost works. This wouldn't be a huge issue, except now we're going to have to add a million redundant operator definitions to it ("why can't I multiply budgets?") so it's going to become a maintenance burden.

[lebedev.ri]

Then please fix the InstrouctionCost comparison operators to do the right thing to avoid most of the changes in this diff in the first place.

[ctetreau]

In D92238#2520415, @lebedev.ri wrote:

Then please fix the InstrouctionCost comparison operators to do the right thing to avoid most of the changes in this diff in the first place.

InstructionCost does do the right thing. InstructionCost is isomorphic to Optional<int>, except in 99% of cases, you want an invalid cost to be considered greater than any other cost. Does int do the wrong thing because std::numeric_limits<int>::max() + 1 < std::numeric_limits<int>::max()?

I don't think that Cost < 0 || Cost.isInvalid() is an unreasonable change to request. Especially considering all the other code that uses InstructionCost that is written much more naturally due to the total ordering. An isExceeded() free function can be added to this file to ensure it's done correctly.

llvm/include/llvm/Support/InstructionCost.h

238 ↗

(On Diff #318949)

I guess this doesn't work either. You have to do the isInvalid() call

[lebedev.ri]

Guess this is stuck, then.

[ctetreau]

llvm/include/llvm/Support/InstructionCost.h

238 ↗

(On Diff #318949)

Things that are logically "budgets" can be initialized with negative numbers, and have costs added to them, rather than subtracted. Then Cost > 0 can be used and will work correctly.

[sdesmalen]

An attempt to meet in the middle:

• No need for InstructionBudget class.
• No need for any changes to return BudgetRemaining < 0 in isHighCostExpansionHelper.

Instead, the check for BudgetRemaining.isValid() is now done at the end in isHighCostExpansion. If the helper function found the budget has exceeded, or if the resulting budget has been invalidated, the expansion is considered high cost.

Is this a step in the right direction?

[lebedev.ri]

In D92238#2521055, @sdesmalen wrote:

An attempt to meet in the middle:

• No need for InstructionBudget class.
• No need for any changes to return BudgetRemaining < 0 in isHighCostExpansionHelper.

Instead, the check for BudgetRemaining.isValid() is now done at the end in isHighCostExpansion. If the helper function found the budget has exceeded, or if the resulting budget has been invalidated, the expansion is considered high cost.

Is this a step in the right direction?

Won't we now not stop as soon as the budget is invalidated?
Perhaps if InstructionCost::operator< can't be fixed, it needs to be ripped out.
I'm sorry, i don't have useful feedback other than "the new abstraction is bad".

[ctetreau]

In D92238#2521095, @lebedev.ri wrote:

Won't we now not stop as soon as the budget is invalidated?
Perhaps if InstructionCost::operator< can't be fixed, it needs to be ripped out.
I'm sorry, i don't have useful feedback other than "the new abstraction is bad".

I agree that it's strictly worse in that the early return is not triggered if anything returns an invalid cost. I disagree that InstructionCost::operator< needs to be fixed or that "the new abstraction is bad".

The new abstraction has a well-defined semantics that is useful in the common case of deciding which of two costs is preferable. Throughout the codebase, we decide if thingA should be done over thingB by asking if the cost of thingA is less than the cost of thingB. If one of those things has an invalid cost, then it makes sense to consider the thing with a valid cost as less than it. If they are both invalid, then it doesn't really matter which one we pick.

I'm curious: what's the objection to having the budget be a negative number, and adding costs instead of subtracting? We could also accumulate costs in a separate variable and compare Cost < Budget, which will also do the right thing in the face of an invalid cost infecting the running total. In this case, Budget doesn't even need to be an InstructionCost.

[sdesmalen]

In D92238#2521095, @lebedev.ri wrote:

Won't we now not stop as soon as the budget is invalidated?

That should be no different than before, isHighCostExpansionHelper is not recursive. The function that calls it, isHighCostExpansion, will still stop immediately when the cost is invalidated.

Perhaps if InstructionCost::operator< can't be fixed, it needs to be ripped out.
I'm sorry, i don't have useful feedback other than "the new abstraction is bad".

The part of the new abstraction that I believe may require more thought is the signedness of InstructionCost. If Cost would be unsigned, the code is already expected to deal with wrapping values due to underflow and InstructionCost::operator< would work as expected. There would be no need to add additional checks for validity. Another possibility would be to return Invalid when there is an underflow, because if there is a negative cost at any point this is probably a bug. That would still mean changing the compare "Budget < 0" in this file, but there are other ways to write it such as @ctetreau's suggests.

Is this something to consider changing at some point?

I'm curious: what's the objection to having the budget be a negative number, and adding costs instead of subtracting?

Personally I find this a bit counter-intuitive, because a budget is always a positive value that is subtracted from.

We could also accumulate costs in a separate variable and compare Cost < Budget, which will also do the right thing in the face of an invalid cost infecting the running total. In this case, Budget doesn't even need to be an InstructionCost.

Yes, I quite like that suggestion.

[sdesmalen]

• Refactored patch to use incrementing Cost value, to be compared with Budget.

There probably isn't much I can do in this patch to take away all @lebedev.ri 's concerns regarding the comparison operators, but hopefully this patch is now written in such a way that it is clear and acceptable to everyone. I'd be happy to follow this up with possible changes to the nature of InstructionCost itself, like I suggested earlier, if that is helpful.

[ctetreau]

Aside from the change to the assert, this version looks good to me.

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

241–251

If the scale factor to the budget changes, then this assert will be wrong. If TCC_Basic is not 1, then the assert is already wrong.

[sdesmalen]

Fixed issue with assert, now uses ScaledBudget.

[sdesmalen]

Hi @lebedev.ri, when we introduced InstructionCost, we left open the possibility of changing the class based on new insights on how costs are used in TargetTransformInfo and IR passes.

In this patch you pointed out a valid concern. It is indeed unexpected to need to have to check isValid() when wanting to know if a budget has gone negative after subtracting a cost, because even though InstructionCost is signed, subtracting Invalid currently results in Invalid, not -Invalid.

I had a chat with @ctetreau offline about changing InstructionCost by making the cost unsigned (i.e. a cost can only ever be positive), or possibly restricting the class's operators that it can only ever be incrementing, so that sorting Invalid after a valid cost value would no longer be strange/unexpected. I think this would address your fundamental concerns with the abstraction. There is some investigation required on what exactly the right solution would be, but I'd be happy to follow this up separately.

For the use-case in this patch (SCEVExpander), I think we now have a suitable solution that makes sense and is future proof for whatever changes we may make to InstructionCost. If you have no strong objections to the current approach (separate Cost and Budget variables), are you happy to accept the patch so that we can move this forward?

[ctetreau]

This version LGTM

[RKSimon]

@lebedev.ri ping

[lebedev.ri]

In D92238#2541226, @RKSimon wrote:

@lebedev.ri ping

I wasn't reachable since 02.02.2021, ~back now.

Let me just complain about the interface next time i use it an obviously fall into it's traps.

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

 
 #ifndef LLVM_TRANSFORMS_UTILS_SCALAREVOLUTIONEXPANDER_H
 #define LLVM_TRANSFORMS_UTILS_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"
+#include "llvm/Support/InstructionCost.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);
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<SCEVOperand, 8> Worklist;
     SmallPtrSet<const SCEV *, 8> Processed;
-    int BudgetRemaining = Budget * TargetTransformInfo::TCC_Basic;
+    InstructionCost Cost = 0;
+    unsigned ScaledBudget = Budget * TargetTransformInfo::TCC_Basic;
     Worklist.emplace_back(-1, -1, Expr);
     while (!Worklist.empty()) {
       const SCEVOperand WorkItem = Worklist.pop_back_val();
-      if (isHighCostExpansionHelper(WorkItem, L, *At, BudgetRemaining,
-                                    *TTI, Processed, Worklist))
+      if (isHighCostExpansionHelper(WorkItem, L, *At, Cost, ScaledBudget, *TTI,
+                                    Processed, Worklist))
         return true;
     }
-    assert(BudgetRemaining >= 0 && "Should have returned from inner loop.");
+    assert(Cost <= ScaledBudget && "Should have returned from inner loop.");
     return false;
   }
 
   /// Return the induction variable increment's IV operand.
   Instruction *getIVIncOperand(Instruction *IncV, Instruction *InsertPos,
                                bool allowScale);
 
   /// Utility for hoisting an IV increment.
   bool hoistIVInc(Instruction *IncV, Instruction *InsertPos);
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 SCEVOperand &WorkItem, Loop *L, const Instruction &At,
-    int &BudgetRemaining, const TargetTransformInfo &TTI,
+  bool isHighCostExpansionHelper(const SCEVOperand &WorkItem, Loop *L,
+                                 const Instruction &At, InstructionCost &Cost,
+                                 unsigned Budget,
+                                 const TargetTransformInfo &TTI,
                                  SmallPtrSetImpl<const SCEV *> &Processed,
                                  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

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;
 }
 
-template<typename T> static int costAndCollectOperands(
+template<typename T> static InstructionCost costAndCollectOperands(
   const SCEVOperand &WorkItem, const TargetTransformInfo &TTI,
   TargetTransformInfo::TargetCostKind CostKind,
   SmallVectorImpl<SCEVOperand> &Worklist) {
 
   const T *S = cast<T>(WorkItem.S);
-  int Cost = 0;
+  InstructionCost Cost = 0;
   // Object to help map SCEV operands to expanded IR instructions.
   struct OperationIndices {
     OperationIndices(unsigned Opc, size_t min, size_t max) :
       Opcode(Opc), MinIdx(min), MaxIdx(max) { }
     unsigned Opcode;
     size_t MinIdx;
     size_t MaxIdx;
   };
 
   // Collect the operations of all the instructions that will be needed to
   // expand the SCEVExpr. This is so that when we come to cost the operands,
   // we know what the generated user(s) will be.
   SmallVector<OperationIndices, 2> Operations;
 
-  auto CastCost = [&](unsigned Opcode) {
+  auto CastCost = [&](unsigned Opcode) -> InstructionCost {
     Operations.emplace_back(Opcode, 0, 0);
     return TTI.getCastInstrCost(Opcode, S->getType(),
                                 S->getOperand(0)->getType(),
                                 TTI::CastContextHint::None, CostKind);
   };
 
   auto ArithCost = [&](unsigned Opcode, unsigned NumRequired,
-                       unsigned MinIdx = 0, unsigned MaxIdx = 1) {
+                       unsigned MinIdx = 0,
+                       unsigned MaxIdx = 1) -> InstructionCost {
     Operations.emplace_back(Opcode, MinIdx, MaxIdx);
     return NumRequired *
       TTI.getArithmeticInstrCost(Opcode, S->getType(), CostKind);
   };
 
-  auto CmpSelCost = [&](unsigned Opcode, unsigned NumRequired,
-                        unsigned MinIdx, unsigned MaxIdx) {
+  auto CmpSelCost = [&](unsigned Opcode, unsigned NumRequired, unsigned MinIdx,
+                        unsigned MaxIdx) -> InstructionCost {
     Operations.emplace_back(Opcode, MinIdx, MaxIdx);
     Type *OpType = S->getOperand(0)->getType();
     return NumRequired * TTI.getCmpSelInstrCost(
                              Opcode, OpType, CmpInst::makeCmpResultType(OpType),
                              CmpInst::BAD_ICMP_PREDICATE, CostKind);
   };
 
   switch (S->getSCEVType()) {
case scAddRecExpr: {
     int NumNonZeroDegreeNonOneTerms =
       llvm::count_if(S->operands(), [](const SCEV *Op) {
                       auto *SConst = dyn_cast<SCEVConstant>(Op);
                       return !SConst || SConst->getAPInt().ugt(1);
                     });
 
     // 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,
-                            /*MinIdx*/1, /*MaxIdx*/1);
+    InstructionCost AddCost = ArithCost(Instruction::Add, NumTerms - 1,
+                                        /*MinIdx*/ 1, /*MaxIdx*/ 1);
     // Here, *each* one of those will require a multiplication.
-    int MulCost = ArithCost(Instruction::Mul, NumNonZeroDegreeNonOneTerms);
+    InstructionCost MulCost =
+        ArithCost(Instruction::Mul, NumNonZeroDegreeNonOneTerms);
     Cost = AddCost + MulCost;
 
     // What is the degree of this polynominal?
     int PolyDegree = S->getNumOperands() - 1;
     assert(PolyDegree >= 1 && "Should be at least affine.");
 
     // The final term will be:
     //   Op_{PolyDegree} * x ^ {PolyDegree}
for (auto SCEVOp : enumerate(S->operands())) {
       Worklist.emplace_back(CostOp.Opcode, OpIdx, SCEVOp.value());
     }
   }
   return Cost;
 }
 
 bool SCEVExpander::isHighCostExpansionHelper(
     const SCEVOperand &WorkItem, Loop *L, const Instruction &At,
-    int &BudgetRemaining, const TargetTransformInfo &TTI,
+    InstructionCost &Cost, unsigned Budget, const TargetTransformInfo &TTI,
     SmallPtrSetImpl<const SCEV *> &Processed,
     SmallVectorImpl<SCEVOperand> &Worklist) {
-  if (BudgetRemaining < 0)
+  if (Cost > Budget)
     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 (!isa<SCEVConstant>(S) && !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.
case scUnknown:
     // Assume to be zero-cost.
     return false;
   case scConstant: {
     // Only evalulate the costs of constants when optimizing for size.
     if (CostKind != TargetTransformInfo::TCK_CodeSize)
       return 0;
     const APInt &Imm = cast<SCEVConstant>(S)->getAPInt();
     Type *Ty = S->getType();
-    BudgetRemaining -= TTI.getIntImmCostInst(
+    Cost += TTI.getIntImmCostInst(
         WorkItem.ParentOpcode, WorkItem.OperandIdx, Imm, Ty, CostKind);
-    return BudgetRemaining < 0;
+    return Cost > Budget;
   }
   case scTruncate:
   case scPtrToInt:
   case scZeroExtend:
   case scSignExtend: {
-    int Cost =
+    Cost +=
         costAndCollectOperands<SCEVCastExpr>(WorkItem, TTI, CostKind, Worklist);
-    BudgetRemaining -= Cost;
     return false; // Will answer upon next entry into this function.
   }
   case scUDivExpr: {
     // 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 =
+    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.
   }
   case scAddExpr:
   case scMulExpr:
   case scUMaxExpr:
   case scSMaxExpr:
   case scUMinExpr:
   case scSMinExpr: {
     assert(cast<SCEVNAryExpr>(S)->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.
-    int Cost =
+    Cost +=
         costAndCollectOperands<SCEVNAryExpr>(WorkItem, TTI, CostKind, Worklist);
-    BudgetRemaining -= Cost;
-    return BudgetRemaining < 0;
+    return Cost > Budget;
   }
   case scAddRecExpr: {
     assert(cast<SCEVAddRecExpr>(S)->getNumOperands() >= 2 &&
            "Polynomial should be at least linear");
-    BudgetRemaining -= costAndCollectOperands<SCEVAddRecExpr>(
+    Cost += costAndCollectOperands<SCEVAddRecExpr>(
         WorkItem, TTI, CostKind, Worklist);
-    return BudgetRemaining < 0;
+    return Cost > Budget;
   }
   }
   llvm_unreachable("Unknown SCEV kind!");
 }
 
 Value *SCEVExpander::expandCodeForPredicate(const SCEVPredicate *Pred,
                                             Instruction *IP) {
   assert(IP);