[Transformations] Enable dynamic decomposition BTS and STB ops (openvinotoolkit#15179)

* Add dynamism for BatchToSpace conversion

* Extend dynamism for BatchToSpace conversion

Only block input needs to be const 'cse axes_order_const is freaky

* Enhace dynamism for BatchToSpace conversion

Block input need not be const now.

* Add dynamism for STB by elements conversion

* Remove const need for crops for BTS by_elements

* temp for review

* Try to fix output tensor overwrite

* Make test to reproduce invalid shape inference

* Reproduce the error with template plugin

* Fix code style

* Fix 0D inputs issue

* Remove 0D shape parts before Concat

* Apply nested namespaces

* Enable non-constant STB Block input

* Fix BTS runtime info

* Fix STB by elems runtime info

* Add dynamism for STB conversion

* Add BTS dynamic data test

* Add STB dynamic data test

* Reduce STB concats

* Add tests naming

* Edit

* style

* Consider other block element types

* Enhance type test

* Use opset10 only

* Check block shape

Author: t-jankowski

src/common/transformations/src/transformations/op_conversions/convert_batch_to_space.cpp

@@ -4,224 +4,216 @@
 
 #include "transformations/op_conversions/convert_batch_to_space.hpp"
 
+#include <algorithm>
+#include <climits>
 #include <memory>
 #include <ngraph/pattern/op/wrap_type.hpp>
 #include <ngraph/rt_info.hpp>
-#include <openvino/opsets/opset3.hpp>
+#include <openvino/opsets/opset10.hpp>
 #include <vector>
 
 #include "itt.hpp"
 
+using namespace std;
+using namespace ov::opset10;
+using namespace ov::element;
+
 void ov::pass::ConvertBatchToSpace::convert_batch_to_space() {
     MATCHER_SCOPE(ConvertBatchToSpace_convert_batch_to_space);
-    auto batch_to_space = ngraph::pattern::wrap_type<ov::opset3::BatchToSpace>();
-    matcher_pass_callback callback = [](pattern::Matcher& m) {
-        auto batch_to_space = std::dynamic_pointer_cast<ov::opset3::BatchToSpace>(m.get_match_root());
-        if (!batch_to_space) {
+    const auto batch_to_space = pattern::wrap_type<BatchToSpace>();
+    matcher_pass_callback callback = [this](pattern::Matcher& m) {
+        const auto batch_to_space = dynamic_pointer_cast<BatchToSpace>(m.get_match_root());
+        if (!batch_to_space || transformation_callback(batch_to_space)) {
             return false;
         }
 
-        NodeVector new_ops;
-        auto data = batch_to_space->input_value(0);
-        auto block = batch_to_space->input_value(1);
-        auto crops_begin = batch_to_space->input_value(2);
-        auto crops_end = batch_to_space->input_value(3);
+        NodeRegistry rg;
+        const auto data = batch_to_space->input_value(0);
+        const auto block = batch_to_space->input_value(1);
+        const auto crops_begin = batch_to_space->input_value(2);
+        const auto crops_end = batch_to_space->input_value(3);
 
-        if (data.get_partial_shape().is_dynamic()) {
-            return false;
+        const auto data_shape_rank = data.get_partial_shape().rank();
+        if (data_shape_rank.is_dynamic()) {
+            return false;  // because StridedSlice masks are std::vector
         }
-        const auto& data_shape = data.get_shape();
 
-        const auto block_const = std::dynamic_pointer_cast<opset3::Constant>(block.get_node_shared_ptr());
-        const auto crops_begin_const = std::dynamic_pointer_cast<opset3::Constant>(crops_begin.get_node_shared_ptr());
-        const auto crops_end_const = std::dynamic_pointer_cast<opset3::Constant>(crops_end.get_node_shared_ptr());
-
-        if (!block_const || !crops_begin_const || !crops_end_const) {
+        if (block.get_partial_shape().is_dynamic() || block.get_shape().size() == 0) {
             return false;
         }
-
-        const std::vector<int64_t>& block_values = block_const->cast_vector<int64_t>();
-        const std::vector<int64_t>& crops_end_values = crops_end_const->cast_vector<int64_t>();
+        const auto block_length = static_cast<int64_t>(block.get_shape()[0]);
 
         // First we have to disperse the data from batch, then rearrange them
         // so as appropriate chunks of data where close to their destination place.
-        // Finally squeeze data from respective dimensions.ss
-        std::vector<int64_t> dispersed_shape;
-        int64_t b_dim_divider = 1;
-        for (const auto& el : block_values) {
-            b_dim_divider *= el;
-        }
+        // Finally squeeze data from respective dimensions
+
+        const auto zero = rg.make<Constant>(i64, Shape{1}, 0);
+        const auto shape_of_data = rg.make<ShapeOf>(data, block.get_element_type());
+        const auto batch = rg.make<Gather>(shape_of_data, zero, zero);
+        const auto block_prod = rg.make<ReduceProd>(block, zero);
+        const auto batch_div = rg.make<Divide>(batch, block_prod);
 
         //   note: B_0 is expected to be 1.
         //      x' = reshape(`data`, [B_1, ..., B_{N - 1}, batch / (B_1 * ... B_{N - 1}), D_1, D_2, ...,
         //      D_{N - 1}]),
         //      where B_i = block_shape[i]
-        dispersed_shape.insert(dispersed_shape.begin(), block_values.begin() + 1, block_values.end());
-        dispersed_shape.push_back(data_shape.at(0) / b_dim_divider);
-        for (size_t i = 1; i < data_shape.size(); ++i) {
-            dispersed_shape.push_back(data_shape.at(i));
-        }
-
-        const auto out_pattern_1 =
-            opset3::Constant::create(element::i64, Shape{dispersed_shape.size()}, dispersed_shape);
+        const auto one = rg.make<Constant>(i64, Shape{1}, 1);
+        const auto end = rg.make<Constant>(i64, Shape{1}, block_length);
+        const auto block_tail = rg.make<Slice>(block, one, end, one);
+        const auto data_shape_tail = rg.make<Slice>(shape_of_data, one, end, one);
+        const auto dispersed_shape = rg.make<Concat>(OutputVector{block_tail, batch_div, data_shape_tail}, 0);
         const bool special_zero = false;
-        std::shared_ptr<Node> flat_node = std::make_shared<ov::opset3::Reshape>(data, out_pattern_1, special_zero);
-        new_ops.push_back(flat_node);
+        shared_ptr<Node> flat_node = rg.make<Reshape>(data, dispersed_shape, special_zero);
+
         // calculate axes to transpose
         //      x'' = transpose(x', [N, N + 1, 0, N + 2, 1, ..., N + N - 1, N - 1])
-        std::vector<size_t> axes_order{block_values.size() - 1};
-        for (size_t i = 0; i < block_values.size() - 1; ++i) {
-            axes_order.push_back(i + block_values.size());
+        vector<int64_t> axes_order{block_length - 1};
+        for (int64_t i = 0; i < block_length - 1; ++i) {
+            axes_order.push_back(i + block_length);
             axes_order.push_back(i);
         }
+        const auto axes_order_const = rg.make<Constant>(i64, Shape{axes_order.size()}, axes_order);
+        flat_node = rg.make<Transpose>(flat_node, axes_order_const);
 
-        const auto axes_order_const =
-            opset3::Constant::create(element::i64,
-                                     Shape{axes_order.size()},
-                                     std::vector<int64_t>(axes_order.begin(), axes_order.end()));
-        flat_node = std::make_shared<ov::opset3::Transpose>(flat_node, axes_order_const);
-        new_ops.push_back(flat_node);
         //   x''' = reshape(x'', [batch / (B_1 * ... * B_{N - 1}), D_1 * B_1, D_2 * B_2, ... , D_{N - 1}
         //   * B_{N - 1}])
-        std::vector<int64_t> squeezed_shape;
-        squeezed_shape.push_back(data_shape.at(0) / b_dim_divider);
-        for (size_t i = 1; i < block_values.size(); ++i) {
-            squeezed_shape.push_back(data_shape.at(i) * block_values.at(i));
-        }
-
-        const auto out_pattern_2 = opset3::Constant::create(element::i64, Shape{squeezed_shape.size()}, squeezed_shape);
-        flat_node = std::make_shared<opset3::Reshape>(flat_node, out_pattern_2, special_zero);
-        new_ops.push_back(flat_node);
+        const auto squeezed_shape_tail = rg.make<Multiply>(block_tail, data_shape_tail);
+        const auto squeezed_shape = rg.make<Concat>(OutputVector{batch_div, squeezed_shape_tail}, 0);
+        flat_node = rg.make<Reshape>(flat_node, squeezed_shape, special_zero);
 
         //    Crop the start and end of dimensions according to `crops_begin`, `crops_end` to produce
         //    the output of shape:
         //    note: `crops_begin[0], crops_end[0]` are expected to be 0.
         //    `y = [batch / (B_1 * ... * B_{N - 1}), crop(D_1 * B_1, crops_begin[1], crops_end[1]),
         //          crop(D_2 * B_2, crops_begin[2], crops_end[2]), ... ,
         //          crop(D_{N - 1} * B_{N - 1}, crops_begin[N - 1], crops_end[N - 1])]`
-        std::vector<int64_t> upperbounds_values;
-        auto flat_node_shape = flat_node->get_shape();
-        for (size_t i = 0; i < flat_node_shape.size(); ++i) {
-            upperbounds_values.push_back(flat_node_shape.at(i) - crops_end_values.at(i));
-        }
-
-        const auto upperbounds = opset3::Constant::create(crops_end.get_element_type(),
-                                                          Shape{upperbounds_values.size()},
-                                                          upperbounds_values);
+        const auto shape_of_flat_node = rg.make<ShapeOf>(flat_node, crops_end.get_element_type());
+        const auto upperbounds = rg.make<Subtract>(shape_of_flat_node, crops_end);
 
-        std::vector<int64_t> begin_mask(data_shape.size(), 0);
-        std::vector<int64_t> end_mask(data_shape.size(), 0);
-        flat_node =
-            std::make_shared<opset3::StridedSlice>(flat_node, crops_begin_const, upperbounds, begin_mask, end_mask);
-        new_ops.push_back(flat_node);
+        const auto begin_mask = vector<int64_t>(data_shape_rank.get_length(), 0);
+        const auto& end_mask = begin_mask;
+        flat_node = rg.make<StridedSlice>(flat_node, crops_begin, upperbounds, begin_mask, end_mask);
 
         flat_node->set_friendly_name(batch_to_space->get_friendly_name());
-        ngraph::copy_runtime_info(batch_to_space, new_ops);
-        ngraph::replace_node(batch_to_space, flat_node);
+        copy_runtime_info(batch_to_space, rg.get());
+        replace_node(batch_to_space, flat_node);
         return true;
     };
 
-    auto m = std::make_shared<ngraph::pattern::Matcher>(batch_to_space, matcher_name);
+    const auto m = make_shared<pattern::Matcher>(batch_to_space, matcher_name);
     this->register_matcher(m, callback);
 }
 
 void ov::pass::ConvertBatchToSpace::convert_batch_to_space_by_elements() {
     MATCHER_SCOPE(ConvertBatchToSpace_convert_batch_to_space_by_elements);
-    auto batch_to_space = ngraph::pattern::wrap_type<ov::opset3::BatchToSpace>();
+    const auto batch_to_space = pattern::wrap_type<BatchToSpace>();
     matcher_pass_callback callback = [this](pattern::Matcher& m) {
-        auto batch_to_space = std::dynamic_pointer_cast<ov::opset3::BatchToSpace>(m.get_match_root());
-        if (!batch_to_space) {
+        const auto batch_to_space = dynamic_pointer_cast<BatchToSpace>(m.get_match_root());
+        if (!batch_to_space || transformation_callback(batch_to_space)) {
             return false;
         }
 
-        auto data = batch_to_space->input_value(0);
+        const auto data = batch_to_space->input_value(0);
 
-        if (data.get_partial_shape().is_dynamic()) {
-            return false;
+        const auto data_shape_rank = data.get_partial_shape().rank();
+        if (data_shape_rank.is_dynamic()) {
+            return false;  // because StridedSlice masks are std::vector
         }
-        auto data_shape = data.get_shape();
 
-        if (transformation_callback(batch_to_space) && (data_shape.size() == 4 || data_shape.size() == 5)) {
-            return false;
-        }
-        auto block = batch_to_space->input_value(1);
-        auto crops_begin = batch_to_space->input_value(2);
-        auto crops_end = batch_to_space->input_value(3);
-
-        const auto block_const = ov::as_type_ptr<opset3::Constant>(block.get_node_shared_ptr());
-        const auto crops_begin_const = ov::as_type_ptr<opset3::Constant>(crops_begin.get_node_shared_ptr());
-        const auto crops_end_const = ov::as_type_ptr<opset3::Constant>(crops_end.get_node_shared_ptr());
-
-        const std::vector<int64_t>& block_values = block_const->cast_vector<int64_t>();
-        const std::vector<int64_t>& crops_end_values = crops_end_const->cast_vector<int64_t>();
-
-        std::vector<int64_t> dispersed_shape(1);
-        dispersed_shape.insert(dispersed_shape.end(), data_shape.begin(), data_shape.end());
-        std::vector<size_t> axes_order(block_values.size() + 1);
-        std::vector<int64_t> squeezed_shape(data_shape.begin(), data_shape.end());
-        if (squeezed_shape.size() > block_values.size()) {
+        const auto block = batch_to_space->input_value(1);
+        const auto crops_begin = batch_to_space->input_value(2);
+        const auto crops_end = batch_to_space->input_value(3);
+
+        if (block.get_partial_shape().is_dynamic() || block.get_shape().size() == 0) {
             return false;
         }
-
-        NodeVector new_ops;
-
-        std::shared_ptr<Node> flat_node = data.get_node_shared_ptr();
-        for (size_t block_idx = 1; block_idx < block_values.size(); ++block_idx) {
-            dispersed_shape[0] = block_values[block_idx];
-            dispersed_shape[1] /= block_values[block_idx];
-            const auto out_pattern_1 =
-                opset3::Constant::create(element::i64, Shape{dispersed_shape.size()}, dispersed_shape);
-            const bool special_zero = false;
-            flat_node = std::make_shared<ov::opset3::Reshape>(flat_node, out_pattern_1, special_zero);
-            new_ops.push_back(flat_node);
-
-            size_t val = 1;
-            for (size_t axis_idx = 0; axis_idx <= block_values.size(); ++axis_idx) {
-                if ((block_idx + 1) == axis_idx) {
+        const auto block_length = static_cast<int64_t>(block.get_shape()[0]);
+
+        NodeRegistry rg;
+        const auto zero = rg.make<Constant>(i64, Shape{1}, 0);
+        const auto one = rg.make<Constant>(i64, Shape{1}, 1);
+        const auto two = rg.make<Constant>(i64, Shape{1}, 2);
+        const auto int_max = rg.make<Constant>(i64, Shape{1}, INT_MAX);
+
+        const auto shape_of_data = rg.make<ShapeOf>(data, block.get_element_type());
+        const auto et_zero = rg.make<Constant>(block.get_element_type(), Shape{1}, 0);
+        shared_ptr<Node> dispersed_shape = rg.make<Concat>(OutputVector{et_zero, shape_of_data}, 0);
+        shared_ptr<Node> squeezed_shape = shape_of_data;
+
+        shared_ptr<Node> flat_node = data.get_node_shared_ptr();
+
+        const auto make_concat = [&](OutputVector nodes) {
+            nodes.erase(remove_if(nodes.begin(),
+                                  nodes.end(),
+                                  [](const Output<Node>& n) {
+                                      return n.get_partial_shape().is_static() && n.get_shape().size() > 0 &&
+                                             n.get_shape()[0] == 0;
+                                  }),
+                        nodes.end());
+            return rg.make<Concat>(nodes, 0);
+        };
+
+        shared_ptr<Node> div;
+        for (int64_t b_idx = 1; b_idx < block_length; ++b_idx) {
+            const auto block_index = rg.make<Constant>(i64, Shape{1}, b_idx);
+            const auto block_index_next = rg.make<Constant>(i64, Shape{1}, b_idx + 1);
+            const auto block_value = rg.make<Gather>(block, block_index, zero);
+
+            // dispersed_shape[0] = block[b_idx];
+            // dispersed_shape[1] /= block[b_idx];
+            if (!div) {
+                const auto batch = rg.make<Gather>(shape_of_data, zero, zero);
+                div = rg.make<Divide>(batch, block_value);
+            } else {
+                div = rg.make<Divide>(div, block_value);
+            }
+            auto ds_tail = rg.make<Slice>(dispersed_shape, two, int_max, one);
+            dispersed_shape = make_concat({block_value, div, ds_tail});
+            constexpr auto special_zero = false;
+            flat_node = rg.make<Reshape>(flat_node, dispersed_shape, special_zero);
+
+            vector<int64_t> axes_order(block_length + 1);
+            int64_t val = 1;
+            for (int64_t axis_idx = 0; axis_idx <= block_length; ++axis_idx) {
+                if ((b_idx + 1) == axis_idx) {
                     axes_order[axis_idx] = 0;
                 } else {
                     axes_order[axis_idx] = val;
                     val++;
                 }
             }
-
-            const auto axes_order_const =
-                ov::opset3::Constant::create(element::i64,
-                                             Shape{axes_order.size()},
-                                             std::vector<int64_t>(axes_order.begin(), axes_order.end()));
-            flat_node = std::make_shared<ov::opset3::Transpose>(flat_node, axes_order_const);
-            new_ops.push_back(flat_node);
-
-            squeezed_shape[0] = dispersed_shape[1];
-            squeezed_shape[block_idx] *= block_values[block_idx];
-            dispersed_shape[block_idx + 1] = squeezed_shape[block_idx];
-            const auto out_pattern_2 =
-                opset3::Constant::create(element::i64, Shape{squeezed_shape.size()}, squeezed_shape);
-            flat_node = std::make_shared<ov::opset3::Reshape>(flat_node, out_pattern_2, special_zero);
-            new_ops.push_back(flat_node);
+            const auto axes_order_const = rg.make<Constant>(i64, Shape{axes_order.size()}, axes_order);
+            flat_node = rg.make<Transpose>(flat_node, axes_order_const);
+
+            // squeezed_shape[0] = dispersed_shape[1];
+            // squeezed_shape[b_idx] *= block[b_idx];
+            const auto sq_slice = rg.make<Slice>(squeezed_shape, one, block_index, one);
+            const auto sq_bidx_dim = rg.make<Gather>(squeezed_shape, block_index, zero);
+            const auto sq_mul = rg.make<Multiply>(sq_bidx_dim, block_value);
+            const auto sq_shape_tail = rg.make<Slice>(squeezed_shape, block_index_next, int_max, one);
+            squeezed_shape.reset();
+            squeezed_shape = make_concat({div, sq_slice, sq_mul, sq_shape_tail});
+            flat_node = rg.make<Reshape>(flat_node, squeezed_shape, special_zero);
+
+            // dispersed_shape[b_idx + 1] = squeezed_shape[b_idx];
+            const auto ds_front = rg.make<Slice>(dispersed_shape, zero, block_index_next, one);
+            ds_tail = rg.make<Slice>(dispersed_shape, rg.make<Constant>(i64, Shape{1}, b_idx + 2), int_max, one);
+            dispersed_shape = make_concat({ds_front, sq_mul, ds_tail});
         }
 
-        std::vector<int64_t> upperbounds_values;
-        auto flat_node_shape = flat_node->get_shape();
-        for (size_t i = 0; i < flat_node_shape.size(); ++i) {
-            upperbounds_values.push_back(flat_node_shape.at(i) - crops_end_values.at(i));
-        }
-        const auto upperbounds = opset3::Constant::create(crops_end.get_element_type(),
-                                                          Shape{upperbounds_values.size()},
-                                                          upperbounds_values);
+        const auto shape_of_flat_node = rg.make<ShapeOf>(flat_node, crops_end.get_element_type());
+        const auto upperbounds = rg.make<Subtract>(shape_of_flat_node, crops_end);
 
-        std::vector<int64_t> begin_mask(data_shape.size(), 0);
-        std::vector<int64_t> end_mask(data_shape.size(), 0);
-        flat_node =
-            std::make_shared<opset3::StridedSlice>(flat_node, crops_begin_const, upperbounds, begin_mask, end_mask);
-        new_ops.push_back(flat_node);
+        const auto begin_mask = vector<int64_t>(data_shape_rank.get_length(), 0);
+        const auto& end_mask = begin_mask;
+        flat_node = rg.make<StridedSlice>(flat_node, crops_begin, upperbounds, begin_mask, end_mask);
 
         flat_node->set_friendly_name(batch_to_space->get_friendly_name());
-        ngraph::copy_runtime_info(batch_to_space, new_ops);
-        ngraph::replace_node(batch_to_space, flat_node);
+        copy_runtime_info(batch_to_space, rg.get());
+        replace_node(batch_to_space, flat_node);
         return true;
     };
 
-    auto m = std::make_shared<ngraph::pattern::Matcher>(batch_to_space, matcher_name);
+    const auto m = make_shared<pattern::Matcher>(batch_to_space, matcher_name);
     this->register_matcher(m, callback);
 }