[SCEVExpander] Support cost evaluation of several SCEVs with same budget

[mkazantsev]

Summary

This is a follow-up from discussion in D138412. Sometimes we want to evaluate
the cost of expansion of several SCEVs together with same budget. For example,
if one of them is a bit above cheap limit, and the second one is free, then
we still want to expand. Checking each of them with "cheap" limit is a bit more
pessimistic.

[lebedev.ri]

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

226

I was thinking Budget should be a non-const-ref,
and we should just subtract the computed cost from it
if we return false from this function.
Would that be worse?

[lebedev.ri]

(Thank you for looking into this!)

[mkazantsev]

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

226

I think it should be fine, let's try it out.

[mkazantsev]

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

226

Let's leave it as is. The issue I'm seeing is that the estimate is done in terms of scaled budget, and there is no clear translation of scaled budget back into the normal one (division is lossy).

[lebedev.ri]

While i do think that the suggestion would be useful,
this is better than nothing

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

226

As far as i'm concerned TargetTransformInfo::TCC_Basic is always 1,
so you could just assert that the division is lossless.

[lebedev.ri]

(Thank you!)

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

SmallVector<Instruction *, 32> getAllInsertedInstructions() const {
     }
 
     return Result;
   }
 
   /// Return true for expressions that can't be evaluated at runtime
   /// within given \b Budget.
   ///
-  /// At is a parameter which specifies point in code where user is going to
-  /// expand this expression. Sometimes this knowledge can lead to
+  /// \p At is a parameter which specifies point in code where user is going to
+  /// expand these expressions. Sometimes this knowledge can lead to
   /// a less pessimistic cost estimation.
-  bool isHighCostExpansion(const SCEV *Expr, Loop *L, unsigned Budget,
-                           const TargetTransformInfo *TTI,
+  bool isHighCostExpansion(ArrayRef<const SCEV *> Exprs, 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;
     InstructionCost Cost = 0;
     unsigned ScaledBudget = Budget * TargetTransformInfo::TCC_Basic;
+    for (auto *Expr : Exprs)
       Worklist.emplace_back(-1, -1, Expr);
     while (!Worklist.empty()) {
       const SCEVOperand WorkItem = Worklist.pop_back_val();
       if (isHighCostExpansionHelper(WorkItem, L, *At, Cost, ScaledBudget, *TTI,
                                     Processed, Worklist))
         return true;
     }
     assert(Cost <= ScaledBudget && "Should have returned from inner loop.");
     return false;

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

bool SimplifyIndvar::makeIVComparisonInvariant(ICmpInst *ICmp,
   if (!LIP)
     return false;
   ICmpInst::Predicate InvariantPredicate = LIP->Pred;
   const SCEV *InvariantLHS = LIP->LHS;
   const SCEV *InvariantRHS = LIP->RHS;
 
   // Do not generate something ridiculous.
   auto *PHTerm = Preheader->getTerminator();
-  if (Rewriter.isHighCostExpansion(InvariantLHS, L, SCEVCheapExpansionBudget,
-                                   TTI, PHTerm))
-    return false;
-  if (Rewriter.isHighCostExpansion(InvariantRHS, L, SCEVCheapExpansionBudget,
-                                   TTI, PHTerm))
+  if (Rewriter.isHighCostExpansion({ InvariantLHS, InvariantRHS }, L,
+                                   2 * SCEVCheapExpansionBudget, TTI, PHTerm))
     return false;
   auto *NewLHS =
       Rewriter.expandCodeFor(InvariantLHS, IVOperand->getType(), PHTerm);
   auto *NewRHS =
       Rewriter.expandCodeFor(InvariantRHS, IVOperand->getType(), PHTerm);
   LLVM_DEBUG(dbgs() << "INDVARS: Simplified comparison: " << *ICmp << '\n');
   ICmp->setPredicate(InvariantPredicate);
   ICmp->setOperand(0, NewLHS);