Introduce NodePtr (apache#9)

tqchen · May 26, 2018 · 512bac9 · 512bac9
1 parent bc82192
commit 512bac9
Show file tree

Hide file tree

Showing 6 changed files with 6 additions and 8 deletions.
diff --git a/nnvm/include/nnvm/base.h b/nnvm/include/nnvm/base.h
@@ -8,6 +8,7 @@
 
 #include <dmlc/base.h>
 #include <dmlc/any.h>
+#include <dmlc/memory.h>
 #include <dmlc/logging.h>
 #include <dmlc/registry.h>
 #include <dmlc/array_view.h>

diff --git a/nnvm/include/nnvm/graph.h b/nnvm/include/nnvm/graph.h
@@ -81,7 +81,7 @@ void PostOrderDFSVisit(const std::vector<GNode>& heads,
 template<typename FVisit>
 inline void DFSVisit(const std::vector<NodeEntry>& heads,
                      FVisit fvisit) {
-  typedef const std::shared_ptr<Node>* GNode;
+  typedef const NodePtr* GNode;
   std::vector<GNode> head_nodes(heads.size());
   std::transform(heads.begin(), heads.end(), head_nodes.begin(),
                  [](const NodeEntry& e)->GNode {

diff --git a/nnvm/include/nnvm/node.h b/nnvm/include/nnvm/node.h
@@ -20,8 +20,7 @@ class Node;
 
 /*!
  * \brief we always used NodePtr for a reference pointer
- *  to the node, so this alias can be changed in case we need
- *  even faster graph composition than 3M ops/sec.
+ *  to the node, so this alias can be changed in case.
  *
  *  By default, NodePtr is a std::shared_ptr of node
  */

diff --git a/nnvm/src/core/graph_attr_types.cc b/nnvm/src/core/graph_attr_types.cc
@@ -14,7 +14,7 @@ IndexedGraph::IndexedGraph(const Graph &g) {
   std::vector<size_t> inputs_rptr{0}, control_rptr{0};
 
   DFSVisit(g.outputs, [this, &inputs_rptr, &control_rptr]
-             (const std::shared_ptr<nnvm::Node>& n) {
+             (const NodePtr& n) {
       CHECK_LT(nodes_.size(), std::numeric_limits<uint32_t>::max());
       uint32_t nid = static_cast<uint32_t>(nodes_.size());
       // nodes_

diff --git a/nnvm/src/core/node.cc b/nnvm/src/core/node.cc
@@ -12,7 +12,7 @@ Node::~Node() {
     // explicit deletion via DFS
     // this is used to avoid stackoverflow caused by chain of deletions
     std::vector<Node*> stack{this};
-    std::vector<std::shared_ptr<Node> > to_delete;
+    std::vector<NodePtr> to_delete;
     while (!stack.empty()) {
       Node* n = stack.back();
       stack.pop_back();
@@ -37,8 +37,6 @@ Node::~Node() {
 }
 
 NodePtr Node::Create() {
-  // NOTE: possible change to thread local memory pool
-  // via std::allocate_shared instead for faster allocation.
   return std::make_shared<Node>();
 }
 

diff --git a/nnvm/src/test_main.cc b/nnvm/src/test_main.cc
@@ -35,7 +35,7 @@ void test_node_speed() {
   auto add = nnvm::Op::Get("add");
   double tstart = dmlc::GetTime();
   size_t rep = 1000;
-  size_t n = 100;
+  size_t n = 1000;
   for (size_t t = 0; t < rep; ++t) {
     nnvm::Symbol s = nnvm::Symbol::CreateVariable("x");
     for (size_t i = 0; i < n; ++i) {