[GitHub] [tvm] areusch commented on a change in pull request #8014: [AOT] Name mangling in AOT

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

areusch commented on a change in pull request #8014:
URL: https://github.com/apache/tvm/pull/8014#discussion_r631378115



##########
File path: python/tvm/micro/model_library_format.py
##########
@@ -35,7 +35,7 @@ class UnsupportedInModelLibraryFormatError(Exception):
     """Raised when export_model_library_format does not support the given Module tree."""
 
 
-def _populate_codegen_dir(mod, codegen_dir: str):
+def _populate_codegen_dir(mod, codegen_dir: str, model_name: str = None):

Review comment:
       can you add model_name to the parameters docstring?

##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -503,16 +506,35 @@ Pass PartitionGraph() {
     return partitioning::RemoveDefaultAnnotations(m);
   };
 
-  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> part_func =
-      [=](IRModule m, PassContext pc) { return partitioning::Partitioner(m).Partition(); };
+  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> part_func = [=](IRModule m,
+                                                                            PassContext pc) {
+    return partitioning::Partitioner(m, mangle_fn).Partition();
+  };
 
   auto flatten_tuples_pass = CreateModulePass(flatten_tuples, 0, "FlattenNestedTuples", {});
   auto remove_default_pass = CreateModulePass(remove_defaults, 0, "RemoveDefaultAnnotations", {});
   auto partition_pass = CreateModulePass(part_func, 0, "PartitionGraph", {});
   return Sequential({flatten_tuples_pass, remove_default_pass, partition_pass, InferType()});
 }
 
+Pass PartitionGraph() {
+  // Default version. All the function signatures will be "COMPILER_ID"

Review comment:
       I think this comment is confusing because it uses "function signatures" and `COMPILER_ID` which isn't obvious as a placeholder (i think?) unless you also read the following comment. Could you clarify?

##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -503,16 +506,35 @@ Pass PartitionGraph() {
     return partitioning::RemoveDefaultAnnotations(m);
   };
 
-  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> part_func =
-      [=](IRModule m, PassContext pc) { return partitioning::Partitioner(m).Partition(); };
+  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> part_func = [=](IRModule m,
+                                                                            PassContext pc) {
+    return partitioning::Partitioner(m, mangle_fn).Partition();
+  };
 
   auto flatten_tuples_pass = CreateModulePass(flatten_tuples, 0, "FlattenNestedTuples", {});
   auto remove_default_pass = CreateModulePass(remove_defaults, 0, "RemoveDefaultAnnotations", {});
   auto partition_pass = CreateModulePass(part_func, 0, "PartitionGraph", {});
   return Sequential({flatten_tuples_pass, remove_default_pass, partition_pass, InferType()});
 }
 
+Pass PartitionGraph() {
+  // Default version. All the function signatures will be "COMPILER_ID"
+  auto mangle_fn = [](String name) { return name; };
+  return PartitionGraphCommon(mangle_fn);
+}
+
+Pass PartitionGraphWithModName(String mod_name) {

Review comment:
       it seems like it may be a better design to create a "mangle" pass. curious what your thoughts are on this?

##########
File path: src/target/source/codegen_c_host.cc
##########
@@ -61,8 +62,14 @@ void CodeGenCHost::AddFunction(const PrimFunc& f) {
   CodeGenC::AddFunction(f);
 }
 
-void CodeGenCHost::LinkParameters(Map<String, LinkedParam> params) {
-  PrintFuncPrefix();
+void CodeGenCHost::LinkParameters(Map<String, LinkedParam> params, bool static_linking) {
+  // TODO(giuseros): remove this once we don't need an external function call to
+  // retrieve the parameters.
+  if (static_linking) {

Review comment:
       can you add documentation explaining why we need this and what it does?

##########
File path: python/tvm/micro/model_library_format.py
##########
@@ -55,17 +55,19 @@ def _populate_codegen_dir(mod, codegen_dir: str):
 
     mod_indices = {"lib": 0, "src": 0}
     host_codegen_dir = os.path.join(codegen_dir, "host")
+    lib_name = f"{model_name}_lib" if model_name else "lib"
+
     for dso_mod in dso_modules:
         if dso_mod.type_key == "c":
             index = mod_indices["src"]
             mod_indices["src"] += 1
             parent_dir = os.path.join(host_codegen_dir, "src")
-            file_name = os.path.join(parent_dir, f"lib{index}.c")
+            file_name = os.path.join(parent_dir, f"{lib_name}{index}.c")

Review comment:
       i think we should rev the model library format version

##########
File path: python/tvm/relay/build_module.py
##########
@@ -115,6 +117,9 @@ def build(self, mod, target=None, target_host=None, params=None, executor="graph
             - If "graph" is specified, then the graph_executor will be used
             - If "aot" is specified, then the aot_executor will be used
 
+        mod_name: Optional[str]
+            The module name we will build

Review comment:
       are there restrictions on this? shall we spell them out? do we need to validate the model name to ensure it mangles correctly or otherwise apply some mangling?

##########
File path: tests/python/relay/aot/aot_test_utils.py
##########
@@ -35,6 +35,13 @@
 from tvm.micro import export_model_library_format
 
 
+def mangle_name(mod_name, name):

Review comment:
       should this become a util function rather than a test function?

##########
File path: src/target/source/codegen_c_host.cc
##########
@@ -386,7 +393,7 @@ runtime::Module BuildCHost(IRModule mod, Target target) {
     // symbols are available to tvm_run_func
     auto fun_name = std::string(kv.first->name_hint);
     const bool is_aot_executor_fn =
-        (fun_name.rfind(::tvm::runtime::symbol::tvm_run_func_prefix, 0) == 0);
+        (fun_name.rfind(::tvm::runtime::symbol::tvm_run_func_suffix) != std::string::npos);

Review comment:
       perhaps it's worth identifying this with an attribute, since it doesn't have a fixed name?




-- 
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.

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