[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

4441

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

4441

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

lib/Analysis/ScalarEvolution.cpp

const SCEV *ScalarEvolution::createNodeForGEP(GEPOperator *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) {
-  if (const SCEVConstant *C = dyn_cast<SCEVConstant>(S))
+  class MinTrailingZerosVisitor
+      : public CachingSCEVVisitor<MinTrailingZerosVisitor, uint32_t> {
+  private:
+    const ScalarEvolution &SE;
+
+  public:
+    MinTrailingZerosVisitor(const ScalarEvolution &SE) : SE(SE) {}
+
+    uint32_t visitConstant(const SCEVConstant *C) {
       return C->getAPInt().countTrailingZeros();
+    }
 
-  if (const SCEVTruncateExpr *T = dyn_cast<SCEVTruncateExpr>(S))
-    return std::min(GetMinTrailingZeros(T->getOperand()),
-                    (uint32_t)getTypeSizeInBits(T->getType()));
+    uint32_t visitTruncateExpr(const SCEVTruncateExpr *T) {
+      return std::min(visit(T->getOperand()),
+                      (uint32_t)SE.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;
+    uint32_t visitZeroExtendExpr(const SCEVZeroExtendExpr *E) {
+      uint32_t OpRes = visit(E->getOperand());
+      return OpRes == SE.getTypeSizeInBits(E->getOperand()->getType())
+                 ? SE.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;
+    uint32_t visitSignExtendExpr(const SCEVSignExtendExpr *E) {
+      uint32_t OpRes = visit(E->getOperand());
+      return OpRes == SE.getTypeSizeInBits(E->getOperand()->getType())
+                 ? SE.getTypeSizeInBits(E->getType())
+                 : OpRes;
     }
 
-  if (const SCEVAddExpr *A = dyn_cast<SCEVAddExpr>(S)) {
+    uint32_t visitAddExpr(const SCEVAddExpr *A) {
       // The result is the min of all operands results.
-    uint32_t MinOpRes = GetMinTrailingZeros(A->getOperand(0));
+      uint32_t MinOpRes = visit(A->getOperand(0));
       for (unsigned i = 1, e = A->getNumOperands(); MinOpRes && i != e; ++i)
-      MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(A->getOperand(i)));
+        MinOpRes = std::min(MinOpRes, visit(A->getOperand(i)));
       return MinOpRes;
     }
 
-  if (const SCEVMulExpr *M = dyn_cast<SCEVMulExpr>(S)) {
+    uint32_t visitMulExpr(const SCEVMulExpr *M) {
       // The result is the sum of all operands results.
-    uint32_t SumOpRes = GetMinTrailingZeros(M->getOperand(0));
-    uint32_t BitWidth = getTypeSizeInBits(M->getType());
+      uint32_t SumOpRes = visit(M->getOperand(0));
+      uint32_t BitWidth = SE.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 + visit(M->getOperand(i)), BitWidth);
       return SumOpRes;
     }
 
-  if (const SCEVAddRecExpr *A = dyn_cast<SCEVAddRecExpr>(S)) {
+    uint32_t visitUDivExpr(const SCEVUDivExpr *Expr) { return 0; }
+
+    uint32_t visitAddRecExpr(const SCEVAddRecExpr *A) {
       // The result is the min of all operands results.
-    uint32_t MinOpRes = GetMinTrailingZeros(A->getOperand(0));
+      uint32_t MinOpRes = visit(A->getOperand(0));
       for (unsigned i = 1, e = A->getNumOperands(); MinOpRes && i != e; ++i)
-      MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(A->getOperand(i)));
+        MinOpRes = std::min(MinOpRes, visit(A->getOperand(i)));
       return MinOpRes;
     }
 
-  if (const SCEVSMaxExpr *M = dyn_cast<SCEVSMaxExpr>(S)) {
+    uint32_t visitSMaxExpr(const SCEVSMaxExpr *M) {
       // The result is the min of all operands results.
-    uint32_t MinOpRes = GetMinTrailingZeros(M->getOperand(0));
+      uint32_t MinOpRes = visit(M->getOperand(0));
       for (unsigned i = 1, e = M->getNumOperands(); MinOpRes && i != e; ++i)
-      MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(M->getOperand(i)));
+        MinOpRes = std::min(MinOpRes, visit(M->getOperand(i)));
       return MinOpRes;
     }
 
-  if (const SCEVUMaxExpr *M = dyn_cast<SCEVUMaxExpr>(S)) {
+    uint32_t visitUMaxExpr(const SCEVUMaxExpr *M) {
       // The result is the min of all operands results.
-    uint32_t MinOpRes = GetMinTrailingZeros(M->getOperand(0));
+      uint32_t MinOpRes = visit(M->getOperand(0));
       for (unsigned i = 1, e = M->getNumOperands(); MinOpRes && i != e; ++i)
-      MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(M->getOperand(i)));
+        MinOpRes = std::min(MinOpRes, visit(M->getOperand(i)));
       return MinOpRes;
     }
 
-  if (const SCEVUnknown *U = dyn_cast<SCEVUnknown>(S)) {
+    uint32_t visitUnknown(const SCEVUnknown *U) {
       // For a SCEVUnknown, ask ValueTracking.
-    unsigned BitWidth = getTypeSizeInBits(U->getType());
+      unsigned BitWidth = SE.getTypeSizeInBits(U->getType());
       APInt Zeros(BitWidth, 0), Ones(BitWidth, 0);
-    computeKnownBits(U->getValue(), Zeros, Ones, getDataLayout(), 0, &AC,
-                     nullptr, &DT);
+      computeKnownBits(U->getValue(), Zeros, Ones, SE.getDataLayout(), 0,
+                       &SE.AC, nullptr, &SE.DT);
       return Zeros.countTrailingOnes();
     }
+  };
 
-  // SCEVUDivExpr
-  return 0;
+  return MinTrailingZerosVisitor(*this).visit(S);
 }
 
 /// 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);
 

test/Analysis/ScalarEvolution/pr18606_min_zeroes.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
+}