[OpenMP] Enable omp_get_num_devices() on Windows

[hbae]

Summary

This patch enables omp_get_num_devices() and omp_get_initial_device() on
Windows by providing an alternative to dlsym on Windows, and proposes to
add a new libomptarget entry, __tgt_get_num_devices().

[hbae]

Applied similar change to omp_get_initial_device().

[AndreyChurbanov]

If I add GetProcAddress result conversion to (void *) I get linker error if built by clang or gcc. E.g. for clang:

C:\Progs\mingw-w64\x86_64-8.1.0-posix-seh-rt_v6-rev0\mingw64\bin\ld.exe: CMakeFiles\omp.dir/objects.a(z_Windows_NT_util.cpp.obj):z_Windows_NT_util.:(.text+0x1dfa): undefined reference to `EnumProcessModules'
C:\Progs\mingw-w64\x86_64-8.1.0-posix-seh-rt_v6-rev0\mingw64\bin\ld.exe: CMakeFiles\omp.dir/objects.a(z_Windows_NT_util.cpp.obj):z_Windows_NT_util.:(.text+0x1e26): undefined reference to `EnumProcessModules'
clang: error: linker command failed with exit code 1 (use -v to see invocation)
make[2]: * [runtime/src/libomp.dll] Error 1
make[1]: * [runtime/src/CMakeFiles/omp.dir/all] Error 2
make: *** [all] Error 2

That indicates missing -lpsapi on link line.

openmp/runtime/src/z_Windows_NT_util.cpp

1652

clang complains on this line:

C:\tmp\llvm\trunk\runtime\src\z_Windows_NT_util.cpp:1652:12: error: assigning to 'void *' from 'FARPROC'

  (aka 'long long (*)()') converts between void pointer and function pointer
proc = GetProcAddress(modules[i], name);
       ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

I guess the result should be converted to (void *) before assignment.

Similar complaint from gcc:

C:/tmp/llvm/trunk/runtime/src/z_Windows_NT_util.cpp:1652:26: error: invalid conversion from 'FARPROC' {aka 'long long int (*)()'} to 'void*' [-fpermissive]
1652 | proc = GetProcAddress(modules[i], name);

|            ~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~
|                          |
|                          FARPROC {aka long long int (*)()}

[protze.joachim]

Two comments, that are not directly related to this patch, but could be addressed while we are here.

openmp/runtime/src/kmp_ftn_entry.h

967–968

Why is there no KMP_EXPAND_NAME macro applied to this function?

971–972

Conceptually, we could really just call the omp_get_num_devices function from above.
I would at least expect that the functions are called in a consistent order.

[hbae]

Addressed issues reported by reviewers.

[AndreyChurbanov]

LGTM

openmp/runtime/cmake/LibompHandleFlags.cmake

 # Library flags
 function(libomp_get_libflags libflags)
   set(libflags_local)
   libomp_append(libflags_local "${CMAKE_THREAD_LIBS_INIT}")
   libomp_append(libflags_local "${LIBOMP_HWLOC_LIBRARY}" LIBOMP_USE_HWLOC)
   if(${IA32})
     libomp_append(libflags_local -lirc_pic LIBOMP_HAVE_IRC_PIC_LIBRARY)
   endif()
+  if(MINGW)
+    libomp_append(libflags_local -lpsapi LIBOMP_HAVE_LIBPSAPI)
+  endif()
   if(LIBOMP_HAVE_SHM_OPEN_WITH_LRT)
     libomp_append(libflags_local -lrt)
   endif()
   if(${CMAKE_SYSTEM_NAME} MATCHES "DragonFly|FreeBSD")
     libomp_append(libflags_local "-Wl,--no-as-needed" LIBOMP_HAVE_AS_NEEDED_FLAG)
     libomp_append(libflags_local "-lm")
     libomp_append(libflags_local "-Wl,--as-needed" LIBOMP_HAVE_AS_NEEDED_FLAG)
   elseif(${CMAKE_SYSTEM_NAME} MATCHES "NetBSD")

openmp/runtime/src/kmp_ftn_entry.h

__kmp_entry_thread()->th.th_current_task->td_icvs.default_device =
       KMP_DEREF arg;
 #endif
 }
 
 // Get number of NON-HOST devices.
 // 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;
   }
 #endif // KMP_MIC || KMP_OS_DARWIN || KMP_OS_WINDOWS || defined(KMP_STUB)
 }
 
 // This function always returns true when called on host device.
 // Compiler/libomptarget should handle when it is called inside target region.
 int FTN_STDCALL KMP_EXPAND_NAME(FTN_IS_INITIAL_DEVICE)(void) KMP_WEAK_ATTRIBUTE_EXTERNAL;
 int FTN_STDCALL KMP_EXPAND_NAME(FTN_IS_INITIAL_DEVICE)(void) {
   return 1; // This is the host
 }
 
 // libomptarget, if loaded, provides this function
-int FTN_STDCALL FTN_GET_INITIAL_DEVICE(void) KMP_WEAK_ATTRIBUTE_EXTERNAL;
-int FTN_STDCALL FTN_GET_INITIAL_DEVICE(void) {
-#if KMP_MIC || KMP_OS_DARWIN || KMP_OS_WINDOWS || defined(KMP_STUB)
-  // same as omp_get_num_devices()
-  return 0;
-#else
-  int (*fptr)();
-  if ((*(void **)(&fptr) = dlsym(RTLD_NEXT, "omp_get_initial_device"))) {
-    return (*fptr)();
-  } else { // liboffload & libomptarget don't exist
-    // same as omp_get_num_devices()
-    return 0;
-  }
-#endif
+int FTN_STDCALL KMP_EXPAND_NAME(FTN_GET_INITIAL_DEVICE)(void)
+    KMP_WEAK_ATTRIBUTE_EXTERNAL;
+int FTN_STDCALL KMP_EXPAND_NAME(FTN_GET_INITIAL_DEVICE)(void) {
+  return KMP_EXPAND_NAME(FTN_GET_NUM_DEVICES)();
 }
 
 #if defined(KMP_STUB)
 // Entries for stubs library
 // As all *target* functions are C-only parameters always passed by value
 void *FTN_STDCALL FTN_TARGET_ALLOC(size_t size, int device_num) { return 0; }
 
 void FTN_STDCALL FTN_TARGET_FREE(void *device_ptr, int device_num) {}
#else
   return __kmp_max_task_priority;
 #endif
 }
 
 // This function will be defined in libomptarget. When libomptarget is not
 // loaded, we assume we are on the host and return KMP_HOST_DEVICE.
 // Compiler/libomptarget will handle this if called inside target.
 int FTN_STDCALL FTN_GET_DEVICE_NUM(void) KMP_WEAK_ATTRIBUTE_EXTERNAL;
-int FTN_STDCALL FTN_GET_DEVICE_NUM(void) { return FTN_GET_INITIAL_DEVICE(); }
+int FTN_STDCALL FTN_GET_DEVICE_NUM(void) {
+  return KMP_EXPAND_NAME(FTN_GET_INITIAL_DEVICE)();
+}
 
 // Compiler will ensure that this is only called from host in sequential region
 int FTN_STDCALL FTN_PAUSE_RESOURCE(kmp_pause_status_t kind, int device_num) {
 #ifdef KMP_STUB
   return 1; // just fail
 #else
-  if (device_num == FTN_GET_INITIAL_DEVICE())
+  if (device_num == KMP_EXPAND_NAME(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
 }
 
 // Compiler will ensure that this is only called from host in sequential region
 int FTN_STDCALL FTN_PAUSE_RESOURCE_ALL(kmp_pause_status_t kind) {
 #ifdef KMP_STUB
   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
 }
 
 // Returns the maximum number of nesting levels supported by implementation
 int FTN_STDCALL FTN_GET_SUPPORTED_ACTIVE_LEVELS(void) {
 #ifdef KMP_STUB

 // OMP_4.5 versioned symbols
 KMP_VERSION_SYMBOL(FTN_GET_MAX_TASK_PRIORITY, 45, "OMP_4.5");
 KMP_VERSION_SYMBOL(FTN_GET_NUM_PLACES, 45, "OMP_4.5");
 KMP_VERSION_SYMBOL(FTN_GET_PLACE_NUM_PROCS, 45, "OMP_4.5");
 KMP_VERSION_SYMBOL(FTN_GET_PLACE_PROC_IDS, 45, "OMP_4.5");
 KMP_VERSION_SYMBOL(FTN_GET_PLACE_NUM, 45, "OMP_4.5");
 KMP_VERSION_SYMBOL(FTN_GET_PARTITION_NUM_PLACES, 45, "OMP_4.5");
 KMP_VERSION_SYMBOL(FTN_GET_PARTITION_PLACE_NUMS, 45, "OMP_4.5");
-// KMP_VERSION_SYMBOL(FTN_GET_INITIAL_DEVICE, 45, "OMP_4.5");
+KMP_VERSION_SYMBOL(FTN_GET_INITIAL_DEVICE, 45, "OMP_4.5");
 
 // OMP_5.0 versioned symbols
 // KMP_VERSION_SYMBOL(FTN_GET_DEVICE_NUM, 50, "OMP_5.0");
 // KMP_VERSION_SYMBOL(FTN_PAUSE_RESOURCE, 50, "OMP_5.0");
 // KMP_VERSION_SYMBOL(FTN_PAUSE_RESOURCE_ALL, 50, "OMP_5.0");
 // KMP_VERSION_SYMBOL(FTN_GET_SUPPORTED_ACTIVE_LEVELS, 50, "OMP_5.0");
 // KMP_VERSION_SYMBOL(FTN_FULFILL_EVENT, 50, "OMP_5.0");
 
 #endif // KMP_USE_VERSION_SYMBOLS
 
 #ifdef __cplusplus
 } // extern "C"
 #endif // __cplusplus
 
 // end of file //

openmp/runtime/src/kmp_os.h

 }
 
 template <typename T>
 bool __kmp_atomic_compare_store_rel(std::atomic<T> *p, T expected, T desired) {
   return p->compare_exchange_strong(
       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"

openmp/runtime/src/z_Windows_NT_util.cpp

 #include "kmp_wait_release.h"
 
 /* This code is related to NtQuerySystemInformation() function. This function
    is used in the Load balance algorithm for OMP_DYNAMIC=true to find the
    number of running threads in the system. */
 
 #include <ntsecapi.h> // UNICODE_STRING
 #include <ntstatus.h>
+#include <psapi.h>
+#pragma comment(lib, "psapi.lib")
 
 enum SYSTEM_INFORMATION_CLASS {
   SystemProcessInformation = 5
 }; // SYSTEM_INFORMATION_CLASS
 
 struct CLIENT_ID {
   HANDLE UniqueProcess;
   HANDLE UniqueThread;
if (buffer != NULL) {
     KMP_INTERNAL_FREE(buffer);
   }
 
   glb_running_threads = running_threads;
 
   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 = (void *)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");
 }
 
 void __kmp_do_initialize_hidden_helper_threads() {
   KMP_ASSERT(0 && "Hidden helper task is not supported on Windows");
 }