Add stack fused op for pdpd use (#3965)

* The stack operator support for a engine:
1) Added the package of stack tests
2) Added stack fused operator
3) Modified cmake file to support stack operator

* 1) Support for a v0 version

* Stack.cpp file is modified
1) Removed unused function

* The stack.cpp file is modified:
1) Empty line was removed: [line: 72]

* Modification of test/backend/fused_op.in.cpp file
1) [line 1340], fall back to the previous version of comment

* Code refactoring:
1) Comment message is different, adjusted to get much more details about stack operator
2) Seralize.cpp file is adjusted to work with stack operator

* Added:
1)generated_adjoints has not implemented yet.
2)int64_t -> size_t

* The variable stored negative value was removed from the code

* Source files: code-format check

* Added included

* fix error

* set output type

src/ngraph/CMakeLists.txt

@@ -409,6 +409,8 @@ set (SRC
     op/fused/scale_shift.hpp
     op/fused/scatter_nd.cpp
     op/fused/scatter_nd.hpp
+    op/fused/stack.cpp
+    op/fused/stack.hpp
     op/fused/selu.cpp
     op/fused/selu.hpp
     op/fused/shuffle_channels.cpp

src/ngraph/op/fused/stack.cpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+#include <memory>
+#include <numeric>
+
+#include "matmul.hpp"
+#include "ngraph/builder/reshape.hpp"
+#include "ngraph/op/concat.hpp"
+#include "ngraph/op/fused/stack.hpp"
+#include "ngraph/op/reshape.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+constexpr NodeTypeInfo op::Stack::type_info;
+
+op::Stack::Stack(const OutputVector& args, int64_t axis)
+    : FusedOp(OutputVector{args})
+    , m_axis(axis)
+{
+    constructor_validate_and_infer_types();
+}
+
+op::Stack::Stack(const NodeVector& args, int64_t axis)
+    : Stack(as_output_vector(args), axis)
+{
+}
+
+shared_ptr<Node> op::Stack::copy_with_new_args(const NodeVector& new_args) const
+{
+    return make_shared<Stack>(new_args, m_axis);
+}
+
+void op::Stack::generate_adjoints(autodiff::Adjoints& adjoints, const OutputVector& deltas)
+{
+    ngraph_error("Not yet implemented");
+}
+
+void op::Stack::pre_validate_and_infer_types()
+{
+    bool is_input_dynamic = false;
+
+    for (size_t i = 0; i < get_input_size(); ++i)
+    {
+        if (get_input_partial_shape(i).is_dynamic())
+        {
+            is_input_dynamic = true;
+            break;
+        }
+    }
+
+    if (is_input_dynamic)
+    {
+        set_output_type(0, get_input_element_type(0), PartialShape::dynamic());
+    }
+}
+
+NodeVector op::Stack::decompose_op() const
+{
+    auto axis = get_axis();
+    std::vector<std::shared_ptr<ngraph::Node>> args;
+    PartialShape inputs_shape_scheme{PartialShape::dynamic()};
+    for (size_t i = 0; i < get_input_size(); ++i)
+    {
+        PartialShape this_input_shape = get_input_partial_shape(i);
+        NODE_VALIDATION_CHECK(
+            this,
+            PartialShape::merge_into(inputs_shape_scheme, this_input_shape),
+            "Argument shapes are inconsistent; they must have the same rank, and must have ",
+            "equal dimension everywhere except on the concatenation axis (axis ",
+            axis,
+            ").");
+    }
+
+    for (size_t i = 0; i < get_input_size(); ++i)
+    {
+        auto data = input_value(i);
+        auto data_shape = data.get_shape();
+        axis = (axis < 0) ? axis + data_shape.size() + 1 : axis;
+        data_shape.insert(data_shape.begin() + axis, 1);
+        std::vector<size_t> input_order(data_shape.size() - 1);
+        std::iota(std::begin(input_order), std::end(input_order), 0);
+        args.push_back(std::make_shared<op::Reshape>(data, AxisVector(input_order), data_shape));
+    }
+    auto concat = std::make_shared<op::Concat>(args, axis);
+    return {concat};
+}

src/ngraph/op/fused/stack.hpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#pragma once
+
+#include "ngraph/node.hpp"
+#include "ngraph/op/op.hpp"
+#include "ngraph/op/util/fused_op.hpp"
+
+namespace ngraph
+{
+    namespace op
+    {
+        namespace v0
+        {
+            /// \brief Operator performing Stack.
+            class Stack : public ngraph::op::util::FusedOp
+            {
+            public:
+                NGRAPH_API
+                static constexpr NodeTypeInfo type_info{"Stack", 0};
+                const NodeTypeInfo& get_type_info() const override { return type_info; }
+                Stack() = default;
+
+                /// \brief Constructs a stack operation.
+                ///
+                /// \param args The outputs producing the input tensors.
+                /// \param axis The axis in the result array along which the input arrays are
+                /// stacked.
+                Stack(const OutputVector& args, int64_t axis);
+
+                /// \brief Constructs a stack operation.
+                ///
+                /// \param args The nodes producing the input tensors.
+                /// \param axis The axis in the result array along which the input arrays are
+                /// stacked.
+                Stack(const NodeVector& args, int64_t axis);
+
+                void generate_adjoints(autodiff::Adjoints& adjoints,
+                                       const OutputVector& deltas) override;
+
+                virtual void pre_validate_and_infer_types() override;
+
+                virtual NodeVector decompose_op() const override;
+
+                virtual std::shared_ptr<Node>
+                    copy_with_new_args(const NodeVector& new_args) const override;
+
+                /// \return The stack axis
+                int64_t get_axis() const { return m_axis; }
+                void set_axis(int64_t axis) { m_axis = axis; }
+            private:
+                int64_t m_axis;
+            };
+        }
+        using v0::Stack;
+    } // namespace op
+} // namespace ngraph

src/ngraph/op/op_version_tbl.hpp

@@ -233,6 +233,7 @@ NGRAPH_OP(Split, ngraph::op::v0, 0)
 NGRAPH_OP(Sqrt, ngraph::op, 0)
 NGRAPH_OP(SquaredDifference, ngraph::op::v0, 0)
 NGRAPH_OP(Squeeze, ngraph::op::v0, 0)
+NGRAPH_OP(Stack, ngraph::op::v0, 0)
 NGRAPH_OP(StopGradient, ngraph::op::v0, 0)
 NGRAPH_OP(StridedSlice, ngraph::op::v1, 1)
 NGRAPH_OP(Subtract, ngraph::op::v0, 0)

src/ngraph/ops.hpp

@@ -118,6 +118,7 @@
 #include "ngraph/op/fused/split.hpp"
 #include "ngraph/op/fused/squared_difference.hpp"
 #include "ngraph/op/fused/squeeze.hpp"
+#include "ngraph/op/fused/stack.hpp"
 #include "ngraph/op/fused/unsqueeze.hpp"
 #include "ngraph/op/gather.hpp"
 #include "ngraph/op/gather_nd.hpp"

src/ngraph/opsets/opset0_tbl.hpp

@@ -199,6 +199,7 @@ NGRAPH_OP(Sqrt, ngraph::op)
 NGRAPH_OP(SquaredDifference, ngraph::op)
 NGRAPH_OP(Squeeze, ngraph::op)
 NGRAPH_OP(StopGradient, ngraph::op)
+NGRAPH_OP(Stack, ngraph::op)
 NGRAPH_OP(Subtract, ngraph::op)
 NGRAPH_OP(Sum, ngraph::op)
 NGRAPH_OP(Tan, ngraph::op)

src/ngraph/runtime/generic_cpu/gcpu_executable.hpp

@@ -1646,6 +1646,7 @@ private:
         case OP_TYPEID::Split:
         case OP_TYPEID::SquaredDifference:
         case OP_TYPEID::Squeeze:
+        case OP_TYPEID::Stack:
         // Tensor Iterator not yet supported
         case OP_TYPEID::TensorIterator:
         case OP_TYPEID::Tile:

src/ngraph/runtime/interpreter/int_executable.hpp

@@ -1888,6 +1888,7 @@ private:
         case OP_TYPEID::SpaceToDepth:
         case OP_TYPEID::SquaredDifference:
         case OP_TYPEID::Squeeze:
+        case OP_TYPEID::Stack:
         case OP_TYPEID::Unsqueeze:
         // Tensor Iterator not yet supported
         case OP_TYPEID::TensorIterator:

src/ngraph/serializer.cpp

@@ -2646,6 +2646,11 @@ shared_ptr<Node> JSONDeserializer::deserialize_node(json node_js)
                 args[1],
                 args[2],
                 read_auto_broadcast(node_js, "auto_broadcast", op::AutoBroadcastType::NUMPY));
+        }
+        case OP_TYPEID::Stack:
+        {
+            auto axis = node_js.at("axis").get<size_t>();
+            node = make_shared<op::Stack>(static_cast<OutputVector>(args), axis);
             break;
         }
         case OP_TYPEID::Selu:
@@ -4507,6 +4512,12 @@ json JSONSerializer::serialize_node(const Node& n)
         node["reduction_axes"] = tmp->get_reduction_axes();
         break;
     }
+    case OP_TYPEID::Stack:
+    {
+        auto tmp = static_cast<const op::Stack*>(&n);
+        node["axis"] = tmp->get_axis();
+        break;
+    }
     case OP_TYPEID::ReduceSum_v1:
     {
         auto tmp = static_cast<const op::v1::ReduceSum*>(&n);

test/backend/fused_op.in.cpp

@@ -46,6 +46,94 @@ using namespace ngraph;
 
 static string s_manifest = "${MANIFEST}";
 
+NGRAPH_TEST(${BACKEND_NAME}, stack_matrix_rowise)
+{
+    Shape shape_a{2, 2};
+    auto A = make_shared<op::Parameter>(element::f32, shape_a);
+    Shape shape_b{2, 2};
+    auto B = make_shared<op::Parameter>(element::f32, shape_b);
+    Shape shape_c{2, 2};
+    auto C = make_shared<op::Parameter>(element::f32, shape_c);
+    Shape shape_r{3, 2, 2};
+    auto f = make_shared<Function>(make_shared<op::Stack>(NodeVector{A, B, C}, 0),
+                                   ParameterVector{A, B, C});
+
+    auto backend = runtime::Backend::create("${BACKEND_NAME}");
+
+    // Create some tensors for input/output
+    auto a = backend->create_tensor(element::f32, shape_a);
+    copy_data(a, vector<float>{2, 4, 8, 8});
+    auto b = backend->create_tensor(element::f32, shape_b);
+    copy_data(b, vector<float>{1, 2, 4, 8});
+    auto c = backend->create_tensor(element::f32, shape_c);
+    copy_data(c, vector<float>{2, 3, 5, 7});
+    auto result = backend->create_tensor(element::f32, shape_r);
+
+    auto handle = backend->compile(f);
+    handle->call_with_validate({result}, {a, b, c});
+    EXPECT_TRUE(test::all_close_f((vector<float>{2, 4, 8, 8, 1, 2, 4, 8, 2, 3, 5, 7}),
+                                  read_vector<float>(result),
+                                  MIN_FLOAT_TOLERANCE_BITS));
+}
+
+NGRAPH_TEST(${BACKEND_NAME}, stack_matrix_colwise)
+{
+    Shape shape_a{2, 2};
+    auto A = make_shared<op::Parameter>(element::f32, shape_a);
+    Shape shape_b{2, 2};
+    auto B = make_shared<op::Parameter>(element::f32, shape_b);
+    Shape shape_c{2, 2};
+    auto C = make_shared<op::Parameter>(element::f32, shape_c);
+    Shape shape_r{2, 3, 2};
+    auto f = make_shared<Function>(make_shared<op::Stack>(NodeVector{A, B, C}, 1),
+                                   ParameterVector{A, B, C});
+
+    auto backend = runtime::Backend::create("${BACKEND_NAME}");
+
+    // Create some tensors for input/output
+    auto a = backend->create_tensor(element::f32, shape_a);
+    copy_data(a, vector<float>{2, 4, 8, 8});
+    auto b = backend->create_tensor(element::f32, shape_b);
+    copy_data(b, vector<float>{1, 2, 4, 8});
+    auto c = backend->create_tensor(element::f32, shape_c);
+    copy_data(c, vector<float>{2, 3, 5, 7});
+    auto result = backend->create_tensor(element::f32, shape_r);
+
+    auto handle = backend->compile(f);
+    handle->call_with_validate({result}, {a, b, c});
+    EXPECT_TRUE(test::all_close_f((vector<float>{2, 4, 1, 2, 2, 3, 8, 8, 4, 8, 5, 7}),
+                                  read_vector<float>(result),
+                                  MIN_FLOAT_TOLERANCE_BITS));
+}
+
+NGRAPH_TEST(${BACKEND_NAME}, stack_negative_axis)
+{
+    auto pshape_a = PartialShape::dynamic();
+    auto A = make_shared<op::Parameter>(element::f32, pshape_a);
+    auto pshape_b = PartialShape::dynamic();
+    auto B = make_shared<op::Parameter>(element::f32, pshape_b);
+    auto pshape_c = PartialShape::dynamic();
+    auto C = make_shared<op::Parameter>(element::f32, pshape_c);
+    auto pshape_r = PartialShape::dynamic();
+    auto f = make_shared<Function>(make_shared<op::Stack>(NodeVector{A, B, C}, -1),
+                                   ParameterVector{A, B, C});
+
+    auto backend = runtime::Backend::create("${BACKEND_NAME}", true);
+
+    // Create some tensors for input/output
+    auto a = backend->create_tensor(element::f32, Shape{2, 2});
+    copy_data(a, vector<float>{2, 4, 8, 16});
+    auto b = backend->create_tensor(element::f32, Shape{2, 2});
+    copy_data(b, vector<float>{1, 2, 4, 8});
+    auto c = backend->create_tensor(element::f32, Shape{2, 2});
+    copy_data(c, vector<float>{2, 3, 5, 7});
+    auto result = backend->create_dynamic_tensor(element::f32, PartialShape::dynamic());
+    auto handle = backend->compile(f);
+    handle->call_with_validate({result}, {a, b, c});
+    ASSERT_EQ(result->get_shape(), (Shape{2, 2, 3}));
+    EXPECT_TRUE(test::all_close_f((vector<float>{2, 1, 2, 4, 2, 3, 8, 4, 5, 16, 8, 7}),
+                                  read_vector<float>(result)));
+}
 NGRAPH_TEST(${BACKEND_NAME}, elu)
 {
     auto A = make_shared<op::Parameter>(element::f32, Shape{3, 2});