[OpenMP][Offload] Automatically map indirect function pointers (#71462)

We already have all the information to automatically map function pointers that have been declared as `indirect` declare target by the user. This is just enabling and testing the functionality by looking through the one level of indirection.
llvm · Nov 7, 2023 · 2d739f1 · 2d739f1
1 parent 7b9d73c
commit 2d739f1
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 30 deletions.
diff --git a/openmp/libomptarget/src/omptarget.cpp b/openmp/libomptarget/src/omptarget.cpp
@@ -191,38 +191,50 @@ static int initLibrary(DeviceTy &Device) {
                                *EntryDeviceEnd = TargetTable->EntriesEnd;
            CurrDeviceEntry != EntryDeviceEnd;
            CurrDeviceEntry++, CurrHostEntry++) {
-        if (CurrDeviceEntry->size != 0) {
-          // has data.
-          assert(CurrDeviceEntry->size == CurrHostEntry->size &&
-                 "data size mismatch");
-
-          // Fortran may use multiple weak declarations for the same symbol,
-          // therefore we must allow for multiple weak symbols to be loaded from
-          // the fat binary. Treat these mappings as any other "regular"
-          // mapping. Add entry to map.
-          if (Device.getTgtPtrBegin(HDTTMap, CurrHostEntry->addr,
-                                    CurrHostEntry->size))
-            continue;
-
-          DP("Add mapping from host " DPxMOD " to device " DPxMOD
-             " with size %zu"
-             "\n",
-             DPxPTR(CurrHostEntry->addr), DPxPTR(CurrDeviceEntry->addr),
-             CurrDeviceEntry->size);
-          HDTTMap->emplace(new HostDataToTargetTy(
-              (uintptr_t)CurrHostEntry->addr /*HstPtrBase*/,
-              (uintptr_t)CurrHostEntry->addr /*HstPtrBegin*/,
-              (uintptr_t)CurrHostEntry->addr +
-                  CurrHostEntry->size /*HstPtrEnd*/,
-              (uintptr_t)CurrDeviceEntry->addr /*TgtAllocBegin*/,
-              (uintptr_t)CurrDeviceEntry->addr /*TgtPtrBegin*/,
-              false /*UseHoldRefCount*/, CurrHostEntry->name,
-              true /*IsRefCountINF*/));
-
-          // Notify about the new mapping.
-          if (Device.notifyDataMapped(CurrHostEntry->addr, CurrHostEntry->size))
+        if (CurrDeviceEntry->size == 0)
+          continue;
+
+        assert(CurrDeviceEntry->size == CurrHostEntry->size &&
+               "data size mismatch");
+
+        // Fortran may use multiple weak declarations for the same symbol,
+        // therefore we must allow for multiple weak symbols to be loaded from
+        // the fat binary. Treat these mappings as any other "regular"
+        // mapping. Add entry to map.
+        if (Device.getTgtPtrBegin(HDTTMap, CurrHostEntry->addr,
+                                  CurrHostEntry->size))
+          continue;
+
+        void *CurrDeviceEntryAddr = CurrDeviceEntry->addr;
+
+        // For indirect mapping, follow the indirection and map the actual
+        // target.
+        if (CurrDeviceEntry->flags & OMP_DECLARE_TARGET_INDIRECT) {
+          AsyncInfoTy AsyncInfo(Device);
+          void *DevPtr;
+          Device.retrieveData(&DevPtr, CurrDeviceEntryAddr, sizeof(void *),
+                              AsyncInfo);
+          if (AsyncInfo.synchronize() != OFFLOAD_SUCCESS)
             return OFFLOAD_FAIL;
+          CurrDeviceEntryAddr = DevPtr;
         }
+
+        DP("Add mapping from host " DPxMOD " to device " DPxMOD " with size %zu"
+           ", name \"%s\"\n",
+           DPxPTR(CurrHostEntry->addr), DPxPTR(CurrDeviceEntry->addr),
+           CurrDeviceEntry->size, CurrDeviceEntry->name);
+        HDTTMap->emplace(new HostDataToTargetTy(
+            (uintptr_t)CurrHostEntry->addr /*HstPtrBase*/,
+            (uintptr_t)CurrHostEntry->addr /*HstPtrBegin*/,
+            (uintptr_t)CurrHostEntry->addr + CurrHostEntry->size /*HstPtrEnd*/,
+            (uintptr_t)CurrDeviceEntryAddr /*TgtAllocBegin*/,
+            (uintptr_t)CurrDeviceEntryAddr /*TgtPtrBegin*/,
+            false /*UseHoldRefCount*/, CurrHostEntry->name,
+            true /*IsRefCountINF*/));
+
+        // Notify about the new mapping.
+        if (Device.notifyDataMapped(CurrHostEntry->addr, CurrHostEntry->size))
+          return OFFLOAD_FAIL;
       }
     }
   }

diff --git a/openmp/libomptarget/test/offloading/indirect_fp_mapping.c b/openmp/libomptarget/test/offloading/indirect_fp_mapping.c
@@ -0,0 +1,37 @@
+// RUN: %libomptarget-compile-generic -fopenmp-version=51
+// RUN: %libomptarget-run-generic | %fcheck-generic
+// RUN: %libomptarget-compileopt-generic -fopenmp-version=51
+// RUN: %libomptarget-run-generic | %fcheck-generic
+
+#include <stdio.h>
+
+int square(int x) { return x * x; }
+#pragma omp declare target indirect to(square)
+
+typedef int (*fp_t)(int);
+
+int main() {
+  int i = 17, r;
+
+  fp_t fp = &square;
+  // CHECK: host: &square =
+  printf("host: &square = %p\n", fp);
+
+#pragma omp target map(from : fp)
+  fp = &square;
+  // CHECK: device: &square = [[DEV_FP:.*]]
+  printf("device: &square = %p\n", fp);
+
+  fp_t fp1 = square;
+  fp_t fp2 = 0;
+#pragma omp target map(from : fp2)
+  fp2 = fp1;
+  // CHECK: device: fp2 = [[DEV_FP]]
+  printf("device: fp2 = %p\n", fp2);
+
+#pragma omp target map(from : r)
+  { r = fp1(i); }
+
+  // CHECK: 17*17 = 289
+  printf("%i*%i = %i\n", i, i, r);
+}