Skip to content

Commit

Permalink
Create struct LoweredExprPort
Browse files Browse the repository at this point in the history
  • Loading branch information
a-sidorova committed Mar 23, 2023
1 parent 13ff3a8 commit c25c2e2
Show file tree
Hide file tree
Showing 14 changed files with 186 additions and 119 deletions.
53 changes: 39 additions & 14 deletions src/common/snippets/include/snippets/lowered_expr.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,30 @@ class IOLoweredExpr : public LoweredExpr {
};

using LoweredExprPtr = std::shared_ptr<LoweredExpr>;
using LoweredExprPort = std::pair<LoweredExprPtr, size_t>;

struct LoweredExprPort {
enum Type {
Input,
Output
};

LoweredExprPort() = default;

static LoweredExprPort make_input(const LoweredExprPtr& expr, size_t port);
static LoweredExprPort make_output(const LoweredExprPtr& expr, size_t port);

LoweredExprPtr m_expr = nullptr;
size_t m_port = 0;
Type m_type = Type::Input;

private:
LoweredExprPort(const LoweredExprPtr& expr, size_t port, Type type);
};

bool operator==(const LoweredExprPort& lhs, const LoweredExprPort& rhs);
bool operator!=(const LoweredExprPort& lhs, const LoweredExprPort& rhs);
bool operator<(const LoweredExprPort& lhs, const LoweredExprPort& rhs);

class LoweredExprIR {
public:
using container = std::list<LoweredExprPtr>;
Expand All @@ -101,8 +124,10 @@ class LoweredExprIR {
LoweredExprPtr get_expr_by_node(const std::shared_ptr<Node>& n) const;
LoweredExprPort get_expr_by_output(const TensorDescriptorPtr& n) const;
const std::set<LoweredExprPort>& get_exprs_by_input(const TensorDescriptorPtr& n) const;
void replace_input(const LoweredExprPort& expr_port, TensorDescriptorPtr to);
void replace_input(const LoweredExprPort& expr_port, const TensorDescriptorPtr& to);
void replace_input(const LoweredExprPtr& expr, size_t port, const TensorDescriptorPtr& to);
void replace_output(const LoweredExprPort& expr_port, const TensorDescriptorPtr& to);
void replace_output(const LoweredExprPtr& expr, size_t port, const TensorDescriptorPtr& to);
exprIt insert(constExprIt pos, const ov::NodeVector& nodes);
exprIt insert(constExprIt pos, const std::shared_ptr<Node>& n);
exprIt insert(constExprIt pos, container::value_type&& value);
Expand Down Expand Up @@ -168,20 +193,20 @@ class LoweredExprIR {
const std::map<size_t, LoweredLoopInfoPtr>& get_map() const;

static void skipped_mark(LoweredExprIR::constExprIt loop_begin_pos,
LoweredExprIR::constExprIt loop_end_pos,
size_t loop_depth);
LoweredExprIR::constExprIt loop_end_pos,
size_t loop_depth);
void mark_loop(LoweredExprIR& linear_ir,
LoweredExprIR::constExprIt loop_begin_pos,
LoweredExprIR::constExprIt loop_end_pos,
size_t loop_depth, size_t vector_size);
LoweredExprIR::constExprIt loop_begin_pos,
LoweredExprIR::constExprIt loop_end_pos,
size_t loop_depth, size_t vector_size);
void mark_loop(LoweredExprIR& linear_ir,
LoweredExprIR::constExprIt loop_begin_pos,
LoweredExprIR::constExprIt loop_end_pos,
size_t idx,
size_t work_amount,
size_t work_amount_increment,
const std::vector<LoweredExprPort>& entries,
const std::vector<LoweredExprPort>& exits);
LoweredExprIR::constExprIt loop_begin_pos,
LoweredExprIR::constExprIt loop_end_pos,
size_t idx,
size_t work_amount,
size_t work_amount_increment,
const std::vector<LoweredExprPort>& entries,
const std::vector<LoweredExprPort>& exits);

static void get_loop_bounds(const LoweredExprIR& linear_ir,
const std::vector<LoweredExprPort>& entries,
Expand Down
75 changes: 58 additions & 17 deletions src/common/snippets/src/lowered_expr.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,37 @@ IOLoweredExpr::IOLoweredExpr(const std::shared_ptr<ov::opset1::Result>& res, int
m_outputs = {};
}

LoweredExprPort::LoweredExprPort(const LoweredExprPtr& expr, size_t port, Type type) : m_expr(expr), m_port(port), m_type(type) {
if (type == Type::Input) {
OPENVINO_ASSERT(port < expr->get_inputs().size(), "The input port must be less than input count");
} else if (type == Type::Output) {
OPENVINO_ASSERT(port < expr->get_outputs().size(), "The output port must be less than output count");
}
}

LoweredExprPort LoweredExprPort::make_input(const LoweredExprPtr& expr, size_t port) {
return LoweredExprPort(expr, port, Type::Input);
}
LoweredExprPort LoweredExprPort::make_output(const LoweredExprPtr& expr, size_t port) {
return LoweredExprPort(expr, port, Type::Output);
}

bool operator==(const LoweredExprPort& lhs, const LoweredExprPort& rhs) {
if (&lhs == &rhs)
return true;
OPENVINO_ASSERT(lhs.m_type == rhs.m_type, "Incorrect comparison: Ports are from different types!");
return lhs.m_expr == rhs.m_expr && lhs.m_port == rhs.m_port;
}

bool operator!=(const LoweredExprPort& lhs, const LoweredExprPort& rhs) {
return !(lhs == rhs);
}

bool operator<(const LoweredExprPort& lhs, const LoweredExprPort& rhs) {
OPENVINO_ASSERT(lhs.m_type == rhs.m_type, "Incorrect comparison: Ports are from different types!");
return (lhs.m_expr < rhs.m_expr) || (lhs.m_expr == rhs.m_expr && lhs.m_port < rhs.m_port);
}

LoweredExprIR::LoweredExprIR(const std::shared_ptr<ov::Model>& model, LoweringConfig config)
: m_io_lowered_ops{}, m_config{std::move(config)}, m_loop_manager(std::make_shared<LoweredLoopManager>()) {
constExprIt scalar_pos = m_lowered_ops.begin();
Expand Down Expand Up @@ -253,9 +284,14 @@ const std::set<LoweredExprPort>& LoweredExprIR::get_exprs_by_input(const TensorD
return found->second;
}

void LoweredExprIR::replace_input(const LoweredExprPort& expr_port, TensorDescriptorPtr to) {
const auto& expr = expr_port.first;
const auto port = expr_port.second;
void LoweredExprIR::replace_input(const LoweredExprPtr& expr, size_t port, const TensorDescriptorPtr& to) {
replace_input(LoweredExprPort::make_input(expr, port), to);
}

void LoweredExprIR::replace_input(const LoweredExprPort& expr_port, const TensorDescriptorPtr& to) {
const auto& expr = expr_port.m_expr;
const auto port = expr_port.m_port;
OPENVINO_ASSERT(expr_port.m_type == LoweredExprPort::Type::Input, "Failed to replace: target input port must have Input type");
OPENVINO_ASSERT(port < expr->m_inputs.size(), "Failed to replace: target input port must be less than input count!");
const auto from = expr->m_inputs[port];
auto found = m_input2expression_map.find(from);
Expand All @@ -272,9 +308,14 @@ void LoweredExprIR::replace_input(const LoweredExprPort& expr_port, TensorDescri
expr->replace_input(port, std::move(to));
}

void LoweredExprIR::replace_output(const LoweredExprPtr& expr, size_t port, const TensorDescriptorPtr& to) {
replace_output(LoweredExprPort::make_output(expr, port), to);
}

void LoweredExprIR::replace_output(const LoweredExprPort& expr_port, const TensorDescriptorPtr& to) {
const auto& expr = expr_port.first;
const auto port = expr_port.second;
const auto& expr = expr_port.m_expr;
const auto port = expr_port.m_port;
OPENVINO_ASSERT(expr_port.m_type == LoweredExprPort::Type::Output, "Failed to replace: target output port must have Output type");
OPENVINO_ASSERT(port < expr->m_outputs.size(), "Failed to replace: target output port must be less than output count!");
const auto from = expr->m_outputs[port];
auto found = m_output2expression_map.find(from);
Expand All @@ -300,12 +341,12 @@ void LoweredExprIR::register_expression(const LoweredExprPtr& expr) {
}
for (size_t i = 0; i < expr->m_outputs.size(); ++i) {
const auto& out = expr->m_outputs[i];
m_output2expression_map[out] = { expr, i };
m_output2expression_map[out] = LoweredExprPort::make_output(expr, i);
}

for (size_t i = 0; i < expr->m_inputs.size(); ++i) {
const auto& in = expr->m_inputs[i];
const auto expr_port = LoweredExprPort{expr, i};
const auto expr_port = LoweredExprPort::make_input(expr, i);
const auto& res = m_input2expression_map.insert({in, std::set<LoweredExprPort>{expr_port}});
// If input is already in the map => add ExprPtr to the mapped set
if (!res.second) {
Expand All @@ -328,7 +369,7 @@ void LoweredExprIR::unregister_expression(const LoweredExprPtr& expr) {
if (users.size() == 1)
m_input2expression_map.erase(found);
else
users.erase({expr, in_port});
users.erase(LoweredExprPort::make_input(expr, in_port));
}
++in_port;
}
Expand Down Expand Up @@ -435,7 +476,7 @@ void LoweredExprIR::LoweredLoopManager::get_loop_bounds(const LoweredExprIR& lin
size_t loop_id) {
OPENVINO_ASSERT(!entries.empty(), "Loop must have entry points");
OPENVINO_ASSERT(!exits.empty(), "Loop must have entry points");
loop_begin_pos = std::find(linear_ir.begin(), linear_ir.end(), entries.front().first);
loop_begin_pos = std::find(linear_ir.begin(), linear_ir.end(), entries.front().m_expr);
OPENVINO_ASSERT(loop_begin_pos != linear_ir.end(), "Loop begin hasn't been found!");

// Some operations in Loop can be before first entry points: Scalars, VectorBuffer.
Expand All @@ -447,7 +488,7 @@ void LoweredExprIR::LoweredLoopManager::get_loop_bounds(const LoweredExprIR& lin
}

// At the moment all Loops must have exit points
loop_end_pos = std::next(std::find(loop_begin_pos, linear_ir.end(), exits.back().first));
loop_end_pos = std::next(std::find(loop_begin_pos, linear_ir.end(), exits.back().m_expr));
OPENVINO_ASSERT(loop_end_pos != linear_ir.end(), "Loop end hasn't been found!");
}

Expand All @@ -465,19 +506,19 @@ void LoweredExprIR::LoweredLoopManager::get_io_loop_ports(LoweredExprIR& linear_

for (size_t in_port = 0; in_port < inputs.size(); ++in_port) {
const auto in_td = inputs[in_port];
const auto parent_expr = linear_ir.get_expr_by_output(in_td).first;
const auto parent_expr = linear_ir.get_expr_by_output(in_td).m_expr;
if (!ov::is_type<opset1::Constant>(parent_expr->get_node()) &&
std::find(loop_begin_pos, expr_it, parent_expr) == expr_it) {
entries.push_back({expr, in_port});
entries.push_back(LoweredExprPort::make_input(expr, in_port));
}
}

for (size_t out_port = 0; out_port < outputs.size(); ++out_port) {
const auto out_td = outputs[out_port];
const auto consumer_exprs = linear_ir.get_exprs_by_input(out_td);
for (const auto& conumer_expr : consumer_exprs) {
if (std::find(expr_it, loop_end_pos, conumer_expr.first) == loop_end_pos) {
exits.push_back({expr, out_port});
if (std::find(expr_it, loop_end_pos, conumer_expr.m_expr) == loop_end_pos) {
exits.push_back(LoweredExprPort::make_output(expr, out_port));
break;
}
}
Expand Down Expand Up @@ -522,8 +563,8 @@ void LoweredExprIR::LoweredLoopManager::mark_loop(LoweredExprIR& linear_ir,
std::vector<size_t> loop_layout;
std::vector<size_t> loop_tensor(1, 1); // Scalar
for (const auto& exit_point : loop_exit_points) {
const auto expr = exit_point.first;
const auto port = exit_point.second;
const auto expr = exit_point.m_expr;
const auto port = exit_point.m_port;
const auto out_td = expr->get_outputs()[port];
const auto out_tensor = out_td->get_tensor();
const auto out_layout = out_td->get_layout();
Expand All @@ -534,7 +575,7 @@ void LoweredExprIR::LoweredLoopManager::mark_loop(LoweredExprIR& linear_ir,
}

for (const auto& entry_point : loop_entry_points) {
const auto expr = entry_point.first;
const auto expr = entry_point.m_expr;
const auto out_td = expr->get_outputs().front();
const auto out_subtensor = out_td->get_subtensor();
if (loop_subtensor.empty())
Expand Down
10 changes: 5 additions & 5 deletions src/common/snippets/src/pass/lowered/assign_registers.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -82,19 +82,19 @@ bool AssignRegisters::run(LoweredExprIR& linear_ir) {
// We should manually set the one vector register for VectorBuffer and Max/Sum output to simulate a accumulator
// TODO [96351]: We should rewrite accumulator pattern using another way
const auto input_td = expr->get_inputs()[0];
const auto& input_expr = linear_ir.get_expr_by_output(input_td).first;
const auto& input_expr = linear_ir.get_expr_by_output(input_td).m_expr;
const auto& input_expr_input_tds = input_expr->get_inputs();
for (const auto& td : input_expr_input_tds) {
if (ov::is_type<op::VectorBuffer>(linear_ir.get_expr_by_output(td).first->get_node())) {
if (ov::is_type<op::VectorBuffer>(linear_ir.get_expr_by_output(td).m_expr->get_node())) {
manually_assigned_vecs[td] = static_cast<Reg>(accumulator_reg);
}
}
const auto output_td = expr->get_outputs()[0];
manually_assigned_vecs[input_td] = static_cast<Reg>(accumulator_reg);
manually_assigned_vecs[output_td] = static_cast<Reg>(accumulator_reg);
for (const auto& child_expr_input : linear_ir.get_exprs_by_input(output_td)) {
if (ov::is_type<op::BroadcastMove>(child_expr_input.first->get_node())) {
manually_assigned_vecs[child_expr_input.first->get_outputs()[0]] =
if (ov::is_type<op::BroadcastMove>(child_expr_input.m_expr->get_node())) {
manually_assigned_vecs[child_expr_input.m_expr->get_outputs()[0]] =
static_cast<Reg>(accumulator_reg);
}
}
Expand Down Expand Up @@ -198,7 +198,7 @@ bool AssignRegisters::run(LoweredExprIR& linear_ir) {
continue;
for (const auto& out : expr->get_outputs()) {
for (const auto& child_expr_input : linear_ir.get_exprs_by_input(out)) {
const auto& child_expr = child_expr_input.first;
const auto& child_expr = child_expr_input.m_expr;
auto child_it = linear_ir.begin();
std::advance(child_it, n);
size_t k = n;
Expand Down
Loading

0 comments on commit c25c2e2

Please sign in to comment.