def tune_each_task( mod, target, config, runner, work_dir, params, ): extracted_tasks = ms.extract_task_from_relay(mod, target, params) database = ms.database.JSONDatabase( path_workload=os.path.join(work_dir, "default_database_workload.json"), path_tuning_record=os.path.join(work_dir, "default_database_tuning_record.json"), ) for task in extracted_tasks: # pylint: disable=protected-access tune_context = ms.tune.Parse._tune_context( tune_context=None, mod=ms.tune.Parse._mod(task.dispatched[0]), target=target, config=config, task_name=task.task_name, space_generator=None, sch_rules=None, postprocs=None, mutator_probs=None, num_threads=os.cpu_count(), ) task_scheduler = ms.tune.Parse._task_scheduler( None, [tune_context], task_weights=[1.0], builder=ms.tune.Parse._builder(None), runner=ms.tune.Parse._runner(runner), database=database, max_trials=config.max_trials_per_task, cost_model=ms.tune.Parse._cost_model(None), measure_callbacks=ms.tune.Parse._callbacks(None), ) # pylint: enable=protected-access task_scheduler.tune() with target, ms.ApplyHistoryBest(database): with PassContext( opt_level=3, config={"relay.backend.use_meta_schedule": True}, ): return relay_build(mod, target=target, params=params)
def tune_relay( mod: IRModule, target: Union[str, Target], config: TuneConfig, work_dir: str, *, params: Optional[Dict[str, NDArray]] = None, builder: Optional[Builder] = None, runner: Optional[Runner] = None, database: Optional[Database] = None, cost_model: Optional[CostModel] = None, measure_callbacks: Optional[List[MeasureCallback]] = None, space: Optional[FnSpaceGenerator] = None, sch_rules: Optional[FnScheduleRule] = None, postprocs: Optional[FnPostproc] = None, mutator_probs: Optional[FnMutatorProb] = None, num_threads: Optional[int] = None, ) -> Module: """Tune a TIR IRModule with a given target. Parameters ---------- mod : IRModule The module to tune. target : Union[str, Target] The target to tune for. config : TuneConfig The search strategy config. params : Optional[Dict[str, tvm.runtime.NDArray]] The associated parameters of the program task_name : str The name of the task. work_dir : Optional[str] The working directory to save intermediate results. builder : Optional[Builder] The builder to use. runner : Optional[Runner] The runner to use. database : Optional[Database] The database to use. measure_callbacks : Optional[List[MeasureCallback]] The callbacks used during tuning. Returns ------- lib : Module The built runtime module for the given relay workload. """ # pylint: disable=import-outside-toplevel from tvm.relay import build as relay_build from .relay_integration import extract_task_from_relay # pylint: disable=protected-access, enable=import-outside-toplevel target = Parse._target(target) # pylint: enable=protected-access, # parse the tuning contexts extracted_tasks = extract_task_from_relay(mod, target, params) database = tune_extracted_tasks( extracted_tasks, config, work_dir, builder=builder, runner=runner, database=database, cost_model=cost_model, measure_callbacks=measure_callbacks, space=space, sch_rules=sch_rules, postprocs=postprocs, mutator_probs=mutator_probs, num_threads=num_threads, ) with target, autotvm_silencer(), ApplyHistoryBest(database): with PassContext( opt_level=3, config={"relay.backend.use_meta_schedule": True}, ): return relay_build(mod, target=target, params=params)
def tune_relay( mod: Union[RelayFunc, IRModule], target: Union[str, Target], config: SearchStrategyConfig, work_dir: str, *, params: Optional[Dict[str, NDArray]] = None, builder: Optional[Builder] = None, runner: Optional[Runner] = None, database: Optional[Database] = None, cost_model: Optional[CostModel] = None, measure_callbacks: Optional[List[MeasureCallback]] = None, task_scheduler: Optional[TaskScheduler] = None, space: Optional[FnSpaceGenerator] = None, sch_rules: Optional[FnScheduleRule] = None, postprocs: Optional[FnPostproc] = None, mutator_probs: Optional[FnMutatorProb] = None, num_threads: Optional[int] = None, ) -> Module: """Tune a TIR IRModule with a given target. Parameters ---------- mod : Union[RelayFunc, IRModule] The module to tune. target : Union[str, Target] The target to tune for. config : SearchStrategyConfig The search strategy config. params : Optional[Dict[str, tvm.runtime.NDArray]] The associated parameters of the program task_name : str The name of the task. work_dir : Optional[str] The working directory to save intermediate results. builder : Optional[Builder] The builder to use. runner : Optional[Runner] The runner to use. database : Optional[Database] The database to use. measure_callbacks : Optional[List[MeasureCallback]] The callbacks used during tuning. f_tune_context : Optional[TYPE_F_TUNE_CONTEXT] The function to create TuneContext. f_task_scheduler : Optional[TYPE_F_TASK_SCHEDULER] The function to create TaskScheduler. Returns ------- lib : Module The built runtime module for the given relay workload. """ logger.info("Working directory: %s", work_dir) extracted_tasks = extract_task_from_relay(mod, target, params) database = tune_extracted_tasks( extracted_tasks, target, config, work_dir, builder=builder, runner=runner, database=database, cost_model=cost_model, measure_callbacks=measure_callbacks, task_scheduler=task_scheduler, space=space, sch_rules=sch_rules, postprocs=postprocs, mutator_probs=mutator_probs, num_threads=num_threads, ) with ApplyHistoryBest(database): with tvm.transform.PassContext( opt_level=3, config={"relay.backend.use_meta_schedule": True}, ): return relay_build(mod, target=target, params=params)