From cfb297c92e4ac707b7b810864c20f10aa0817bcb Mon Sep 17 00:00:00 2001 From: Chris Sidebottom Date: Thu, 11 Nov 2021 22:17:16 +0000 Subject: [PATCH] Correct Graph Executor Python API --- python/tvm/autotvm/task/relay_integration.py | 2 +- python/tvm/relay/backend/graph_executor_codegen.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py index a7d998c33bc6..d2861cfe8ad6 100644 --- a/python/tvm/autotvm/task/relay_integration.py +++ b/python/tvm/autotvm/task/relay_integration.py @@ -46,7 +46,7 @@ def _lower(mod, target, params): with vta.build_config(opt_level=3, disabled_pass={"AlterOpLayout"}): mod, _ = relay.optimize(mod, target, params) grc = graph_executor_codegen.GraphExecutorCodegen(None, target) - grc.codegen(mod["main"]) + grc.codegen(mod, mod["main"]) return compiler = relay.vm.VMCompiler() diff --git a/python/tvm/relay/backend/graph_executor_codegen.py b/python/tvm/relay/backend/graph_executor_codegen.py index 58717a0ab482..02e6c4f61af2 100644 --- a/python/tvm/relay/backend/graph_executor_codegen.py +++ b/python/tvm/relay/backend/graph_executor_codegen.py @@ -64,11 +64,13 @@ def _setup(self, mod, target): tgts[_expr.IntImm("int32", 0)] = Target(target) self._init(mod, tgts) - def codegen(self, func): + def codegen(self, ir_module, func): """Compile a single function into a graph. Parameters ---------- + ir_module: tvm.ir.Module + The module to compile func: tvm.relay.Expr The function to compile. @@ -82,7 +84,7 @@ def codegen(self, func): Additional constant parameters. """ default_mod_name = mangle_module_name("default") - self._codegen(func, default_mod_name) + self._codegen(ir_module, func, default_mod_name) graph_json = self._get_graph_json() lowered_func = self._get_irmodule() param_names = self._list_params_name()