Lunderberg commented on code in PR #16184:
URL: https://github.com/apache/tvm/pull/16184#discussion_r1432176139
##########
src/tir/ir/specialize.cc:
##########
@@ -140,16 +141,54 @@ class PrimFuncSpecializer : public StmtExprMutator {
}
}
+ Stmt VisitStmt_(const DeclBufferNode* op) final {
+ // Visit the buffer before delegating to StmtExprMutator, so the
+ // buffer's replacement will be defined before the point of use.
+ Var old_buffer_var = op->buffer->data;
+ Buffer new_buf = MutateAllocBuffer(op->buffer);
+
+ auto node = Downcast<DeclBuffer>(StmtExprMutator::VisitStmt_(op));
+
+ if (!new_buf.same_as(node->buffer)) {
+ node.CopyOnWrite()->buffer = new_buf;
+ }
+
+ // If the buffer variable is begin remapped to an expression, we
Review Comment:
Thank you, and fixed.
--
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]
