areusch commented on a change in pull request #7533:
URL: https://github.com/apache/tvm/pull/7533#discussion_r586855463
##########
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:
I think the internals here are all subject to being rewritten when we
broaden support. right now i'm just bolting this on until it's better
understood what kind of standardized data structure should be returned from
GraphPlanMemory. I agree with your insight, though--directly building this from
the memory planner output would be simpler and easier to maintain. I think as
we move to add e.g. tensor pinning, we can revisit this.
----------------------------------------------------------------
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:
[email protected]
