[GitHub] [tvm] manupa-arm commented on a change in pull request #8926: Fix incorrect AOT Memory Planning

[GitBox <gi...@apache.org>]

manupa-arm commented on a change in pull request #8926:
URL: https://github.com/apache/tvm/pull/8926#discussion_r702754213



##########
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(

Review comment:
       Hi @mbs-octoml @electriclilies ,
   
   Ideally, we should run this when we have the the whole of IRModule in TIR which is a state currently only present in the AoT flow -- barring the ability to represent non-scalar constants in TIR (https://github.com/apache/tvm-rfcs/pull/22).
   
   Once, all the executors could could start working based on TIR-only IRModule (after having done many progressive lowerings -- [a relevant discussion](https://github.com/apache/tvm-rfcs/pull/9#discussion_r688253505)), I think we could move UpdateFunctionMetadata (in the same time turning it to a pass) to the tail end of the pipeline.
   
   In the mean time, we can still have the correct workspace sizes updated. For AoT, that should be the TIR main function just prior to the codegen and not the relay one used before.
   
   
   
   




-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org