[OpenMP] Add omp_get_device_num() and update several other device API…

… functions

Add omp_get_device_num() function for 5.0 which returns the number of the device
the current thread is running on. Also, did some cleanup and updating of device
API functions to make them into weak functions that should be replaced with
libomptarget functions when libomptarget is present.

Patch by Terry Wilmarth

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

llvm-svn: 347799

openmp/runtime/src/dllexports

@@ -509,9 +509,7 @@ kmp_set_warnings_off                        780
     omp_is_initial_device                   869
     omp_set_default_device                  879
     omp_get_default_device                  880
-    %ifdef stub
-        omp_get_num_devices                 881
-    %endif
+    omp_get_num_devices                     881
 %endif # OMP_40
 
 # OpenMP 45
@@ -526,8 +524,8 @@ kmp_set_warnings_off                        780
     omp_get_place_num                       876
     omp_get_partition_num_places            877
     omp_get_partition_place_nums            878
+    omp_get_initial_device                  882
     %ifdef stub
-        omp_get_initial_device              882
         omp_target_alloc                    883
         omp_target_free                     884
         omp_target_is_present               885
@@ -546,6 +544,7 @@ kmp_set_disp_num_buffers                    890
     omp_get_default_allocator               893
     omp_alloc                               894
     omp_free                                895
+    omp_get_device_num                      896
 
     OMP_NULL_ALLOCATOR                     DATA
     omp_default_mem_alloc                  DATA

openmp/runtime/src/include/45/omp.h.var

@@ -195,4 +195,3 @@
 #   endif
 
 #endif /* __OMP_H */
-

openmp/runtime/src/include/45/omp_lib.f.var

@@ -562,6 +562,7 @@
 !dec$ attributes alias:'OMP_GET_TEAM_NUM' :: omp_get_team_num
 !dec$ attributes alias:'OMP_GET_CANCELLATION' :: omp_get_cancellation
 !dec$ attributes alias:'OMP_IS_INITIAL_DEVICE' :: omp_is_initial_device
+!dec$ attributes alias:'OMP_GET_INITIAL_DEVICE' :: omp_get_initial_device
 !dec$ attributes alias:'OMP_GET_MAX_TASK_PRIORITY' :: omp_get_max_task_priority
 
 !dec$ attributes alias:'omp_init_lock' :: omp_init_lock
@@ -642,6 +643,7 @@
 !dec$ attributes alias:'_OMP_GET_TEAM_NUM' :: omp_get_team_num
 !dec$ attributes alias:'_OMP_GET_CANCELLATION' :: omp_get_cancellation
 !dec$ attributes alias:'_OMP_IS_INITIAL_DEVICE' :: omp_is_initial_device
+!dec$ attributes alias:'_OMP_GET_INITIAL_DEVICE' :: omp_get_initial_device
 !dec$ attributes alias:'_OMP_GET_MAX_TASK_PRIORTY' :: omp_get_max_task_priority
 
 !dec$ attributes alias:'_omp_init_lock' :: omp_init_lock
@@ -725,6 +727,7 @@
 !dec$ attributes alias:'omp_get_team_num_'::omp_get_team_num
 !dec$ attributes alias:'omp_get_cancellation_'::omp_get_cancellation
 !dec$ attributes alias:'omp_is_initial_device_'::omp_is_initial_device
+!dec$ attributes alias:'omp_get_initial_device_'::omp_get_initial_device
 !dec$ attributes alias:'omp_get_max_task_priority_'::omp_get_max_task_priority
 
 !dec$ attributes alias:'omp_init_lock_'::omp_init_lock
@@ -800,10 +803,14 @@
 !dec$ attributes alias:'_omp_get_proc_bind_' :: omp_get_proc_bind
 !dec$ attributes alias:'_omp_get_wtime_'::omp_get_wtime
 !dec$ attributes alias:'_omp_get_wtick_'::omp_get_wtick
+!dec$ attributes alias:'_omp_get_default_device_'::omp_get_default_device
+!dec$ attributes alias:'_omp_set_default_device_'::omp_set_default_device
+!dec$ attributes alias:'_omp_get_num_devices_'::omp_get_num_devices
 !dec$ attributes alias:'_omp_get_num_teams_'::omp_get_num_teams
 !dec$ attributes alias:'_omp_get_team_num_'::omp_get_team_num
 !dec$ attributes alias:'_omp_get_cancellation_'::omp_get_cancellation
 !dec$ attributes alias:'_omp_is_initial_device_'::omp_is_initial_device
+!dec$ attributes alias:'_omp_get_initial_device_'::omp_get_initial_device
 !dec$ attributes alias:'_omp_get_max_task_priorty_'::omp_get_max_task_priority
 
 !dec$ attributes alias:'_omp_init_lock_'::omp_init_lock
@@ -852,4 +859,3 @@
 !dec$ endif
 
       end module omp_lib
-

openmp/runtime/src/include/45/omp_lib.h.var

@@ -523,6 +523,7 @@
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_default_device
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_set_default_device
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_is_initial_device
+!DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_initial_device
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_num_devices
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_num_teams
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_team_num
@@ -595,6 +596,7 @@
 !$omp declare target(omp_get_default_device )
 !$omp declare target(omp_set_default_device )
 !$omp declare target(omp_is_initial_device )
+!$omp declare target(omp_get_initial_device )
 !$omp declare target(omp_get_num_devices )
 !$omp declare target(omp_get_num_teams )
 !$omp declare target(omp_get_team_num )
@@ -641,4 +643,3 @@
 !$omp declare target(omp_init_nest_lock_with_hint )
 !DIR$ ENDIF
 !DIR$ ENDIF
-

openmp/runtime/src/include/50/omp.h.var

@@ -141,6 +141,9 @@
     extern int   __KAI_KMPC_CONVENTION  omp_target_associate_ptr(void *, void *, size_t, size_t, int);
     extern int   __KAI_KMPC_CONVENTION  omp_target_disassociate_ptr(void *, int);
 
+    /* OpenMP 5.0 */
+    extern int   __KAI_KMPC_CONVENTION  omp_get_device_num (void);
+
     /* kmp API functions */
     extern int    __KAI_KMPC_CONVENTION  kmp_get_stacksize          (void);
     extern void   __KAI_KMPC_CONVENTION  kmp_set_stacksize          (int);
@@ -249,4 +252,3 @@
 #   endif
 
 #endif /* __OMP_H */
-

openmp/runtime/src/include/50/omp_lib.f.var

@@ -278,6 +278,11 @@
             integer (kind=omp_integer_kind) omp_get_initial_device
           end function omp_get_initial_device
 
+          function omp_get_device_num()
+            use omp_lib_kinds
+            integer (kind=omp_integer_kind) omp_get_device_num
+          end function omp_get_device_num
+
           subroutine omp_init_lock(svar)
 !DIR$ IF(__INTEL_COMPILER.GE.1400)
 !DIR$ attributes known_intrinsic :: omp_init_lock
@@ -592,7 +597,9 @@
 !dec$ attributes alias:'OMP_GET_TEAM_NUM' :: omp_get_team_num
 !dec$ attributes alias:'OMP_GET_CANCELLATION' :: omp_get_cancellation
 !dec$ attributes alias:'OMP_IS_INITIAL_DEVICE' :: omp_is_initial_device
+!dec$ attributes alias:'OMP_GET_INITIAL_DEVICE' :: omp_get_initial_device
 !dec$ attributes alias:'OMP_GET_MAX_TASK_PRIORITY' :: omp_get_max_task_priority
+!dec$ attributes alias:'OMP_GET_DEVICE_NUM' :: omp_get_device_num
 !dec$ attributes alias:'OMP_CONTROL_TOOL' :: omp_control_tool
 
 !dec$ attributes alias:'omp_init_lock' :: omp_init_lock
@@ -673,7 +680,9 @@
 !dec$ attributes alias:'_OMP_GET_TEAM_NUM' :: omp_get_team_num
 !dec$ attributes alias:'_OMP_GET_CANCELLATION' :: omp_get_cancellation
 !dec$ attributes alias:'_OMP_IS_INITIAL_DEVICE' :: omp_is_initial_device
+!dec$ attributes alias:'_OMP_GET_INITIAL_DEVICE' :: omp_get_initial_device
 !dec$ attributes alias:'_OMP_GET_MAX_TASK_PRIORTY' :: omp_get_max_task_priority
+!dec$ attributes alias:'_OMP_GET_DEVICE_NUM' :: omp_get_device_num
 !dec$ attributes alias:'_OMP_CONTROL_TOOL' :: omp_control_tool
 
 !dec$ attributes alias:'_omp_init_lock' :: omp_init_lock
@@ -757,7 +766,9 @@
 !dec$ attributes alias:'omp_get_team_num_'::omp_get_team_num
 !dec$ attributes alias:'omp_get_cancellation_'::omp_get_cancellation
 !dec$ attributes alias:'omp_is_initial_device_'::omp_is_initial_device
+!dec$ attributes alias:'omp_get_initial_device_'::omp_get_initial_device
 !dec$ attributes alias:'omp_get_max_task_priority_'::omp_get_max_task_priority
+!dec$ attributes alias:'omp_get_device_num_'::omp_get_device_num
 
 !dec$ attributes alias:'omp_init_lock_'::omp_init_lock
 !dec$ attributes alias:'omp_init_lock_with_hint_'::omp_init_lock_with_hint
@@ -833,12 +844,16 @@
 !dec$ attributes alias:'_omp_get_proc_bind_' :: omp_get_proc_bind
 !dec$ attributes alias:'_omp_get_wtime_'::omp_get_wtime
 !dec$ attributes alias:'_omp_get_wtick_'::omp_get_wtick
+!dec$ attributes alias:'_omp_get_default_device_'::omp_get_default_device
+!dec$ attributes alias:'_omp_set_default_device_'::omp_set_default_device
+!dec$ attributes alias:'_omp_get_num_devices_'::omp_get_num_devices
 !dec$ attributes alias:'_omp_get_num_teams_'::omp_get_num_teams
 !dec$ attributes alias:'_omp_get_team_num_'::omp_get_team_num
 !dec$ attributes alias:'_omp_get_cancellation_'::omp_get_cancellation
 !dec$ attributes alias:'_omp_is_initial_device_'::omp_is_initial_device
+!dec$ attributes alias:'_omp_get_initial_device_'::omp_get_initial_device
 !dec$ attributes alias:'_omp_get_max_task_priorty_'::omp_get_max_task_priority
-
+!dec$ attributes alias:'_omp_get_device_num_'::omp_get_device_num
 !dec$ attributes alias:'_omp_init_lock_'::omp_init_lock
 !dec$ attributes alias:'_omp_init_lock_with_hint_'::omp_init_lock_with_hint
 !dec$ attributes alias:'_omp_destroy_lock_'::omp_destroy_lock
@@ -886,4 +901,3 @@
 !dec$ endif
 
       end module omp_lib
-

openmp/runtime/src/include/50/omp_lib.f90.var

@@ -294,6 +294,11 @@
             integer (kind=omp_integer_kind) omp_get_initial_device
           end function omp_get_initial_device
 
+          function omp_get_device_num() bind(c)
+            use omp_lib_kinds
+            integer (kind=omp_integer_kind) omp_get_device_num
+          end function omp_get_device_num
+
           subroutine omp_init_lock(svar) bind(c)
 !DIR$ IF(__INTEL_COMPILER.GE.1400)
 !DIR$ attributes known_intrinsic :: omp_init_lock

openmp/runtime/src/include/50/omp_lib.h.var

@@ -327,6 +327,11 @@
           integer (kind=omp_integer_kind) omp_get_initial_device
         end function omp_get_initial_device
 
+        function omp_get_device_num() bind(c)
+          import
+          integer (kind=omp_integer_kind) omp_get_device_num
+        end function omp_get_device_num
+
         subroutine omp_init_lock(svar) bind(c)
 !DIR$ IF(__INTEL_COMPILER.GE.1400)
 !DIR$ attributes known_intrinsic :: omp_init_lock
@@ -623,7 +628,9 @@
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_default_device
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_set_default_device
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_is_initial_device
+!DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_initial_device
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_num_devices
+!DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_device_num
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_num_teams
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_team_num
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_init_lock
@@ -696,7 +703,9 @@
 !$omp declare target(omp_get_default_device )
 !$omp declare target(omp_set_default_device )
 !$omp declare target(omp_is_initial_device )
+!$omp declare target(omp_get_initial_device )
 !$omp declare target(omp_get_num_devices )
+!$omp declare target(omp_get_device_num )
 !$omp declare target(omp_get_num_teams )
 !$omp declare target(omp_get_team_num )
 !$omp declare target(omp_init_lock )

openmp/runtime/src/kmp.h

@@ -3943,6 +3943,13 @@ extern kmp_target_offload_kind_t __kmp_target_offload;
 extern int __kmpc_get_target_offload();
 #endif
 
+#if OMP_40_ENABLED
+// Constants used in libomptarget
+#define KMP_DEVICE_DEFAULT -1 // This is libomptarget's default device.
+#define KMP_HOST_DEVICE -10 // This is what it is in libomptarget, go figure.
+#define KMP_DEVICE_ALL -11 // This is libomptarget's "all devices".
+#endif // OMP_40_ENABLED
+
 #ifdef __cplusplus
 }
 #endif

openmp/runtime/src/kmp_ftn_entry.h

@@ -860,39 +860,40 @@ void FTN_STDCALL KMP_EXPAND_NAME(FTN_SET_DEFAULT_DEVICE)(int KMP_DEREF arg) {
 #endif
 }
 
-#if KMP_MIC || KMP_OS_DARWIN || defined(KMP_STUB)
-
-int FTN_STDCALL FTN_GET_NUM_DEVICES(void) { return 0; }
-
-#endif // KMP_MIC || KMP_OS_DARWIN || defined(KMP_STUB)
-
-#if !KMP_OS_LINUX
-
-int FTN_STDCALL KMP_EXPAND_NAME(FTN_IS_INITIAL_DEVICE)(void) { return 1; }
-
+// 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;
+int FTN_STDCALL KMP_EXPAND_NAME(FTN_GET_NUM_DEVICES)(void) {
+#if KMP_MIC || KMP_OS_DARWIN || KMP_OS_WINDOWS || defined(KMP_STUB)
+  return 0;
 #else
-
-// This internal function is used when the entry from the offload library
-// is not found.
-int _Offload_get_device_number(void) KMP_WEAK_ATTRIBUTE;
-
-int FTN_STDCALL KMP_EXPAND_NAME(FTN_IS_INITIAL_DEVICE)(void) {
-  if (_Offload_get_device_number) {
-    return _Offload_get_device_number() == -1;
-  } else {
-    return 1;
+  int (*fptr)();
+  if ((*(void **)(&fptr) = dlsym(RTLD_DEFAULT, "_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)
 }
 
-#endif // ! KMP_OS_LINUX
+// This function always returns true when called on host device.
+// Compilier/libomptarget should handle when it is called inside target region.
+int FTN_STDCALL KMP_EXPAND_NAME(FTN_IS_INITIAL_DEVICE)(void) KMP_WEAK_ATTRIBUTE;
+int FTN_STDCALL KMP_EXPAND_NAME(FTN_IS_INITIAL_DEVICE)(void) {
+  return 1; // This is the host
+}
 
 #endif // OMP_40_ENABLED
 
-#if OMP_45_ENABLED && defined(KMP_STUB)
-// OpenMP 4.5 entries for stubs library
+#if OMP_45_ENABLED
+// OpenMP 4.5 entries
 
-int FTN_STDCALL FTN_GET_INITIAL_DEVICE(void) { return -1; }
+// libomptarget, if loaded, provides this function
+int FTN_STDCALL FTN_GET_INITIAL_DEVICE(void) KMP_WEAK_ATTRIBUTE;
+int FTN_STDCALL FTN_GET_INITIAL_DEVICE(void) { return KMP_HOST_DEVICE; }
 
+#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; }
 
@@ -923,7 +924,8 @@ int FTN_STDCALL FTN_TARGET_ASSOCIATE_PTR(void *host_ptr, void *device_ptr,
 int FTN_STDCALL FTN_TARGET_DISASSOCIATE_PTR(void *host_ptr, int device_num) {
   return -1;
 }
-#endif // OMP_45_ENABLED && defined(KMP_STUB)
+#endif // defined(KMP_STUB)
+#endif // OMP_45_ENABLED
 
 #ifdef KMP_STUB
 typedef enum { UNINIT = -1, UNLOCKED, LOCKED } kmp_stub_lock_t;
@@ -1234,6 +1236,14 @@ int FTN_STDCALL KMP_EXPAND_NAME(FTN_GET_MAX_TASK_PRIORITY)(void) {
 }
 #endif
 
+#if OMP_50_ENABLED
+// 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;
+int FTN_STDCALL FTN_GET_DEVICE_NUM(void) { return KMP_HOST_DEVICE; }
+#endif // OMP_50_ENABLED
+
 // GCC compatibility (versioned symbols)
 #ifdef KMP_USE_VERSION_SYMBOLS
 
@@ -1317,6 +1327,7 @@ KMP_VERSION_SYMBOL(FTN_GET_CANCELLATION, 40, "OMP_4.0");
 KMP_VERSION_SYMBOL(FTN_GET_DEFAULT_DEVICE, 40, "OMP_4.0");
 KMP_VERSION_SYMBOL(FTN_SET_DEFAULT_DEVICE, 40, "OMP_4.0");
 KMP_VERSION_SYMBOL(FTN_IS_INITIAL_DEVICE, 40, "OMP_4.0");
+KMP_VERSION_SYMBOL(FTN_GET_NUM_DEVICES, 40, "OMP_4.0");
 #endif /* OMP_40_ENABLED */
 
 #if OMP_45_ENABLED
@@ -1328,10 +1339,12 @@ 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");
 #endif
 
 #if OMP_50_ENABLED
 // OMP_5.0 versioned symbols
+// KMP_VERSION_SYMBOL(FTN_GET_DEVICE_NUM, 50, "OMP_5.0");
 #endif
 
 #endif // KMP_USE_VERSION_SYMBOLS