[MetaSchedule] Refactor ScheduleRule Attributes (#13195)

Author: junrushao

include/tvm/meta_schedule/schedule/cpu/.gitignore

@@ -16,37 +16,53 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-#ifndef TVM_META_SCHEDULE_SCHEDULE_RULE_AUTO_BIND_H_
-#define TVM_META_SCHEDULE_SCHEDULE_RULE_AUTO_BIND_H_
+#ifndef TVM_META_SCHEDULE_SCHEDULE_CUDA_THREAD_BIND_H_
+#define TVM_META_SCHEDULE_SCHEDULE_CUDA_THREAD_BIND_H_
 
-#include "../utils.h"
+#include <tvm/tir/schedule/schedule.h>
+
+#include <algorithm>
+#include <limits>
+#include <utility>
 
 namespace tvm {
 namespace meta_schedule {
 
 /*!
- * \brief Bind the given block if it is not bound to blockIdx or threadIdx.
+ * \brief Given candidates of thread_extents, make a sampler that use `sch->SampleCategorical`
+ * to return a random thread extent.
+ * \param sch The schedule
+ * \param thread_extents The candidate thread extents.
+ * \return A sampler that returns a random thread extent.
+ */
+std::function<tir::ExprRV(int64_t)> MakeFactorSampler(tir::Schedule sch,
+                                                      Array<Integer> thread_extents);
+
+/*!
+ * \brief Bind blockIdx.x and threadIdx.x to the given loop
  * \param sch The schedule.
- * \param block The block to be bound.
+ * \param loop The loop to be bound.
  * \param max_threadblocks The maximum number of threadblocks allowed.
- * \param max_threads The maximum number of threads allowed.
+ * \param max_threads_per_block The maximum number of threads allowed.
  * \param get_factor A function that returns the tiling factor.
  */
-void BindBlockThreadIdx(const tir::Schedule& sch, const tir::BlockRV& block,
-                        int64_t max_threadblocks, int64_t max_threads_per_block,
-                        std::function<tir::ExprRV(int64_t max_extent)> get_factor);
+Array<tir::LoopRV> BindSpatialLoop(tir::Schedule sch, tir::LoopRV loop,  //
+                                   int64_t max_threadblocks, int64_t max_threads_per_block,
+                                   std::function<tir::ExprRV(int64_t)> get_factor = nullptr);
 
 /*!
- * \brief Given candidates of thread_extents, make a sampler that use `sch->SampleCategorical`
- * to return a random thread extent.
- * \param sch The schedule
- * \param thread_extents The candidate thread extents.
- * \return A sampler that returns a random thread extent.
+ * \brief Bind the given block if it is not bound to blockIdx or threadIdx.
+ * \param sch The schedule.
+ * \param block The block to be bound.
+ * \param max_threadblocks The maximum number of threadblocks allowed.
+ * \param max_threads_per_block The maximum number of threads allowed.
+ * \param get_factor A function that returns the tiling factor.
  */
-std::function<tir::ExprRV(int64_t max_extent)> MakeFactorSampler(tir::Schedule sch,
-                                                                 Array<Integer> thread_extents);
+void BindBlockThreadIdx(tir::Schedule sch, tir::BlockRV block,  //
+                        int64_t max_threadblocks, int64_t max_threads_per_block,
+                        std::function<tir::ExprRV(int64_t max_extent)> get_factor = nullptr);
 
 }  // namespace meta_schedule
 }  // namespace tvm
 
-#endif  // TVM_META_SCHEDULE_SCHEDULE_RULE_AUTO_BIND_H_
+#endif  // TVM_META_SCHEDULE_SCHEDULE_CUDA_THREAD_BIND_H_

src/meta_schedule/schedule_rule/auto_bind.h renamed to include/tvm/meta_schedule/schedule/cuda/thread_bind.h

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_META_SCHEDULE_SCHEDULE_GENERIC_WINOGRAD_H_
+#define TVM_META_SCHEDULE_SCHEDULE_GENERIC_WINOGRAD_H_
+
+#include <tvm/tir/schedule/schedule.h>
+
+namespace tvm {
+namespace meta_schedule {
+
+/*!
+ * \brief Get the producer block of a given block.
+ * If there is a constant winograd transform matrix, inline it.
+ * \return The only producer block.
+ */
+tir::BlockRV GetWinogradProducerAndInlineConst(tir::Schedule sch, tir::BlockRV block);
+
+}  // namespace meta_schedule
+}  // namespace tvm
+
+#endif  // TVM_META_SCHEDULE_SCHEDULE_GENERIC_WINOGRAD_H_

include/tvm/meta_schedule/schedule/generic/winograd.h

@@ -99,6 +99,14 @@ class ScheduleRule : public runtime::ObjectRef {
    * \return The cloned schedule rule.
    */
   using FClone = runtime::TypedPackedFunc<ScheduleRule()>;
+  /*!
+   * \brief Create a rule that applies customized rules registered using block attribute
+   * `schedule_rule`. The rule will be dispatched according to target keys.
+   * \return The created schedule rule.
+   */
+  TVM_DLL static ScheduleRule ApplyCustomRule();
+  /*! \brief Check if the rule is `ApplyCustomRule` */
+  TVM_DLL static bool IsApplyCustomRule(const ScheduleRule& rule);
   /*!
    * \brief Create an auto-inline rule that inlines spatial blocks if it satisfies some conditions
    * \param into_producer If allows to inline a block into its producer

include/tvm/meta_schedule/schedule/x86/.gitignore

@@ -26,6 +26,7 @@
     postproc,
     relay_integration,
     runner,
+    schedule,
     schedule_rule,
     search_strategy,
     space_generator,
@@ -41,10 +42,7 @@
 from .mutator import Mutator
 from .postproc import Postproc
 from .profiler import Profiler
-from .relay_integration import (
-    is_meta_schedule_dispatch_enabled,
-    is_meta_schedule_enabled,
-)
+from .relay_integration import is_meta_schedule_enabled
 from .runner import Runner
 from .schedule_rule import ScheduleRule
 from .search_strategy import MeasureCandidate, SearchStrategy

include/tvm/meta_schedule/schedule_rule.h

@@ -377,7 +377,7 @@ def compile_relay(
     mod, target, params, pass_config, executor = _normalize_params(
         mod, target, params, pass_config, executor
     )
-    pass_config.setdefault("relay.backend.use_meta_schedule_dispatch", target.kind.name != "cuda")
+    pass_config.setdefault("relay.backend.use_meta_schedule_dispatch", True)
     with Profiler.timeit("PostTuningCompilation"):
         with target, _autotvm_silencer(), database:
             with transform.PassContext(
@@ -404,17 +404,3 @@ def is_meta_schedule_enabled() -> bool:
         "relay.backend.use_meta_schedule",
         False,
     )
-
-
-def is_meta_schedule_dispatch_enabled() -> bool:
-    """Return whether the meta-schedule dispatch is enabled.
-
-    Returns
-    -------
-    enabled: bool
-        Whether the meta schedule is enabled
-    """
-    return transform.PassContext.current().config.get(
-        "relay.backend.use_meta_schedule_dispatch",
-        False,
-    )

python/tvm/meta_schedule/__init__.py

@@ -0,0 +1,18 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Per-block schedule rules in MetaSchedule"""
+from . import cpu, cuda, generic, x86

python/tvm/meta_schedule/relay_integration.py

@@ -0,0 +1,17 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Per-block schedule rules in MetaSchedule for target key 'cpu'"""

python/tvm/meta_schedule/schedule/__init__.py

@@ -0,0 +1,17 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Per-block schedule rules in MetaSchedule for target key 'cuda'"""