[SCEV] Cache results during GetMinTrailingZeros query

[igor-laevsky]

Summary

This fixes bug 18606.

In certain cases we may get exponentially sized SCEV's in the form of:
S0 = InitialValue
S1 = mul S0 S0
S2 = mul S1 S1
S3 = mul S2 S2

Naively visiting every expression will result in the exponential complexity. However by adding caching layer we can simplify this and achieve linear complexity.

Note that caching doesn't protect us from other forms of the exponentially sized SCEV's. In order to not run infinite compiles in those cases I will add separate threshold in the follow up change.

[efriedma]

Would it make sense to add a global GetMinTrailingZeros cache instead?

[sanjoy]

In D29759#672573, @efriedma wrote:

Would it make sense to add a global GetMinTrailingZeros cache instead?

That's actually what I had in mind -- please make the cache global to the SCEV class.

[igor-laevsky]

Store cache inside of the ScalarEvolution pass.

[sanjoy]

lgtm, but please make the Impl a private member function of ScalarEvolution.

lib/Analysis/ScalarEvolution.cpp

4440 ↗

(On Diff #87993)

Extracting out the Impl function as a static function is somewhat odd -- why did you just keep it as a normal member function?

[igor-laevsky]

lib/Analysis/ScalarEvolution.cpp

4440 ↗

(On Diff #87993)

Idea was to hide implementation details, but both options seem fine to me.

llvm/trunk/include/llvm/Analysis/ScalarEvolution.h

private:
   /// Set to true by isLoopBackedgeGuardedByCond when we're walking the set of
   /// conditions dominating the backedge of a loop.
   bool WalkingBEDominatingConds;
 
   /// Set to true by isKnownPredicateViaSplitting when we're trying to prove a
   /// predicate by splitting it into a set of independent predicates.
   bool ProvingSplitPredicate;
 
+  /// Memoized values for the GetMinTrailingZeros
+  DenseMap<const SCEV *, uint32_t> MinTrailingZerosCache;
+
+  /// Private helper method for the GetMinTrailingZeros method
+  uint32_t GetMinTrailingZerosImpl(const SCEV *S);
+
   /// Information about the number of loop iterations for which a loop exit's
   /// branch condition evaluates to the not-taken path.  This is a temporary
   /// pair of exact and max expressions that are eventually summarized in
   /// ExitNotTakenInfo and BackedgeTakenInfo.
   struct ExitLimit {
     const SCEV *ExactNotTaken; // The exit is not taken exactly this many times
     const SCEV *MaxNotTaken; // The exit is not taken at most this many times
     bool MaxOrZero; // Not taken either exactly MaxNotTaken or zero times

llvm/trunk/lib/Analysis/ScalarEvolution.cpp

if (!GEP->getSourceElementType()->isSized())
     return getUnknown(GEP);
 
   SmallVector<const SCEV *, 4> IndexExprs;
   for (auto Index = GEP->idx_begin(); Index != GEP->idx_end(); ++Index)
     IndexExprs.push_back(getSCEV(*Index));
   return getGEPExpr(GEP, IndexExprs);
 }
 
-uint32_t
-ScalarEvolution::GetMinTrailingZeros(const SCEV *S) {
+uint32_t ScalarEvolution::GetMinTrailingZerosImpl(const SCEV *S) {
   if (const SCEVConstant *C = dyn_cast<SCEVConstant>(S))
     return C->getAPInt().countTrailingZeros();
 
   if (const SCEVTruncateExpr *T = dyn_cast<SCEVTruncateExpr>(S))
     return std::min(GetMinTrailingZeros(T->getOperand()),
                     (uint32_t)getTypeSizeInBits(T->getType()));
 
   if (const SCEVZeroExtendExpr *E = dyn_cast<SCEVZeroExtendExpr>(S)) {
     uint32_t OpRes = GetMinTrailingZeros(E->getOperand());
-    return OpRes == getTypeSizeInBits(E->getOperand()->getType()) ?
-             getTypeSizeInBits(E->getType()) : OpRes;
+    return OpRes == getTypeSizeInBits(E->getOperand()->getType())
+               ? getTypeSizeInBits(E->getType())
+               : OpRes;
   }
 
   if (const SCEVSignExtendExpr *E = dyn_cast<SCEVSignExtendExpr>(S)) {
     uint32_t OpRes = GetMinTrailingZeros(E->getOperand());
-    return OpRes == getTypeSizeInBits(E->getOperand()->getType()) ?
-             getTypeSizeInBits(E->getType()) : OpRes;
+    return OpRes == getTypeSizeInBits(E->getOperand()->getType())
+               ? getTypeSizeInBits(E->getType())
+               : OpRes;
   }
 
   if (const SCEVAddExpr *A = dyn_cast<SCEVAddExpr>(S)) {
     // The result is the min of all operands results.
     uint32_t MinOpRes = GetMinTrailingZeros(A->getOperand(0));
     for (unsigned i = 1, e = A->getNumOperands(); MinOpRes && i != e; ++i)
       MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(A->getOperand(i)));
     return MinOpRes;
   }
 
   if (const SCEVMulExpr *M = dyn_cast<SCEVMulExpr>(S)) {
     // The result is the sum of all operands results.
     uint32_t SumOpRes = GetMinTrailingZeros(M->getOperand(0));
     uint32_t BitWidth = getTypeSizeInBits(M->getType());
     for (unsigned i = 1, e = M->getNumOperands();
          SumOpRes != BitWidth && i != e; ++i)
-      SumOpRes = std::min(SumOpRes + GetMinTrailingZeros(M->getOperand(i)),
-                          BitWidth);
+      SumOpRes =
+          std::min(SumOpRes + GetMinTrailingZeros(M->getOperand(i)), BitWidth);
     return SumOpRes;
   }
 
   if (const SCEVAddRecExpr *A = dyn_cast<SCEVAddRecExpr>(S)) {
     // The result is the min of all operands results.
     uint32_t MinOpRes = GetMinTrailingZeros(A->getOperand(0));
     for (unsigned i = 1, e = A->getNumOperands(); MinOpRes && i != e; ++i)
       MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(A->getOperand(i)));
computeKnownBits(U->getValue(), Zeros, Ones, getDataLayout(), 0, &AC,
                      nullptr, &DT);
     return Zeros.countTrailingOnes();
   }
 
   // SCEVUDivExpr
   return 0;
 }
 
+uint32_t ScalarEvolution::GetMinTrailingZeros(const SCEV *S) {
+  auto I = MinTrailingZerosCache.find(S);
+  if (I != MinTrailingZerosCache.end())
+    return I->second;
+
+  uint32_t Result = GetMinTrailingZerosImpl(S);
+  auto InsertPair = MinTrailingZerosCache.insert({S, Result});
+  assert(InsertPair.second && "Should insert a new key");
+  return InsertPair.first->second;
+}
+
 /// Helper method to assign a range to V from metadata present in the IR.
 static Optional<ConstantRange> GetRangeFromMetadata(Value *V) {
   if (Instruction *I = dyn_cast<Instruction>(V))
     if (MDNode *MD = I->getMetadata(LLVMContext::MD_range))
       return getConstantRangeFromMetadata(*MD);
 
   return None;
 }

 }
 
 ScalarEvolution::ScalarEvolution(ScalarEvolution &&Arg)
     : F(Arg.F), HasGuards(Arg.HasGuards), TLI(Arg.TLI), AC(Arg.AC), DT(Arg.DT),
       LI(Arg.LI), CouldNotCompute(std::move(Arg.CouldNotCompute)),
       ValueExprMap(std::move(Arg.ValueExprMap)),
       PendingLoopPredicates(std::move(Arg.PendingLoopPredicates)),
       WalkingBEDominatingConds(false), ProvingSplitPredicate(false),
+      MinTrailingZerosCache(std::move(Arg.MinTrailingZerosCache)),
       BackedgeTakenCounts(std::move(Arg.BackedgeTakenCounts)),
       PredicatedBackedgeTakenCounts(
           std::move(Arg.PredicatedBackedgeTakenCounts)),
       ConstantEvolutionLoopExitValue(
           std::move(Arg.ConstantEvolutionLoopExitValue)),
       ValuesAtScopes(std::move(Arg.ValuesAtScopes)),
       LoopDispositions(std::move(Arg.LoopDispositions)),
       LoopPropertiesCache(std::move(Arg.LoopPropertiesCache)),

 void ScalarEvolution::forgetMemoizedResults(const SCEV *S) {
   ValuesAtScopes.erase(S);
   LoopDispositions.erase(S);
   BlockDispositions.erase(S);
   UnsignedRanges.erase(S);
   SignedRanges.erase(S);
   ExprValueMap.erase(S);
   HasRecMap.erase(S);
+  MinTrailingZerosCache.erase(S);
 
   auto RemoveSCEVFromBackedgeMap =
       [S, this](DenseMap<const Loop *, BackedgeTakenInfo> &Map) {
         for (auto I = Map.begin(), E = Map.end(); I != E;) {
           BackedgeTakenInfo &BEInfo = I->second;
           if (BEInfo.hasOperand(S, this)) {
             BEInfo.clear();
             Map.erase(I++);

llvm/trunk/test/Analysis/ScalarEvolution/pr18606-min-zeros.ll

+; RUN: opt -S -indvars < %s | FileCheck %s
+
+; CHECK: @test
+; CHECK: %5 = add i32 %local_6_, %local_0_
+; CEHCK: %37 = mul i32 %36, %36
+
+define i32 @test(i32, i32) {
+bci_0:
+  br label %bci_30
+
+bci_68:                                           ; preds = %bci_45
+  %local_6_.lcssa = phi i32 [ %local_6_, %bci_45 ]
+  %.lcssa1.lcssa = phi i32 [ %37, %bci_45 ]
+  %.lcssa.lcssa = phi i32 [ 34, %bci_45 ]
+  %2 = add i32 %local_6_.lcssa, 262
+  %3 = add i32 %2, %.lcssa1.lcssa
+  %4 = add i32 %3, %.lcssa.lcssa
+  ret i32 %4
+
+bci_30:                                           ; preds = %bci_45, %bci_0
+  %local_0_ = phi i32 [ %0, %bci_0 ], [ %38, %bci_45 ]
+  %local_6_ = phi i32 [ 2, %bci_0 ], [ %39, %bci_45 ]
+  %5 = add i32 %local_6_, %local_0_
+  br label %bci_45
+
+bci_45:                                           ; preds = %bci_30
+  %6 = mul i32 %5, %5
+  %7 = mul i32 %6, %6
+  %8 = mul i32 %7, %7
+  %9 = mul i32 %8, %8
+  %10 = mul i32 %9, %9
+  %11 = mul i32 %10, %10
+  %12 = mul i32 %11, %11
+  %13 = mul i32 %12, %12
+  %14 = mul i32 %13, %13
+  %15 = mul i32 %14, %14
+  %16 = mul i32 %15, %15
+  %17 = mul i32 %16, %16
+  %18 = mul i32 %17, %17
+  %19 = mul i32 %18, %18
+  %20 = mul i32 %19, %19
+  %21 = mul i32 %20, %20
+  %22 = mul i32 %21, %21
+  %23 = mul i32 %22, %22
+  %24 = mul i32 %23, %23
+  %25 = mul i32 %24, %24
+  %26 = mul i32 %25, %25
+  %27 = mul i32 %26, %26
+  %28 = mul i32 %27, %27
+  %29 = mul i32 %28, %28
+  %30 = mul i32 %29, %29
+  %31 = mul i32 %30, %30
+  %32 = mul i32 %31, %31
+  %33 = mul i32 %32, %32
+  %34 = mul i32 %33, %33
+  %35 = mul i32 %34, %34
+  %36 = mul i32 %35, %35
+  %37 = mul i32 %36, %36
+  %38 = add i32 %37, -11
+  %39 = add i32 %local_6_, 1
+  %40 = icmp sgt i32 %39, 76
+  br i1 %40, label %bci_68, label %bci_30
+}