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
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
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_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