1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
|
#include "caffe2/core/module.h"
#include "caffe2/core/logging.h"
#ifndef _WIN32
#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) {
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());
}
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;
}
#ifdef _WIN32
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* 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 // _WIN32
}
} // namespace caffe2
|
