From 43e0b2f7f98299679807aaf1ffb13cce2b5f5ce3 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 7 Apr 2022 18:25:14 +0900 Subject: [PATCH] rewrite_vnni -> rewrite_tensorize --- python/tvm/meta_schedule/postproc/__init__.py | 2 +- .../{rewrite_vnni.py => rewrite_tensorize.py} | 10 +++++----- python/tvm/meta_schedule/tune.py | 1 - .../{rewrite_vnni.cc => rewrite_tensorize.cc} | 16 ++++++++-------- .../multi_level_tiling_with_intrin.cc | 4 ++-- src/tir/schedule/analysis/analysis.cc | 2 +- 6 files changed, 17 insertions(+), 18 deletions(-) rename python/tvm/meta_schedule/postproc/{rewrite_vnni.py => rewrite_tensorize.py} (77%) rename src/meta_schedule/postproc/{rewrite_vnni.cc => rewrite_tensorize.cc} (88%) diff --git a/python/tvm/meta_schedule/postproc/__init__.py b/python/tvm/meta_schedule/postproc/__init__.py index 0b95960b313f5..39113bb90011a 100644 --- a/python/tvm/meta_schedule/postproc/__init__.py +++ b/python/tvm/meta_schedule/postproc/__init__.py @@ -22,4 +22,4 @@ from .rewrite_reduction_block import RewriteReductionBlock from .rewrite_unbound_block import RewriteUnboundBlock from .verify_gpu_code import VerifyGPUCode -from .rewrite_vnni import RewriteVNNI +from .rewrite_tensorize import RewriteTensorize diff --git a/python/tvm/meta_schedule/postproc/rewrite_vnni.py b/python/tvm/meta_schedule/postproc/rewrite_tensorize.py similarity index 77% rename from python/tvm/meta_schedule/postproc/rewrite_vnni.py rename to python/tvm/meta_schedule/postproc/rewrite_tensorize.py index b4de67184f61b..90cc756a692a3 100644 --- a/python/tvm/meta_schedule/postproc/rewrite_vnni.py +++ b/python/tvm/meta_schedule/postproc/rewrite_tensorize.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""A postprocessor that tensorize VNNI related components.""" +"""A postprocessor that tensorize related components.""" from tvm._ffi.registry import register_object from .. import _ffi_api @@ -22,11 +22,11 @@ import tvm.tir.tensor_intrin -@register_object("meta_schedule.RewriteVNNI") -class RewriteVNNI(Postproc): - """A postprocessor that tensorize VNNI related components.""" +@register_object("meta_schedule.RewriteTensorize") +class RewriteTensorize(Postproc): + """A postprocessor that tensorize related components.""" def __init__(self) -> None: self.__init_handle_by_constructor__( - _ffi_api.PostprocRewriteVNNI, # type: ignore # pylint: disable=no-member + _ffi_api.PostprocRewriteTensorize, # type: ignore # pylint: disable=no-member ) diff --git a/python/tvm/meta_schedule/tune.py b/python/tvm/meta_schedule/tune.py index ecf70053a5d10..86157e0fb32e8 100644 --- a/python/tvm/meta_schedule/tune.py +++ b/python/tvm/meta_schedule/tune.py @@ -214,7 +214,6 @@ def _postproc() -> List[Postproc]: M.DisallowDynamicLoop(), M.RewriteParallelVectorizeUnroll(), M.RewriteReductionBlock(), - M.RewriteVNNI(), ] @staticmethod diff --git a/src/meta_schedule/postproc/rewrite_vnni.cc b/src/meta_schedule/postproc/rewrite_tensorize.cc similarity index 88% rename from src/meta_schedule/postproc/rewrite_vnni.cc rename to src/meta_schedule/postproc/rewrite_tensorize.cc index 442c38b058d9b..f77e3d014c916 100644 --- a/src/meta_schedule/postproc/rewrite_vnni.cc +++ b/src/meta_schedule/postproc/rewrite_tensorize.cc @@ -27,7 +27,7 @@ using tir::LoopRV; using BlockPosition = std::tuple; -class RewriteVNNINode : public PostprocNode { +class RewriteTensorizeNode : public PostprocNode { public: // Inherited from PostprocNode void InitializeWithTuneContext(const TuneContext& context) final {} @@ -37,8 +37,8 @@ class RewriteVNNINode : public PostprocNode { void VisitAttrs(tvm::AttrVisitor* v) {} - static constexpr const char* _type_key = "meta_schedule.RewriteVNNI"; - TVM_DECLARE_FINAL_OBJECT_INFO(RewriteVNNINode, PostprocNode); + static constexpr const char* _type_key = "meta_schedule.RewriteTensorize"; + TVM_DECLARE_FINAL_OBJECT_INFO(RewriteTensorizeNode, PostprocNode); }; void CollectPostProcTasks(const tir::Schedule& sch, const String& func_name, @@ -60,7 +60,7 @@ void CollectPostProcTasks(const tir::Schedule& sch, const String& func_name, }); } -bool RewriteVNNINode::Apply(const tir::Schedule& sch) { +bool RewriteTensorizeNode::Apply(const tir::Schedule& sch) { std::vector tasks; for (const auto& kv : sch->mod()->functions) { GlobalVar g_var = kv.first; @@ -88,13 +88,13 @@ bool RewriteVNNINode::Apply(const tir::Schedule& sch) { return true; } -Postproc RewriteVNNI() { - ObjectPtr n = make_object(); +Postproc RewriteTensorize() { + ObjectPtr n = make_object(); return Postproc(n); } -TVM_REGISTER_NODE_TYPE(RewriteVNNINode); -TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteVNNI").set_body_typed(RewriteVNNI); +TVM_REGISTER_NODE_TYPE(RewriteTensorizeNode); +TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteTensorize").set_body_typed(RewriteTensorize); } // namespace meta_schedule } // namespace tvm diff --git a/src/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc b/src/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc index 3af8937d885c2..ba85629ab8049 100644 --- a/src/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc +++ b/src/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc @@ -35,10 +35,10 @@ class MultiLevelTilingWithIntrinNode : public MultiLevelTilingNode { } public: + String intrin_name; + static constexpr const char* _type_key = "meta_schedule.MultiLevelTilingWithIntrin"; TVM_DECLARE_FINAL_OBJECT_INFO(MultiLevelTilingWithIntrinNode, MultiLevelTilingNode); - - String intrin_name; }; ScheduleRule ScheduleRule::MultiLevelTilingWithIntrin( diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index e2d8dc76969df..b358a1461da6b 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -2096,7 +2096,7 @@ Optional GetTensorizeLoopMapping(const tir::ScheduleState& self, break; } } - if (block_loop == nullptr) { + if (desc_loop == nullptr) { return NullOpt; } // Step 4.3. Check divisibility of loop extents