Yep, this is possible to mix different ways of scheduling (manual, handcrafted search space, auto-generated space) with Relax. To do this, you may do task extraction, and customize it, and then plug it back.
Example:
extracted_tasks = ms.relax_integration.extract_tasks(mod, target, params=None)
tasks, task_weights = inject_schedule(extracted_tasks)
^^^^^^^^^^^^^^^
customize the way we schedule
database = ms.tune_tasks(
tasks=tasks,
task_weights=task_weights,
work_dir=WORK_DIR,
max_trials_global=500,
num_trials_per_iter=32,
)
relax_exec = ms.relax_integration.compile_relax(
database,
mod=mod,
target=target,
params=None,
)
A typical way of injecting schedules looks like:
def inject_schedule(extracted_tasks, work_dir):
tasks = []
task_weights = []
for task, logger, rand_state in zip(
extracted_tasks,
ms.logging.get_loggers_from_work_dir(work_dir, [t.task_name for t in extracted_tasks]),
ms.utils.fork_seed(None, n=len(extracted_tasks)),
):
if task.task_name == "XXX":
# handcrafted search space
# or handcrafted schedule without search
space = ms.space_generator.ScheduleFn(
sch_fn = my_schedule_function,
sch_rules=[],
postprocs=[],
)
else:
# auto-generated search space
space = "post-order-apply"
tasks.append(
ms.TuneContext(
mod=task.dispatched[0],
target=task.target,
space_generator=space,
task_name=task.task_name,
logger=logger,
rand_state=rand_state,
).clone()
)
task_weights.append(task.weight)
return tasks, task_weights