[offload][openMP] Add __omp_get_device_ptr_if_present api to offload

[skc7]

This PR adds "__omp_get_device_ptr_if_present" to offload libomptraget.

"__omp_get_device_ptr_if_present" returns a device pointer for data that is mapped to the specified target device. If the data is not present on the device, it returns the original pointer unchanged.

pre-requisite for PR #140523

[mjklemm]

Please do not use assert, but make this a hard failure with exit(1) or return 1. With the assert the test will depend on whether assertions have been enabled, and the test should correctly test for all possible build configurations.

[skc7]

Updated the tests to return 1 on failure, following what others tests in offload do.

[mjklemm]

This routine is not part of the OpenMP API. Please use a different prefix. If the routine only is only to be called from generated code, it may be prudent to prefix it with __ to put into the compiler's "namespace".

[skc7]

Name changed to __omp_get_device_ptr_if_present. Thanks

[llvmbot]

@llvm/pr-subscribers-offload

Author: Chaitanya (skc7)

Changes

This PR adds "omp_get_device_ptr_if_present" to offload libomptraget.

"omp_get_device_ptr_if_present" returns a device pointer for data that is mapped to the specified target device. If the data is not present on the device, it returns the original pointer unchanged.

pre-requisite for PR #140523

---

Full diff: https://github.com/llvm/llvm-project/pull/153146.diff

3 Files Affected:

• (modified) offload/libomptarget/OpenMP/API.cpp (+17)
• (modified) offload/libomptarget/exports (+1)
• (added) offload/test/api/omp_get_device_ptr_if_present.c (+127)

diff --git a/offload/libomptarget/OpenMP/API.cpp b/offload/libomptarget/OpenMP/API.cpp
index 4576f9bd06121..6023b7d691763 100644
--- a/offload/libomptarget/OpenMP/API.cpp
+++ b/offload/libomptarget/OpenMP/API.cpp
@@ -683,3 +683,20 @@ EXTERN void *omp_get_mapped_ptr(const void *Ptr, int DeviceNum) {
 
   return TPR.TargetPointer;
 }
+
+EXTERN void *__omp_get_device_ptr_if_present(void *Ptr, int DeviceNum) {
+  if (!Ptr) {
+    return nullptr;
+  }
+  // Validate device number
+  if (DeviceNum < 0 || DeviceNum >= omp_get_num_devices()) {
+    return Ptr;
+  }
+  // If not present on the device, return the Ptr.
+  if (!omp_target_is_present(Ptr, DeviceNum)) {
+    return Ptr;
+  }
+  // Get the mapped device pointer
+  void *DevicePtr = omp_get_mapped_ptr(Ptr, DeviceNum);
+  return DevicePtr ? DevicePtr : Ptr;
+}
diff --git a/offload/libomptarget/exports b/offload/libomptarget/exports
index 2406776c1fb5f..f4dc2c99cb6ad 100644
--- a/offload/libomptarget/exports
+++ b/offload/libomptarget/exports
@@ -40,6 +40,7 @@ VERS1.0 {
     omp_get_num_devices;
     omp_get_device_num;
     omp_get_initial_device;
+    omp_get_device_ptr_if_present;
     omp_target_alloc;
     omp_target_free;
     omp_target_is_present;
diff --git a/offload/test/api/omp_get_device_ptr_if_present.c b/offload/test/api/omp_get_device_ptr_if_present.c
new file mode 100644
index 0000000000000..e80b054da00ec
--- /dev/null
+++ b/offload/test/api/omp_get_device_ptr_if_present.c
@@ -0,0 +1,127 @@
+// RUN: %libomptarget-compile-run-and-check-generic
+// REQUIRES: unified_shared_memory
+
+#include <omp.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define N 1024
+#define OFFSET 16
+
+int main(int argc, char *argv[]) {
+  // Check if we have any devices
+  if (omp_get_num_devices() == 0) {
+    printf("SKIP: No devices available\n");
+    return 0;
+  }
+
+  int *host_data = (int *)malloc(sizeof(int) * N);
+  if (!host_data) {
+    printf("FAIL: Memory allocation failed\n");
+    return 1;
+  }
+
+  int device_num = omp_get_default_device();
+
+  // Initialize data
+  for (int i = 0; i < N; i++) {
+    host_data[i] = i;
+  }
+
+  // Test 1: NULL pointer should return NULL
+  void *result = __omp_get_device_ptr_if_present(NULL, device_num);
+  if (result != NULL) {
+    printf("FAIL: NULL input should return NULL\n");
+    free(host_data);
+    return 1;
+  }
+
+  // Test 2: Invalid device number should return original pointer
+  result = __omp_get_device_ptr_if_present(host_data, -1);
+  if (result != host_data) {
+    printf("FAIL: Invalid device should return original pointer\n");
+    free(host_data);
+    return 1;
+  }
+
+  result =
+      __omp_get_device_ptr_if_present(host_data, omp_get_num_devices() + 1);
+  if (result != host_data) {
+    printf("FAIL: Out of range device should return original pointer\n");
+    free(host_data);
+    return 1;
+  }
+
+  // Test 3: Unmapped pointer should return original pointer
+  result = __omp_get_device_ptr_if_present(host_data, device_num);
+  if (result != host_data) {
+    printf("FAIL: Unmapped pointer should return original\n");
+    free(host_data);
+    return 1;
+  }
+
+  // Map data to device
+#pragma omp target enter data map(to : host_data[:N])
+
+  // Test 4: Mapped pointer should return device pointer
+  result = __omp_get_device_ptr_if_present(host_data, device_num);
+  void *expected_device_ptr = omp_get_mapped_ptr(host_data, device_num);
+  if (result != expected_device_ptr) {
+    printf("FAIL: Should return device pointer for mapped data\n");
+#pragma omp target exit data map(delete : host_data[:N])
+    free(host_data);
+    return 1;
+  }
+
+  // Test 5: Pointer with offset should also work
+  result = __omp_get_device_ptr_if_present(host_data + OFFSET, device_num);
+  expected_device_ptr = omp_get_mapped_ptr(host_data + OFFSET, device_num);
+  if (result != expected_device_ptr) {
+    printf("FAIL: Should handle offset correctly\n");
+#pragma omp target exit data map(delete : host_data[:N])
+    free(host_data);
+    return 1;
+  }
+
+  // Test 6: Verify device pointer works in target region
+  void *ptr_from_device = NULL;
+  void *ptr_from_api = __omp_get_device_ptr_if_present(host_data, device_num);
+
+#pragma omp target map(from : ptr_from_device)
+  { ptr_from_device = host_data; }
+
+  if (ptr_from_device != ptr_from_api) {
+    printf("FAIL: Device pointer should match in target region\n");
+#pragma omp target exit data map(delete : host_data[:N])
+    free(host_data);
+    return 1;
+  }
+
+  // Unmap data
+#pragma omp target exit data map(delete : host_data[:N])
+
+  // Test 7: After unmapping, should return original pointer again
+  result = __omp_get_device_ptr_if_present(host_data, device_num);
+  if (result != host_data) {
+    printf("FAIL: After unmapping should return original pointer\n");
+    free(host_data);
+    return 1;
+  }
+
+  // Test 8: Already device pointer scenario
+  void *device_alloc = omp_target_alloc(sizeof(int) * N, device_num);
+  if (device_alloc) {
+    result = __omp_get_device_ptr_if_present(device_alloc, device_num);
+    if (result != device_alloc) {
+      printf("FAIL: Device-allocated pointer should return as-is\n");
+      omp_target_free(device_alloc, device_num);
+      free(host_data);
+      return 1;
+    }
+    omp_target_free(device_alloc, device_num);
+  }
+
+  free(host_data);
+  printf("PASS\n");
+  return 0;
+}

[skc7]

Hi @mjklemm
I'm re-evaluating if this new utility is actually required. Given a host pointer, making a call to omp_get_mapped_ptr from mlir, will do the work for a given valid device number(queried from omp.target) in #140523
Closing this PR now. Will re-open this if really required.