allow overwriting url to tophub repo

tum-ei-eda · Jun 27, 2022 · 357d7f4 · 357d7f4
1 parent 3401cc2
commit 357d7f4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/mlonmcu/flow/tvm/backend/backend.py b/mlonmcu/flow/tvm/backend/backend.py
@@ -58,6 +58,7 @@ class TVMBackend(Backend):
         "use_tuning_results": False,
         "tvmc_extra_args": [],  # Currently compile subcommand only!
         "tvmc_custom_script": None,
+        "tophub_url": None,  # See https://github.com/apache/tvm/blob/1115fd9bc261619ffa0539746ae0aebc46232dc6/python/tvm/autotvm/tophub.py#L38
         **{("autotuning_" + key): value for key, value in TVMTuner.DEFAULTS.items()},
     }
 

diff --git a/mlonmcu/flow/tvm/backend/python_utils.py b/mlonmcu/flow/tvm/backend/python_utils.py
@@ -19,9 +19,11 @@
 import os
 
 
-def prepare_python_environment(pythonpath, tvm_build_dir, tvm_configs_dir):
+def prepare_python_environment(pythonpath, tvm_build_dir, tvm_configs_dir, tophub_url=None):
     env = os.environ.copy()
     env["PYTHONPATH"] = str(pythonpath)
     env["TVM_LIBRARY_PATH"] = str(tvm_build_dir)
     env["TVM_CONFIGS_JSON_DIR"] = str(tvm_configs_dir)
+    if tophub_url:
+        env["TOPHUB_LOCATION"] = tophub_url
     return env