Mousius commented on a change in pull request #8926:
URL: https://github.com/apache/tvm/pull/8926#discussion_r702172139
##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -656,6 +651,20 @@ class AOTExecutorCodegen : public MixedModeVisitor {
auto storage_rewrite = tir::transform::StorageRewrite();
mod_run = storage_rewrite(mod_run);
+ // The workspace for main function should be calculated after performing
storage_rewrite for
+ // the top level TIR function.
+ auto workspace_byte_alignment =
+
target_host_->GetAttr<Integer>("workspace-byte-alignment").value_or(16);
+ Integer main_workspace_size = CalculateWorkspaceBytes(
+
Downcast<tir::PrimFunc>(mod_run->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix)),
+ workspace_byte_alignment);
+
+ Optional<backend::FunctionInfo> main_func_info =
+ lowered_mod->GetAttr<backend::FunctionInfo>("main_func_info");
+ ICHECK(main_func_info) << "The attribute \"main_func_info\" should be set
at this point.";
+ main_func_info.value()->workspace_sizes.Set(target_host_,
main_workspace_size);
+ function_metadata_.Set(runtime::symbol::tvm_module_main,
main_func_info.value());
Review comment:
Ideally it'd only be computed once, the issue is that the
`StorageRewrite` phase optimises allocations so changes the overall shape of
the workspace memory (in the case of mobilenet it shrinks drastically). I'm
unsure what the correct fix is as AOT runs `StorageRewrite` just above this
but other executors run it as part of `driver_api.cc` later
(https://github.com/apache/tvm/blob/c6f62aafc91e2600ed7772597fd4238c924c2a1b/src/driver/driver_api.cc#L242)
This does speak to @mbs-octoml's comment though, the workspace sizes for the
other executors are likely also wrong - we just don't use them as yet.
--
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]
