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
13namespace at {
14
15
16#ifndef C10_MOBILE
17#ifndef _WIN32
18
19// Unix
20
21static 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}
28DynamicLibrary::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
43void* DynamicLibrary::sym(const char* name) {
44 AT_ASSERT(handle);
45 return checkDL(dlsym(handle, name));
46}
47
48DynamicLibrary::~DynamicLibrary() {
49 if (!handle || leak_handle) {
50 return;
51 }
52 dlclose(handle);
53}
54
55#else
56
57// Windows
58
59DynamicLibrary::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
91void* 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
100DynamicLibrary::~DynamicLibrary() {
101 if (!handle || leak_handle) {
102 return;
103 }
104 FreeLibrary((HMODULE)handle);
105}
106
107#endif
108#endif
109
110} // namespace at
111