[SCEV][IndVars] Always provide insertion point to the SCEVExpander::isHighCostExpansion()

[lebedev.ri]

Summary

This addresses the llvm/test/Transforms/IndVarSimplify/elim-extend.ll @nestedIV regression from D73728

[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. Hope it doesn't end up exploding. :)

llvm/include/llvm/Analysis/ScalarEvolutionExpander.h

184

Should we assert that provided At is not null, or there are other places where it'd be OK?

[lebedev.ri]

In D73777#1890718, @mkazantsev wrote:

LGTM. Hope it doesn't end up exploding. :)

Thank you so much for reviewing this patch set!
We'll find out whether this explodes soon :)

llvm/include/llvm/Analysis/ScalarEvolutionExpander.h

184

Hm, such an assertion does not fail as of check-llvm, so let's do that.

[lebedev.ri]

Rebased, addressed review notes.

llvm/include/llvm/Analysis/ScalarEvolutionExpander.h

void clear() {
       InsertedValues.clear();
       InsertedPostIncValues.clear();
       ChainedPhis.clear();
     }
 
     /// Return true for expressions that can't be evaluated at runtime
     /// within given \b Budget.
     ///
-    /// At is an optional parameter which specifies point in code where user is
-    /// going to expand this expression. Sometimes this knowledge can lead to a
-    /// more accurate cost estimation.
+    /// At is a parameter which specifies point in code where user is going to
+    /// expand this expression. Sometimes this knowledge can lead to
+    /// a less pessimistic cost estimation.
     bool isHighCostExpansion(const SCEV *Expr, Loop *L, unsigned Budget,
                              const TargetTransformInfo *TTI,
-                             const Instruction *At = nullptr) {
+                             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.
       SmallPtrSet<const SCEV *, 8> Processed;
       int BudgetRemaining = Budget * TargetTransformInfo::TCC_Basic;
-      return isHighCostExpansionHelper(Expr, L, At, BudgetRemaining, *TTI,
+      return isHighCostExpansionHelper(Expr, L, *At, BudgetRemaining, *TTI,
                                        Processed);
     }
 
     /// This method returns the canonical induction variable of the specified
     /// type for the specified loop (inserting one if there is none).  A
     /// canonical induction variable starts at zero and steps by one on each
     /// iteration.
     PHINode *getOrInsertCanonicalInductionVariable(const Loop *L, Type *Ty);
#endif
     Optional<ScalarEvolution::ValueOffsetPair>
     getRelatedExistingExpansion(const SCEV *S, const Instruction *At, Loop *L);
 
   private:
     LLVMContext &getContext() const { return SE.getContext(); }
 
     /// Recursive helper function for isHighCostExpansion.
     bool isHighCostExpansionHelper(const SCEV *S, Loop *L,
-                                   const Instruction *At, int &BudgetRemaining,
+                                   const Instruction &At, int &BudgetRemaining,
                                    const TargetTransformInfo &TTI,
                                    SmallPtrSetImpl<const SCEV *> &Processed);
 
     /// 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);

llvm/lib/Analysis/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;
 }
 
 bool SCEVExpander::isHighCostExpansionHelper(
-    const SCEV *S, Loop *L, const Instruction *At, int &BudgetRemaining,
+    const SCEV *S, Loop *L, const Instruction &At, int &BudgetRemaining,
     const TargetTransformInfo &TTI, SmallPtrSetImpl<const SCEV *> &Processed) {
   if (BudgetRemaining < 0)
     return true; // Already run out of budget, give up.
 
   // 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 (At && getRelatedExistingExpansion(S, At, L))
+  if (getRelatedExistingExpansion(S, &At, L))
     return false; // Consider the expression to be free.
 
   switch (S->getSCEVType()) {
   case scUnknown:
   case scConstant:
     return false; // Assume to be zero-cost.
   }
 
if (auto *SC = dyn_cast<SCEVConstant>(UDivExpr->getRHS())) {
       }
     }
 
     // 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.
 
-    BasicBlock *ExitingBB = L->getExitingBlock();
-    if (At || ExitingBB) {
-      if (!At)
-        At = &ExitingBB->back();
-
-      // 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.
-    }
+            SE.getAddExpr(S, SE.getConstant(S->getType(), 1)), &At, L))
 
     // 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);
   }

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

if (predicateLoopExits(L, Rewriter)) {
     Changed = true;
     // Given we've changed exit counts, notify SCEV
     SE->forgetLoop(L);
   }
 
   // If we have a trip count expression, rewrite the loop's exit condition
   // using it.
   if (!DisableLFTR) {
+    BasicBlock *PreHeader = L->getLoopPreheader();
+    BranchInst *PreHeaderBR = cast<BranchInst>(PreHeader->getTerminator());
+
     SmallVector<BasicBlock*, 16> ExitingBlocks;
     L->getExitingBlocks(ExitingBlocks);
     for (BasicBlock *ExitingBB : ExitingBlocks) {
       // Can't rewrite non-branch yet.
       if (!isa<BranchInst>(ExitingBB->getTerminator()))
         continue;
 
       // If our exitting block exits multiple loops, we can only rewrite the
for (BasicBlock *ExitingBB : ExitingBlocks) {
 
       PHINode *IndVar = FindLoopCounter(L, ExitingBB, ExitCount, SE, DT);
       if (!IndVar)
         continue;
 
       // Avoid high cost expansions.  Note: This heuristic is questionable in
       // that our definition of "high cost" is not exactly principled.
       if (Rewriter.isHighCostExpansion(ExitCount, L, SCEVCheapExpansionBudget,
-                                       TTI))
+                                       TTI, PreHeaderBR))
         continue;
 
       // Check preconditions for proper SCEVExpander operation. SCEV does not
       // express SCEVExpander's dependencies, such as LoopSimplify. Instead
       // any pass that uses the SCEVExpander must do it. This does not work
       // well for loop passes because SCEVExpander makes assumptions about
       // all loops, while LoopPassManager only forces the current loop to be
       // simplified.

llvm/test/Transforms/IndVarSimplify/elim-extend.ll

 ; CHECK-NEXT:    br label [[INNERLOOP:%.*]]
 ; CHECK:       innerloop:
 ; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[TMP3]], [[INNERLOOP_PREHEADER]] ], [ [[INDVARS_IV_NEXT:%.*]], [[INNERLOOP]] ]
 ; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add nsw i64 [[INDVARS_IV]], 1
 ; CHECK-NEXT:    [[ADR2:%.*]] = getelementptr i8, i8* [[ADDRESS]], i64 [[INDVARS_IV]]
 ; CHECK-NEXT:    store i8 0, i8* [[ADR2]]
 ; CHECK-NEXT:    [[ADR3:%.*]] = getelementptr i8, i8* [[ADDRESS]], i64 [[INDVARS_IV_NEXT]]
 ; CHECK-NEXT:    store i8 0, i8* [[ADR3]]
-; CHECK-NEXT:    [[INNERCMP:%.*]] = icmp sgt i64 [[TMP0]], [[INDVARS_IV_NEXT]]
-; CHECK-NEXT:    br i1 [[INNERCMP]], label [[INNERLOOP]], label [[INNEREXIT:%.*]]
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i64 [[INDVARS_IV_NEXT]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[EXITCOND]], label [[INNERLOOP]], label [[INNEREXIT:%.*]]
 ; CHECK:       innerexit:
 ; CHECK-NEXT:    [[TMP4:%.*]] = trunc i64 [[TMP0]] to i32
 ; CHECK-NEXT:    br label [[OUTERMERGE]]
 ; CHECK:       outermerge:
 ; CHECK-NEXT:    [[INNERCOUNT_MERGE]] = phi i32 [ [[TMP4]], [[INNEREXIT]] ], [ [[INNERCOUNT]], [[INNERPREHEADER]] ]
 ; CHECK-NEXT:    [[ADR4:%.*]] = getelementptr i8, i8* [[ADDRESS]], i64 [[INDVARS_IV1]]
 ; CHECK-NEXT:    store i8 0, i8* [[ADR4]]
 ; CHECK-NEXT:    [[OFS5:%.*]] = sext i32 [[INNERCOUNT_MERGE]] to i64
 ; CHECK-NEXT:    [[ADR5:%.*]] = getelementptr i8, i8* [[ADDRESS]], i64 [[OFS5]]
 ; CHECK-NEXT:    store i8 0, i8* [[ADR5]]
 ; CHECK-NEXT:    [[INDVARS_IV_NEXT2]] = add nuw nsw i64 [[INDVARS_IV1]], 1
-; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i64 [[INDVARS_IV_NEXT2]], [[WIDE_TRIP_COUNT]]
-; CHECK-NEXT:    br i1 [[EXITCOND]], label [[OUTERLOOP]], label [[RETURN:%.*]]
+; CHECK-NEXT:    [[EXITCOND4:%.*]] = icmp ne i64 [[INDVARS_IV_NEXT2]], [[WIDE_TRIP_COUNT]]
+; CHECK-NEXT:    br i1 [[EXITCOND4]], label [[OUTERLOOP]], label [[RETURN:%.*]]
 ; CHECK:       return:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %limitdec = add i32 %limit, -1
   br label %outerloop
 
 ; Eliminate %ofs1 after widening outercount.