[OpenMP] Handle function calls from `libomp` to `libomptarget` correctly

[tianshilei1992]

Summary

D132005 introduced function calls from libomp to libomptarget if offloading
is enabled. However, the external function declaration may not always work. For
example, it causes a link error on macOS. Currently it is guarded properly by
a macro, but in order to get OpenMP target offloading working on macOS, it has
to be handled correctly. This patch applies the same idea of how we support
target memory extension by using function pointer indirect call for that function.

[jdoerfert]

openmp/runtime/src/kmp_tasking.cpp

28

should we check/assert this is non-null, or do it later?

[tianshilei1992]

openmp/runtime/src/kmp_tasking.cpp

28

I'll add an assertion since it is supposed to be there.

[jdoerfert]

LG, with assertion

[tianshilei1992]

rebase and fix comments

[tianshilei1992]

openmp/runtime/src/kmp_tasking.cpp

1811–1812

We actually want to do the assert here.

[tianshilei1992]

fix runtime error

openmp/runtime/src/kmp.h

 extern void ___kmpc_free(int gtid, void *ptr, omp_allocator_handle_t al);
 
 extern void __kmp_init_memkind();
 extern void __kmp_fini_memkind();
 extern void __kmp_init_target_mem();
 
 /* ------------------------------------------------------------------------ */
 
+#if ENABLE_LIBOMPTARGET
+extern void __kmp_init_target_task();
+#endif
+
+/* ------------------------------------------------------------------------ */
+
 #define KMP_UINT64_MAX                                                         \
   (~((kmp_uint64)1 << ((sizeof(kmp_uint64) * (1 << 3)) - 1)))
 
 #define KMP_MIN_NTH 1
 
 #ifndef KMP_MAX_NTH
 #if defined(PTHREAD_THREADS_MAX) && PTHREAD_THREADS_MAX < INT_MAX
 #define KMP_MAX_NTH PTHREAD_THREADS_MAX

openmp/runtime/src/kmp_runtime.cpp

 static void __kmp_init_allocator() {
   __kmp_init_memkind();
   __kmp_init_target_mem();
 }
 static void __kmp_fini_allocator() { __kmp_fini_memkind(); }
 
 /* ------------------------------------------------------------------------ */
 
+#if ENABLE_LIBOMPTARGET
+static void __kmp_init_omptarget() {
+  __kmp_init_target_task();
+}
+#endif
+
+/* ------------------------------------------------------------------------ */
+
 #if KMP_DYNAMIC_LIB
 #if KMP_OS_WINDOWS
 
 BOOL WINAPI DllMain(HINSTANCE hInstDLL, DWORD fdwReason, LPVOID lpReserved) {
   //__kmp_acquire_bootstrap_lock( &__kmp_initz_lock );
 
   switch (fdwReason) {
 

 #endif
 #if OMPD_SUPPORT
   __kmp_env_dump();
   ompd_init();
 #endif
 
   __kmp_validate_locks();
 
+#if ENABLE_LIBOMPTARGET
+  /* Initialize functions from libomptarget */
+  __kmp_init_omptarget();
+#endif
+
   /* Initialize internal memory allocator */
   __kmp_init_allocator();
 
   /* Register the library startup via an environment variable or via mapped
      shared memory file and check to see whether another copy of the library is
      already registered. Since forked child process is often terminated, we
      postpone the registration till middle initialization in the child */
   if (__kmp_need_register_serial)

openmp/runtime/src/kmp_tasking.cpp

 #include "kmp_wait_release.h"
 #include "kmp_taskdeps.h"
 
 #if OMPT_SUPPORT
 #include "ompt-specific.h"
 #endif
 
 #if ENABLE_LIBOMPTARGET
-// Declaration of synchronization function from libomptarget.
-extern "C" void __tgt_target_nowait_query(void **) KMP_WEAK_ATTRIBUTE_INTERNAL;
+static void (*tgt_target_nowait_query)(void **);
+
+void __kmp_init_target_task() {
+  *(void **)(&tgt_target_nowait_query) = KMP_DLSYM("__tgt_target_nowait_query");
+}
 #endif
 
 /* forward declaration */
 static void __kmp_enable_tasking(kmp_task_team_t *task_team,
                                  kmp_info_t *this_thr);
 static void __kmp_alloc_task_deque(kmp_info_t *thread,
                                    kmp_thread_data_t *thread_data);
 static int __kmp_realloc_task_threads_data(kmp_info_t *thread,
#if USE_ITT_BUILD && USE_ITT_NOTIFY
     KMP_FSYNC_ACQUIRED(taskdata); // acquired self (new task)
 #endif
 
 #if ENABLE_LIBOMPTARGET
     if (taskdata->td_target_data.async_handle != NULL) {
       // If we have a valid target async handle, that means that we have already
       // executed the task routine once. We must query for the handle completion
       // instead of re-executing the routine.
-      __tgt_target_nowait_query(&taskdata->td_target_data.async_handle);
+      KMP_ASSERT(tgt_target_nowait_query);
+      tgt_target_nowait_query(&taskdata->td_target_data.async_handle);
     } else
 #endif
     if (task->routine != NULL) {
 #ifdef KMP_GOMP_COMPAT
       if (taskdata->td_flags.native) {
         ((void (*)(void *))(*(task->routine)))(task->shareds);
       } else
 #endif /* KMP_GOMP_COMPAT */