wehub-resource-sync
291 lines
11 KiB
C++
291 lines
11 KiB
C++
// Copyright (c) 2024 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.
|
|
|
|
#pragma once
|
|
|
|
#include <functional>
|
|
#include <unordered_set>
|
|
#include <variant>
|
|
#include <vector>
|
|
#include "glog/logging.h"
|
|
#include "paddle/cinn/operator_fusion/fusion_tracker/tracker.h"
|
|
#include "paddle/cinn/operator_fusion/pir_graph_analyzing/loop_axis_mapping.h"
|
|
#include "paddle/cinn/operator_fusion/pir_graph_analyzing/loop_transform_analysis.h"
|
|
#include "paddle/cinn/operator_fusion/utils.h"
|
|
#include "paddle/pir/include/core/operation.h"
|
|
|
|
namespace cinn::fusion {
|
|
|
|
enum class PatternType {
|
|
Trivial = 0,
|
|
Reduce,
|
|
ReduceTree,
|
|
ReduceTreePlusTrivial,
|
|
Anchor,
|
|
Horizontal,
|
|
Unsupported = -1,
|
|
};
|
|
|
|
struct PatternContent {
|
|
explicit PatternContent(pir::Operation* op) : op(op) {}
|
|
pir::Operation* op;
|
|
bool operator==(const PatternContent& other) const { return op == other.op; }
|
|
};
|
|
|
|
struct PatternBase {
|
|
explicit PatternBase(const std::string& id, const FusionTrackerPtr& tracker)
|
|
: id_(id), tracker_(tracker) {}
|
|
explicit PatternBase(const std::string& id,
|
|
const FusionTrackerPtr& tracker,
|
|
const std::vector<pir::Operation*>& ops)
|
|
: id_(id), tracker_(tracker), ops_(ops) {}
|
|
std::string id_;
|
|
std::string id() const { return id_; }
|
|
std::vector<pir::Operation*> ops_;
|
|
std::vector<pir::Operation*> ops() const { return ops_; }
|
|
FusionTrackerPtr tracker_;
|
|
void update_tracker() const {}
|
|
LoopAxisMapping loop_axis_mapping_;
|
|
LoopAxisMapping loop_axis_mapping() const { return loop_axis_mapping_; }
|
|
void set_loop_axis_mapping(const LoopAxisMapping& loop_axis_mapping) {
|
|
loop_axis_mapping_ = loop_axis_mapping;
|
|
}
|
|
};
|
|
|
|
#define DEFINE_PATTERN_STATIC_ATTR(pattern) \
|
|
static PatternType type() { return PatternType::pattern; } \
|
|
static std::string UniqueId() { \
|
|
static std::atomic<int64_t> counter = 0; \
|
|
return std::string(#pattern) + "_" + std::to_string(++counter); \
|
|
}
|
|
|
|
struct TrivialPattern : public PatternBase {
|
|
explicit TrivialPattern(const std::vector<pir::Operation*>& ops,
|
|
pir::Operation* sink_op,
|
|
const FusionTrackerPtr& tracker)
|
|
: PatternBase(UniqueId(), tracker, ops), sink_op_(sink_op) {}
|
|
DEFINE_PATTERN_STATIC_ATTR(Trivial);
|
|
pir::Operation* sink_op_;
|
|
pir::Operation* sink_op() const { return sink_op_; }
|
|
};
|
|
|
|
struct ReducePattern : public PatternBase {
|
|
explicit ReducePattern(const std::vector<pir::Operation*>& ops,
|
|
const FusionTrackerPtr& tracker)
|
|
: PatternBase(UniqueId(), tracker, ops) {}
|
|
DEFINE_PATTERN_STATIC_ATTR(Reduce);
|
|
pir::Operation* GetReduceOp() const { return ops_.back(); }
|
|
};
|
|
|
|
struct ReduceTreePattern : public PatternBase {
|
|
explicit ReduceTreePattern(const std::vector<ReduceTreePattern>& children,
|
|
const ReducePattern& root,
|
|
const FusionTrackerPtr& tracker)
|
|
: PatternBase(UniqueId(), tracker), children_(children), root_(root) {
|
|
cur_id_ = id_;
|
|
}
|
|
DEFINE_PATTERN_STATIC_ATTR(ReduceTree);
|
|
|
|
std::vector<pir::Operation*> ops() const {
|
|
std::vector<pir::Operation*> result{root_.ops()};
|
|
for (const auto& child : children_) {
|
|
result = UniqueConcatVector(result, child.ops());
|
|
}
|
|
return result;
|
|
}
|
|
const ReducePattern& GetRootPattern() const { return root_; }
|
|
const std::vector<ReduceTreePattern>& children() const { return children_; }
|
|
std::vector<ReduceTreePattern>& children() { return children_; }
|
|
void InsertChild(const ReduceTreePattern& child) {
|
|
children_.push_back(child);
|
|
}
|
|
std::vector<ReducePattern> FlattenReducePattern() const {
|
|
std::vector<ReducePattern> result{root_};
|
|
for (const auto& child : children_) {
|
|
result = ConcatVector(result, child.FlattenReducePattern());
|
|
}
|
|
return result;
|
|
}
|
|
|
|
mutable std::string cur_id_;
|
|
std::string cur_id() const { return cur_id_; }
|
|
void reset_cur_id(std::string id) const { cur_id_ = id; }
|
|
|
|
void update_tracker() const {
|
|
const std::string& root_name = GetRootPattern().id();
|
|
std::vector<std::string> names;
|
|
UpdateTrackerImpl(
|
|
root_name, *this, std::vector<size_t>(), this->tracker_, &names);
|
|
tracker_->append(std::make_shared<CombineInstr>(names, cur_id()));
|
|
}
|
|
|
|
void UpdateTrackerImpl(const std::string root_name,
|
|
const ReduceTreePattern& root,
|
|
const std::vector<size_t>& fake_reduce_iter_idx,
|
|
FusionTrackerPtr tracker,
|
|
std::vector<std::string>* names) const {
|
|
// Apply a brunch of tracker to get a output_name of ReduceTreePattern.
|
|
// names and trackers collect all the needed fusion nodes.
|
|
for (const auto& child : root.children()) {
|
|
auto origin_child_id = child.cur_id();
|
|
auto new_child_id = GetNewTmpId(origin_child_id);
|
|
child.reset_cur_id(new_child_id);
|
|
tracker->append(
|
|
std::make_shared<TmpTransformInstr>(origin_child_id,
|
|
root_name,
|
|
new_child_id,
|
|
root.cur_id(),
|
|
fake_reduce_iter_idx));
|
|
UpdateTrackerImpl(
|
|
new_child_id, child, fake_reduce_iter_idx, tracker, names);
|
|
}
|
|
names->push_back(root.cur_id());
|
|
}
|
|
|
|
private:
|
|
std::vector<ReduceTreePattern> children_;
|
|
ReducePattern root_;
|
|
};
|
|
|
|
struct ReduceTreePlusTrivialPattern : public PatternBase {
|
|
explicit ReduceTreePlusTrivialPattern(const ReduceTreePattern& tree,
|
|
const TrivialPattern& sink_trivial,
|
|
const FusionTrackerPtr& tracker)
|
|
: PatternBase(UniqueId(), tracker),
|
|
tree(tree),
|
|
sink_trivial(sink_trivial) {}
|
|
DEFINE_PATTERN_STATIC_ATTR(ReduceTreePlusTrivial);
|
|
|
|
ReduceTreePattern tree;
|
|
TrivialPattern sink_trivial;
|
|
std::vector<size_t> fake_reduce_iter_idx;
|
|
|
|
std::vector<pir::Operation*> ops() const {
|
|
return UniqueConcatVector(tree.ops(), sink_trivial.ops());
|
|
}
|
|
|
|
void update_tracker() const {
|
|
const std::string& root_name = id();
|
|
const std::string& origin_tree_id = tree.cur_id();
|
|
const std::string& new_tree_id = GetNewTmpId(origin_tree_id);
|
|
tree.reset_cur_id(new_tree_id);
|
|
std::vector<std::string> names;
|
|
tracker_->append(std::make_shared<TmpTransformInstr>(origin_tree_id,
|
|
sink_trivial.id(),
|
|
new_tree_id,
|
|
root_name,
|
|
fake_reduce_iter_idx));
|
|
tree.UpdateTrackerImpl(
|
|
new_tree_id, tree, fake_reduce_iter_idx, this->tracker_, &names);
|
|
names.push_back(root_name);
|
|
// optimize the loop range of R + T for speed up.
|
|
tracker_->append(std::make_shared<TrivialLoopAlignInstr>(
|
|
new_tree_id, root_name, root_name, fake_reduce_iter_idx));
|
|
// collect all the Expr and represent the root_name.
|
|
tracker_->append(std::make_shared<CombineInstr>(names, root_name));
|
|
}
|
|
};
|
|
|
|
struct AnchorPattern : public PatternBase {
|
|
explicit AnchorPattern(const std::vector<pir::Operation*>& ops,
|
|
const FusionTrackerPtr& tracker,
|
|
const LoopAxisMapping& loop_axis_mapping)
|
|
: PatternBase(UniqueId(), tracker, ops) {
|
|
set_loop_axis_mapping(loop_axis_mapping);
|
|
}
|
|
DEFINE_PATTERN_STATIC_ATTR(Anchor);
|
|
};
|
|
|
|
struct UnsupportedPattern : public PatternBase {
|
|
explicit UnsupportedPattern(const std::vector<pir::Operation*>& ops,
|
|
const FusionTrackerPtr& tracker)
|
|
: PatternBase(UniqueId(), tracker, ops) {}
|
|
DEFINE_PATTERN_STATIC_ATTR(Unsupported);
|
|
};
|
|
|
|
using StmtPattern = std::variant<TrivialPattern,
|
|
ReducePattern,
|
|
ReduceTreePattern,
|
|
ReduceTreePlusTrivialPattern,
|
|
AnchorPattern,
|
|
UnsupportedPattern>;
|
|
|
|
static PatternType GetPatternType(const StmtPattern& s) {
|
|
return std::visit([](const auto& impl) { return impl.type(); }, s);
|
|
}
|
|
|
|
static std::string GetPatternId(const StmtPattern& s) {
|
|
return std::visit([](const auto& impl) { return impl.id(); }, s);
|
|
}
|
|
|
|
static FusionTrackerPtr GetFusionTracker(const StmtPattern& s) {
|
|
return std::visit([](const auto& impl) { return impl.tracker_; }, s);
|
|
}
|
|
|
|
static std::vector<pir::Operation*> GetOpsInPattern(
|
|
const StmtPattern& pattern) {
|
|
return std::visit([](const auto& impl) { return impl.ops(); }, pattern);
|
|
}
|
|
|
|
static LoopAxisMapping GetPatternLoopAxisMapping(const StmtPattern& s) {
|
|
return std::visit([](const auto& impl) { return impl.loop_axis_mapping(); },
|
|
s);
|
|
}
|
|
|
|
static std::unordered_set<pir::Value> GetPatternInputValuesIncludeInner(
|
|
const StmtPattern& A) {
|
|
std::unordered_set<pir::Value> result;
|
|
for (const auto& op : GetOpsInPattern(A)) {
|
|
for (const auto& value : op->operands()) {
|
|
result.insert(value.source());
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
static std::unordered_set<pir::Value> GetPatternOutputValuesIncludedInner(
|
|
const StmtPattern& A) {
|
|
std::unordered_set<pir::Value> result;
|
|
for (const auto& op : GetOpsInPattern(A)) {
|
|
for (const auto& value : op->results()) {
|
|
result.insert(value);
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
static std::unordered_set<pir::Value> GetPatternInputValues(
|
|
const StmtPattern& A) {
|
|
auto all_input_values = GetPatternInputValuesIncludeInner(A);
|
|
for (const auto& value : GetPatternOutputValuesIncludedInner(A)) {
|
|
all_input_values.erase(value);
|
|
}
|
|
VLOG(4) << "GetPatternInputValues: " << all_input_values.size();
|
|
return all_input_values;
|
|
}
|
|
|
|
static void PatternUpdateTracker(const StmtPattern& pattern) {
|
|
return std::visit([](const auto& impl) { impl.update_tracker(); }, pattern);
|
|
}
|
|
|
|
static std::string StmtPatternDebugStr(const StmtPattern& stmt) {
|
|
std::stringstream ss;
|
|
auto all_ops = GetOpsInPattern(stmt);
|
|
ss << "StmtPattern, size " << all_ops.size() << " :\n";
|
|
ss << OpsDebugStr(all_ops);
|
|
return ss.str();
|
|
}
|
|
|
|
} // namespace cinn::fusion
|