[SCEV] SCEVExpander::isHighCostExpansionHelper(): cost-model polynomial recurrence

[lebedev.ri]

Summary

So, i wouldn't call this *obviously* correct,
but i think i got it right this time :)

Roughly, we have

Op0*x^0 + Op1*x^1 + Op2*x^2 ...

where Op_{n} * x^{n} is called term, and n the degree of term.

Due to the way they are stored internally in SCEVAddRecExpr,
i believe we can have Op_{n} to be 0, so we should not charge for those.

I think it is most straight-forward to count the cost in 4 steps:

1. First, count it the same way we counted scAddExpr, but be sure to skip terms with zero constants. Much like with add expr we will have one less addition than number of terms.
2. Each non-constant term (term degree >= 1) requires a multiplication between the Op_{n} and x^{n}. But again, only charge for it if it is required - Op_{n} must not be 0 (no term) or 1 (no multiplication needed), and obviously don't charge constant terms (x^0 == 1).
3. We must charge for all the x^0..x^{poly_degree} themselves. Since x^{poly_degree} is x * x * ... * x, i.e. poly_degree x'es multiplied, for final poly_degree term we again require poly_degree-1 multiplications. Note that all the x^{0}..x^{poly_degree-1} will be computed for the free along the way there.
4. And finally, the operands themselves.

Here, much like with add/mul exprs, we really don't look for preexisting instructions..

[lebedev.ri]

NFC, reorder branches for simple further patch.

[lebedev.ri]

Hmm, though this is still a bit too pessimistic (high) cost if we the same operand (coeffient) more than once, consider:

C0 + C1*x + C1*x*x

currently will be counted as 2x add + 3x mul, but C1*x part can be reused,
so this is actually just 2x add + 2x mul:

C0 + C1*x*(1 + x)

I'm struggling to find a counter-example where such grouping would result in too optimistic (low) cost though.

[lebedev.ri]

Rebased.

[lebedev.ri]

Ping @mkazantsev / @reames
Thanks

[lebedev.ri]

Ping @reames / @mkazantsev
Thanks

[lebedev.ri]

Ping @reames / @mkazantsev
Please do indicate if there is any way i can help move this process along.
Thanks

[mkazantsev]

LGTM

llvm/lib/Analysis/ScalarEvolutionExpander.cpp

2246

This variable name is somewhat misleading, it should clearly say that it only counts non-one coefficients.

2259

Maybe also assert that the last operand is not zero (trailing zeros should be thrown away on SCEV creation).

[lebedev.ri]

In D73741#1890707, @mkazantsev wrote:

LGTM

Thank you for the review.
Here i will also add a FIXME to investigate whether this is too pessimistic.

llvm/lib/Analysis/ScalarEvolutionExpander.cpp

2259

Good idea.

[lebedev.ri]

Rebased, review notes addressed.

[lebedev.ri]

Actually upload updated patch - fixes got committed into a followup commit locally.

llvm/lib/Analysis/ScalarEvolutionExpander.cpp

if (auto *UDivExpr = dyn_cast<SCEVUDivExpr>(S)) {
     // Need to count the cost of this UDiv.
     BudgetRemaining -= TTI.getOperationCost(Instruction::UDiv, S->getType());
     return isHighCostExpansionHelper(UDivExpr->getLHS(), L, At, BudgetRemaining,
                                      TTI, Processed) ||
            isHighCostExpansionHelper(UDivExpr->getRHS(), L, At, BudgetRemaining,
                                      TTI, Processed);
   }
 
+  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.getOperationCost(Instruction::Add, OpType);
+    int MulCost = TTI.getOperationCost(Instruction::Mul, OpType);
+
+    // 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(); });
+    assert(NumTerms >= 1 && "Polynominal should have at least one term.");
+    assert(!(*std::prev(NAry->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);
+                       });
+    // Here, *each* one of those will require a multiplication.
+    BudgetRemaining -= MulCost * NumNonZeroDegreeNonOneTerms;
+    if (BudgetRemaining < 0)
+      return true;
+
+    // What is the degree of this polynominal?
+    int PolyDegree = NAry->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;
+
+    // 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;
+  }
+
   if (S->getSCEVType() == scAddExpr || S->getSCEVType() == scMulExpr) {
     const SCEVNAryExpr *NAry = dyn_cast<SCEVNAryExpr>(S);
 
     unsigned Opcode;
     switch (S->getSCEVType()) {
     case scAddExpr:
       Opcode = Instruction::Add;
       break;
if (S->getSCEVType() == scAddExpr || S->getSCEVType() == scMulExpr) {
     return BudgetRemaining < 0;
   }
 
   // HowManyLessThans uses a Max expression whenever the loop is not guarded by
   // the exit condition.
   if (isa<SCEVMinMaxExpr>(S))
     return true;
 
-  // Recurse past nary expressions, which commonly occur in the
-  // BackedgeTakenCount. They may already exist in program code, and if not,
-  // they are not too expensive rematerialize.
-  if (const SCEVNAryExpr *NAry = dyn_cast<SCEVNAryExpr>(S)) {
-    for (auto *Op : NAry->operands())
-      if (isHighCostExpansionHelper(Op, L, At, BudgetRemaining, TTI, Processed))
-        return true;
-  }
-
   // If we haven't recognized an expensive SCEV pattern, assume it's an
   // expression produced by program code.
   return false;
 }
 
 Value *SCEVExpander::expandCodeForPredicate(const SCEVPredicate *Pred,
                                             Instruction *IP) {
   assert(IP);