fix tsort crash if rerun (#141)

add tensor buffer size unit test

fix warnings when building mkldnn test

get tensor is_output set. Add Manager state and Function to pass manager

use pass manager internal state

src/ngraph/CMakeLists.txt

@@ -26,6 +26,7 @@ set (SRC
     pass/liveness.cpp
     pass/manager.cpp
     pass/memory_layout.cpp
+    pass/memory_visualize.cpp
     pass/pass.cpp
     pass/propagate_types.cpp
     pass/topological_sort.cpp

src/ngraph/descriptor/tensor.cpp

@@ -13,35 +13,54 @@
 // ----------------------------------------------------------------------------
 
 #include "ngraph/descriptor/tensor.hpp"
+#include "ngraph/descriptor/tensor_view.hpp"
 #include "ngraph/node.hpp"
 
-namespace ngraph
+using namespace ngraph;
+using namespace ngraph::descriptor;
+
+Tensor::Tensor(const element::Type& element_type,
+        PrimaryTensorView*   primary_tensor_view,
+        const Node*          parent,
+        size_t               value_index)
+    : m_element_type(element_type)
+    , m_primary_tensor_view(primary_tensor_view)
+    , m_is_output{parent->is_output()}
+    , m_is_input{parent->is_parameter()}
+    , m_is_persistent{false}
+    , m_name{parent->get_node_id() + "_" + std::to_string(value_index)}
+    , m_next_view_id{0}
 {
-    namespace descriptor
+    size_t size = 1;
+    for (size_t s : primary_tensor_view->get_tensor_view_type()->get_shape())
     {
-        Tensor::Tensor(const element::Type& element_type,
-                       PrimaryTensorView*   primary_tensor_view,
-                       const Node*          parent,
-                       size_t               value_index)
-            : m_element_type(element_type)
-            , m_primary_tensor_view(primary_tensor_view)
-            , m_is_output{false}
-            , m_is_input{parent->is_parameter()}
-            , m_is_persistent{false}
-            , m_name{parent->get_node_id() + "_" + std::to_string(value_index)}
-            , m_next_view_id{0}
-        {
-        }
-
-        std::string Tensor::get_next_view_name()
-        {
-            return m_name + "_TV" + std::to_string(m_next_view_id++);
-        }
-
-        std::ostream& operator<<(std::ostream& out, const Tensor& tensor)
-        {
-            out << "Tensor(" << tensor.get_name() << ")";
-            return out;
-        }
+        size *= s;
     }
+    m_size = size * m_element_type.size();
+}
+
+std::string Tensor::get_next_view_name()
+{
+    return m_name + "_TV" + std::to_string(m_next_view_id++);
+}
+
+size_t Tensor::size() const
+{
+    return m_size;
+}
+
+void Tensor::set_pool_offset(size_t offset)
+{
+    m_pool_offset = offset;
+}
+
+size_t Tensor::get_pool_offset() const
+{
+    return m_pool_offset;
+}
+
+std::ostream& descriptor::operator<<(std::ostream& out, const Tensor& tensor)
+{
+    out << "Tensor(" << tensor.get_name() << ")";
+    return out;
 }

src/ngraph/descriptor/tensor.hpp

@@ -31,37 +31,44 @@ namespace ngraph
     {
         class TensorView;
         class PrimaryTensorView;
+        class Tensor;
+    }
+}
+
+class ngraph::descriptor::Tensor
+{
+    friend class PrimaryTensorView;
 
-        class Tensor
-        {
-            friend class PrimaryTensorView;
+private:
+    Tensor(const Tensor&) = delete;
+    Tensor& operator=(const Tensor&) = delete;
 
-        private:
-            Tensor(const Tensor&) = delete;
-            Tensor& operator=(const Tensor&) = delete;
+    Tensor(const element::Type& element_type,
+           PrimaryTensorView*   tensor_view,
+           const Node*          parent,
+           size_t               value_index);
 
-            Tensor(const element::Type& element_type,
-                   PrimaryTensorView*   tensor_view,
-                   const Node*          parent,
-                   size_t               value_index);
+    std::string get_next_view_name();
 
-            std::string get_next_view_name();
+public:
+    bool               is_output() const { return m_is_output; }
+    bool               is_input() const { return m_is_input; }
+    bool               is_persistent() const { return m_is_persistent; }
+    const std::string& get_name() const { return m_name; }
+    size_t             size() const;
+    void               set_pool_offset(size_t);
+    size_t             get_pool_offset() const;
 
-        public:
-            bool               is_output() const { return m_is_output; }
-            bool               is_input() const { return m_is_input; }
-            bool               is_persistent() const { return m_is_persistent; }
-            const std::string& get_name() const { return m_name; }
-        protected:
-            const element::Type& m_element_type;
-            PrimaryTensorView*   m_primary_tensor_view;
-            bool                 m_is_output;
-            bool                 m_is_input;
-            bool                 m_is_persistent;
-            std::string          m_name;
-            size_t               m_next_view_id;
-        };
+    friend std::ostream& operator<<(std::ostream&, const Tensor&);
 
-        std::ostream& operator<<(std::ostream&, const Tensor&);
-    }
-}
+protected:
+    const element::Type& m_element_type;
+    PrimaryTensorView*   m_primary_tensor_view;
+    bool                 m_is_output;
+    bool                 m_is_input;
+    bool                 m_is_persistent;
+    std::string          m_name;
+    size_t               m_next_view_id;
+    size_t               m_size;
+    size_t               m_pool_offset;
+};

src/ngraph/descriptor/tensor_view.cpp

@@ -17,6 +17,16 @@
 using namespace ngraph;
 using namespace descriptor;
 
+PrimaryTensorView::PrimaryTensorView(const std::shared_ptr<const TensorViewType>& tensor_view_type,
+    const Node* parent, size_t value_index)
+    : TensorView(tensor_view_type)
+    , m_tensor(tensor_view_type->get_element_type(), this, parent, value_index)
+{
+    // Set the name in the parent TensorView.
+    // This can't be done until after the m_tensor is constructed.
+    m_name = m_tensor.get_next_view_name();
+}
+
 const Tensor& PrimaryTensorView::get_tensor() const
 {
     return m_tensor;

src/ngraph/descriptor/tensor_view.hpp

@@ -29,7 +29,7 @@ namespace ngraph
         class TensorViewLayout;
 
         // Describes a view of an instantiated tensor
-        class TensorView : public std::enable_shared_from_this<TensorView>
+        class TensorView
         {
             TensorView(const TensorView&) = delete;
             TensorView& operator=(const TensorView&) = delete;
@@ -73,14 +73,7 @@ namespace ngraph
         {
         public:
             PrimaryTensorView(const std::shared_ptr<const TensorViewType>& tensor_view_type,
-                const Node* parent, size_t value_index)
-                : TensorView(tensor_view_type)
-                , m_tensor(tensor_view_type->get_element_type(), this, parent, value_index)
-            {
-                // Set the name in the parent TensorView.
-                // This can't be done until after the m_tensor is constructed.
-                m_name = m_tensor.get_next_view_name();
-            }
+                const Node* parent, size_t value_index);
 
             virtual const Tensor& get_tensor() const override;
             virtual Tensor&       get_tensor() override;

src/ngraph/function.hpp

@@ -14,6 +14,11 @@
 
 #pragma once
 
+#include <initializer_list>
+#include <memory>
+#include <vector>
+#include <string>
+
 #include "ngraph/descriptor/tensor_view.hpp"
 #include "ngraph/node.hpp"
 #include "ngraph/ops/op.hpp"

src/ngraph/node.cpp

@@ -23,6 +23,7 @@ Node::Node(const std::vector<shared_ptr<Node>>& arguments, shared_ptr<ValueType>
     : m_arguments(arguments)
     , m_value_type(value_type)
     , m_instance_id(m_next_instance_id++)
+    , m_is_output(false)
 {
     // Add this node as a user of each argument.
     for (auto node : m_arguments)
@@ -81,6 +82,16 @@ bool Node::is_parameter() const
     return dynamic_cast<const op::Parameter*>(this) != nullptr;
 }
 
+bool Node::is_output() const
+{
+    return m_is_output;
+}
+
+void Node::set_is_output()
+{
+    m_is_output = true;
+}
+
 std::string Node::get_node_id() const
 {
     stringstream ss;

src/ngraph/node.hpp

@@ -98,6 +98,8 @@ namespace ngraph
         void set_value_type_checked(const std::shared_ptr<ValueType>& value_type);
 
         bool is_parameter() const;
+        bool is_output() const;
+        void set_is_output();
 
         size_t               get_instance_id() const { return m_instance_id; }
         friend std::ostream& operator<<(std::ostream&, const Node&);
@@ -120,5 +122,6 @@ namespace ngraph
         static size_t                   m_next_instance_id;
         std::vector<descriptor::Input>  m_inputs;
         std::vector<descriptor::Output> m_outputs;
+        bool                            m_is_output;
     };
 }

src/ngraph/pass/assign_tensors.cpp

@@ -20,6 +20,7 @@
 #include "ngraph/log.hpp"
 #include "ngraph/ngraph.hpp"
 #include "ngraph/pass/propagate_types.hpp"
+#include "ngraph/pass/manager.hpp"
 
 using namespace std;
 using namespace ngraph;
@@ -30,6 +31,13 @@ bool pass::AssignTensors::run_on_call_list(std::list<Node*>& node_list)
     {
         try
         {
+            // We need to set the nodes is_output state prior to call assign_tensors
+            // so that the output state can be passes to the constructed tensors.
+            if (node == get_state().get_function()->get_result().get())
+            {
+                node->set_is_output();
+            }
+
             node->assign_tensors();
         }
         catch (exception& e)

src/ngraph/pass/dump_sorted.cpp

@@ -18,8 +18,9 @@
 #include "ngraph/ngraph.hpp"
 #include "ngraph/util.hpp"
 
-using namespace ngraph;
 using namespace std;
+using namespace ngraph;
+using namespace ngraph::descriptor;
 
 pass::DumpSorted::DumpSorted(const string& output_file)
     : m_output_file{output_file}
@@ -35,7 +36,7 @@ bool pass::DumpSorted::run_on_call_list(list<Node*>& nodes)
         {
             out << node->get_node_id() << "(";
             vector<string> inputs;
-            for (const descriptor::Input& input : node->get_inputs())
+            for (const Input& input : node->get_inputs())
             {
                 inputs.push_back(input.get_tensor().get_name());
             }
@@ -43,29 +44,25 @@ bool pass::DumpSorted::run_on_call_list(list<Node*>& nodes)
             out << ") -> ";
 
             vector<string> outputs;
-            for (const descriptor::Output& output : node->get_outputs())
+            for (const Output& output : node->get_outputs())
             {
                 outputs.push_back(output.get_tensor().get_name());
             }
             out << join(outputs);
             out << "\n";
 
-            for (const descriptor::Tensor* tensor : node->liveness_live_list)
+
+            for (const Tensor* tensor : node->liveness_live_list)
+            {
+                out << "    L " << tensor->get_name() << "\n";
+            }
+            for (const Tensor* tensor : node->liveness_new_list)
+            {
+                out << "    N " << tensor->get_name() << "\n";
+            }
+            for (const Tensor* tensor : node->liveness_free_list)
             {
-                out << "    ";
-                if (contains(node->liveness_new_list, tensor))
-                {
-                    out << "N ";
-                }
-                else if (contains(node->liveness_free_list, tensor))
-                {
-                    out << "F ";
-                }
-                else
-                {
-                    out << "L ";
-                }
-                out << tensor->get_name() << "\n";
+                out << "    F " << tensor->get_name() << "\n";
             }
         }
     }

src/ngraph/pass/liveness.cpp

@@ -25,57 +25,56 @@
 
 using namespace std;
 using namespace ngraph;
+using namespace ngraph::descriptor;
 
 bool pass::Liveness::run_on_call_list(list<Node*>& ops)
 {
-    unordered_set<descriptor::Tensor*> currently_live;
+    unordered_set<Tensor*> currently_live;
 
     for(auto it=ops.rbegin(); it!=ops.rend(); it++)
     {
-        Node& exop = **it;
-        exop.liveness_live_list.clear();
-        exop.liveness_new_list.clear();
-        exop.liveness_free_list.clear();
-        unordered_set<descriptor::Tensor*> input_tensor_decls;
-        for (auto input_decl : exop.get_inputs())
+        Node* node = *it;
+        node->liveness_live_list.clear();
+        node->liveness_new_list.clear();
+        node->liveness_free_list.clear();
+        unordered_set<Tensor*> input_tensor_decls;
+        for (auto input_decl : node->get_inputs())
         {
-            descriptor::Tensor& tensor = input_decl.get_tensor();
+            Tensor& tensor = input_decl.get_tensor();
             if (is_temporary(tensor))
             {
                 input_tensor_decls.insert(&tensor);
             }
         }
 
-        unordered_set<descriptor::Tensor*> output_tensor_decls;
-        for (auto output_decl : exop.get_outputs())
+        unordered_set<Tensor*> output_tensor_decls;
+        for (auto output_decl : node->get_outputs())
         {
-            descriptor::Tensor& tensor = output_decl.get_tensor();
+            Tensor& tensor = output_decl.get_tensor();
             if (is_temporary(tensor))
             {
                 output_tensor_decls.insert(&tensor);
             }
         }
 
-        unordered_set<descriptor::Tensor*> free_tensor_decls;
-        unordered_set<descriptor::Tensor*> new_tensor_decls;
-        unordered_set<descriptor::Tensor*> all_tensor_decls = input_tensor_decls;
+        unordered_set<Tensor*> free_tensor_decls;
+        unordered_set<Tensor*> new_tensor_decls;
+        unordered_set<Tensor*> all_tensor_decls = input_tensor_decls;
+        all_tensor_decls.insert(output_tensor_decls.begin(), output_tensor_decls.end());
 
-        for (auto decls : {input_tensor_decls, output_tensor_decls})
+        for (Tensor* tensor_decl : all_tensor_decls)
         {
-            for (descriptor::Tensor* tensor_decl : decls)
+            if (!contains(currently_live, tensor_decl))
             {
-                if (!contains(currently_live, tensor_decl))
-                {
-                    // this is the last node that value is seen in
-                    // delete it at the end of the op
-                    currently_live.insert(tensor_decl);
-                    free_tensor_decls.insert(tensor_decl);
-                }
+                // this is the last node that value is seen in
+                // delete it at the end of the op
+                currently_live.insert(tensor_decl);
+                free_tensor_decls.insert(tensor_decl);
             }
         }
 
-        exop.liveness_live_list = currently_live;
-        for (descriptor::Tensor* output_decl : output_tensor_decls)
+        node->liveness_live_list = currently_live;
+        for (Tensor* output_decl : output_tensor_decls)
         {
             if (contains(currently_live, output_decl))
             {
@@ -83,33 +82,34 @@ bool pass::Liveness::run_on_call_list(list<Node*>& ops)
                 currently_live.erase(output_decl);
             }
         }
-        exop.liveness_free_list = free_tensor_decls;
-        exop.liveness_new_list = new_tensor_decls;
+        node->liveness_free_list = free_tensor_decls;
+        node->liveness_new_list = new_tensor_decls;
     }
 
     // Anything marked as output must remain live for the remainder of the graph
     // Add outputs to live_list and remove from free_list
-    unordered_set<descriptor::Tensor*> outputs;
-    unordered_set<descriptor::Tensor*> seen;
-    for (Node* exop : ops)
+    unordered_set<Tensor*> outputs;
+    unordered_set<Tensor*> seen;
+    for (Node* node : ops)
     {
-        for (descriptor::Tensor* tensor : exop->liveness_live_list)
+        for (Tensor* tensor : node->liveness_live_list)
         {
             if (tensor->is_output())
             {
                 outputs.insert(tensor);
             }
         }
-        for (descriptor::Tensor* tensor : outputs)
+        for (Tensor* tensor : outputs)
         {
-            exop->liveness_live_list.insert(tensor);
-            exop->liveness_free_list.erase(tensor);
+            NGRAPH_INFO << "found output";
+            node->liveness_live_list.insert(tensor);
+            node->liveness_free_list.erase(tensor);
 
-            if (contains(exop->liveness_new_list, tensor))
+            if (contains(node->liveness_new_list, tensor))
             {
                 if (contains(seen, tensor))
                 {
-                    exop->liveness_new_list.erase(tensor);
+                    node->liveness_new_list.erase(tensor);
                 }
                 else
                 {
@@ -119,7 +119,7 @@ bool pass::Liveness::run_on_call_list(list<Node*>& ops)
         }
     }
 
-    validate_liveness(ops);
+    // validate_liveness(ops);
     return false;
 }
 
@@ -141,11 +141,12 @@ void pass::Liveness::check_dependencies(
     }
 }
 
-bool pass::Liveness::is_temporary(const descriptor::Tensor& tensor)
+bool pass::Liveness::is_temporary(const Tensor& tensor)
 {
     return
         tensor.is_persistent() == false
         && tensor.is_input() == false
+        && tensor.is_output() == false
         ;
         // && tensor.is_constant() == false
         // && tensor.is_compile_only() == false;
@@ -153,20 +154,20 @@ bool pass::Liveness::is_temporary(const descriptor::Tensor& tensor)
 
 void pass::Liveness::validate_liveness(const list<Node*>& ops)
 {
-    unordered_set<descriptor::Tensor*> dead_tensors;
-    for (const Node* exop : ops)
+    unordered_set<Tensor*> dead_tensors;
+    for (const Node* node : ops)
     {
-        auto active = exop->liveness_live_list;
-        active.insert(exop->liveness_new_list.begin(), exop->liveness_new_list.end());
-        active.insert(exop->liveness_free_list.begin(), exop->liveness_free_list.end());
-        for (const descriptor::Tensor* tensor : active)
+        auto active = node->liveness_live_list;
+        active.insert(node->liveness_new_list.begin(), node->liveness_new_list.end());
+        active.insert(node->liveness_free_list.begin(), node->liveness_free_list.end());
+        for (const Tensor* tensor : active)
         {
             if (contains(dead_tensors, tensor))
             {
                 throw runtime_error("Liveness: Dead tensors intersect active tensors");
             }
         }
-        dead_tensors.insert(exop->liveness_free_list.begin(), exop->liveness_free_list.end());
+        dead_tensors.insert(node->liveness_free_list.begin(), node->liveness_free_list.end());
     }
 }
 

src/ngraph/pass/manager.cpp

@@ -18,8 +18,40 @@
 #include "ngraph/log.hpp"
 #include "ngraph/pass/manager.hpp"
 #include "ngraph/node.hpp"
+#include "ngraph/function.hpp"
 
 using namespace std;
+using namespace ngraph;
+
+Function* ngraph::pass::ManagerState::get_function()
+{
+    return m_function;
+}
+
+void ngraph::pass::ManagerState::set_function(Function* func)
+{
+    m_function = func;
+}
+
+size_t ngraph::pass::ManagerState::get_temporary_pool_size()
+{
+    return m_temporary_pool_size;
+}
+
+void ngraph::pass::ManagerState::set_temporary_pool_size(size_t size)
+{
+    m_temporary_pool_size = size;
+}
+
+std::list<Node*>& ngraph::pass::ManagerState::get_call_graph()
+{
+    return m_call_graph;
+}
+
+const std::list<Node*>& ngraph::pass::ManagerState::get_call_graph() const
+{
+    return m_call_graph;
+}
 
 ngraph::pass::Manager::Manager()
 {
@@ -53,24 +85,33 @@ void ngraph::pass::Manager::register_pass(std::shared_ptr<CallBase> p)
     m_call_passes.push_back(p);
 }
 
-void ngraph::pass::Manager::run_passes(std::shared_ptr<Node> nodes)
+void ngraph::pass::Manager::run_passes(shared_ptr<Function> func)
 {
+    run_passes(func.get());
+}
+
+void ngraph::pass::Manager::run_passes(Function* func)
+{
+    m_state.set_function(func);
     for (shared_ptr<TreeBase> p : m_tree_passes)
     {
-        p->run_on_tree(nodes);
-        if (p->call_graph_produced())
-        {
-            m_sorted_list = p->get_call_graph();
-        }
+        p->set_state(get_state());
+        p->run_on_tree(func->get_result());
     }
 
     for (shared_ptr<CallBase>& p : m_call_passes)
     {
-        p->run_on_call_list(m_sorted_list);
+        p->set_state(get_state());
+        p->run_on_call_list(get_state().get_call_graph());
     }
 }
 
-const std::list<ngraph::Node*>& ngraph::pass::Manager::get_sorted_list() const
+const std::list<ngraph::Node*>& ngraph::pass::Manager::get_call_graph() const
+{
+    return m_state.get_call_graph();
+}
+
+ngraph::pass::ManagerState& ngraph::pass::Manager::get_state()
 {
-    return m_sorted_list;
+    return m_state;
 }

src/ngraph/pass/manager.hpp

@@ -15,6 +15,8 @@
 #pragma once
 
 #include <vector>
+#include <memory>
+#include <list>
 
 #include "ngraph/pass/call_pass.hpp"
 #include "ngraph/pass/tree_pass.hpp"
@@ -24,11 +26,31 @@ namespace ngraph
     namespace pass
     {
         class Manager;
+        class ManagerState;
     }
 
     class Node;
+    class Function;
 }
 
+class ngraph::pass::ManagerState
+{
+public:
+    Function* get_function();
+    void      set_function(Function*);
+
+    size_t get_temporary_pool_size();
+    void   set_temporary_pool_size(size_t);
+
+    std::list<Node*>& get_call_graph();
+    const std::list<Node*>& get_call_graph() const;
+
+private:
+    Function*        m_function = nullptr;
+    size_t           m_temporary_pool_size = 0;
+    std::list<Node*> m_call_graph;
+};
+
 class ngraph::pass::Manager
 {
 public:
@@ -40,12 +62,15 @@ public:
     void register_pass(std::shared_ptr<TreeBase>);
     void register_pass(std::shared_ptr<CallBase>);
 
-    void run_passes(std::shared_ptr<Node> nodes);
+    void run_passes(Function*);
+    void run_passes(std::shared_ptr<Function>);
+
+    const std::list<Node*>& get_call_graph() const;
 
-    const std::list<Node*>& get_sorted_list() const;
+    ManagerState& get_state();
 
 private:
     std::vector<std::shared_ptr<TreeBase>> m_tree_passes;
     std::vector<std::shared_ptr<CallBase>> m_call_passes;
-    std::list<Node*>                       m_sorted_list;
+    ManagerState                           m_state;
 };

src/ngraph/pass/memory_layout.cpp

@@ -18,17 +18,32 @@
 #include "ngraph/log.hpp"
 #include "ngraph/ngraph.hpp"
 #include "ngraph/pass/liveness.hpp"
+#include "ngraph/pass/manager.hpp"
 #include "ngraph/pass/memory_layout.hpp"
 #include "ngraph/log.hpp"
+#include "ngraph/util.hpp"
 
 using namespace std;
 using namespace ngraph;
+using namespace ngraph::descriptor;
 
 bool pass::MemoryLayout::run_on_call_list(std::list<Node*>& node_list)
 {
-    for (Node* node : node_list)
+    MemoryManager mm;
+    for (const Node* node : node_list)
     {
+        for (Tensor* tensor : node->liveness_new_list)
+        {
+            size_t offset = mm.allocate(tensor->size());
+            tensor->set_pool_offset(offset);
+        }
+        for (const Tensor* tensor : node->liveness_free_list)
+        {
+            mm.free(tensor->get_pool_offset());
+        }
     }
+    get_state().set_temporary_pool_size(mm.max_allocated());
+
     return false;
 }
 

src/ngraph/pass/memory_visualize.cpp

+// ----------------------------------------------------------------------------
+// Copyright 2017 Nervana Systems Inc.
+// 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
+// ----------------------------------------------------------------------------
+
+#include <fstream>
+#include <unordered_set>
+#include <unordered_map>
+#include <algorithm>
+
+#include "memory_visualize.hpp"
+#include "ngraph/descriptor/tensor.hpp"
+#include "ngraph/node.hpp"
+#include "ngraph/util.hpp"
+
+using namespace std;
+using namespace ngraph;
+using namespace ngraph::descriptor;
+
+pass::MemoryVisualize::MemoryVisualize(const string& filename)
+    : m_filename{filename}
+{
+}
+
+bool pass::MemoryVisualize::run_on_call_list(list<Node*>& _nodes)
+{
+    const list<Node*> nodes = _nodes;
+    ofstream file(m_filename);
+    {
+        file << "<!DOCTYPE html>\n<html>\n";
+        file << "<head>\n";
+        file << "    <style>\n";
+        file << "        th, td {\n";
+        file << "            border-bottom: 1px solid #ddd;\n";
+        file << "            width: 200px;\n";
+        file << "        }\n";
+        file << "        table, td, th {\n";
+        // file << "            border: 1px solid #ddd;\n";
+        // file << "            text-align: left;\n";
+        file << "        }\n";
+        file << "        table {\n";
+        file << "            border-collapse: collapse;\n";
+        // file << "            width: 100%;\n";
+        file << "        }\n";
+        // file << "        tr:hover {background-color: #f5f5f5}\n";
+        file << "        tr:nth-child(even) {background-color: #f2f2f2}\n";
+        file << "    </style>\n";
+        file << "</head>\n";
+
+        file << "<body>\n";
+        unordered_set<descriptor::Tensor*> tensors;
+        size_t temp_max_size = 0;
+        for (Node* node : nodes)
+        {
+            tensors.insert(node->liveness_live_list.begin(), node->liveness_live_list.end());
+        }
+        for (descriptor::Tensor* tensor : tensors)
+        {
+            if (tensor->is_persistent() == false)
+            {
+                temp_max_size += tensor->size();
+            }
+        }
+
+        // file << "<table>\n";
+        // file << "<tr><td>Persistent Memory Footprint</td><td align=\"right\">";
+        // file << computation_decl.exop_block.persistent_size() << "</td></tr>\n";
+        // file << "<tr><td>Temporary Memory Footprint</td><td align=\"right\">";
+        // file << computation_decl.exop_block.memory_footprint() << "</td></tr>\n";
+        // file << "<tr><td>Max temporary Memory Footprint</td><td align=\"right\">";
+        // file << temp_max_size << "</td></tr>\n";
+        // file << "</table>\n";
+
+        file << "<hr>\n";
+        draw_tensor_weight(file, nodes);
+        // file << "<hr>\n";
+        // draw_op_influence(file);
+        file << "<hr>\n";
+        draw_histogram(file, nodes);
+        // file << "<hr>\n";
+        file << "</body>\n</html>\n";
+    }
+    return false;
+}
+
+void pass::MemoryVisualize::check_dependencies(const vector<shared_ptr<CallBase>>& deps) const
+{
+}
+
+const Node* pass::MemoryVisualize::find_largest_op(const list<Node*>& nodes)
+{
+    const Node* largest_op = nullptr;
+    size_t largest_size = 0;
+    for (const Node* exop : nodes)
+    {
+        size_t size = 0;
+        for (const Tensor* tensor : exop->liveness_live_list)
+        {
+            size += tensor->size();
+        }
+        if (size > largest_size)
+        {
+            largest_size = size;
+            largest_op = exop;
+        }
+    }
+    return largest_op;
+}
+
+void pass::MemoryVisualize::draw_tensor_weight(ostream& file, const list<Node*>& nodes)
+{
+    const Node* largest_op = find_largest_op(nodes);
+
+    if (largest_op)
+    {
+        unordered_set<Tensor*> largest_live;
+        for (Tensor* tensor : largest_op->liveness_live_list)
+        {
+            largest_live.insert(tensor);
+        }
+
+        unordered_map<const Tensor*, size_t> age_list;
+        vector<const Tensor*> tensor_set;
+        unordered_map<const Tensor*, const Node*> generator_op;
+        file << "<table>\n";
+        file << "    <tr>";
+        file << "<th align=\"left\">tensor</th>";
+        file << "<th align=\"right\">size</th>";
+        file << "<th align=\"right\">age</th>";
+        file << "<th align=\"right\">generator weight</th>";
+        file << "</tr>\n";
+        size_t i = 0;
+        for (const Node* exop : nodes)
+        {
+            for (const Tensor* tensor : exop->liveness_new_list)
+            {
+                age_list[tensor] = i;
+                generator_op[tensor] = exop;
+            }
+            for (const Tensor* tensor : exop->liveness_free_list)
+            {
+                size_t start = age_list[tensor];
+                age_list[tensor] = (i - start);
+                tensor_set.push_back(tensor);
+            }
+            i++;
+        }
+        sort(tensor_set.begin(), tensor_set.end(), [](const Tensor* t1, const Tensor* t2)
+        {
+            return t1->size() < t2->size();
+        });
+        for (const Tensor* tensor : tensor_set)
+        {
+            int generator_weight = compute_op_weight(generator_op[tensor]);
+            if (contains(largest_live, tensor))
+            {
+                file << "    <tr style=\"background-color: #f0c0f0\">";
+            }
+            else
+            {
+                file << "    <tr>";
+            }
+            file << "<td>" << tensor->get_name() << "</td>";
+            file << "<td align=\"right\">" << tensor->size() << "</td>";
+            file << "<td align=\"right\">" << age_list[tensor] << "</td>";
+            file << "<td align=\"right\">" << generator_weight << "/td>";
+            file << "</tr>\n";
+        }
+
+        file << "</table>\n";
+    }
+}
+
+void pass::MemoryVisualize::draw_histogram(ostream& file, const list<Node*>& nodes)
+{
+    size_t stroke_width = 14;
+    size_t text_offset = 4;
+    size_t offset = 200;
+    size_t width = 1000;
+    size_t scale = width - offset;
+    size_t line_spacing = stroke_width * 1.5;
+    size_t line_count = 0;
+    for (const Node* node : nodes)
+    {
+        (void)node;
+        line_count += 1;
+    }
+    size_t height = line_count * line_spacing + stroke_width;
+    size_t memory_footprint = max<size_t>(1, MemoryVisualize::memory_footprint(nodes));
+
+    file << "<svg viewBox=\"0 0 " << width << " " << height << "\">\n";
+    size_t y = 0;
+    for (const Node* node : nodes)
+    {
+        float usage = float(MemoryVisualize::memory_usage(node));
+        float footprint = float(MemoryVisualize::memory_footprint(node));
+        y += line_spacing;
+        size_t x1 = offset;
+        size_t x2 = ((usage / memory_footprint) * scale) + offset;
+        file << "<text x=\"" << 0 << "\" y=\"" << y + text_offset << "\" fill=\"" << "black" << "\">" << node->get_node_id() << "</text>\n";
+        file << "<line x1=\"" << x1 << "\" y1=\"" << y << "\" x2=\"" << x2 << "\" y2=\"" << y << "\"";
+        file << " style=\"stroke:forestgreen;stroke-width:" << stroke_width << "\" />\n";
+        x1 = x2;
+        x2 = ((footprint / memory_footprint) * scale) + offset;
+        file << "<line x1=\"" << x1 << "\" y1=\"" << y << "\" x2=\"" << x2 << "\" y2=\"" << y << "\"";
+        file << " style=\"stroke:firebrick;stroke-width:" << stroke_width << "\" />\n";
+    }
+    file << "</svg>\n";
+}
+
+void pass::MemoryVisualize::draw_op_influence(ostream& file, const list<Node*>& nodes)
+{
+    file << "<table>\n";
+    file << "    <tr>";
+    file << "<th align=\"left\">op</th>";
+    file << "<th align=\"right\">influence</th>";
+    file << "</tr>\n";
+    for (const Node* exop : nodes)
+    {
+        int weight = compute_op_weight(exop);
+        file << "    <tr>";
+        file << "<td>" << exop->get_node_id() << "</td>";
+        file << "<td align=\"right\">" << weight << "</td>";
+        file << "</tr>\n";
+    }
+}
+
+int pass::MemoryVisualize::compute_op_weight(const Node* exop)
+{
+    int mass = 0;
+    // for input_decl in exop.input_decls:
+    //     tensor = input_decl.source_output_decl.tensor
+    //     if tensor.is_persistent is False:
+    //         mass += tensor->size()
+    // for output_decl in exop.output_decls:
+    //     tensor = output_decl.tensor
+    //     if tensor.is_persistent is False:
+    //         mass -= tensor->size()
+    for (const Tensor* tensor : exop->liveness_new_list)
+    {
+        if (tensor->is_persistent() == false)
+        {
+            mass += tensor->size();
+        }
+    }
+    for (const Tensor* tensor : exop->liveness_free_list)
+    {
+        if (tensor->is_persistent() == false)
+        {
+            mass -= tensor->size();
+        }
+    }
+    return mass;
+}
+
+size_t pass::MemoryVisualize::memory_usage(const Node* node)
+{
+    return 0;
+}
+
+size_t pass::MemoryVisualize::memory_footprint(const Node* node)
+{
+    return 0;
+}
+
+size_t pass::MemoryVisualize::memory_footprint(const std::list<Node*>& nodes)
+{
+    return 0;
+}

src/ngraph/pass/memory_visualize.hpp

+// ----------------------------------------------------------------------------
+// Copyright 2017 Nervana Systems Inc.
+// 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
+// ----------------------------------------------------------------------------
+
+#pragma once
+
+#include <limits>
+#include <list>
+#include <iostream>
+
+#include "ngraph/pass/call_pass.hpp"
+
+namespace ngraph
+{
+    namespace pass
+    {
+        class MemoryVisualize;
+    }
+    class Node;
+}
+
+class ngraph::pass::MemoryVisualize : public CallBase
+{
+public:
+    MemoryVisualize(const std::string& filename);
+    virtual bool run_on_call_list(std::list<Node*>&) override;
+
+    void check_dependencies(const std::vector<std::shared_ptr<CallBase>>&) const override;
+
+private:
+    const Node* find_largest_op(const std::list<Node*>& nodes);
+    void draw_tensor_weight(std::ostream& file, const std::list<Node*>& nodes);
+    void draw_histogram(std::ostream& file, const std::list<Node*>& nodes);
+    void draw_op_influence(std::ostream& file, const std::list<Node*>& nodes);
+    int compute_op_weight(const Node* exop);
+
+    static size_t memory_usage(const Node*);
+    static size_t memory_footprint(const Node*);
+    static size_t memory_footprint(const std::list<Node*>&);
+    
+    const std::string m_filename;
+};

src/ngraph/pass/pass.cpp

@@ -13,3 +13,14 @@
 // ----------------------------------------------------------------------------
 
 #include "ngraph/pass/pass.hpp"
+#include "ngraph/pass/manager.hpp"
+
+ngraph::pass::ManagerState& ngraph::pass::Base::get_state()
+{
+    return *m_state;
+}
+
+void ngraph::pass::Base::set_state(ManagerState& state)
+{
+    m_state = &state;
+}

src/ngraph/pass/pass.hpp

@@ -19,11 +19,19 @@ namespace ngraph
     namespace pass
     {
         class Base;
+        class Manager;
+        class ManagerState;
     }
 }
 
 class ngraph::pass::Base
 {
+    friend class Manager;
 public:
+protected:
+    ManagerState& get_state();
+    void set_state(ManagerState&);
+
 private:
+    ManagerState* m_state;
 };

src/ngraph/pass/topological_sort.cpp

@@ -17,6 +17,7 @@
 
 #include "ngraph/log.hpp"
 #include "ngraph/node.hpp"
+#include "ngraph/pass/manager.hpp"
 #include "ngraph/pass/topological_sort.hpp"
 #include "ngraph/util.hpp"
 
@@ -25,6 +26,8 @@ using namespace std;
 
 bool ngraph::pass::TopologicalSort::run_on_tree(std::shared_ptr<Node> p)
 {
+    list<Node*>& sorted_list = get_state().get_call_graph();
+    sorted_list.clear();
     deque<Node*> independent_nodes;
     unordered_map<Node*, size_t> node_depencency_count;
 
@@ -39,7 +42,7 @@ bool ngraph::pass::TopologicalSort::run_on_tree(std::shared_ptr<Node> p)
     while (independent_nodes.size() > 0)
     {
         auto independent_node = independent_nodes.front();
-        m_sorted_list.push_back(independent_node);
+        sorted_list.push_back(independent_node);
         independent_nodes.pop_front();
 
         for (auto user : independent_node->users())
@@ -55,8 +58,3 @@ bool ngraph::pass::TopologicalSort::run_on_tree(std::shared_ptr<Node> p)
 
     return false;
 }
-
-std::list<Node*> ngraph::pass::TopologicalSort::get_call_graph() const
-{
-    return m_sorted_list;
-}

src/ngraph/pass/topological_sort.hpp

@@ -33,10 +33,4 @@ class ngraph::pass::TopologicalSort : public TreeBase
 public:
     TopologicalSort() {}
     bool run_on_tree(std::shared_ptr<Node>) override;
-
-    bool             call_graph_produced() const override { return true; }
-    std::list<Node*> get_call_graph() const override;
-
-private:
-    std::list<Node*> m_sorted_list;
 };

src/ngraph/pass/tree_pass.hpp

@@ -37,10 +37,6 @@ public:
     // return true if changes were made to the tree
     virtual bool run_on_tree(std::shared_ptr<Node>) = 0;
 
-    virtual bool             call_graph_produced() const { return false; }
-    virtual std::list<Node*> get_call_graph() const { return std::list<Node*>(); }
     // derived class throws exception if its dependencies have not been met
     virtual void check_dependencies(const std::vector<std::shared_ptr<TreeBase>>&) const {}
-private:
-    std::list<Node*> m_sorted_list;
 };

src/ngraph/runtime/external_function.cpp

@@ -24,6 +24,7 @@
 #include "ngraph/node.hpp"
 #include "ngraph/ops/add.hpp"
 #include "ngraph/ops/multiply.hpp"
+#include "ngraph/pass/manager.hpp"
 #include "ngraph/pass/topological_sort.hpp"
 #include "ngraph/runtime/eigen/add.hpp"
 #include "ngraph/runtime/external_function.hpp"
@@ -97,9 +98,11 @@ void ExternalFunction::compile()
 
     // This will be replaced with the pass manager
     // Get the ordered list of ops in execution order
-    pass::TopologicalSort ts;
-    ts.run_on_tree(m_function->get_result());
-    auto nodes = ts.get_call_graph();
+    pass::Manager pass_manager;
+    auto          topological_sort = make_shared<pass::TopologicalSort>();
+    pass_manager.register_pass(topological_sort);
+    pass_manager.run_passes(m_function);
+    auto nodes = pass_manager.get_call_graph();
     // Types
     for (auto node : nodes)
     {

test/CMakeLists.txt

@@ -43,7 +43,7 @@ set (SRC
 )
 
 if(MKLDNN_INCLUDE_DIR)
-    include_directories(${MKLDNN_INCLUDE_DIR})
+    include_directories(SYSTEM ${MKLDNN_INCLUDE_DIR})
     link_directories(${MKLDNN_LIB_DIR})
     set(SRC ${SRC} mkldnn.cpp)
 endif()

test/pass_liveness.cpp

@@ -55,10 +55,9 @@ TEST(pass, liveness)
     pass_manager.register_pass(liveness);
     pass_manager.register_pass(dump_sorted);
 
-    auto   graph      = make_test_graph();
-    size_t node_count = get_node_count(graph);
-    pass_manager.run_passes(graph);
-    auto sorted = pass_manager.get_sorted_list();
+    shared_ptr<Function>   func      = make_test_graph();
+    pass_manager.run_passes(func.get());
+    auto sorted = pass_manager.get_call_graph();
 
     // for (const Node* node : sorted)
     // {

test/pass_manager.cpp

@@ -41,9 +41,9 @@ TEST(pass_manager, add)
     pass_manager.register_pass(assign_tensors);
 
     auto   graph      = make_test_graph();
-    size_t node_count = get_node_count(graph);
-    pass_manager.run_passes(graph);
-    auto sorted = pass_manager.get_sorted_list();
+    size_t node_count = get_node_count(graph->get_result());
+    pass_manager.run_passes(graph.get());
+    auto sorted = pass_manager.get_call_graph();
     EXPECT_EQ(node_count, sorted.size());
     EXPECT_TRUE(validate_list(sorted));
 }

test/pass_memory_layout.cpp

@@ -20,6 +20,14 @@
 #include "gtest/gtest.h"
 
 #include "ngraph/ngraph.hpp"
+#include "ngraph/pass/liveness.hpp"
+#include "ngraph/pass/assign_tensors.hpp"
+#include "ngraph/pass/manager.hpp"
+#include "ngraph/pass/propagate_types.hpp"
+#include "ngraph/pass/topological_sort.hpp"
+#include "ngraph/pass/liveness.hpp"
+#include "ngraph/pass/visualize_tree.hpp"
+#include "ngraph/pass/dump_sorted.hpp"
 #include "ngraph/pass/memory_layout.hpp"
 #include "test_tools.hpp"
 
@@ -196,3 +204,28 @@ TEST(memory_manager, memory_align)
     EXPECT_EQ(64, mm.allocate(4));
     EXPECT_EQ(128, mm.allocate(4));
 }
+
+TEST(memory_layout, basic)
+{
+    string dump_file = "memory_layout.txt";
+    pass::Manager pass_manager;
+    auto          topological_sort = make_shared<pass::TopologicalSort>();
+    auto          propagate_types  = make_shared<pass::PropagateTypes>();
+    auto          assign_tensors   = make_shared<pass::AssignTensors>();
+    auto          liveness         = make_shared<pass::Liveness>();
+    auto          memory_layout    = make_shared<pass::MemoryLayout>();
+    auto          dump_sorted      = make_shared<pass::DumpSorted>(dump_file);
+
+    pass_manager.register_pass(topological_sort);
+    pass_manager.register_pass(propagate_types);
+    pass_manager.register_pass(assign_tensors);
+    pass_manager.register_pass(liveness);
+    pass_manager.register_pass(memory_layout);
+    pass_manager.register_pass(dump_sorted);
+
+    auto graph = make_test_graph();
+    pass_manager.run_passes(graph);
+    auto sorted = pass_manager.get_call_graph();
+    size_t temporary_pool_size = pass_manager.get_state().get_temporary_pool_size();
+    EXPECT_EQ(12, temporary_pool_size);
+}

test/tensor.cpp

@@ -16,27 +16,71 @@
 #include <sstream>
 #include <string>
 #include <vector>
+#include <memory>
 
 #include "gtest/gtest.h"
+#include "ngraph/ngraph.hpp"
+#include "ngraph/pass/assign_tensors.hpp"
+#include "ngraph/pass/manager.hpp"
+#include "ngraph/pass/propagate_types.hpp"
+#include "ngraph/pass/liveness.hpp"
+#include "ngraph/pass/topological_sort.hpp"
+#include "ngraph/function.hpp"
+#include "test_tools.hpp"
 
 using namespace std;
+using namespace ngraph;
+using namespace ngraph::descriptor;
 
-TEST(tensor, test)
+TEST(tensor, size)
 {
-    class test
+    pass::Manager pass_manager;
+    auto          topological_sort = make_shared<pass::TopologicalSort>();
+    auto          propagate_types  = make_shared<pass::PropagateTypes>();
+    auto          assign_tensors   = make_shared<pass::AssignTensors>();
+    auto          liveness         = make_shared<pass::Liveness>();
+
+    pass_manager.register_pass(topological_sort);
+    pass_manager.register_pass(propagate_types);
+    pass_manager.register_pass(assign_tensors);
+    pass_manager.register_pass(liveness);
+
+    {
+        auto arg0 = make_shared<op::Parameter>(element::Float32::element_type(), Shape{2, 3});
+        auto add  = make_shared<op::Add>(arg0, arg0);
+        auto f0   = make_shared<Function>(add, op::Parameters{arg0});
+
+        pass_manager.run_passes(f0);
+
+        auto outputs = arg0->get_outputs();
+        ASSERT_EQ(1, outputs.size());
+        Tensor& output = outputs[0].get_tensor();
+        EXPECT_EQ(2 * 3 * 4, output.size());
+    }
+
     {
-    public:
-        test(int i)
-            : value{i}
-        {
-        }
+        auto arg0 = make_shared<op::Parameter>(element::Float32::element_type(), Shape{});
+        auto add  = make_shared<op::Add>(arg0, arg0);
+        auto f0   = make_shared<Function>(add, op::Parameters{arg0});
+
+        pass_manager.run_passes(f0);
 
-        int value;
-    };
-    map<int, shared_ptr<test>> test_map;
+        auto outputs = arg0->get_outputs();
+        ASSERT_EQ(1, outputs.size());
+        Tensor& output = outputs[0].get_tensor();
+        EXPECT_EQ(1 * 4, output.size());
+    }
+
+    {
+        auto arg0 = make_shared<op::Parameter>(element::Float32::element_type(), Shape{1});
+        auto add  = make_shared<op::Add>(arg0, arg0);
+        auto f0   = make_shared<Function>(add, op::Parameters{arg0});
 
-    test_map[1] = make_shared<test>(2);
+        pass_manager.run_passes(f0);
 
-    EXPECT_NE(nullptr, test_map[1]);
-    EXPECT_EQ(nullptr, test_map[2]);
+        auto outputs = arg0->get_outputs();
+        ASSERT_EQ(1, outputs.size());
+        Tensor& output = outputs[0].get_tensor();
+        EXPECT_EQ(1 * 4, output.size());
+    }
 }

test/test_tools.cpp

@@ -53,7 +53,7 @@ bool validate_list(const list<Node*>& nodes)
     return rc;
 }
 
-shared_ptr<Node> make_test_graph()
+shared_ptr<Function> make_test_graph()
 {
     auto arg_0 = make_shared<op::Parameter>(element::Float32::element_type(), Shape{});
     auto arg_1 = make_shared<op::Parameter>(element::Float32::element_type(), Shape{});
@@ -71,7 +71,9 @@ shared_ptr<Node> make_test_graph()
 
     auto r0 = make_shared<op::Add>(t3, t4);
 
-    return r0;
+    auto f0 = make_shared<Function>(r0, op::Parameters{arg_0, arg_1, arg_2, arg_3, arg_4, arg_5});
+
+    return f0;
 }
 
 size_t get_node_count(std::shared_ptr<Node> n)

test/test_tools.hpp

@@ -20,8 +20,9 @@
 namespace ngraph
 {
     class Node;
+    class Function;
 }
 
 bool validate_list(const std::list<ngraph::Node*>& nodes);
-std::shared_ptr<ngraph::Node> make_test_graph();
+std::shared_ptr<ngraph::Function> make_test_graph();
 size_t get_node_count(std::shared_ptr<ngraph::Node> n);

test/topological_sort.cpp

@@ -20,6 +20,7 @@
 #include "gtest/gtest.h"
 
 #include "ngraph/ngraph.hpp"
+#include "ngraph/pass/manager.hpp"
 #include "ngraph/pass/topological_sort.hpp"
 #include "ngraph/visualize.hpp"
 #include "ngraph/util.hpp"
@@ -62,9 +63,11 @@ TEST(topological_sort, basic)
     // Visualize vz;
     // vz.add(r0);
     // vz.save_dot("test.png");
-    pass::TopologicalSort ts;
-    ts.run_on_tree(r0);
-    auto sorted_list = ts.get_call_graph();
+    pass::Manager pass_manager;
+    auto          topological_sort = make_shared<pass::TopologicalSort>();
+    pass_manager.register_pass(topological_sort);
+    pass_manager.run_passes(f0);
+    auto sorted_list = pass_manager.get_call_graph();
 
     size_t node_count = get_node_count(r0);
 
@@ -99,18 +102,24 @@ TEST(benchmark, topological_sort)
     stopwatch timer;
     // x[i+1] = tanh(dot(W,x[i])+b)
     shared_ptr<Node> result;
+    vector<shared_ptr<op::Parameter>> args;
     result = make_shared<op::Parameter>(element::Float32::element_type(), Shape{1});
     for (int i=0; i<1000000; i++)
     {
         auto in_1 = make_shared<op::Parameter>(element::Float32::element_type(), Shape{1});
         auto in_2 = make_shared<op::Parameter>(element::Float32::element_type(), Shape{1});
+        args.push_back(in_1);
+        args.push_back(in_2);
         result = make_cell(result, in_1, in_2);
     }
+    auto f0 = make_shared<Function>(result, args);
 
     timer.start();
-    pass::TopologicalSort ts;
-    ts.run_on_tree(result);
-    auto sorted_list = ts.get_call_graph();
+    pass::Manager pass_manager;
+    auto          topological_sort = make_shared<pass::TopologicalSort>();
+    pass_manager.register_pass(topological_sort);
+    pass_manager.run_passes(f0);
+    auto sorted_list = pass_manager.get_call_graph();
     timer.stop();
     NGRAPH_INFO << "topological sort took " << timer.get_milliseconds() << "ms";