[GitHub] [tvm] mbrookhart commented on a change in pull request #7558: [RELAY] Modify some passes to not stack overflow on many lets.

[GitBox <gi...@apache.org>]

mbrookhart commented on a change in pull request #7558:
URL: https://github.com/apache/tvm/pull/7558#discussion_r585087701



##########
File path: src/relay/transforms/dead_code.cc
##########
@@ -77,16 +83,26 @@ class Eliminator : private ExprMutator {
 
   Expr VisitExpr_(const VarNode* op) final {
     Var v = GetRef<Var>(op);
-    return (expr_map_.count(v) == 0 || HasLet(v)) ? v : VisitExpr(expr_map_[v]);
+    return (expr_map_.count(v) == 0 || HasLet(v)) ? v : VisitExpr(expr_map_.at(v));
   }
 
   Expr VisitExpr_(const LetNode* op) final {
-    Var v = op->var;
-    if (HasLet(v)) {
-      return Let(v, VisitExpr(op->value), VisitExpr(op->body));
-    } else {
-      return VisitExpr(op->body);
-    }
+    auto pre_visit = [this](const LetNode* op) { Expr value = this->VisitExpr(op->value); };
+    auto post_visit = [this](const LetNode* op) {
+      // Rely on the Memoizer to cache pre-visit values
+      Expr value = this->VisitExpr(op->value);
+      // Visit body and cache the op
+      Expr body = this->VisitExpr(op->body);
+      auto expr = GetRef<Expr>(op);
+      Var v = op->var;
+      if (HasLet(v)) {
+        this->memo_[expr] = Let(v, value, body);
+      } else {
+        this->memo_[expr] = body;
+      }
+    };
+    ExpandANormalForm(op, pre_visit, post_visit);
+    return memo_.at(GetRef<Expr>(op));

Review comment:
       `[]` instead of `.at` for performance?

##########
File path: src/relay/transforms/dead_code.cc
##########
@@ -77,16 +83,26 @@ class Eliminator : private ExprMutator {
 
   Expr VisitExpr_(const VarNode* op) final {
     Var v = GetRef<Var>(op);
-    return (expr_map_.count(v) == 0 || HasLet(v)) ? v : VisitExpr(expr_map_[v]);
+    return (expr_map_.count(v) == 0 || HasLet(v)) ? v : VisitExpr(expr_map_.at(v));

Review comment:
       go back to `[]` instead of `.at` for performance?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org