manupa-arm commented on a change in pull request #8096:
URL: https://github.com/apache/tvm/pull/8096#discussion_r659942293
##########
File path: src/tir/transforms/legalize_packed_calls.cc
##########
@@ -60,30 +60,41 @@ class PackedCallLegalizer : public StmtExprMutator {
if (call) {
if (call->op.same_as(builtin::tvm_call_cpacked())) {
Array<PrimExpr> packed_args{call->args[0]};
+ std::vector<tir::Var> tvm_values;
for (unsigned i = 1; i < call->args.size(); i++) {
// No need to pack inputs of the prim_func
if (inputs_[call->args[i]] == true) {
packed_args.push_back(call->args[i]);
} else {
// Pack the argument inside a TVMValue
- auto sid_array = tir::Var("tvm_value", DataType::Handle());
- tir::Stmt set_struct_stmt = tir::Evaluate(
+ std::stringstream ss;
+ ss << "tvm_value_" << tvm_value_index_++;
+ auto sid_array = tir::Var(ss.str(), DataType::Handle());
+ tvm_values.push_back(sid_array);
+
+ new_stmts.push_back(tir::Evaluate(
tvm::tir::Call(DataType::Handle(),
tvm::tir::builtin::tvm_struct_set(),
- {sid_array, 0, tir::builtin::kArrData,
call->args[i]}));
- new_stmts.push_back(LetStmt(sid_array, StackAlloca("array", 1),
set_struct_stmt));
+ {sid_array, 0, tir::builtin::kArrData,
call->args[i]})));
packed_args.push_back(sid_array);
}
}
- // Finally, evaluate the packed call and return a sequential statement
+ // Evaluate the packed call
new_stmts.push_back(tir::Evaluate(tir::Call(call->dtype, call->op,
packed_args)));
- return tir::SeqStmt(new_stmts);
+ tir::Stmt call_stmt = tir::SeqStmt(new_stmts);
+
+ // Allocate the TVMValues on the stack and define the variables
+ for (auto v : tvm_values) {
+ call_stmt = LetStmt(v, StackAlloca("array", 1), call_stmt);
Review comment:
yes, that sounds good! (not sure we want a seperate pool but I can see
them pooled to 'a' workspace buffer)
--
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]
