[MetaSchedule] Developer Ergonomics Enhancement II

Follow-up of apache#11622, per discussion with @Kathryn-cat

- [x] Allow using a string `"default"` in `TuneContext` to quickly specify a set of target-specific
rules
- [x] Enhance detection of `ScheduleFn` in `TuneContext` to make it easier for users to quickly try
out template-driven scheduling on TIR.

Next PR:
- Add `TuneContext.tune` to allow directly tuning without task scheduler.

Co-Authored-By: Kathryn (Jinqi) Chen <[email protected]>

Author: junrushao

python/tvm/meta_schedule/space_generator/__init__.py

@@ -19,7 +19,7 @@
 Meta Schedule design space generators that generates design
 space for generation of measure candidates.
 """
-from .space_generator import SpaceGenerator, PySpaceGenerator
-from .space_generator_union import SpaceGeneratorUnion
-from .schedule_fn import ScheduleFn
 from .post_order_apply import PostOrderApply
+from .schedule_fn import SCH_FN_TYPE, ScheduleFn
+from .space_generator import PySpaceGenerator, SpaceGenerator
+from .space_generator_union import SpaceGeneratorUnion

python/tvm/meta_schedule/space_generator/schedule_fn.py

@@ -30,18 +30,17 @@
 if TYPE_CHECKING:
     from ..tune_context import TuneContext
 
+SCH_FN_TYPE = Union[  # pylint: disable=invalid-name
+    Callable[[Schedule], None],  # No output
+    Callable[[Schedule], Schedule],  # Single output
+    Callable[[Schedule], List[Schedule]],  # Multiple outputs
+]
+
 
 @derived_object
 class ScheduleFn(PySpaceGenerator):
     """A design space generator with design spaces specified by a schedule function."""
 
-    # Multiple cases of schedule functions supported
-    SCH_FN_TYPE = Union[
-        Callable[[Schedule], None],  # No output
-        Callable[[Schedule], Schedule],  # Single output
-        Callable[[Schedule], List[Schedule]],  # Multiple outputs
-    ]
-
     def __init__(self, sch_fn: SCH_FN_TYPE):
         """Constructor.
 

python/tvm/meta_schedule/tune.py

@@ -88,7 +88,7 @@ class TuneConfig(NamedTuple):
     search_strategy_config: Optional[Dict[str, Any]] = None
     logger_config: Optional[Dict[str, Any]] = None
 
-    def create_strategy(self, **kwargs):
+    def create_strategy(self):
         """Create search strategy from configuration"""
         cls_tbl = {
             "evolutionary": EvolutionarySearch,
@@ -111,7 +111,6 @@ def create_strategy(self, **kwargs):
         return cls_tbl[self.strategy](
             num_trials_per_iter=self.num_trials_per_iter,
             max_trials_per_task=max_trials_per_task,
-            **kwargs,
             **config,
         )
 

python/tvm/meta_schedule/tune_context.py

@@ -17,7 +17,7 @@
 """Meta Schedule tuning context."""
 
 import logging
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 from tvm import IRModule
 from tvm._ffi import register_object
@@ -36,7 +36,8 @@
     from .runner import RunnerResult
     from .schedule_rule import ScheduleRule
     from .search_strategy import MeasureCandidate, SearchStrategy
-    from .space_generator import SpaceGenerator
+    from .space_generator import SCH_FN_TYPE, SpaceGenerator
+    from .tune import TuneConfig
 
 
 @register_object("meta_schedule.TuneContext")
@@ -54,16 +55,24 @@ class TuneContext(Object):
         The workload to be optimized.
     target : Optional[Target] = None
         The target to be optimized for.
-    space_generator : Optional[SpaceGenerator] = None
+    space_generator : Union[None, SCH_FN_TYPE, SpaceGenerator] = None
         The design space generator.
-    search_strategy : Optional[SearchStrategy] = None
+    search_strategy : Union[None, TuneConfig, SearchStrategy] = None
         The search strategy.
-    sch_rules: Optional[List[ScheduleRule]] = None,
+        if None, the strategy is left blank.
+        If TuneConfig, the strategy is initialized with the TuneConfig.create_strategy().
+    sch_rules: Union[None, str, List[ScheduleRule]] = None,
         The schedule rules.
-    postprocs: Optional[List[Postproc"]] = None,
+        If None, use an empty list of rules.
+        if "default", use target-default rules.
+    postprocs: Union[None, str, List[Postproc"]] = None,
         The postprocessors.
-    mutator_probs: Optional[Dict[Mutator, float]]
+        If None, use an empty list of rules.
+        if "default", use target-default rules.
+    mutator_probs: Union[None, str, Dict[Mutator, float]]
         Mutators and their probability mass.
+        If None, use an empty list of rules.
+        if "default", use target-default rules.
     task_name : Optional[str] = None
         The name of the tuning task.
     logger : logging.Logger
@@ -99,24 +108,53 @@ def __init__(
         mod: Optional[IRModule] = None,
         *,
         target: Optional[Target] = None,
-        space_generator: Optional["SpaceGenerator"] = None,
-        search_strategy: Optional["SearchStrategy"] = None,
-        sch_rules: Optional[List["ScheduleRule"]] = None,
-        postprocs: Optional[List["Postproc"]] = None,
-        mutator_probs: Optional[Dict["Mutator", float]] = None,
+        space_generator: Union[None, "SCH_FN_TYPE", "SpaceGenerator"] = None,
+        search_strategy: Union[None, "SearchStrategy", "TuneConfig"] = None,
+        sch_rules: Union[None, str, List["ScheduleRule"]] = None,
+        postprocs: Union[None, str, List["Postproc"]] = None,
+        mutator_probs: Union[None, str, Dict["Mutator", float]] = None,
         task_name: str = "main",
         logger: Optional[logging.Logger] = None,
         rand_state: int = -1,
         num_threads: Optional[int] = None,
     ):
+        # pylint: disable=import-outside-toplevel
+        from . import default_config
+        from .space_generator import ScheduleFn
+        from .tune import TuneConfig
+
+        # pylint: enable=import-outside-toplevel
         if isinstance(mod, PrimFunc):
             mod = IRModule.from_expr(mod)
-        if num_threads is None:
-            num_threads = cpu_count()
+        if callable(space_generator):
+            space_generator = ScheduleFn(space_generator)
+        if isinstance(search_strategy, TuneConfig):
+            search_strategy = search_strategy.create_strategy()
+        if isinstance(sch_rules, str):
+            if sch_rules == "default":
+                if target is None:
+                    raise ValueError("target is required when sch_rules is 'default'")
+                sch_rules = default_config.schedule_rules(None, target)
+            else:
+                raise ValueError("sch_rules should be a list of ScheduleRule or 'default'")
+        if isinstance(postprocs, str):
+            if postprocs == "default":
+                if target is None:
+                    raise ValueError("target is required when postprocs is 'default'")
+                postprocs = default_config.postproc(None, target)
+            else:
+                raise ValueError("postprocs should be a list of Postproc or 'default'")
+        if isinstance(mutator_probs, str):
+            if mutator_probs == "default":
+                if target is None:
+                    raise ValueError("target is required when mutator_probs is 'default'")
+                mutator_probs = default_config.mutator_probs(None, target)
         if logger is None:
             self.logger = logging.getLogger(__name__)
         else:
             self.logger = None
+        if num_threads is None:
+            num_threads = cpu_count()
         self.__init_handle_by_constructor__(
             _ffi_api.TuneContext,  # type: ignore # pylint: disable=no-member
             mod,
@@ -131,9 +169,6 @@ def __init__(
             rand_state,
             num_threads,
         )
-
-    def initialize(self):
-        """Initialize the tuning context"""
         _ffi_api.TuneContextInitialize(self)  # type: ignore # pylint: disable=no-member
 
     def generate_design_space(self) -> List[Schedule]:

tests/python/unittest/test_meta_schedule_custom_rule_winograd_cpu.py

@@ -173,7 +173,6 @@ def test_conv2d_winograd_cpu():
             target,
         ),
     )
-    context.initialize()
     post_order_apply = context.space_generator
     (sch,) = post_order_apply.generate_design_space(mod)
     decisions = dict(

tests/python/unittest/test_meta_schedule_custom_rule_winograd_cuda.py

@@ -290,7 +290,6 @@ def test_conv2d_winograd_cuda():
             None, Target("cuda")
         ),
     )
-    context.initialize()
     post_order_apply = context.space_generator
     (sch,) = post_order_apply.generate_design_space(mod)
     decisions = dict(

tests/python/unittest/test_meta_schedule_integration.py

@@ -221,7 +221,7 @@ def traverse(t):
         mod,
         target="llvm",
         params=params,
-        filter_func=filter_func,
+        te_filter_func=filter_func,
     )
     expected_task_names = [
         "fused_" + s

tests/python/unittest/test_meta_schedule_mutator_mutate_compute_location.py

@@ -69,7 +69,6 @@ def _make_mutator(target: Target) -> Mutator:
             MutateComputeLocation(): 1.0,
         },
     )
-    ctx.initialize()
     return list(ctx.mutator_probs.keys())[0]
 
 

tests/python/unittest/test_meta_schedule_mutator_mutate_parallel.py

@@ -87,7 +87,6 @@ def _make_mutator(target: Target, max_jobs_per_core: int) -> Mutator:
             MutateParallel(max_jobs_per_core): 1.0,
         },
     )
-    ctx.initialize()
     return list(ctx.mutator_probs.keys())[0]
 
 

tests/python/unittest/test_meta_schedule_mutator_mutate_thread_binding.py

@@ -70,7 +70,6 @@ def _make_mutator(target: Target) -> Mutator:
             MutateThreadBinding(): 1.0,
         },
     )
-    ctx.initialize()
     return list(ctx.mutator_probs.keys())[0]