[Openmp-commits] [PATCH] D102890: [libomptarget] Fix a bug whereby firstprivates are not copied over to the device

[George Rokos via Phabricator via Openmp-commits]

grokos created this revision.
grokos added reviewers: ABataev, jdoerfert.
grokos added a project: OpenMP.
grokos requested review of this revision.
Herald added a subscriber: sstefan1.

After D86781 <https://reviews.llvm.org/D86781>, the check for the `TO` flag when processing firstprivates is missing. As a result, sometimes the device copy of a firstprivate never gets initialized.

More specifically, D86781 <https://reviews.llvm.org/D86781> tried to force lambda structs to be allocated immediately by marking them as a non-firstprivate, so that `PrivateArgumentManagerTy::addArg` allocates memory for them immediately. However, calling `addArg` with `IsFirstPrivate=false` makes the function skip initializing the device copy. Whether an argument is firstprivate and whether we need to allocate memory immediately are not synonyms, so this patch introduces one more control variable for immediate allocation and sets it apart from initialization.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D102890

Files:
  openmp/libomptarget/src/omptarget.cpp
  openmp/libomptarget/test/mapping/lambda_by_value.cpp


Index: openmp/libomptarget/test/mapping/lambda_by_value.cpp
===================================================================
--- /dev/null
+++ openmp/libomptarget/test/mapping/lambda_by_value.cpp
@@ -0,0 +1,34 @@
+// RUN: %libomptarget-compilexx-run-and-check-generic
+
+#include <stdio.h>
+#include <stdint.h>
+
+// CHECK: before: [[V1:111]] [[V2:222]] [[PX:0x[^ ]+]] [[PY:0x[^ ]+]]
+// CHECK: lambda: [[V1]] [[V2]] [[PX_TGT:0x[^ ]+]] 0x{{.*}}
+// CHECK: tgt   : [[V2]] [[PX_TGT]] 1
+// CHECK: out   : [[V2]] [[V2]] [[PX]] [[PY]]
+
+int main() {
+  int x[10];
+  long y[8];
+  x[1] = 111;
+  y[1] = 222;
+
+  auto lambda = [&x, y]() {
+    printf("lambda: %d %ld %p %p\n", x[1], y[1], &x[0], &y[0]);
+    x[1] = y[1];
+  };
+
+  printf("before: %d %ld %p %p\n", x[1], y[1], &x[0], &y[0]);
+
+  intptr_t xp = (intptr_t) &x[0];
+#pragma omp target firstprivate(xp)
+  {
+    lambda();
+    printf("tgt   : %d %p %d\n", x[1], &x[0], (&x[0] != (int*) xp));
+  }
+  printf("out   : %d %ld %p %p\n", x[1], y[1], &x[0], &y[0]);
+
+  return 0;
+}
+
Index: openmp/libomptarget/src/omptarget.cpp
===================================================================
--- openmp/libomptarget/src/omptarget.cpp
+++ openmp/libomptarget/src/omptarget.cpp
@@ -1127,11 +1127,13 @@
   /// Add a private argument
   int addArg(void *HstPtr, int64_t ArgSize, int64_t ArgOffset,
              bool IsFirstPrivate, void *&TgtPtr, int TgtArgsIndex,
-             const map_var_info_t HstPtrName = nullptr) {
+             const map_var_info_t HstPtrName = nullptr,
+             const bool AllocImmediately = false) {
     // If the argument is not first-private, or its size is greater than a
     // predefined threshold, we will allocate memory and issue the transfer
     // immediately.
-    if (ArgSize > FirstPrivateArgSizeThreshold || !IsFirstPrivate) {
+    if (ArgSize > FirstPrivateArgSizeThreshold || !IsFirstPrivate ||
+        AllocImmediately) {
       TgtPtr = Device.allocData(ArgSize, HstPtr);
       if (!TgtPtr) {
         DP("Data allocation for %sprivate array " DPxMOD " failed.\n",
@@ -1148,6 +1150,7 @@
 #endif
       // If first-private, copy data from host
       if (IsFirstPrivate) {
+        DP("Submitting firstprivate data to the device.\n");
         int Ret = Device.submitData(TgtPtr, HstPtr, ArgSize, AsyncInfo);
         if (Ret != OFFLOAD_SUCCESS) {
           DP("Copying data to device failed, failed.\n");
@@ -1325,13 +1328,16 @@
       TgtBaseOffset = 0;
     } else if (ArgTypes[I] & OMP_TGT_MAPTYPE_PRIVATE) {
       TgtBaseOffset = (intptr_t)HstPtrBase - (intptr_t)HstPtrBegin;
-      // Can be marked for optimization if the next argument(s) do(es) not
-      // depend on this one.
-      const bool IsFirstPrivate =
-          (I >= ArgNum - 1 || !(ArgTypes[I + 1] & OMP_TGT_MAPTYPE_MEMBER_OF));
+      const bool IsFirstPrivate = (ArgTypes[I] & OMP_TGT_MAPTYPE_TO);
+      // If there is a next argument and it depends on the current one, we need
+      // to allocate the private memory immediately. If this is not the case,
+      // then the argument can be marked for optimization and packed with the
+      // other privates.
+      const bool AllocImmediately =
+          (I < ArgNum - 1 && (ArgTypes[I + 1] & OMP_TGT_MAPTYPE_MEMBER_OF));
       Ret = PrivateArgumentManager.addArg(
           HstPtrBegin, ArgSizes[I], TgtBaseOffset, IsFirstPrivate, TgtPtrBegin,
-          TgtArgs.size(), HstPtrName);
+          TgtArgs.size(), HstPtrName, AllocImmediately);
       if (Ret != OFFLOAD_SUCCESS) {
         REPORT("Failed to process %sprivate argument " DPxMOD "\n",
                (IsFirstPrivate ? "first-" : ""), DPxPTR(HstPtrBegin));


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D102890.346887.patch
Type: text/x-patch
Size: 3675 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/openmp-commits/attachments/20210520/97eab219/attachment.bin>