wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
ff97fb6869c2eeecb7b4c09af41ce69f43f2c574
paddlepaddle--paddle/paddle/fluid/pir/transforms/sub_graph_detector.cc
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

1062 lines
37 KiB
C++
Raw Blame History

// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed 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.
#include "paddle/fluid/pir/transforms/sub_graph_detector.h"
#include <memory>
#include <climits>
#include <iterator>
#include <queue>
#include <regex>
#include <set>
#include <string>
#include <unordered_map>
#ifdef PADDLE_WITH_CINN
#include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
#include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
#include "paddle/cinn/hlir/framework/pir/utils.h"
#include "paddle/cinn/utils/string.h"
#endif
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
#include "paddle/fluid/pir/dialect/operator/trait/inplace.h"
#include "paddle/fluid/pir/utils/general_functions.h"
#include "paddle/pir/include/core/builder.h"
#include "paddle/pir/include/core/builtin_op.h"
#include "paddle/pir/include/core/op_trait.h"
#include "paddle/pir/include/dialect/control_flow/ir/cf_dialect.h"
#include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"
#include "paddle/pir/include/pass/pass.h"
#include "paddle/pir/include/pass/pass_registry.h"
#include "paddle/common/flags.h"
#include "paddle/common/macros.h"
#ifdef PADDLE_WITH_DNNL
#include "paddle/fluid/pir/dialect/operator/ir/onednn_op.h"
#include "paddle/fluid/pir/dialect/operator/ir/op_onednn_dialect.h"
#include "paddle/fluid/pir/dialect/operator/trait/onednn.h"
#endif
COMMON_DECLARE_bool(merge_all_horizontal_groups);
REGISTER_FILE_SYMBOLS(sub_graph_detector);
namespace pir {
std::vector<Operation*> InverselyTopologicalSort(Block* block) {
std::vector<Operation*> sort_ops;
std::unordered_map<Operation*, size_t> pending_count;
// step 1: initialize pending_cout for defined op
for (auto& op : *block) {
if (pending_count.find(&op) == pending_count.end()) {
pending_count[&op] = 0;
}
for (auto operand : GetUsedExternalValue(op)) {
if (!operand || !operand.defining_op()) {
continue;
}
auto* defined_op = operand.defining_op();
if (pending_count.find(defined_op) != pending_count.end()) {
++pending_count[defined_op];
} else {
pending_count[defined_op] = 1;
}
}
}
std::queue<Operation*> queue;
for (auto& op : *block) {
if (pending_count[&op] == 0) {
queue.push(&op);
}
}
while (!queue.empty()) {
auto* op = queue.front();
queue.pop();
sort_ops.push_back(op);
for (auto operand : GetUsedExternalValue(*op)) {
if (!operand || !operand.defining_op()) {
continue;
}
auto* defined_op = operand.defining_op();
--pending_count[defined_op];
if (defined_op && pending_count[defined_op] == 0 &&
defined_op->GetParent() == block) {
queue.push(defined_op);
}
}
}
PADDLE_ENFORCE_EQ(
block->size(),
sort_ops.size(),
common::errors::InvalidArgument("sort_ops.size() must be equal to "
"block.size(), but received %d != %d",
block->size(),
sort_ops.size()));
return sort_ops;
}
std::vector<Operation*> GetProducerOpsReverseSort(
Operation* op, const std::unordered_map<Operation*, int>& op2index) {
std::unordered_set<Operation*> producers;
std::vector<Operation*> vec_res;
for (auto operand : GetUsedExternalValue(*op)) {
if (!operand || !operand.defining_op()) {
continue;
}
auto* source_op = operand.defining_op();
if (source_op && !producers.count(source_op) &&
source_op->GetParent() == op->GetParent()) {
producers.insert(source_op);
PADDLE_ENFORCE(
op2index.count(source_op),
common::errors::PreconditionNotMet("source op MUST in op2index map"));
vec_res.emplace_back(source_op);
}
}
std::sort(
vec_res.begin(), vec_res.end(), [&op2index](Operation* a, Operation* b) {
return op2index.at(a) > op2index.at(b);
});
return vec_res;
}
std::vector<Operation*> GetProducerOps(Operation* op) {
std::vector<Operation*> producers;
for (auto operand : GetUsedExternalValue(*op)) {
if (!operand || !operand.defining_op()) {
continue;
}
auto* source_op = operand.defining_op();
if (source_op && source_op->GetParent() == op->GetParent()) {
producers.push_back(source_op);
}
}
return producers;
}
std::vector<Operation*> GetConsumerOps(
Operation* op, const std::unordered_map<Operation*, int>& op2index) {
std::vector<Operation*> consumers;
for (auto& result : op->results()) {
for (auto it = result.use_begin(); it != result.use_end(); ++it) {
auto parent_op = it->owner();
while (parent_op) {
if (op2index.count(parent_op)) {
consumers.push_back(parent_op);
break;
}
parent_op = parent_op->GetParentOp();
}
}
}
return consumers;
}
std::vector<std::pair<Value, Value>> GetInplaceValues(Operation* op) {
if (!op->HasInterface<paddle::dialect::OpYamlInfoInterface>()) return {};
auto op_info =
op->dyn_cast<paddle::dialect::OpYamlInfoInterface>().GetOpInfo();
auto input_info_list = std::get<0>(op_info);
auto output_info_list = std::get<2>(op_info);
auto inplace_info_map = std::get<3>(op_info).inplace;
std::unordered_map<std::string, Value> input_name_value;
std::unordered_map<std::string, Value> output_name_value;
for (size_t i = 0; i < input_info_list.size(); ++i) {
input_name_value[input_info_list[i].name] = op->operand_source(i);
}
for (size_t i = 0; i < output_info_list.size(); ++i) {
output_name_value[output_info_list[i].name] = op->result(i);
}
std::vector<std::pair<Value, Value>> inplace_values;
for (const auto& [out, in] : inplace_info_map) {
inplace_values.emplace_back(output_name_value[out], input_name_value[in]);
}
return inplace_values;
}
bool IsSideEffectButNotInplaceOp(Operation* op) {
return op->HasTrait<SideEffectTrait>() &&
!op->HasTrait<paddle::dialect::InplaceTrait>();
}
static std::string OpsDebugStr(std::vector<Operation*> ops) {
std::stringstream ss;
IrPrinter printer(ss);
for (const auto* op : ops) {
printer.PrintOperation(*op);
ss << "{" << op->id() << "}\n";
}
return ss.str();
}
struct SubGraph : public std::enable_shared_from_this<SubGraph> {
using SubGraphPtr = std::shared_ptr<SubGraph>;
SubGraph() = delete;
SubGraph(Operation* op, int index, bool subst)
: substitute(subst), topo_index(index), id(UniqueId()) {
ops.push_back(op);
}
void Merge(const SubGraphPtr& other);
static size_t UniqueId() {
static std::atomic<size_t> counter{0};
return counter++;
}
template <typename V>
static std::string JointName(const V& subgraphs) {
std::stringstream ss;
for (const auto& subgraph : subgraphs) {
ss << subgraph->name() << ", ";
}
auto str = ss.str();
return str.empty() ? str : str.substr(0, str.size() - 2);
}
std::string DebugStr(bool print_ops = false) const {
std::stringstream ss;
ss << "=========================================\n";
ss << name() << " (substitute=" << substitute << ", "
<< "index=" << topo_index << ", "
<< "size=" << ops.size() << ")\n";
if (print_ops) ss << OpsDebugStr(ops);
ss << "upstream: " << JointName(upstreams);
ss << "\ndownstream: " << JointName(downstreams);
return ss.str();
}
std::string name() const {
return std::string("Subgraph_") + std::to_string(id);
}
struct CompareById {
bool operator()(const SubGraphPtr& lhs, const SubGraphPtr& rhs) const {
// sort by reverse order of topo id
return lhs->id > rhs->id;
}
};
struct CompareByTopo {
bool operator()(const SubGraphPtr& lhs, const SubGraphPtr& rhs) const {
// sort by topo index
return lhs->topo_index > rhs->topo_index;
}
};
std::vector<Operation*> ops;
std::set<SubGraphPtr, CompareById> upstreams;
std::set<SubGraphPtr, CompareById> downstreams;
bool substitute; // whether this subgraph can be merged
int topo_index;
size_t id;
};
using SubGraphPtr = std::shared_ptr<SubGraph>;
void SubGraph::Merge(const SubGraphPtr& other) {
// Merge other subgraph into this subgraph:
// Inherit its upstreams/downstreams and ops
SubGraphPtr self = shared_from_this();
for (const auto& upstream : other->upstreams) {
if (upstream == self) continue;
upstream->downstreams.erase(other);
upstream->downstreams.insert(self);
upstreams.insert(upstream);
}
for (const auto& downstream : other->downstreams) {
if (downstream == self) continue;
downstream->upstreams.erase(other);
downstream->upstreams.insert(self);
downstreams.insert(downstream);
}
upstreams.erase(other);
downstreams.erase(other);
ops.insert(ops.begin(), other->ops.begin(), other->ops.end());
}
bool HasSinkRoute(const SubGraphPtr& source, const SubGraphPtr& target) {
if (source == target) return true;
std::unordered_set<SubGraphPtr> visited;
std::queue<SubGraphPtr> queue;
queue.push(source);
visited.insert(source);
while (!queue.empty()) {
SubGraphPtr cur = queue.front();
queue.pop();
if (cur == target) return true;
if (cur->topo_index > target->topo_index) continue;
for (const auto& subgraph : cur->downstreams) {
if (visited.count(subgraph)) continue;
queue.push(subgraph);
visited.insert(subgraph);
}
}
return false;
}
bool HasLiftRoute(const SubGraphPtr& source, const SubGraphPtr& target) {
if (source == target) return true;
std::unordered_set<SubGraphPtr> visited;
std::queue<SubGraphPtr> queue;
queue.push(source);
visited.insert(source);
while (!queue.empty()) {
SubGraphPtr cur = queue.front();
queue.pop();
if (cur == target) return true;
if (source->topo_index < target->topo_index) continue;
for (const auto& subgraph : cur->upstreams) {
if (visited.count(subgraph)) continue;
queue.push(subgraph);
visited.insert(subgraph);
}
}
return false;
}
bool HasRoute(const SubGraphPtr& up, const SubGraphPtr& down) {
return HasSinkRoute(up, down) || HasLiftRoute(down, up);
}
bool CanFuseUpstream2Downstream(const SubGraphPtr& upstream,
const SubGraphPtr& downstream) {
PADDLE_ENFORCE(upstream->downstreams.count(downstream) &&
downstream->upstreams.count(upstream),
::common::errors::InvalidArgument(
"Subgraphs to be fused must have direct relationship."));
auto up_downstreams = upstream->downstreams;
up_downstreams.erase(downstream);
auto down_upstreams = downstream->upstreams;
down_upstreams.erase(upstream);
if (up_downstreams.empty() || down_upstreams.empty()) return true;
for (const auto& subgraph : up_downstreams) {
if (HasSinkRoute(subgraph, downstream)) return false;
}
for (const auto& subgraph : down_upstreams) {
if (HasLiftRoute(subgraph, upstream)) return false;
}
return true;
}
std::optional<std::string> DetectCirclesInSubgraphs(
const std::vector<SubGraphPtr>& subgraph_list) {
std::set<SubGraphPtr, SubGraph::CompareById> subgraph_set(
subgraph_list.begin(), subgraph_list.end());
std::unordered_map<SubGraphPtr, size_t> in_degree;
std::unordered_map<SubGraphPtr, size_t> out_degree;
for (const auto& subgraph : subgraph_set) {
in_degree[subgraph] = subgraph->upstreams.size();
out_degree[subgraph] = subgraph->downstreams.size();
}
// Recursively remove nodes with in_degree or out_degree = 0
bool erase_flag = true;
while (erase_flag) {
erase_flag = false;
for (const auto& subgraph : subgraph_list) {
if (subgraph_set.count(subgraph) == 0) continue;
if (in_degree[subgraph] == 0) {
for (const auto& downstream : subgraph->downstreams) {
in_degree[downstream]--;
}
subgraph_set.erase(subgraph);
erase_flag = true;
continue;
}
if (out_degree[subgraph] == 0) {
for (const auto& upstream : subgraph->upstreams) {
out_degree[upstream]--;
}
subgraph_set.erase(subgraph);
erase_flag = true;
continue;
}
}
}
if (subgraph_set.empty()) return std::nullopt;
// If subgraph_set is not empty, there are circles in the subgraphs.
auto circle_size = subgraph_set.size();
std::stringstream ss;
ss << "Circles detected in subgraphs (size=" << circle_size << "): \n";
for (const auto& subgraph : subgraph_set) {
ss << subgraph->DebugStr() << "\n";
}
return std::make_optional(ss.str());
}
class SubgraphDetector {
public:
SubgraphDetector(Block* block, const OpClassifier& classifier);
void SubgraphFusion();
std::vector<GroupOpsVec> BuildGroups();
private:
void ReorderIndexOfSubgraphs();
bool MergeSource2Target(const SubGraphPtr& source, const SubGraphPtr& target);
void FallbackSubGraphFusion(const SubGraphPtr& source,
const SubGraphPtr& target,
const SubGraph& source_back,
const SubGraph& target_back);
void InitInplaceOpsOrder(const std::vector<Operation*>& inplace_ops);
bool CheckSideEffectOpsOrder() {
int last_index = INT_MIN;
for (const auto& op : side_effect_ops_) {
auto subgraph = GetOpSubgraph(op);
if (subgraph->topo_index < last_index) return false;
last_index = subgraph->topo_index;
}
for (const auto& ops : inplace_ops_order_) {
last_index = INT_MIN;
for (const auto& op : ops) {
auto subgraph = GetOpSubgraph(op);
if (subgraph->topo_index < last_index) return false;
last_index = subgraph->topo_index;
}
}
return true;
}
SubGraphPtr GetOpSubgraph(Operation* op) {
PADDLE_ENFORCE(
op2subgraph_.count(op),
::common::errors::InvalidArgument(
"Can not find op in op2subgraph_: \n%s", OpsDebugStr({op})));
return op2subgraph_.at(op);
}
std::vector<SubGraphPtr> GetSubgraphList() {
std::unordered_set<SubGraphPtr> subgraph_set;
std::vector<SubGraphPtr> subgraph_list;
for (const auto& op : sort_ops_) {
SubGraphPtr subgraph = GetOpSubgraph(op);
if (subgraph_set.count(subgraph)) continue;
subgraph_set.insert(subgraph);
subgraph_list.push_back(subgraph);
}
return subgraph_list;
}
std::unordered_map<Operation*, int> op2index_;
std::vector<Operation*> sort_ops_;
std::vector<Operation*> side_effect_ops_;
std::vector<std::vector<Operation*>> inplace_ops_order_;
std::unordered_map<Operation*, SubGraphPtr> op2subgraph_;
std::unordered_set<int> subgraph_index_set_;
};
void SubgraphDetector::ReorderIndexOfSubgraphs() {
// After merging subgraphs with direct relation, brother subgraphs with
// indirect relation may not be detected by index order. So we need to
// reorder the index of subgraphs.
using SubGraphQueue = std::priority_queue<SubGraphPtr,
std::vector<SubGraphPtr>,
SubGraph::CompareByTopo>;
SubGraphQueue queue; // min heap
std::unordered_map<SubGraphPtr, int> in_degree;
for (auto it = sort_ops_.rbegin(); it != sort_ops_.rend(); ++it) {
auto subgraph = GetOpSubgraph(*it);
if (in_degree.count(subgraph)) continue;
in_degree[subgraph] = subgraph->upstreams.size();
if (in_degree[subgraph] == 0) queue.push(subgraph);
}
subgraph_index_set_.clear();
int index = 0;
while (!queue.empty()) {
auto subgraph = queue.top();
queue.pop();
subgraph->topo_index = index++;
subgraph_index_set_.insert(subgraph->topo_index);
for (const auto& downstream : subgraph->downstreams) {
in_degree[downstream]--;
if (in_degree[downstream] == 0) queue.push(downstream);
}
}
}
bool SubgraphDetector::MergeSource2Target(const SubGraphPtr& source,
const SubGraphPtr& target) {
VLOG(6) << "Merge source: " << source->DebugStr();
VLOG(6) << "Merge target: " << target->DebugStr();
SubGraph source_back = *source;
SubGraph target_back = *target;
target->Merge(source);
for (const auto& op : source->ops) {
op2subgraph_[op] = target;
}
const auto& update_topo_index = [&]() -> void {
int max_index = std::max(source->topo_index, target->topo_index);
int min_index = std::min(source->topo_index, target->topo_index);
auto merged = target;
// Check if merged subgraph and its related subgraphs
// satisfy the topological order condition.
int upstream_max_index = -1, downstream_min_index = INT_MAX;
for (const auto& upstream : merged->upstreams) {
upstream_max_index = std::max(upstream->topo_index, upstream_max_index);
}
for (const auto& downstream : merged->downstreams) {
downstream_min_index =
std::min(downstream->topo_index, downstream_min_index);
}
// 1. If satisfy the topological order after merging, just set max_index
VLOG(6) << "Check if satisfy the topological order after merging";
if (min_index > upstream_max_index && max_index < downstream_min_index) {
merged->topo_index = max_index;
subgraph_index_set_.erase(min_index);
return;
}
// 2. If not satisfy the order, find a index between upstream_max_index
// and downstream_min_index while not in subgraph_index_set_.
VLOG(6) << "Try to find a valid index not in subgraph_index_set_";
for (int i = upstream_max_index + 1; i < downstream_min_index; ++i) {
if (!subgraph_index_set_.count(i)) {
merged->topo_index = i;
subgraph_index_set_.erase(min_index);
subgraph_index_set_.erase(max_index);
subgraph_index_set_.insert(i);
return;
}
}
// 3. If can not find a valid index, reorder topo index of all subgraphs.
VLOG(6) << "Reorder topo index of all subgraphs";
merged->topo_index = max_index;
ReorderIndexOfSubgraphs();
};
update_topo_index();
if (CheckSideEffectOpsOrder()) {
VLOG(6) << "Merged subgraph: " << target->DebugStr();
return true;
} else {
FallbackSubGraphFusion(source, target, source_back, target_back);
return false;
}
}
void SubgraphDetector::FallbackSubGraphFusion(const SubGraphPtr& source,
const SubGraphPtr& target,
const SubGraph& source_back,
const SubGraph& target_back) {
const auto fall_back_subgraph = [](const SubGraphPtr& subgraph,
const SubGraph& back) {
subgraph->ops = back.ops;
subgraph->upstreams = back.upstreams;
subgraph->downstreams = back.downstreams;
subgraph->topo_index = back.topo_index;
};
// 1. Update source and target subgraph
subgraph_index_set_.erase(target->topo_index);
subgraph_index_set_.insert(source_back.topo_index);
subgraph_index_set_.insert(target_back.topo_index);
fall_back_subgraph(source, source_back);
fall_back_subgraph(target, target_back);
for (const auto& op : source->ops) {
op2subgraph_[op] = source;
}
// 2. Update source's upstreams and downstreams
for (const auto& upstream : source->upstreams) {
if (upstream == target) continue;
upstream->downstreams.insert(source);
if (target->upstreams.count(upstream)) continue;
upstream->downstreams.erase(target);
}
for (const auto& downstream : source->downstreams) {
if (downstream == target) continue;
downstream->upstreams.insert(source);
if (target->downstreams.count(downstream)) continue;
downstream->upstreams.erase(target);
}
// 3. Check topo index and update
const auto need_reorder_topo_index = [&]() {
for (const auto& up : source->upstreams)
if (up->topo_index >= source->topo_index) return true;
for (const auto& down : source->downstreams)
if (down->topo_index <= source->topo_index) return true;
for (const auto& up : target->upstreams)
if (up->topo_index >= target->topo_index) return true;
for (const auto& down : target->downstreams)
if (down->topo_index <= target->topo_index) return true;
return false;
};
if (need_reorder_topo_index()) ReorderIndexOfSubgraphs();
VLOG(6) << "After fallback subgraph fusion: "
<< "\n source: " << source->DebugStr()
<< "\n target: " << target->DebugStr();
}
void SubgraphDetector::InitInplaceOpsOrder(
const std::vector<Operation*>& inplace_ops) {
std::unordered_map<Value, Value> inplace_map;
const auto& get_inplace_root_value = [&inplace_map](const Value& value) {
Value root = value;
std::unordered_set<Value> visited;
while (inplace_map.count(root)) {
if (visited.count(root)) break;
visited.insert(root);
root = inplace_map.at(root);
}
return root;
};
std::vector<std::set<Value>> inplace_values_sets;
for (const auto& op : inplace_ops) {
auto output_input_values = GetInplaceValues(op);
std::set<Value> inplace_input_values;
for (const auto& output_input_value : output_input_values) {
inplace_input_values.insert(
get_inplace_root_value(output_input_value.second));
inplace_map[output_input_value.first] = output_input_value.second;
}
inplace_values_sets.push_back(inplace_input_values);
}
std::set<Value> shared_inplace_values_set;
for (size_t i = 0; i < inplace_values_sets.size(); ++i) {
for (size_t j = i + 1; j < inplace_values_sets.size(); ++j) {
std::set_intersection(inplace_values_sets[i].begin(),
inplace_values_sets[i].end(),
inplace_values_sets[j].begin(),
inplace_values_sets[j].end(),
std::inserter(shared_inplace_values_set,
shared_inplace_values_set.begin()));
}
}
std::vector<std::vector<Operation*>> inplace_ops_order;
std::vector<Value> shared_inplace_values;
for (const auto& shared_value : shared_inplace_values_set) {
inplace_ops_order.emplace_back();
shared_inplace_values.push_back(shared_value);
for (size_t i = 0; i < inplace_values_sets.size(); ++i) {
if (inplace_values_sets[i].count(shared_value)) {
inplace_ops_order.back().push_back(inplace_ops[i]);
}
}
}
// If a value is inplaced by multiple ops, the order of ops which use this
// value after different inplace op also needs to be considered together.
for (size_t i = 0; i < inplace_ops_order.size(); ++i) {
auto only_inplace_ops = inplace_ops_order[i];
auto inplace_root_value = shared_inplace_values[i];
std::unordered_set<Operation*> ordered_ops_set(only_inplace_ops.begin(),
only_inplace_ops.end());
for (const auto& inplace_op : only_inplace_ops) {
Value output_inplace_value;
auto output_input_values = GetInplaceValues(inplace_op);
for (const auto& [output_value, _unused] : output_input_values) {
if (get_inplace_root_value(output_value) == inplace_root_value) {
output_inplace_value = output_value;
break;
}
}
if (output_inplace_value.use_empty()) continue;
for (auto use_iter = output_inplace_value.use_begin();
use_iter != output_inplace_value.use_end();
++use_iter) {
auto user_op = use_iter.owner();
if (ordered_ops_set.count(user_op)) continue;
ordered_ops_set.insert(user_op);
}
}
// Sort by origin order in blocks
std::vector<Operation*> ordered_ops(ordered_ops_set.begin(),
ordered_ops_set.end());
std::sort(ordered_ops.begin(),
ordered_ops.end(),
[this](const auto& lhs, const auto& rhs) {
return this->op2index_.at(lhs) < this->op2index_.at(rhs);
});
this->inplace_ops_order_.push_back(ordered_ops);
}
}
SubgraphDetector::SubgraphDetector(Block* block,
const OpClassifier& classifier) {
// init sort_ops_ in reverse topo order and op2index_ in topo order
std::vector<Operation*> inplace_ops;
int index = 0;
for (auto& op : *block) {
sort_ops_.push_back(&op);
op2index_[&op] = index++;
if (IsSideEffectButNotInplaceOp(&op)) {
side_effect_ops_.push_back(&op);
}
if (op.HasTrait<paddle::dialect::InplaceTrait>()) {
inplace_ops.push_back(&op);
}
}
std::reverse(sort_ops_.begin(), sort_ops_.end());
InitInplaceOpsOrder(inplace_ops);
// construct subgraphs and upstream/downstream relation
std::vector<SubGraphPtr> subgraph_list;
for (const auto& op : sort_ops_) {
bool substitute = classifier(*op);
auto subgraph = std::make_shared<SubGraph>(op, op2index_[op], substitute);
op2subgraph_[op] = subgraph;
subgraph_index_set_.insert(op2index_[op]);
subgraph_list.push_back(subgraph);
}
for (const auto& op : sort_ops_) {
auto subgraph = op2subgraph_[op];
for (const auto& producer : GetProducerOps(op)) {
if (!op2subgraph_.count(producer)) continue;
subgraph->upstreams.insert(op2subgraph_[producer]);
op2subgraph_[producer]->downstreams.insert(subgraph);
}
for (const auto& consumer : GetConsumerOps(op, op2index_)) {
if (!op2subgraph_.count(consumer)) continue;
subgraph->downstreams.insert(op2subgraph_[consumer]);
op2subgraph_[consumer]->upstreams.insert(subgraph);
}
}
VLOG(6) << "Subgraphs before building groups: ";
for (const auto& subgraph : subgraph_list) {
VLOG(6) << subgraph->DebugStr();
}
auto circle_info = DetectCirclesInSubgraphs(subgraph_list);
if (circle_info) {
PADDLE_THROW(::common::errors::PreconditionNotMet(
"Before building groups: %s", circle_info.value()));
}
}
void SubgraphDetector::SubgraphFusion() {
// Two subgraphs can be merged only if they have no route except direct
// connection between them (brother subgraphs should have no any route),
// otherwise a circle will be formed after merging them.
VLOG(4) << "Merge subgraphs with direct relation";
for (const auto& op : sort_ops_) {
auto downstream = GetOpSubgraph(op);
if (!downstream->substitute) continue;
for (const auto& producer : GetProducerOpsReverseSort(op, op2index_)) {
auto upstream = GetOpSubgraph(producer);
if (upstream == downstream || !upstream->substitute) continue;
if (CanFuseUpstream2Downstream(upstream, downstream)) {
MergeSource2Target(upstream, downstream);
}
}
}
VLOG(4) << "Merge brother subgraphs with same upstream";
for (const auto& op : sort_ops_) {
auto subgraph = GetOpSubgraph(op);
if (!subgraph->substitute) continue;
for (auto producer : GetProducerOpsReverseSort(op, op2index_)) {
if (GetOpSubgraph(producer) == subgraph) continue;
for (auto consumer : GetConsumerOps(producer, op2index_)) {
auto brother = GetOpSubgraph(consumer);
if (brother == subgraph || !brother->substitute) continue;
if (!HasRoute(subgraph, brother) && !HasRoute(brother, subgraph)) {
MergeSource2Target(brother, subgraph);
}
}
}
}
auto subgraph_list = GetSubgraphList();
VLOG(4) << "Merge non-related subgraphs (size=" << subgraph_list.size()
<< ")";
if (subgraph_list.size() > 2048 && !FLAGS_merge_all_horizontal_groups) return;
for (size_t i = 0; i < subgraph_list.size(); ++i) {
auto lhs = subgraph_list[i];
if (!lhs->substitute) continue;
for (size_t j = i + 1; j < subgraph_list.size();) {
auto rhs = subgraph_list[j];
if (lhs == rhs || !rhs->substitute || HasRoute(lhs, rhs) ||
HasRoute(rhs, lhs)) {
++j;
continue;
}
SubGraph lhs_back = *lhs;
SubGraph rhs_back = *rhs;
if (MergeSource2Target(rhs, lhs)) {
subgraph_list.erase(subgraph_list.begin() + j);
} else {
++j;
}
}
}
}
std::vector<GroupOpsVec> SubgraphDetector::BuildGroups() {
// 1. Get subgraph list in topo order
auto subgraph_list = GetSubgraphList();
std::reverse(subgraph_list.begin(), subgraph_list.end());
VLOG(6) << "Subgraphs after building groups: ";
for (const auto& subgraph : subgraph_list) {
VLOG(6) << subgraph->DebugStr();
}
auto circle_info = DetectCirclesInSubgraphs(subgraph_list);
if (circle_info) {
PADDLE_THROW(::common::errors::PreconditionNotMet(
"After building groups: %s", circle_info.value()));
}
// 2. Build group ops in subgraph which can be substituted
std::vector<GroupOpsVec> groups;
for (const auto& subgraph : subgraph_list) {
if (!subgraph->substitute) {
continue;
}
// sort group ops by natural increasing index.
std::vector<Operation*> group_ops(subgraph->ops.begin(),
subgraph->ops.end());
std::sort(
group_ops.begin(), group_ops.end(), [this](Operation* a, Operation* b) {
return this->op2index_.at(a) < this->op2index_.at(b);
});
groups.push_back(group_ops);
}
return groups;
}
std::vector<GroupOpsVec> DetectSubGraphs(Block* block,
const OpClassifier& classifier) {
auto subgraph_detector = SubgraphDetector(block, classifier);
subgraph_detector.SubgraphFusion();
return subgraph_detector.BuildGroups();
}
std::vector<Value> AnalysisOutputs(const GroupOpsVec& group_ops,
bool at_least_one_output) {
// Get output by ud chain
std::unordered_set<Operation*> op_set(group_ops.begin(), group_ops.end());
std::vector<Value> outputs;
for (auto* op : group_ops) {
for (size_t i = 0; i < op->num_results(); ++i) {
auto result = op->result(i);
for (auto use_iter = result.use_begin(); use_iter != result.use_end();
++use_iter) {
if (!op_set.count(use_iter->owner())) {
outputs.push_back(result);
break;
}
}
}
}
// NOTE: If all value are not used outside, we mark last op's results
// as outputs. But keep in mind that is risky.
if (at_least_one_output && outputs.size() == 0) {
for (size_t i = 0; i < group_ops.back()->num_results(); ++i) {
outputs.push_back(group_ops.back()->result(i));
}
}
return outputs;
}
namespace {
struct IncrementalOrder {
bool operator()(const Operation* lhs, const Operation* rhs) const {
PADDLE_ENFORCE_EQ(lhs->GetParent() == rhs->GetParent(),
true,
common::errors::PreconditionNotMet(
"lhs and rhs should have same parent block."));
auto lhs_iter = lhs->operator Block::ConstIterator();
auto rhs_iter = rhs->operator Block::ConstIterator();
auto end_iter = lhs->GetParent()->end();
while (lhs_iter != end_iter) {
lhs_iter++;
if (lhs_iter == rhs_iter) return true;
if (lhs_iter == end_iter) return false;
}
PADDLE_ENFORCE_EQ(
false,
true,
common::errors::InvalidArgument("rhs is not reachable from lhs."));
return false;
}
};
std::unordered_set<Operation*> GetUpstreamOpsAfterPosition(
const Operation* position_op,
const Block* block,
Operation* op,
std::unordered_set<Operation*>* visited_ops) {
std::unordered_set<Operation*> ops;
const auto& IsInBlock = [](const Operation* src_op, const Block* block) {
for (auto& item : *block) {
if (src_op->id() == item.id()) return true;
}
return false;
};
std::vector<Value> op_inputs = GetUsedExternalValue(*op);
for (auto value : op_inputs) {
if (!value || !value.defining_op()) continue;
Operation* defining_op = value.defining_op();
if (visited_ops->count(defining_op)) continue;
visited_ops->insert(defining_op);
if (!IsInBlock(defining_op, block)) continue;
if (IncrementalOrder()(defining_op, position_op)) continue;
ops.insert(defining_op);
auto recursive_ops = GetUpstreamOpsAfterPosition(
position_op, block, defining_op, visited_ops);
ops.insert(recursive_ops.begin(), recursive_ops.end());
}
return ops;
}
} // namespace
void MoveUpstreamOpBeforeGroup(const GroupOpsVec& group_ops,
Block* block,
Operation* insert_point_op) {
const auto moved_ops = [&]() {
std::set<Operation*, IncrementalOrder> ops_set;
std::unordered_set<Operation*> visited_ops;
for (auto& op : group_ops) {
auto upstream_ops =
GetUpstreamOpsAfterPosition(insert_point_op, block, op, &visited_ops);
ops_set.insert(upstream_ops.begin(), upstream_ops.end());
}
return ops_set;
}();
for (auto& op : moved_ops) {
if (op == insert_point_op) continue;
VLOG(4) << "Move " << op->id() << " " << op->name() << " before "
<< insert_point_op->id() << " " << insert_point_op->name();
op->MoveTo(block, insert_point_op->operator Block::Iterator());
}
}
Operation* FindInsertPoint(const GroupOpsVec& group_ops,
const std::vector<Value>& outputs) {
// Regard last op as insert position if there are no downstream ops between in
// group_ops.
Operation* first_op = group_ops.front();
Operation* insert_point_op = group_ops.back();
auto order_info = [&]() -> std::unordered_map<const Operation*, int64_t> {
std::unordered_map<const Operation*, int64_t> map;
// initialize the position index with block size by default.
auto block = insert_point_op->GetParent();
int64_t order = 0;
for (auto& op : *block) {
map[&op] = order++;
}
return map;
}();
for (auto* op : group_ops) {
if (order_info.at(op) > order_info.at(insert_point_op)) {
insert_point_op = op;
}
if (order_info.at(op) < order_info.at(first_op)) {
first_op = op;
}
}
auto begin = first_op->operator Block::ConstIterator();
auto end = ++(insert_point_op->operator Block::ConstIterator());
const std::unordered_set<Value> outputs_set(outputs.begin(), outputs.end());
const std::unordered_set<const Operation*> group_ops_set(group_ops.begin(),
group_ops.end());
const auto& IsDownstreamOp = [&](const Operation* op) -> bool {
if (group_ops_set.find(op) != group_ops_set.end()) return false;
for (auto& value : GetUsedExternalValue(*op)) {
if (outputs_set.find(value) != outputs_set.end()) {
return true;
}
}
return false;
};
// Find first downstream op as final insert position.
for (; begin != end; ++begin) {
if (IsDownstreamOp(begin)) {
insert_point_op = begin;
break;
}
}
return insert_point_op;
}
void ReplaceWithGroupOp(Block* block,
const GroupOpsVec& group_ops,
bool at_least_one_output) { // NOLINT
IrContext* ctx = IrContext::Instance();
#ifdef PADDLE_WITH_CINN
ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
#endif
#ifdef PADDLE_WITH_DNNL
ctx->GetOrRegisterDialect<paddle::dialect::OneDNNOperatorDialect>();
#endif
Builder builder = Builder(ctx, block);
const std::vector<Value> outputs =
AnalysisOutputs(group_ops, at_least_one_output);
// step 1: Analysis and insert group op before insert_point.
auto* insert_point = FindInsertPoint(group_ops, outputs);
MoveUpstreamOpBeforeGroup(group_ops, block, insert_point);
builder.set_insertion_point(insert_point);
VLOG(6) << "Insert GroupOp after " << insert_point->name();
// step 2: Replace the old op with GroupOp.
#ifdef PADDLE_WITH_CINN
auto new_group_op = [&]() -> cinn::dialect::GroupOp {
std::vector<Type> output_types;
for (auto& value : outputs) output_types.emplace_back(value.type());
auto group_op = builder.Build<cinn::dialect::GroupOp>(output_types);
for (auto op : group_ops) {
op->MoveTo(group_op.block(), group_op.block()->end());
}
return group_op;
}();
#else
auto new_group_op = [&]() -> GroupOp {
std::vector<Type> output_types;
for (auto& value : outputs) output_types.emplace_back(value.type());
auto group_op = builder.Build<GroupOp>(output_types);
for (auto op : group_ops) {
op->MoveTo(group_op.block(), group_op.block()->end());
}
return group_op;
}();
#endif
// step 3: Replace outputs of inner ops
const std::vector<Value> group_outs = new_group_op->results();
std::unordered_set<Operation*> inner_ops(group_ops.begin(), group_ops.end());
for (size_t i = 0; i < outputs.size(); ++i) {
outputs[i].ReplaceUsesWithIf(group_outs[i], [&inner_ops](OpOperand op) {
return !inner_ops.count(op.owner());
});
}
// step 4: Insert YieldOp for outputs
builder.SetInsertionPointToBlockEnd(new_group_op.block());
builder.Build<YieldOp>(outputs);
}
} // namespace pir
Reference in New Issue View Git Blame Copy Permalink