[SCEV] Memoize visitMulExpr results in SCEVRewriteVisitor. Fix PR18606

[lihuang]

Summary

[SCEV] Memoize visitMulExpr results in SCEVRewriteVisitor. Fix PR18606

When SCEVRewriteVisitor traverses the SCEV DAG, it may visit the same SCEV multiple times if this SCEV is referenced by multiple other SCEVs. This has exponential time complexity in the worst case. Memoizing the results will avoid re-visiting the same SCEV.

It's legal to memoize the rewrite results because SCEVRewriteVisitor is state-less, i.e., it should return the same rewrite result each time for the same input SCEV. Currently only memoize results of visitMulExpr because other functions are much cheaper in the worst case.

This patch depends on D25794, and they together fix PR18606.

[sanjoy]

include/llvm/Analysis/ScalarEvolutionExpressions.h

552 ↗

(On Diff #75262)

Design-wise I'm not too happy to do see this happen only for multiplications -- have you considered over-riding the visit function and implementing the general logic there?

585 ↗

(On Diff #75262)

s/it/It/

592 ↗

(On Diff #75262)

I'd use .insert.

[lihuang]

include/llvm/Analysis/ScalarEvolutionExpressions.h

552 ↗

(On Diff #75262)

Override the visit function in SCEVRewriteVisitor class. This function checks the map first, then forward the input to SCEVVisitor::visit.

[sanjoy]

lgtm with nits

include/llvm/Analysis/ScalarEvolutionExpressions.h

540 ↗

(On Diff #75373)

Please document the fact that this only implements "pure" transforms.

560 ↗

(On Diff #75373)

Just to be pedantic, I'd assert here that you did in fact insert a new entry.

test/Analysis/ScalarEvolution/pr18606.ll

3 ↗

(On Diff #75373)

Please add a basic output sanity check.

[lihuang]

include/llvm/Analysis/ScalarEvolutionExpressions.h

540 ↗

(On Diff #75373)

Sorry, but by "pure" transform do you mean identity transform? So what about:

/// Recursively visits a SCEV expression and re-writes an identical SCEV.

560 ↗

(On Diff #75373)

what about:

assert(RewriteResults.insert({S, Result}).second && "Should insert a new entry");

test/Analysis/ScalarEvolution/pr18606.ll

3 ↗

(On Diff #75373)

Okay.

[sanjoy]

include/llvm/Analysis/ScalarEvolutionExpressions.h

540 ↗

(On Diff #75373)

I meant a transform that maps equal inputs to equal outputs. For instance, I cannot use this framework to implement a transform that maps all SCEVConstant s to getConstant(Counter++) since then the caching scheme won't work.

560 ↗

(On Diff #75373)

sgtm

[lihuang]

include/llvm/Analysis/ScalarEvolutionExpressions.h

540 ↗

(On Diff #75373)

I see, thanks. What about:

/// This visitor recursively visits a SCEV expression and re-writes it. The result from each visit is cached, so it will return the same SCEV for the same input.

[sanjoy]

include/llvm/Analysis/ScalarEvolutionExpressions.h

540 ↗

(On Diff #75373)

sgtm

llvm/trunk/include/llvm/Analysis/ScalarEvolutionExpressions.h

namespace llvm {
 
   /// Use SCEVTraversal to visit all nodes in the given expression tree.
   template<typename SV>
   void visitAll(const SCEV *Root, SV& Visitor) {
     SCEVTraversal<SV> T(Visitor);
     T.visitAll(Root);
   }
 
-  /// Recursively visits a SCEV expression and re-writes it.
+  /// This visitor recursively visits a SCEV expression and re-writes it.
+  /// The result from each visit is cached, so it will return the same
+  /// SCEV for the same input.
   template<typename SC>
   class SCEVRewriteVisitor : public SCEVVisitor<SC, const SCEV *> {
   protected:
     ScalarEvolution &SE;
+    // Memoize the result of each visit so that we only compute once for
+    // the same input SCEV. This is to avoid redundant computations when
+    // a SCEV is referenced by multiple SCEVs. Without memoization, this
+    // visit algorithm would have exponential time complexity in the worst
+    // case, causing the compiler to hang on certain tests.
+    DenseMap<const SCEV *, const SCEV *> RewriteResults;
+
   public:
     SCEVRewriteVisitor(ScalarEvolution &SE) : SE(SE) {}
 
+    const SCEV *visit(const SCEV *S) {
+      auto It = RewriteResults.find(S);
+      if (It != RewriteResults.end())
+        return It->second;
+      auto *Result = SCEVVisitor<SC, const SCEV *>::visit(S);
+      assert(RewriteResults.insert({S, Result}).second &&
+             "Should insert a new entry");
+      return Result;
+    }
+
     const SCEV *visitConstant(const SCEVConstant *Constant) {
       return Constant;
     }
 
     const SCEV *visitTruncateExpr(const SCEVTruncateExpr *Expr) {
       const SCEV *Operand = ((SC*)this)->visit(Expr->getOperand());
       return SE.getTruncateExpr(Operand, Expr->getType());
     }

llvm/trunk/test/Analysis/ScalarEvolution/pr18606.ll

+; RUN: opt -S -indvars < %s | FileCheck %s
+
+; CHECK: @main
+; CHECK: %mul.lcssa5 = phi i32 [ %a.promoted4, %entry ], [ %mul.30, %for.body3 ]
+; CEHCK: %mul = mul nsw i32 %mul.lcssa5, %mul.lcssa5
+; CHECK: %mul.30 = mul nsw i32 %mul.29, %mul.29
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+@a = local_unnamed_addr global i32 0, align 4
+@b = local_unnamed_addr global i32 0, align 4
+
+; Function Attrs: norecurse nounwind uwtable
+define i32 @main() local_unnamed_addr {
+entry:
+  %a.promoted4 = load i32, i32* @a, align 4
+  br label %for.cond1.preheader
+
+for.cond1.preheader:                              ; preds = %entry, %for.body3
+  %mul.lcssa5 = phi i32 [ %a.promoted4, %entry ], [ %mul.30, %for.body3 ]
+  %i.03 = phi i32 [ 0, %entry ], [ %inc5, %for.body3 ]
+  br label %for.body3
+
+for.body3:                                        ; preds = %for.cond1.preheader
+  %mul = mul nsw i32 %mul.lcssa5, %mul.lcssa5
+  %mul.1 = mul nsw i32 %mul, %mul
+  %mul.2 = mul nsw i32 %mul.1, %mul.1
+  %mul.3 = mul nsw i32 %mul.2, %mul.2
+  %mul.4 = mul nsw i32 %mul.3, %mul.3
+  %mul.5 = mul nsw i32 %mul.4, %mul.4
+  %mul.6 = mul nsw i32 %mul.5, %mul.5
+  %mul.7 = mul nsw i32 %mul.6, %mul.6
+  %mul.8 = mul nsw i32 %mul.7, %mul.7
+  %mul.9 = mul nsw i32 %mul.8, %mul.8
+  %mul.10 = mul nsw i32 %mul.9, %mul.9
+  %mul.11 = mul nsw i32 %mul.10, %mul.10
+  %mul.12 = mul nsw i32 %mul.11, %mul.11
+  %mul.13 = mul nsw i32 %mul.12, %mul.12
+  %mul.14 = mul nsw i32 %mul.13, %mul.13
+  %mul.15 = mul nsw i32 %mul.14, %mul.14
+  %mul.16 = mul nsw i32 %mul.15, %mul.15
+  %mul.17 = mul nsw i32 %mul.16, %mul.16
+  %mul.18 = mul nsw i32 %mul.17, %mul.17
+  %mul.19 = mul nsw i32 %mul.18, %mul.18
+  %mul.20 = mul nsw i32 %mul.19, %mul.19
+  %mul.21 = mul nsw i32 %mul.20, %mul.20
+  %mul.22 = mul nsw i32 %mul.21, %mul.21
+  %mul.23 = mul nsw i32 %mul.22, %mul.22
+  %mul.24 = mul nsw i32 %mul.23, %mul.23
+  %mul.25 = mul nsw i32 %mul.24, %mul.24
+  %mul.26 = mul nsw i32 %mul.25, %mul.25
+  %mul.27 = mul nsw i32 %mul.26, %mul.26
+  %mul.28 = mul nsw i32 %mul.27, %mul.27
+  %mul.29 = mul nsw i32 %mul.28, %mul.28
+  %mul.30 = mul nsw i32 %mul.29, %mul.29
+  %inc5 = add nuw nsw i32 %i.03, 1
+  %exitcond = icmp ne i32 %inc5, 10
+  br i1 %exitcond, label %for.cond1.preheader, label %for.end6
+
+for.end6:                                         ; preds = %for.body3
+  %mul.lcssa.lcssa = phi i32 [ %mul.30, %for.body3 ]
+  %inc.lcssa.lcssa = phi i32 [ 31, %for.body3 ]
+  store i32 %mul.lcssa.lcssa, i32* @a, align 4
+  store i32 %inc.lcssa.lcssa, i32* @b, align 4
+  ret i32 0
+}