Buffer operation separation

Author: v-Golubev

src/common/snippets/include/snippets/op/buffer.hpp

@@ -5,6 +5,7 @@
 #pragma once
 
 #include "openvino/op/op.hpp"
+#include "snippets/shape_inference/shape_inference.hpp"
 
 namespace ov {
 namespace snippets {
@@ -13,13 +14,10 @@ namespace op {
 /**
  * @interface Buffer
  * @brief This is a base class for memory storage.
- *        If Buffer has a parent, the operation is for intermediate data storage - IntermediateMemory type.
- *        Otherwise, the operation is for allocation of new empty memory with shape `m_shape` - NewMemory type
  *        Notes:
  *               - All buffers with the same ID in a graph have the same memory pointer. So if we have a few buffers,
  *                 each the corresponding MemoryAccess op for Buffer should have offset for common memory pointer of this Buffer
  *               - Buffer should be a single consumer for operation output port
- * @param m_type - type of Buffer: IntermediateMemory/NewMemory
  * @param m_shape - output allocation shape for Buffer with type NewMemory
  * @param m_offset - offset in common Buffer scratchpad
  * @param m_id - Buffer ID in common Buffer system
@@ -29,39 +27,65 @@ class Buffer : public ov::op::Op {
 public:
     OPENVINO_OP("Buffer", "SnippetsOpset");
     Buffer() = default;
-    Buffer(const ov::Shape& shape, ov::element::Type element_type = ov::element::u8, size_t id = 0);
-    Buffer(const ov::Output<ov::Node>& arg, const ov::Shape& shape, size_t id = 0);
-    Buffer(const ov::Output<ov::Node>& arg, int32_t allocation_rank = -1, size_t id = 0);
+    Buffer(const OutputVector& arguments, const ov::Shape& shape, size_t id, ov::element::Type element_type = ov::element::u8);
 
     bool visit_attributes(AttributeVisitor& visitor) override;
-    void validate_and_infer_types() override;
-    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
-
-    enum Type {
-        NewMemory,
-        IntermediateMemory
-    };
 
     size_t get_id() const { return m_id; }
-    Type get_type() const { return m_type; }
     int64_t get_offset() const { return m_offset; }
     void set_id(size_t id) { m_id = id; }
     const ov::Shape& get_allocation_shape() const { return m_shape; }
     void set_allocation_shape(const ov::Shape& allocation_shape) { m_shape = allocation_shape; }
     void set_offset(int64_t offset) { m_offset = offset; }
     size_t get_byte_size() const;
 
-    void set_element_type(ov::element::Type element_type);
-
-    bool is_intermediate_memory() const { return m_type == Type::IntermediateMemory; }
-    bool is_new_memory() const { return m_type == Type::NewMemory; }
-
-private:
-    Type m_type = Type::IntermediateMemory;
+protected:
     ov::Shape m_shape = {};
-    int64_t m_offset = 0;
     size_t m_id = 0;  // Default ID - 0. All Buffers are from the same set
     ov::element::Type m_element_type = ov::element::u8;  // u8 - default 1 byte
+    int64_t m_offset = 0;
+};
+
+/**
+ * @interface IntermediateMemoryBuffer
+ * @brief Represents an intermediate memory storage operation. It always has a parent.
+ * @ingroup snippets
+ *
+ */
+class IntermediateMemoryBuffer : public Buffer {
+public:
+    OPENVINO_OP("IntermediateMemoryBuffer", "SnippetsOpset", Buffer);
+    IntermediateMemoryBuffer(const ov::Output<ov::Node>& arg, const ov::Shape& shape, size_t id = 0);
+    IntermediateMemoryBuffer(const ov::Output<ov::Node>& arg, int32_t allocation_rank = -1, size_t id = 0);
+
+    void validate_and_infer_types() override;
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
+
+private:
+    ov::Shape compute_shape_from_allocation_rank(const ov::Output<ov::Node>& arg, int32_t allocation_rank);
+};
+
+/**
+ * @interface NewMemoryBuffer
+ * @brief Represents a new empty memory for allocation with specified shape. It has no parent operations.
+ * @ingroup snippets
+ *
+ */
+class NewMemoryBuffer : public Buffer {
+public:
+    OPENVINO_OP("NewMemoryBuffer", "SnippetsOpset", Buffer);
+    NewMemoryBuffer(const ov::Shape& shape, size_t id = 0, ov::element::Type element_type = ov::element::u8);
+
+    void validate_and_infer_types() override;
+    void set_element_type(ov::element::Type element_type);
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
+
+    class ShapeInfer : public IShapeInferSnippets {
+        ov::Shape m_shape;
+    public:
+        explicit ShapeInfer(const std::shared_ptr<ov::Node>& n);
+        Result infer(const std::vector<VectorDimsRef>& input_shapes) override;
+    };
 };
 
 } // namespace op

src/common/snippets/src/generator.cpp

@@ -76,7 +76,8 @@ Generator::opRegType Generator::get_op_reg_type(const std::shared_ptr<Node>& op)
         std::dynamic_pointer_cast<op::LoopBegin>(op) ||
         std::dynamic_pointer_cast<op::LoopEnd>(op) ||
         std::dynamic_pointer_cast<op::Brgemm>(op) ||
-        std::dynamic_pointer_cast<op::Buffer>(op) ||
+        std::dynamic_pointer_cast<op::IntermediateMemoryBuffer>(op) ||
+        std::dynamic_pointer_cast<op::NewMemoryBuffer>(op) ||
         std::dynamic_pointer_cast<op::RankNormalization>(op))
         return gpr2gpr;
     else if (std::dynamic_pointer_cast<snippets::op::Load>(op) ||

src/common/snippets/src/lowered/pass/allocate_buffers.cpp

@@ -20,20 +20,18 @@ void AllocateBuffers::propagate_offset(const LinearIR& linear_ir, const Expressi
     buffer->set_offset(static_cast<int64_t>(offset));
 
     // Propagate to up: in Store. Buffer can have only one Store
-    {
-        if (buffer->is_intermediate_memory()) {
-            OPENVINO_ASSERT(buffer_expr->get_input_port_connectors().size() == 1, "Buffer with intermediate memory must have one parent");
-            const auto& parent_output = buffer_expr->get_input_port_connector(0)->get_source();
-            const auto& parent_expr = parent_output.get_expr();
-            const auto port = parent_output.get_index();
-            const auto& parent_node = parent_expr->get_node();
-            auto memory_access = ov::as_type_ptr<ov::snippets::op::MemoryAccess>(parent_node);
-            if (memory_access && memory_access->is_memory_access_output_port(port)) {
-                memory_access->set_output_offset(offset, port);
-            } else {
-                OPENVINO_THROW(
-                        "Buffer::set_offset() was called when Buffer didn't have the corresponding MemoryAccess op for offset propagation");
-            }
+    if (ov::is_type<op::IntermediateMemoryBuffer>(buffer)) {
+        OPENVINO_ASSERT(buffer_expr->get_input_port_connectors().size() == 1, "Buffer with intermediate memory must have one parent");
+        const auto& parent_output = buffer_expr->get_input_port_connector(0)->get_source();
+        const auto& parent_expr = parent_output.get_expr();
+        const auto port = parent_output.get_index();
+        const auto& parent_node = parent_expr->get_node();
+        auto memory_access = ov::as_type_ptr<ov::snippets::op::MemoryAccess>(parent_node);
+        if (memory_access && memory_access->is_memory_access_output_port(port)) {
+            memory_access->set_output_offset(offset, port);
+        } else {
+            OPENVINO_THROW(
+                    "Buffer::set_offset() was called when Buffer didn't have the corresponding MemoryAccess op for offset propagation");
         }
     }
     // Propagate to down: in Load. Buffer can have several Load
@@ -89,7 +87,7 @@ bool AllocateBuffers::run(LinearIR& linear_ir) {
                 continue;
             }
 
-            if (buffer->is_intermediate_memory()) {
+            if (ov::is_type<op::IntermediateMemoryBuffer>(buffer)) {
                 const auto& parent_expr = expr->get_input_port_connector(0)->get_source().get_expr();
                 const auto& parent_node = parent_expr->get_node();
                 // Full MemoryAccess ops need new memory. Previous logic is to check for parent isn't Loop
@@ -142,6 +140,7 @@ bool AllocateBuffers::run(LinearIR& linear_ir) {
                 allocated_buffers.insert(expr);
                 prev_data_size = current_data_size;
             } else {
+                OPENVINO_ASSERT(ov::is_type<op::NewMemoryBuffer>(buffer), "NewMemoryBuffer is expected");
                 if (!new_memory_buffer_allocated) {
                     allocate(buffer, *expr_it, buffer_size);
                     new_memory_buffer_allocated = true;

src/common/snippets/src/lowered/pass/assign_registers.cpp

@@ -64,7 +64,7 @@ bool AssignRegisters::run(LinearIR& linear_ir) {
         } else if (const auto& buffer = ov::as_type_ptr<op::Buffer>(op)) {
             const auto buffer_id = buffer->get_id();
             // All buffers have one common data pointer
-            if (buffer->is_intermediate_memory()) {
+            if (ov::is_type<op::IntermediateMemoryBuffer>(buffer)) {
                 manually_assigned_gprs[expr->get_input_port_connector(0)] =
                         static_cast<Reg>(num_results + num_parameters + buffer_id);
             }

src/common/snippets/src/lowered/pass/insert_buffers.cpp

@@ -155,7 +155,7 @@ void InsertBuffers::insertion(LinearIR& linear_ir, const LinearIR::constExprIt&
                                                                    parent_loops,
                                                                    parent_expr_output,
                                                                    m_buffer_allocation_rank);
-            const auto buffer = std::make_shared<op::Buffer>(parent->output(parent_port), allocation_shape);
+            const auto buffer = std::make_shared<op::IntermediateMemoryBuffer>(parent->output(parent_port), allocation_shape);
             PortDescriptorUtils::set_port_descriptor_ptr(buffer->output(0), parent_expr_output.get_descriptor_ptr()->clone());
             // Output connector is automatically filled from PortDescriptor
             const auto buffer_expr = linear_ir.create_expression(buffer, {input_connector});
@@ -248,7 +248,7 @@ void InsertBuffers::insertion(LinearIR& linear_ir, const LinearIR::constExprIt&
                                                                    current_loops,
                                                                    *exit_port,
                                                                    m_buffer_allocation_rank);
-            auto buffer = std::make_shared<op::Buffer>(node->output(port_idx), allocation_shape);
+            auto buffer = std::make_shared<op::IntermediateMemoryBuffer>(node->output(port_idx), allocation_shape);
             PortDescriptorUtils::set_port_descriptor_ptr(buffer->output(0), exit_port->get_descriptor_ptr()->clone());
             // We cannot insert Node output connector on Buffer output because not all consumers of Node needs Buffer
             //  Example:

src/common/snippets/src/lowered/pass/insert_load_store.cpp

@@ -122,9 +122,9 @@ bool InsertLoadStore::run(LinearIR& linear_ir) {
             modified |= insert_load(linear_ir, expr_it);
         } else if (ov::is_type<ov::op::v0::Result>(node)) {
             modified |= insert_store(linear_ir, expr_it);
-        } else if (auto buffer = ov::as_type_ptr<op::Buffer>(node)) {
+        } else if (ov::is_type<op::Buffer>(node)) {
             modified |= insert_load(linear_ir, expr_it);
-            if (buffer->is_intermediate_memory())
+            if (ov::is_type<op::IntermediateMemoryBuffer>(node))
                 modified |= insert_store(linear_ir, expr_it);
         }
     }

src/common/snippets/src/op/buffer.cpp

@@ -13,79 +13,96 @@ namespace ov {
 namespace snippets {
 namespace op {
 
+Buffer::Buffer(const OutputVector& arguments, const ov::Shape& shape, size_t id, ov::element::Type element_type)
+    : Op(arguments), m_shape(shape), m_id(id), m_element_type(std::move(element_type)), m_offset(0) {
+    constructor_validate_and_infer_types();
+}
+
+bool Buffer::visit_attributes(AttributeVisitor& visitor) {
+    INTERNAL_OP_SCOPE(Buffer_visit_attributes);
+    visitor.on_attribute("allocation_shape", m_shape);
+    visitor.on_attribute("offset", m_offset);
+    visitor.on_attribute("id", m_id);
+    visitor.on_attribute("element_type", m_element_type);
+    return true;
+}
 
-Buffer::Buffer(const ov::Shape& shape, ov::element::Type element_type, size_t id)
-    : Op(), m_type(Type::NewMemory), m_shape(shape), m_offset(0), m_id(id), m_element_type(std::move(element_type)) {
+size_t Buffer::get_byte_size() const {
+    const auto shape = get_allocation_shape();
+    return ov::shape_size(shape) * get_element_type().size();
+}
+
+IntermediateMemoryBuffer::IntermediateMemoryBuffer(const ov::Output<ov::Node>& arg, const ov::Shape& shape, size_t id)
+    : Buffer({arg}, shape, id) {
     constructor_validate_and_infer_types();
 }
 
-Buffer::Buffer(const ov::Output<ov::Node>& arg, const ov::Shape& shape, size_t id)
-    : Op({arg}), m_type(Type::IntermediateMemory), m_shape(shape), m_offset(0), m_id(id) {
+IntermediateMemoryBuffer::IntermediateMemoryBuffer(const ov::Output<ov::Node>& arg, int32_t allocation_rank, size_t id)
+    : Buffer({arg}, compute_shape_from_allocation_rank(arg, allocation_rank), id) {
     constructor_validate_and_infer_types();
 }
 
-Buffer::Buffer(const ov::Output<ov::Node>& arg, int32_t allocation_rank, size_t id)
-    : Op({arg}), m_type(Type::IntermediateMemory), m_offset(0), m_id(id) {
+ov::Shape IntermediateMemoryBuffer::compute_shape_from_allocation_rank(const ov::Output<ov::Node>& arg, int32_t allocation_rank) {
     const auto& pshape = arg.get_partial_shape();
     OPENVINO_ASSERT(pshape.is_static(), "Buffer supports only static input shape");
     const auto shape = pshape.get_shape();
     const auto normalize_rank = utils::normalize_rank(static_cast<int32_t>(allocation_rank), shape.size());
     const auto offset = static_cast<int32_t>(shape.size()) - normalize_rank;
-    m_shape = {shape.begin() + offset, shape.end()};
-    constructor_validate_and_infer_types();
-}
-
-bool Buffer::visit_attributes(AttributeVisitor& visitor) {
-    INTERNAL_OP_SCOPE(Buffer_visit_attributes);
-    visitor.on_attribute("allocation_shape", m_shape);
-    visitor.on_attribute("offset", m_offset);
-    visitor.on_attribute("id", m_id);
-    visitor.on_attribute("element_type", m_element_type);
-    return true;
+    return ov::Shape{shape.begin() + offset, shape.end()};
 }
 
-void Buffer::validate_and_infer_types() {
+void IntermediateMemoryBuffer::validate_and_infer_types() {
     INTERNAL_OP_SCOPE(Buffer_validate_and_infer_types);
     ov::PartialShape output_shape;
-    if (m_type == Type::NewMemory) {
-        OPENVINO_ASSERT(get_input_size() == 0, "Buffer with new allocated memory must to not have arguments!");
-        output_shape = m_shape;
-    } else if (m_type == Type::IntermediateMemory) {
-        m_element_type = get_input_element_type(0);
-        output_shape = get_input_partial_shape(0);
-    } else {
-        OPENVINO_THROW("Buffer supports only the following types: NewMemory and IntermediateMemory");
-    }
+    m_element_type = get_input_element_type(0);
+    output_shape = get_input_partial_shape(0);
     set_output_type(0, m_element_type, output_shape);
 }
 
-std::shared_ptr<Node> Buffer::clone_with_new_inputs(const OutputVector& new_args) const {
+std::shared_ptr<Node> IntermediateMemoryBuffer::clone_with_new_inputs(const OutputVector& new_args) const {
     INTERNAL_OP_SCOPE(Buffer_clone_with_new_inputs);
     check_new_args_count(this, new_args);
-    std::shared_ptr<op::Buffer> new_buffer = nullptr;
-    if (m_type == Type::NewMemory) {
-        new_buffer = std::make_shared<Buffer>(m_shape, m_element_type, m_id);
-    } else if (m_type == Type::IntermediateMemory) {
-        new_buffer = std::make_shared<Buffer>(new_args.at(0), m_shape, m_id);
-    } else {
-        OPENVINO_THROW("Buffer supports only the following types: NewMemory and IntermediateMemory");
-    }
-    new_buffer->m_offset = m_offset;
+    auto new_buffer = std::make_shared<IntermediateMemoryBuffer>(new_args.at(0), m_shape, m_id);
+    new_buffer->set_offset(m_offset);
     return new_buffer;
 }
 
-size_t Buffer::get_byte_size() const {
-    const auto shape = get_allocation_shape();
-    return ov::shape_size(shape) * get_element_type().size();
+NewMemoryBuffer::NewMemoryBuffer(const ov::Shape& shape, size_t id, ov::element::Type element_type)
+    : Buffer({}, shape, id) {
+    constructor_validate_and_infer_types();
 }
 
-void Buffer::set_element_type(ov::element::Type element_type) {
-    OPENVINO_ASSERT(is_new_memory(), "Only Buffer with NewMemory can change his output precision!");
+void NewMemoryBuffer::validate_and_infer_types() {
+    INTERNAL_OP_SCOPE(Buffer_validate_and_infer_types);
+    OPENVINO_ASSERT(get_input_size() == 0, "Buffer with new allocated memory mustn't have arguments!");
+    set_output_type(0, m_element_type, m_shape);
+}
+
+std::shared_ptr<Node> NewMemoryBuffer::clone_with_new_inputs(const OutputVector& new_args) const {
+    INTERNAL_OP_SCOPE(Buffer_clone_with_new_inputs);
+    check_new_args_count(this, new_args);
+    auto new_buffer = std::make_shared<NewMemoryBuffer>(m_shape, m_id, m_element_type);
+    new_buffer->set_offset(m_offset);
+    return new_buffer;
+}
+
+void NewMemoryBuffer::set_element_type(ov::element::Type element_type) {
     m_element_type = std::move(element_type);
     // Apply the change
     validate_and_infer_types();
 }
 
+NewMemoryBuffer::ShapeInfer::ShapeInfer(const std::shared_ptr<ov::Node>& n) {
+    const auto& buffer = ov::as_type_ptr<NewMemoryBuffer>(n);
+    OPENVINO_ASSERT(buffer, "Got invalid node in NewMemoryBuffer::ShapeInfer");
+    m_shape = buffer->get_shape();
+}
+
+IShapeInferSnippets::Result NewMemoryBuffer::ShapeInfer::infer(const std::vector<VectorDimsRef>& input_shapes) {
+    OPENVINO_ASSERT(input_shapes.size() == 1, "Got unexpected number of input shapes");
+    return {{m_shape}, ShapeInferStatus::success};
+}
+
 } // namespace op
 } // namespace snippets
 } // namespace ov

src/common/snippets/src/shape_inference/shape_inference.cpp

@@ -39,7 +39,7 @@ const IShapeInferSnippetsFactory::TRegistry IShapeInferSnippetsFactory::registry
         SHAPE_INFER_PREDEFINED(op::ConvertSaturation, PassThroughShapeInfer),
         SHAPE_INFER_PREDEFINED(op::Load, PassThroughShapeInfer),
         SHAPE_INFER_PREDEFINED(op::Store, PassThroughShapeInfer),
-        SHAPE_INFER_PREDEFINED(op::Buffer, PassThroughShapeInfer),
+        SHAPE_INFER_PREDEFINED(op::IntermediateMemoryBuffer, PassThroughShapeInfer),
         SHAPE_INFER_PREDEFINED(op::Fill, PassThroughShapeInfer),
         SHAPE_INFER_PREDEFINED(ov::op::v0::Parameter, PassThroughShapeInfer),
         // Note: We should remove Softmax shape infers after the decomposition activity,
@@ -66,6 +66,7 @@ const IShapeInferSnippetsFactory::TRegistry IShapeInferSnippetsFactory::registry
         SHAPE_INFER_OP_SPECIFIC(op::RankNormalization),
         SHAPE_INFER_OP_SPECIFIC(op::BroadcastLoad),
         SHAPE_INFER_OP_SPECIFIC(op::BroadcastMove),
+        SHAPE_INFER_OP_SPECIFIC(op::NewMemoryBuffer),
 };
 #undef SHAPE_INFER_OP_SPECIFIC_EXTERNAL
 #undef SHAPE_INFER_OP_SPECIFIC

src/common/snippets/tests/src/lowering_utils.cpp

@@ -42,7 +42,8 @@ DummyTargetMachine::DummyTargetMachine(const std::vector<ov::Node::type_info_t>&
     jitters[ov::snippets::op::LoopBegin::get_type_info_static()] = dummy_functor;
     jitters[ov::snippets::op::LoopEnd::get_type_info_static()] = dummy_functor;
     jitters[ov::snippets::op::Brgemm::get_type_info_static()] = dummy_functor;
-    jitters[ov::snippets::op::Buffer::get_type_info_static()] = dummy_functor;
+    jitters[ov::snippets::op::IntermediateMemoryBuffer::get_type_info_static()] = dummy_functor;
+    jitters[ov::snippets::op::NewMemoryBuffer::get_type_info_static()] = dummy_functor;
     jitters[ov::snippets::op::VectorBuffer::get_type_info_static()] = dummy_functor;
     jitters[ov::snippets::op::Fill::get_type_info_static()] = dummy_functor;