zhuwenxi commented on a change in pull request #7619:
URL: https://github.com/apache/tvm/pull/7619#discussion_r603074292
##########
File path: src/tir/transforms/lower_tvm_builtin.cc
##########
@@ -48,30 +48,54 @@ inline PrimExpr StackAlloca(std::string type, size_t num) {
// These information are needed during codegen.
class BuiltinLower : public StmtExprMutator {
public:
- Stmt Build(Stmt stmt) {
- stack_shape_ = Var("stack_shape", DataType::Handle());
- stack_array_ = Var("stack_array", DataType::Handle());
- stack_value_ = Var("stack_value", DataType::Handle());
- stack_tcode_ = Var("stack_tcode", DataType::Handle());
+ // Record stack frame for existing scope.
+ struct AllocaScope {
+ Var stack_shape_ = Var("stack_shape", DataType::Handle());
+ Var stack_array_ = Var("stack_array", DataType::Handle());
+ Var stack_value_ = Var("stack_value", DataType::Handle());
+ Var stack_tcode_ = Var("stack_tcode", DataType::Handle());
+
+ int64_t max_shape_stack_{-1};
+ uint64_t max_array_stack_{0};
+ uint64_t max_arg_stack_{0};
+
+ int64_t run_shape_stack_{-1};
+ uint64_t run_array_stack_{0};
+ uint64_t run_arg_stack_{0};
+ };
+
+ Stmt Build(Stmt stmt) { return this->RealizeAlloca(stmt); }
+
+ // Allcoate stack frames, only at parallel-for or root.
+ Stmt RealizeAlloca(Stmt stmt) {
Review comment:
Done.
##########
File path: src/tir/transforms/lower_tvm_builtin.cc
##########
@@ -48,30 +48,54 @@ inline PrimExpr StackAlloca(std::string type, size_t num) {
// These information are needed during codegen.
class BuiltinLower : public StmtExprMutator {
public:
- Stmt Build(Stmt stmt) {
- stack_shape_ = Var("stack_shape", DataType::Handle());
- stack_array_ = Var("stack_array", DataType::Handle());
- stack_value_ = Var("stack_value", DataType::Handle());
- stack_tcode_ = Var("stack_tcode", DataType::Handle());
+ // Record stack frame for existing scope.
+ struct AllocaScope {
+ Var stack_shape_ = Var("stack_shape", DataType::Handle());
+ Var stack_array_ = Var("stack_array", DataType::Handle());
+ Var stack_value_ = Var("stack_value", DataType::Handle());
+ Var stack_tcode_ = Var("stack_tcode", DataType::Handle());
+
+ int64_t max_shape_stack_{-1};
+ uint64_t max_array_stack_{0};
+ uint64_t max_arg_stack_{0};
+
+ int64_t run_shape_stack_{-1};
+ uint64_t run_array_stack_{0};
+ uint64_t run_arg_stack_{0};
+ };
+
+ Stmt Build(Stmt stmt) { return this->RealizeAlloca(stmt); }
+
+ // Allcoate stack frames, only at parallel-for or root.
+ Stmt RealizeAlloca(Stmt stmt) {
+ alloca_scope_.emplace_back();
stmt = this->VisitStmt(stmt);
- // create a shape var if any shape is made (including scalar shapes)
- if (max_shape_stack_ != -1) {
- stmt = LetStmt(stack_shape_, StackAlloca("shape", max_shape_stack_),
stmt);
+ auto& scope = alloca_scope_.back();
Review comment:
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.
For queries about this service, please contact Infrastructure at:
[email protected]
