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

Reland https://github.com/llvm/llvm-project/pull/164392 with Fortran support moved to follow-up PR

offload/include/OpenMP/omp.h

@@ -30,6 +30,14 @@
 
 extern "C" {
 
+/// Definitions
+///{
+
+// See definition in OpenMP (omp.h.var/omp_lib.(F90|h).var)
+#define omp_invalid_device -2
+
+///}
+
 /// Type declarations
 ///{
 

offload/include/omptarget.h

@@ -272,6 +272,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

@@ -403,6 +403,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/kmp_ftn_cdecl.cpp

@@ -29,6 +29,36 @@ char const __kmp_version_ftncdecl[] =
 #define FTN_STDCALL /* no stdcall */
 #include "kmp_ftn_os.h"
 #include "kmp_ftn_entry.h"
+
+// FIXME: this is a hack to get the UID functions working for C.
+// It will be moved and also made available for Fortran in a follow-up patch.
+extern "C" {
+const char *FTN_STDCALL omp_get_uid_from_device(int device_num)
+    KMP_WEAK_ATTRIBUTE_EXTERNAL;
+const char *FTN_STDCALL omp_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 omp_get_device_from_uid(const char *device_uid)
+    KMP_WEAK_ATTRIBUTE_EXTERNAL;
+int FTN_STDCALL omp_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
+}
+}
 #else
                        "no";
 #endif /* KMP_FTN_ENTRIES */

openmp/runtime/test/api/omp_device_uid.c

@@ -0,0 +1,77 @@
+// RUN: %libomp-compile-and-run 2>&1 | FileCheck %s
+// Linking fails for icc 18
+// UNSUPPORTED: icc-18
+
+#include <omp_testsuite.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