Recover: [Pass] Separate ApplyHistoryBest from tuning passes (apache#226)

Author: sunggg

python/tvm/meta_schedule/relax_integration.py

@@ -22,11 +22,12 @@
 
 # isort: on
 
-from tvm._ffi import get_global_func
+from tvm._ffi import get_global_func, register_func
 from tvm.ir import IRModule
 from tvm.ir.transform import PassContext
 from tvm.runtime import NDArray
 from tvm.target import Target
+from tvm.tir.expr import IntImm
 
 from .builder import Builder
 from .cost_model import CostModel
@@ -223,6 +224,94 @@ def tune_relax(
     )
 
 
+@register_func("tvm.meta_schedule.tune_relax")
+def _tune_relax(
+    mod: Union[IRModule, "relax.Function"],
+    params: Dict[str, NDArray],
+    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 = "gradient",
+    space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+    strategy: SearchStrategy.SearchStrategyType = "evolutionary",
+    seed: Optional[int] = None,
+) -> Database:
+    """Interface with tuning api to tune a Relax program.
+
+    Parameters
+    ----------
+    mod : Union[IRModule, relax.Function]
+        The module or function to tune
+    params : Optional[Dict[str, tvm.runtime.NDArray]]
+        The associated parameters of the program
+    target : Union[Target, str]
+        The compilation target
+    work_dir : str
+        The working directory to store the tuning records
+    max_trials_global : int
+        The maximum number of trials to run
+    max_trials_per_task : Optional[int]
+        The maximum number of trials to run for each task
+    num_trials_per_iter : int
+        The number of trials to run per iteration
+    builder : BuilderType
+        The builder to use
+    runner : RunnerType
+        The runner to use
+    database : DatabaseType
+        The database to use
+    cost_model : CostModelType
+        The cost model to use
+    measure_callbacks : CallbackListType
+        The measure callbacks to use
+    task_scheduler : TaskSchedulerType
+        The task scheduler to use
+    space : SpaceGeneratorType
+        The space generator to use
+    strategy : SearchStrategyType
+        The search strategy to use
+    seed : Optional[int]
+        The random seed
+
+    Returns
+    -------
+    ret_mod : IRModule
+        IRModule
+    """
+    if isinstance(max_trials_global, IntImm):
+        max_trials_global = int(max_trials_global)
+
+    tune_relax(
+        mod,
+        params,
+        target,
+        work_dir,
+        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,
+        space=space,
+        strategy=strategy,
+        seed=seed,
+    )
+    # Return original IRModule
+    # This pass only makes optimization decision
+    return mod
+
+
 def compile_relax(
     database: Database,
     mod: IRModule,

python/tvm/meta_schedule/tir_integration.py

@@ -22,7 +22,9 @@
 
 # isort: on
 from tvm import ir, tir
+from tvm._ffi import register_func
 from tvm.target import Target
+from tvm.tir.expr import IntImm
 
 from .builder import Builder
 from .cost_model import CostModel
@@ -128,6 +130,93 @@ def tune_tir(
     )
 
 
+@register_func("tvm.meta_schedule.tune_tir")
+def _tune_tir(
+    mod: Union[ir.IRModule, tir.PrimFunc],
+    target: Union[str, Target],
+    work_dir: str,
+    max_trials_global: int,
+    *,
+    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",
+    space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+    strategy: SearchStrategy.SearchStrategyType = "evolutionary",
+    task_name: str = "main",
+    num_threads: Union[Literal["physical", "logical"], int] = "physical",
+    seed: Optional[int] = None,
+) -> Database:
+    """Interface with tuning api to tune a TIR program.
+
+    Parameters
+    ----------
+    mod : Union[ir.IRModule, tir.PrimFunc]
+        The TIR function 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.
+    num_trials_per_iter : int
+        The number of trials to run per iteration
+    builder : Builder.BuilderType
+        The builder.
+    runner : Runner.RunnerType
+        The runner.
+    database : Database.DatabaseType
+        The database.
+    cost_model : CostModel.CostModelType
+        The cost model.
+    measure_callbacks : MeasureCallback.CallbackListType
+        The measure callbacks.
+    task_scheduler : TaskScheduler.TaskSchedulerType
+        The task scheduler.
+    space : SpaceGenerator.SpaceGeneratorType
+        The space generator.
+    strategy : SearchStrategy.SearchStrategyType
+        The search strategy.
+    task_name : str
+        The name of the task.
+    num_threads : Union[Literal["physical", "logical"], int]
+        The number of threads to use.
+    seed : Optional[int]
+        The seed for the random number generator.
+
+    Returns
+    -------
+    ret_mod : IRModule
+        IRModule
+    """
+    if isinstance(max_trials_global, IntImm):
+        max_trials_global = int(max_trials_global)
+    tune_tir(
+        mod,
+        target,
+        work_dir,
+        max_trials_global,
+        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,
+        space=space,
+        strategy=strategy,
+        task_name=task_name,
+        num_threads=num_threads,
+        seed=seed,
+    )
+    # Return original IRModule
+    # This pass only makes optimization decision
+    return mod
+
+
 def compile_tir(
     database: Database,
     mod: Union[ir.IRModule, tir.PrimFunc],

python/tvm/meta_schedule/tune_context.py

@@ -24,7 +24,7 @@
 # isort: on
 
 from tvm import IRModule
-from tvm._ffi import register_object
+from tvm._ffi import register_object, register_func
 from tvm.runtime import Object
 from tvm.target import Target
 from tvm.tir import PrimFunc, Schedule
@@ -41,6 +41,7 @@
     from .space_generator import SpaceGenerator
 
 
+@register_func("tvm.meta_schedule.normalize_mod")
 def _normalize_mod(mod: Union[PrimFunc, IRModule]) -> IRModule:
     """Normalize the input to an IRModule"""
     if isinstance(mod, PrimFunc):

python/tvm/relax/transform/transform.py

@@ -23,6 +23,9 @@
 
 import numpy as np
 import tvm.ir
+from tvm.runtime import NDArray
+from tvm.target import Target
+from tvm.meta_schedule.database import PyDatabase
 
 from . import _ffi_api
 
@@ -283,15 +286,58 @@ def FuseTIR() -> tvm.ir.transform.Pass:
     return _ffi_api.FuseTIR()
 
 
-def MetaScheduleApplyDatabase() -> tvm.ir.transform.Pass:
+def MetaScheduleApplyDatabase(
+    work_dir: Optional[str] = None,
+) -> tvm.ir.transform.Pass:
     """Apply the best schedule from tuning database.
-
+    work_dir : Optional[str]
+       work directory to deduce default database if database is not provided
+       (it will be ignored when an user passes database)
     Returns
     -------
     ret : tvm.transform.Pass
-        The registered pass for tir fusion.
+        The registered pass
+    """
+    return _ffi_api.MetaScheduleApplyDatabase(work_dir)
+
+
+def MetaScheduleTuneTIR(
+    work_dir: str,
+    max_trials_global: int,
+) -> tvm.ir.transform.Pass:
+    """Tune TIR with MetaSchedule.
+    Parameters
+    ----------
+    work_dir: str
+       work directory
+    max_trials_gloabl: int
+       maximum number of total trials allowed for tuning
+    Returns
+    -------
+    ret: tvm.ir.transform.Pass
+    """
+    return _ffi_api.MetaScheduleTuneTIR(work_dir, max_trials_global)
+
+
+def MetaScheduleTuneIRMod(
+    params: Dict[str, NDArray],
+    work_dir: str,
+    max_trials_global: int,
+) -> tvm.ir.transform.Pass:
+    """Tune Relax IRModule with MetaSchedule.
+    Parameters
+    ----------
+    params: Dict[str, NDArray]
+       model params
+    work_dir: str
+       work directory
+    max_trials_gloabl: int
+       maximum number of total trials allowed for tuning
+    Returns
+    -------
+    ret: tvm.ir.transform.Pass
     """
-    return _ffi_api.MetaScheduleApplyDatabase()
+    return _ffi_api.MetaScheduleTuneIRMod(params, work_dir, max_trials_global)
 
 
 def _wrap_class_function_pass(pass_cls, pass_info):

src/relax/backend/task_extraction.cc

@@ -57,7 +57,10 @@ class TaskExtractor : public ExprVisitor {
 
  private:
   explicit TaskExtractor(IRModule mod, Target target)
-      : mod_(std::move(mod)), target_(std::move(target)) {}
+      : mod_(std::move(mod)), target_(std::move(target)) {
+    normalize_mod_func_ = runtime::Registry::Get("tvm.meta_schedule.normalize_mod");
+    ICHECK(normalize_mod_func_) << "Normalization function is not found.";
+  }
 
   void VisitExpr_(const CallNode* call) final {
     static const Op& call_tir_op = Op::Get("relax.call_tir");
@@ -76,7 +79,7 @@ class TaskExtractor : public ExprVisitor {
       return;
     }
 
-    IRModule tir_mod({{global_var, func}});
+    IRModule tir_mod = (*normalize_mod_func_)(func);
     ExtractedTask task(/*task_name=*/global_var->name_hint,  //
                        /*mod=*/tir_mod,                      //
                        /*target=*/target_,                   //
@@ -90,6 +93,7 @@ class TaskExtractor : public ExprVisitor {
   Target target_;
   Array<ExtractedTask> tasks_;
   std::unordered_map<tir::PrimFunc, ExtractedTask, StructuralHash, StructuralEqual> func2task_;
+  const runtime::PackedFunc* normalize_mod_func_;
 };
 
 TVM_REGISTER_GLOBAL("relax.backend.MetaScheduleExtractTask")