def test_meta_schedule_custom_search_space():
mod = MatmulCustomized
context = TuneContext(
mod=mod,
target=Target("llvm"),
task_name="Custom Search Space Task",
sch_rules=[],
)
post_order_apply = PostOrderApply()
post_order_apply.initialize_with_tune_context(context)
post_order_apply.generate_design_space(mod)
called = False
def custom_search_space_func(sch: Schedule, _: BlockRV) -> List[Schedule]:
nonlocal called
called = True
return [sch]
register_func("tvm.meta_schedule.test.custom_search_space",
custom_search_space_func)
post_order_apply.generate_design_space(mod)
assert called
def test_tune_relay_manual_tir_vnni():
manual_tir_common(do_tune=False)
# pylint: disable=W0105
"""
We can inject and apply a custom TIR scheduling to a TE compute of interest, using
the "schedule_rule" annotation. For example, in topi/x86/dense.py we have the following
declaration for int8 dense targeting the VNNI instruction.
C = te.compute(
...
attrs={"schedule_rule": "meta_schedule.dense_vnni"},
)
When the meta scheduler encounters a TensorIR block with the "schedule_rule" annotation,
it looks up the packed func registry for a function that is associated with the given schedule
rule key ("meta_schedule.dense_vnni" in this example). The signature of such custom schedule
functions must be
(tir.schedule.Schedule, tir.schedule.BlockRV) -> [tir.schedule.Schedule].
The BlockRV argument corresponds to the TE compute annotated with "schedule_rule".
The relevant code is in meta_schedule/space_generator/post_order_apply.cc.
"""
# pylint: enable=W0105
def schedule_rule_dense_vnni(sch: Schedule, dense_block: BlockRV):
schedule_dense(dense_block, None, True, sch)
return [sch]
register_func("meta_schedule.dense_vnni", schedule_rule_dense_vnni)
manual_tir_common(do_tune=True)