[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.

clang/docs/OpenMPSupport.rst

 
 .. _OpenMP implementation details:
 
 OpenMP 5.0 Implementation Details
 =================================
 
 The following table provides a quick overview over various OpenMP 5.0 features
 and their implementation status. Please post on the
-`Discourse forums (Runtimes - OpenMP category)`_ for more 
+`Discourse forums (Runtimes - OpenMP category)`_ for more
 information or if you want to help with the
 implementation.
 
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 |Category                      | Feature                                                      | Status                   | Reviews                                                               |
 +==============================+==============================================================+==========================+=======================================================================+
 | loop                         | support != in the canonical loop form                        | :good:`done`             | D54441                                                                |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+

 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 
 
 OpenMP 5.1 Implementation Details
 =================================
 
 The following table provides a quick overview over various OpenMP 5.1 features
 and their implementation status, as defined in the technical report 8 (TR8).
-Please post on the 
-`Discourse forums (Runtimes - OpenMP category)`_ for more 
+Please post on the
+`Discourse forums (Runtimes - OpenMP category)`_ for more
 information or if you want to help with the
 implementation.
 
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 |Category                      | Feature                                                      | Status                   | Reviews                                                               |
 +==============================+==============================================================+==========================+=======================================================================+
 | atomic                       | 'compare' clause on atomic construct                         | :good:`done`             | D120290, D120007, D118632, D120200, D116261, D118547, D116637         |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | atomic                       | 'fail' clause on atomic construct                            | :part:`worked on`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | base language                | C++ attribute specifier syntax                               | :good:`done`             | D105648                                                               |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | 'present' map type modifier                                  | :good:`done`             | D83061, D83062, D84422                                                |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | 'present' motion modifier                                    | :good:`done`             | D84711, D84712                                                        |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | 'present' in defaultmap clause                               | :good:`done`             | D92427                                                                |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | map clause reordering reordering based on 'present' modifier | :none:`unclaimed`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | device-specific environment variables                        | :none:`unclaimed`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | omp_target_is_accessible routine                             | :none:`unclaimed`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
-| device                       | omp_get_mapped_ptr routine                                   | :none:`unclaimed`        |                                                                       |
+| device                       | omp_get_mapped_ptr routine                                   | :none:`done`             |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | new async target memory copy routines                        | :none:`unclaimed`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | thread_limit clause on target construct                      | :none:`worked on`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | has_device_addr clause on target construct                   | :none:`unclaimed`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | device                       | iterators in map clause or motion clauses                    | :none:`unclaimed`        |                                                                       |

 | task                         | inoutset in depend clause                                    | :none:`unclaimed`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 | task                         | nowait clause on taskwait                                    | :part:`worked on`        |                                                                       |
 +------------------------------+--------------------------------------------------------------+--------------------------+-----------------------------------------------------------------------+
 
 OpenMP Extensions
 =================
 
-The following table provides a quick overview over various OpenMP 
+The following table provides a quick overview over various OpenMP
 extensions and their implementation status.  These extensions are not
 currently defined by any standard, so links to associated LLVM
 documentation are provided.  As these extensions mature, they will be
 considered for standardization. Please post on the
-`Discourse forums (Runtimes - OpenMP category)`_ to provide feedback. 
+`Discourse forums (Runtimes - OpenMP category)`_ to provide feedback.
 
 +------------------------------+-----------------------------------------------------------------------------------+--------------------------+--------------------------------------------------------+
 |Category                      | Feature                                                                           | Status                   | Reviews                                                |
 +==============================+===================================================================================+==========================+========================================================+
 | atomic extension             | `'atomic' strictly nested within 'teams'                                          | :good:`prototyped`       | D126323                                                |
 |                              | <https://openmp.llvm.org/docs/openacc/OpenMPExtensions.html#atomicWithinTeams>`_  |                          |                                                        |
 +------------------------------+-----------------------------------------------------------------------------------+--------------------------+--------------------------------------------------------+
 | device extension             | `'ompx_hold' map type modifier                                                    | :good:`prototyped`       | D106509, D106510                                       |
 |                              | <https://openmp.llvm.org/docs/openacc/OpenMPExtensions.html#ompx-hold>`_          |                          |                                                        |
 +------------------------------+-----------------------------------------------------------------------------------+--------------------------+--------------------------------------------------------+
 
 .. _Discourse forums (Runtimes - OpenMP category): https://discourse.llvm.org/c/runtimes/openmp/35

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, returning Ptr " DPxMOD ".\n",
+           DeviceNum, DPxPTR(Ptr));
+    return const_cast<void *>(Ptr);
+  }
+
+  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;
+}