diff --git a/openmp/runtime/src/kmp_ftn_entry.h b/openmp/runtime/src/kmp_ftn_entry.h --- a/openmp/runtime/src/kmp_ftn_entry.h +++ b/openmp/runtime/src/kmp_ftn_entry.h @@ -941,13 +941,15 @@ // libomptarget, if loaded, provides this function in api.cpp. int FTN_STDCALL KMP_EXPAND_NAME(FTN_GET_NUM_DEVICES)(void) KMP_WEAK_ATTRIBUTE_EXTERNAL; int FTN_STDCALL KMP_EXPAND_NAME(FTN_GET_NUM_DEVICES)(void) { -#if KMP_MIC || KMP_OS_DARWIN || KMP_OS_WINDOWS || defined(KMP_STUB) +#if KMP_MIC || KMP_OS_DARWIN || defined(KMP_STUB) return 0; #else int (*fptr)(); - if ((*(void **)(&fptr) = dlsym(RTLD_DEFAULT, "_Offload_number_of_devices"))) { + if ((*(void **)(&fptr) = KMP_DLSYM("__tgt_get_num_devices"))) { + return (*fptr)(); + } else if ((*(void **)(&fptr) = KMP_DLSYM_NEXT("omp_get_num_devices"))) { return (*fptr)(); - } else if ((*(void **)(&fptr) = dlsym(RTLD_NEXT, "omp_get_num_devices"))) { + } else if ((*(void **)(&fptr) = KMP_DLSYM("_Offload_number_of_devices"))) { return (*fptr)(); } else { // liboffload & libomptarget don't exist return 0; @@ -970,7 +972,7 @@ return 0; #else int (*fptr)(); - if ((*(void **)(&fptr) = dlsym(RTLD_NEXT, "omp_get_initial_device"))) { + if ((*(void **)(&fptr) = KMP_DLSYM_NEXT("omp_get_initial_device"))) { return (*fptr)(); } else { // liboffload & libomptarget don't exist // same as omp_get_num_devices() @@ -1331,12 +1333,10 @@ if (device_num == FTN_GET_INITIAL_DEVICE()) return __kmpc_pause_resource(kind); else { -#if !KMP_OS_WINDOWS int (*fptr)(kmp_pause_status_t, int); - if ((*(void **)(&fptr) = dlsym(RTLD_DEFAULT, "tgt_pause_resource"))) + if ((*(void **)(&fptr) = KMP_DLSYM("tgt_pause_resource"))) return (*fptr)(kind, device_num); else -#endif return 1; // just fail if there is no libomptarget } #endif @@ -1348,11 +1348,9 @@ return 1; // just fail #else int fails = 0; -#if !KMP_OS_WINDOWS int (*fptr)(kmp_pause_status_t, int); - if ((*(void **)(&fptr) = dlsym(RTLD_DEFAULT, "tgt_pause_resource"))) + if ((*(void **)(&fptr) = KMP_DLSYM("tgt_pause_resource"))) fails = (*fptr)(kind, KMP_DEVICE_ALL); // pause devices -#endif fails += __kmpc_pause_resource(kind); // pause host return fails; #endif diff --git a/openmp/runtime/src/kmp_os.h b/openmp/runtime/src/kmp_os.h --- a/openmp/runtime/src/kmp_os.h +++ b/openmp/runtime/src/kmp_os.h @@ -1077,6 +1077,16 @@ expected, desired, std::memory_order_release, std::memory_order_relaxed); } +// Symbol lookup on Linux/Windows +#if KMP_OS_WINDOWS +extern void *__kmp_lookup_symbol(const char *name); +#define KMP_DLSYM(name) __kmp_lookup_symbol(name) +#define KMP_DLSYM_NEXT(name) nullptr +#else +#define KMP_DLSYM(name) dlsym(RTLD_DEFAULT, name) +#define KMP_DLSYM_NEXT(name) dlsym(RTLD_NEXT, name) +#endif + #endif /* KMP_OS_H */ // Safe C API #include "kmp_safe_c_api.h" diff --git a/openmp/runtime/src/z_Windows_NT_util.cpp b/openmp/runtime/src/z_Windows_NT_util.cpp --- a/openmp/runtime/src/z_Windows_NT_util.cpp +++ b/openmp/runtime/src/z_Windows_NT_util.cpp @@ -23,6 +23,8 @@ #include // UNICODE_STRING #include +#include +#pragma comment(lib, "psapi.lib") enum SYSTEM_INFORMATION_CLASS { SystemProcessInformation = 5 @@ -1632,6 +1634,29 @@ return running_threads; } //__kmp_get_load_balance() +// Find symbol from the loaded modules +void *__kmp_lookup_symbol(const char *name) { + HANDLE process = GetCurrentProcess(); + DWORD needed; + HMODULE *modules = nullptr; + if (!EnumProcessModules(process, modules, 0, &needed)) + return nullptr; + DWORD num_modules = needed / sizeof(HMODULE); + modules = (HMODULE *)malloc(num_modules * sizeof(HMODULE)); + if (!EnumProcessModules(process, modules, needed, &needed)) { + free(modules); + return nullptr; + } + void *proc = nullptr; + for (uint32_t i = 0; i < num_modules; i++) { + proc = GetProcAddress(modules[i], name); + if (proc) + break; + } + free(modules); + return proc; +} + // Functions for hidden helper task void __kmp_hidden_helper_worker_thread_wait() { KMP_ASSERT(0 && "Hidden helper task is not supported on Windows");