From 37a1897df2c4bc067d28b2d8fd49a9bb2d70146a Mon Sep 17 00:00:00 2001 From: Tristan Konolige Date: Tue, 25 Jan 2022 15:37:47 -0800 Subject: [PATCH 1/2] [AUTOTVM] Use opt level 3 when extracting tasks Autotvm was implicitly ignoring opt_level when extracting tasks because pass opt_level is a thread local variable and extraction happens in a new thread. Not having opt_level 3 causes alter op layout to not fire, which in turn prevents tuning from finding all possible kernels. --- python/tvm/autotvm/task/relay_integration.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py index d2861cfe8ad6..1afd8788f80f 100644 --- a/python/tvm/autotvm/task/relay_integration.py +++ b/python/tvm/autotvm/task/relay_integration.py @@ -34,7 +34,7 @@ # TODO(moreau89) find a more elegant way to lower for VTAs -def _lower(mod, target, params): +def _lower(mod, target, params, opt_level=3): """Helper to lower VTA properly.""" # pylint: disable=import-outside-toplevel from tvm import relay @@ -43,16 +43,17 @@ def _lower(mod, target, params): if hasattr(target, "device_name") and target.device_name == "vta": import vta - with vta.build_config(opt_level=3, disabled_pass={"AlterOpLayout"}): + with vta.build_config(opt_level=opt_level, disabled_pass={"AlterOpLayout"}): mod, _ = relay.optimize(mod, target, params) grc = graph_executor_codegen.GraphExecutorCodegen(None, target) grc.codegen(mod, mod["main"]) return - compiler = relay.vm.VMCompiler() - if params: - compiler.set_params(params) - compiler.lower(mod, target=target) + with tvm.transform.PassContext(opt_level=opt_level): + compiler = relay.vm.VMCompiler() + if params: + compiler.set_params(params) + compiler.lower(mod, target=target) def extract_from_program(mod, params, target, target_host=None, ops=None): From 6ff5106cc2bb84a11da139397810f02abc3c7734 Mon Sep 17 00:00:00 2001 From: Tristan Konolige Date: Mon, 31 Jan 2022 13:58:53 -0800 Subject: [PATCH 2/2] disable alter op layout --- python/tvm/autotvm/task/relay_integration.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py index 1afd8788f80f..04ca333a5ea8 100644 --- a/python/tvm/autotvm/task/relay_integration.py +++ b/python/tvm/autotvm/task/relay_integration.py @@ -49,7 +49,9 @@ def _lower(mod, target, params, opt_level=3): grc.codegen(mod, mod["main"]) return - with tvm.transform.PassContext(opt_level=opt_level): + # Alter op layout code has been written expecting that tuning is applied + # without it, so we disable AlterOpLayout to maintain that behavior. + with tvm.transform.PassContext(opt_level=opt_level, disabled_pass={"AlterOpLayout"}): compiler = relay.vm.VMCompiler() if params: compiler.set_params(params)