diff --git a/apps/relax_examples/e2e_auto_tir.py b/apps/relax_examples/e2e_auto_tir.py index c5dc6151f6..92cda16f79 100644 --- a/apps/relax_examples/e2e_auto_tir.py +++ b/apps/relax_examples/e2e_auto_tir.py @@ -191,18 +191,21 @@ def main(): relax_mod = apply_opt_before_tuning(relay_mod, params, target=ARGS.target) assert isinstance(relax_mod, tvm.IRModule) - executable = ms.tune_relax( + db = ms.relax_integration.tune_relax( mod=relax_mod, target=ARGS.target, - config=ms.TuneConfig( - strategy="evolutionary", - num_trials_per_iter=64, - max_trials_per_task=ARGS.num_trials, - max_trials_global=ARGS.num_trials, - ), + params=params, + num_trials_per_iter=64, + max_trials_per_task=ARGS.num_trials, + max_trials_global=ARGS.num_trials, runner=get_runner(), work_dir=ARGS.work_dir, - num_threads=os.cpu_count(), + ) + executable = ms.relax_integration.compile_relax( + db, + mod=relax_mod, + target=ARGS.target, + params=params, ) for input_name, input_shape in input_info.items(): diff --git a/python/tvm/relax/testing/relay_translator.py b/python/tvm/relax/testing/relay_translator.py index 77968de0e2..7c49631709 100644 --- a/python/tvm/relax/testing/relay_translator.py +++ b/python/tvm/relax/testing/relay_translator.py @@ -28,6 +28,7 @@ from tvm.relay.backend.te_compiler import select_implementation from tvm.runtime import NDArray from tvm.target import Target +from tvm.meta_schedule.relay_integration import _autotvm_silencer def from_relay( @@ -196,7 +197,7 @@ def visit_func(node): # Since optimization passes and OpStrategy are highly context-dependent, # we match the exact same context with `extract_task_from_relay()` env - with tvm.transform.PassContext( + with _autotvm_silencer(), tvm.transform.PassContext( opt_level=opt_level, config=pass_config, disabled_pass=disabled_pass, diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py index 548cd84479..456a7a211a 100644 --- a/python/tvm/relax/transform/transform.py +++ b/python/tvm/relax/transform/transform.py @@ -24,9 +24,6 @@ import numpy as np import tvm.ir from tvm.runtime import NDArray -from tvm.target import Target -from tvm.meta_schedule.database import PyDatabase - from . import _ffi_api