[MetaSchedule] Enhance tune_tir to tune IRModule of TIR Collections (#14784)

This PR enhances the user-facing API `tune_tir` that makes it
more convenient to feed in an IRModule consists of multiple TIRs and
customize the search space for them. This is widely used in our MLC-LLM
project where we wanted to customize the search space for some
quantization-related operators.

An example usecase:

```python

from tvm import meta_schedule as ms

ms.tir_integration.tune_tir(
    ...
    space="cuda",  <========== by default, the space is "cuda" rather than "cuda-tensorcore"
    special_space={
        "fused_decode1_fused_matmul2_add1_gelu": sch_fused_decode_gemv,
        "decode": sch_decode,
    },
)
```

Author: junrushao

python/tvm/contrib/torch/as_torch.py

@@ -33,6 +33,7 @@
 
 import torch
 import torch.utils.dlpack
+
 import tvm
 from tvm import meta_schedule as ms
 from tvm.target.target import Target
@@ -66,7 +67,6 @@ def tune(
         task_scheduler: ms.TaskScheduler.TaskSchedulerType = "round-robin",
         space: ms.SpaceGenerator.SpaceGeneratorType = "post-order-apply",
         strategy: ms.SearchStrategy.SearchStrategyType = "replay-trace",
-        task_name: str = "main",
         num_tuning_cores: Union[Literal["physical", "logical"], int] = "physical",
         seed: Optional[int] = None,
     ) -> None:
@@ -99,7 +99,6 @@ def tune(
                 task_scheduler=task_scheduler,
                 space=space,
                 strategy=strategy,
-                task_name=task_name,
                 num_tuning_cores=num_tuning_cores,
                 seed=seed,
             )

python/tvm/meta_schedule/space_generator/space_generator.py

@@ -126,6 +126,8 @@ def create_schedule_fn(
             return PostOrderApply(*args, **kwargs)
         if kind == "union":
             return SpaceGeneratorUnion(*args, **kwargs)
+        if isinstance(kind, str):
+            return PostOrderApply(sch_rules=kind, postprocs=kind, mutator_probs=kind)
         raise ValueError(f"Unknown SpaceGenerator: {kind}")
 
 

python/tvm/meta_schedule/testing/tune_te.py

@@ -135,7 +135,6 @@ def main():
                 adaptive_training=ARGS.adaptive_training,
             ),
             strategy=ms.search_strategy.EvolutionarySearch(),
-            task_name=ARGS.workload,
         )
 
     print("Tuning Time:")

python/tvm/meta_schedule/tir_integration.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """MetaSchedule-TIR integration"""
-from typing import Optional, Union
+from typing import List, Mapping, Optional, Tuple, Union
 
 # isort: off
 from typing_extensions import Literal
@@ -38,37 +38,41 @@
 from .utils import fork_seed
 
 
-def tune_tir(
+def tune_tir(  # pylint: disable=too-many-locals
     mod: Union[ir.IRModule, tir.PrimFunc],
     target: Union[str, Target],
     work_dir: str,
     max_trials_global: int,
     *,
+    max_trials_per_task: Optional[int] = None,
     num_trials_per_iter: int = 64,
     builder: Builder.BuilderType = "local",
     runner: Runner.RunnerType = "local",
     database: Database.DatabaseType = "json",
     cost_model: CostModel.CostModelType = "xgb",
     measure_callbacks: MeasureCallback.CallbackListType = "default",
-    task_scheduler: TaskScheduler.TaskSchedulerType = "round-robin",
+    task_scheduler: TaskScheduler.TaskSchedulerType = "gradient",
     space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
     strategy: SearchStrategy.SearchStrategyType = "evolutionary",
-    task_name: str = "main",
     num_tuning_cores: Union[Literal["physical", "logical"], int] = "physical",
     seed: Optional[int] = None,
+    module_equality: str = "structural",
+    special_space: Optional[Mapping[str, SpaceGenerator.SpaceGeneratorType]] = None,
 ) -> Database:
-    """Tune a TIR function.
+    """Tune a TIR function or an IRModule of TIR functions.
 
     Parameters
     ----------
     mod : Union[ir.IRModule, tir.PrimFunc]
-        The TIR function to tune.
+        The TIR IRModule to tune.
     target : Union[str, Target]
         The target to tune for.
     work_dir : str
         The working directory.
     max_trials_global : int
         The maximum number of trials to run globally.
+    max_trials_per_task : Optional[int]
+        The maximum number of trials to run per task.
     num_trials_per_iter : int
         The number of trials to run per iteration
     builder : Builder.BuilderType
@@ -87,44 +91,69 @@ def tune_tir(
         The space generator.
     strategy : SearchStrategy.SearchStrategyType
         The search strategy.
-    task_name : str
-        The name of the task.
     num_tuning_cores : Union[Literal["physical", "logical"], int]
         The number of CPU cores to use during tuning.
     seed : Optional[int]
         The seed for the random number generator.
+    module_equality : Optional[str]
+        A string to specify the module equality testing and hashing method.
+    special_space : Optional[Mapping[str, SpaceGenerator.SpaceGeneratorType]]
+        A mapping from task name to a special space generator for that task.
 
     Returns
     -------
     database : Database
         The database with all tuning records
     """
-    (logger,) = get_loggers_from_work_dir(work_dir, [task_name])
-    (seed,) = fork_seed(seed, n=1)
-    return tune_tasks(
-        tasks=[
+    if isinstance(mod, tir.PrimFunc):
+        mod = _normalize_mod(mod)
+
+    named_tasks: List[Tuple[str, tir.PrimFunc]] = []
+    for gv, func in mod.functions.items():  # pylint: disable=invalid-name
+        if isinstance(func, tir.PrimFunc):
+            named_tasks.append((gv.name_hint, func))
+    named_tasks.sort(key=lambda x: x[0])
+
+    task_names = [x for x, _ in named_tasks]
+    tasks: List[TuneContext] = []
+    for task_name, task_func, logger, rand_state in zip(
+        task_names,
+        [x for _, x in named_tasks],
+        get_loggers_from_work_dir(work_dir, task_names),
+        fork_seed(seed, n=len(named_tasks)),
+    ):
+        if special_space and task_name in special_space:
+            task_space = special_space[task_name]
+        else:
+            task_space = space
+        if task_space is None:
+            continue
+        tasks.append(
             TuneContext(
-                mod=mod,
+                mod=task_func,
                 target=target,
-                space_generator=space,
+                space_generator=task_space,
                 search_strategy=strategy,
                 task_name=task_name,
-                logger=logger,
-                rand_state=seed,
+                rand_state=rand_state,
                 num_threads=num_tuning_cores,
+                logger=logger,
             ).clone()
-        ],
-        task_weights=[1.0],
+        )
+    return tune_tasks(
+        tasks=tasks,
+        task_weights=[1.0] * len(tasks),
         work_dir=work_dir,
         max_trials_global=max_trials_global,
-        max_trials_per_task=max_trials_global,
+        max_trials_per_task=max_trials_per_task,
         num_trials_per_iter=num_trials_per_iter,
         builder=builder,
         runner=runner,
         database=database,
         cost_model=cost_model,
         measure_callbacks=measure_callbacks,
         task_scheduler=task_scheduler,
+        module_equality=module_equality,
     )