[Snippets] Transpose support

src/common/snippets/include/snippets/generator.hpp

@@ -84,15 +84,15 @@ class Schedule {
      * @param f can this kernel be linearided to 1D range
      * @param p pointer to generated code
      */
-    Schedule(const Shape& ws, bool f, code p) : work_size(ws), is_flat(f), ptr(p) {}
+    Schedule(const ov::PartialShape& ws, bool f, code p) : work_size(ws), is_flat(f), ptr(p) {}
     /**
      * @brief Returns callable instanse of code pointer
      */
     template<typename K> K get_callable() const {
         return reinterpret_cast<K>(const_cast<unsigned char*>(ptr));
     }
 
-    Shape work_size {};
+    ov::PartialShape work_size {};
     bool is_flat {false};
     code ptr {nullptr};
 };
@@ -123,7 +123,7 @@ class Generator {
      * @brief gets target machine
      * @return pointer to constant target machine
      */
-    std::shared_ptr<const TargetMachine> get_target_machine() const { return target; }
+    std::shared_ptr<const TargetMachine> get_target_machine() const;
 
 protected:
     std::shared_ptr<TargetMachine> target;

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

@@ -21,25 +21,14 @@ class BroadcastLoad : public BroadcastMove {
 public:
     OPENVINO_OP("BroadcastLoad", "SnippetsOpset", ngraph::snippets::op::BroadcastMove);
 
-    BroadcastLoad(const Output<Node>& x, Shape output_shape);
+    BroadcastLoad(const Output<Node>& x, ov::PartialShape output_shape);
     BroadcastLoad() = default;
 
     bool visit_attributes(AttributeVisitor& visitor) override;
 
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
 
     void validate_and_infer_types() override;
-
-    void set_broadcast_info(const Shape& bct) {
-        broadcast_info = bct;
-    }
-
-    bool is_broadcast(size_t idx) {
-        return broadcast_info[idx] == 1;
-    }
-
-private:
-    Shape broadcast_info;
 };
 
 } // namespace op

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

@@ -19,7 +19,7 @@ class BroadcastMove : public ngraph::op::Op {
 public:
     OPENVINO_OP("BroadcastMove", "SnippetsOpset");
 
-    BroadcastMove(const Output<Node>& x, Shape output_shape);
+    BroadcastMove(const Output<Node>& x, ov::PartialShape output_shape);
     BroadcastMove() = default;
 
     bool visit_attributes(AttributeVisitor& visitor) override;
@@ -28,12 +28,9 @@ class BroadcastMove : public ngraph::op::Op {
 
     void validate_and_infer_types() override;
 
-    OPENVINO_SUPPRESS_DEPRECATED_START
-    bool evaluate(const HostTensorVector& output_values, const HostTensorVector& input_values) const override;
-    OPENVINO_SUPPRESS_DEPRECATED_END
 
 protected:
-    Shape output_shape;
+    ov::PartialShape output_shape;
 };
 
 } // namespace op

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

@@ -20,13 +20,14 @@ class Kernel : public ngraph::op::Op {
 public:
     OPENVINO_OP("Kernel", "SnippetsOpset");
 
-    Kernel(const std::vector<std::pair<std::shared_ptr<ngraph::snippets::Emitter>, ngraph::snippets::RegInfo>>& region);
+    Kernel(std::vector<AllocatedEmitter> region, std::shared_ptr<const ov::Model> m);
     Kernel() = default;
 
-    std::vector<std::pair<std::shared_ptr<ngraph::snippets::Emitter>, ngraph::snippets::RegInfo>> region;
+    std::vector<AllocatedEmitter> region;
+    const std::shared_ptr<const ov::Model> model;
 
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& inputs) const override {
-        return std::make_shared<Kernel>(region);
+        return std::make_shared<Kernel>(region, model);
     }
     const void *compile_params = nullptr;
 };

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

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <ngraph/op/op.hpp>
+#include "snippets/op/memory_access.hpp"
 
 namespace ngraph {
 namespace snippets {
@@ -17,31 +18,28 @@ namespace op {
  *        Default value is "1" - to load one element
  * @ingroup snippets
  */
-class Load : public ngraph::op::Op {
+class Load : public MemoryAccess {
 public:
     OPENVINO_OP("Load", "SnippetsOpset");
 
     Load(const Output<Node>& x, const size_t count = 1lu);
     Load() = default;
 
-    size_t get_count() const { return m_count; }
-
-    void set_count(const size_t count) { m_count = count; }
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
+};
 
+// todo: LoadReshape is just Load (and mapped on LoadEmitter), it just allows to keep correct shape propagation
+//  when we decompose Transpose to Load and Store.
+class LoadReshape : public Load {
+public:
+    OPENVINO_OP("LoadReshape", "SnippetsOpset");
+    LoadReshape(const Output<Node>& x, size_t count = 1lu, std::vector<size_t> order = {});
     bool visit_attributes(AttributeVisitor& visitor) override;
-
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
-
     void validate_and_infer_types() override;
-
-    OPENVINO_SUPPRESS_DEPRECATED_START
-    bool evaluate(const HostTensorVector& output_values, const HostTensorVector& input_values) const override;
-    OPENVINO_SUPPRESS_DEPRECATED_END
-
-protected:
-    size_t m_count = 0lu;
+private:
+    std::vector<size_t> m_order;
 };
-
 } // namespace op
 } // namespace snippets
 } // namespace ngraph

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

@@ -0,0 +1,111 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "ngraph/op/op.hpp"
+#include "snippets/emitter.hpp"
+#include "ngraph/op/parameter.hpp"
+
+namespace ngraph {
+namespace snippets {
+namespace op {
+
+/**
+ * @interface LoopBase
+ * @brief Base class for LoopBegin and LoopEnd
+ * @ingroup snippets
+ */
+class LoopBase : public ngraph::op::Op {
+public:
+    OPENVINO_OP("LoopBase", "SnippetsOpset");
+    LoopBase(const std::vector<Output<Node>>& args, size_t work_amount, size_t increment);
+    LoopBase() = delete;
+    bool visit_attributes(AttributeVisitor& visitor) override;
+    size_t get_work_amount() const;
+    size_t get_increment() const;
+    bool get_evaluate_once() const;
+
+protected:
+    size_t work_amount;
+    size_t wa_increment;
+    bool evaluate_once; // true if the Loop is executed only once, used to skip setting and testing the loop counter
+};
+class LoopEnd;
+/**
+ * @interface LoopBegin
+ * @brief Marks the start of the Loop region.
+ *        Number of outputs always equals to the number of inputs (bypassed values) + 1 (edge to the corresponding LoopEnd)
+ * @param args - vector of input values, they are passed directly to output.
+ * @ingroup snippets
+ */
+class LoopBegin : public LoopBase {
+    friend LoopEnd;
+
+public:
+    OPENVINO_OP("LoopBegin", "SnippetsOpset");
+    explicit LoopBegin(const OutputVector& args);
+    LoopBegin() = delete;
+    void validate_and_infer_types() override;
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& inputs)  const override;
+    std::shared_ptr<LoopEnd> get_loop_end();
+    // begin_address and input_regs are needed to communicate information between LoopBegin and LoopEnd emitters
+    const uint8_t* begin_address;
+    std::vector<size_t> input_regs;
+
+private:
+    void validate_and_infer_types_except_LoopEnd();
+    LoopBegin(const std::vector<Output<Node>>& args, size_t work_amount, size_t wa_increment);
+};
+
+/**
+ * @interface LoopEnd
+ * @brief Marks the end of the Loop region and defines the loop properties.
+ *        Number of outputs always equals to the number of inputs (bypassed values) - 1 (edge to the corresponding LoopEnd)
+ * @param args vector of input values + LoopBegin, all values except for the LoopBegin are passed directly to output.
+ * @param work_amount total number of evaluations to be processed by the loop
+ * @param increment number of evaluations processed in one iteration of the loop.
+ * @param apply_increment describes which data pointers attributed to the loop should be incremented on every iteration.
+ * should be used when Loop is connected to Parameters and/or Results. If apply_increment[i] == true then i-th i/o data
+ * pointer will be incremented by work_amount*data_size on every iteration.
+ * @param ptr_increments specifies i/o pointer increment performed on every iteration. This is an alternative to
+ * apply_increments, which enables more flexibility.
+ * @param finalization_offsets pointer increments that are be applied to i/o pointers before exiting the loop
+ * @ingroup snippets
+ */
+class LoopEnd : public LoopBase {
+public:
+    OPENVINO_OP("LoopEnd", "SnippetsOpset");
+    LoopEnd(const std::vector<Output<Node>>& args, size_t work_amount, size_t wa_increment,
+              std::vector<bool> apply_increment, std::vector<int64_t> finalization_offsets);
+    LoopEnd(const std::vector<Output<Node>>& args, size_t work_amount, size_t wa_increment,
+            std::vector<int64_t> ptr_increments, std::vector<int64_t> finalization_offsets);
+    LoopEnd() = delete;
+    std::shared_ptr<LoopBegin> get_loop_begin();
+    void validate_and_infer_types() override;
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& inputs)  const override;
+    const std::vector<int64_t>& get_finalization_offsets() const;
+    const std::vector<int64_t>& get_ptr_increments() const;
+    void set_finalization_offsets(std::vector<int64_t> offsets);
+    void set_ptr_increments(std::vector<int64_t> new_ptr_increments);
+    // update_ptr_increments resets non-zero increments to the new_increments. It's used when wa_increment is
+    // updated and we need to refresh ptr increments accordingly while respecting the broadcasting pattern
+    void update_ptr_increments(int64_t new_increment);
+    void set_work_amount(size_t new_work_amount);
+    void set_increment(size_t new_increment);
+    void set_evaluate_once(bool once);
+    // Used to propagate information about Loop structure, needed to simplify some optimizations. For example,
+    // to skip pointer increments when outer Loop is empty, and work_amount == vector_size (one inner vector Loop)
+    // true by default, the optimizations enabled if it's false;
+    bool has_outer_loop;
+
+private:
+    std::vector<int64_t> ptr_increments;
+    std::vector<int64_t> finalization_offsets;
+    size_t loop_io_size;
+};
+
+} // namespace op
+} // namespace snippets
+} // namespace ngraph

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

@@ -0,0 +1,39 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <ngraph/op/op.hpp>
+
+namespace ngraph {
+namespace snippets {
+namespace op {
+
+/**
+ * @interface MemoryAccess
+ * @brief This is a base class for memory access operations (like Load and Store).
+ *        It provides universal set/get interface to manipulate the number
+ *        of elements accessed during one operation call ("count").
+ *        Default "count" value is "1" - it means to load/store one element
+ * @ingroup snippets
+ */
+
+class MemoryAccess : public ngraph::op::Op {
+public:
+    OPENVINO_OP("MemoryAccess", "SnippetsOpset");
+
+    size_t get_count() const;
+    void set_count(size_t count);
+    bool visit_attributes(AttributeVisitor& visitor) override;
+    void validate_and_infer_types() override;
+
+protected:
+    explicit MemoryAccess(const Output<Node>& x, size_t count = 1lu);
+    MemoryAccess() = default;
+    size_t m_count = 0lu;
+};
+
+} // namespace op
+} // namespace snippets
+} // namespace ngraph

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

@@ -34,6 +34,7 @@ class Scalar  : public ov::op::v0::Constant {
 
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
     void validate_and_infer_types() override;
+    bool visit_attributes(AttributeVisitor& visitor) override;
 };
 
 } // namespace op

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

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <ngraph/op/op.hpp>
+#include "snippets/op/memory_access.hpp"
 
 namespace ngraph {
 namespace snippets {
@@ -17,29 +18,14 @@ namespace op {
  *        Default value is "1" - to store one element
  * @ingroup snippets
  */
-class Store : public ngraph::op::Op {
+class Store : public MemoryAccess {
 public:
     OPENVINO_OP("Store", "SnippetsOpset");
 
     Store(const Output<Node>& x, const size_t count = 1lu);
     Store() = default;
 
-    size_t get_count() const { return m_count; }
-
-    void set_count(const size_t count) { m_count = count; }
-
-    bool visit_attributes(AttributeVisitor& visitor) override;
-
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
-
-    void validate_and_infer_types() override;
-
-    OPENVINO_SUPPRESS_DEPRECATED_START
-    bool evaluate(const HostTensorVector& output_values, const HostTensorVector& input_values) const override;
-    OPENVINO_SUPPRESS_DEPRECATED_END
-
-protected:
-    size_t m_count = 0lu;
 };
 
 } // namespace op