Merge pull request #246 from NervanaSystems/jmenon/mem_leaks

CPU: Define ownerships and lifetimes for all entities

src/ngraph/codegen/compiler.cpp

@@ -82,7 +82,7 @@ void Compiler::set_precompiled_header_source(const std::string& source)
 std::unique_ptr<llvm::Module> Compiler::compile(const std::string& source)
 {
     lock_guard<mutex> lock(m_mutex);
-    return s_static_compiler.compile(source);
+    return s_static_compiler.compile(compiler_action, source);
 }
 
 static std::string GetExecutablePath(const char* Argv0)
@@ -112,11 +112,10 @@ StaticCompiler::StaticCompiler()
     args.push_back(m_source_name.c_str());
 
     // Prepare DiagnosticEngine
-    DiagnosticOptions DiagOpts;
-    TextDiagnosticPrinter* textDiagPrinter = new clang::TextDiagnosticPrinter(errs(), &DiagOpts);
-    IntrusiveRefCntPtr<clang::DiagnosticIDs> pDiagIDs;
-    DiagnosticsEngine* pDiagnosticsEngine =
-        new DiagnosticsEngine(pDiagIDs, &DiagOpts, textDiagPrinter);
+    IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts = new DiagnosticOptions();
+    TextDiagnosticPrinter* textDiagPrinter = new clang::TextDiagnosticPrinter(errs(), &*DiagOpts);
+    IntrusiveRefCntPtr<DiagnosticIDs> DiagID(new DiagnosticIDs());
+    DiagnosticsEngine DiagEngine(DiagID, &*DiagOpts, textDiagPrinter);
 
     // Create and initialize CompilerInstance
     m_compiler = std::unique_ptr<CompilerInstance>(new CompilerInstance());
@@ -124,7 +123,7 @@ StaticCompiler::StaticCompiler()
 
     // Initialize CompilerInvocation
     CompilerInvocation::CreateFromArgs(
-        m_compiler->getInvocation(), &args[0], &args[0] + args.size(), *pDiagnosticsEngine);
+        m_compiler->getInvocation(), &args[0], &args[0] + args.size(), DiagEngine);
 
     // Infer the builtin include path if unspecified.
     if (m_compiler->getHeaderSearchOpts().UseBuiltinIncludes &&
@@ -209,7 +208,7 @@ StaticCompiler::StaticCompiler()
     CGO.OmitLeafFramePointer = 1;
     CGO.VectorizeLoop = 1;
     CGO.VectorizeSLP = 1;
-    CGO.CXAAtExit = 0;
+    CGO.CXAAtExit = 1;
 
     if (m_debuginfo_enabled)
     {
@@ -305,7 +304,9 @@ void StaticCompiler::use_cached_files()
     }
 }
 
-std::unique_ptr<llvm::Module> StaticCompiler::compile(const string& source)
+std::unique_ptr<llvm::Module>
+    StaticCompiler::compile(std::unique_ptr<clang::CodeGenAction>& compiler_action,
+                            const string& source)
 {
     if (!m_precompiled_header_valid && m_precomiled_header_source.empty() == false)
     {
@@ -325,11 +326,11 @@ std::unique_ptr<llvm::Module> StaticCompiler::compile(const string& source)
     m_compiler->getInvocation().getPreprocessorOpts().addRemappedFile(m_source_name, buffer.get());
 
     // Create and execute action
-    CodeGenAction* compilerAction = new EmitCodeGenOnlyAction();
+    compiler_action.reset(new EmitCodeGenOnlyAction());
     std::unique_ptr<llvm::Module> rc;
-    if (m_compiler->ExecuteAction(*compilerAction) == true)
+    if (m_compiler->ExecuteAction(*compiler_action) == true)
     {
-        rc = compilerAction->takeModule();
+        rc = compiler_action->takeModule();
     }
 
     buffer.release();
@@ -359,4 +360,5 @@ void StaticCompiler::generate_pch(const string& source)
     buffer.release();
 
     m_compiler->getInvocation().getPreprocessorOpts().clearRemappedFiles();
+    delete compilerAction;
 }

src/ngraph/codegen/compiler.hpp

@@ -18,6 +18,8 @@
 #include <memory>
 #include <string>
 
+#include <clang/CodeGen/CodeGenAction.h>
+
 #include <llvm/ExecutionEngine/MCJIT.h> // forces JIT to link in
 #include <llvm/ExecutionEngine/SectionMemoryManager.h>
 #include <llvm/Option/Arg.h>
@@ -53,8 +55,9 @@ public:
     ~Compiler();
     void set_precompiled_header_source(const std::string& source);
     std::unique_ptr<llvm::Module> compile(const std::string& source);
-
+    std::unique_ptr<clang::CodeGenAction>& get_compiler_action() { return compiler_action; }
 private:
+    std::unique_ptr<clang::CodeGenAction> compiler_action;
 };
 
 class ngraph::codegen::StaticCompiler : public llvm::SectionMemoryManager
@@ -70,7 +73,9 @@ public:
         m_precomiled_header_source = source;
     }
     void add_header_search_path(const std::string& path);
-    std::unique_ptr<llvm::Module> compile(const std::string& source);
+
+    std::unique_ptr<llvm::Module> compile(std::unique_ptr<clang::CodeGenAction>& compiler_action,
+                                          const std::string& source);
     void generate_pch(const std::string& source);
 
 private:

src/ngraph/codegen/execution_engine.cpp

@@ -34,11 +34,11 @@ bool codegen::ExecutionEngine::add_module(std::unique_ptr<llvm::Module>& module)
     {
         if (!m_execution_engine)
         {
-            m_execution_engine = llvm::EngineBuilder(move(module))
-                                     .setEngineKind(llvm::EngineKind::JIT)
-                                     .setOptLevel(llvm::CodeGenOpt::Aggressive)
-                                     .setErrorStr(&m_jit_error)
-                                     .create();
+            m_execution_engine.reset(llvm::EngineBuilder(move(module))
+                                         .setEngineKind(llvm::EngineKind::JIT)
+                                         .setOptLevel(llvm::CodeGenOpt::Aggressive)
+                                         .setErrorStr(&m_jit_error)
+                                         .create());
 
             if (!m_execution_engine)
             {

src/ngraph/codegen/execution_engine.hpp

@@ -46,7 +46,7 @@ public:
     }
 
 private:
-    llvm::ExecutionEngine* m_execution_engine;
+    std::unique_ptr<llvm::ExecutionEngine> m_execution_engine;
     std::string m_jit_error;
 
     template <typename signature>

src/ngraph/runtime/cpu/call_frame.cpp

@@ -20,8 +20,10 @@
 using namespace std;
 using namespace ngraph::runtime::cpu;
 
-CallFrame::CallFrame(EntryPoint compiled_function)
-    : m_compiled_function(compiled_function)
+CallFrame::CallFrame(std::shared_ptr<ExternalFunction> external_function,
+                     EntryPoint compiled_function)
+    : m_external_function(external_function)
+    , m_compiled_function(compiled_function)
 {
 }
 

src/ngraph/runtime/cpu/call_frame.hpp

@@ -31,6 +31,7 @@ namespace ngraph
         namespace cpu
         {
             class CallFrame;
+            class ExternalFunction;
 
             using EntryPoint_t = void(void** inputs, void** outputs);
 
@@ -40,7 +41,8 @@ namespace ngraph
             class CallFrame : public ngraph::runtime::CallFrame
             {
             public:
-                CallFrame(EntryPoint compiled_function);
+                CallFrame(std::shared_ptr<ExternalFunction> external_function,
+                          EntryPoint compiled_function);
 
                 /// @brief Invoke the function with values matching the signature of the function.
                 ///
@@ -55,6 +57,7 @@ namespace ngraph
                                  const std::vector<std::shared_ptr<TensorView>>& outputs);
 
             protected:
+                std::shared_ptr<ExternalFunction> m_external_function;
                 EntryPoint m_compiled_function;
             };
         }

src/ngraph/runtime/cpu/external_function.cpp

@@ -198,6 +198,13 @@ using namespace ngraph::runtime::cpu::eigen;
 )";
     string pch_header_source = TU.get_code();
 
+    // The "dso_handle" symbol is required by __cxa_atexit()
+    // which is enabled because the JIT uses it as the default mechanism
+    // to register cleanup handlers. We use it, and not atexit(), because
+    // atexit() happens too late, when the JIT is no longer alive
+
+    TU << "void *__dso_handle = 0;\n\n";
+
     TU << "// Declare all functions\n";
     for (shared_ptr<Function> f : pass_manager.get_state().get_functions())
     {
@@ -312,19 +319,20 @@ using namespace ngraph::runtime::cpu::eigen;
     out << code;
     out.close();
 
-    codegen::Compiler compiler;
-    codegen::ExecutionEngine execution_engine;
+    compiler.reset(new codegen::Compiler());
+    execution_engine.reset(new codegen::ExecutionEngine());
+
+    compiler->set_precompiled_header_source(pch_header_source);
 
-    compiler.set_precompiled_header_source(pch_header_source);
+    auto llvm_module = compiler->compile(code);
 
-    auto llvm_module = compiler.compile(code);
     if (llvm_module == nullptr)
     {
         throw runtime_error("function failed to compile");
     }
-    execution_engine.add_module(llvm_module);
-    execution_engine.finalize();
-    m_compiled_function = execution_engine.find_function<EntryPoint_t>(function_name);
+    execution_engine->add_module(llvm_module);
+    execution_engine->finalize();
+    m_compiled_function = execution_engine->find_function<EntryPoint_t>(function_name);
     assert(m_compiled_function);
 
     m_is_compiled = true;
@@ -341,5 +349,5 @@ shared_ptr<ngraph::runtime::CallFrame> ExternalFunction::make_call_frame()
         compile();
     }
 
-    return make_shared<ngraph::runtime::cpu::CallFrame>(m_compiled_function);
+    return make_shared<ngraph::runtime::cpu::CallFrame>(shared_from_this(), m_compiled_function);
 }

src/ngraph/runtime/cpu/external_function.hpp

@@ -21,6 +21,7 @@
 #include <unordered_map>
 
 #include "ngraph/codegen/compiler.hpp"
+#include "ngraph/codegen/execution_engine.hpp"
 #include "ngraph/function.hpp"
 #include "ngraph/runtime/cpu/call_frame.hpp"
 #include "ngraph/runtime/external_function.hpp"
@@ -36,9 +37,6 @@ namespace ngraph
             class Emitter;
             class CallFrame;
 
-            using FunctionMap =
-                std::unordered_map<std::shared_ptr<Function>, std::shared_ptr<ExternalFunction>>;
-
             using OpFunction = std::function<void(Emitter*,
                                                   const ngraph::Node*,
                                                   const std::vector<TensorViewInfo>& inputs,
@@ -46,7 +44,8 @@ namespace ngraph
 
             using OpMap = std::unordered_map<std::type_index, OpFunction>;
 
-            class ExternalFunction : public ngraph::runtime::ExternalFunction
+            class ExternalFunction : public ngraph::runtime::ExternalFunction,
+                                     public std::enable_shared_from_this<ExternalFunction>
             {
             public:
                 ExternalFunction(const std::shared_ptr<ngraph::Function>& function,
@@ -57,6 +56,10 @@ namespace ngraph
                 void compile();
 
                 EntryPoint m_compiled_function;
+
+            private:
+                std::unique_ptr<codegen::Compiler> compiler;
+                std::unique_ptr<codegen::ExecutionEngine> execution_engine;
             };
         }
     }