Skip to content

Commit

Permalink
[Snippets] Made as static class
Browse files Browse the repository at this point in the history
  • Loading branch information
a-sidorova committed Aug 28, 2023
1 parent 8843f5d commit 8a6c3ef
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 78 deletions.
45 changes: 10 additions & 35 deletions src/common/snippets/include/snippets/lowered/buffer_manager.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,48 +19,23 @@ namespace lowered {
*/
class BufferManager {
public:
BufferManager(const lowered::LinearIR& linear_ir);
BufferManager() = default;

/**
* @brief allocate optimal memory size using MemorySolver
* @return size of common memory blob
*/
int64_t allocate();
static int64_t allocate(const lowered::LinearIR& linear_ir);

private:
using BufferSystem = std::vector<ExpressionPtr>;
using BufferCluster = std::set<ExpressionPtr>;
using BufferClusters = std::vector<BufferCluster>;

/**
* @brief init Buffers as graph edges and other subgraph around the Buffers as Nodes using enumeration
* Parameter
* |--- LoopBegin Parameter
* | LoadReshape <Edge> <- already allocated (since after Parameter). Skip
* | Store ---> Node <- (LoopBegin,...,LoopEnd)
* |--- LoopEnd <Edge> <- Buffer. Intermediate memory (edge)
* Buffer Node <- (LoopBegin,...,LoopEnd)
* |--- LoopBegin ...
* | ...
*/
void init_clusters(const lowered::LinearIR& linear_ir);
/**
* @brief init boxes for MemorySolver
*/
void init_boxes();
/**
* @brief Default initialization: makes non-inplace Buffers and sets different Buffer IDs, offsets
*/
void initialization();
/**
* @brief Set offset to Buffer op and propagate it to the connected memory access ops
*/
void set_buffer_offset(const ExpressionPtr& buffer_expr, const size_t offset) const;
static BufferSystem create_buffer_system(const lowered::LinearIR& linear_ir);
static size_t init_default(const BufferSystem& buffer_system);
static BufferClusters init_clusters(const BufferSystem& buffer_system);
static std::vector<MemorySolver::Box> init_boxes(const BufferClusters& buffer_clusters);
static void set_buffer_offset(const ExpressionPtr& buffer_expr, const size_t offset);

BufferClusters buffer_clusters = {};
std::vector<MemorySolver::Box> boxes = {};
size_t m_scratchpad_size = 0;
bool m_enable_optimizations = true;
constexpr static int64_t alignment = 32; // 32 bytes
constexpr static bool m_enable_optimizations = true;
constexpr static int64_t m_alignment = 32; // 32 bytes
};

} // namespace lowered
Expand Down
86 changes: 45 additions & 41 deletions src/common/snippets/src/lowered/buffer_manager.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -17,69 +17,72 @@ namespace ov {
namespace snippets {
namespace lowered {

BufferManager::BufferManager(const lowered::LinearIR& linear_ir) {
// Initialize edges and nodes
init_clusters(linear_ir);
}

int64_t BufferManager::allocate() {
initialization();
int64_t BufferManager::allocate(const lowered::LinearIR& linear_ir) {
OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::BufferManager::allocate")
const auto buffer_system = create_buffer_system(linear_ir);
auto scratchpad_size = init_default(buffer_system);

if (m_enable_optimizations) {
// Initialize boxes for MemorySolver
init_boxes();
const auto buffer_clusters = init_clusters(buffer_system);
const auto boxes = init_boxes(buffer_clusters);

MemorySolver staticMemSolver(boxes);
m_scratchpad_size = static_cast<size_t>(staticMemSolver.solve()) * alignment;
scratchpad_size = static_cast<size_t>(staticMemSolver.solve()) * m_alignment; // alignment in byte

// Set offsets for Buffers (edges)
// Set offsets for Buffers
for (auto& box : boxes) {
for (auto& buffer : buffer_clusters[box.id]) {
int64_t offset = staticMemSolver.getOffset(box.id);
set_buffer_offset(buffer, offset * alignment); // alignment in byte
set_buffer_offset(buffer, offset * m_alignment); // alignment in byte
}
}
}

return m_scratchpad_size;
return scratchpad_size;
}

void BufferManager::initialization() {
OV_ITT_SCOPED_TASK(ov::pass::itt::domains::SnippetsTransform, "Snippets::BufferManager::initialization")
BufferManager::BufferSystem BufferManager::create_buffer_system(const lowered::LinearIR& linear_ir) {
int64_t order = 0;
BufferSystem system;
for (const auto& expr : linear_ir) {
const auto op = expr->get_node();
if (const auto buffer = ov::as_type_ptr<op::Buffer>(op)) {
system.push_back(expr);
}
ov::snippets::pass::SetTopologicalOrder(op, order++);
}
return system;
}

size_t BufferManager::init_default(const BufferSystem& buffer_system) {
size_t buffer_id = 0;
size_t buffer_offset = 0;
for (const auto& buffer_expr : buffer_system) {
const auto node = buffer_expr->get_node();
const auto buffer = ov::as_type_ptr<ov::snippets::op::Buffer>(node);
if (!buffer)
continue;

for (const auto& cluster : buffer_clusters) {
for (const auto& buffer_expr : cluster) {
const auto node = buffer_expr->get_node();
const auto buffer = ov::as_type_ptr<ov::snippets::op::Buffer>(node);
if (!buffer)
continue;

const auto byte_size = buffer->get_byte_size();
set_buffer_offset(buffer_expr, buffer_offset);
buffer->set_id(buffer_id);
const auto byte_size = buffer->get_byte_size();
set_buffer_offset(buffer_expr, buffer_offset);
buffer->set_id(buffer_id);

buffer_offset += byte_size;
buffer_id++;
}
buffer_offset += byte_size;
buffer_id++;
}
m_scratchpad_size = buffer_offset;
return buffer_offset;
}

void BufferManager::init_clusters(const LinearIR& linear_ir) {
int64_t order = 0;
for (const auto& expr : linear_ir) {
const auto op = expr->get_node();
if (const auto buffer = ov::as_type_ptr<op::Buffer>(op)) {
buffer_clusters.push_back(BufferCluster{expr}); // TODO: Add support of inplace
}
ov::snippets::pass::SetTopologicalOrder(op, order++);
BufferManager::BufferClusters BufferManager::init_clusters(const BufferSystem& buffer_system) {
// TODO: Add support of inplace
BufferClusters buffer_clusters;
for (const auto& buffer_expr : buffer_system) {
buffer_clusters.push_back(BufferCluster{buffer_expr});
}
return buffer_clusters;
}

void BufferManager::init_boxes() {
std::vector<MemorySolver::Box> BufferManager::init_boxes(const BufferClusters& buffer_clusters) {
std::vector<MemorySolver::Box> boxes;
const auto count = static_cast<int>(buffer_clusters.size());
for (int i = 0; i < count; i++) {
MemorySolver::Box box = { std::numeric_limits<int>::max(), 0, 0, i };
Expand Down Expand Up @@ -127,12 +130,13 @@ void BufferManager::init_boxes() {
box.finish = std::max(e_finish, box.finish);
}

box.size = utils::div_up(box_size, alignment);
box.size = utils::div_up(box_size, m_alignment);
boxes.push_back(box);
}
return boxes;
}

void BufferManager::set_buffer_offset(const ExpressionPtr& buffer_expr, const size_t offset) const {
void BufferManager::set_buffer_offset(const ExpressionPtr& buffer_expr, const size_t offset) {
// If Buffer has offset We set this offset in the connected MemoryAccess ops
// to correctly read and write data because all Buffers have the common data pointer on buffer scratchpad

Expand Down
3 changes: 1 addition & 2 deletions src/common/snippets/src/op/subgraph.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -569,8 +569,7 @@ void snippets::op::Subgraph::control_flow_transformations(lowered::LinearIR& lin

target_pipeline.run(linear_ir);

lowered::BufferManager buffer_manager(linear_ir);
m_buffer_scratchpad = buffer_manager.allocate();
m_buffer_scratchpad = lowered::BufferManager::allocate(linear_ir);

lowered::pass::PassPipeline final_pipeline;
final_pipeline.register_pass<lowered::pass::PropagateLayout>();
Expand Down

0 comments on commit 8a6c3ef

Please sign in to comment.