Mousius commented on a change in pull request #8926:
URL: https://github.com/apache/tvm/pull/8926#discussion_r702332909
##########
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:
I've just had another look and I think the other function metadata is
likely fine, as `LowerSchedule` is called from `te_compiler.cc`:
https://github.com/apache/tvm/blob/0744641389e6bf55fc5f2d46b9c08cf60aa14b75/src/relay/backend/te_compiler.cc#L252
Which then calls `CreatePassList`, so for any other function except the AOT
main, the TE Compiler runs the `StorageRewrite`, and then the metadata is
updated afterwards:
https://github.com/apache/tvm/blob/0744641389e6bf55fc5f2d46b9c08cf60aa14b75/src/relay/backend/te_compiler.cc#L431
So the AOT main information is a special case we'll need to think about
:thinking: but at least the other functions should be correct :smile_cat:
--
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]
