From 6b11cd2f9b2bf98a77922e1508ba7cf10eb6b70f Mon Sep 17 00:00:00 2001 From: Giuseppe Rossini Date: Mon, 29 Mar 2021 20:15:59 +0100 Subject: [PATCH 1/3] AOT] Remove lookup parameter function in AOT This PR aims at removing the function call to extract the parameters within the AOT main function by introducing a tir::lookup_param builtin. This has different benefits: - In AOT we now only use the v_handle field - We save cycles by not calling an intermediate function to extract local parameters - We reduce code size, since we don't need to pack a call to extract parameters and we don't need to produce the lookup_param function anymore within the compilation unit Change-Id: I36c2f0724a79606424a4374f4f5cd669bb2a8a55 --- include/tvm/tir/builtin.h | 8 +++++ src/relay/backend/aot_executor_codegen.cc | 25 +++------------ src/target/source/codegen_c.cc | 5 +++ src/target/source/codegen_c_host.cc | 37 ++++++++++++++--------- src/target/source/codegen_c_host.h | 1 + src/tir/op/builtin.cc | 4 +++ 6 files changed, 44 insertions(+), 36 deletions(-) diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h index aab5d662d49c..61280d33f1df 100644 --- a/include/tvm/tir/builtin.h +++ b/include/tvm/tir/builtin.h @@ -278,6 +278,14 @@ TVM_DLL const Op& tvm_struct_get(); */ TVM_DLL const Op& tvm_struct_set(); +/*! + * \brief See pseudo code + * Type lookup_param(String param_name) { + * return __tvm_param__param_name; + * } + */ +TVM_DLL const Op& lookup_param(); + /*! * \brief See pesudo code * diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc index 573206c782f6..a48a8a847c7c 100644 --- a/src/relay/backend/aot_executor_codegen.cc +++ b/src/relay/backend/aot_executor_codegen.cc @@ -152,40 +152,23 @@ class AOTExecutorCodegen : public ExprVisitor { * \return Variable that represents the DLTensor associated with the parameters */ tir::Var PackParam(Expr expr) { - // TODO(giuseros): Using call_extern to call into lookup_linked_param. This is because the - // builtin::ret is not supported yet in the c target. Once return is supported we can use - // tvm_call_packed_lowered(). int param_sid = param_storage_ids_[params_by_expr_[expr]]; - auto lookup_linked_param_fn = tir::StringImm(::tvm::runtime::symbol::tvm_lookup_linked_param); auto param_array = te::Var(MakeString("param_", param_sid, "_array"), DataType::Handle()); // Compose the lookup_call using a local stack Array lookup_call; - auto param_var = te::Var(MakeString("param_", param_sid, "_value"), DataType::Handle()); - auto ret_var = te::Var("ret_value", DataType::Handle()); - auto ret_code = te::Var("ret_value", DataType::Handle()); - - lookup_call.push_back(tir::Evaluate( - tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(), - {param_var, 0, tir::builtin::kTVMValueContent, ConstInt32(param_sid)}))); - lookup_call.push_back(tir::Evaluate( - tvm::tir::Call(DataType::Handle(), tir::builtin::call_extern(), - {lookup_linked_param_fn, param_var, 0, 0, ret_var, ret_code, 0}))); - auto ret_var_handle = tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_get(), - {ret_var, 0, tir::builtin::kTVMValueContent}); - // Set the param to the value returned by lookup_call + auto param_handle = tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::lookup_param(), + {tir::StringImm(params_by_expr_[expr])}); + tvm::PrimExpr set_param_array = tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(), - {param_array, 0, tir::builtin::kArrData, ret_var_handle}); + {param_array, 0, tir::builtin::kArrData, param_handle}); lookup_call.push_back(tir::Evaluate(set_param_array)); tir::Stmt lookup_body = tir::SeqStmt(lookup_call); // Allocate the DLTensors on the stack - lookup_body = tir::LetStmt(param_var, StackAlloca("arg_value", 1), lookup_body); - lookup_body = tir::LetStmt(ret_var, StackAlloca("arg_value", 1), lookup_body); - lookup_body = tir::LetStmt(ret_code, StackAlloca("arg_value", 1), lookup_body); lookup_body = tir::LetStmt(param_array, StackAlloca("arg_value", 1), lookup_body); stmts_.push_back(lookup_body); return param_array; diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc index 1627b6003391..d4d0e54c6db4 100644 --- a/src/target/source/codegen_c.cc +++ b/src/target/source/codegen_c.cc @@ -662,6 +662,11 @@ void CodeGenC::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT(*) os << " != "; this->PrintExpr(op->args[0], os); os << ")"; + } else if (op->op.same_as(builtin::lookup_param())) { + ICHECK_EQ(op->args.size(), 1); + const StringImmNode* str = op->args[0].as(); + ICHECK(str != nullptr); + os << "__tvm_param__" << str->value; } else { LOG(FATAL) << "Unresolved call " << op->op; } diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc index 0bfbade23f01..8b15a4170dcc 100644 --- a/src/target/source/codegen_c_host.cc +++ b/src/target/source/codegen_c_host.cc @@ -61,20 +61,7 @@ void CodeGenCHost::AddFunction(const PrimFunc& f) { CodeGenC::AddFunction(f); } -void CodeGenCHost::LinkParameters(Map params) { - PrintFuncPrefix(); - stream << " " << tvm::runtime::symbol::tvm_lookup_linked_param - << "(void* args, int* arg_type_ids, int num_args, void* out_ret_value, " - << "int* out_ret_tcode, void* resource_handle) {\n"; - ICHECK_EQ(GetUniqueName(tvm::runtime::symbol::tvm_lookup_linked_param), - tvm::runtime::symbol::tvm_lookup_linked_param) - << "builtin PackedFunc name already taken: " << tvm::runtime::symbol::tvm_lookup_linked_param; - stream << " switch (((int64_t*) args)[0]) {\n" - << " default:\n" - << " out_ret_tcode[0] = " << kTVMNullptr << ";\n" - << " return 0;\n"; - - function_names_.push_back(tvm::runtime::symbol::tvm_lookup_linked_param); +void CodeGenCHost::DeclareParameters(Map params) { for (auto kv : params) { decl_stream << "\n" << "#ifdef __cplusplus\n" @@ -93,6 +80,24 @@ void CodeGenCHost::LinkParameters(Map params) { << "#ifdef __cplusplus\n" << "} // extern \"C\"\n" << "#endif\n"; + } +} + +void CodeGenCHost::LinkParameters(Map params) { + PrintFuncPrefix(); + stream << " " << tvm::runtime::symbol::tvm_lookup_linked_param + << "(void* args, int* arg_type_ids, int num_args, void* out_ret_value, " + << "int* out_ret_tcode, void* resource_handle) {\n"; + ICHECK_EQ(GetUniqueName(tvm::runtime::symbol::tvm_lookup_linked_param), + tvm::runtime::symbol::tvm_lookup_linked_param) + << "builtin PackedFunc name already taken: " << tvm::runtime::symbol::tvm_lookup_linked_param; + stream << " switch (((int64_t*) args)[0]) {\n" + << " default:\n" + << " out_ret_tcode[0] = " << kTVMNullptr << ";\n" + << " return 0;\n"; + + function_names_.push_back(tvm::runtime::symbol::tvm_lookup_linked_param); + for (auto kv : params) { stream << " case " << kv.second->id << ":\n" << " ((uint64_t*)out_ret_value)[0] = (uint64_t) (uintptr_t) " << ::tvm::runtime::symbol::tvm_param_prefix << kv.first << ";\n" @@ -398,12 +403,14 @@ runtime::Module BuildCHost(IRModule mod, Target target) { cg.AddFunction(f); } - if (could_have_linked_params) { + if (could_have_linked_params && !aot_executor_fn.defined()) { ICHECK(found_linked_params) << "-link-params given but none found"; + cg.DeclareParameters(linked_params); cg.LinkParameters(linked_params); } if (aot_executor_fn.defined()) { + cg.DeclareParameters(linked_params); cg.AddFunction(aot_executor_fn); } diff --git a/src/target/source/codegen_c_host.h b/src/target/source/codegen_c_host.h index 2ee31b8c7e0e..e54d78030ed9 100644 --- a/src/target/source/codegen_c_host.h +++ b/src/target/source/codegen_c_host.h @@ -43,6 +43,7 @@ class CodeGenCHost final : public CodeGenC { void AddFunction(const PrimFunc& f); /*! \brief Add linked parameters, if they are present. */ + void DeclareParameters(Map params); void LinkParameters(Map params); void PrintType(DataType t, std::ostream& os) final; // NOLINT(*) diff --git a/src/tir/op/builtin.cc b/src/tir/op/builtin.cc index f3ab78f89bec..f0ca04cbd5fd 100644 --- a/src/tir/op/builtin.cc +++ b/src/tir/op/builtin.cc @@ -155,6 +155,10 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_struct_set) .set_num_inputs(4) .set_attr("TCallEffectKind", Integer(CallEffectKind::kUpdateState)); +TIR_DEFINE_BUILTIN_FUNC(lookup_param) + .set_num_inputs(4) + .set_attr("TCallEffectKind", Integer(CallEffectKind::kUpdateState)); + TIR_DEFINE_BUILTIN_FUNC(tvm_throw_last_error) .set_num_inputs(0) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); From c7109ca886f34f8fbecf062c692a4266dd56f630 Mon Sep 17 00:00:00 2001 From: Giuseppe Rossini Date: Tue, 18 May 2021 22:55:24 +0100 Subject: [PATCH 2/3] addressing comments Change-Id: I83ba0189f559d310b5a80fe0bcc4d601b490d21a --- src/target/source/codegen_c_host.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc index 8b15a4170dcc..03fef4709b5e 100644 --- a/src/target/source/codegen_c_host.cc +++ b/src/target/source/codegen_c_host.cc @@ -409,7 +409,7 @@ runtime::Module BuildCHost(IRModule mod, Target target) { cg.LinkParameters(linked_params); } - if (aot_executor_fn.defined()) { + if (could_have_linked_params && aot_executor_fn.defined()) { cg.DeclareParameters(linked_params); cg.AddFunction(aot_executor_fn); } From 3d968b1cf5a1150d012401749264ee0022c4713b Mon Sep 17 00:00:00 2001 From: Giuseppe Rossini Date: Wed, 19 May 2021 09:52:05 +0100 Subject: [PATCH 3/3] retrigger CI Change-Id: I84ab4a526d1284ded41fe95636e94c15412f6b28