From 05448f73d0b4d971d1e252e6d920bf256a7b8d24 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 10 Nov 2022 20:14:35 +0900 Subject: [PATCH 1/4] [MetaSchedule] Allow ignoring schedule_rule annotation --- src/meta_schedule/space_generator/post_order_apply.cc | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index 491af6e28f77..bc92fe5df28c 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -141,12 +141,6 @@ class PostOrderApplyNode : public SpaceGeneratorNode { stack.emplace_back(sch, blocks); continue; } - if (!ScheduleRule::IsApplyCustomRule(sch_rule)) { - if (tir::GetAnn(sch->GetSRef(block_rv), "schedule_rule").defined()) { - stack.emplace_back(sch, blocks); - continue; - } - } Array applied = sch_rule->Apply(sch, /*block=*/block_rv); for (const tir::Schedule& sch : applied) { stack.emplace_back(sch, blocks); From 8384056ea078b7c79e7d1b4612613ccf753206ec Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 11 Nov 2022 05:02:44 +0900 Subject: [PATCH 2/4] Revert "[MetaSchedule] Allow ignoring schedule_rule annotation" This reverts commit e738b16cc6d05e39966bf409ae8d424873449a40. --- src/meta_schedule/space_generator/post_order_apply.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index bc92fe5df28c..491af6e28f77 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -141,6 +141,12 @@ class PostOrderApplyNode : public SpaceGeneratorNode { stack.emplace_back(sch, blocks); continue; } + if (!ScheduleRule::IsApplyCustomRule(sch_rule)) { + if (tir::GetAnn(sch->GetSRef(block_rv), "schedule_rule").defined()) { + stack.emplace_back(sch, blocks); + continue; + } + } Array applied = sch_rule->Apply(sch, /*block=*/block_rv); for (const tir::Schedule& sch : applied) { stack.emplace_back(sch, blocks); From c6e876513b1713b64b98a3a3d83fc27d070bb6e7 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 11 Nov 2022 05:10:35 +0900 Subject: [PATCH 3/4] Unannotate schedule_rule if schedule func is not found --- src/meta_schedule/schedule_rule/apply_custom_rule.cc | 1 + tests/python/integration/test_auto_tensorize.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/meta_schedule/schedule_rule/apply_custom_rule.cc b/src/meta_schedule/schedule_rule/apply_custom_rule.cc index b133cd30c1f4..8b92d85b25c6 100644 --- a/src/meta_schedule/schedule_rule/apply_custom_rule.cc +++ b/src/meta_schedule/schedule_rule/apply_custom_rule.cc @@ -54,6 +54,7 @@ class ApplyCustomRuleNode : public ScheduleRuleNode { os << "\n " << GetCustomRuleName(ann.value(), key); } LOG(WARNING) << os.str(); + sch->Unannotate(block_rv, "schedule_rule"); } } return {sch}; diff --git a/tests/python/integration/test_auto_tensorize.py b/tests/python/integration/test_auto_tensorize.py index 8c06e147c01f..572da53b34fd 100644 --- a/tests/python/integration/test_auto_tensorize.py +++ b/tests/python/integration/test_auto_tensorize.py @@ -31,6 +31,7 @@ from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN SCH_RULES_FOR_VNNI = [ + ms.schedule_rule.ApplyCustomRule(), ms.schedule_rule.AutoInline( into_producer=False, into_consumer=True, From 7cd67aa766a8d73139b7246ae20a68b761402899 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 11 Nov 2022 05:51:43 +0900 Subject: [PATCH 4/4] lint --- src/meta_schedule/schedule_rule/apply_custom_rule.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/meta_schedule/schedule_rule/apply_custom_rule.cc b/src/meta_schedule/schedule_rule/apply_custom_rule.cc index 8b92d85b25c6..011e96e895d2 100644 --- a/src/meta_schedule/schedule_rule/apply_custom_rule.cc +++ b/src/meta_schedule/schedule_rule/apply_custom_rule.cc @@ -54,7 +54,7 @@ class ApplyCustomRuleNode : public ScheduleRuleNode { os << "\n " << GetCustomRuleName(ann.value(), key); } LOG(WARNING) << os.str(); - sch->Unannotate(block_rv, "schedule_rule"); + sch->Unannotate(block_rv, "schedule_rule"); } } return {sch};