[Metaschedule] Make custom schedule_rule registration optional (#10975)

See the discussion in #10793 (comment) for the context. Now I'm doing auto-tensorization on VNNI, I do need to be able to switch on / off `schedule_rule` freely.
apache · Apr 13, 2022 · dbfab5c · dbfab5c
1 parent 11b8cd3
commit dbfab5c
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 6 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,33 @@ 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")) {
+        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<tir::Schedule> 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
@@ -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