Making a module option for Caffe2

Summary:
This will help releasing models that are using Caffe2 but have their own operator implementations and extensions. More detailed docs to arrive later. Let's see what contbuild says.
Closes https://github.com/caffe2/caffe2/pull/1378

Differential Revision: D6155045

Pulled By: Yangqing

fbshipit-source-id: 657a4c8de2f8e095bad5ed5db5b3e476b2a877e1

CMakeLists.txt

@@ -46,6 +46,11 @@ option(USE_ZMQ "Use ZMQ" OFF)
 # ---[ CMake scripts + modules
 list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules)
 
+# ---[ CMake build directories
+set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
+set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
+set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
+
 enable_testing()
 
 # External projects
@@ -77,9 +82,6 @@ if(CMAKE_COMPILER_IS_GNUCXX AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 4.8.0
   list(APPEND Caffe2_DEPENDENCY_LIBS gcc_s gcc)
 endif()
 
-# ---[ Set output directories
-set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "binaries")
-
 # ---[ Build flags
 set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -std=c11")
 if(NOT MSVC)
@@ -174,3 +176,7 @@ if (NOT ANDROID)
       FILE Caffe2Targets.cmake
       COMPONENT dev)
 endif()
+
+
+# ---[ Modules
+add_subdirectory(modules/module_test)

caffe2/core/macros.h.in

@@ -25,6 +25,7 @@ static_assert(
    CAFFE2_VERSION_PATCH)
 
 #cmakedefine CAFFE2_ANDROID
+#cmakedefine CAFFE2_BUILD_SHARED_LIBS
 #cmakedefine CAFFE2_FORCE_FALLBACK_CUDA_MPI
 #cmakedefine CAFFE2_HAS_MKL_DNN
 #cmakedefine CAFFE2_HAS_MKL_SGEMM_PACK

caffe2/core/module.cc

@@ -0,0 +1,138 @@
+/**
+ * Copyright (c) 2016-present, Facebook, 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
+ * limitations under the License.
+ */
+
+#include "caffe2/core/logging.h"
+#include "caffe2/core/module.h"
+
+#ifndef _MSC_VER
+#include <dlfcn.h>
+#endif
+
+namespace caffe2 {
+
+static std::mutex& gModuleChangeMutex() {
+  static std::mutex m_;
+  return m_;
+}
+
+static CaffeMap<string, const ModuleSchema*>& MutableCurrentModules() {
+  static CaffeMap<string, const ModuleSchema*> module_schema_map_;
+  return module_schema_map_;
+}
+
+// Note(jiayq): I am not sure whether the module handles are going to be used
+// as C2 uses modules via registration, but let's keep the handles at least.
+static CaffeMap<string, void*> CurrentModuleHandles() {
+  static CaffeMap<string, void*> module_handle_map_;
+  return module_handle_map_;
+}
+
+const CaffeMap<string, const ModuleSchema*>& CurrentModules() {
+  return MutableCurrentModules();
+}
+
+ModuleSchema::ModuleSchema(const char* name, const char* description)
+    : name_(name), description_(description) {
+  std::lock_guard<std::mutex> guard(gModuleChangeMutex());
+  MutableCurrentModules().emplace(name, this);
+}
+
+bool HasModule(const string& name) {
+ auto& modules = CurrentModules();
+ return (modules.find(name) != modules.end());
+}
+
+#ifdef _MSC_VER
+
+void LoadModule(const string& name, const string& filename) {
+  CAFFE_ENFORCE(!HasModule(name),
+    "On Windows, LoadModule is currently not supported yet and you should "
+    "use static linking for any module that you intend to use.");
+}
+
+#else
+
+void LoadModule(const string& name, const string& filename) {
+  CAFFE_ENFORCE(
+      name.size() > 0 || filename.size() > 0,
+      "You must provide at least one of name and filename.");
+  if (name.size() && HasModule(name)) {
+    VLOG(1) << "Module " << name << " already present. Skip loading."; 
+    return;
+  }
+  void* handle = nullptr;
+  if (filename.size()) {
+    handle = dlopen(
+        filename.c_str(), RTLD_NOW | RTLD_GLOBAL);
+    CAFFE_ENFORCE(handle != nullptr,
+      "Cannot load module ",
+      name,
+      " (with given filename ",
+      filename,
+      "), are you sure it is correct?");
+  } else {
+    string inferred_name = string("lib") + name + ".so";
+    handle = dlopen(
+        inferred_name.c_str(), RTLD_NOW | RTLD_GLOBAL);
+#ifdef __APPLE__
+    // For apple, we will also try the dylib extension.
+    if (!handle) {
+      string inferred_name = string("lib") + name + ".dylib";
+      handle = dlopen(
+          inferred_name.c_str(), RTLD_NOW | RTLD_GLOBAL | RTLD_NODELETE);
+    }
+#endif
+    CAFFE_ENFORCE(handle != nullptr,
+        "Cannot load module ",
+        name,
+        " (with inferred filename ",
+        inferred_name,
+        "), are you sure it is in the dynamic linker search path?");
+  }
+  // After the module is loaded, we should check if it actually has the
+  // intended module name. If not, it might be that the module file name
+  // and the module name are inconsistent.
+  if (name.size()) {
+    string module_name_check = "gCaffe2ModuleSanityCheck" + name;
+    CAFFE_ENFORCE(dlsym(handle, module_name_check.c_str()),
+        "The loaded module ",
+        name,
+        " did not pass the module name sanity check. Is it built with the "
+        "right configs? Make sure the file name and the CAFFE2_MODULE name "
+        "are consistent.");
+    // After it passes the dlopen and dlsym check, we should add it to the
+    // current handles.
+    std::lock_guard<std::mutex> guard(gModuleChangeMutex());
+    CurrentModuleHandles()[name] = handle;
+  } else {
+    // If not, we issue a warning that one is recommended to use explicit
+    // module name.
+    LOG(WARNING)
+        << "Module file " << filename
+        << " was loaded without a proper module name. It is recommended "
+           "that one load a model with an explicit module name in addition "
+           "to the filename.";
+    // As a contingency, we will store the current module handle with the
+    // filename.
+    std::lock_guard<std::mutex> guard(gModuleChangeMutex());
+    CurrentModuleHandles()[filename] = handle;
+  }
+}
+
+#endif // _MSC_VER
+
+}  // namespace caffe2
+

caffe2/core/module.h

@@ -0,0 +1,91 @@
+/**
+ * Copyright (c) 2016-present, Facebook, 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
+ * limitations under the License.
+ */
+
+/**
+ * A global dictionary that holds information about what Caffe2 modules have
+ * been loaded in the current runtime, and also utility functions to load
+ * modules.
+ */
+#ifndef CAFFE2_CORE_MODULE_H_
+#define CAFFE2_CORE_MODULE_H_
+
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <functional>
+#include <memory>
+#include <mutex>
+
+#include "caffe2/core/common.h"
+#include "caffe2/core/typeid.h"
+
+namespace caffe2 {
+
+/**
+ * A module schema that can be used to store specific information about
+ * different modules. Currently, we only store the name and a simple
+ * description of what this module does.
+ */
+class ModuleSchema {
+ public:
+  ModuleSchema(const char* name, const char* description);
+
+ private:
+  const char* name_;
+  const char* description_;
+};
+
+
+/**
+ * @brief Current Modules present in the Caffe2 runtime.
+ * Returns:
+ *   map: a map of modules and (optionally) their description. The key is the
+ *       module name, and the value is the description for that module. The
+ *       module name is recommended to be the part that constitutes the trunk
+ *       of the dynamic library: for example, a module called
+ *       libcaffe2_db_rocksdb.so should have the name "caffe2_db_rocksdb". The
+ *       reason we do not use "lib" is because it's somewhat redundant, and
+ *       the reason we do not include ".so" is for cross-platform compatibility
+ *       on platforms like mac os.
+ */
+const CaffeMap<string, const ModuleSchema*>& CurrentModules();
+
+/**
+ * @brief Checks whether a module is already present in the current binary.
+ */
+bool HasModule(const string& name);
+
+/**
+ * @brief Load a module.
+ * Inputs:
+ *   name: a module name or a path name.
+ *       It is recommended that you use the name of the module, and leave the
+ *       full path option to only experimental modules.
+ *   filename: (optional) a filename that serves as a hint to load the module.
+ */
+void LoadModule(const string& name, const string& filename="");
+
+
+#define CAFFE2_MODULE(name, description)                                    \
+  extern "C" {                                                              \
+    const bool gCaffe2ModuleSanityCheck##name() { return true; }            \
+  }                                                                         \
+  namespace {                                                               \
+    static ::caffe2::ModuleSchema module_schema_##name(#name, description); \
+  }
+
+}  // namespace caffe2
+#endif  // CAFFE2_CORE_MODULE_H_

caffe2/core/module_test.cc

@@ -0,0 +1,93 @@
+/**
+ * Copyright (c) 2016-present, Facebook, 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
+ * limitations under the License.
+ */
+
+#include <iostream>
+#include <memory>
+
+#include "caffe2/core/module.h"
+#include "caffe2/core/operator.h"
+#include <gtest/gtest.h>
+#include "caffe2/core/logging.h"
+
+// An explicitly defined module, testing correctness when we statically link a
+// module
+CAFFE2_MODULE(caffe2_module_test_static, "Static module for testing.");
+
+namespace caffe2 {
+
+class Caffe2ModuleTestStaticDummyOp : public OperatorBase {
+ public:
+  using OperatorBase::OperatorBase;
+  bool Run(int /* unused */ /*stream_id*/) override {
+    return true;
+  }
+  virtual string type() {
+    return "base";
+  }
+};
+
+REGISTER_CPU_OPERATOR(
+  Caffe2ModuleTestStaticDummy, Caffe2ModuleTestStaticDummyOp);
+OPERATOR_SCHEMA(Caffe2ModuleTestStaticDummy);
+
+TEST(ModuleTest, StaticModule) {
+  const string name = "caffe2_module_test_static";
+  const auto& modules = CurrentModules();
+  EXPECT_EQ(modules.count(name), 1);
+  EXPECT_TRUE(HasModule(name));
+
+  // LoadModule should not raise an error, since the module is already present.
+  LoadModule(name);
+  // Even a non-existing path should not cause error.
+  LoadModule(name, "/does/not/exist.so");
+  EXPECT_EQ(modules.count(name), 1);
+  EXPECT_TRUE(HasModule(name));
+
+  // The module will then introduce the Caffe2ModuleTestStaticDummyOp.
+  OperatorDef op_def;
+  Workspace ws;
+  op_def.set_type("Caffe2ModuleTestStaticDummy");
+  unique_ptr<OperatorBase> op = CreateOperator(op_def, &ws);
+  EXPECT_NE(nullptr, op.get());
+}
+
+#ifdef CAFFE2_BUILD_SHARED_LIBS
+TEST(ModuleTest, DynamicModule) {
+  const string name = "caffe2_module_test_dynamic";
+  const auto& modules = CurrentModules();
+  EXPECT_EQ(modules.count(name), 0);
+  EXPECT_FALSE(HasModule(name));
+
+  // Before loading, we should not be able to create the op.
+  OperatorDef op_def;
+  Workspace ws;
+  op_def.set_type("Caffe2ModuleTestDynamicDummy");
+  EXPECT_THROW(
+      CreateOperator(op_def, &ws),
+      EnforceNotMet);
+
+  // LoadModule should load the proper module.
+  LoadModule(name);
+  EXPECT_EQ(modules.count(name), 1);
+  EXPECT_TRUE(HasModule(name));
+
+  // The module will then introduce the Caffe2ModuleTestDynamicDummyOp.
+  unique_ptr<OperatorBase> op_after_load = CreateOperator(op_def, &ws);
+  EXPECT_NE(nullptr, op_after_load.get());
+}
+#endif
+
+}  // namespace caffe2

cmake/Dependencies.cmake

@@ -127,9 +127,9 @@ if(BUILD_TEST)
   caffe2_include_directories(${PROJECT_SOURCE_DIR}/third_party/googletest/googletest/include)
 
   # We will not need to test benchmark lib itself.
-  set(BENCHMARK_ENABLE_TESTING OFF)
+  set(BENCHMARK_ENABLE_TESTING OFF CACHE BOOL "Disable benchmark testing as we don't need it.")
   # We will not need to install benchmark since we link it statically. 
-  set(BENCHMARK_ENABLE_INSTALL OFF)
+  set(BENCHMARK_ENABLE_INSTALL OFF CACHE BOOL "Disable benchmark install to avoid overwriting vendor install.")
   add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/benchmark)
   caffe2_include_directories(${PROJECT_SOURCE_DIR}/third_party/benchmark/include)
 

cmake/MiscCheck.cmake

@@ -113,3 +113,6 @@ if (USE_ASAN)
   set(CMAKE_MODULE_LINKER_FLAGS "${CMAKE_MODULE_LINKER_FLAGS} ${CAFFE2_ASAN_FLAG}")
   set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} ${CAFFE2_ASAN_FLAG}")
 endif()
+
+# ---[ Create CAFFE2_BUILD_SHARED_LIBS for macros.h.in usage.
+set(CAFFE2_BUILD_SHARED_LIBS ${BUILD_SHARED_LIBS})

modules/module_test/CMakeLists.txt

@@ -0,0 +1,6 @@
+add_library(
+    caffe2_module_test_dynamic SHARED
+    ${CMAKE_CURRENT_SOURCE_DIR}/module_test_dynamic.cc)
+
+target_link_libraries(caffe2_module_test_dynamic caffe2)
+install(TARGETS caffe2_module_test_dynamic DESTINATION lib)

modules/module_test/module_test_dynamic.cc

@@ -0,0 +1,41 @@
+/**
+ * Copyright (c) 2016-present, Facebook, 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
+ * limitations under the License.
+ */
+
+#include "caffe2/core/module.h"
+#include "caffe2/core/operator.h"
+
+// An explicitly defined module, testing correctness when we dynamically link a
+// module
+CAFFE2_MODULE(caffe2_module_test_dynamic, "Dynamic module for testing.");
+
+namespace caffe2 {
+
+class Caffe2ModuleTestDynamicDummyOp : public OperatorBase {
+ public:
+  using OperatorBase::OperatorBase;
+  bool Run(int /* unused */ /*stream_id*/) override {
+    return true;
+  }
+  virtual string type() {
+    return "base";
+  }
+};
+
+REGISTER_CPU_OPERATOR(
+  Caffe2ModuleTestDynamicDummy, Caffe2ModuleTestDynamicDummyOp);
+OPERATOR_SCHEMA(Caffe2ModuleTestDynamicDummy);
+
+} // namespace caffe2