fix tvmrt and tvmcg fot latest tvm

tum-ei-eda · Jul 19, 2022 · 5aaeb13 · 5aaeb13
1 parent f896da1
commit 5aaeb13
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/mlonmcu/flow/tvm/backend/tvmcg.py b/mlonmcu/flow/tvm/backend/tvmcg.py
@@ -42,7 +42,12 @@ class TVMCGBackend(TVMRTBackend):
 
     def get_max_workspace_size_from_metadata(self, metadata):
         max_workspace = 0
-        for op in metadata["memory"]["functions"]["operator_functions"]:
+        if "modules" in metadata:
+            temp = metadata["modules"]["default"]["memory"]
+        else:
+            # backwards compatibility
+            temp = metadata["memory"]
+        for op in temp["functions"]["operator_functions"]:
             max_workspace = max(
                 max_workspace,
                 op["workspace"][0]["workspace_size_bytes"] if len(op["workspace"]) > 0 else 0,
@@ -65,7 +70,8 @@ def generate_code(self):
                     with open(mlf_path / "metadata.json") as handle:
                         metadata = json.load(handle)
                     tvmcg_exe = self.config["utvmcg.exe"]
-                    graph_json_file = mlf_path / "executor-config" / "graph" / "graph.json"
+                    # graph_json_file = mlf_path / "executor-config" / "graph" / "graph.json"
+                    graph_json_file = mlf_path / "executor-config" / "graph" / "default.graph"
                     params_bin_file = mlf_path / "parameters" / "default.params"
                     max_workspace_size = self.get_max_workspace_size_from_metadata(metadata)
                     args = []

diff --git a/mlonmcu/flow/tvm/backend/tvmrt.py b/mlonmcu/flow/tvm/backend/tvmrt.py
@@ -60,7 +60,7 @@ def get_tvmc_compile_args(self, out):
 
     def get_graph_and_params_from_mlf(self, path):
         graph = None
-        with open(Path(path) / "executor-config" / "graph" / "graph.json", "r") as handle:
+        with open(Path(path) / "executor-config" / "graph" / "default.graph", "r") as handle:
             graph = handle.read()
         params = None
         with open(Path(path) / "parameters" / "default.params", "rb") as handle: