Reland: [OpenMP] Implement omp_get_uid_from_device() / omp_get_device_from_uid()

offload/include/OpenMP/omp.h

@@ -30,6 +30,13 @@
 
 extern "C" {
 
+/// Definitions
+///{
+
+#define omp_invalid_device -2
+
+///}
+
 /// Type declarations
 ///{
 

offload/include/omptarget.h

@@ -274,6 +274,8 @@ extern "C" {
 void ompx_dump_mapping_tables(void);
 int omp_get_num_devices(void);
 int omp_get_device_num(void);
+int omp_get_device_from_uid(const char *DeviceUid);
+const char *omp_get_uid_from_device(int DeviceNum);
 int omp_get_initial_device(void);
 void *omp_target_alloc(size_t Size, int DeviceNum);
 void omp_target_free(void *DevicePtr, int DeviceNum);

offload/libomptarget/OpenMP/API.cpp

@@ -40,6 +40,8 @@ EXTERN void ompx_dump_mapping_tables() {
 using namespace llvm::omp::target::ompt;
 #endif
 
+using GenericDeviceTy = llvm::omp::target::plugin::GenericDeviceTy;
+
 void *targetAllocExplicit(size_t Size, int DeviceNum, int Kind,
                           const char *Name);
 void targetFreeExplicit(void *DevicePtr, int DeviceNum, int Kind,
@@ -68,6 +70,62 @@ EXTERN int omp_get_device_num(void) {
   return HostDevice;
 }
 
+static inline bool is_initial_device_uid(const char *DeviceUid) {
+  return strcmp(DeviceUid, GenericPluginTy::getHostDeviceUid()) == 0;
+}
+
+EXTERN int omp_get_device_from_uid(const char *DeviceUid) {
+  TIMESCOPE();
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
+
+  if (!DeviceUid) {
+    DP("Call to omp_get_device_from_uid returning omp_invalid_device\n");
+    return omp_invalid_device;
+  }
+  if (is_initial_device_uid(DeviceUid)) {
+    DP("Call to omp_get_device_from_uid returning initial device number %d\n",
+       omp_get_initial_device());
+    return omp_get_initial_device();
+  }
+
+  int DeviceNum = omp_invalid_device;
+
+  auto ExclusiveDevicesAccessor = PM->getExclusiveDevicesAccessor();
+  for (const DeviceTy &Device : PM->devices(ExclusiveDevicesAccessor)) {
+    const char *Uid = Device.RTL->getDevice(Device.RTLDeviceID).getDeviceUid();
+    if (Uid && strcmp(DeviceUid, Uid) == 0) {
+      DeviceNum = Device.DeviceID;
+      break;
+    }
+  }
+
+  DP("Call to omp_get_device_from_uid returning %d\n", DeviceNum);
+  return DeviceNum;
+}
+
+EXTERN const char *omp_get_uid_from_device(int DeviceNum) {
+  TIMESCOPE();
+  OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
+
+  if (DeviceNum == omp_invalid_device) {
+    DP("Call to omp_get_uid_from_device returning nullptr\n");
+    return nullptr;
+  }
+  if (DeviceNum == omp_get_initial_device()) {
+    DP("Call to omp_get_uid_from_device returning initial device UID\n");
+    return GenericPluginTy::getHostDeviceUid();
+  }
+
+  auto DeviceOrErr = PM->getDevice(DeviceNum);
+  if (!DeviceOrErr)
+    FATAL_MESSAGE(DeviceNum, "%s", toString(DeviceOrErr.takeError()).c_str());
+
+  const char *Uid =
+      DeviceOrErr->RTL->getDevice(DeviceOrErr->RTLDeviceID).getDeviceUid();
+  DP("Call to omp_get_uid_from_device returning %s\n", Uid);
+  return Uid;
+}
+
 EXTERN int omp_get_initial_device(void) {
   TIMESCOPE();
   OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));

offload/libomptarget/exports

@@ -40,6 +40,8 @@ VERS1.0 {
     omp_get_mapped_ptr;
     omp_get_num_devices;
     omp_get_device_num;
+    omp_get_device_from_uid;
+    omp_get_uid_from_device;
     omp_get_initial_device;
     omp_target_alloc;
     omp_target_free;

offload/test/api/omp_device_uid.c

@@ -0,0 +1,76 @@
+// RUN: %libomptarget-compile-run-and-check-generic
+
+#include <omp.h>
+#include <stdio.h>
+#include <string.h>
+
+int test_omp_device_uid(int device_num) {
+  const char *device_uid = omp_get_uid_from_device(device_num);
+  if (device_uid == NULL) {
+    printf("FAIL for device %d: omp_get_uid_from_device returned NULL\n",
+           device_num);
+    return 0;
+  }
+
+  int device_num_from_uid = omp_get_device_from_uid(device_uid);
+  if (device_num_from_uid != device_num) {
+    printf(
+        "FAIL for device %d: omp_get_device_from_uid returned %d (UID: %s)\n",
+        device_num, device_num_from_uid, device_uid);
+    return 0;
+  }
+
+  if (device_num == omp_get_initial_device())
+    return 1;
+
+  int success = 1;
+
+// Note that the following code may be executed on the host if the host is the
+// device
+#pragma omp target map(tofrom : success) device(device_num)
+  {
+    int device_num = omp_get_device_num();
+
+    // omp_get_uid_from_device() in the device runtime is a dummy function
+    // returning NULL
+    const char *device_uid = omp_get_uid_from_device(device_num);
+
+    // omp_get_device_from_uid() in the device runtime is a dummy function
+    // returning omp_invalid_device.
+    int device_num_from_uid = omp_get_device_from_uid(device_uid);
+
+    // Depending on whether we're executing on the device or the host, we either
+    // got NULL as the device UID or the correct device UID.  Consequently,
+    // omp_get_device_from_uid() either returned omp_invalid_device or the
+    // correct device number (aka omp_get_initial_device()).
+    if (device_uid ? device_num_from_uid != device_num
+                   : device_num_from_uid != omp_invalid_device) {
+      printf("FAIL for device %d (target): omp_get_device_from_uid returned %d "
+             "(UID: %s)\n",
+             device_num, device_num_from_uid, device_uid);
+      success = 0;
+    }
+  }
+
+  return success;
+}
+
+int main() {
+  int num_devices = omp_get_num_devices();
+  int num_failed = 0;
+  // (also test initial device aka num_devices)
+  for (int i = 0; i < num_devices + 1; i++) {
+    if (!test_omp_device_uid(i)) {
+      printf("FAIL for device %d\n", i);
+      num_failed++;
+    }
+  }
+  if (num_failed) {
+    printf("FAIL\n");
+    return 1;
+  }
+  printf("PASS\n");
+  return 0;
+}
+
+// CHECK: PASS

openmp/device/include/DeviceTypes.h

@@ -21,6 +21,9 @@ template <typename T> using Constant = __gpu_constant T;
 template <typename T> using Local = __gpu_local T;
 template <typename T> using Global = __gpu_local T;
 
+// See definition in OpenMP (omp.h.var/omp_lib.(F90|h).var)
+#define omp_invalid_device -2
+
 enum omp_proc_bind_t {
   omp_proc_bind_false = 0,
   omp_proc_bind_true = 1,

openmp/device/include/Interface.h

@@ -130,6 +130,10 @@ int omp_get_num_devices(void);
 
 int omp_get_device_num(void);
 
+int omp_get_device_from_uid(const char *DeviceUid);
+
+const char *omp_get_uid_from_device(int DeviceNum);
+
 int omp_get_num_teams(void);
 
 int omp_get_team_num();

openmp/device/src/State.cpp

@@ -423,6 +423,12 @@ int omp_get_num_devices(void) { return config::getNumDevices(); }
 
 int omp_get_device_num(void) { return config::getDeviceNum(); }
 
+int omp_get_device_from_uid(const char *DeviceUid) {
+  return omp_invalid_device;
+}
+
+const char *omp_get_uid_from_device(int DeviceNum) { return nullptr; }
+
 int omp_get_num_teams(void) { return mapping::getNumberOfBlocksInKernel(); }
 
 int omp_get_team_num() { return mapping::getBlockIdInKernel(); }

openmp/runtime/src/dllexports

@@ -544,6 +544,8 @@ kmp_set_disp_num_buffers                    890
     omp_get_devices_all_allocator           819
     omp_get_memspace_num_resources          820
     omp_get_submemspace                     821
+    omp_get_device_from_uid                 822
+    omp_get_uid_from_device                 823
     %ifndef stub
         __kmpc_set_default_allocator
         __kmpc_get_default_allocator

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

@@ -536,6 +536,11 @@
 
     /* OpenMP 5.2 */
     extern int __KAI_KMPC_CONVENTION omp_in_explicit_task(void);
+    #define omp_invalid_device -2
+
+    /* OpenMP 6.0 */
+    extern int   __KAI_KMPC_CONVENTION  omp_get_device_from_uid(const char *DeviceUid);
+    extern const char *   __KAI_KMPC_CONVENTION  omp_get_uid_from_device(int DeviceNum);
 
     /* LLVM Extensions */
     extern void *llvm_omp_target_dynamic_shared_alloc(void);

openmp/runtime/src/include/omp_lib.F90.var

@@ -215,6 +215,8 @@
 
         integer (kind=omp_interop_kind), parameter, public :: omp_interop_none = 0
 
+        integer (kind=omp_integer_kind), parameter, public :: omp_invalid_device = -2
+
         interface
 
 !         ***
@@ -417,6 +419,20 @@
             integer (kind=omp_integer_kind) omp_get_device_num
           end function omp_get_device_num
 
+          function omp_get_uid_from_device(device_num) bind(c)
+            use, intrinsic :: iso_c_binding, only: c_ptr
+            use omp_lib_kinds
+            integer (kind=omp_integer_kind), value :: device_num
+            type(c_ptr) omp_get_uid_from_device
+          end function omp_get_uid_from_device
+
+          function omp_get_device_from_uid(device_uid) bind(c)
+            use, intrinsic :: iso_c_binding, only: c_ptr
+            use omp_lib_kinds
+            type(c_ptr), value :: device_uid
+            integer (kind=omp_integer_kind) omp_get_device_from_uid
+          end function omp_get_device_from_uid
+
           function omp_pause_resource(kind, device_num) bind(c)
             use omp_lib_kinds
             integer (kind=omp_pause_resource_kind), value :: kind
@@ -1099,6 +1115,8 @@
         public :: omp_is_initial_device
         public :: omp_get_initial_device
         public :: omp_get_device_num
+        public :: omp_get_uid_from_device
+        public :: omp_get_device_from_uid
         public :: omp_pause_resource
         public :: omp_pause_resource_all
         public :: omp_get_supported_active_levels

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

@@ -291,6 +291,9 @@
       integer(kind=omp_interop_kind)omp_interop_none
       parameter(omp_interop_none=0)
 
+      integer(kind=omp_integer_kind)omp_invalid_device
+      parameter(omp_invalid_device=-2)
+
       interface
 
 !       ***
@@ -486,6 +489,20 @@
           integer (kind=omp_integer_kind) omp_get_device_num
         end function omp_get_device_num
 
+        function omp_get_uid_from_device(device_num) bind(c)
+          import
+          use, intrinsic :: iso_c_binding, only : c_ptr
+          integer (kind=omp_integer_kind), value :: device_num
+          type(c_ptr) omp_get_uid_from_device
+        end function omp_get_uid_from_device
+
+        function omp_get_device_from_uid(device_uid) bind(c)
+          import
+          use, intrinsic :: iso_c_binding, only : c_ptr
+          type(c_ptr), value :: device_uid
+          integer (kind=omp_integer_kind) omp_get_device_from_uid
+        end function omp_get_device_from_uid
+
         function omp_pause_resource(kind, device_num) bind(c)
           import
           integer (kind=omp_pause_resource_kind), value :: kind
@@ -1159,6 +1176,8 @@
 !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_uid_from_device
+!DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_device_from_uid
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_pause_resource
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_pause_resource_all
 !DIR$ ATTRIBUTES OFFLOAD:MIC :: omp_get_supported_active_levels
@@ -1242,6 +1261,8 @@
 !$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_uid_from_device )
+!$omp declare target(omp_get_device_from_uid )
 !$omp declare target(omp_pause_resource )
 !$omp declare target(omp_pause_resource_all )
 !$omp declare target(omp_get_supported_active_levels )

openmp/runtime/src/kmp_ftn_entry.h

@@ -1543,13 +1543,38 @@ int FTN_STDCALL KMP_EXPAND_NAME(FTN_GET_MAX_TASK_PRIORITY)(void) {
 #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.
+// These functions will be defined in libomptarget. When libomptarget is not
+// loaded, we assume we are on the host.
 // 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 KMP_EXPAND_NAME(FTN_GET_INITIAL_DEVICE)();
 }
+const char *FTN_STDCALL FTN_GET_UID_FROM_DEVICE(int device_num)
+    KMP_WEAK_ATTRIBUTE_EXTERNAL;
+const char *FTN_STDCALL FTN_GET_UID_FROM_DEVICE(int device_num) {
+#if KMP_OS_DARWIN || KMP_OS_WASI || defined(KMP_STUB)
+  return nullptr;
+#else
+  const char *(*fptr)(int);
+  if ((*(void **)(&fptr) = KMP_DLSYM_NEXT("omp_get_uid_from_device")))
+    return (*fptr)(device_num);
+  // Returns the same string as used by libomptarget
+  return "HOST";
+#endif
+}
+int FTN_STDCALL FTN_GET_DEVICE_FROM_UID(const char *device_uid)
+    KMP_WEAK_ATTRIBUTE_EXTERNAL;
+int FTN_STDCALL FTN_GET_DEVICE_FROM_UID(const char *device_uid) {
+#if KMP_OS_DARWIN || KMP_OS_WASI || defined(KMP_STUB)
+  return omp_invalid_device;
+#else
+  int (*fptr)(const char *);
+  if ((*(void **)(&fptr) = KMP_DLSYM_NEXT("omp_get_device_from_uid")))
+    return (*fptr)(device_uid);
+  return KMP_EXPAND_NAME(FTN_GET_INITIAL_DEVICE)();
+#endif
+}
 
 // Compiler will ensure that this is only called from host in sequential region
 int FTN_STDCALL KMP_EXPAND_NAME(FTN_PAUSE_RESOURCE)(kmp_pause_status_t kind,