def test_meta_schedule_task_scheduler_override_next_task_id_only(): # pylint: disable=invalid-name num_trials_per_iter = 6 max_trials_per_task = 101 tasks = [ TuneContext( MatmulModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_matmul), search_strategy=ReplayTrace(num_trials_per_iter, max_trials_per_task), task_name="Matmul", rand_state=42, ), TuneContext( MatmulReluModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_matmul), search_strategy=ReplayTrace(num_trials_per_iter, max_trials_per_task), task_name="MatmulRelu", rand_state=0xDEADBEEF, ), TuneContext( BatchMatmulModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_batch_matmul), search_strategy=ReplayTrace(num_trials_per_iter, max_trials_per_task), task_name="BatchMatmul", rand_state=0x114514, ), ] database = DummyDatabase() scheduler = MyTaskScheduler( tasks, DummyBuilder(), DummyRunner(), database, measure_callbacks=[ measure_callback.AddToDatabase(), ], max_trials=max_trials_per_task * len(tasks), ) scheduler.tune() assert len(database) == max_trials_per_task * len(tasks) for task in tasks: assert (len( database.get_top_k( database.commit_workload(task.mod), 100000, )) == max_trials_per_task)
def test_meta_schedule_task_scheduler_multiple(): num_trials_per_iter = 6 max_trials_per_task = 101 tasks = [ TuneContext( MatmulModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_matmul), search_strategy=ReplayTrace(num_trials_per_iter, max_trials_per_task), task_name="Matmul", rand_state=42, ), TuneContext( MatmulReluModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_matmul), search_strategy=ReplayTrace(num_trials_per_iter, max_trials_per_task), task_name="MatmulRelu", rand_state=0xDEADBEEF, ), TuneContext( BatchMatmulModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_batch_matmul), search_strategy=ReplayTrace(num_trials_per_iter, max_trials_per_task), task_name="BatchMatmul", rand_state=0x114514, ), ] database = DummyDatabase() round_robin = RoundRobin( tasks, [1.0], DummyBuilder(), DummyRunner(), database, measure_callbacks=[measure_callback.AddToDatabase()], max_trials=max_trials_per_task * len(tasks), ) round_robin.tune() assert len(database) == max_trials_per_task * len(tasks) for task in tasks: assert (len( database.get_top_k( database.commit_workload(task.mod), 100000, )) == max_trials_per_task)
def test_meta_schedule_task_scheduler_single(): num_trials_per_iter = 3 max_trials_per_task = 10 sch_fn = ScheduleFn(sch_fn=_schedule_matmul) replay = ReplayTrace(num_trials_per_iter, max_trials_per_task) task = TuneContext( MatmulModule, target=tvm.target.Target("llvm"), space_generator=sch_fn, search_strategy=replay, task_name="Test", rand_state=42, ) database = DummyDatabase() round_robin = RoundRobin( [task], [1.0], DummyBuilder(), DummyRunner(), database, measure_callbacks=[measure_callback.AddToDatabase()], max_trials=max_trials_per_task, ) round_robin.tune() assert len(database) == max_trials_per_task
def test_meta_schedule_task_scheduler_multiple(): num_trials_per_iter = 6 num_trials_total = 101 tasks = [ TuneContext( MatmulModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_matmul), search_strategy=ReplayTrace(num_trials_per_iter, num_trials_total), task_name="Matmul", rand_state=42, ), TuneContext( MatmulReluModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_matmul), search_strategy=ReplayTrace(num_trials_per_iter, num_trials_total), task_name="MatmulRelu", rand_state=0xDEADBEEF, ), TuneContext( BatchMatmulModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_batch_matmul), search_strategy=ReplayTrace(num_trials_per_iter, num_trials_total), task_name="BatchMatmul", rand_state=0x114514, ), ] database = DummyDatabase() round_robin = RoundRobin(tasks, DummyBuilder(), DummyRunner(), database) round_robin.tune() assert len(database) == num_trials_total * len(tasks) print(database.workload_reg) for task in tasks: assert len(database.get_top_k(database.commit_workload(task.mod), 1e9)) == num_trials_total
def test_meta_schedule_task_scheduler_single(): num_trials_per_iter = 3 num_trials_total = 10 sch_fn = ScheduleFn(sch_fn=_schedule_matmul) replay = ReplayTrace(num_trials_per_iter, num_trials_total) task = TuneContext( MatmulModule, target=tvm.target.Target("llvm"), space_generator=sch_fn, search_strategy=replay, task_name="Test", rand_state=42, ) database = DummyDatabase() round_robin = RoundRobin([task], DummyBuilder(), DummyRunner(), database) round_robin.tune() assert len(database) == num_trials_total
def test_meta_schedule_task_scheduler_override_next_task_id_only(): class MyTaskScheduler(PyTaskScheduler): done = set() def next_task_id(self) -> int: while len(self.done) != len(tasks): x = random.randint(0, len(tasks) - 1) task = tasks[x] if not task.is_stopped: """Calling base func via following route: Python side: PyTaskScheduler does not have `_is_task_running` Call TaskScheduler's `is_task_running`, which calls ffi C++ side: The ffi calls TaskScheduler's `is_task_running` But it is overridden in PyTaskScheduler PyTaskScheduler checks if the function is overridden in python If not, it returns the TaskScheduler's vtable, calling TaskScheduler::IsTaskRunning """ if self._is_task_running(x): # Same Here self._join_running_task(x) return x else: self.done.add(x) return -1 num_trials_per_iter = 6 num_trials_total = 101 tasks = [ TuneContext( MatmulModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_matmul), search_strategy=ReplayTrace(num_trials_per_iter, num_trials_total), task_name="Matmul", rand_state=42, ), TuneContext( MatmulReluModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_matmul), search_strategy=ReplayTrace(num_trials_per_iter, num_trials_total), task_name="MatmulRelu", rand_state=0xDEADBEEF, ), TuneContext( BatchMatmulModule, target=tvm.target.Target("llvm"), space_generator=ScheduleFn(sch_fn=_schedule_batch_matmul), search_strategy=ReplayTrace(num_trials_per_iter, num_trials_total), task_name="BatchMatmul", rand_state=0x114514, ), ] database = DummyDatabase() scheduler = MyTaskScheduler(tasks, DummyBuilder(), DummyRunner(), database) scheduler.tune() assert len(database) == num_trials_total * len(tasks) for task in tasks: assert len(database.get_top_k(database.commit_workload(task.mod), 1e9)) == num_trials_total
def test_meta_schedule_replay_trace(): num_trials_per_iter = 7 num_trials_total = 20 (example_sch,) = ScheduleFn(sch_fn=_schedule_matmul).generate_design_space(Matmul) replay = ReplayTrace(num_trials_per_iter=num_trials_per_iter, num_trials_total=num_trials_total) tune_context = TuneContext(mod=Matmul) replay.initialize_with_tune_context(tune_context) num_trials_each_round: List[int] = [] replay.pre_tuning([example_sch]) while True: candidates = replay.generate_measure_candidates() if candidates is None: break num_trials_each_round.append(len(candidates)) runner_results: List[RunnerResult] = [] for candidate in candidates: assert _is_trace_equal(candidate.sch, example_sch) runner_results.append(RunnerResult(run_secs=[0.5, 0.4, 0.3], error_msg=None)) replay.notify_runner_results(runner_results) replay.post_tuning() assert num_trials_each_round == [7, 7, 6]