[Snippets] Added support of Port Descriptor

src/common/snippets/include/snippets/lowered/expression.hpp

@@ -4,58 +4,29 @@
 
 #pragma once
 
-#include <list>
-
 #include <openvino/core/node.hpp>
 #include <openvino/opsets/opset1.hpp>
 
-#include "snippets/tensor_descriptor.hpp"
 #include "snippets/emitter.hpp"
 #include "snippets/target_machine.hpp"
+#include "snippets/lowered/tensor.hpp"
+#include "snippets/lowered/expression_port.hpp"
 
 
 namespace ngraph {
 namespace snippets {
 namespace lowered {
 
 class LinearIR;
-class Expression;
-using ExpressionPtr = std::shared_ptr<Expression>;
-
-class ExpressionPort {
-    friend class Expression;
-
-public:
-    enum Type {
-        Input,
-        Output
-    };
-
-    ExpressionPort() = default;
-
-    Type get_type() const { return m_type; }
-
-    ExpressionPtr expr = nullptr;
-    size_t port = 0;
-
-private:
-    ExpressionPort(const ExpressionPtr& expr, size_t port, Type type);
-
-    Type m_type = Type::Input;
-};
 
 class Expression : public std::enable_shared_from_this<Expression> {
     friend class LinearIR;
+    friend class ExpressionPort;
 
 public:
     static size_t LOOP_NULL_ID;
 
     Expression() = default;
-    explicit Expression(const std::shared_ptr<Node>& n);
-    // The ctor fills outputs automatically from rt_info and/or tensor shapes
-    explicit Expression(const std::shared_ptr<Node>& n, std::vector<TensorDescriptorPtr> inputs);
-    explicit Expression(const std::shared_ptr<Node>& n, std::vector<TensorDescriptorPtr> inputs, std::vector<TensorDescriptorPtr> outputs);
-
     virtual ~Expression() = default;
 
     std::shared_ptr<Node> get_node() const;
@@ -64,53 +35,66 @@ class Expression : public std::enable_shared_from_this<Expression> {
     RegInfo get_reg_info() const { return  m_reg_info; }
     void set_reg_info(RegInfo rinfo) { m_reg_info = std::move(rinfo); }
 
-    const std::vector<TensorDescriptorPtr>& get_inputs() { return m_inputs; }
-    const std::vector<TensorDescriptorPtr>& get_outputs() { return m_outputs; }
+    const TensorPtr& get_input_tensor(size_t i) const;
+    const TensorPtr& get_output_tensor(size_t i) const;
+    std::vector<TensorPtr> get_input_tensors() const { return m_input_tensors; }
+    std::vector<TensorPtr> get_output_tensors() const { return m_output_tensors; }
+
+    const PortDescriptorPtr& get_input_port_descriptor(size_t i) const;
+    const PortDescriptorPtr& get_output_port_descriptor(size_t i) const;
+    std::vector<PortDescriptorPtr> get_input_port_descriptors() const { return m_input_port_descriptors; }
+    std::vector<PortDescriptorPtr> get_output_port_descriptors() const { return m_output_port_descriptors; }
+
+    size_t get_input_count() const { return m_input_tensors.size(); }
+    size_t get_output_count() const { return m_output_tensors.size(); }
 
     std::vector<size_t> get_loop_ids() const { return m_loop_ids; }
     void set_loop_ids(const std::vector<size_t>& loops) { m_loop_ids = loops; }
     void set_loop_id(size_t id, size_t idx);
     void remove_loop_id(size_t id);
-    bool is_outside_loop() const { return m_is_outside_loop; }
 
+    void validate() const;
     void init_emitter(const std::shared_ptr<const TargetMachine>& target);
 
-    ExpressionPort input_port(size_t i);
-    ExpressionPort output_port(size_t i);
+    ExpressionPort get_input_port(size_t i);
+    ExpressionPort get_output_port(size_t i);
 
 protected:
-    void replace_input(size_t port, TensorDescriptorPtr to);
-    void replace_output(size_t port, TensorDescriptorPtr to);
+    // Note: The constructor and tensor initialization are private since an expression can be created only by Linear IR.
+    //       These methods must be used only by Linear IR builder of expressions!
+    explicit Expression(const std::shared_ptr<Node>& n);
+
+    void replace_input(size_t port, TensorPtr to);
 
     std::shared_ptr<Node> m_source_node{nullptr};
     std::shared_ptr<Emitter> m_emitter{nullptr};
-    std::vector<TensorDescriptorPtr> m_inputs;
-    std::vector<TensorDescriptorPtr> m_outputs;
+    std::vector<TensorPtr> m_input_tensors{};
+    std::vector<TensorPtr> m_output_tensors{};
+    std::vector<PortDescriptorPtr> m_input_port_descriptors{};
+    std::vector<PortDescriptorPtr> m_output_port_descriptors{};
     RegInfo m_reg_info{{}, {}};
     // The order Loops identifies: Outer ---> Inner
     std::vector<size_t> m_loop_ids;
-    bool m_is_outside_loop = false;
 };
+using ExpressionPtr = std::shared_ptr<Expression>;
 
 class IOExpression : public Expression {
+    friend class LinearIR;
+
 public:
     enum class io_type {INPUT, OUTPUT, UNDEFINED};
 
-    IOExpression(const std::shared_ptr<ov::opset1::Parameter>& n, int64_t index);
-    IOExpression(const std::shared_ptr<ov::opset1::Result>& n, int64_t index, std::vector<TensorDescriptorPtr> inputs);
-
     int64_t get_index() const  { return m_index; }
     io_type get_type() const { return m_type; }
 
 private:
+    explicit IOExpression(const std::shared_ptr<ov::opset1::Parameter>& n, int64_t index);
+    explicit IOExpression(const std::shared_ptr<ov::opset1::Result>& n, int64_t index);
+
     int64_t m_index = -1;
     io_type m_type = io_type::UNDEFINED;
 };
 
-bool operator==(const ExpressionPort& lhs, const ExpressionPort& rhs);
-bool operator!=(const ExpressionPort& lhs, const ExpressionPort& rhs);
-bool operator<(const ExpressionPort& lhs, const ExpressionPort& rhs);
-
 } // namespace lowered
 } // namespace snippets
 } // namespace ngraph

src/common/snippets/include/snippets/lowered/expression_factory.hpp

@@ -0,0 +1,55 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "linear_ir.hpp"
+
+#include "snippets/snippets_isa.hpp"
+
+namespace ngraph {
+namespace snippets {
+namespace lowered {
+
+class LinearIR::ExpressionFactory {
+public:
+    template<class... Args>
+    static ExpressionPtr build(const std::shared_ptr<Node>& n, Args&&... params) {
+        if (const auto par = ov::as_type_ptr<ov::op::v0::Parameter>(n)) {
+            return create(par, params...);
+        } else if (const auto res = ov::as_type_ptr<ov::op::v0::Result>(n)) {
+            return create(res, params...);
+        } else if (const auto loop_begin = ov::as_type_ptr<op::LoopBegin>(n)) {
+            return create(loop_begin, params...);
+        } else if (const auto loop_end = ov::as_type_ptr<op::LoopEnd>(n)) {
+            return create(loop_end, params...);
+        }
+        return create(n, params...);
+    }
+
+private:
+    /* -- Default Builders - initialize input tensors from parents and create new output tensors themselves */
+    static ExpressionPtr create(const std::shared_ptr<ngraph::op::v0::Parameter>& par, const LinearIR& linear_ir,
+                                const std::shared_ptr<ov::Model>& model);
+    static ExpressionPtr create(const std::shared_ptr<ngraph::op::v0::Result>& res, const LinearIR& linear_ir,
+                                const std::shared_ptr<ov::Model>& model);
+    static ExpressionPtr create(const std::shared_ptr<ov::Node>& n, const LinearIR& linear_ir,
+                                const std::shared_ptr<ov::Model>& model);
+
+    /* -- Input Builders - get input tensors from method parameters and create new output tensors themselves */
+    static ExpressionPtr create(const std::shared_ptr<op::LoopBegin>& n, const std::vector<TensorPtr>& inputs);
+    static ExpressionPtr create(const std::shared_ptr<op::LoopEnd>& n, const std::vector<TensorPtr>& inputs);
+    static ExpressionPtr create(const std::shared_ptr<ov::Node>& n, const std::vector<TensorPtr>& inputs);
+
+    // Creates inputs for expression using parent output tensors
+    static void create_expression_inputs(const LinearIR& linear_ir, const ExpressionPtr& expr);
+    // Creates new output tensors
+    static void create_expression_outputs(const ExpressionPtr& expr);
+    // The method verifies of input tensors to availability of the expression as consumer and add it if missed
+    static void init_expression_inputs(const ExpressionPtr& expr, const std::vector<TensorPtr>& inputs);
+};
+
+} // namespace lowered
+} // namespace snippets
+} // namespace ngraph

src/common/snippets/include/snippets/lowered/expression_port.hpp

@@ -0,0 +1,51 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "port_descriptor.hpp"
+
+
+namespace ngraph {
+namespace snippets {
+namespace lowered {
+
+class Tensor;
+class Expression;
+class ExpressionPort {
+public:
+    enum Type {
+        Input,
+        Output
+    };
+
+    ExpressionPort() = default;
+    explicit ExpressionPort(const std::shared_ptr<Expression>& expr, Type type, size_t port);
+
+    const std::shared_ptr<Expression>& get_expr() const { return m_expr; }
+    Type get_type() const { return m_type; }
+    size_t get_index() const { return m_port_index; }
+
+    const PortDescriptorPtr& get_descriptor_ptr() const;
+    const std::shared_ptr<Tensor>& get_tensor_ptr() const;
+    // Returns connected ports to the current:
+    //  - Input port returns one source (parent) port
+    //  - Output port returns all consumer ports (children)
+    std::set<ExpressionPort> get_connected_ports() const;
+
+    friend bool operator==(const ExpressionPort& lhs, const ExpressionPort& rhs);
+    friend bool operator!=(const ExpressionPort& lhs, const ExpressionPort& rhs);
+    friend bool operator<(const ExpressionPort& lhs, const ExpressionPort& rhs);
+
+private:
+    std::shared_ptr<Expression> m_expr;
+    Type m_type = Type::Output;
+    size_t m_port_index = 0;
+};
+} // namespace lowered
+} // namespace snippets
+} // namespace ngraph

src/common/snippets/include/snippets/lowered/linear_ir.hpp

@@ -18,12 +18,12 @@ class Config {
     bool m_save_lowered_code = false;
     // True if we should check runtime info for nodes to call specific needed transformations
     bool m_need_fill_tail_register = false;
-    bool m_explicit_loop_insertion = false;
     ov::PartialShape m_master_shape{};
     size_t m_loop_depth = 1;
 };
 
 class LinearIR {
+    class ExpressionFactory;
 public:
     using container = std::list<ExpressionPtr>;
     using io_container = std::list<std::shared_ptr<IOExpression>>;
@@ -33,21 +33,18 @@ class LinearIR {
     LinearIR() = default;
     explicit LinearIR(const std::shared_ptr<ov::Model>& m, Config config = {});
 
-    LinearIR deep_copy() const;
+    ExpressionPtr create_expression(const std::shared_ptr<Node>& n, const std::vector<TensorPtr>& inputs);
+
     static LinearIR::container deep_copy_range(LinearIR::container::const_iterator begin, LinearIR::container::const_iterator end);
 
     const container& get_ops() const {return m_lowered_ops; }
     const io_container& get_IO_ops() const {return m_io_lowered_ops; }
     Config get_config() {return m_config; }
 
-    ExpressionPtr get_expr_by_node(const std::shared_ptr<Node>& n) const;
-    ExpressionPort get_expr_by_output(const TensorDescriptorPtr& n) const;
-    const std::set<ExpressionPort>& get_exprs_by_input(const TensorDescriptorPtr& n) const;
+    const ExpressionPtr& get_expr_by_node(const std::shared_ptr<Node>& n) const;
 
-    void replace_input(const ExpressionPort& expr_port, const TensorDescriptorPtr& to);
-    void replace_input(const ExpressionPtr& expr, size_t port, const TensorDescriptorPtr& to);
-    void replace_output(const ExpressionPort& expr_port, const TensorDescriptorPtr& to);
-    void replace_output(const ExpressionPtr& expr, size_t port, const TensorDescriptorPtr& to);
+    void replace_input(const std::set<ExpressionPort>& consumers, const TensorPtr& to);
+    void replace_input(const ExpressionPort& expr_port, const TensorPtr& to);
 
     /**
     * @brief Move an expression from the position "from" to the position immediately before "to".
@@ -88,26 +85,21 @@ class LinearIR {
     void init_emitters(const std::shared_ptr<TargetMachine>& target);
     void serialize(const std::string& xml, const std::string& bin);
 
-    static ov::NodeVector get_ordered_ops(const std::shared_ptr<ov::Model>& model);
-
     class LoopManager;
     using LoopManagerPtr = std::shared_ptr<LoopManager>;
 
     const LoopManagerPtr& get_loop_manager() const { return m_loop_manager; }
 
 private:
-    void register_expression(const ExpressionPtr& expr);
-    // Like register_expression, but doesn't allow Parameter or Result registration. You can do it only through ctor
-    void register_regular_expression(const ExpressionPtr& expr);
+    static ov::NodeVector get_ordered_ops(const std::shared_ptr<ov::Model>& model);
+    // Default ctor - can be called only from Linear IR initialization as default way
+    ExpressionPtr create_expression(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model = nullptr);
+
+    void register_expression(const ExpressionPtr& expr, bool io_allowed = false);
     void unregister_expression(const ExpressionPtr& expr);
 
     container m_lowered_ops{};
     std::unordered_map<std::shared_ptr<Node>, std::shared_ptr<Expression>> m_node2expression_map;
-    // Expression must be uniquely identified by an output, so there can't be expressions that have the same output
-    std::unordered_map<TensorDescriptorPtr, ExpressionPort> m_output2expression_map;
-    // At the same time, several expressions can have the same input if they are connected to the same parent
-    // E.g. LoopEnd will always have the same input as a Load inside the loop (since it has to increment the same reg)
-    std::unordered_map<TensorDescriptorPtr, std::set<ExpressionPort>> m_input2expression_map;
     io_container m_io_lowered_ops;
     Config m_config{};
     LoopManagerPtr m_loop_manager = nullptr;

src/common/snippets/include/snippets/lowered/loop_manager.hpp

@@ -9,7 +9,7 @@
 #include <openvino/core/node.hpp>
 #include <openvino/opsets/opset1.hpp>
 
-#include "snippets/tensor_descriptor.hpp"
+#include "port_descriptor.hpp"
 
 namespace ngraph {
 namespace snippets {
@@ -43,15 +43,10 @@ class LinearIR::LoopManager {
     size_t get_loop_count() const { return m_map.size(); }
     const std::map<size_t, LoopInfoPtr>& get_map() const;
 
-    static void skipped_mark(LinearIR::constExprIt loop_begin_pos,
-                             LinearIR::constExprIt loop_end_pos,
-                             size_t loop_depth);
-    void mark_loop(LinearIR& linear_ir,
-                   LinearIR::constExprIt loop_begin_pos,
+    void mark_loop(LinearIR::constExprIt loop_begin_pos,
                    LinearIR::constExprIt loop_end_pos,
                    size_t loop_depth, size_t vector_size);
-    void mark_loop(LinearIR& linear_ir,
-                   LinearIR::constExprIt loop_begin_pos,
+    void mark_loop(LinearIR::constExprIt loop_begin_pos,
                    LinearIR::constExprIt loop_end_pos,
                    size_t idx,
                    size_t work_amount,
@@ -74,8 +69,7 @@ class LinearIR::LoopManager {
     static void exprs_marking(LinearIR::constExprIt loop_begin_pos,
                               LinearIR::constExprIt loop_end_pos,
                               size_t loop_id, size_t idx);
-    static void get_io_loop_ports(LinearIR& linear_ir,
-                                  LinearIR::constExprIt loop_begin_pos,
+    static void get_io_loop_ports(LinearIR::constExprIt loop_begin_pos,
                                   LinearIR::constExprIt loop_end_pos,
                                   std::vector<ExpressionPort>& entries,
                                   std::vector<ExpressionPort>& exits);

src/common/snippets/include/snippets/lowered/pass/fuse_loops.hpp

@@ -27,15 +27,13 @@ class FuseLoops : public Transformation {
 private:
     static bool can_be_fused(const LinearIR::LoopManager::LoopInfoPtr& loop_current,
                              const LinearIR::LoopManager::LoopInfoPtr& loop_target);
-    static bool fuse_upper_into_current(LinearIR& linear_ir, const LinearIR::LoopManagerPtr& loop_manager,
-                                        const ExpressionPort& current_entry_point, const ExpressionPort& target_exit_point,
+    static bool fuse_upper_into_current(LinearIR& linear_ir, const LinearIR::LoopManagerPtr& loop_manager, const ExpressionPort& current_entry_point,
                                         size_t current_loop_id, size_t target_loop_id, size_t dim_idx,
                                         LinearIR::constExprIt& current_loop_begin_pos, LinearIR::constExprIt& current_loop_end_pos);
-    static bool fuse_lower_into_current(LinearIR& linear_ir, const LinearIR::LoopManagerPtr& loop_manager,
-                                        const ExpressionPort& current_entry_point, const ExpressionPort& target_exit_point,
+    static bool fuse_lower_into_current(LinearIR& linear_ir, const LinearIR::LoopManagerPtr& loop_manager, const ExpressionPort& current_entry_point,
                                         size_t current_loop_id, size_t target_loop_id, size_t dim_idx,
                                         LinearIR::constExprIt& current_loop_begin_pos, LinearIR::constExprIt& current_loop_end_pos);
-    static void fuse_points(LinearIR& linear_ir, std::vector<ExpressionPort>& exit_points, std::vector<ExpressionPort>& entry_points,
+    static void fuse_points(std::vector<ExpressionPort>& exit_points, std::vector<ExpressionPort>& entry_points,
                             LinearIR::constExprIt loop_begin_pos, LinearIR::constExprIt loop_end_pos);
 };
 

src/common/snippets/include/snippets/lowered/pass/insert_buffers.hpp

@@ -5,7 +5,6 @@
 #pragma once
 
 #include "transformation.hpp"
-#include "snippets/tensor_descriptor.hpp"
 
 namespace ngraph {
 namespace snippets {