[Snippets] Added Softmax support

[Snippets] Added support for Reshape around Softmax

applied comment part

Added config parameter to disable MHA ops tokenization

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

@@ -116,6 +116,8 @@ class Generator {
     /**
      * @brief virtual method any specific implementation should implement
      * @param m model in canonical for for table-based code generation
+     * @param config config with transformation and optimization parameters
+     * @param compile_params parameters for generated code
      * @return pointer to generated code
      */
     code generate(std::shared_ptr<ov::Model>& m, const SubgraphConfig& config, const void* compile_params = nullptr);

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

@@ -0,0 +1,43 @@
+// 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 Buffer
+ * @brief The operation is for intermediate data storage
+ *        Notes:
+ *               - All buffers in a graph have the same memory pointer. So if we have a few buffers,
+ *                 each buffer should have its own offset for common memory
+ *               - If Buffer is an input for operation output, this Buffer should be a single consumer for this port
+ * @ingroup snippets
+ */
+class Buffer : public ngraph::op::Op {
+public:
+    OPENVINO_OP("Buffer", "SnippetsOpset");
+    BWDCMP_RTTI_DECLARATION;
+
+    Buffer(const Output<Node>& x);
+    Buffer() = default;
+
+    size_t get_offset() const { return m_offset; }
+    void set_offset(const size_t offset);
+
+    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;
+
+private:
+    size_t m_offset = 0lu;
+};
+
+} // namespace op
+} // namespace snippets
+} // namespace ngraph

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

@@ -0,0 +1,47 @@
+// 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 Fill
+ * @brief Generated in Tail Loop vector representation in code generation step for cases when we should
+ *        refill regsiters by special numbers.
+ *        For example, for cases with ReduceMax or ReduceSum in Softmax
+ *        Where:
+ *          -  offset - is value shift for filling
+ *          - fill_value - hexadecimal filling value
+ * @ingroup snippets
+ */
+class Fill : public ngraph::op::Op {
+public:
+    OPENVINO_OP("Fill", "SnippetsOpset");
+
+    Fill(const Output<Node>& x, const size_t offset, const uint32_t fill_value = 0x0);
+    Fill() = default;
+
+    size_t get_offset() const { return m_offset; }
+    uint32_t get_fill_value() const { return m_fill_value; }
+
+    void set_offset(const size_t offset) { m_offset = offset; }
+    void set_fill_value(const uint32_t fill_value) { m_fill_value = fill_value; }
+
+    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;
+
+protected:
+    size_t m_offset = 0lu;
+    uint32_t m_fill_value = 0x0;
+};
+
+} // namespace op
+} // namespace snippets
+} // namespace ngraph

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

@@ -0,0 +1,32 @@
+// 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 HorizonMax
+ * @brief The operation calculates a horizon maximum of a vector register
+ * @ingroup snippets
+ */
+class HorizonMax : public ngraph::op::Op {
+public:
+    OPENVINO_OP("HorizonMax", "SnippetsOpset");
+
+    HorizonMax(const Output<Node>& x);
+    HorizonMax() = default;
+
+    bool visit_attributes(AttributeVisitor& visitor) override { return true;}
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
+    void validate_and_infer_types() override;
+};
+
+} // namespace op
+} // namespace snippets
+} // namespace ngraph

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

@@ -0,0 +1,32 @@
+// 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 HorizonSum
+ * @brief The operation calculates a horizon sum of a vector register
+ * @ingroup snippets
+ */
+class HorizonSum : public ngraph::op::Op {
+public:
+    OPENVINO_OP("HorizonSum", "SnippetsOpset");
+
+    HorizonSum(const Output<Node>& x);
+    HorizonSum() = default;
+
+    bool visit_attributes(AttributeVisitor& visitor) override { return true;}
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
+    void validate_and_infer_types() override;
+};
+
+} // namespace op
+} // namespace snippets
+} // namespace ngraph

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

@@ -13,16 +13,16 @@ namespace op {
 
 /**
  * @interface Load
- * @brief Generated by Canonicalization step where explicit instructions should be emitted for data loading
- *        where number of elements to load is determined by "count"
- *        Default value is "1" - to load one element
+ * @brief Generated during Lowering stage (convert_to_snippets_dialect) where explicit instructions should be emitted for data loading
+ *        where number of elements to load is determined by "count" (Default value is "1" - to load one element)
+ *        and memory offset for loading is determined by "offset" (Default value is "0" - to load starting from the first element)
  * @ingroup snippets
  */
 class Load : public MemoryAccess {
 public:
     OPENVINO_OP("Load", "SnippetsOpset");
 
-    Load(const Output<Node>& x, const size_t count = 1lu);
+    Load(const Output<Node>& x, const size_t count = 1lu, const size_t offset = 0lu);
     Load() = default;
 
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
@@ -33,7 +33,7 @@ class Load : public MemoryAccess {
 class LoadReshape : public Load {
 public:
     OPENVINO_OP("LoadReshape", "SnippetsOpset");
-    LoadReshape(const Output<Node>& x, size_t count = 1lu, std::vector<size_t> order = {});
+    LoadReshape(const Output<Node>& x, size_t count = 1lu, const size_t offset = 0lu, 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;

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

@@ -23,14 +23,17 @@ class MemoryAccess : public ngraph::op::Op {
     OPENVINO_OP("MemoryAccess", "SnippetsOpset");
 
     size_t get_count() const;
-    void set_count(size_t count);
+    size_t get_offset() const;
+    void set_count(const size_t count);
+    void set_offset(const size_t offset);
     bool visit_attributes(AttributeVisitor& visitor) override;
     void validate_and_infer_types() override;
 
 protected:
-    explicit MemoryAccess(const Output<Node>& x, size_t count = 1lu);
+    explicit MemoryAccess(const Output<Node>& x, size_t count = 1lu, size_t offset = 0lu);
     MemoryAccess() = default;
     size_t m_count = 0lu;
+    size_t m_offset = 0lu;
 };
 
 } // namespace op

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

@@ -13,16 +13,16 @@ namespace op {
 
 /**
  * @interface Store
- * @brief Generated by Canonicalization step where explicit instructions should be emitted for data storing
- *        where number of elements to store is determined by "count"
- *        Default value is "1" - to store one element
+ * @brief Generated during Lowering stage (convert_to_snippets_dialect) where explicit instructions should be emitted for data storing
+ *        where number of elements to store is determined by "count" (Default value is "1" - to store one element)
+ *        and memory offset for storing is determined by "offset" (Default value is "0" - to store starting at start memory ptr)
  * @ingroup snippets
  */
 class Store : public MemoryAccess {
 public:
     OPENVINO_OP("Store", "SnippetsOpset");
 
-    Store(const Output<Node>& x, const size_t count = 1lu);
+    Store(const Output<Node>& x, const size_t count = 1lu, const size_t offset = 0lu);
     Store() = default;
 
     std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;

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

@@ -90,21 +90,13 @@ class Subgraph : public ngraph::op::Op {
         return m_generator;
     }
 
-    size_t get_non_scalar_constants_count() const {
-        return m_non_scalar_constants_count;
-    }
-
-    bool is_quantized() const {
-        return config.m_is_quantized;
-    }
-
-    bool has_type_relaxed_ops() const {
-        return config.m_has_type_relaxed_ops;
-    }
-
-    bool has_domain_sensitive_ops() const {
-        return config.m_has_domain_sensitive_ops;
-    }
+    // Return common memory size for all buffers in body
+    size_t get_buffer_scratchpad_size() const;
+    size_t get_virtual_port_count() const { return m_virtual_port_count; }
+    bool is_buffer_needed() const { return m_buffer_needed; }
+    bool is_quantized() const { return config.m_is_quantized; }
+    bool has_type_relaxed_ops() const { return config.m_has_type_relaxed_ops; }
+    bool has_domain_sensitive_ops() const { return config.m_has_domain_sensitive_ops; }
 
     size_t tileRank = 0; // set by plugin to facilitate scheduling
 
@@ -121,7 +113,8 @@ class Subgraph : public ngraph::op::Op {
     // plugin sets generator for a snippet to some specific generator.
     // it's going to be replaced with Jitters table later
     void set_generator(std::shared_ptr<ngraph::snippets::Generator> generator);
-    void set_non_scalar_constants_count(const size_t count);
+    void set_virtual_port_count(const size_t count);
+    void buffer_needed(const bool need);
 
     void print() const;
     void print_statistics(bool verbose);
@@ -136,11 +129,14 @@ class Subgraph : public ngraph::op::Op {
     void align_element_types(const BlockedShapeVector& outputShapes, const BlockedShapeVector& inputShapes);
     void convert_to_snippet_dialect();
     void init_config();
-    // Count of potentional non-scalar Consants that will be created after some tranformations
-    // At the moment it's relevant only for FakeQuantize decomposition
-    // NOTE: To avoid overheads in each calcution of this count (for example, in validate_and_type_infer()),
+    // Count of Subgraph virtual ports:
+    //  - Potential non-scalar Constants that will be created after some transformations (At the moment it's relevant only for FakeQuantize decomposition)
+    // Need Buffer op or not
+    //  - Buffers. All Buffers are considered as one common additional virtual port. So we cannot summarize them as potential non-scalar Constants
+    // NOTE: To avoid overheads in each calculation of this count (for example, in validate_and_type_infer()),
     //       we should MANUALLY calculate it where it needed.
-    size_t m_non_scalar_constants_count = 0;
+    size_t m_virtual_port_count = 0;
+    bool m_buffer_needed = false;
     Shape exec_domain = {};
     std::shared_ptr<ov::Model> m_body = nullptr;
     std::shared_ptr<ngraph::snippets::Generator> m_generator = nullptr;

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

@@ -0,0 +1,34 @@
+// 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 VectorBuffer
+ * @brief The operation is for intermediate data storage in vector register
+ * @ingroup snippets
+ */
+class VectorBuffer : public ngraph::op::Op {
+public:
+    OPENVINO_OP("VectorBuffer", "SnippetsOpset");
+
+    VectorBuffer(const ov::element::Type element_type = ov::element::f32);
+
+    bool visit_attributes(AttributeVisitor& visitor) override { return true;}
+    std::shared_ptr<Node> clone_with_new_inputs(const OutputVector& new_args) const override;
+    void validate_and_infer_types() override;
+
+private:
+    ov::element::Type m_element_type;
+};
+
+} // namespace op
+} // namespace snippets
+} // namespace ngraph

src/common/snippets/include/snippets/pass/insert_buffer.hpp

@@ -0,0 +1,27 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <ngraph/pass/graph_rewrite.hpp>
+#include <ngraph/pattern/matcher.hpp>
+
+namespace ngraph {
+namespace snippets {
+namespace pass {
+
+/**
+ * @interface InsertBuffer
+ * @brief The pass inserts Buffers on Inputs and Outputs of special operations [Softmax, Transpose] is it's needed
+ * @ingroup snippets
+ */
+class InsertBuffer: public ngraph::pass::MatcherPass {
+public:
+    InsertBuffer();
+};
+
+
+}  // namespace pass
+}  // namespace snippets
+}  // namespace ngraph

src/common/snippets/include/snippets/pass/insert_load_store.hpp

@@ -13,7 +13,7 @@ namespace pass {
 
 /**
  * @interface InsertLoad
- * @brief Inserts explicit load instruction after each parameter.
+ * @brief Inserts explicit load instruction after each parameter and buffer.
  * The pass is used to convert model to a canonical form for code generation
  * @ingroup snippets
  */
@@ -24,7 +24,7 @@ class InsertLoad: public ngraph::pass::MatcherPass {
 
 /**
  * @interface InsertStore
- * @brief Inserts explicit store instruction before each result.
+ * @brief Inserts explicit store instruction before each result and buffer.
  * The pass is used to convert model to a canonical form for code generation
  * @ingroup snippets
  */

src/common/snippets/include/snippets/pass/insert_loops.hpp

@@ -19,13 +19,17 @@ namespace pass {
 class InsertLoops: public ngraph::pass::FunctionPass {
 public:
     OPENVINO_RTTI("InsertLoops", "0");
-    InsertLoops(ov::PartialShape master_shape, size_t loop_depth, size_t vector_size);
+    InsertLoops(ov::PartialShape master_shape, size_t loop_depth, size_t vector_size, bool is_optimized = true);
     bool run_on_model(const std::shared_ptr<ngraph::Function>& m) override;
 
+    static std::vector<bool> calculate_inner_apply_increments(const ov::PartialShape& master, const std::vector<ov::PartialShape>& shapes);
+    static std::vector<bool> calculate_outer_apply_increments(const std::vector<ov::PartialShape>& shapes);
+    static std::vector<int64_t> calculate_finalization_offsets(const ov::PartialShape& master, const std::vector<ov::PartialShape>& shapes);
 private:
     ov::PartialShape m_master_shape;
     size_t m_loop_depth;
     size_t m_vector_size;
+    bool m_is_optimized;
 };
 
 }  // namespace pass