Make cost estimation in RewriteLoopExitValues smarter

[igor-laevsky]

Summary

In the recent change (http://reviews.llvm.org/D10782) Sanjoy updated cost estimation for the RewriteLoopExitValues to consider complicated SCEV expressions as cheap if there already exists llvm-ir expression which computes same value.

However if such expression is wrapped in a several other expressions we will still consider it as high cost.
For example this is considered high cost:

{-2,+,1} + (zext (<expr>)), where <expr> is high cost expression already defined in the IR.

In this change I updated cost estimation to evaluate such nested expressions as cheap. Turns out it is fairly important on one of our benchmarks (~20% performance).

In order to do that I have added new method "TryToGetValueForSCEVAt" into SCEVExpander, which tries to look for pre-existing llvm ir value. This method is used inside "ExpandSCEVIfNeeded" and inside "isHighCostExpansionHelper"

It would be natural to add "InsertedExpressions" lookup into this method, but I am not sure if there is really need for this. Also it is possible to use this method during "SCEVExpander::expand", but I am worried about having to much impact on the unrelated areas with no good reason.

[sanjoy]

This needs test cases.

include/llvm/Analysis/ScalarEvolutionExpander.h

121

Nit: spacing around =.

121

Please document what At does here.

197

Nit: "available"

197

Nit: I'd use "LLVM IR"

201

Nit: "exhaustive".

202

Nit: "didn't find"

203

I'd call this getExistingSCEVExpansion, making it clear that we're looking for a value that already exists.

lib/Analysis/ScalarEvolutionExpander.cpp

1823

The way we're using this now in RewriteLoopExitValues, S is not an induction variable, but is an operand to the loop's backedge comparison which happens to be equal to the trip count of the loop.

I'd just leave in the first sentence in the comment, and delete the rest.

1894

Should this division case also use TryToGetValueForSCEVAt? The search logic here is basically identical, except the SCEV subtraction in the end (and perhaps we should change TryToGetValueForSCEVAt to look at exiting blocks too?).

[igor-laevsky]

Thanks for the comments!

include/llvm/Analysis/ScalarEvolutionExpander.h

203

Renamed to tryToGetExistingExpansion. I'd like to leave tryTo in the name to emphasise that function can sometimes fail to find any value.

lib/Analysis/ScalarEvolutionExpander.cpp

1894

In TryToGetValueForSCEVAt we are looking for:
ExitCond == S, where ExitCond is LHS or RHS
On success we return ExitCond, because it is existing LLVM IR value for S.

For division we are looking for:
(ExitCond - 1) == S, where ExitCond is LHS or RHS
If we move this subtraction logic into TryToGetValueForSCEVAt we will not be able to return anything since we actually didn't find any pre-existing LLVM IR value.

I suspect we can completely remove division special handling if:

1. (ExitCond - 1) == S <=> ExitCond == (S + 1) (or other non AddRec expression with S)
2. isHighCostExpansionHelper can recursively look into expressions. I.e IsHighCost(a + b) == IsHighCost(a) || isHighCost(b) (this is true right now)
3. We will update TryToGetValueForSCEVAt to look into exiting blocks

However I am afraid to break something, so I will do it in separate change. (But first I will check that division case is covered by existing tests.)

[sanjoy]

LGTM with some comments inline.

include/llvm/Analysis/ScalarEvolutionExpander.h

203

How about findExistingExpansion then?

lib/Analysis/ScalarEvolutionExpander.cpp

1853

Minor: I'm not a native English speaker, but are there some prepositions missing here? I'd have expected "If we can find an existing value for this SCEV available at At then consider the expression cheap".

include/llvm/Analysis/ScalarEvolutionExpander.h

void clear() {
       InsertedExpressions.clear();
       InsertedValues.clear();
       InsertedPostIncValues.clear();
       ChainedPhis.clear();
     }
 
     /// \brief Return true for expressions that may incur non-trivial cost to
     /// evaluate at runtime.
-    bool isHighCostExpansion(const SCEV *Expr, Loop *L) {
+    bool isHighCostExpansion(const SCEV *Expr, Loop *L,
+                             const Instruction *At=nullptr) {
       SmallPtrSet<const SCEV *, 8> Processed;
-      return isHighCostExpansionHelper(Expr, L, Processed);
+      return isHighCostExpansionHelper(Expr, L, At, Processed);
     }
 
     /// \brief 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
     /// \brief Return true if the specified instruction was inserted by the code
     /// rewriter.  If so, the client should not modify the instruction.
     bool isInsertedInstruction(Instruction *I) const {
       return InsertedValues.count(I) || InsertedPostIncValues.count(I);
     }
 
     void setChainedPhi(PHINode *PN) { ChainedPhis.insert(PN); }
 
+    /// \brief Try to find llvm-ir value for 'S' avaliable at the point 'At'.
+    // 'L' is a hint which tells in which loop to look for the suitable value.
+    // On success return value which is equivalent to the expanded 'S' at point
+    // 'At'. Return nullptr if value was not found.
+    // Note that this function does not perform exaustive search. I.e if it
+    // didn't found any value it does not mean that there is no such value.
+    Value *TryToGetValueForSCEVAt(const SCEV *S, const Instruction *At,
+                                  Loop *L);
+
   private:
     LLVMContext &getContext() const { return SE.getContext(); }
 
     /// \brief Recursive helper function for isHighCostExpansion.
     bool isHighCostExpansionHelper(const SCEV *S, Loop *L,
+                                   const Instruction *At,
                                    SmallPtrSetImpl<const SCEV *> &Processed);
 
     /// \brief Insert the specified binary operator, doing a small amount
     /// of work to avoid inserting an obviously redundant operation.
     Value *InsertBinop(Instruction::BinaryOps Opcode, Value *LHS, Value *RHS);
 
     /// \brief Arrange for there to be a cast of V to Ty at IP, reusing an
     /// existing cast if a suitable one exists, moving an existing cast if a

lib/Analysis/ScalarEvolutionExpander.cpp

if (OrigPhiRef->getType() != Phi->getType()) {
       NewIV = Builder.CreateTruncOrBitCast(OrigPhiRef, Phi->getType(), IVName);
     }
     Phi->replaceAllUsesWith(NewIV);
     DeadInsts.emplace_back(Phi);
   }
   return NumElim;
 }
 
+Value *SCEVExpander::TryToGetValueForSCEVAt(const SCEV *S,
+                                            const Instruction *At, Loop *L) {
+  using namespace llvm::PatternMatch;
+
+  SmallVector<BasicBlock *, 4> Latches;
+  L->getLoopLatches(Latches);
+
+  // Look for suitable value in simple conditions at the loop latches. This is
+  // based on the fact that 'S' is probably an induction variable evaluated
+  // after loop 'L'.
+  for (BasicBlock *BB : Latches) {
+    ICmpInst::Predicate Pred;
+    Instruction *LHS, *RHS;
+    BasicBlock *TrueBB, *FalseBB;
+
+    if (!match(BB->getTerminator(),
+               m_Br(m_ICmp(Pred, m_Instruction(LHS), m_Instruction(RHS)),
+                    TrueBB, FalseBB)))
+      continue;
+
+    if (SE.getSCEV(LHS) == S && SE.DT->dominates(LHS, At))
+      return LHS;
+
+    if (SE.getSCEV(RHS) == S && SE.DT->dominates(RHS, At))
+      return RHS;
+  }
+
+  // There is potential to make this significantly smarter, but this simple
+  // heuristic already gets some interesting cases.
+
+  // Can not find suitable value.
+  return nullptr;
+}
+
 bool SCEVExpander::isHighCostExpansionHelper(
-    const SCEV *S, Loop *L, SmallPtrSetImpl<const SCEV *> &Processed) {
+    const SCEV *S, Loop *L, const Instruction *At,
+    SmallPtrSetImpl<const SCEV *> &Processed) {
+
+  // If we can find value for this scev avaliable at the point "At"
+  // consider expression cheap.
+  if (At && TryToGetValueForSCEVAt(S, At, L) != nullptr)
+    return false;
 
   // Zero/One operand expressions
   switch (S->getSCEVType()) {
   case scUnknown:
   case scConstant:
     return false;
   case scTruncate:
-    return isHighCostExpansionHelper(cast<SCEVTruncateExpr>(S)->getOperand(), L,
-                                     Processed);
+    return isHighCostExpansionHelper(cast<SCEVTruncateExpr>(S)->getOperand(),
+                                     L, At, Processed);
   case scZeroExtend:
     return isHighCostExpansionHelper(cast<SCEVZeroExtendExpr>(S)->getOperand(),
-                                     L, Processed);
+                                     L, At, Processed);
   case scSignExtend:
     return isHighCostExpansionHelper(cast<SCEVSignExtendExpr>(S)->getOperand(),
-                                     L, Processed);
+                                     L, At, Processed);
   }
 
   if (!Processed.insert(S).second)
     return false;
 
   if (auto *UDivExpr = dyn_cast<SCEVUDivExpr>(S)) {
     // If the divisor is a power of two and the SCEV type fits in a native
     // integer, consider the divison cheap irrespective of whether it occurs in
     // the user code since it can be lowered into a right shift.
     if (auto *SC = dyn_cast<SCEVConstant>(UDivExpr->getRHS()))
       if (SC->getValue()->getValue().isPowerOf2()) {
         const DataLayout &DL =
             L->getHeader()->getParent()->getParent()->getDataLayout();
         unsigned Width = cast<IntegerType>(UDivExpr->getType())->getBitWidth();
         return DL.isIllegalInteger(Width);
       }
 
     // 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, assume the former consider UDivExpr expensive to
     // compute.
     BasicBlock *ExitingBB = L->getExitingBlock();
     if (!ExitingBB)
       return true;
 
     BranchInst *ExitingBI = dyn_cast<BranchInst>(ExitingBB->getTerminator());
     if (!ExitingBI || !ExitingBI->isConditional())
       return true;
 
     ICmpInst *OrigCond = dyn_cast<ICmpInst>(ExitingBI->getCondition());
if (isa<SCEVSMaxExpr>(S) || isa<SCEVUMaxExpr>(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 (SCEVNAryExpr::op_iterator I = NAry->op_begin(), E = NAry->op_end();
          I != E; ++I) {
-      if (isHighCostExpansionHelper(*I, L, Processed))
+      if (isHighCostExpansionHelper(*I, L, At, Processed))
         return true;
     }
   }
 
   // If we haven't recognized an expensive SCEV pattern, assume it's an
   // expression produced by program code.
   return false;
 }

lib/Transforms/Scalar/IndVarSimplify.cpp

private:
     void RewriteLoopExitValues(Loop *L, SCEVExpander &Rewriter);
 
     Value *LinearFunctionTestReplace(Loop *L, const SCEV *BackedgeTakenCount,
                                      PHINode *IndVar, SCEVExpander &Rewriter);
 
     void SinkUnusedInvariants(Loop *L);
 
     Value *ExpandSCEVIfNeeded(SCEVExpander &Rewriter, const SCEV *S, Loop *L,
-                              Instruction *InsertPt, Type *Ty,
-                              bool &IsHighCostExpansion);
+                              Instruction *InsertPt, Type *Ty);
   };
 }
 
 char IndVarSimplify::ID = 0;
 INITIALIZE_PASS_BEGIN(IndVarSimplify, "indvars",
                 "Induction Variable Simplification", false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
struct RewritePhi {
 
   RewritePhi(PHINode *P, unsigned I, Value *V, bool H, bool S)
       : PN(P), Ith(I), Val(V), HighCost(H), SafePhi(S) {}
 };
 }
 
 Value *IndVarSimplify::ExpandSCEVIfNeeded(SCEVExpander &Rewriter, const SCEV *S,
                                           Loop *L, Instruction *InsertPt,
-                                          Type *ResultTy,
-                                          bool &IsHighCostExpansion) {
-  using namespace llvm::PatternMatch;
-
-  if (!Rewriter.isHighCostExpansion(S, L)) {
-    IsHighCostExpansion = false;
-    return Rewriter.expandCodeFor(S, ResultTy, InsertPt);
-  }
-
+                                          Type *ResultTy) {
   // Before expanding S into an expensive LLVM expression, see if we can use an
-  // already existing value as the expansion for S.  There is potential to make
-  // this significantly smarter, but this simple heuristic already gets some
-  // interesting cases.
-
-  SmallVector<BasicBlock *, 4> Latches;
-  L->getLoopLatches(Latches);
-
-  for (BasicBlock *BB : Latches) {
-    ICmpInst::Predicate Pred;
-    Instruction *LHS, *RHS;
-    BasicBlock *TrueBB, *FalseBB;
-
-    if (!match(BB->getTerminator(),
-               m_Br(m_ICmp(Pred, m_Instruction(LHS), m_Instruction(RHS)),
-                    TrueBB, FalseBB)))
-      continue;
-
-    if (SE->getSCEV(LHS) == S && DT->dominates(LHS, InsertPt)) {
-      IsHighCostExpansion = false;
-      return LHS;
-    }
-
-    if (SE->getSCEV(RHS) == S && DT->dominates(RHS, InsertPt)) {
-      IsHighCostExpansion = false;
-      return RHS;
-    }
-  }
+  // already existing value as the expansion for S.
+  if (Value *RetValue = Rewriter.TryToGetValueForSCEVAt(S, InsertPt, L))
+    return RetValue;
 
   // We didn't find anything, fall back to using SCEVExpander.
-  assert(Rewriter.isHighCostExpansion(S, L) && "this should not have changed!");
-  IsHighCostExpansion = true;
   return Rewriter.expandCodeFor(S, ResultTy, InsertPt);
 }
 
 //===----------------------------------------------------------------------===//
 // RewriteLoopExitValues - Optimize IV users outside the loop.
 // As a side effect, reduces the amount of IV processing within the loop.
 //===----------------------------------------------------------------------===//
 
while ((PN = dyn_cast<PHINode>(BBI++))) {
               if (Opc != Instruction::Call && Opc != Instruction::Ret)
                 NumSoftExternalUses++;
             }
           }
           if (NumUses <= 6 && NumHardInternalUses && !NumSoftExternalUses)
             continue;
         }
 
-        bool HighCost = false;
-        Value *ExitVal = ExpandSCEVIfNeeded(Rewriter, ExitValue, L, Inst,
-                                            PN->getType(), HighCost);
+        bool HighCost = Rewriter.isHighCostExpansion(ExitValue, L, Inst);
+        Value *ExitVal =
+            ExpandSCEVIfNeeded(Rewriter, ExitValue, L, Inst, PN->getType());
 
         DEBUG(dbgs() << "INDVARS: RLEV: AfterLoopVal = " << *ExitVal << '\n'
                      << "  LoopVal = " << *Inst << "\n");
 
         if (!isValidRewrite(Inst, ExitVal)) {
           DeadInsts.push_back(ExitVal);
           continue;
         }