From 2da23bd86905c3abec857a28b6135d069ea9515d Mon Sep 17 00:00:00 2001
From: Pedro Larroy <pedro.larroy.lists@gmail.com>
Date: Fri, 8 Feb 2019 19:18:50 +0100
Subject: [PATCH] Optimize move semantics of NodeEntry reducing copies of
 shared_ptr which causes atomic contention (#2576)

---
 nnvm/include/nnvm/node.h        | 25 +++++++++++++++++++++++--
 nnvm/src/core/node.cc           |  4 ----
 nnvm/src/core/symbolic.cc       |  8 ++++----
 nnvm/src/pass/correct_layout.cc |  4 ++--
 4 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/nnvm/include/nnvm/node.h b/nnvm/include/nnvm/node.h
index ae782f04965e..4527401f45a3 100644
--- a/nnvm/include/nnvm/node.h
+++ b/nnvm/include/nnvm/node.h
@@ -30,6 +30,18 @@ using NodePtr = std::shared_ptr<Node>;
 
 /*! \brief an entry that represents output data from a node */
 struct NodeEntry {
+  NodeEntry(NodePtr node, uint32_t index, uint32_t version):
+    node(std::move(node)),
+    index(index),
+    version(version)
+  {}
+
+  NodeEntry():
+    node(),
+    index(),
+    version()
+  {}
+
   /*! \brief the source node of this data */
   NodePtr node;
   /*! \brief index of output from the source. */
@@ -113,6 +125,11 @@ struct NodeAttrs {
  */
 class NNVM_DLL Node {
  public:
+  Node() = default;
+  Node(const Op* op, const std::string& name) {
+    this->attrs.op = op;
+    this->attrs.name = name;
+  }
   /*! \brief The attributes in the node. */
   NodeAttrs attrs;
   /*! \brief inputs to this node */
@@ -142,7 +159,10 @@ class NNVM_DLL Node {
    * \brief create a new empty shared_ptr of Node.
    * \return a created empty node.
    */
-  static NodePtr Create();
+  template<class ...Args>
+  static NodePtr Create(Args&&... args) {
+    return std::make_shared<Node>(std::forward<Args>(args)...);
+  }
 };
 
 /*!
@@ -167,13 +187,14 @@ inline NodeEntry MakeNode(
     p->attrs.op->attr_parser(&(p->attrs));
   }
   p->inputs = std::move(inputs);
-  return NodeEntry{p, 0, 0};
+  return NodeEntry(p, 0, 0);
 }
 
 // implementation of functions.
 inline const Op* Node::op() const {
   return this->attrs.op;
 }
+
 inline bool Node::is_variable() const {
   return this->op() == nullptr;
 }
diff --git a/nnvm/src/core/node.cc b/nnvm/src/core/node.cc
index b5b5ec22fd3f..fe04466a87d1 100644
--- a/nnvm/src/core/node.cc
+++ b/nnvm/src/core/node.cc
@@ -37,8 +37,4 @@ Node::~Node() {
   }
 }
 
-NodePtr Node::Create() {
-  return std::make_shared<Node>();
-}
-
 }  // namespace nnvm
diff --git a/nnvm/src/core/symbolic.cc b/nnvm/src/core/symbolic.cc
index 680a1900c008..938053aae688 100644
--- a/nnvm/src/core/symbolic.cc
+++ b/nnvm/src/core/symbolic.cc
@@ -601,8 +601,8 @@ Symbol Symbol::CreateFunctor(const Op* op,
   if (fnum_vis_output.count(n->op())) {
     nout = fnum_vis_output[n->op()](n->attrs);
   }
-  for (uint32_t i = 0; i < nout; ++i) {
-    s.outputs.emplace_back(NodeEntry{n, i, 0});
+  for (size_t i = 0; i < nout; i++) {
+    s.outputs.emplace_back(n, i, 0);
   }
   return s;
 }
@@ -618,7 +618,7 @@ Symbol Symbol::CreateFunctor(const NodeAttrs& attrs) {
     nout = fnum_vis_output[n->op()](n->attrs);
   }
   for (uint32_t i = 0; i < nout; ++i) {
-    s.outputs.emplace_back(NodeEntry{n, i, 0});
+    s.outputs.emplace_back(n, i, 0);
   }
   return s;
 }
@@ -633,7 +633,7 @@ Symbol Symbol::CreateGroup(const std::vector<Symbol> &symbols) {
 
 Symbol Symbol::CreateVariable(const std::string& name) {
   Symbol s;
-  s.outputs.emplace_back(NodeEntry{CreateVariableNode(name), 0, 0});
+  s.outputs.emplace_back(CreateVariableNode(name), 0, 0);
   return s;
 }
 
diff --git a/nnvm/src/pass/correct_layout.cc b/nnvm/src/pass/correct_layout.cc
index cd088257d1b0..cba12514c63a 100644
--- a/nnvm/src/pass/correct_layout.cc
+++ b/nnvm/src/pass/correct_layout.cc
@@ -114,10 +114,10 @@ nnvm::Graph CorrectLayout(nnvm::Graph src) {
         nnvm::NodePtr tnode = CreateLayoutTransformNode(produce, request);
         tnode->attrs.name = idx[e.node_id].source->attrs.name + "_" + request.name();
         tnode->inputs.emplace_back(new_node->inputs[i]);
-        nnvm::NodeEntry tnode_output{tnode, 0, 0};
+        nnvm::NodeEntry tnode_output(std::move(tnode), 0, 0);
         new_node->inputs[i] = tnode_output;
         // layout produced by LayoutTransformNode
-        new_layouts[tnode.get()] = {request};
+        new_layouts[tnode_output.node.get()] = {request};
       } else if (!produce.defined()) {
         // do reverse infer
         new_layouts[in.get()][e.index] = request;