From 2ce206699f10b03b9611c4683018f7e0c70c7eb5 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 7 Apr 2022 20:48:17 +0900 Subject: [PATCH] allow missing schedule_rule in post order apply --- .../space_generator/post_order_apply.cc | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index cae42bee4fe46..1b761a63538bd 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -136,19 +136,30 @@ 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")) { + bool has_schedule_rule = ann.defined() && runtime::Registry::Get(ann.value()) != nullptr; + + if (ann.defined() && !has_schedule_rule) { + LOG(WARNING) << "Custom schedule rule not found, ignoring schedule_rule annotation: " + << ann.value(); + } + + if ((!sch_rule.defined() && !has_schedule_rule) || + (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 { + } else if (has_schedule_rule) { const runtime::PackedFunc* f = runtime::Registry::Get(ann.value()); CHECK(f) << "ValueError: Custom schedule rule not found: " << ann.value(); applied = (*f)(sch, block_rv); } + for (const tir::Schedule& sch : applied) { stack.emplace_back(sch, blocks); }