allow missing schedule_rule in post order apply

masahi · Apr 7, 2022 · 2ce2066 · 2ce2066
1 parent 3a69353
commit 2ce2066
Showing 1 changed file with 13 additions and 2 deletions.
diff --git 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<String> ann = tir::GetAnn<String>(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<tir::Schedule> 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);
         }