From c6d6a069020d37d6b98743248a1cd20cc7d6c91a Mon Sep 17 00:00:00 2001
From: sung <sunggg@umich.edu>
Date: Mon, 17 Oct 2022 19:07:01 -0700
Subject: [PATCH] Reenable autotvm silencer; fix e2e_auto_tir.py; fix lint.

Co-authored-by: YuchenJin <yuchenj@cs.washington.edu>
---
 apps/relax_examples/e2e_auto_tir.py          | 19 +++++++++++--------
 python/tvm/relax/testing/relay_translator.py |  3 ++-
 python/tvm/relax/transform/transform.py      |  3 ---
 3 files changed, 13 insertions(+), 12 deletions(-)

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