manupa-arm commented on a change in pull request #7533:
URL: https://github.com/apache/tvm/pull/7533#discussion_r585371881
##########
File path: python/tvm/runtime/module.py
##########
@@ -370,6 +370,33 @@ def export_library(self, file_name, fcompile=None,
addons=None, workspace_dir=No
return fcompile(file_name, files, **kwargs)
+ def export_model_library_format(self, codegen_dir: str):
Review comment:
I understand that this does not support non-DSO-exportable models yet.
I think it should be better to throw a not implemented error if the current
runtime.Module contains those rather than providing the illusion of successful
export of Model Library format. WDYT?
##########
File path: python/tvm/relay/backend/graph_runtime_factory.py
##########
@@ -56,6 +67,85 @@ def __init__(self, graph_json_str, libmod, libmod_name,
params):
def export_library(self, file_name, fcompile=None, addons=None, **kwargs):
return self.module.export_library(file_name, fcompile, addons,
**kwargs)
+ def _build_memory_map(self):
+ graph = json.loads(self.graph_json)
+
+ seen_storage_ids = set()
+ memory_map = []
+ for node_id, storage_id in enumerate(graph["attrs"]["storage_id"][1]):
Review comment:
See my comment below -- this could be simplified if we can have
relay.Expr --> sid Map somewhere accessible and use that to create the json
later while this function also being another consumer of that map rather than
parsing the json and extracting size information out of it. WDYT?
##########
File path: python/tvm/relay/backend/graph_runtime_factory.py
##########
@@ -56,6 +67,85 @@ def __init__(self, graph_json_str, libmod, libmod_name,
params):
def export_library(self, file_name, fcompile=None, addons=None, **kwargs):
return self.module.export_library(file_name, fcompile, addons,
**kwargs)
+ def _build_memory_map(self):
+ graph = json.loads(self.graph_json)
+
+ seen_storage_ids = set()
+ memory_map = []
+ for node_id, storage_id in enumerate(graph["attrs"]["storage_id"][1]):
+ if storage_id in seen_storage_ids:
+ continue
+
+ seen_storage_ids.add(storage_id)
+ num_elements = 1
+ for dim in graph["attrs"]["shape"][1][storage_id]:
+ num_elements *= dim
+
+ dltype = graph["attrs"]["dltype"][1][storage_id]
+ m = re.match(r"^[a-zA-Z]+([0-9]+)$", dltype)
+ assert m, f"Exported graph contains unknown dltype {dltype}"
+
+ elem_bits = int(m.group(1))
+
+ map_entry = {
+ "storage_id": storage_id,
+ "size_bytes": (num_elements * elem_bits + 7) // 8,
+ }
+ if node_id in graph["arg_nodes"]:
+ map_entry["input_binding"] = graph["nodes"][node_id]["name"]
+
+ memory_map.append(map_entry)
+
+ return memory_map
+
+ def export_model_library_format(self, file_name):
Review comment:
[Clarification] Would it be the expectation that we would need to
implement a simliar function in aot_runtime_factory (or whatever the runtime
family categorization we finally agree) sometime later ?
Im asking this because the current implementation seems to have a dependency
on the presense of the json (which is truly derived from running relay graph
plan memory on main function). So based on the anwser to the above question,
Yes -- then this is fine as the graph runtime is always going to have json
and aot runtime implement this differently.
No -- then we might need to capture the sids prior to the creation of json.
Having said that, I would personally prefer to use relay.Expr --> sid map to
generate the memory map.
WDYT ?
----------------------------------------------------------------
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:
us...@infra.apache.org
