[Snippets] Convolution support #2

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

@@ -0,0 +1,24 @@
+// Copyright (C) ß2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <ngraph/ngraph.hpp>
+#include <ngraph/pass/graph_rewrite.hpp>
+#include <ngraph/pattern/matcher.hpp>
+
+
+namespace ngraph {
+namespace snippets {
+namespace pass {
+
+class Markup : public ngraph::pass::FunctionPass {
+public:
+    OPENVINO_RTTI("ConvolutionDecomposition", "0");
+    bool run_on_model(const std::shared_ptr<ngraph::Function>& m) override;
+};
+
+}  // namespace pass
+}  // namespace snippets
+}  // namespace ngraph

src/common/snippets/include/snippets/roi_backprop/convolution.hpp

@@ -0,0 +1,27 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <openvino/op/gather.hpp>
+#include <openvino/core/partial_shape.hpp>
+
+namespace ov {
+namespace op {
+namespace v1 {
+template <class ShapeType>
+void roi_backprop(
+        const Convolution* op,
+        const std::vector<ShapeType>& input_shapes,
+        std::vector<ShapeType>& roi_shapes,
+        std::vector<ov::Shape>& strides) {
+    NODE_VALIDATION_CHECK(op, (input_shapes.size() == 2ul) && (roi_shapes.size() == 1));
+
+    // TODO: just to test
+    //
+}
+
+}  // namespace v1
+}  // namespace op
+}  // namespace ov

src/common/snippets/include/snippets/roi_backprop/gather_roi_backprop.hpp

@@ -0,0 +1,39 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <openvino/op/gather.hpp>
+#include <openvino/core/partial_shape.hpp>
+
+namespace ov {
+namespace op {
+namespace v8 {
+template <class ShapeType>
+void roi_backprop(
+        const Gather* op,
+        const std::vector<ShapeType>& input_shapes,
+        std::vector<ShapeType>& roi_shapes,
+        std::vector<ov::Shape>& strides) {
+   NODE_VALIDATION_CHECK(op, input_shapes.size() == 3);
+
+    roi_shapes.resize(3ul);
+
+    auto& roi_data = roi_shapes[0];
+    auto& data_shape = input_shapes[0];
+    if (auto constant = ov::as_type_ptr<ov::opset1::Constant>(op->get_input_node_shared_ptr(2))) {
+        const auto axis_value = constant->cast_vector<int64_t>()[0];
+        roi_data = ov::Shape(data_shape.size(), 1);
+        roi_data[axis_value] = data_shape[axis_value];
+    } else {
+        roi_data = data_shape;
+    }
+
+    roi_shapes[1] = input_shapes[1];
+    roi_shapes[2] = input_shapes[2];
+}
+
+}  // namespace v8
+}  // namespace op
+}  // namespace ov

src/common/snippets/include/snippets/roi_backprop/max_pool.hpp

@@ -0,0 +1,44 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <openvino/op/gather.hpp>
+#include <openvino/core/partial_shape.hpp>
+
+namespace ov {
+namespace op {
+namespace v1 {
+template <class ShapeType>
+void roi_backprop(
+        const MaxPool* op,
+        const std::vector<ShapeType>& input_shapes,
+        std::vector<ShapeType>& roi_shapes,
+        std::vector<ov::Shape>& strides) {
+    NODE_VALIDATION_CHECK(op, (input_shapes.size() == 1ul) && (roi_shapes.size() == 1));
+    NODE_VALIDATION_CHECK(op, input_shapes[0].rank() == roi_shapes[0].rank());
+
+    // TODO: roi backpropagation: dynamic shape can be used
+    // TODO: roi backpropagation: params are ignored
+
+    const auto kernel = op->get_kernel();
+    auto roi_shape = roi_shapes[0];
+    auto shape_offset = input_shapes[0].size() - (kernel.size() + 2ul);
+    for (auto i = 2ul; i < roi_shape.size() - shape_offset; ++i) {
+        roi_shape[i] = roi_shape[i] * kernel[i - 2ul];
+    }
+    roi_shapes[0] = roi_shape;
+
+    const auto op_strides = op->get_strides();
+    auto strides0 = strides[0];
+    auto strides0_offset = input_shapes[0].size() - (op_strides.size() + 2ul);
+    for (auto i = 2ul; i < strides0.size() - strides0_offset; ++i) {
+        strides0[i] = strides0[i] * op_strides[i - 2ul];
+    }
+    strides[0] = strides0;
+}
+
+}  // namespace v1
+}  // namespace op
+}  // namespace ov

src/common/snippets/include/snippets/roi_backprop/roi_backprop.hpp

@@ -0,0 +1,45 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <map>
+
+#include <ngraph/runtime/host_tensor.hpp>
+#include <openvino/core/core.hpp>
+#include <openvino/core/node.hpp>
+
+namespace ov {
+namespace snippets {
+
+class ROIBackprop {
+public:
+    std::vector<ov::PartialShape> shapes;
+    std::vector<ov::Shape> strides;
+};
+
+void roi_backprop(ov::Node* op,
+                  const std::vector<ov::PartialShape>& input_shapes,
+                  const std::vector<ov::PartialShape>& cur_roi,
+                  const std::vector<ov::Shape>& cur_roi_strides,
+                  std::vector<ov::PartialShape>& new_roi,
+                  std::vector<ov::Shape>& new_roi_strides);
+
+using roi_map = std::map<ov::Node*, ROIBackprop>;
+roi_map get_roi_from_function(const std::shared_ptr<ov::Model>& m, const std::vector<ov::PartialShape>& start_roi);
+
+class BaseROIBackprop {
+public:
+    BaseROIBackprop(std::shared_ptr<ov::Node> node) : node(node) {}
+    virtual ROIBackprop infer_roi(
+            const std::vector<ov::PartialShape>& input_shapes,
+            const std::vector<ov::PartialShape>& cur_roi,
+            const std::vector<ov::Shape>& cur_strides) = 0;
+
+protected:
+    std::shared_ptr<ov::Node> node;
+};
+
+std::shared_ptr<BaseROIBackprop> make_roi_backprop(const std::shared_ptr<ngraph::Node>& op);
+
+}   // namespace snippets
+}   // namespace ov

src/common/snippets/include/snippets/roi_backprop/utils.hpp

@@ -0,0 +1,21 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+#pragma once
+
+#include <openvino/core/validation_util.hpp>
+#include <openvino/opsets/opset1.hpp>
+
+template <class OpType, class ShapeType>
+void transparent_roi_backprop(const OpType* op,
+                              const std::vector<ShapeType>& input_shapes,
+                              const std::vector<ShapeType>& cur_roi,
+                              const std::vector<ov::Shape>& cur_strides,
+                              std::vector<ShapeType>& new_roi,
+                              std::vector<ov::Shape>& new_strides) {
+    NODE_VALIDATION_CHECK(op, cur_roi.size() == 1, "Incorrect number of current roi shapes");
+    for (auto& roi_shape : new_roi)
+        roi_shape = cur_roi[0];
+
+    new_strides = cur_strides;
+}

src/common/snippets/src/generator.cpp

@@ -34,9 +34,13 @@ auto ngraph::snippets::getRegisters(std::shared_ptr<ngraph::Node>& n) -> ngraph:
         auto rt = input.get_source_output().get_node_shared_ptr()->get_rt_info();
         auto it_rt = rt.find("reginfo");
         if (it_rt != rt.end()) {
-            for (auto& reg : it_rt->second.as<std::vector<size_t>>()) {
-                rin.push_back(reg);
+            const auto& registers = it_rt->second.as<std::vector<size_t>>();
+            const auto output_index = input.get_source_output().get_index();
+            if (registers.size() <= output_index) {
+                throw ov::Exception("unexcepted registers count");
             }
+
+            rin.push_back(registers[output_index]);
         }
     }
     return std::make_pair(rin, rout);
@@ -77,10 +81,20 @@ ngraph::snippets::code ngraph::snippets::Generator::generate(std::shared_ptr<ov:
 
     // scalar tile
     auto m_scalar = ov::clone_model(*m.get());
+
+#ifdef CPU_DEBUG_CAPS
+    ov::pass::VisualizeTree("svg/snippets.generator.2.svg").run_on_model(m_scalar);
+#endif
+
     ngraph::pass::Manager mng;
     mng.register_pass<ngraph::snippets::pass::SetScalarCountForLoad>();
     mng.register_pass<ngraph::snippets::pass::SetScalarCountForStore>();
     mng.run_passes(m_scalar);
+
+#ifdef CPU_DEBUG_CAPS
+    ov::pass::VisualizeTree("svg/snippets.generator.3.svg").run_on_model(m);
+#endif
+
     OV_ITT_TASK_NEXT(GENERATE, "::ScalarTile_get")
     std::vector<AllocatedEmitter> scalar_lowered;
     for (auto n : m_scalar->get_ordered_ops()) {
@@ -90,8 +104,11 @@ ngraph::snippets::code ngraph::snippets::Generator::generate(std::shared_ptr<ov:
     // wrapping into tiles1D
     //todo: in, out, and io_last_dims should derive naturally from the graph representation
     const auto& vector_tile = std::make_shared<ngraph::snippets::op::Tile>(lowered, target->get_lanes(), in, out, io_last_dims, io_data_sizes);
-    const auto& vector_region = std::make_pair(target->get(ngraph::snippets::op::Tile::get_type_info_static())(vector_tile),
-                                   std::make_pair(std::vector<size_t>{}, std::vector<size_t>{}));
+    const auto& vector_region = std::make_pair(
+        target->get(ngraph::snippets::op::Tile::get_type_info_static())(vector_tile),
+        std::make_pair(
+            std::vector<size_t>{target->get_lanes()},
+            std::vector<size_t>{}));
     const auto& scalar_tile = std::make_shared<ngraph::snippets::op::Tile>(scalar_lowered, 1, in, out, io_last_dims, io_data_sizes);
     const auto& scalar_region = std::make_pair(target->get(ngraph::snippets::op::Tile::get_type_info_static())(scalar_tile),
                     std::make_pair(std::vector<size_t>{}, std::vector<size_t>{}));
@@ -115,9 +132,5 @@ ngraph::snippets::code ngraph::snippets::Generator::generate(std::shared_ptr<ov:
         op.first->emit_data();
     }
     OV_ITT_TASK_NEXT(GENERATE, "::GetSnippet")
-
-#ifdef CPU_DEBUG_CAPS
-    ov::pass::VisualizeTree("svg/snippets.generator.2.svg").run_on_model(m);
-#endif
     return target->get_snippet();
 }

src/common/snippets/src/op/subgraph.cpp

@@ -16,6 +16,8 @@
 #include "snippets/pass/vector_to_scalar.hpp"
 #include "snippets/pass/transform_convert.hpp"
 #include "snippets/pass/align_element_type.hpp"
+#include "snippets/pass/convolution_decomposition.hpp"
+#include "snippets/pass/markup.hpp"
 #include "snippets/utils.hpp"
 
 #include "transformations/common_optimizations/nop_elimination.hpp"
@@ -206,7 +208,17 @@ Shape snippets::op::Subgraph::canonicalize(const BlockedShapeVector& outputShape
     };
     Shape baseShape;
     AxisVector baseOrder;
-    std::tie(baseShape, baseOrder, std::ignore) = getMaxRankBlockedShape(inputShapes);
+
+    // TODO: improve later
+    if (inputShapes.size() == 1ull) {
+        std::tie(baseShape, baseOrder, std::ignore) = getMaxRankBlockedShape(inputShapes);
+    } else {
+        auto inputShapesWithoutSpecificConstant = (inputShapes.size() < 5ull) ? 
+            BlockedShapeVector{inputShapes[0], inputShapes[2]} : 
+            BlockedShapeVector{inputShapes[0], inputShapes[2], inputShapes[4]};
+        std::tie(baseShape, baseOrder, std::ignore) = getMaxRankBlockedShape(inputShapesWithoutSpecificConstant);
+    }
+
     const auto baseRank = baseShape.size();
     const bool baseIsBlocked = baseOrder.size() != std::set<size_t>(baseOrder.begin(), baseOrder.end()).size();
     for (size_t i = 0; i < inputShapes.size(); i++) {
@@ -215,6 +227,13 @@ Shape snippets::op::Subgraph::canonicalize(const BlockedShapeVector& outputShape
         AxisVector inOrder;
         element::Type inType;
         std::tie(inShape, inOrder, inType) = blockedShape;
+
+        // TODO: will be improved later
+        if ((i == 1ul) || (i == 3ul)) {
+            body_ptr()->replace_parameter(i, std::make_shared<opset1::Parameter>(inType, inShape));
+            continue;
+        }
+
         const auto inRank = inShape.size();
         NODE_VALIDATION_CHECK(this, inRank <= baseRank, "Input rank can't be larger than output rank in snippets.");
         if (inRank < baseRank) {
@@ -236,9 +255,10 @@ Shape snippets::op::Subgraph::canonicalize(const BlockedShapeVector& outputShape
                                   "Snippets canonicalization got input shapes of equal ranks but different layouts, which is not supported");
         }
         ov::PartialShape tmpPShape(baseShape);
-        NODE_VALIDATION_CHECK(this,
-                              PartialShape::broadcast_merge_into(tmpPShape, inShape, ::ngraph::op::AutoBroadcastType::NUMPY),
-                              "Failed to create broadcastable shapes in snippets canonicalization");
+        // TODO: for debug: Convolution
+        //NODE_VALIDATION_CHECK(this,
+        //                      PartialShape::broadcast_merge_into(tmpPShape, inShape, ::ngraph::op::AutoBroadcastType::NUMPY),
+        //                      "Failed to create broadcastable shapes in snippets canonicalization");
         const auto paramShape = body_ptr()->get_parameters()[i]->get_shape();
         const auto paramType =  body_ptr()->get_parameters()[i]->get_element_type();
         if (paramShape.size() != inShape.size() || !equal(paramShape.begin(), paramShape.end(), inShape.begin()))
@@ -339,6 +359,11 @@ void snippets::op::Subgraph::align_element_types(const BlockedShapeVector& outpu
 void snippets::op::Subgraph::convert_to_snippet_dialect() {
     INTERNAL_OP_SCOPE(Subgraph);
     OV_ITT_SCOPED_TASK(ngraph::pass::itt::domains::SnippetsTransform, "Snippets::convert_to_snippet_dialect")
+
+#ifdef CPU_DEBUG_CAPS
+    ov::pass::VisualizeTree("svg/snippets.convert_to_snippet_dialect.1.svg").run_on_model(body_ptr());
+#endif
+
     auto skip_matching_domain = [](const std::shared_ptr<const ov::Node>& n) -> bool {
         return n->get_input_shape(0).back() != 1;
     };
@@ -378,6 +403,27 @@ void snippets::op::Subgraph::convert_to_snippet_dialect() {
         set_callback<ngraph::snippets::pass::SetScalarCountForStore>(skip_matching_domain);
     }
     manager.run_passes(body_ptr());
+
+#ifdef CPU_DEBUG_CAPS
+    ov::pass::VisualizeTree("svg/snippets.convert_to_snippet_dialect.2.svg").run_on_model(body_ptr());
+    // ov::pass::Serialize("svg/snippets.convert_to_snippet_dialect.2.xml",
+    // "svg/snippets.convert_to_snippet_dialect.2.bin").run_on_model(m_body);
+#endif
+
+    // TODO: will be fixed later: ConvolutionDecomposition will be moved upper by execution flow
+    {
+        // TODO: use the the same manager
+        ngraph::pass::Manager manager;
+        manager.register_pass<snippets::pass::Markup>();
+        manager.register_pass<snippets::pass::ConvolutionDecomposition>();
+        manager.run_passes(body_ptr());
+    }
+
+#ifdef CPU_DEBUG_CAPS
+    ov::pass::VisualizeTree("svg/snippets.convert_to_snippet_dialect.3.svg").run_on_model(body_ptr());
+    // ov::pass::Serialize("svg/snippets.convert_to_snippet_dialect.3.xml",
+    // "svg/snippets.convert_to_snippet_dialect.3.bin").run_on_model(m_body);
+#endif
 }
 
 snippets::Schedule snippets::op::Subgraph::generate(const BlockedShapeVector& output_shapes,