[Snippets] Fixed IdentifyBuffers for nested loop case (#21470)
This commit is contained in:
1 parent
bf00569ae1
commit
7e8a491276
3 files changed
+93
-65
No files matched your search
@@ -60,6 +60,7 @@ public:
|
||||
|
||||
private:
|
||||
using BufferPool = std::vector<ExpressionPtr>;
|
||||
using BufferMap = std::map<ExpressionPtr, ShiftPtrParams>;
|
||||
|
||||
/**
|
||||
* @brief Get Buffer Index in Buffer set
|
||||
@@ -105,6 +106,19 @@ private:
|
||||
*/
|
||||
static bool are_adjacent(const std::pair<ExpressionPtr, ShiftPtrParams>& lhs,
|
||||
const std::pair<ExpressionPtr, ShiftPtrParams>& rhs);
|
||||
|
||||
/**
|
||||
* @brief Find all buffers that are connected to the current LoopEnd
|
||||
* @param loop_end_expr expression of the target LoopEnd
|
||||
* @return buffer map [buffer expr -> ShiftDataPtrs]
|
||||
*/
|
||||
static BufferMap get_buffer_loop_neighbours(const ExpressionPtr& loop_end_expr);
|
||||
/**
|
||||
* @brief Find all buffers that are inside the current Loop.
|
||||
* @param loop_end_it expression iterator in LinearIR of the target LoopEnd
|
||||
* @return set of inner buffers
|
||||
*/
|
||||
static BufferMap get_buffer_loop_inside(const LinearIR::constExprIt& loop_end_it);
|
||||
};
|
||||
|
||||
} // namespace pass
|
||||
|
||||
@@ -65,10 +65,14 @@ void IdentifyBuffers::update_adj_matrix(const std::pair<ExpressionPtr, ShiftPtrP
|
||||
const std::pair<ExpressionPtr, ShiftPtrParams>& rhs,
|
||||
const BufferPool& buffers,
|
||||
std::vector<bool>& adj) {
|
||||
const auto size = buffers.size();
|
||||
const auto lhs_idx = get_buffer_idx(lhs.first, buffers);
|
||||
const auto rhs_idx = get_buffer_idx(rhs.first, buffers);
|
||||
// Already adjacent - skip
|
||||
if (adj[index(size, rhs_idx, lhs_idx)])
|
||||
return;
|
||||
|
||||
if (are_adjacent(lhs, rhs)) {
|
||||
const auto size = buffers.size();
|
||||
const auto lhs_idx = get_buffer_idx(lhs.first, buffers);
|
||||
const auto rhs_idx = get_buffer_idx(rhs.first, buffers);
|
||||
adj[index(size, rhs_idx, lhs_idx)] = adj[index(size, lhs_idx, rhs_idx)] = true;
|
||||
}
|
||||
}
|
||||
@@ -84,77 +88,87 @@ std::vector<bool> IdentifyBuffers::create_adjacency_matrix(const LinearIR& linea
|
||||
|
||||
for (auto expr_it = linear_ir.cbegin(); expr_it != linear_ir.cend(); expr_it++) {
|
||||
const auto &expr = *expr_it;
|
||||
const auto& loop_end = ov::as_type_ptr<op::LoopEnd>(expr->get_node());
|
||||
if (!loop_end)
|
||||
if (!ov::is_type<op::LoopEnd>(expr->get_node()))
|
||||
continue;
|
||||
|
||||
const auto input_count = loop_end->get_input_num();
|
||||
const auto output_count = loop_end->get_output_num();
|
||||
|
||||
const auto& ptr_increments = loop_end->get_ptr_increments();
|
||||
const auto& finalization_offsets = loop_end->get_finalization_offsets();
|
||||
const auto& data_sizes = loop_end->get_element_type_sizes();
|
||||
|
||||
// Buffer -> <ptr increment, finalization_offsets>
|
||||
std::map<ExpressionPtr, ShiftPtrParams> buffer_neighbours;
|
||||
|
||||
for (size_t i = 0; i < input_count; ++i) {
|
||||
const auto& parent_output = expr->get_input_port_connector(i)->get_source().get_expr();
|
||||
if (ov::is_type<op::Buffer>(parent_output->get_node())) {
|
||||
if (buffer_neighbours.count(parent_output) > 0) {
|
||||
OPENVINO_ASSERT(buffer_neighbours[parent_output].ptr_increment == ptr_increments[i] &&
|
||||
buffer_neighbours[parent_output].finalization_offset == finalization_offsets[i],
|
||||
"Invalid data pointer shifts: If Buffer has several consumers, this consumers must have the same shifts or zero");
|
||||
continue;
|
||||
}
|
||||
buffer_neighbours[parent_output] = { data_sizes[i], ptr_increments[i], finalization_offsets[i] };
|
||||
}
|
||||
}
|
||||
for (size_t i = input_count; i < input_count + output_count; ++i) {
|
||||
// The consumers of the corresponding Store ops
|
||||
const auto consumer_inputs = expr->get_input_port_connector(i)->get_consumers();
|
||||
size_t buffer_count = 0;
|
||||
size_t loop_count = 0;
|
||||
for (const auto& consumer_input : consumer_inputs) {
|
||||
const auto& child_expr = consumer_input.get_expr();
|
||||
if (ov::is_type<op::Buffer>(child_expr->get_node())) {
|
||||
buffer_neighbours[child_expr] = { data_sizes[i], ptr_increments[i], finalization_offsets[i] };
|
||||
buffer_count++;
|
||||
} else if (ov::is_type<op::LoopEnd>(child_expr->get_node())) {
|
||||
loop_count++;
|
||||
}
|
||||
}
|
||||
if (buffer_count > 0) {
|
||||
OPENVINO_ASSERT((buffer_count == 1) && (buffer_count + loop_count == consumer_inputs.size()),
|
||||
"Loop output must have not more than 1 Buffer");
|
||||
}
|
||||
}
|
||||
|
||||
// Buffers which are connected to the current Loop but without ptr shifts and Buffers which are inside this Loop - must be adjacent because
|
||||
// after each Loop iteration GPR will be shifted using ptr increment of Buffer outside. But Buffers inside have the same GPR - it means that
|
||||
// Buffers inside will work with shifted memory.
|
||||
const auto loop_begin = loop_end->get_loop_begin();
|
||||
for (auto it = std::reverse_iterator<LinearIR::constExprIt>(expr_it); (*it)->get_node() != loop_begin; ++it) {
|
||||
const auto& inner_expr = *it;
|
||||
if (ov::is_type<op::Buffer>(inner_expr->get_node())) {
|
||||
// To make Buffers adjacent, we set value "INT64_MAX" for data ptr shifts params for inner Buffers,
|
||||
// since outer Buffers (and other any Buffers) cannot have this value in shifting because of semantic of Loop op.
|
||||
// Thus, inner and outer Buffers have always different data shift ptr params -> they're adjacent
|
||||
if (buffer_neighbours.count(inner_expr) == 0)
|
||||
buffer_neighbours[inner_expr] = { INT64_MAX, INT64_MAX, INT64_MAX };
|
||||
}
|
||||
}
|
||||
|
||||
for (auto buffer_it = buffer_neighbours.begin(); buffer_it != buffer_neighbours.end(); ++buffer_it) {
|
||||
for (auto neighbour_it = std::next(buffer_it); neighbour_it != buffer_neighbours.end(); ++neighbour_it) {
|
||||
const auto buffer_loop_neighbours = get_buffer_loop_neighbours(expr);
|
||||
const auto buffers_loop_inside = get_buffer_loop_inside(expr_it);
|
||||
for (auto buffer_it = buffer_loop_neighbours.cbegin(); buffer_it != buffer_loop_neighbours.cend(); ++buffer_it) {
|
||||
// If Buffers, that are connected to the same Loop, have not proportionally ptr shift params for this Loop - these Buffers are adjacent
|
||||
for (auto neighbour_it = std::next(buffer_it); neighbour_it != buffer_loop_neighbours.cend(); ++neighbour_it) {
|
||||
update_adj_matrix(*buffer_it, *neighbour_it, pool, adj);
|
||||
}
|
||||
// Buffers which are connected to the current Loop with zero ptr shifts and Buffers which are inside this Loop - must be adjacent:
|
||||
// after each the Loop iteration GPR will be shifted using ptr increment of Buffer outside.
|
||||
// But if inner Buffers have the same GPR - it means that these Buffers will work with shifted memory.
|
||||
for (auto inner_it = buffers_loop_inside.cbegin(); inner_it != buffers_loop_inside.cend(); ++inner_it) {
|
||||
update_adj_matrix(*buffer_it, *inner_it, pool, adj);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return adj;
|
||||
}
|
||||
|
||||
IdentifyBuffers::BufferMap IdentifyBuffers::get_buffer_loop_neighbours(const ExpressionPtr& loop_end_expr) {
|
||||
const auto& loop_end = ov::as_type_ptr<op::LoopEnd>(loop_end_expr->get_node());
|
||||
const auto input_count = loop_end->get_input_num();
|
||||
const auto output_count = loop_end->get_output_num();
|
||||
|
||||
const auto& ptr_increments = loop_end->get_ptr_increments();
|
||||
const auto& finalization_offsets = loop_end->get_finalization_offsets();
|
||||
const auto& data_sizes = loop_end->get_element_type_sizes();
|
||||
|
||||
BufferMap buffer_neighbours;
|
||||
for (size_t i = 0; i < input_count; ++i) {
|
||||
const auto& parent_output = loop_end_expr->get_input_port_connector(i)->get_source().get_expr();
|
||||
if (ov::is_type<op::Buffer>(parent_output->get_node())) {
|
||||
if (buffer_neighbours.count(parent_output) > 0) {
|
||||
OPENVINO_ASSERT(buffer_neighbours[parent_output].ptr_increment == ptr_increments[i] &&
|
||||
buffer_neighbours[parent_output].finalization_offset == finalization_offsets[i],
|
||||
"Invalid data pointer shifts: If Buffer has several consumers, this consumers must have the same shifts or zero");
|
||||
continue;
|
||||
}
|
||||
buffer_neighbours[parent_output] = { data_sizes[i], ptr_increments[i], finalization_offsets[i] };
|
||||
}
|
||||
}
|
||||
for (size_t i = input_count; i < input_count + output_count; ++i) {
|
||||
// The consumers of the corresponding Store ops
|
||||
const auto consumer_inputs = loop_end_expr->get_input_port_connector(i)->get_consumers();
|
||||
size_t buffer_count = 0;
|
||||
size_t loop_count = 0;
|
||||
for (const auto& consumer_input : consumer_inputs) {
|
||||
const auto& child_expr = consumer_input.get_expr();
|
||||
if (ov::is_type<op::Buffer>(child_expr->get_node())) {
|
||||
buffer_neighbours[child_expr] = { data_sizes[i], ptr_increments[i], finalization_offsets[i] };
|
||||
buffer_count++;
|
||||
} else if (ov::is_type<op::LoopEnd>(child_expr->get_node())) {
|
||||
loop_count++;
|
||||
}
|
||||
}
|
||||
if (buffer_count > 0) {
|
||||
OPENVINO_ASSERT((buffer_count == 1) && (buffer_count + loop_count == consumer_inputs.size()),
|
||||
"Loop output must have not more than 1 Buffer");
|
||||
}
|
||||
}
|
||||
return buffer_neighbours;
|
||||
}
|
||||
|
||||
IdentifyBuffers::BufferMap IdentifyBuffers::get_buffer_loop_inside(const LinearIR::constExprIt& loop_end_it) {
|
||||
const auto& loop_end = ov::as_type_ptr<op::LoopEnd>((*loop_end_it)->get_node());
|
||||
const auto loop_begin = loop_end->get_loop_begin();
|
||||
BufferMap inner_buffers;
|
||||
for (auto it = std::reverse_iterator<LinearIR::constExprIt>(loop_end_it); (*it)->get_node() != loop_begin; ++it) {
|
||||
const auto& inner_expr = *it;
|
||||
if (ov::is_type<op::Buffer>(inner_expr->get_node())) {
|
||||
// Set default zero values since it's not used for adjacency definition in case with Buffers in Loop
|
||||
if (inner_buffers.count(inner_expr) == 0)
|
||||
inner_buffers[inner_expr] = { 0, 0, 0 };
|
||||
}
|
||||
}
|
||||
return inner_buffers;
|
||||
}
|
||||
|
||||
auto IdentifyBuffers::coloring(BufferPool& buffers, std::vector<bool>& adj) -> std::map<size_t, BufferPool> {
|
||||
size_t color = 0;
|
||||
std::map<size_t, BufferPool> color_groups;
|
||||
|
||||
+1
-1
@@ -195,7 +195,7 @@ INSTANTIATE_TEST_SUITE_P(smoke_Snippets_BufferAllocation_MHAOptimizedWSplit, MHA
|
||||
::testing::Values(true),
|
||||
::testing::Values(true),
|
||||
::testing::Values(90112),
|
||||
::testing::Values(5)),
|
||||
::testing::Values(3)),
|
||||
BufferAllocationCPUTest::getTestCaseName);
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(smoke_Snippets_BufferAllocation_MHANotOptimizedWOSplit, MHABF16AMXBufferAllocationTest,
|
||||
|
||||
Reference in new issue
Block a user