1 #include <c10/util/Exception.h>
2 #include <c10/util/Unicode.h>
3 #include <ATen/DynamicLibrary.h>
4 #include <ATen/Utils.h>
5
6 #ifndef _WIN32
7 #include <dlfcn.h>
8 #include <libgen.h>
9 #else
10 #include <c10/util/win32-headers.h>
11 #endif
12
13 namespace at {
14
15
16 #ifndef C10_MOBILE
17 #ifndef _WIN32
18
19 // Unix
20
checkDL(void * x)21 static void* checkDL(void* x) {
22 if (!x) {
23 TORCH_CHECK_WITH(DynamicLibraryError, false, "Error in dlopen or dlsym: ", dlerror());
24 }
25
26 return x;
27 }
DynamicLibrary(const char * name,const char * alt_name,bool leak_handle_)28 DynamicLibrary::DynamicLibrary(const char* name, const char* alt_name, bool leak_handle_): leak_handle(leak_handle_), handle(dlopen(name, RTLD_LOCAL | RTLD_NOW)) {
29 if (!handle) {
30 if (alt_name) {
31 handle = dlopen(alt_name, RTLD_LOCAL | RTLD_NOW);
32 if (!handle) {
33 TORCH_CHECK_WITH(DynamicLibraryError, false, "Error in dlopen for library ", name, "and ", alt_name);
34 }
35 } else {
36 TORCH_CHECK_WITH(DynamicLibraryError, false, "Error in dlopen: ", dlerror());
37 }
38 }
39 }
40
sym(const char * name)41 void* DynamicLibrary::sym(const char* name) {
42 AT_ASSERT(handle);
43 return checkDL(dlsym(handle, name));
44 }
45
~DynamicLibrary()46 DynamicLibrary::~DynamicLibrary() {
47 if (!handle || leak_handle) {
48 return;
49 }
50 dlclose(handle);
51 }
52
53 #else
54
55 // Windows
56
57 DynamicLibrary::DynamicLibrary(const char* name, const char* alt_name, bool leak_handle_): leak_handle(leak_handle_) {
58 // NOLINTNEXTLINE(hicpp-signed-bitwise)
59 HMODULE theModule;
60 bool reload = true;
61 auto wname = c10::u8u16(name);
62 // Check if LOAD_LIBRARY_SEARCH_DEFAULT_DIRS is supported
63 if (GetProcAddress(GetModuleHandleW(L"KERNEL32.DLL"), "AddDllDirectory") != NULL) {
64 theModule = LoadLibraryExW(
65 wname.c_str(),
66 NULL,
67 LOAD_LIBRARY_SEARCH_DEFAULT_DIRS);
68 if (theModule != NULL || (GetLastError() != ERROR_MOD_NOT_FOUND)) {
69 reload = false;
70 }
71 }
72
73 if (reload) {
74 theModule = LoadLibraryW(wname.c_str());
75 }
76
77 if (theModule) {
78 handle = theModule;
79 } else {
80 char buf[256];
81 DWORD dw = GetLastError();
82 FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
83 NULL, dw, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
84 buf, (sizeof(buf) / sizeof(char)), NULL);
85 TORCH_CHECK_WITH(DynamicLibraryError, false, "error in LoadLibrary for ", name, ". WinError ", dw, ": ", buf);
86 }
87 }
88
89 void* DynamicLibrary::sym(const char* name) {
90 AT_ASSERT(handle);
91 FARPROC procAddress = GetProcAddress((HMODULE)handle, name);
92 if (!procAddress) {
93 TORCH_CHECK_WITH(DynamicLibraryError, false, "error in GetProcAddress");
94 }
95 return (void*)procAddress;
96 }
97
98 DynamicLibrary::~DynamicLibrary() {
99 if (!handle || leak_handle) {
100 return;
101 }
102 FreeLibrary((HMODULE)handle);
103 }
104
105 #endif
106 #endif
107
108 } // namespace at
109