electriclilies commented on a change in pull request #8974:
URL: https://github.com/apache/tvm/pull/8974#discussion_r708585393
##########
File path: src/relay/backend/te_compiler.h
##########
@@ -158,6 +149,16 @@ void UpdateFunctionMetadata(Function relay_func,
*/
Target GetTargetFromInteger(DLDeviceType dev_type, TargetMap targets);
+/*!
+ * \brief Update the "main" control function's metadata
+ *
+ * \param mod The module
+ * \param targets Map of targets
+ * \return function_infos Function info for each function in the module
+ */
+backend::FunctionInfo UpdateMainWorkspaceSize(const IRModule& mod,
tec::TargetMap targets,
+ Map<Expr, backend::StorageInfo>
storage_info_map);
+
Review comment:
Do we want UpdateMainWorkspaceSize to be in the header file?
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -659,8 +666,11 @@ class AOTExecutorCodegen : public MixedModeVisitor {
Downcast<tir::PrimFunc>(mod_run->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix)),
workspace_byte_alignment);
+ lowered_mod = WithAttr(lowered_mod, "main_func_info", func_info);
+
Optional<backend::FunctionInfo> main_func_info =
lowered_mod->GetAttr<backend::FunctionInfo>("main_func_info");
+
Review comment:
Like I said above, can you attach the func_info right before LowerTEPass
is called? And we can then remove the check about whether the attribute is on
the module.
##########
File path: src/relay/backend/utils.h
##########
@@ -146,6 +146,17 @@ struct LoweredOutput {
runtime::Metadata metadata;
};
+/*!
+ * \brief This class is needed to avoid a GCC 5 bug that prevents maps
containing enums from being
+ compiled. If i386 GCC version is increased, we can remove it.
+ */
+struct EnumClassHash {
+ template <typename T>
+ std::size_t operator()(T t) const {
+ return static_cast<std::size_t>(t);
+ }
+};
Review comment:
Thanks for making this change!
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -221,8 +221,16 @@ class GraphExecutorCodegen : public
backend::MemoizedExprTranslator<std::vector<
device_context_map.insert({expr, dev});
}
+ backend::FunctionInfo func_info;
+
+ if (memory_plan_.defined()) {
+ // TODO(@electriclilies, @jroesch): remove UpdateMainWorkspaceSize
+ func_info =
+ relay::tec::UpdateMainWorkspaceSize(mod, targets_,
memory_plan_->expr_to_storage_info);
+ }
+
Review comment:
Same comment as in aot_executor_codegen -- can we put the func_info on
the module here instead of after LowerTEPass is called and delete the check for
main_func_info being set?
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -583,8 +583,15 @@ class AOTExecutorCodegen : public MixedModeVisitor {
// performing the preexisting AOT executor code generation phase.
IRModule mod = IRModule::FromExpr(func);
+ backend::FunctionInfo func_info;
+
+ if (memory_plan.defined()) {
+ // TODO(@electriclilies, @jroesch): remove UpdateMainWorkspaceSize
+ func_info = tec::UpdateMainWorkspaceSize(mod, targets_,
memory_plan->expr_to_storage_info);
+ }
Review comment:
Can you just put the `func_info` on the mod here before passing the
module into LowerTE? Then you don't need to re-extract it later, and also the
logic surrounding func_info is all in one place. (LowerTEPass should preserve
all attributes on modules passed into it)
--
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]
