[Offload][OpenMP] Accept omp_initial_device for runtime calls (#205290)

In OpenMP v5.2, the constant `omp_initial_device` was introduced, which
was defined as a value of `-1`.
Support for this value was added to Clang and Flang in
09cd294482 and
3c14034c55.

However, runtime calls are not aware of this constant. As such,
execution of a runtime function, passing `omp_initial_device` as the
`device_num`, would abort with e.g.,:

```
omptarget fatal error -1: "invalid value" device number '-1' out of range, only 1 devices available
```

To fix this issue, also extend the host device checks in the API
functions to also check for `omp_initial_device`.
For testing, extend API tests using `omp_get_initial_device()` to also
test `omp_initial_device` and extend `omp_device_uid.c` to check if the
UID for `omp_initial_device` and `omp_get_initial_device()` matches.

Partially addresses https://github.com/llvm/llvm-project/issues/192980
(see
https://github.com/llvm/llvm-project/issues/192980#issuecomment-4404241994)

---------

Signed-off-by: Jan André Reuter <j.reuter@fz-juelich.de>
Signed-off-by: Jan André Reuter <jan@zyten.de>

offload/libomptarget/OpenMP/API.cpp

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "../private.h"
 #include "PluginManager.h"
 #include "device.h"
 #include "omptarget.h"
@@ -114,7 +115,7 @@ EXTERN const char *omp_get_uid_from_device(int DeviceNum) {
     ODBG(ODT_Interface) << "Call to " << __func__ << " returning nullptr";
     return nullptr;
   }
-  if (DeviceNum == omp_get_initial_device()) {
+  if (isInitialDevice(DeviceNum)) {
     ODBG(ODT_Interface) << "Call to " << __func__
                         << " returning initial device UID";
     return GenericPluginTy::getHostDeviceUid();
@@ -141,7 +142,7 @@ EXTERN int omp_get_initial_device(void) {
 EXTERN size_t omp_get_gprivate_limit(int DeviceNum, omp_access_t AccessGroup) {
   TIMESCOPE();
   OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
-  if (DeviceNum == omp_get_initial_device())
+  if (isInitialDevice(DeviceNum))
     return 0;
 
   if (AccessGroup != omp_access_cgroup)
@@ -230,7 +231,7 @@ EXTERN int omp_target_is_present(const void *Ptr, int DeviceNum) {
     return false;
   }
 
-  if (DeviceNum == omp_get_initial_device()) {
+  if (isInitialDevice(DeviceNum)) {
     ODBG(ODT_Interface) << "Call to " << __func__ << " on host, returning true";
     return true;
   }
@@ -267,7 +268,7 @@ EXTERN int omp_target_is_accessible(const void *Ptr, size_t Size,
     return false;
   }
 
-  if (DeviceNum == omp_get_initial_device() || DeviceNum == -1) {
+  if (isInitialDevice(DeviceNum)) {
     ODBG(ODT_Interface) << "Call to " << __func__ << " on host, returning true";
     return true;
   }
@@ -307,13 +308,12 @@ EXTERN int omp_target_memcpy(void *Dst, const void *Src, size_t Length,
   void *SrcAddr = (char *)const_cast<void *>(Src) + SrcOffset;
   void *DstAddr = (char *)Dst + DstOffset;
 
-  if (SrcDevice == omp_get_initial_device() &&
-      DstDevice == omp_get_initial_device()) {
+  if (isInitialDevice(SrcDevice) && isInitialDevice(DstDevice)) {
     ODBG(ODT_Interface) << "copy from host to host";
     const void *P = memcpy(DstAddr, SrcAddr, Length);
     if (P == NULL)
       Rc = OFFLOAD_FAIL;
-  } else if (SrcDevice == omp_get_initial_device()) {
+  } else if (isInitialDevice(SrcDevice)) {
     ODBG(ODT_Interface) << "copy from host to device";
     auto DstDeviceOrErr = PM->getDevice(DstDevice);
     if (!DstDeviceOrErr)
@@ -321,7 +321,7 @@ EXTERN int omp_target_memcpy(void *Dst, const void *Src, size_t Length,
                     toString(DstDeviceOrErr.takeError()).c_str());
     AsyncInfoTy AsyncInfo(*DstDeviceOrErr);
     Rc = DstDeviceOrErr->submitData(DstAddr, SrcAddr, Length, AsyncInfo);
-  } else if (DstDevice == omp_get_initial_device()) {
+  } else if (isInitialDevice(DstDevice)) {
     ODBG(ODT_Interface) << "copy from device to host";
     auto SrcDeviceOrErr = PM->getDevice(SrcDevice);
     if (!SrcDeviceOrErr)
@@ -473,7 +473,7 @@ EXTERN void *omp_target_memset(void *Ptr, int ByteVal, size_t NumBytes,
     return Ptr;
   }
 
-  if (DeviceNum == omp_get_initial_device()) {
+  if (isInitialDevice(DeviceNum)) {
     ODBG(ODT_Interface) << "filling memory on host via memset";
     memset(Ptr, ByteVal, NumBytes); // ignore return value, memset() cannot fail
   } else {
@@ -671,7 +671,7 @@ EXTERN int omp_target_associate_ptr(const void *HostPtr, const void *DevicePtr,
     return OFFLOAD_FAIL;
   }
 
-  if (DeviceNum == omp_get_initial_device()) {
+  if (isInitialDevice(DeviceNum)) {
     REPORT() << __func__ << ": no association possible on the host";
     return OFFLOAD_FAIL;
   }
@@ -704,7 +704,7 @@ EXTERN int omp_target_disassociate_ptr(const void *HostPtr, int DeviceNum) {
     return OFFLOAD_FAIL;
   }
 
-  if (DeviceNum == omp_get_initial_device()) {
+  if (isInitialDevice(DeviceNum)) {
     REPORT() << __func__ << ": no association possible on the host";
     return OFFLOAD_FAIL;
   }
@@ -736,7 +736,7 @@ EXTERN void *omp_get_mapped_ptr(const void *Ptr, int DeviceNum) {
   }
 
   int NumDevices = omp_get_initial_device();
-  if (DeviceNum == NumDevices) {
+  if (isInitialDevice(DeviceNum)) {
     ODBG(ODT_Interface) << "Device " << DeviceNum
                         << " is initial device, returning Ptr " << Ptr;
     return const_cast<void *>(Ptr);

offload/libomptarget/interface.cpp

@@ -67,7 +67,7 @@ bool checkDevice(int64_t &DeviceID, ident_t *Loc) {
     return true;
   }
 
-  if (DeviceID == omp_get_initial_device()) {
+  if (isInitialDevice(static_cast<int>(DeviceID))) {
     ODBG(ODT_Device) << "Device is host (" << DeviceID
                      << "), returning as if offload is disabled";
     return true;

offload/libomptarget/omptarget.cpp

@@ -211,7 +211,7 @@ void *targetAllocExplicit(size_t Size, int DeviceNum, int Kind,
 
   void *Rc = NULL;
 
-  if (DeviceNum == omp_get_initial_device()) {
+  if (isInitialDevice(DeviceNum)) {
     Rc = malloc(Size);
     ODBG(ODT_Interface) << Name << " returns host ptr " << Rc;
     return Rc;
@@ -236,7 +236,7 @@ void targetFreeExplicit(void *DevicePtr, int DeviceNum, int Kind,
     return;
   }
 
-  if (DeviceNum == omp_get_initial_device()) {
+  if (isInitialDevice(DeviceNum)) {
     free(DevicePtr);
     ODBG(ODT_Interface) << Name << " deallocated host ptr";
     return;

offload/libomptarget/private.h

@@ -17,6 +17,7 @@
 #include "Shared/SourceInfo.h"
 
 #include "OpenMP/InternalTypes.h"
+#include "OpenMP/omp.h"
 
 #include "device.h"
 #include "omptarget.h"
@@ -85,4 +86,14 @@ printKernelArguments(const ident_t *Loc, const int64_t DeviceId,
   }
 }
 
+////////////////////////////////////////////////////////////////////////////////
+/// Checks if the passed device is the initial device (i.e., host device)
+/// While the device number is defined as the value of the total number of
+/// host devices (i.e. omp_get_initial_device()), the alias omp_initial_device
+/// can be used as well (Ref. OpenMP v5.2, p.23 l.33 - p.24 l.6).
+static inline bool isInitialDevice(const int &DeviceNum) {
+  return DeviceNum == omp_get_initial_device() ||
+         DeviceNum == omp_initial_device;
+}
+
 #endif

offload/test/api/omp_device_uid.c

@@ -65,6 +65,14 @@ int main() {
       num_failed++;
     }
   }
+  // test that uid for omp_initial_device matches uid
+  // for omp_get_num_devices(), i.e. the host device
+  if (omp_get_uid_from_device(omp_initial_device) !=
+      omp_get_uid_from_device(num_devices)) {
+    printf("FAIL for matching uid for omp_initial_device and "
+           "omp_get_initial_device()");
+    num_failed++;
+  }
   if (num_failed) {
     printf("FAIL\n");
     return 1;

offload/test/api/omp_host_pinned_memory.c

@@ -7,10 +7,10 @@
 void *llvm_omp_target_alloc_host(size_t, int);
 void llvm_omp_target_free_host(void *, int);
 
-int main() {
+void run_test(int hostDev) {
   const int N = 64;
   const int device = omp_get_default_device();
-  const int host = omp_get_initial_device();
+  int host = hostDev;
 
   int *hst_ptr = llvm_omp_target_alloc_host(N * sizeof(int), device);
 
@@ -27,7 +27,13 @@ int main() {
     sum += hst_ptr[i];
 
   llvm_omp_target_free_host(hst_ptr, device);
-  // CHECK: PASS
   if (sum == N)
     printf("PASS\n");
 }
+
+int main() {
+  // CHECK: PASS
+  run_test(omp_get_initial_device());
+  // CHECK: PASS
+  run_test(omp_initial_device);
+}

offload/test/api/omp_target_memcpy_async1.c

@@ -6,9 +6,9 @@
 #include <omp.h>
 #include <stdlib.h>
 
-int main() {
+int run_test(int hostDev) {
   int d = omp_get_default_device();
-  int id = omp_get_initial_device();
+  int id = hostDev;
   int q[128], i;
   void *p;
 
@@ -46,3 +46,7 @@ int main() {
 
   return 0;
 }
+
+int main() {
+  return run_test(omp_get_initial_device()) && run_test(omp_initial_device);
+}

offload/test/api/omp_target_memcpy_async2.c

@@ -4,9 +4,9 @@
 #include <omp.h>
 #include <stdlib.h>
 
-int main() {
+int run_test(int hostDev) {
   int d = omp_get_default_device();
-  int id = omp_get_initial_device();
+  int id = hostDev;
   int a[128], b[64], c[32], e[16], q[128], i;
   void *p;
 
@@ -71,3 +71,7 @@ int main() {
 
   return 0;
 }
+
+int main() {
+  return run_test(omp_get_initial_device()) && run_test(omp_initial_device);
+}

offload/test/api/omp_target_memcpy_rect_async1.c

@@ -6,9 +6,9 @@
 
 #define NUM_DIMS 3
 
-int main() {
+int run_test(int hostDev) {
   int d = omp_get_default_device();
-  int id = omp_get_initial_device();
+  int id = hostDev;
   int q[128], q2[128], i;
   void *p;
 
@@ -64,3 +64,7 @@ int main() {
   omp_target_free(p, d);
   return 0;
 }
+
+int main() {
+  return run_test(omp_get_initial_device()) && run_test(omp_initial_device);
+}

offload/test/api/omp_target_memcpy_rect_async2.c

@@ -5,9 +5,9 @@
 
 #define NUM_DIMS 3
 
-int main() {
+int run_test(int hostDev) {
   int d = omp_get_default_device();
-  int id = omp_get_initial_device();
+  int id = hostDev;
   int a[128], b[64], c[128], e[16], q[128], i;
   void *p;
 
@@ -87,3 +87,7 @@ int main() {
   omp_target_free(p, d);
   return 0;
 }
+
+int main() {
+  return run_test(omp_get_initial_device()) && run_test(omp_initial_device);
+}

offload/test/api/omp_target_memset.c

@@ -4,9 +4,9 @@
 #include <omp.h>
 #include <stdlib.h>
 
-int main() {
+int run_test(int hostDev) {
   int d = omp_get_default_device();
-  int id = omp_get_initial_device();
+  int id = hostDev;
   int q[128], i;
   void *p;
   void *result;
@@ -43,3 +43,7 @@ int main() {
 
   return 0;
 }
+
+int main() {
+  return run_test(omp_get_initial_device()) && run_test(omp_initial_device);
+}