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 | |
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 | } |
28 | DynamicLibrary::DynamicLibrary(const char* name, const char* alt_name, bool leak_handle_): leak_handle(leak_handle_) { |
29 | // NOLINTNEXTLINE(hicpp-signed-bitwise) |
30 | handle = dlopen(name, RTLD_LOCAL | RTLD_NOW); |
31 | if (!handle) { |
32 | if (alt_name) { |
33 | handle = dlopen(alt_name, RTLD_LOCAL | RTLD_NOW); |
34 | if (!handle) { |
35 | TORCH_CHECK_WITH(DynamicLibraryError, false, "Error in dlopen for library " , name, "and " , alt_name); |
36 | } |
37 | } else { |
38 | TORCH_CHECK_WITH(DynamicLibraryError, false, "Error in dlopen: " , dlerror()); |
39 | } |
40 | } |
41 | } |
42 | |
43 | void* DynamicLibrary::sym(const char* name) { |
44 | AT_ASSERT(handle); |
45 | return checkDL(dlsym(handle, name)); |
46 | } |
47 | |
48 | DynamicLibrary::~DynamicLibrary() { |
49 | if (!handle || leak_handle) { |
50 | return; |
51 | } |
52 | dlclose(handle); |
53 | } |
54 | |
55 | #else |
56 | |
57 | // Windows |
58 | |
59 | DynamicLibrary::DynamicLibrary(const char* name, const char* alt_name, bool leak_handle_): leak_handle(leak_handle_) { |
60 | // NOLINTNEXTLINE(hicpp-signed-bitwise) |
61 | HMODULE theModule; |
62 | bool reload = true; |
63 | auto wname = c10::u8u16(name); |
64 | // Check if LOAD_LIBRARY_SEARCH_DEFAULT_DIRS is supported |
65 | if (GetProcAddress(GetModuleHandleW(L"KERNEL32.DLL" ), "AddDllDirectory" ) != NULL) { |
66 | theModule = LoadLibraryExW( |
67 | wname.c_str(), |
68 | NULL, |
69 | LOAD_LIBRARY_SEARCH_DEFAULT_DIRS); |
70 | if (theModule != NULL || (GetLastError() != ERROR_MOD_NOT_FOUND)) { |
71 | reload = false; |
72 | } |
73 | } |
74 | |
75 | if (reload) { |
76 | theModule = LoadLibraryW(wname.c_str()); |
77 | } |
78 | |
79 | if (theModule) { |
80 | handle = theModule; |
81 | } else { |
82 | char buf[256]; |
83 | DWORD dw = GetLastError(); |
84 | FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, |
85 | NULL, dw, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), |
86 | buf, (sizeof(buf) / sizeof(char)), NULL); |
87 | TORCH_CHECK_WITH(DynamicLibraryError, false, "error in LoadLibrary for " , name, ". WinError " , dw, ": " , buf); |
88 | } |
89 | } |
90 | |
91 | void* DynamicLibrary::sym(const char* name) { |
92 | AT_ASSERT(handle); |
93 | FARPROC procAddress = GetProcAddress((HMODULE)handle, name); |
94 | if (!procAddress) { |
95 | TORCH_CHECK_WITH(DynamicLibraryError, false, "error in GetProcAddress" ); |
96 | } |
97 | return (void*)procAddress; |
98 | } |
99 | |
100 | DynamicLibrary::~DynamicLibrary() { |
101 | if (!handle || leak_handle) { |
102 | return; |
103 | } |
104 | FreeLibrary((HMODULE)handle); |
105 | } |
106 | |
107 | #endif |
108 | #endif |
109 | |
110 | } // namespace at |
111 | |