def _lower(mod, target, params): """Helper to lower VTA properly.""" # pylint: disable=import-outside-toplevel from tvm import relay from tvm.relay.backend import graph_executor_codegen if hasattr(target, "device_name") and target.device_name == "vta": import vta 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"]) return # default case # Try graph codegen first to extract autotvm tasks. # If failed to compile, then fallback to use VM compiler. # TODO: Currently VM compiler is likely to stack overflow for large models. try: opt_mod, _ = relay.optimize(mod, target, params) grc = graph_executor_codegen.GraphExecutorCodegen(None, target) grc.codegen(opt_mod["main"]) except tvm.TVMError as e: print("Get errors with GraphExecutorCodegen for task extraction. " "Fallback to VMCompiler. Error details:\n%s" % str(e)) compiler = relay.vm.VMCompiler() if params: compiler.set_params(params) compiler.lower(mod, target=target)
def call_all_topi_funcs(mod, params, target): """Call all TOPI compute to extract auto_scheduler tasks in a Relay program""" # pylint: disable=import-outside-toplevel from tvm import relay from tvm.relay.backend import graph_executor_codegen # Turn off AutoTVM config not found warnings old_autotvm_silent = autotvm.GLOBAL_SCOPE.silent autotvm.GLOBAL_SCOPE.silent = True with transform.PassContext( opt_level=3, config={ "relay.backend.use_auto_scheduler": True, "relay.backend.disable_compile_engine_cache": True, }, disabled_pass={"AutoSchedulerLayoutRewrite"}, ): try: opt_mod, _ = relay.optimize(mod, target, params) grc = graph_executor_codegen.GraphExecutorCodegen(None, target) grc.codegen(opt_mod["main"]) except tvm.TVMError: print("Get errors with GraphExecutorCodegen for task extraction. " "Fallback to VMCompiler.") compiler = relay.vm.VMCompiler() if params: compiler.set_params(params) mod = tvm.IRModule.from_expr(mod) if isinstance( mod, relay.Function) else mod compiler.lower(mod, target) autotvm.GLOBAL_SCOPE.silent = old_autotvm_silent
def build_graph(mod, target): target = relay.build_module.build_target_by_device_type_map(target) target, target_host = tvm.target.Target.check_and_update_host_consist( target) mod, _ = relay.optimize(mod, target, None) grc = graph_executor_codegen.GraphExecutorCodegen(None, target) _, lowered_funcs, _ = grc.codegen(mod, mod["main"]) _ = relay.backend._backend.build(lowered_funcs, target, target_host)
def _lower(mod, target, params): """Helper to lower VTA properly.""" # pylint: disable=import-outside-toplevel from tvm import relay from tvm.relay.backend import graph_executor_codegen if hasattr(target, "device_name") and target.device_name == "vta": import vta 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"]) return compiler = relay.vm.VMCompiler() if params: compiler.set_params(params) compiler.lower(mod, target=target)
def _lower(mod, target, params, opt_level=3): """Helper to lower VTA properly.""" # pylint: disable=import-outside-toplevel from tvm import relay from tvm.relay.backend import graph_executor_codegen if hasattr(target, "device_name") and target.device_name == "vta": import vta 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 # 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) compiler.lower(mod, target=target)
def build_graph(mod, target): target, target_host = tvm.target.Target.canon_target_and_host(target) mod, _ = relay.optimize(mod, target) grc = graph_executor_codegen.GraphExecutorCodegen(None, target) _, lowered_funcs, _ = grc.codegen(mod, mod["main"]) _ = relay.backend._backend.build(lowered_funcs, target)