diff --git a/python/tvm/meta_schedule/__init__.py b/python/tvm/meta_schedule/__init__.py index 2a69d3c69610..3612bb81a6bc 100644 --- a/python/tvm/meta_schedule/__init__.py +++ b/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 diff --git a/python/tvm/meta_schedule/search_strategy/__init__.py b/python/tvm/meta_schedule/search_strategy/__init__.py index 174672235b42..2046067d6c00 100644 --- a/python/tvm/meta_schedule/search_strategy/__init__.py +++ b/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 diff --git a/python/tvm/meta_schedule/search_strategy/evolutionary_search.py b/python/tvm/meta_schedule/search_strategy/evolutionary_search.py index c302d570c2aa..20d0b33378e3 100644 --- a/python/tvm/meta_schedule/search_strategy/evolutionary_search.py +++ b/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, - ) diff --git a/python/tvm/meta_schedule/search_strategy/replay_func.py b/python/tvm/meta_schedule/search_strategy/replay_func.py index ef1fd07527bd..d89e2b133cde 100644 --- a/python/tvm/meta_schedule/search_strategy/replay_func.py +++ b/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) diff --git a/python/tvm/meta_schedule/search_strategy/replay_trace.py b/python/tvm/meta_schedule/search_strategy/replay_trace.py index ec4fa88b5f3e..70461d65f776 100644 --- a/python/tvm/meta_schedule/search_strategy/replay_trace.py +++ b/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) diff --git a/python/tvm/meta_schedule/tune.py b/python/tvm/meta_schedule/tune.py index c65e92aec3c7..86157e0fb32e 100644 --- a/python/tvm/meta_schedule/tune.py +++ b/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,11 +40,7 @@ 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 @@ -52,11 +48,6 @@ 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.""" diff --git a/tests/python/unittest/test_meta_schedule_tune_relay.py b/tests/python/unittest/test_meta_schedule_tune_relay.py index 389f6c3719aa..63fbac6748c7 100644 --- a/tests/python/unittest/test_meta_schedule_tune_relay.py +++ b/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. diff --git a/tests/python/unittest/test_meta_schedule_tune_te.py b/tests/python/unittest/test_meta_schedule_tune_te.py index e0a7a8190419..f58ebf34787e 100644 --- a/tests/python/unittest/test_meta_schedule_tune_te.py +++ b/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, ) diff --git a/tests/python/unittest/test_meta_schedule_tune_tir.py b/tests/python/unittest/test_meta_schedule_tune_tir.py index 6a80d895dfdc..5ac6a24a423a 100644 --- a/tests/python/unittest/test_meta_schedule_tune_tir.py +++ b/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: