[VPU][NGraph] Extend TopK K propagation (#2943)

* [VPU][NGraph] Support ShapeOf and Gather in TopK K propagation

* [VPU] Save calculated K value

* [VPU][Tests] Introduces tests

* [Tests] Review fixes

inference-engine/src/vpu/common/include/vpu/ngraph/operations/static_shape_topk.hpp

@@ -58,6 +58,7 @@ public:
 
 protected:
     int64_t m_axis;
+    int64_t m_maximumK;
     uint64_t m_normalized_axis;
     Mode m_mode;
     SortType m_sort;

inference-engine/src/vpu/common/src/ngraph/operations/static_shape_topk.cpp

@@ -18,6 +18,7 @@ ngraph::vpu::op::StaticShapeTopK::StaticShapeTopK(
         const element::Type& index_element_type)
         : Op{{data, k}}
         , m_axis{axis}
+        , m_maximumK{-1}
         , m_normalized_axis{0}
         , m_mode{as_enum<Mode>(mode)}
         , m_sort{as_enum<SortType>(sort)}
@@ -34,6 +35,7 @@ ngraph::vpu::op::StaticShapeTopK::StaticShapeTopK(
         const ngraph::element::Type &index_element_type)
         : Op{{data, k}}
         , m_axis{axis}
+        , m_maximumK{-1}
         , m_normalized_axis{0}
         , m_mode{mode}
         , m_sort{sort}
@@ -95,11 +97,12 @@ void ngraph::vpu::op::StaticShapeTopK::validate_and_infer_types() {
         const auto is_max_value_calculated = max_k.first;
         const auto calculated_max_value = max_k.second;
         if (is_max_value_calculated) {
-            output_shape[m_normalized_axis] = calculated_max_value;
-        } else {
-            output_shape[m_normalized_axis] = -1;
+            m_maximumK = calculated_max_value;
         }
     }
+
+    output_shape[m_normalized_axis] = m_maximumK;
+
     NODE_VALIDATION_CHECK(this, output_shape.is_static(),
             "StaticShapeTopK output shape is not fully defined: ", output_shape);
 

inference-engine/tests/functional/plugin/myriad/subgraph_tests/topk_k_propagation.cpp

@@ -7,7 +7,7 @@
 #include <ngraph/type/element_type.hpp>
 #include <ngraph/function.hpp>
 #include <ngraph_functions/utils/ngraph_helpers.hpp>
-#include <ngraph/opsets/opset4.hpp>
+#include <ngraph/opsets/opset5.hpp>
 #include <vpu/ngraph/operations/dynamic_shape_resolver.hpp>
 #include <vpu/ngraph/operations/static_shape_topk.hpp>
 #include <vpu/ngraph/transformations/dynamic_to_static_shape_topk.hpp>
@@ -15,34 +15,14 @@
 
 namespace {
 
-class DynamicToStaticTopKPropagation : public CommonTestUtils::TestsCommon,
-        public testing::WithParamInterface<int64_t> {
-public:
-    void SetUp() override {
+class DynamicToStaticTopKPropagationBase : public CommonTestUtils::TestsCommon,
+                                           public testing::WithParamInterface<int64_t> {
+protected:
+    void validate(const ngraph::Function& function) const {
         const auto& k = GetParam();
 
-        const auto data = std::make_shared<ngraph::opset4::Parameter>(ngraph::element::i64, ngraph::Shape{1000});
-        const auto realK = std::make_shared<ngraph::opset4::Parameter>(ngraph::element::i32, ngraph::Shape{1});
-        const auto maxK = ngraph::opset4::Constant::create(ngraph::element::i32, {1}, {k});
-
-        const auto concat = std::make_shared<ngraph::opset4::Concat>(ngraph::OutputVector{realK, maxK}, 0);
-
-        const auto reduceMin = std::make_shared<ngraph::opset4::ReduceMin>(concat, ngraph::opset4::Constant::create(ngraph::element::i32, {1}, {0}), false);
-        const auto builtSubgraph = buildSubgraph(reduceMin);
-
-        const auto dsr = std::make_shared<ngraph::vpu::op::DynamicShapeResolver>(data, realK);
-        const auto topK = std::make_shared<ngraph::opset4::TopK>(dsr, builtSubgraph, 0, "max", "value");
-
-        ngraph::ResultVector results{std::make_shared<ngraph::opset4::Result>(topK->output(0)),
-                                     std::make_shared<ngraph::opset4::Result>(topK->output(1))};
-        const auto function = std::make_shared<ngraph::Function>(results, ngraph::ParameterVector{data, realK}, "TopKPropagationOfK");
-        topK->set_output_type(0, dsr->get_input_element_type(0), ngraph::PartialShape::dynamic(1));
-
-        const auto transformations = vpu::Transformations{{topK->type_info, vpu::dynamicToStaticShapeTopK}};
-        ASSERT_NO_THROW(vpu::DynamicToStaticShape(transformations).run_on_function(function));
-
         ngraph::ResultVector processedResults;
-        ASSERT_NO_THROW(processedResults = function->get_results());
+        ASSERT_NO_THROW(processedResults = function.get_results());
         EXPECT_EQ(processedResults.size(), 2);
 
         const auto topKOutPartialShape = processedResults[0]->get_input_partial_shape(0);
@@ -53,55 +33,133 @@ public:
         EXPECT_EQ(topKOutShape[0], k);
     }
 
-protected:
     virtual std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const {
         return node;
     }
 };
 
+class DynamicToStaticTopKPropagationConcatBased : public DynamicToStaticTopKPropagationBase {
+public:
+    void SetUp() override {
+        const auto& k = GetParam();
+        ngraph::Shape inputShape{upperBoundK};
+
+        const auto data = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, inputShape);
+        const auto dataShape = ngraph::opset5::Constant::create(ngraph::element::i32, ngraph::Shape{inputShape.size()}, inputShape);
+        const auto resultK = ngraph::opset5::Constant::create(ngraph::element::i32, {1}, {k});
+
+        const auto concat = std::make_shared<ngraph::opset5::Concat>(ngraph::OutputVector{dataShape, resultK}, 0);
+
+        const auto reduceMin = std::make_shared<ngraph::opset5::ReduceMin>(concat, ngraph::opset5::Constant::create(ngraph::element::i32, {1}, {0}), false);
+        const auto builtSubgraph = buildSubgraph(reduceMin);
+
+        const auto dsr = std::make_shared<ngraph::vpu::op::DynamicShapeResolver>(data, dataShape);
+        const auto topK = std::make_shared<ngraph::opset5::TopK>(dsr, builtSubgraph, 0, "max", "value");
+
+        ngraph::ResultVector results{std::make_shared<ngraph::opset5::Result>(topK->output(0)),
+                                     std::make_shared<ngraph::opset5::Result>(topK->output(1))};
+        const auto function = std::make_shared<ngraph::Function>(results, ngraph::ParameterVector{data}, "TopKPropagationOfK");
+
+        const auto transformations = vpu::Transformations{{topK->type_info, vpu::dynamicToStaticShapeTopK}};
+        ASSERT_NO_THROW(vpu::DynamicToStaticShape(transformations).run_on_function(function));
+        validate(*function);
+    }
+
+protected:
+    static constexpr int64_t upperBoundK = 1000;
+};
+
 const std::vector<int64_t> kVec = {0, 10, 100, 200, 500};
 
-TEST_P(DynamicToStaticTopKPropagation, KPropagation) {
+TEST_P(DynamicToStaticTopKPropagationConcatBased, KPropagation) {
 }
 
-INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagation, ::testing::ValuesIn(kVec));
+INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConcatBased, ::testing::ValuesIn(kVec));
 
-class DynamicToStaticTopKPropagationReshape : public DynamicToStaticTopKPropagation {
+class DynamicToStaticTopKPropagationConcatReshape : public DynamicToStaticTopKPropagationConcatBased {
 protected:
     std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const override {
-        return std::make_shared<ngraph::opset4::Reshape>(node, ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}), false);
+        return std::make_shared<ngraph::opset5::Reshape>(node, ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}), false);
     }
 };
 
-TEST_P(DynamicToStaticTopKPropagationReshape, KPropagation) {
+TEST_P(DynamicToStaticTopKPropagationConcatReshape, KPropagation) {
 }
 
-INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationReshape, ::testing::ValuesIn(kVec));
+INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConcatReshape, ::testing::ValuesIn(kVec));
 
-class DynamicToStaticTopKPropagationSqueezeUnsqueeze : public DynamicToStaticTopKPropagation {
+class DynamicToStaticTopKPropagationConcatSqueezeUnsqueeze : public DynamicToStaticTopKPropagationConcatBased {
 protected:
     std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const override {
-        const auto unsqueeze = std::make_shared<ngraph::opset4::Unsqueeze>(node, ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}));
-        return std::make_shared<ngraph::opset4::Squeeze>(unsqueeze, ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}));
+        const auto unsqueeze = std::make_shared<ngraph::opset5::Unsqueeze>(
+            node,
+            ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}));
+        return std::make_shared<ngraph::opset5::Squeeze>(unsqueeze, ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}));
     }
 };
 
-TEST_P(DynamicToStaticTopKPropagationSqueezeUnsqueeze, KPropagation) {
+TEST_P(DynamicToStaticTopKPropagationConcatSqueezeUnsqueeze, KPropagation) {
 }
 
-INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationSqueezeUnsqueeze, ::testing::ValuesIn(kVec));
+INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConcatSqueezeUnsqueeze, ::testing::ValuesIn(kVec));
 
-class DynamicToStaticTopKPropagationConvert : public DynamicToStaticTopKPropagation {
+class DynamicToStaticTopKPropagationConcatConvert : public DynamicToStaticTopKPropagationConcatBased {
 protected:
     std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const override {
-        const auto convert = std::make_shared<ngraph::opset4::Convert>(node, ngraph::element::i32);
-        return std::make_shared<ngraph::opset4::Convert>(convert, ngraph::element::i64);
+        const auto convert = std::make_shared<ngraph::opset5::Convert>(node, ngraph::element::i32);
+        return std::make_shared<ngraph::opset5::Convert>(convert, ngraph::element::i64);
     }
 };
 
-TEST_P(DynamicToStaticTopKPropagationConvert, KPropagation) {
+TEST_P(DynamicToStaticTopKPropagationConcatConvert, KPropagation) {
 }
 
-INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConvert, ::testing::ValuesIn(kVec));
+INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConcatConvert, ::testing::ValuesIn(kVec));
+
+class DynamicToStaticTopKPropagationShapeOfBased : public DynamicToStaticTopKPropagationBase {
+public:
+    void SetUp() override {
+        const auto& k = GetParam();
+
+        const auto upperBoundData = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, ngraph::Shape{upperBoundK});
+        const auto realData = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i32, ngraph::Shape{static_cast<size_t>(k)});
+
+        const auto dsr = std::make_shared<ngraph::vpu::op::DynamicShapeResolver>(
+            upperBoundData,
+            ngraph::opset5::Constant::create(ngraph::element::i32, ngraph::Shape{1}, {k}));
+
+        const auto shapeOf = std::make_shared<ngraph::opset5::ShapeOf>(realData);
+        const auto builtSubgraph = buildSubgraph(shapeOf);
+
+        const auto topK = std::make_shared<ngraph::opset5::TopK>(dsr, builtSubgraph, 0, "max", "value");
+
+        ngraph::ResultVector results{std::make_shared<ngraph::opset5::Result>(topK->output(0)),
+                                     std::make_shared<ngraph::opset5::Result>(topK->output(1))};
+
+        const auto function = std::make_shared<ngraph::Function>(results, ngraph::ParameterVector{upperBoundData, realData}, "TopKPropagationOfK");
+
+        const auto transformations = vpu::Transformations{{topK->type_info, vpu::dynamicToStaticShapeTopK}};
+        ASSERT_NO_THROW(vpu::DynamicToStaticShape(transformations).run_on_function(function));
+        validate(*function);
+    }
+
+protected:
+    static constexpr int64_t upperBoundK = 1000;
+};
+
+class DynamicToStaticTopKPropagationShapeOfGather : public DynamicToStaticTopKPropagationShapeOfBased {
+protected:
+    std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const override {
+        return std::make_shared<ngraph::opset5::Gather>(
+            node,
+            ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0}),
+            ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0}));
+    }
+};
+
+TEST_P(DynamicToStaticTopKPropagationShapeOfGather, KPropagation) {
+}
+
+INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationShapeOfGather, ::testing::ValuesIn(kVec));
 
 }  // namespace

ngraph/core/src/validation_util.cpp

@@ -19,9 +19,11 @@
 #include "ngraph/evaluator.hpp"
 #include "ngraph/op/concat.hpp"
 #include "ngraph/op/convert.hpp"
+#include "ngraph/op/gather.hpp"
 #include "ngraph/op/min.hpp"
 #include "ngraph/op/minimum.hpp"
 #include "ngraph/op/reshape.hpp"
+#include "ngraph/op/shape_of.hpp"
 #include "ngraph/op/squeeze.hpp"
 #include "ngraph/op/unsqueeze.hpp"
 #include "ngraph/runtime/host_tensor.hpp"
@@ -1084,6 +1086,54 @@ namespace
         // Noting we can do
         return {MaxValue(data.m_value)};
     }
+
+    vector<MaxValue> exec_shape_of(Node* node, vector<MaxValue>& inputs)
+    {
+        const auto& inputPS = node->get_input_partial_shape(0);
+        std::vector<uint64_t> shapeDims;
+        for (size_t i = 0; i < inputPS.rank().get_length(); i++)
+        {
+            if (inputPS[i].is_static())
+            {
+                shapeDims.push_back(inputPS[i].get_length());
+            }
+            else
+            {
+                shapeDims.push_back(std::numeric_limits<uint64_t>::max());
+            }
+        }
+
+        return {MaxValue(shapeDims, 0)};
+    }
+
+    vector<MaxValue> exec_gather(Node* node, vector<MaxValue>& inputs)
+    {
+        auto gather = as_type<op::v1::Gather>(node);
+
+        const auto& indices =
+            as_type_ptr<op::v0::Constant>(node->input_value(1).get_node_shared_ptr());
+        const auto& axis =
+            as_type_ptr<op::v0::Constant>(node->input_value(2).get_node_shared_ptr());
+
+        if (!indices || !axis)
+        {
+            return {MaxValue()};
+        }
+
+        if (gather->get_axis() != 0)
+        {
+            return {MaxValue()};
+        }
+
+        const auto& indicesVec = indices->cast_vector<int64_t>();
+        if (indicesVec.size() != 1 || indicesVec[0] >= inputs[0].m_slices.size())
+        {
+            return {MaxValue()};
+        }
+
+        return {MaxValue(inputs[0].m_slices[indicesVec[0]])};
+    }
+
     vector<MaxValue> exec_nop(Node* node, vector<MaxValue>& inputs) { return {inputs.at(0)}; }
 }
 
@@ -1093,12 +1143,14 @@ pair<bool, uint64_t> ngraph::maximum_value(const Output<Node>& value)
         {op::v0::Concat::type_info, exec_concat},
         {op::v0::Constant::type_info, exec_constant},
         {op::v0::Convert::type_info, exec_nop},
+        {op::v1::Gather::type_info, exec_gather},
         {op::v0::Minimum::type_info, exec_minimum},
         {op::v1::Minimum::type_info, exec_minimum},
         {op::v1::ReduceMin::type_info, exec_reduce_min},
+        {op::v1::Reshape::type_info, exec_nop},
+        {op::v3::ShapeOf::type_info, exec_shape_of},
         {op::v0::Squeeze::type_info, exec_nop},
-        {op::v0::Unsqueeze::type_info, exec_nop},
-        {op::v1::Reshape::type_info, exec_nop}};
+        {op::v0::Unsqueeze::type_info, exec_nop}};
     Evaluator<MaxValue>::value_map value_map;
     Evaluator<MaxValue> evaluator(handlers, value_map);
     auto val = evaluator.evaluate(value);