Memory Leak with External Function (#568)

* add unit test for resource deallocation

fix leak

* cleanup

src/ngraph/runtime/cpu/cpu_external_function.cpp

@@ -249,6 +249,10 @@ runtime::cpu::CPU_ExternalFunction::CPU_ExternalFunction(
 {
 }
 
+runtime::cpu::CPU_ExternalFunction::~CPU_ExternalFunction()
+{
+}
+
 void runtime::cpu::CPU_ExternalFunction::compile()
 {
     if (m_is_compiled)
@@ -256,14 +260,16 @@ void runtime::cpu::CPU_ExternalFunction::compile()
         return;
     }
 
-    m_mkldnn_emitter.reset(new MKLDNNEmitter(shared_from_this()));
+    string function_name = m_function->get_name();
+
+    m_mkldnn_emitter.reset(new MKLDNNEmitter());
 
     ngraph::pass::Manager pass_manager;
 
     pass_manager.register_pass<ngraph::pass::CoreFusion>();
     pass_manager.register_pass<runtime::cpu::pass::CPUFusion>();
-    pass_manager.register_pass<runtime::cpu::pass::CPUAssignment>(shared_from_this());
-    pass_manager.register_pass<runtime::cpu::pass::CPULayout>(shared_from_this());
+    pass_manager.register_pass<runtime::cpu::pass::CPUAssignment>(this);
+    pass_manager.register_pass<runtime::cpu::pass::CPULayout>(this);
     pass_manager.register_pass<ngraph::pass::Liveness>();
     pass_manager.register_pass<ngraph::pass::MemoryLayout>(s_memory_pool_alignment);
 

src/ngraph/runtime/cpu/cpu_external_function.hpp

@@ -75,6 +75,7 @@ namespace ngraph
             public:
                 CPU_ExternalFunction(const std::shared_ptr<ngraph::Function>& function,
                                      bool release_function = true);
+                ~CPU_ExternalFunction();
                 std::shared_ptr<ngraph::runtime::CallFrame> make_call_frame();
 
                 const LayoutDescriptorPtrs& get_parameter_layout_descriptors();

src/ngraph/runtime/cpu/cpu_manager.cpp

@@ -30,7 +30,8 @@ std::shared_ptr<ngraph::runtime::Backend> runtime::cpu::CPU_Manager::allocate_ba
 std::shared_ptr<ngraph::runtime::ExternalFunction>
     runtime::cpu::CPU_Manager::compile(const std::shared_ptr<ngraph::Function>& fun)
 {
-    return std::make_shared<CPU_ExternalFunction>(fun);
+    auto rc = std::make_shared<CPU_ExternalFunction>(fun);
+    return rc;
 }
 
 ngraph::runtime::Manager::Factory runtime::cpu::CPU_Manager::factory =

src/ngraph/runtime/cpu/mkldnn_emitter.cpp

@@ -27,18 +27,18 @@ using namespace ngraph::runtime::cpu;
 
 const std::vector<mkldnn::primitive*>& MKLDNNEmitter::get_mkldnn_primitives() const
 {
-    return mkldnn_primitives;
+    return m_mkldnn_primitives;
 }
 
 size_t MKLDNNEmitter::insert_primitive(mkldnn::primitive* primitive)
 {
-    mkldnn_primitives.emplace_back(primitive);
-    return (mkldnn_primitives.size() - 1);
+    m_mkldnn_primitives.emplace_back(primitive);
+    return (m_mkldnn_primitives.size() - 1);
 }
 
 const std::vector<size_t>& MKLDNNEmitter::get_primitive_deps(size_t index) const
 {
-    return primitive_deps.at(index);
+    return m_primitive_deps.at(index);
 }
 
 mkldnn::memory::desc MKLDNNEmitter::build_memory_descriptor(const TensorViewWrapper& tvw,
@@ -96,11 +96,11 @@ size_t MKLDNNEmitter::build_convolution_forward(const mkldnn::memory::desc& inpu
           mkldnn::memory::dims(padding_above.begin(), padding_above.end()),
           mkldnn::padding_kind::zero},
          mkldnn_utils::global_cpu_engine},
-        *mkldnn_primitives[input_data_index],
-        *mkldnn_primitives[weights_index],
-        *mkldnn_primitives[result_index]));
+        *m_mkldnn_primitives[input_data_index],
+        *m_mkldnn_primitives[weights_index],
+        *m_mkldnn_primitives[result_index]));
 
-    primitive_deps[conv_index] = {input_data_index, weights_index, result_index};
+    m_primitive_deps[conv_index] = {input_data_index, weights_index, result_index};
     return conv_index;
 }
 
@@ -129,11 +129,11 @@ size_t MKLDNNEmitter::build_convolution_forward(const mkldnn::memory::desc& inpu
           mkldnn::memory::dims(padding_above.begin(), padding_above.end()),
           mkldnn::padding_kind::zero},
          mkldnn_utils::global_cpu_engine},
-        *mkldnn_primitives[input_data_index],
-        *mkldnn_primitives[weights_index],
-        *mkldnn_primitives[result_index]));
+        *m_mkldnn_primitives[input_data_index],
+        *m_mkldnn_primitives[weights_index],
+        *m_mkldnn_primitives[result_index]));
 
-    primitive_deps[conv_index] = {input_data_index, weights_index, result_index};
+    m_primitive_deps[conv_index] = {input_data_index, weights_index, result_index};
     return conv_index;
 }
 
@@ -151,16 +151,16 @@ size_t MKLDNNEmitter::build_elementwise_add(
     size_t input1_data_index = build_memory_primitive(input1_data_desc);
     size_t result_index = build_memory_primitive(result_desc);
 
-    inputs_primitive.push_back(*mkldnn_primitives[input0_data_index]);
-    inputs_primitive.push_back(*mkldnn_primitives[input1_data_index]);
+    inputs_primitive.push_back(*m_mkldnn_primitives[input0_data_index]);
+    inputs_primitive.push_back(*m_mkldnn_primitives[input1_data_index]);
 
     // elementwise sum primtive descriptor
     mkldnn::sum::primitive_desc sum_pd =
         mkldnn::sum::primitive_desc(result_desc, scale_vector, inputs_pd);
     // sum primitive
     size_t add_index = insert_primitive(
-        new mkldnn::sum(sum_pd, inputs_primitive, *mkldnn_primitives[result_index]));
+        new mkldnn::sum(sum_pd, inputs_primitive, *m_mkldnn_primitives[result_index]));
 
-    primitive_deps[add_index] = {input0_data_index, input1_data_index, result_index};
+    m_primitive_deps[add_index] = {input0_data_index, input1_data_index, result_index};
     return add_index;
 }

src/ngraph/runtime/cpu/mkldnn_emitter.hpp

@@ -37,11 +37,7 @@ namespace ngraph
             class MKLDNNEmitter
             {
             public:
-                MKLDNNEmitter(std::shared_ptr<CPU_ExternalFunction> ef)
-                    : external_function(ef)
-                {
-                }
-
+                MKLDNNEmitter() {}
                 const std::vector<mkldnn::primitive*>& get_mkldnn_primitives() const;
 
                 size_t insert_primitive(mkldnn::primitive* primitive);
@@ -77,10 +73,9 @@ namespace ngraph
                     const std::vector<mkldnn::memory::primitive_desc>& input_pd);
 
             private:
-                std::shared_ptr<CPU_ExternalFunction> external_function;
-                std::vector<mkldnn::primitive*> mkldnn_primitives;
-                std::vector<mkldnn::stream> mkldnn_streams;
-                std::unordered_map<size_t, std::vector<size_t>> primitive_deps;
+                std::vector<mkldnn::primitive*> m_mkldnn_primitives;
+                std::vector<mkldnn::stream> m_mkldnn_streams;
+                std::unordered_map<size_t, std::vector<size_t>> m_primitive_deps;
             };
         }
     }

src/ngraph/runtime/cpu/pass/cpu_assignment.cpp

@@ -256,7 +256,7 @@ bool runtime::cpu::pass::CPUAssignment::run_on_call_graph(
         auto handler = s_dispatcher.find(TI(n));
         if (handler != s_dispatcher.end())
         {
-            handler->second(m_external_function.get(), node.get());
+            handler->second(m_external_function, node.get());
         }
     }
 

src/ngraph/runtime/cpu/pass/cpu_assignment.hpp

@@ -39,7 +39,7 @@ namespace ngraph
                 class CPUAssignment : public ngraph::pass::CallGraphPass
                 {
                 public:
-                    CPUAssignment(std::shared_ptr<CPU_ExternalFunction> external_function)
+                    CPUAssignment(CPU_ExternalFunction* external_function)
                         : m_external_function(external_function)
                     {
                     }
@@ -56,7 +56,7 @@ namespace ngraph
                     }
 
                 private:
-                    std::shared_ptr<CPU_ExternalFunction> m_external_function;
+                    CPU_ExternalFunction* m_external_function;
                 };
             }
         }

src/ngraph/runtime/cpu/pass/cpu_layout.cpp

@@ -719,11 +719,11 @@ bool runtime::cpu::pass::CPULayout::run_on_call_graph(const std::list<std::share
         auto handler = s_dispatcher.find(TI(n));
         if (handler != s_dispatcher.end())
         {
-            handler->second(m_external_function.get(), node);
+            handler->second(m_external_function, node);
         }
         else
         {
-            set_default_layouts(m_external_function.get(), node);
+            set_default_layouts(m_external_function, node);
         }
     }
 

src/ngraph/runtime/cpu/pass/cpu_layout.hpp

@@ -39,7 +39,7 @@ namespace ngraph
                 class CPULayout : public ngraph::pass::CallGraphPass
                 {
                 public:
-                    CPULayout(std::shared_ptr<CPU_ExternalFunction> external_function)
+                    CPULayout(CPU_ExternalFunction* external_function)
                         : m_external_function(external_function)
                     {
                     }
@@ -52,7 +52,7 @@ namespace ngraph
                                std::shared_ptr<ngraph::Node> node);
 
                 private:
-                    std::shared_ptr<CPU_ExternalFunction> m_external_function;
+                    CPU_ExternalFunction* m_external_function;
                     static std::shared_ptr<Node> insert_input_conversions(
                         CPU_ExternalFunction* external_function,
                         std::shared_ptr<Node>& node,

test/backend_test.in.cpp

@@ -43,6 +43,29 @@ static const vector<element::Type> s_known_element_types = {element::from<float>
                                                             element::from<uint32_t>(),
                                                             element::from<uint64_t>()};
 
+TEST(${BACKEND_NAME}, component_cleanup)
+{
+    shared_ptr<runtime::Backend> backend;
+    shared_ptr<runtime::ExternalFunction> external;
+    shared_ptr<runtime::CallFrame> cf;
+    {
+        Shape shape{2, 2};
+        auto A = make_shared<op::Parameter>(element::f32, shape);
+        auto B = make_shared<op::Parameter>(element::f32, shape);
+        auto f = make_shared<Function>(A + B, op::ParameterVector{A, B});
+
+        auto manager = runtime::Manager::get("${BACKEND_NAME}");
+        external = manager->compile(f);
+        backend = manager->allocate_backend();
+        cf = backend->make_call_frame(external);
+    }
+    EXPECT_EQ(cf.use_count(), 1);
+    cf = nullptr;
+    EXPECT_EQ(backend.use_count(), 1);
+    backend = nullptr;
+    EXPECT_EQ(external.use_count(), 1);
+}
+
 TEST(${BACKEND_NAME}, aliased_output)
 {
     Shape shape{2, 2};
@@ -138,7 +161,7 @@ TEST(${BACKEND_NAME}, abc)
     auto f = make_shared<Function>((A + B) * C, op::ParameterVector{A, B, C});
 
     auto manager = runtime::Manager::get("${BACKEND_NAME}");
-    auto external = manager->compile(f);
+    shared_ptr<runtime::ExternalFunction> external = manager->compile(f);
     auto backend = manager->allocate_backend();
     auto cf = backend->make_call_frame(external);