From a8b6c95b8a50a2ac19d2108b796740a357018b1b Mon Sep 17 00:00:00 2001 From: sunli <466530738@qq.com> Date: Tue, 16 May 2023 07:47:01 +0000 Subject: [PATCH 1/2] using weak ptr --- cinn/hlir/framework/graph.h | 15 ++++++- cinn/hlir/pass/fusion_merge_pass.cc | 69 ++++++++++++++++++----------- cinn/hlir/pass/op_fusion_pass.cc | 3 +- 3 files changed, 57 insertions(+), 30 deletions(-) diff --git a/cinn/hlir/framework/graph.h b/cinn/hlir/framework/graph.h index 484d551508..0ab30f6f46 100644 --- a/cinn/hlir/framework/graph.h +++ b/cinn/hlir/framework/graph.h @@ -89,14 +89,25 @@ class Graph : public cinn::common::Graph { return first.get() == second.get(); } }; + + struct WeakGroupHasher { + size_t operator()(const std::weak_ptr& group) const noexcept { + return std::hash()(reinterpret_cast(group.lock().get())); + } + }; + struct WeakGroupComparator { + bool operator()(const std::weak_ptr& first, const std::weak_ptr& second) const noexcept { + return first.lock().get() == second.lock().get(); + } + }; // input groups - std::unordered_set, SharedGroupHasher, SharedGroupComparator> producer_groups; + std::unordered_set, WeakGroupHasher, WeakGroupComparator> producer_groups; // output grous std::unordered_set, SharedGroupHasher, SharedGroupComparator> consumer_groups; // fused sub-groups, used for fusion merge pass std::vector> fused_sub_groups; // if as sub-group, used for belong groups. - std::unordered_set, SharedGroupHasher, SharedGroupComparator> belong_groups; + std::unordered_set, WeakGroupHasher, WeakGroupComparator> belong_groups; // for op lowering. std::vector input_names; diff --git a/cinn/hlir/pass/fusion_merge_pass.cc b/cinn/hlir/pass/fusion_merge_pass.cc index c7c5a850bc..4182763a8f 100644 --- a/cinn/hlir/pass/fusion_merge_pass.cc +++ b/cinn/hlir/pass/fusion_merge_pass.cc @@ -32,9 +32,14 @@ using common::GraphNode; using Comparator = Graph::Group::SharedGroupComparator; using Hasher = Graph::Group::SharedGroupHasher; +using WeakComparator = Graph::Group::WeakGroupComparator; +using WeakHasher = Graph::Group::WeakGroupHasher; + using GroupPtr = std::shared_ptr; using GroupList = std::vector; +using WeakGroupPtr = std::weak_ptr; + using ConditionFunction = std::function; // Op Fusion Pass which performs Ops fusion, Ops are fused @@ -62,7 +67,8 @@ class FusionMergePassHelper : public FusionHelperBase { VLOG(3) << " Fused Sub-Group -> " << sub_group->group_id; } for (auto& producer : group->producer_groups) { - VLOG(3) << " Producer -> " << producer->group_id; + CHECK(!producer.expired()); + VLOG(3) << " Producer -> " << producer.lock()->group_id; } for (auto& consumer : group->consumer_groups) { VLOG(3) << " Consumer -> " << consumer->group_id; @@ -147,8 +153,9 @@ class FusionMergePassHelper : public FusionHelperBase { bool exist = false; for (auto& producer : group->producer_groups) { - if (fusion_groups_set.count(producer)) { - VLOG(4) << group->group_id << " " << producer->group_id; + CHECK(!producer.expired()); + if (fusion_groups_set.count(producer.lock())) { + VLOG(4) << group->group_id << " " << producer.lock()->group_id; exist = true; break; } @@ -312,10 +319,11 @@ class FusionMergePassHelper : public FusionHelperBase { } // producer group for (auto& producer : consumer->producer_groups) { - fused_group->producer_groups.insert(producer); + CHECK(!producer.expired()); + fused_group->producer_groups.insert(producer.lock()); // update producer's consumer - producer->consumer_groups.erase(consumer); - producer->consumer_groups.insert(fused_group); + producer.lock()->consumer_groups.erase(consumer); + producer.lock()->consumer_groups.insert(fused_group); } // consumer group for (auto& gconsumer : consumer->consumer_groups) { @@ -485,10 +493,11 @@ class FusionMergePassHelper : public FusionHelperBase { // producer groups for (auto& group : producer->producer_groups) { + CHECK(!group.expired()); fused_group->producer_groups.insert(group); // update producer's producer's consumer - group->consumer_groups.erase(producer); - group->consumer_groups.insert(fused_group); + group.lock()->consumer_groups.erase(producer); + group.lock()->consumer_groups.insert(fused_group); } // sub groups @@ -535,11 +544,12 @@ class FusionMergePassHelper : public FusionHelperBase { // producer nodes for (auto& group : consumer->producer_groups) { - if (group.get() != producer.get()) { + CHECK(!group.expired()); + if (group.lock().get() != producer.get()) { fused_group->producer_groups.insert(group); // update consumer's producer's consumer - group->consumer_groups.erase(consumer); - group->consumer_groups.insert(fused_group); + group.lock()->consumer_groups.erase(consumer); + group.lock()->consumer_groups.insert(fused_group); } } // consumer nodes @@ -723,15 +733,16 @@ class FusionMergePassHelper : public FusionHelperBase { auto& candidate = candidates.front(); candidates.pop(); for (auto& producer : candidate->producer_groups) { - if (producer.get() == producer_g.get()) { + CHECK(!producer.expired()); + if (producer.lock().get() == producer_g.get()) { continue; } - if (consumers.count(producer)) { + if (consumers.count(producer.lock())) { return true; } - if (!visited_set.count(producer)) { - visited_set.insert(producer); - candidates.push(producer); + if (!visited_set.count(producer.lock())) { + visited_set.insert(producer.lock()); + candidates.push(producer.lock()); } } } @@ -750,18 +761,19 @@ class FusionMergePassHelper : public FusionHelperBase { auto& candidate = candidates.front(); candidates.pop(); for (auto& producer : candidate->producer_groups) { - if (producer.get() == producer_g.get()) { + CHECK(!producer.expired()); + if (producer.lock().get() == producer_g.get()) { continue; } - if (producer->min_depth > check_upper_depth) { + if (producer.lock()->min_depth > check_upper_depth) { continue; } - if (consumers.count(producer)) { + if (consumers.count(producer.lock())) { return true; } - if (!visited_set.count(producer)) { - visited_set.insert(producer); - candidates.push(producer); + if (!visited_set.count(producer.lock())) { + visited_set.insert(producer.lock()); + candidates.push(producer.lock()); } } } @@ -818,7 +830,8 @@ class FusionMergePassHelper : public FusionHelperBase { if (consumer->belong_groups.size()) { // inset belong group to consumers. for (auto& belong_group : consumer->belong_groups) { - updated_consumers.insert(belong_group); + CHECK(belong_group.expired()); + updated_consumers.insert(belong_group.lock()); } } else { updated_consumers.insert(consumer); @@ -872,16 +885,18 @@ class FusionMergePassHelper : public FusionHelperBase { // update producer and consumer. for (auto& group : fusion_groups_) { - std::unordered_set producers; + std::unordered_set producers; std::unordered_set consumers; for (auto& producer : group->producer_groups) { - CHECK(producer->belong_groups.size()); - producers.insert(*producer->belong_groups.begin()); + CHECK(!producer.expired()); + CHECK(producer.lock()->belong_groups.size()); + producers.insert(*producer.lock()->belong_groups.begin()); } for (auto& consumer : group->consumer_groups) { CHECK(consumer->belong_groups.size()); - consumers.insert(*consumer->belong_groups.begin()); + CHECK(!consumer->belong_groups.begin()->expired()); + consumers.insert(consumer->belong_groups.begin()->lock()); } CHECK_EQ(group->producer_groups.size(), producers.size()); CHECK_EQ(group->consumer_groups.size(), consumers.size()); diff --git a/cinn/hlir/pass/op_fusion_pass.cc b/cinn/hlir/pass/op_fusion_pass.cc index 026f2c6195..7bc22c9e13 100644 --- a/cinn/hlir/pass/op_fusion_pass.cc +++ b/cinn/hlir/pass/op_fusion_pass.cc @@ -450,7 +450,8 @@ void OpFusionPassInternal(Graph* graph) { for (auto& group : graph->fusion_groups) { VLOG(3) << "Group Id : " << group->group_id; for (auto& producer : group->producer_groups) { - VLOG(3) << " producer group -> " << producer->group_id; + CHECK(!producer.expired()); + VLOG(3) << " producer group -> " << producer.lock()->group_id; } for (auto& consumer : group->consumer_groups) { VLOG(3) << " consumer group -> " << consumer->group_id; From 5d3eebeacd127d586ece2370c182090cbf1208ba Mon Sep 17 00:00:00 2001 From: BiynXu <244524405@qq.com> Date: Wed, 17 May 2023 21:48:59 +0800 Subject: [PATCH 2/2] fix bug in UpdateInputToConsumers --- cinn/hlir/pass/fusion_merge_pass.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cinn/hlir/pass/fusion_merge_pass.cc b/cinn/hlir/pass/fusion_merge_pass.cc index 4182763a8f..785d749308 100644 --- a/cinn/hlir/pass/fusion_merge_pass.cc +++ b/cinn/hlir/pass/fusion_merge_pass.cc @@ -830,7 +830,7 @@ class FusionMergePassHelper : public FusionHelperBase { if (consumer->belong_groups.size()) { // inset belong group to consumers. for (auto& belong_group : consumer->belong_groups) { - CHECK(belong_group.expired()); + CHECK(!belong_group.expired()); updated_consumers.insert(belong_group.lock()); } } else {