Merge pull request #208 from NervanaSystems/yixing/dlopen-backend

support for runtime dlopen

src/ngraph/CMakeLists.txt

@@ -125,6 +125,29 @@ if (NGRAPH_CPU_ENABLE AND LLVM_INCLUDE_DIR AND
 endif()
 
 add_library(ngraph SHARED ${SRC})
+
+# Colon separated string for specified runtime plugin loading, this is made explicit s.t. if a
+# plugin is specified at compile time but the corresponding library could not be resolved at run-
+# time, an error will be generated.
+# E.g. assume compiling with Argon and Xpu, then -DRUNTIME_PLUGIN_LIBS="libargon.so:libxpu.so".
+if (DEFINED RUNTIME_PLUGIN_LIBS)
+    target_compile_definitions(ngraph PRIVATE RUNTIME_PLUGIN_LIBS=${RUNTIME_PLUGIN_LIBS})
+else()
+    target_compile_definitions(ngraph PRIVATE RUNTIME_PLUGIN_LIBS="")
+endif()
+
+# This is used to ensure that libngraph.so and libargon.so are in the same directory for dlopen.
+# Effective at build time. Does not affect `make install` logics.
+if (DEFINED COMMON_LIBRARY_OUTPUT_DIRECTORY)
+    set_target_properties(ngraph PROPERTIES
+        LIBRARY_OUTPUT_DIRECTORY ${COMMON_LIBRARY_OUTPUT_DIRECTORY})
+else()
+    set_target_properties(ngraph PROPERTIES
+        LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
+endif()
+message(STATUS "LIBRARY_OUTPUT_DIRECTORY set to: ${COMMON_LIBRARY_OUTPUT_DIRECTORY}")
+
+
 target_include_directories(ngraph PUBLIC "${NGRAPH_INCLUDE_PATH}")
 
 if(NGRAPH_CPU_ENABLE AND LLVM_LINK_LIBS)
@@ -141,6 +164,7 @@ if(NGRAPH_CPU_ENABLE AND MKLDNN_LIB_DIR)
     target_link_libraries(ngraph LINK_PRIVATE mkldnn)
 endif()
 
+
 #-----------------------------------------------------------------------------------------------
 # Installation logic...
 #-----------------------------------------------------------------------------------------------

src/ngraph/runtime/manager.cpp

@@ -12,13 +12,63 @@
 // See the License for the specific language governing permissions and
 // ----------------------------------------------------------------------------
 
+#include <dlfcn.h>
+#include <iostream>
+#include <sstream>
+#include <string>
+
+#include "ngraph/except.hpp"
 #include "ngraph/runtime/manager.hpp"
+#include "ngraph/util.hpp"
 
 using namespace ngraph::runtime;
 
+bool Manager::m_is_factory_map_initialized = false;
+std::shared_ptr<std::vector<void*>> Manager::m_plugin_lib_handles =
+    std::make_shared<std::vector<void*>>(std::vector<void*>());
+
+void Manager::load_plugins(const std::string& runtime_plugin_libs)
+{
+    std::vector<std::string> plugin_lib_paths = ngraph::split(runtime_plugin_libs, ':', false);
+
+    for (auto plugin_lib_path : plugin_lib_paths)
+    {
+        if (plugin_lib_path.size() > 0)
+        {
+            void* lib_handle = dlopen(plugin_lib_path.c_str(), RTLD_NOW);
+            if (lib_handle)
+            {
+                Manager::m_plugin_lib_handles->push_back(lib_handle);
+            }
+            else
+            {
+                throw ngraph_error("Cannot open library " + plugin_lib_path);
+            }
+        }
+    }
+}
+
+// TODO: Should call this function after plugin is not needed anymore.
+void Manager::close_plugins()
+{
+    for (auto lib_handle : *Manager::m_plugin_lib_handles)
+    {
+        dlclose(lib_handle);
+    }
+    Manager::m_plugin_lib_handles->clear();
+}
+
 Manager::FactoryMap& Manager::get_factory_map()
 {
+    // Stores Manager Factories
     static FactoryMap factory_map;
+
+    // Try to load runtime plugins
+    if (!Manager::m_is_factory_map_initialized)
+    {
+        Manager::load_plugins(RUNTIME_PLUGIN_LIBS);
+        Manager::m_is_factory_map_initialized = true;
+    }
     return factory_map;
 }
 
@@ -27,7 +77,7 @@ std::shared_ptr<Manager> Manager::get(const std::string& name)
     return get_factory_map().at(name)(name);
 }
 
-Manager::Factory Manager::register_factory(std::string name, Factory factory)
+Manager::Factory Manager::register_factory(const std::string& name, Factory factory)
 {
     get_factory_map()[name] = factory;
     return factory;

src/ngraph/runtime/manager.hpp

@@ -18,6 +18,7 @@
 #include <map>
 #include <memory>
 #include <string>
+#include <vector>
 
 namespace ngraph
 {
@@ -46,13 +47,23 @@ namespace ngraph
                 compile(const std::shared_ptr<ngraph::Function>& fun) = 0;
 
             using Factory = std::function<std::shared_ptr<Manager>(const std::string&)>;
-            using FactoryMap = std::map<std::string, Factory>;
-
-            static FactoryMap& get_factory_map();
 
             static std::shared_ptr<Manager> get(const std::string& name);
 
-            static Factory register_factory(std::string name, Factory factory);
+            static Factory register_factory(const std::string& name, Factory factory);
+
+        private:
+            static void load_plugins(const std::string& runtime_plugin_libs);
+
+            static void close_plugins();
+
+            static std::shared_ptr<std::vector<void*>> m_plugin_lib_handles;
+
+            static bool m_is_factory_map_initialized;
+
+            using FactoryMap = std::map<std::string, Factory>;
+
+            static FactoryMap& get_factory_map();
         };
     }
 }