You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/01/29 22:00:06 UTC
[GitHub] [tvm] mbrookhart opened a new pull request #7374: [WIP][Relay][Passes] Iterative A-normal Traversals
mbrookhart opened a new pull request #7374:
URL: https://github.com/apache/tvm/pull/7374
On very large models, we can have ASTs that stack overflow because we recursively traverse tens of thousands of chained Let nodes. In A-normal form, we can simply iterate over that chain instead.
These changes find every pass that segfaults with a stack overflow in ONNX SSD-Mobilenet (~20k nodes) and hijacks the visitor to do iterative traversals over chains of lets. This lets me run that model without setting by significantly reducing the stack size.
I also needed to convert FuseOps to mixed mode.
This is still WIP, there's a lot of code copy in here I'd like to clean up and I need to figure out how to test the issues, but I'd like to kick of CI to make sure I didnt' break any edge cases not covered in the pass tests.
Thanks!
cc @jroesch @tqchen @tristan-arm @altanh @zhiics @masahi @kevinthesun
----------------------------------------------------------------
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
[GitHub] [tvm] tmoreau89 merged pull request #7374: [Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
tmoreau89 merged pull request #7374:
URL: https://github.com/apache/tvm/pull/7374
----------------------------------------------------------------
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
[GitHub] [tvm] altanh commented on a change in pull request #7374: [WIP][Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
altanh commented on a change in pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#discussion_r568114858
##########
File path: src/relay/transforms/fold_constant.cc
##########
@@ -92,19 +92,33 @@ class ConstantFolder : public MixedModeMutator {
using MixedModeMutator::VisitExpr_;
Expr VisitExpr_(const LetNode* op) final {
- Expr value = this->Mutate(op->value);
- if (value.as<ConstantNode>()) {
- memo_[op->var] = value;
- return this->Mutate(op->body);
- } else {
- Var var = Downcast<Var>(this->Mutate(op->var));
- Expr body = this->Mutate(op->body);
- if (var.same_as(op->var) && value.same_as(op->value) && body.same_as(op->body)) {
- return GetRef<Expr>(op);
+ auto pre_visit = [this](const LetNode* op) {
+ // Rely on the Memoizer to cache pre-visit values
+ Expr value = this->Mutate(op->value);
+ if (value.as<ConstantNode>()) {
+ memo_[op->var] = value;
Review comment:
explicitly use `this->memo_`
##########
File path: src/relay/transforms/de_duplicate.cc
##########
@@ -61,8 +63,19 @@ Expr DeDup(const Expr& e) {
}
Expr VisitExpr_(const LetNode* op) final {
- Var v = Fresh(op->var);
- return Let(v, VisitExpr(op->value), VisitExpr(op->body));
+ std::unordered_map<Expr, Var, ObjectPtrHash, ObjectPtrEqual> new_vars;
+ auto pre_visit = [this, &new_vars](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ new_vars[expr] = Fresh(op->var);
+ // Rely on the Memoizer to cache pre-visit values
+ VisitExpr(op->value);
+ };
+ auto post_visit = [this, &new_vars](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ memo_[expr] = Let(new_vars[expr], VisitExpr(op->value), VisitExpr(op->body));
Review comment:
explicitly use `this->memo_`
##########
File path: src/relay/transforms/de_duplicate.cc
##########
@@ -61,8 +63,19 @@ Expr DeDup(const Expr& e) {
}
Expr VisitExpr_(const LetNode* op) final {
- Var v = Fresh(op->var);
- return Let(v, VisitExpr(op->value), VisitExpr(op->body));
+ std::unordered_map<Expr, Var, ObjectPtrHash, ObjectPtrEqual> new_vars;
+ auto pre_visit = [this, &new_vars](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ new_vars[expr] = Fresh(op->var);
+ // Rely on the Memoizer to cache pre-visit values
+ VisitExpr(op->value);
Review comment:
explicitly use `this->VisitExpr`
##########
File path: src/relay/transforms/fold_constant.cc
##########
@@ -92,19 +92,33 @@ class ConstantFolder : public MixedModeMutator {
using MixedModeMutator::VisitExpr_;
Expr VisitExpr_(const LetNode* op) final {
- Expr value = this->Mutate(op->value);
- if (value.as<ConstantNode>()) {
- memo_[op->var] = value;
- return this->Mutate(op->body);
- } else {
- Var var = Downcast<Var>(this->Mutate(op->var));
- Expr body = this->Mutate(op->body);
- if (var.same_as(op->var) && value.same_as(op->value) && body.same_as(op->body)) {
- return GetRef<Expr>(op);
+ auto pre_visit = [this](const LetNode* op) {
+ // Rely on the Memoizer to cache pre-visit values
+ Expr value = this->Mutate(op->value);
+ if (value.as<ConstantNode>()) {
+ memo_[op->var] = value;
} else {
- return Let(var, value, body);
+ this->Mutate(op->var);
}
- }
+ };
+ auto post_visit = [this](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ // Rely on the Memoizer to cache pre-visit values
+ Expr value = this->Mutate(op->value);
+ if (value.as<ConstantNode>()) {
+ memo_[expr] = this->Mutate(op->body);
Review comment:
`this->memo_`
##########
File path: src/relay/transforms/fold_constant.cc
##########
@@ -92,19 +92,33 @@ class ConstantFolder : public MixedModeMutator {
using MixedModeMutator::VisitExpr_;
Expr VisitExpr_(const LetNode* op) final {
- Expr value = this->Mutate(op->value);
- if (value.as<ConstantNode>()) {
- memo_[op->var] = value;
- return this->Mutate(op->body);
- } else {
- Var var = Downcast<Var>(this->Mutate(op->var));
- Expr body = this->Mutate(op->body);
- if (var.same_as(op->var) && value.same_as(op->value) && body.same_as(op->body)) {
- return GetRef<Expr>(op);
+ auto pre_visit = [this](const LetNode* op) {
+ // Rely on the Memoizer to cache pre-visit values
+ Expr value = this->Mutate(op->value);
+ if (value.as<ConstantNode>()) {
+ memo_[op->var] = value;
} else {
- return Let(var, value, body);
+ this->Mutate(op->var);
}
- }
+ };
+ auto post_visit = [this](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ // Rely on the Memoizer to cache pre-visit values
+ Expr value = this->Mutate(op->value);
+ if (value.as<ConstantNode>()) {
+ memo_[expr] = this->Mutate(op->body);
+ } else {
+ Var var = Downcast<Var>(this->Mutate(op->var));
+ Expr body = this->Mutate(op->body);
+ if (var.same_as(op->var) && value.same_as(op->value) && body.same_as(op->body)) {
+ memo_[expr] = expr;
+ } else {
+ memo_[expr] = Let(var, value, body);
Review comment:
`this->memo_`
##########
File path: src/relay/transforms/fold_constant.cc
##########
@@ -92,19 +92,33 @@ class ConstantFolder : public MixedModeMutator {
using MixedModeMutator::VisitExpr_;
Expr VisitExpr_(const LetNode* op) final {
- Expr value = this->Mutate(op->value);
- if (value.as<ConstantNode>()) {
- memo_[op->var] = value;
- return this->Mutate(op->body);
- } else {
- Var var = Downcast<Var>(this->Mutate(op->var));
- Expr body = this->Mutate(op->body);
- if (var.same_as(op->var) && value.same_as(op->value) && body.same_as(op->body)) {
- return GetRef<Expr>(op);
+ auto pre_visit = [this](const LetNode* op) {
+ // Rely on the Memoizer to cache pre-visit values
+ Expr value = this->Mutate(op->value);
+ if (value.as<ConstantNode>()) {
+ memo_[op->var] = value;
} else {
- return Let(var, value, body);
+ this->Mutate(op->var);
}
- }
+ };
+ auto post_visit = [this](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ // Rely on the Memoizer to cache pre-visit values
+ Expr value = this->Mutate(op->value);
+ if (value.as<ConstantNode>()) {
+ memo_[expr] = this->Mutate(op->body);
+ } else {
+ Var var = Downcast<Var>(this->Mutate(op->var));
+ Expr body = this->Mutate(op->body);
+ if (var.same_as(op->var) && value.same_as(op->value) && body.same_as(op->body)) {
+ memo_[expr] = expr;
Review comment:
`this->memo_`
##########
File path: src/relay/transforms/fuse_ops.cc
##########
@@ -315,11 +315,20 @@ class IndexedForwardGraph::Creator : private ExprVisitor {
void VisitExpr_(const LetNode* op) final {
// do not fuse through let.
- this->Update(op->var, nullptr, kOpaque);
- this->Update(op->value, nullptr, kOpaque);
- this->Update(op->body, nullptr, kOpaque);
- ExprVisitor::VisitExpr_(op);
- this->AddNode(op);
+ auto pre_visit = [this](const LetNode* op) {
+ // Rely on the Memoizer to cache pre-visit values
+ this->Update(op->var, nullptr, kOpaque);
+ this->Update(op->value, nullptr, kOpaque);
+ this->Update(op->body, nullptr, kOpaque);
+ VisitExpr(op->var);
+ VisitExpr(op->value);
+ };
+ auto post_visit = [this](const LetNode* op) {
+ VisitExpr(op->body);
+ visit_counter_[op] += 1;
+ this->AddNode(op);
Review comment:
explicit `this` in here
----------------------------------------------------------------
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
[GitHub] [tvm] zhiics commented on a change in pull request #7374: [Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#discussion_r568256834
##########
File path: src/relay/transforms/type_infer.cc
##########
@@ -603,7 +611,21 @@ class TypeInferencer::Resolver : public MixedModeMutator, PatternMutator {
Expr Rewrite_(const CallNode* op, const Expr& post) final { return AttachCheckedType(op, post); }
- Expr VisitExpr_(const LetNode* op) final { return AttachCheckedType(op); }
+ Expr VisitExpr_(const LetNode* op) final {
+ auto pre_visit = [this](const LetNode* op) {
+ this->VisitExpr(op->var);
+ this->VisitExpr(op->value);
+ };
+ auto post_visit = [this](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ Var var = Downcast<Var>(this->VisitExpr(op->var));
+ Expr value = this->VisitExpr(op->value);
Review comment:
do we need to visit var and value again in the post?
##########
File path: src/relay/transforms/type_infer.cc
##########
@@ -603,7 +611,21 @@ class TypeInferencer::Resolver : public MixedModeMutator, PatternMutator {
Expr Rewrite_(const CallNode* op, const Expr& post) final { return AttachCheckedType(op, post); }
- Expr VisitExpr_(const LetNode* op) final { return AttachCheckedType(op); }
+ Expr VisitExpr_(const LetNode* op) final {
+ auto pre_visit = [this](const LetNode* op) {
+ this->VisitExpr(op->var);
+ this->VisitExpr(op->value);
+ };
+ auto post_visit = [this](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ Var var = Downcast<Var>(this->VisitExpr(op->var));
+ Expr value = this->VisitExpr(op->value);
Review comment:
i see. thanks.
----------------------------------------------------------------
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
[GitHub] [tvm] zhiics commented on a change in pull request #7374: [Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#discussion_r568256834
##########
File path: src/relay/transforms/type_infer.cc
##########
@@ -603,7 +611,21 @@ class TypeInferencer::Resolver : public MixedModeMutator, PatternMutator {
Expr Rewrite_(const CallNode* op, const Expr& post) final { return AttachCheckedType(op, post); }
- Expr VisitExpr_(const LetNode* op) final { return AttachCheckedType(op); }
+ Expr VisitExpr_(const LetNode* op) final {
+ auto pre_visit = [this](const LetNode* op) {
+ this->VisitExpr(op->var);
+ this->VisitExpr(op->value);
+ };
+ auto post_visit = [this](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ Var var = Downcast<Var>(this->VisitExpr(op->var));
+ Expr value = this->VisitExpr(op->value);
Review comment:
do we need to visit var and value again in the post?
----------------------------------------------------------------
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
[GitHub] [tvm] tmoreau89 merged pull request #7374: [Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
tmoreau89 merged pull request #7374:
URL: https://github.com/apache/tvm/pull/7374
----------------------------------------------------------------
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
[GitHub] [tvm] manupa-arm commented on pull request #7374: [WIP][Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#issuecomment-770751801
Hi @mbrookhart ,
This looks great! In the code clean up, would you be considering factor out let-chain traversal possibly as an utility ?. I see a pattern that the current implementation of Let node is had "pre" and "post" processing stage that is being done prior and after the pushing them to the stack. So I was wondering, whether its possible to expose a generic interface just to provide the pre processing functionality (going down the let chain) and post processing functionality (coming up the let chain). I m not sure how the actual implementation should look like -- do let me know what you have in mind :) .
----------------------------------------------------------------
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
[GitHub] [tvm] tmoreau89 commented on pull request #7374: [Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
tmoreau89 commented on pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#issuecomment-771826972
Thanks @mbrookhart @manupa-arm @zhiics @altanh the PR has been merged!
----------------------------------------------------------------
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
[GitHub] [tvm] zhiics commented on a change in pull request #7374: [Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#discussion_r568342734
##########
File path: src/relay/transforms/type_infer.cc
##########
@@ -603,7 +611,21 @@ class TypeInferencer::Resolver : public MixedModeMutator, PatternMutator {
Expr Rewrite_(const CallNode* op, const Expr& post) final { return AttachCheckedType(op, post); }
- Expr VisitExpr_(const LetNode* op) final { return AttachCheckedType(op); }
+ Expr VisitExpr_(const LetNode* op) final {
+ auto pre_visit = [this](const LetNode* op) {
+ this->VisitExpr(op->var);
+ this->VisitExpr(op->value);
+ };
+ auto post_visit = [this](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ Var var = Downcast<Var>(this->VisitExpr(op->var));
+ Expr value = this->VisitExpr(op->value);
Review comment:
i see. thanks.
----------------------------------------------------------------
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
[GitHub] [tvm] tmoreau89 commented on pull request #7374: [Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
tmoreau89 commented on pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#issuecomment-771826972
Thanks @mbrookhart @manupa-arm @zhiics @altanh the PR has been merged!
----------------------------------------------------------------
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
[GitHub] [tvm] mbrookhart commented on pull request #7374: [WIP][Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#issuecomment-770960941
@manupa-arm Thanks for the suggestion. @altanh and I were talking about doing something similar late last week using higher order functions. I'll let you know what I can figure out.
----------------------------------------------------------------
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
[GitHub] [tvm] mbrookhart commented on a change in pull request #7374: [Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#discussion_r568310278
##########
File path: src/relay/transforms/type_infer.cc
##########
@@ -603,7 +611,21 @@ class TypeInferencer::Resolver : public MixedModeMutator, PatternMutator {
Expr Rewrite_(const CallNode* op, const Expr& post) final { return AttachCheckedType(op, post); }
- Expr VisitExpr_(const LetNode* op) final { return AttachCheckedType(op); }
+ Expr VisitExpr_(const LetNode* op) final {
+ auto pre_visit = [this](const LetNode* op) {
+ this->VisitExpr(op->var);
+ this->VisitExpr(op->value);
+ };
+ auto post_visit = [this](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ Var var = Downcast<Var>(this->VisitExpr(op->var));
+ Expr value = this->VisitExpr(op->value);
Review comment:
We just need to pull the values out of the cache. Instead of maintaining a cache shared by the two lambdas, I'm using the memorization cache in the Mutator. The second time visit is called, it will short circuit and return the previously computed value.
----------------------------------------------------------------
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
[GitHub] [tvm] mbrookhart commented on a change in pull request #7374: [Relay][Passes] Iterative A-normal Traversals
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #7374:
URL: https://github.com/apache/tvm/pull/7374#discussion_r568310278
##########
File path: src/relay/transforms/type_infer.cc
##########
@@ -603,7 +611,21 @@ class TypeInferencer::Resolver : public MixedModeMutator, PatternMutator {
Expr Rewrite_(const CallNode* op, const Expr& post) final { return AttachCheckedType(op, post); }
- Expr VisitExpr_(const LetNode* op) final { return AttachCheckedType(op); }
+ Expr VisitExpr_(const LetNode* op) final {
+ auto pre_visit = [this](const LetNode* op) {
+ this->VisitExpr(op->var);
+ this->VisitExpr(op->value);
+ };
+ auto post_visit = [this](const LetNode* op) {
+ Expr expr = GetRef<Expr>(op);
+ Var var = Downcast<Var>(this->VisitExpr(op->var));
+ Expr value = this->VisitExpr(op->value);
Review comment:
We just need to pull the values out of the cache. Instead of maintaining a cache shared by the two lambdas, I'm using the memorization cache in the Mutator. The second time visit is called, it will short circuit and return the previously computed value.
----------------------------------------------------------------
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