diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index cae42bee4fe4..09c134a101bc 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -136,19 +136,33 @@ class PostOrderApplyNode : public SpaceGeneratorNode { stack.emplace_back(sch, blocks); continue; } + Optional ann = tir::GetAnn(sch->GetSRef(block_rv), "schedule_rule"); - if (ann.defined() == sch_rule.defined() || (ann.defined() && ann.value() == "None")) { + const runtime::PackedFunc* custom_schedule_fn = + ann.defined() ? runtime::Registry::Get(ann.value()) : nullptr; + const bool has_schedule_rule = custom_schedule_fn != nullptr; + + if (ann.defined() && !has_schedule_rule) { + LOG(WARNING) << "Custom schedule rule not found, ignoring schedule_rule annotation: " + << ann.value(); + } + + if ((has_schedule_rule && sch_rule.defined()) || + (!has_schedule_rule && !sch_rule.defined()) || + (ann.defined() && ann.value() == "None")) { stack.emplace_back(sch, blocks); continue; } + Array applied{nullptr}; if (sch_rule.defined()) { applied = sch_rule.value()->Apply(sch, /*block=*/block_rv); } else { - const runtime::PackedFunc* f = runtime::Registry::Get(ann.value()); - CHECK(f) << "ValueError: Custom schedule rule not found: " << ann.value(); - applied = (*f)(sch, block_rv); + ICHECK(custom_schedule_fn) + << "ValueError: Custom schedule rule not found: " << ann.value(); + applied = (*custom_schedule_fn)(sch, block_rv); } + for (const tir::Schedule& sch : applied) { stack.emplace_back(sch, blocks); } diff --git a/tests/python/unittest/test_meta_schedule_post_order_apply.py b/tests/python/unittest/test_meta_schedule_post_order_apply.py index 40bb82f95929..e20da435f972 100644 --- a/tests/python/unittest/test_meta_schedule_post_order_apply.py +++ b/tests/python/unittest/test_meta_schedule_post_order_apply.py @@ -371,8 +371,8 @@ def test_meta_schedule_custom_search_space(): ) post_order_apply = PostOrderApply() post_order_apply.initialize_with_tune_context(context) - with pytest.raises(ValueError, match="Custom schedule rule not found"): - post_order_apply.generate_design_space(mod) + + post_order_apply.generate_design_space(mod) called = False