From 0f3c9b54e60b384728c0c24518b8f2645719275e Mon Sep 17 00:00:00 2001
From: Roman Lebedev <lebedev.ri@gmail.com>
Date: Tue, 25 Feb 2020 21:52:32 +0300
Subject: [PATCH] [SCEV] SCEVExpander::isHighCostExpansionHelper(): cost-model
 polynomial recurrence

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

Reviewers: reames, mkazantsev, wmi, sanjoy

Reviewed By: mkazantsev

Subscribers: hiraditya, javed.absar, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D73741
---
 llvm/lib/Analysis/ScalarEvolutionExpander.cpp | 67 ++++++++++++++++---
 1 file changed, 58 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Analysis/ScalarEvolutionExpander.cpp b/llvm/lib/Analysis/ScalarEvolutionExpander.cpp
index 30634f70229c..6d9901dbe1c2 100644
--- a/llvm/lib/Analysis/ScalarEvolutionExpander.cpp
+++ b/llvm/lib/Analysis/ScalarEvolutionExpander.cpp
@@ -2219,6 +2219,64 @@ bool SCEVExpander::isHighCostExpansionHelper(
                                      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);
 
@@ -2258,15 +2316,6 @@ bool SCEVExpander::isHighCostExpansionHelper(
   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;