[Snippets] Added support of Port Descriptor

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

@@ -4,45 +4,19 @@
 
 #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"
 
 
 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;
@@ -51,11 +25,6 @@ class Expression : public std::enable_shared_from_this<Expression> {
     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 +33,61 @@ 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& input(size_t i) const;
+    const TensorPtr& output(size_t i) const;
+    const std::vector<TensorPtr>& inputs() const { return m_inputs; }
+    const std::vector<TensorPtr>& outputs() const { return m_outputs; }
+    size_t get_input_count() const { return m_inputs.size(); }
+    size_t get_output_count() const { return m_outputs.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 init_emitter(const std::shared_ptr<const TargetMachine>& target);
 
-    ExpressionPort input_port(size_t i);
-    ExpressionPort output_port(size_t i);
+    TensorDescriptor input_port(size_t i);
+    TensorDescriptor 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 init_inputs(const std::vector<TensorPtr>& inputs) { m_inputs = inputs; }
+    void init_outputs(const std::vector<TensorPtr>& outputs) { m_outputs = outputs; }
+
+    // Note: These methods don't control availability of the current expression in this Tensor (as Consumer or Source)
+    void replace_input(size_t port, TensorPtr to);
+    void replace_output(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_inputs;
+    std::vector<TensorPtr> m_outputs;
     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,102 @@
+// Copyright (C) 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "linear_ir.hpp"
+
+namespace ngraph {
+namespace snippets {
+namespace lowered {
+
+class LinearIR::BaseExpressionFactory {
+public:
+    BaseExpressionFactory() = default;
+    BaseExpressionFactory(const LinearIR& linear_ir) : m_linear_ir(linear_ir) {}
+
+    virtual ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model);
+    virtual ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model,
+                                const std::vector<TensorPtr>& inputs);
+    virtual ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model,
+                                const std::vector<TensorPtr>& inputs, const std::vector<TensorPtr>& outputs);
+
+    static std::shared_ptr<LinearIR::BaseExpressionFactory> get(const LinearIR& linear_ir, const std::shared_ptr<Node>& n);
+
+protected:
+    virtual ExpressionPtr create(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model) = 0;
+    // Creates inputs for expression using parent output tensors
+    virtual std::vector<TensorPtr> create_expression_inputs(const ExpressionPtr& expr);
+    // Creates new output tensors
+    virtual std::vector<TensorPtr> create_expression_outputs(const ExpressionPtr& expr);
+    // The method verifies of input tensors to availability of the expression as consumer and add it if missed
+    virtual void validate_inputs(const ExpressionPtr& expr, const std::vector<TensorPtr>& inputs);
+
+    LinearIR m_linear_ir;
+};
+
+class LinearIR::ExpressionFactory : public LinearIR::BaseExpressionFactory {
+public:
+    ExpressionFactory() = default;
+    ExpressionFactory(const LinearIR& linear_ir) : BaseExpressionFactory(linear_ir) {}
+
+    ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model) override;
+    ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model,
+                        const std::vector<TensorPtr>& inputs) override;
+    ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model,
+                        const std::vector<TensorPtr>& inputs, const std::vector<TensorPtr>& outputs) override;
+
+protected:
+    ExpressionPtr create(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model) override;
+};
+
+class LinearIR::ParameterExpressionFactory : public LinearIR::BaseExpressionFactory {
+public:
+    ParameterExpressionFactory() = default;
+    ParameterExpressionFactory(const LinearIR& linear_ir) : BaseExpressionFactory(linear_ir) {}
+
+    ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model) override;
+
+protected:
+    ExpressionPtr create(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model) override;
+};
+
+class LinearIR::ResultExpressionFactory : public LinearIR::BaseExpressionFactory {
+public:
+    ResultExpressionFactory() = default;
+    ResultExpressionFactory(const LinearIR& linear_ir) : BaseExpressionFactory(linear_ir) {}
+
+    ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model) override;
+
+protected:
+    ExpressionPtr create(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model) override;
+};
+
+class LinearIR::LoopBeginExpressionFactory : public LinearIR::BaseExpressionFactory {
+public:
+    LoopBeginExpressionFactory() = default;
+    LoopBeginExpressionFactory(const LinearIR& linear_ir) : BaseExpressionFactory(linear_ir) {}
+
+    ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model,
+                        const std::vector<TensorPtr>& inputs) override;
+
+protected:
+    ExpressionPtr create(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model) override;
+};
+
+class LinearIR::LoopEndExpressionFactory : public LinearIR::BaseExpressionFactory {
+public:
+    LoopEndExpressionFactory() = default;
+    LoopEndExpressionFactory(const LinearIR& linear_ir) : BaseExpressionFactory(linear_ir) {}
+
+    ExpressionPtr build(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model,
+                        const std::vector<TensorPtr>& inputs) override;
+
+protected:
+    ExpressionPtr create(const std::shared_ptr<Node>& n, const std::shared_ptr<ov::Model>& model) override;
+    void validate_inputs(const ExpressionPtr& expr, const std::vector<TensorPtr>& inputs) override;
+};
+
+} // namespace lowered
+} // namespace snippets
+} // namespace ngraph

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

@@ -18,12 +18,17 @@ 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 BaseExpressionFactory;
+    class ExpressionFactory;
+    class ParameterExpressionFactory;
+    class ResultExpressionFactory;
+    class LoopBeginExpressionFactory;
+    class LoopEndExpressionFactory;
 public:
     using container = std::list<ExpressionPtr>;
     using io_container = std::list<std::shared_ptr<IOExpression>>;
@@ -33,6 +38,11 @@ class LinearIR {
     LinearIR() = default;
     explicit LinearIR(const std::shared_ptr<ov::Model>& m, Config config = {});
 
+    ExpressionPtr create_expression(const std::shared_ptr<Node>& n, const std::vector<TensorPtr> inputs,
+                                    const std::shared_ptr<ov::Model>& model = nullptr);
+    ExpressionPtr create_expression(const std::shared_ptr<Node>& n, const std::vector<TensorPtr> inputs, const std::vector<TensorPtr> outputs,
+                                    const std::shared_ptr<ov::Model>& model = nullptr);
+
     LinearIR deep_copy() const;
     static LinearIR::container deep_copy_range(LinearIR::container::const_iterator begin, LinearIR::container::const_iterator end);
 
@@ -41,13 +51,12 @@ class LinearIR {
     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;
 
-    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::vector<TensorDescriptor>& consumers, const TensorPtr& to);
+    void replace_input(const TensorDescriptor& expr_port, const TensorPtr& to);
+    void replace_input(const ExpressionPtr& expr, size_t port, const TensorPtr& to);
+    void replace_output(const TensorDescriptor& expr_port, const TensorPtr& to);
+    void replace_output(const ExpressionPtr& expr, size_t port, const TensorPtr& to);
 
     /**
     * @brief Move an expression from the position "from" to the position immediately before "to".
@@ -96,18 +105,16 @@ class LinearIR {
     const LoopManagerPtr& get_loop_manager() const { return m_loop_manager; }
 
 private:
+    // 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);
     // 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);
     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 "snippets/port_descriptor.hpp"
 
 namespace ngraph {
 namespace snippets {
@@ -23,17 +23,17 @@ class LinearIR::LoopManager {
     public:
         LoopInfo() = default;
         LoopInfo(size_t work_amount, size_t increment,
-                 const std::vector<ExpressionPort>& entries,
-                 const std::vector<ExpressionPort>& exits)
+                 const std::vector<TensorDescriptor>& entries,
+                 const std::vector<TensorDescriptor>& exits)
             : work_amount(work_amount), increment(increment), entry_exprs(entries), exit_exprs(exits) {}
         size_t work_amount = 0;
         size_t increment = 0;
         // The order of entry and exit expressions is important:
         //     - The position before first entry expr is Loop Begin position
         //     - The position after last exit expr is Loop End position
         // Note: Scalars aren't entry expressions but can be before first entry expr in Linear IR
-        std::vector<ExpressionPort> entry_exprs = {};
-        std::vector<ExpressionPort> exit_exprs = {};
+        std::vector<TensorDescriptor> entry_exprs = {};
+        std::vector<TensorDescriptor> exit_exprs = {};
     };
     using LoopInfoPtr = std::shared_ptr<LoopInfo>;
 
@@ -43,29 +43,24 @@ 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,
                    size_t work_amount_increment,
-                   const std::vector<ExpressionPort>& entries,
-                   const std::vector<ExpressionPort>& exits);
+                   const std::vector<TensorDescriptor>& entries,
+                   const std::vector<TensorDescriptor>& exits);
 
     void get_loop_bounds(const LinearIR& linear_ir,
                          size_t loop_id,
                          LinearIR::constExprIt& loop_begin_pos,
                          LinearIR::constExprIt& loop_end_pos) const;
     static void get_loop_bounds(const LinearIR& linear_ir,
-                                const std::vector<ExpressionPort>& entries,
-                                const std::vector<ExpressionPort>& exits,
+                                const std::vector<TensorDescriptor>& entries,
+                                const std::vector<TensorDescriptor>& exits,
                                 LinearIR::constExprIt& loop_begin_pos,
                                 LinearIR::constExprIt& loop_end_pos,
                                 size_t loop_id = Expression::LOOP_NULL_ID);
@@ -74,11 +69,10 @@ 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);
+                                  std::vector<TensorDescriptor>& entries,
+                                  std::vector<TensorDescriptor>& exits);
 
     std::map<size_t, LoopInfoPtr> m_map = {};
     size_t next_id = 0;