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

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.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D149557

3 files changed, 26 insertions, 3 deletions

diff --git a/openmp/runtime/src/kmp.h b/openmp/runtime/src/kmp.h
index 43e723fa7387..bd50d58f0fdc 100644
--- a/openmp/runtime/src/kmp.h
+++ b/openmp/runtime/src/kmp.h
@@ -1100,6 +1100,12 @@ 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)))
 
diff --git a/openmp/runtime/src/kmp_runtime.cpp b/openmp/runtime/src/kmp_runtime.cpp
index 776160289981..a0a72f163e7c 100644
--- a/openmp/runtime/src/kmp_runtime.cpp
+++ b/openmp/runtime/src/kmp_runtime.cpp
@@ -553,6 +553,14 @@ 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
 
@@ -7041,6 +7049,11 @@ static void __kmp_do_serial_initialize(void) {
 
   __kmp_validate_locks();
 
+#if ENABLE_LIBOMPTARGET
+  /* Initialize functions from libomptarget */
+  __kmp_init_omptarget();
+#endif
+
   /* Initialize internal memory allocator */
   __kmp_init_allocator();
 
diff --git a/openmp/runtime/src/kmp_tasking.cpp b/openmp/runtime/src/kmp_tasking.cpp
index 71922ed70800..7e9147eeeebf 100644
--- a/openmp/runtime/src/kmp_tasking.cpp
+++ b/openmp/runtime/src/kmp_tasking.cpp
@@ -22,8 +22,11 @@
 #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 */
@@ -1805,7 +1808,8 @@ static void __kmp_invoke_task(kmp_int32 gtid, kmp_task_t *task,
       // 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) {