Add NodeInput and NodeOutput classes (#2612)

* Add NodeInput and NodeOutput classes

* Review comments

src/ngraph/CMakeLists.txt

@@ -69,6 +69,10 @@ set (SRC
     ngraph_visibility.hpp
     node.cpp
     node.hpp
+    node_input.cpp
+    node_input.hpp
+    node_output.cpp
+    node_output.hpp
     node_vector.hpp
     op/abs.cpp
     op/abs.hpp

src/ngraph/ngraph.hpp

@@ -57,6 +57,8 @@
 #include "ngraph/except.hpp"
 #include "ngraph/function.hpp"
 #include "ngraph/node.hpp"
+#include "ngraph/node_input.hpp"
+#include "ngraph/node_output.hpp"
 #include "ngraph/op/abs.hpp"
 #include "ngraph/op/acos.hpp"
 #include "ngraph/op/add.hpp"

src/ngraph/node_input.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 "ngraph/node_input.hpp"
+#include "ngraph/node_output.hpp"
+
+using namespace ngraph;
+
+NodeOutput NodeInput::get_source_output() const
+{
+    auto& output_descriptor = m_node->get_inputs().at(m_index).get_output();
+    return NodeOutput(output_descriptor.get_node(), output_descriptor.get_index());
+}

src/ngraph/node_input.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/descriptor/input.hpp"
+#include "ngraph/node.hpp"
+
+namespace ngraph
+{
+    class NodeOutput;
+
+    /// \brief A handle for one of a node's inputs.
+    class NodeInput
+    {
+    public:
+        /// \brief Constructs a NodeInput.
+        /// \param node Pointer to the node for the input handle.
+        /// \param index The index of the input.
+        NodeInput(Node* node, size_t index)
+            : m_node(node)
+            , m_index(index)
+        {
+        }
+
+        /// \return A pointer to the node referenced by this input handle.
+        Node* get_node() const { return m_node; }
+        /// \return The index of the input referred to by this input handle.
+        size_t get_index() const { return m_index; }
+        /// \return The element type of the input referred to by this input handle.
+        const element::Type& get_element_type() const
+        {
+            return m_node->get_input_element_type(m_index);
+        }
+        /// \return The shape of the input referred to by this input handle.
+        const Shape& get_shape() const { return m_node->get_input_shape(m_index); }
+        /// \return The partial shape of the input referred to by this input handle.
+        const PartialShape& get_partial_shape() const
+        {
+            return m_node->get_input_partial_shape(m_index);
+        }
+        /// \return A handle to the output that is connected to this input.
+        NodeOutput get_source_output() const;
+
+        bool operator==(const NodeInput& other) const
+        {
+            return m_node == other.m_node && m_index == other.m_index;
+        }
+        bool operator!=(const NodeInput& other) const
+        {
+            return m_node != other.m_node || m_index != other.m_index;
+        }
+        bool operator<(const NodeInput& other) const
+        {
+            return m_node < other.m_node || (m_node == other.m_node && m_index < other.m_index);
+        }
+        bool operator>(const NodeInput& other) const
+        {
+            return m_node > other.m_node || (m_node == other.m_node && m_index > other.m_index);
+        }
+        bool operator<=(const NodeInput& other) const
+        {
+            return m_node <= other.m_node || (m_node == other.m_node && m_index <= other.m_index);
+        }
+        bool operator>=(const NodeInput& other) const
+        {
+            return m_node >= other.m_node || (m_node == other.m_node && m_index >= other.m_index);
+        }
+
+    private:
+        Node* const m_node;
+        const size_t m_index;
+    };
+} // namespace ngraph

src/ngraph/node_output.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 "ngraph/node_output.hpp"

src/ngraph/node_output.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/descriptor/input.hpp"
+#include "ngraph/node.hpp"
+
+namespace ngraph
+{
+    /// \brief A handle for one of a node's outputs.
+    class NodeOutput
+    {
+    public:
+        /// \brief Constructs a NodeOutput.
+        /// \param node A `shared_ptr` to the node for the output handle.
+        /// \param index The index of the output.
+        NodeOutput(const std::shared_ptr<Node>& node, size_t index)
+            : m_node(node)
+            , m_index(index)
+        {
+        }
+
+        /// \brief Constructs a NodeOutput, referencing the zeroth output of the node.
+        /// \param node A `shared_ptr` to the node for the output handle.
+        template <typename T>
+        NodeOutput(const std::shared_ptr<T>& node)
+            : NodeOutput(node, 0)
+        {
+        }
+
+        /// \return A `shared_ptr` to the node referred to by this output handle.
+        const std::shared_ptr<Node>& get_node() const { return m_node; }
+        /// \return The index of the output referred to by this output handle.
+        size_t get_index() const { return m_index; }
+        /// \return The element type of the output referred to by this output handle.
+        const element::Type& get_element_type() const
+        {
+            return m_node->get_output_element_type(m_index);
+        }
+        /// \return The shape of the output referred to by this output handle.
+        const Shape& get_shape() const { return m_node->get_output_shape(m_index); }
+        /// \return The partial shape of the output referred to by this output handle.
+        const PartialShape& get_partial_shape() const
+        {
+            return m_node->get_output_partial_shape(m_index);
+        }
+
+        bool operator==(const NodeOutput& other) const
+        {
+            return m_node == other.m_node && m_index == other.m_index;
+        }
+        bool operator!=(const NodeOutput& other) const
+        {
+            return m_node != other.m_node || m_index != other.m_index;
+        }
+        bool operator<(const NodeOutput& other) const
+        {
+            return m_node < other.m_node || (m_node == other.m_node && m_index < other.m_index);
+        }
+        bool operator>(const NodeOutput& other) const
+        {
+            return m_node > other.m_node || (m_node == other.m_node && m_index > other.m_index);
+        }
+        bool operator<=(const NodeOutput& other) const
+        {
+            return m_node <= other.m_node || (m_node == other.m_node && m_index <= other.m_index);
+        }
+        bool operator>=(const NodeOutput& other) const
+        {
+            return m_node >= other.m_node || (m_node == other.m_node && m_index >= other.m_index);
+        }
+
+    private:
+        const std::shared_ptr<Node> m_node;
+        const size_t m_index;
+    };
+} // namespace ngraph

test/CMakeLists.txt

@@ -45,6 +45,7 @@ set(SRC
     input_output_assign.cpp
     main.cpp
     misc.cpp
+    node_input_output.cpp
     nop_elimination.cpp
     op.cpp
     partial_shape.cpp

test/node_input_output.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 <sstream>
+#include <string>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "ngraph/ngraph.hpp"
+
+using namespace ngraph;
+using namespace std;
+
+TEST(node_input_output, input_create)
+{
+    auto x = make_shared<op::Parameter>(element::f32, Shape{1, 2, 3, 4});
+    auto y = make_shared<op::Parameter>(element::f32, Shape{1, 2, 3, 4});
+    auto add = make_shared<op::Add>(x, y);
+
+    auto add_in_0 = NodeInput(add.get(), 0);
+    auto add_in_1 = NodeInput(add.get(), 1);
+    auto add_in_2 = NodeInput(add.get(), 2);
+
+    EXPECT_EQ(add_in_0.get_node(), add.get());
+    EXPECT_EQ(add_in_0.get_index(), 0);
+    EXPECT_EQ(add_in_0.get_element_type(), element::f32);
+    EXPECT_EQ(add_in_0.get_shape(), (Shape{1, 2, 3, 4}));
+    EXPECT_TRUE(add_in_0.get_partial_shape().same_scheme(PartialShape{1, 2, 3, 4}));
+    EXPECT_EQ(add_in_0.get_source_output(), NodeOutput(x, 0));
+
+    EXPECT_EQ(add_in_1.get_node(), add.get());
+    EXPECT_EQ(add_in_1.get_index(), 1);
+    EXPECT_EQ(add_in_1.get_element_type(), element::f32);
+    EXPECT_EQ(add_in_1.get_shape(), (Shape{1, 2, 3, 4}));
+    EXPECT_TRUE(add_in_1.get_partial_shape().same_scheme(PartialShape{1, 2, 3, 4}));
+    EXPECT_EQ(add_in_1.get_source_output(), NodeOutput(y, 0));
+
+    EXPECT_EQ(add_in_2.get_node(), add.get());
+    EXPECT_EQ(add_in_2.get_index(), 2);
+    EXPECT_THROW(add_in_2.get_element_type(), std::out_of_range);
+    EXPECT_THROW(add_in_2.get_shape(), std::out_of_range);
+    EXPECT_THROW(add_in_2.get_partial_shape(), std::out_of_range);
+    EXPECT_THROW(add_in_2.get_source_output(), std::out_of_range);
+}
+
+TEST(node_input_output, output_create)
+{
+    auto x = make_shared<op::Parameter>(element::f32, Shape{1, 2, 3, 4});
+    auto y = make_shared<op::Parameter>(element::f32, Shape{1, 2, 3, 4});
+    auto add = make_shared<op::Add>(x, y);
+
+    auto add_out_0 = NodeOutput(add, 0);
+    auto add_out_1 = NodeOutput(add, 1);
+
+    EXPECT_EQ(add_out_0.get_node(), add);
+    EXPECT_EQ(add_out_0.get_index(), 0);
+    EXPECT_EQ(add_out_0.get_element_type(), element::f32);
+    EXPECT_EQ(add_out_0.get_shape(), (Shape{1, 2, 3, 4}));
+    EXPECT_TRUE(add_out_0.get_partial_shape().same_scheme(PartialShape{1, 2, 3, 4}));
+
+    EXPECT_EQ(add_out_1.get_node(), add);
+    EXPECT_EQ(add_out_1.get_index(), 1);
+    EXPECT_THROW(add_out_1.get_element_type(), std::out_of_range);
+    EXPECT_THROW(add_out_1.get_shape(), std::out_of_range);
+    EXPECT_THROW(add_out_1.get_partial_shape(), std::out_of_range);
+}