[nGraph] Create Python API support for rt_info (#1696)

ngraph/python/setup.py

@@ -211,6 +211,7 @@ sources = [
     "pyngraph/types/element_type.cpp",
     "pyngraph/types/regmodule_pyngraph_types.cpp",
     "pyngraph/util.cpp",
+    "pyngraph/variant.cpp",
 ]
 
 packages = [

ngraph/python/src/ngraph/__init__.py

@@ -23,6 +23,7 @@ try:
 except DistributionNotFound:
     __version__ = "0.0.0.dev0"
 
+import ngraph.utils.rt_map
 from ngraph.impl import Node
 from ngraph.helpers import function_from_cnn
 

ngraph/python/src/ngraph/impl/__init__.py

@@ -45,5 +45,7 @@ from _pyngraph import CoordinateDiff
 from _pyngraph import AxisSet
 from _pyngraph import AxisVector
 from _pyngraph import Coordinate
+from _pyngraph import VariantString
+from _pyngraph import VariantInt
 
 from _pyngraph import util

ngraph/python/src/ngraph/utils/rt_map.py

@@ -0,0 +1,52 @@
+# ******************************************************************************
+# Copyright 2017-2020 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.
+# ******************************************************************************
+"""Overrides pybind PyRTMap class methods."""
+
+from typing import Any
+
+import _pyngraph
+
+from _pyngraph import Variant
+from ngraph.impl import VariantInt, VariantString
+from ngraph.exceptions import UserInputError
+
+
+def _convert_to_variant(item: Any) -> Variant:
+    """Convert value to Variant class, otherwise throw error."""
+    if isinstance(item, Variant):
+        return item
+    variant_mapping = {
+        int: VariantInt,
+        str: VariantString,
+    }
+
+    new_type = variant_mapping.get(type(item), None)
+
+    if new_type is None:
+        raise UserInputError("Cannot map value to any of registered Variant classes", str(item))
+
+    return new_type(item)
+
+
+binding_copy = _pyngraph.PyRTMap.__setitem__
+
+
+def _setitem(self: _pyngraph.PyRTMap, arg0: str, arg1: Any) -> None:
+    """Override setting values in dictionary."""
+    binding_copy(self, arg0, _convert_to_variant(arg1))
+
+
+_pyngraph.PyRTMap.__setitem__ = _setitem

ngraph/python/src/pyngraph/node.cpp

@@ -14,7 +14,9 @@
 // limitations under the License.
 //*****************************************************************************
 
+#include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
 
 #include "dict_attribute_visitor.hpp"
 #include "ngraph/node.hpp"
@@ -22,12 +24,23 @@
 #include "ngraph/op/divide.hpp"
 #include "ngraph/op/multiply.hpp"
 #include "ngraph/op/subtract.hpp"
+#include "ngraph/variant.hpp"
 #include "pyngraph/node.hpp"
+#include "pyngraph/variant.hpp"
 
 namespace py = pybind11;
 
+using PyRTMap = std::map<std::string, std::shared_ptr<ngraph::Variant>>;
+
+PYBIND11_MAKE_OPAQUE(PyRTMap);
+
 void regclass_pyngraph_Node(py::module m)
 {
+    auto py_map = py::bind_map<PyRTMap>(m, "PyRTMap");
+    py_map.doc() =
+        "ngraph.impl.PyRTMap makes bindings for std::map<std::string, "
+        "std::shared_ptr<ngraph::Variant>>, which can later be used as ngraph::Node::RTMap";
+
     py::class_<ngraph::Node, std::shared_ptr<ngraph::Node>> node(m, "Node", py::dynamic_attr());
     node.doc() = "ngraph.impl.Node wraps ngraph::Node";
     node.def("__add__",
@@ -87,9 +100,15 @@ void regclass_pyngraph_Node(py::module m)
     node.def("outputs",
              (std::vector<ngraph::Output<ngraph::Node>>(ngraph::Node::*)()) &
                  ngraph::Node::outputs);
+    node.def("get_rt_info",
+             (PyRTMap & (ngraph::Node::*)()) & ngraph::Node::get_rt_info,
+             py::return_value_policy::reference_internal);
 
     node.def_property_readonly("shape", &ngraph::Node::get_shape);
     node.def_property_readonly("name", &ngraph::Node::get_name);
+    node.def_property_readonly("rt_info",
+                               (PyRTMap & (ngraph::Node::*)()) & ngraph::Node::get_rt_info,
+                               py::return_value_policy::reference_internal);
     node.def_property(
         "friendly_name", &ngraph::Node::get_friendly_name, &ngraph::Node::set_friendly_name);
 

ngraph/python/src/pyngraph/pyngraph.cpp

@@ -38,6 +38,7 @@
 #include "pyngraph/strides.hpp"
 #include "pyngraph/types/regmodule_pyngraph_types.hpp"
 #include "pyngraph/util.hpp"
+#include "pyngraph/variant.hpp"
 
 namespace py = pybind11;
 
@@ -68,4 +69,7 @@ PYBIND11_MODULE(_pyngraph, m)
     regmodule_pyngraph_op_util(m_op);
     regmodule_pyngraph_passes(m);
     regmodule_pyngraph_util(m);
+    regclass_pyngraph_Variant(m);
+    regclass_pyngraph_VariantWrapper<std::string>(m, std::string("String"));
+    regclass_pyngraph_VariantWrapper<int64_t>(m, std::string("Int"));
 }

ngraph/python/src/pyngraph/variant.cpp

@@ -0,0 +1,31 @@
+//*****************************************************************************
+// Copyright 2017-2020 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 <pybind11/pybind11.h>
+
+#include "ngraph/variant.hpp" // ngraph::Variant
+#include "pyngraph/variant.hpp"
+
+namespace py = pybind11;
+
+void regclass_pyngraph_Variant(py::module m)
+{
+    py::class_<ngraph::Variant, std::shared_ptr<ngraph::Variant>> variant_base(m, "Variant");
+    variant_base.doc() = "ngraph.impl.Variant wraps ngraph::Variant";
+}
+
+template void regclass_pyngraph_VariantWrapper<std::string>(py::module m, std::string typestring);
+template void regclass_pyngraph_VariantWrapper<int64_t>(py::module m, std::string typestring);

ngraph/python/src/pyngraph/variant.hpp

@@ -0,0 +1,50 @@
+//*****************************************************************************
+// Copyright 2017-2020 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 <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <string>
+
+#include "ngraph/variant.hpp" // ngraph::Variant
+
+namespace py = pybind11;
+
+void regclass_pyngraph_Variant(py::module m);
+
+template <typename VT>
+extern void regclass_pyngraph_VariantWrapper(py::module m, std::string typestring)
+{
+    auto pyclass_name = py::detail::c_str((std::string("Variant") + typestring));
+    py::class_<ngraph::VariantWrapper<VT>,
+               std::shared_ptr<ngraph::VariantWrapper<VT>>,
+               ngraph::Variant>
+        variant_wrapper(m, pyclass_name);
+    variant_wrapper.doc() =
+        "ngraph.impl.Variant[typestring] wraps ngraph::VariantWrapper<typestring>";
+
+    variant_wrapper.def(py::init<const VT&>());
+    variant_wrapper.def("get",
+                        (VT & (ngraph::VariantWrapper<VT>::*)()) & ngraph::VariantWrapper<VT>::get);
+    variant_wrapper.def("set", &ngraph::VariantWrapper<VT>::set);
+
+    variant_wrapper.def_property("value",
+                                 (VT & (ngraph::VariantWrapper<VT>::*)()) &
+                                     ngraph::VariantWrapper<VT>::get,
+                                 &ngraph::VariantWrapper<VT>::set);
+}

ngraph/python/tests/test_ngraph/test_basic.py

@@ -20,7 +20,8 @@ import pytest
 
 import ngraph as ng
 from ngraph.exceptions import UserInputError
-from ngraph.impl import Function, PartialShape, Shape
+from ngraph.impl import Function, PartialShape, Shape, Type, VariantInt, VariantString
+from ngraph.impl.op import Parameter
 from tests.runtime import get_runtime
 from tests.test_ngraph.util import run_op_node
 from tests import (xfail_issue_34323,
@@ -31,6 +32,8 @@ from tests import (xfail_issue_34323,
                    xfail_issue_36479,
                    xfail_issue_36480)
 
+from openvino.inference_engine import IENetwork
+
 
 def test_ngraph_function_api():
     shape = [2, 2]
@@ -388,3 +391,41 @@ def test_node_target_inputs_soruce_output():
     assert in_model1.get_node().name == parameter_b.name
     assert np.equal([in_model0.get_shape()], [model.get_output_shape(0)]).all()
     assert np.equal([in_model1.get_shape()], [model.get_output_shape(0)]).all()
+
+
+def test_variants():
+    variant_int = VariantInt(32)
+    variant_str = VariantString("test_text")
+
+    assert variant_int.get() == 32
+    assert variant_str.get() == "test_text"
+
+    variant_int.set(777)
+    variant_str.set("another_text")
+
+    assert variant_int.get() == 777
+    assert variant_str.get() == "another_text"
+
+
+def test_runtime_info():
+    test_shape = PartialShape([1, 3, 22, 22])
+    test_type = Type.f32
+    test_param = Parameter(test_type, test_shape)
+    relu_node = ng.relu(test_param)
+    runtime_info = relu_node.get_rt_info()
+    runtime_info["affinity"] = "test_affinity"
+    relu_node.set_friendly_name("testReLU")
+    runtime_info_after = relu_node.get_rt_info()
+
+    assert runtime_info == runtime_info_after
+
+    params = [test_param]
+    results = [relu_node]
+
+    ng_function = Function(results, params, "testFunc")
+
+    capsule = Function.to_capsule(ng_function)
+    cnn_network = IENetwork(capsule)
+    cnn_layer = cnn_network.layers["testReLU"]
+    assert cnn_layer is not None
+    assert cnn_layer.affinity == "test_affinity"