From f1614b9135b00d402e73ee8c3a0ca02c308172a1 Mon Sep 17 00:00:00 2001
From: "Stepyreva, Evgenya" <evgenya.stepyreva@intel.com>
Date: Thu, 22 Jul 2021 11:59:53 +0300
Subject: [PATCH] set_output_type speedup

---
 .../core/include/ngraph/descriptor/tensor.hpp | 14 +++++-------
 ngraph/core/src/descriptor/tensor.cpp         | 22 ++++++++-----------
 ngraph/core/src/node.cpp                      |  2 +-
 3 files changed, 15 insertions(+), 23 deletions(-)

diff --git a/ngraph/core/include/ngraph/descriptor/tensor.hpp b/ngraph/core/include/ngraph/descriptor/tensor.hpp
index 381e528e531858..9c2612bd0e5e85 100644
--- a/ngraph/core/include/ngraph/descriptor/tensor.hpp
+++ b/ngraph/core/include/ngraph/descriptor/tensor.hpp
@@ -4,9 +4,11 @@
 
 #pragma once
 
+#include <mutex>
 #include <memory>
 #include <string>
 #include <unordered_set>
+#include <atomic>
 
 #include "ngraph/partial_shape.hpp"
 #include "ngraph/shape.hpp"
@@ -73,17 +75,11 @@ namespace ngraph
 
         protected:
             element::Type m_element_type;
-
-            // TODO(amprocte): For now we are maintaining both m_shape and m_partial_shape fields,
-            //    with m_shape possibly being invalid (get_shape will throw an exception if it
-            //    is). This is because get_shape() returns a const reference. I think ideally we
-            //    should refactor so that get_shape returns by value.
-            Shape m_shape;
+            mutable std::atomic<bool> shape_changed;
+            mutable std::mutex shape_mutex;
+            mutable Shape m_shape; // TODO: remove along with get_shape
             PartialShape m_partial_shape;
-            Node* m_node{nullptr};
             HostTensorPtr m_lower_value, m_upper_value;
-            size_t m_node_output_number{0};
-
             std::string m_name;
             std::unordered_set<std::string> m_names;
         };
diff --git a/ngraph/core/src/descriptor/tensor.cpp b/ngraph/core/src/descriptor/tensor.cpp
index 1d8335fee080dc..a1a3b9a33d8ab9 100644
--- a/ngraph/core/src/descriptor/tensor.cpp
+++ b/ngraph/core/src/descriptor/tensor.cpp
@@ -4,7 +4,6 @@
 
 #include "ngraph/descriptor/tensor.hpp"
 #include "ngraph/node.hpp"
-#include "ngraph/runtime/host_tensor.hpp"
 
 using namespace ngraph;
 using namespace std;
@@ -13,9 +12,9 @@ descriptor::Tensor::Tensor(const element::Type& element_type,
                            const PartialShape& pshape,
                            const std::string& name)
     : m_element_type(element_type)
-    , m_shape(pshape.is_static() ? pshape.to_shape() : Shape{})
     , m_partial_shape(pshape)
     , m_name(name)
+    , shape_changed(true)
 {
 }
 
@@ -24,10 +23,8 @@ descriptor::Tensor::Tensor(const element::Type& element_type,
                            Node* node,
                            size_t node_output_number)
     : m_element_type(element_type)
-    , m_shape(pshape.is_static() ? pshape.to_shape() : Shape{})
     , m_partial_shape(pshape)
-    , m_node(node)
-    , m_node_output_number(node_output_number)
+    , shape_changed(true)
 {
 }
 
@@ -46,14 +43,7 @@ void descriptor::Tensor::set_element_type(const element::Type& element_type)
 void descriptor::Tensor::set_partial_shape(const PartialShape& partial_shape)
 {
     m_partial_shape = partial_shape;
-    if (m_partial_shape.is_static())
-    {
-        m_shape = m_partial_shape.to_shape();
-    }
-    else
-    {
-        m_shape = Shape{};
-    }
+    shape_changed = true;
 }
 
 void descriptor::Tensor::invalidate_values()
@@ -82,6 +72,12 @@ const Shape& descriptor::Tensor::get_shape() const
 {
     if (m_partial_shape.is_static())
     {
+        if (shape_changed) {
+            shape_mutex.lock();
+            m_shape = m_partial_shape.to_shape();
+            shape_mutex.unlock();
+            shape_changed = false;
+        }
         return m_shape;
     }
     else
diff --git a/ngraph/core/src/node.cpp b/ngraph/core/src/node.cpp
index 8d23c8f65bbb7e..d600333e900c71 100644
--- a/ngraph/core/src/node.cpp
+++ b/ngraph/core/src/node.cpp
@@ -210,7 +210,7 @@ descriptor::Output& Node::get_output_descriptor(size_t position)
             make_shared<descriptor::Tensor>(element::dynamic, PartialShape::dynamic(), this, i);
         m_outputs.emplace_back(this, i, tensor_descriptor);
     }
-    return m_outputs.at(position);
+    return m_outputs[position];
 }
 
 void Node::set_argument(size_t position, const Output<Node>& argument)