[Meta Schedule] Add customizable search space to PostOrderApply.

src/meta_schedule/space_generator/post_order_apply.cc

@@ -40,7 +40,16 @@ class BlockCollector : public tir::StmtVisitor {
         blocks_to_collect_.clear();
         VisitStmt(func->body);
         for (const String& block_name : blocks_to_collect_) {
-          results_.push_back(sch_->GetBlock(block_name, func_name_));
+          tir::BlockRV block_rv = sch_->GetBlock(block_name, func_name_);
+          // pick out the blocks with annotation for customized search space
+          if (Optional<ObjectRef> custom_sch_rule_name =
+                  sch_->Get(block_rv)->annotations.Get("rule")) {
+            const auto* custom_sch_rule_func =
+                runtime::Registry::Get(Downcast<String>(custom_sch_rule_name.value()));
+            (*custom_sch_rule_func)(sch_, block_rv);
+          } else {
+            results_.push_back(block_rv);
+          }
         }
       }
     }

tests/python/unittest/test_meta_schedule_post_order_apply.py

@@ -29,6 +29,7 @@
 from tvm.script import tir as T
 from tvm.target import Target
 from tvm.tir.schedule import BlockRV, Schedule
+from tvm import register_func
 
 
 # pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,
@@ -50,6 +51,22 @@ def main(a: T.handle, b: T.handle, c: T.handle) -> None:
                 C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
 
 
+@tvm.script.ir_module
+class MatmulCustomized:
+    @T.prim_func
+    def main(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main"})
+        A = T.match_buffer(a, (1024, 1024), "float32")
+        B = T.match_buffer(b, (1024, 1024), "float32")
+        C = T.match_buffer(c, (1024, 1024), "float32")
+        for i, j, k in T.grid(1024, 1024, 1024):
+            with T.block("matmul"):
+                T.block_attr({"rule": "tvm.meta_schedule.test.custom_search_space"})
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vi, vj] = 0.0
+                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
+
 @tvm.script.ir_module
 class DuplicateMatmul:
     @T.prim_func
@@ -338,5 +355,31 @@ def correct_trace(a, b, c, d):
         )
 
 
+def test_meta_schedule_post_order_apply_custom_search_space():
+    class DontCallThisRule(PyScheduleRule):
+        def initialize_with_tune_context(self, tune_context: "TuneContext") -> None:
+            pass
+
+        def apply(self, sch: Schedule, block: BlockRV) -> List[Schedule]:
+            raise Exception("This rule should not be called!")
+
+    @register_func("tvm.meta_schedule.test.custom_search_space")
+    def custom_search_space_func(sch: Schedule, block: BlockRV):
+        raise ValueError("Customized search space triggered!")
+
+    mod = MatmulCustomized
+    context = TuneContext(
+        mod=mod,
+        target=Target("llvm"),
+        task_name="Custom Search Space Task",
+        sch_rules=[DontCallThisRule()],
+    )
+    function_called = False
+    post_order_apply = PostOrderApply()
+    post_order_apply.initialize_with_tune_context(context)
+    with pytest.raises(ValueError, match="Customized search space triggered!"):
+        _ = post_order_apply.generate_design_space(mod)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__] + sys.argv[1:]))