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_measure_callback_fail(): @derived_object class FailingMeasureCallback(PyMeasureCallback): def apply( self, task_scheduler: TaskScheduler, task_id: int, measure_candidates: List[MeasureCandidate], builds: List[BuilderResult], results: List[RunnerResult], ) -> None: raise ValueError("test") measure_callback = FailingMeasureCallback() with pytest.raises(ValueError, match="test"): measure_callback.apply( RoundRobin([], [], DummyBuilder(), DummyRunner(), DummyDatabase(), max_trials=1), 0, [MeasureCandidate(Schedule(Matmul), None)], [BuilderResult("test_build", None)], [RunnerResult([1.0, 2.1], None)], )
def test_meta_schedule_measure_callback(): @derived_object class FancyMeasureCallback(PyMeasureCallback): def apply( self, task_scheduler: TaskScheduler, task_id: int, measure_candidates: List[MeasureCandidate], builds: List[BuilderResult], results: List[RunnerResult], ) -> None: assert len(measure_candidates) == 1 assert_structural_equal(measure_candidates[0].sch.mod, Matmul) assert (len(builds) == 1 and builds[0].error_msg is None and builds[0].artifact_path == "test_build") assert (len(results) == 1 and results[0].error_msg is None and len(results[0].run_secs) == 2) measure_callback = FancyMeasureCallback() measure_callback.apply( RoundRobin([], [], DummyBuilder(), DummyRunner(), DummyDatabase(), max_trials=1), 0, [MeasureCandidate(Schedule(Matmul), None)], [BuilderResult("test_build", None)], [RunnerResult([1.0, 2.1], None)], )
def test_meta_schedule_task_scheduler_NIE(): # pylint: disable=invalid-name @derived_object class NIETaskScheduler(PyTaskScheduler): pass with pytest.raises( TVMError, match="PyTaskScheduler's NextTaskId method not implemented!"): scheduler = NIETaskScheduler([], DummyBuilder(), DummyRunner(), DummyDatabase(), 1) scheduler.next_task_id()
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_avoid_cyclic(): # pylint: disable=invalid-name database = DummyDatabase() scheduler = MyTaskScheduler( [], DummyBuilder(), DummyRunner(), database, measure_callbacks=[ measure_callback.AddToDatabase(), ], ) test = weakref.ref(scheduler) # test if it can be destructed successfully del scheduler assert test() is None