From 40d5193a962230fb68cd11b2f56047d20c4f62d4 Mon Sep 17 00:00:00 2001
From: Muyang Li <lmxyy@users.noreply.github.com>
Date: Mon, 21 Jun 2021 15:42:00 +0800
Subject: [PATCH] [Auto Scheduler] Make the opt_level of task extraction
 adjustable (#8288)

* fix bugs in the auto scheduler record:

* reformat the code

* use the os.path.abspath

* change error to warning

* reformat the warning code

* fix some typos

* fix the port number typo

* fix a typo

* make query_rpc_tracker show the correct port and the customized address

* disable the pycharm reformat

* reformat the code

* make the opt_level of extract_tasks adjustable

* Update rpc_server.py

* fix a typo

* Update tracker.py

* support checking the port and customized address

* reformat the code

* fix a typo
---
 python/tvm/auto_scheduler/relay_integration.py | 18 ++++++++++++++----
 python/tvm/rpc/tracker.py                      |  2 +-
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/python/tvm/auto_scheduler/relay_integration.py b/python/tvm/auto_scheduler/relay_integration.py
index 099502d17d78..0d18bc08e5ed 100644
--- a/python/tvm/auto_scheduler/relay_integration.py
+++ b/python/tvm/auto_scheduler/relay_integration.py
@@ -46,7 +46,7 @@
 logger = logging.getLogger("auto_scheduler")
 
 
-def call_all_topi_funcs(mod, params, target):
+def call_all_topi_funcs(mod, params, target, opt_level=3):
     """Call all TOPI compute to extract auto_scheduler tasks in a Relay program"""
     # pylint: disable=import-outside-toplevel
     from tvm import relay
@@ -57,7 +57,7 @@ def call_all_topi_funcs(mod, params, target):
     autotvm.GLOBAL_SCOPE.silent = True
 
     with transform.PassContext(
-        opt_level=3,
+        opt_level=opt_level,
         config={
             "relay.backend.use_auto_scheduler": True,
             "relay.backend.disable_compile_engine_cache": True,
@@ -91,7 +91,13 @@ def call_all_topi_funcs(mod, params, target):
 
 
 def extract_tasks(
-    mod, params, target, target_host=None, hardware_params=None, include_simple_tasks=False
+    mod,
+    params,
+    target,
+    target_host=None,
+    hardware_params=None,
+    include_simple_tasks=False,
+    opt_level=3,
 ):
     """Extract tuning tasks from a relay program.
 
@@ -109,6 +115,8 @@ def extract_tasks(
         Hardware parameters used for the search tasks
     include_simple_tasks: bool
         Whether to extract simple tasks that do not include complicated ops.
+    opt_level : Optional[int]
+        The optimization level of the task extractions.
 
     Returns
     -------
@@ -132,7 +140,9 @@ def extract_tasks(
     with env:
         # Wrap build call in a new thread to avoid the conflict
         # between python's multiprocessing and tvm's thread pool
-        build_thread = threading.Thread(target=call_all_topi_funcs, args=(mod, params, target))
+        build_thread = threading.Thread(
+            target=call_all_topi_funcs, args=(mod, params, target, opt_level)
+        )
         build_thread.start()
         build_thread.join()
     dispatch_ctx.verbose = old_verbose
diff --git a/python/tvm/rpc/tracker.py b/python/tvm/rpc/tracker.py
index c8ab15bc4a03..74c1f7ac07aa 100644
--- a/python/tvm/rpc/tracker.py
+++ b/python/tvm/rpc/tracker.py
@@ -16,7 +16,7 @@
 # under the License.
 """RPC Tracker, tracks and distributes the TVM RPC resources.
 
-This folder implemements the tracker server logic.
+This folder implements the tracker server logic.
 
 Note
 ----