[MetaSchedule][BugFix] Fix skipped tests

python/tvm/meta_schedule/__init__.py

@@ -15,23 +15,27 @@
 # specific language governing permissions and limitations
 # under the License.
 """Package `tvm.meta_schedule`. The meta schedule infrastructure."""
-from . import arg_info
-from . import database
-from . import builder
-from . import runner
-from . import mutator
-from . import postproc
-from . import schedule_rule
-from . import space_generator
-from . import search_strategy
-from . import integration
-from . import feature_extractor
-from . import cost_model
-from .search_strategy import (
+from . import (
+    arg_info,
+    builder,
+    cost_model,
+    database,
+    feature_extractor,
+    integration,
+    mutator,
+    postproc,
+    runner,
+    schedule_rule,
+    search_strategy,
+    space_generator,
+)
+from .search_strategy import MeasureCandidate
+from .tune import (
     EvolutionarySearchConfig,
-    MeasureCandidate,
     ReplayFuncConfig,
     ReplayTraceConfig,
+    tune_relay,
+    tune_te,
+    tune_tir,
 )
-from .tune import tune_te, tune_tir, tune_relay
 from .tune_context import TuneContext

python/tvm/meta_schedule/search_strategy/__init__.py

@@ -20,7 +20,7 @@
 to generate measure candidates.
 """
 
-from .search_strategy import SearchStrategy, PySearchStrategy, MeasureCandidate
-from .replay_trace import ReplayTrace, ReplayTraceConfig
-from .replay_func import ReplayFunc, ReplayFuncConfig
-from .evolutionary_search import EvolutionarySearch, EvolutionarySearchConfig
+from .evolutionary_search import EvolutionarySearch
+from .replay_func import ReplayFunc
+from .replay_trace import ReplayTrace
+from .search_strategy import MeasureCandidate, PySearchStrategy, SearchStrategy

python/tvm/meta_schedule/search_strategy/evolutionary_search.py

@@ -15,9 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """Evolutionary Search Strategy"""
-
-from typing import NamedTuple
-
 from tvm._ffi import register_object
 
 from .. import _ffi_api
@@ -88,31 +85,3 @@ def __init__(
             genetic_max_fail_count,
             eps_greedy,
         )
-
-
-class EvolutionarySearchConfig(NamedTuple):
-    """Configuration for EvolutionarySearch"""
-
-    num_trials_per_iter: int
-    max_trials_per_task: int
-    max_trials_global: int
-    population_size: int = 2048
-    init_measured_ratio: float = 0.2
-    init_min_unmeasured: int = 50
-    genetic_num_iters: int = 4
-    genetic_mutate_prob: float = 0.85
-    genetic_max_fail_count: int = 10
-    eps_greedy: float = 0.05
-
-    def create_strategy(self) -> EvolutionarySearch:
-        return EvolutionarySearch(
-            num_trials_per_iter=self.num_trials_per_iter,
-            max_trials_per_task=self.max_trials_per_task,
-            population_size=self.population_size,
-            init_measured_ratio=self.init_measured_ratio,
-            init_min_unmeasured=self.init_min_unmeasured,
-            genetic_num_iters=self.genetic_num_iters,
-            genetic_mutate_prob=self.genetic_mutate_prob,
-            genetic_max_fail_count=self.genetic_max_fail_count,
-            eps_greedy=self.eps_greedy,
-        )

python/tvm/meta_schedule/search_strategy/replay_func.py

@@ -15,8 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """Replay Trace Search Strategy"""
-from typing import NamedTuple
-
 from tvm._ffi import register_object
 
 from .. import _ffi_api
@@ -34,7 +32,7 @@ class ReplayFunc(SearchStrategy):
     num_trials_per_iter : int
         Number of trials per iteration.
     max_trials_per_task : int
-        Total number of trials.
+        Total number of trials for one task
     """
 
     num_trials_per_iter: int
@@ -51,14 +49,3 @@ def __init__(
             num_trials_per_iter,
             max_trials_per_task,
         )
-
-
-class ReplayFuncConfig(NamedTuple):
-    """Configuration for ReplayFunc"""
-
-    num_trials_per_iter: int
-    max_trials_per_task: int
-    max_trials_global: int
-
-    def create_strategy(self) -> ReplayFunc:
-        return ReplayFunc(self.num_trials_per_iter, self.max_trials_per_task)

python/tvm/meta_schedule/search_strategy/replay_trace.py

@@ -15,8 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """Replay Trace Search Strategy"""
-from typing import NamedTuple
-
 from tvm._ffi import register_object
 
 from .. import _ffi_api
@@ -34,7 +32,7 @@ class ReplayTrace(SearchStrategy):
     num_trials_per_iter : int
         Number of trials per iteration.
     max_trials_per_task : int
-        Total number of trials.
+        Total number of trials for one task
     """
 
     num_trials_per_iter: int
@@ -47,14 +45,3 @@ def __init__(self, num_trials_per_iter: int, max_trials_per_task: int):
             num_trials_per_iter,
             max_trials_per_task,
         )
-
-
-class ReplayTraceConfig(NamedTuple):
-    """Configuration for ReplayTrace"""
-
-    num_trials_per_iter: int
-    max_trials_per_task: int
-    max_trials_global: int
-
-    def create_strategy(self) -> ReplayTrace:
-        return ReplayTrace(self.num_trials_per_iter, self.max_trials_per_task)

python/tvm/meta_schedule/tune.py

@@ -18,7 +18,7 @@
 # pylint: disable=import-outside-toplevel
 import logging
 import os.path
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, NamedTuple, Optional, Tuple, Union
 
 from tvm._ffi.registry import register_func
 from tvm.ir import IRModule, structural_hash
@@ -40,23 +40,14 @@
 from .postproc import Postproc
 from .runner import LocalRunner, Runner
 from .schedule_rule import ScheduleRule
-from .search_strategy import (
-    EvolutionarySearchConfig,
-    ReplayFuncConfig,
-    ReplayTraceConfig,
-)
+from .search_strategy import EvolutionarySearch, ReplayFunc, ReplayTrace
 from .space_generator import PostOrderApply, SpaceGenerator
 from .task_scheduler import GradientBased, TaskScheduler
 from .tune_context import TuneContext
 from .utils import autotvm_silencer
 
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
-SearchStrategyConfig = Union[
-    ReplayFuncConfig,
-    ReplayTraceConfig,
-    EvolutionarySearchConfig,
-]
 FnSpaceGenerator = Callable[[], SpaceGenerator]
 FnScheduleRule = Callable[[], List[ScheduleRule]]
 FnPostproc = Callable[[], List[Postproc]]
@@ -75,6 +66,107 @@
 ]
 
 
+class ReplayFuncConfig(NamedTuple):
+    """Configuration for ReplayFunc
+
+    Parameters
+    ----------
+    num_trials_per_iter : int
+        Number of trials per iteration.
+    max_trials_per_task : int
+        Total number of trials for one task
+    max_trials_global : int
+        Total number of trials for all tasks in the task scheduler
+    """
+
+    num_trials_per_iter: int
+    max_trials_per_task: int
+    max_trials_global: int
+
+    def create_strategy(self) -> ReplayFunc:
+        return ReplayFunc(self.num_trials_per_iter, self.max_trials_per_task)
+
+
+class ReplayTraceConfig(NamedTuple):
+    """Configuration for ReplayTrace
+
+    Parameters
+    ----------
+    num_trials_per_iter : int
+        Number of trials per iteration.
+    max_trials_per_task : int
+        Total number of trials for one task
+    max_trials_global : int
+        Total number of trials for all tasks in the task scheduler
+    """
+
+    num_trials_per_iter: int
+    max_trials_per_task: int
+    max_trials_global: int
+
+    def create_strategy(self) -> ReplayTrace:
+        return ReplayTrace(self.num_trials_per_iter, self.max_trials_per_task)
+
+
+class EvolutionarySearchConfig(NamedTuple):
+    """Configuration for EvolutionarySearch
+
+    Parameters
+    ----------
+    num_trials_per_iter : int
+        Number of trials per iteration.
+    max_trials_per_task : int
+        Total number of trials.
+    max_trials_global : int
+        Total number of trials for all tasks in the task scheduler
+    population_size : int
+        The initial population of traces from measured samples and randomly generated samples.
+    init_measured_ratio : int
+        The ratio of measured samples in the initial population.
+    init_min_unmeasured : int
+        The minimal size of unmeasured population in the initial sampling.
+    genetic_num_iters : int
+        The number of iterations for genetic algorithm.
+    genetic_mutate_prob : float
+        The probability of mutation.
+    genetic_max_fail_count : int
+        The maximum number to retry mutation.
+    eps_greedy : float
+        The ratio of greedy selected samples in the final picks.
+    """
+
+    num_trials_per_iter: int
+    max_trials_per_task: int
+    max_trials_global: int
+    population_size: int = 2048
+    init_measured_ratio: float = 0.2
+    init_min_unmeasured: int = 50
+    genetic_num_iters: int = 4
+    genetic_mutate_prob: float = 0.85
+    genetic_max_fail_count: int = 10
+    eps_greedy: float = 0.05
+
+    def create_strategy(self) -> EvolutionarySearch:
+        return EvolutionarySearch(
+            num_trials_per_iter=self.num_trials_per_iter,
+            max_trials_per_task=self.max_trials_per_task,
+            population_size=self.population_size,
+            init_measured_ratio=self.init_measured_ratio,
+            init_min_unmeasured=self.init_min_unmeasured,
+            genetic_num_iters=self.genetic_num_iters,
+            genetic_mutate_prob=self.genetic_mutate_prob,
+            genetic_max_fail_count=self.genetic_max_fail_count,
+            eps_greedy=self.eps_greedy,
+        )
+
+
+SearchStrategyConfig = Union[
+    ReplayFuncConfig,
+    ReplayTraceConfig,
+    EvolutionarySearchConfig,
+]
+
+
 class DefaultLLVM:
     """Default tuning configuration for LLVM."""
 

tests/python/unittest/test_meta_schedule_tune_relay.py

@@ -149,6 +149,7 @@ def test_meta_schedule_tune_relay(
             config=ReplayTraceConfig(
                 num_trials_per_iter=32,
                 max_trials_per_task=32,
+                max_trials_global=20000,
             ),
             work_dir=work_dir,
             database=JSONDatabase(
@@ -493,7 +494,8 @@ def manual_tir_common(do_tune=False):
         if do_tune:
             config = ReplayTraceConfig(
                 num_trials_per_iter=64,
-                num_trials_total=64,
+                max_trials_per_task=64,
+                max_trials_global=20000,
             )
             # postprocs=lambda: [] is important to prevent default post processors from
             # tampering with the manual schedule.

tests/python/unittest/test_meta_schedule_tune_te.py

@@ -37,6 +37,7 @@ def test_tune_matmul():
             config=ReplayTraceConfig(
                 num_trials_per_iter=32,
                 max_trials_per_task=32,
+                max_trials_global=32,
             ),
             work_dir=work_dir,
         )

tests/python/unittest/test_meta_schedule_tune_tir.py

@@ -20,10 +20,9 @@
 
 import pytest
 import tvm
-from tvm.meta_schedule import ReplayTraceConfig, postproc, schedule_rule, tune_tir
+from tvm.meta_schedule import ReplayTraceConfig, schedule_rule, tune_tir
 from tvm.meta_schedule.space_generator import PostOrderApply
 from tvm.meta_schedule.testing import te_workload
-from tvm.meta_schedule.tune_context import TuneContext
 from tvm.script import tir as T
 from tvm.target.target import Target
 from tvm.te.operation import create_prim_func
@@ -61,6 +60,7 @@ def test_tune_matmul_cpu():
             config=ReplayTraceConfig(
                 num_trials_per_iter=32,
                 max_trials_per_task=32,
+                max_trials_global=32,
             ),
             work_dir=work_dir,
         )
@@ -80,6 +80,7 @@ def test_tune_matmul_cuda():
             config=ReplayTraceConfig(
                 num_trials_per_iter=32,
                 max_trials_per_task=32,
+                max_trials_global=32,
             ),
             work_dir=work_dir,
         )
@@ -98,6 +99,7 @@ def test_tune_matmul_cuda_tensor_core():
     config = ReplayTraceConfig(
         num_trials_per_iter=32,
         max_trials_per_task=320,
+        max_trials_global=320,
     )
 
     class DefaultTensorCore: