Lunderberg commented on a change in pull request #10558:
URL: https://github.com/apache/tvm/pull/10558#discussion_r825951343
##########
File path: src/tir/transforms/lower_vtcm_alloc.cc
##########
@@ -34,12 +34,11 @@ class VtcmAllocator : public StmtExprMutator {
VtcmAllocator() {}
Stmt VisitStmt_(const AllocateNode* op) final {
- Stmt body = this->VisitStmt(op->body);
std::string storage_scope = GetStorageScope(op->buffer_var);
- Stmt stmt = StmtExprMutator::VisitStmt_(op);
- op = stmt.as<AllocateNode>();
+ Stmt stmt = StmtExprMutator::VisitStmt_(op);
if (IsVtcmStorage(storage_scope)) {
+ Stmt body = this->VisitStmt(op->body);
Review comment:
This might be a silly question, but is this second `VisitStmt` necessary
at all? The earlier `StmtExprMutator::VisitStmt_` walks through the entire
subtree, so I think this walks through the tree twice for every vtcm
allocation. I'd recommend one of the two changes:
* Option A: In the if-block, end by `return LetStmt(...)`, then make an
else-block with `return StmtExprMutator::VisitStmt_(op)`. This performs a
traversal either in `if` or the `else`, but never performs a second traversal.
* Option B: Outside the if-block, replace `Stmt stmt =
StmtExprMutator::VisitStmt_(op)` with `Allocate alloc =
Downcast<Allocate>(StmtExprMutator::VisitStmt_(op))`. In the if-block, remove
the `Stmt body = this->VisitStmt(op->body)`, and replace all instances of `op`
with `alloc`. This uses the results of the first traversal, rather than
performing an additional traversal.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
