[OpenMP] Implement `omp_get_mapped_ptr`

[tianshilei1992]

Summary

This patch implements the function omp_get_mapped_ptr.

Fix #59945.

[jdoerfert]

Fix the initial device case. Use a shared lock if we can already. Otherwise, it'll wait for @carlo.bertolli. LG

openmp/libomptarget/src/api.cpp

335

Otherwise it returns the device pointer, which is ptr if device_num is the value returned by omp_get_initial_device().

340

We need a shared lock here. (and basically everywhere else)

355

And this needs to be (shared) locked too.

@carlo.bertolli will fix these up consistently (see also above).

[tianshilei1992]

rebase and fix comments

[tianshilei1992]

openmp/libomptarget/src/api.cpp

355

Yeah, they will be needed fixed consistently.

openmp/libomptarget/src/api.cpp

 #include "device.h"
 #include "omptarget.h"
 #include "private.h"
 #include "rtl.h"
 
 #include <climits>
 #include <cstdlib>
 #include <cstring>
+#include <mutex>
 
 EXTERN int omp_get_num_devices(void) {
   TIMESCOPE();
   PM->RTLsMtx.lock();
   size_t DevicesSize = PM->Devices.size();
   PM->RTLsMtx.unlock();
 
   DP("Call to omp_get_num_devices returning %zd\n", DevicesSize);
if (!deviceIsReady(DeviceNum)) {
     return OFFLOAD_FAIL;
   }
 
   DeviceTy &Device = *PM->Devices[DeviceNum];
   int Rc = Device.disassociatePtr(const_cast<void *>(HostPtr));
   DP("omp_target_disassociate_ptr returns %d\n", Rc);
   return Rc;
 }
+
+EXTERN void *omp_get_mapped_ptr(const void *Ptr, int DeviceNum) {
+  TIMESCOPE();
+  DP("Call to omp_get_mapped_ptr with ptr " DPxMOD ", device_num %d.\n",
+     DPxPTR(Ptr), DeviceNum);
+
+  if (!Ptr) {
+    REPORT("Call to omp_get_mapped_ptr with nullptr.\n");
+    return nullptr;
+  }
+
+  if (DeviceNum == omp_get_initial_device()) {
+    REPORT("Device %d is initial device. No mapped pointer.\n", DeviceNum);
+    return nullptr;
+  }
+
+  int DevicesSize = omp_get_initial_device();
+  {
+    std::lock_guard<std::mutex> LG(PM->RTLsMtx);
+    DevicesSize = PM->Devices.size();
+  }
+  if (DevicesSize <= DeviceNum) {
+    DP("DeviceNum %d is invalid, returning nullptr.\n", DeviceNum);
+    return nullptr;
+  }
+
+  if (!deviceIsReady(DeviceNum)) {
+    REPORT("Device %d is not ready, returning nullptr.\n", DeviceNum);
+    return nullptr;
+  }
+
+  bool IsLast = false;
+  bool IsHostPtr = false;
+  auto &Device = *PM->Devices[DeviceNum];
+  TargetPointerResultTy TPR =
+      Device.getTgtPtrBegin(const_cast<void *>(Ptr), 1, IsLast,
+                            /*UpdateRefCount=*/false,
+                            /*UseHoldRefCount=*/false, IsHostPtr);
+  if (!TPR.isPresent()) {
+    DP("Ptr " DPxMOD "is not present on device %d, returning nullptr.\n",
+       DPxPTR(Ptr), DeviceNum);
+    return nullptr;
+  }
+
+  DP("omp_get_mapped_ptr returns " DPxMOD ".\n", DPxPTR(TPR.TargetPointer));
+
+  return TPR.TargetPointer;
+}

openmp/libomptarget/src/exports

global:
     __tgt_target_teams_nowait_mapper;
     __tgt_target_kernel;
     __tgt_target_kernel_nowait;
     __tgt_target_nowait_query;
     __tgt_mapper_num_components;
     __tgt_push_mapper_component;
     __kmpc_push_target_tripcount;
     __kmpc_push_target_tripcount_mapper;
+    omp_get_mapped_ptr;
     omp_get_num_devices;
     omp_get_device_num;
     omp_get_initial_device;
     omp_target_alloc;
     omp_target_free;
     omp_target_is_present;
     omp_target_memcpy;
     omp_target_memcpy_rect;

openmp/libomptarget/test/api/omp_get_mapped_ptr.c

+// RUN: %libomptarget-compile-and-run-generic
+
+#include <assert.h>
+#include <omp.h>
+#include <stdlib.h>
+
+#define N 1024
+#define OFFSET 16
+
+int main(int argc, char *argv[]) {
+  int *host_data = (int *)malloc(sizeof(int) * N);
+  void *device_ptr = omp_get_mapped_ptr(host_data, 0);
+
+  assert(device_ptr == NULL && "the pointer should not be mapped right now");
+
+#pragma omp target enter data map(to: host_data[:N])
+
+  device_ptr = omp_get_mapped_ptr(host_data, 0);
+
+  assert(device_ptr && "the pointer should be mapped now");
+
+  void *ptr = NULL;
+
+#pragma omp target map(from: ptr)
+  { ptr = host_data; }
+
+  assert(ptr == device_ptr && "wrong pointer mapping");
+
+  device_ptr = omp_get_mapped_ptr(host_data + OFFSET, 0);
+
+  assert(device_ptr && "the pointer with offset should be mapped");
+
+#pragma omp target map(from: ptr)
+  { ptr = host_data + OFFSET; }
+
+  assert(ptr == device_ptr && "wrong pointer mapping");
+
+  return 0;
+}