From 87c3402d8401e00ed4fffe7dad1258000b663544 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 7 Apr 2022 20:48:17 +0900 Subject: [PATCH 1/6] allow missing schedule_rule in post order apply --- .../space_generator/post_order_apply.cc | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index cae42bee4fe4..1b761a63538b 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -136,19 +136,30 @@ class PostOrderApplyNode : public SpaceGeneratorNode { stack.emplace_back(sch, blocks); continue; } + Optional ann = tir::GetAnn(sch->GetSRef(block_rv), "schedule_rule"); - if (ann.defined() == sch_rule.defined() || (ann.defined() && ann.value() == "None")) { + bool has_schedule_rule = ann.defined() && runtime::Registry::Get(ann.value()) != nullptr; + + if (ann.defined() && !has_schedule_rule) { + LOG(WARNING) << "Custom schedule rule not found, ignoring schedule_rule annotation: " + << ann.value(); + } + + if ((!sch_rule.defined() && !has_schedule_rule) || + (ann.defined() && ann.value() == "None")) { stack.emplace_back(sch, blocks); continue; } + Array applied{nullptr}; if (sch_rule.defined()) { applied = sch_rule.value()->Apply(sch, /*block=*/block_rv); - } else { + } else if (has_schedule_rule) { const runtime::PackedFunc* f = runtime::Registry::Get(ann.value()); CHECK(f) << "ValueError: Custom schedule rule not found: " << ann.value(); applied = (*f)(sch, block_rv); } + for (const tir::Schedule& sch : applied) { stack.emplace_back(sch, blocks); } From 91b7d6cfe0f4ada54943ca75cd48f93861bbbc17 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 12 Apr 2022 18:37:29 +0900 Subject: [PATCH 2/6] update test --- tests/python/unittest/test_meta_schedule_post_order_apply.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/unittest/test_meta_schedule_post_order_apply.py b/tests/python/unittest/test_meta_schedule_post_order_apply.py index 40bb82f95929..e20da435f972 100644 --- 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 From ae87f1a36d2daf58af32de51657987516cf12c3a Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 13 Apr 2022 06:13:45 +0900 Subject: [PATCH 3/6] Skip applying sch_rule when both ann and sch_rule are defined --- src/meta_schedule/space_generator/post_order_apply.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index 1b761a63538b..af67e6334845 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -145,7 +145,7 @@ class PostOrderApplyNode : public SpaceGeneratorNode { << ann.value(); } - if ((!sch_rule.defined() && !has_schedule_rule) || + if ((ann.defined() && sch_rule.defined()) || (!has_schedule_rule && !sch_rule.defined()) || (ann.defined() && ann.value() == "None")) { stack.emplace_back(sch, blocks); continue; From a4651db3d7823307ed52d4132acbe8fedfd71741 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 13 Apr 2022 13:11:40 +0900 Subject: [PATCH 4/6] fixed condition for real --- src/meta_schedule/space_generator/post_order_apply.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index af67e6334845..c72b302e6d33 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -145,7 +145,8 @@ class PostOrderApplyNode : public SpaceGeneratorNode { << ann.value(); } - if ((ann.defined() && sch_rule.defined()) || (!has_schedule_rule && !sch_rule.defined()) || + if ((has_schedule_rule && sch_rule.defined()) || + (!has_schedule_rule && !sch_rule.defined()) || (ann.defined() && ann.value() == "None")) { stack.emplace_back(sch, blocks); continue; From acf75dbc4578ec85fce9f571bfd478bdf938a3ca Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 13 Apr 2022 14:54:14 +0900 Subject: [PATCH 5/6] address comment --- src/meta_schedule/space_generator/post_order_apply.cc | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index c72b302e6d33..81f582162bf7 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -138,7 +138,9 @@ class PostOrderApplyNode : public SpaceGeneratorNode { } Optional ann = tir::GetAnn(sch->GetSRef(block_rv), "schedule_rule"); - bool has_schedule_rule = ann.defined() && runtime::Registry::Get(ann.value()) != nullptr; + 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: " @@ -155,10 +157,9 @@ class PostOrderApplyNode : public SpaceGeneratorNode { Array applied{nullptr}; if (sch_rule.defined()) { applied = sch_rule.value()->Apply(sch, /*block=*/block_rv); - } else if (has_schedule_rule) { - const runtime::PackedFunc* f = runtime::Registry::Get(ann.value()); - CHECK(f) << "ValueError: Custom schedule rule not found: " << ann.value(); - applied = (*f)(sch, block_rv); + } else { + 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) { From 55996fe7a06ee71a86d664576f4bf075cf2d4f6c Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 13 Apr 2022 14:55:08 +0900 Subject: [PATCH 6/6] cpplint --- src/meta_schedule/space_generator/post_order_apply.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc index 81f582162bf7..09c134a101bc 100644 --- a/src/meta_schedule/space_generator/post_order_apply.cc +++ b/src/meta_schedule/space_generator/post_order_apply.cc @@ -158,7 +158,8 @@ class PostOrderApplyNode : public SpaceGeneratorNode { if (sch_rule.defined()) { applied = sch_rule.value()->Apply(sch, /*block=*/block_rv); } else { - ICHECK(custom_schedule_fn) << "ValueError: Custom schedule rule not found: " << ann.value(); + ICHECK(custom_schedule_fn) + << "ValueError: Custom schedule rule not found: " << ann.value(); applied = (*custom_schedule_fn)(sch, block_rv); }