[Openmp-commits] [openmp] d0bc04d - [libomptarget] Fix a bug whereby firstprivates are not copied over to the device

[George Rokos via Openmp-commits]

Author: George Rokos
Date: 2021-05-21T10:52:08-07:00
New Revision: d0bc04d6b91d6bc5f15c981da9d2b911fb578c59

URL: https://github.com/llvm/llvm-project/commit/d0bc04d6b91d6bc5f15c981da9d2b911fb578c59
DIFF: https://github.com/llvm/llvm-project/commit/d0bc04d6b91d6bc5f15c981da9d2b911fb578c59.diff

LOG: [libomptarget] Fix a bug whereby firstprivates are not copied over to the device

The check for the TO flag when processing firstprivates is missing. As a result,
sometimes the device copy of a firstprivate never gets initialized. Currectly we
try 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.

Differential Revision: https://reviews.llvm.org/D102890

Added: 
    openmp/libomptarget/test/mapping/lambda_by_value.cpp

Modified: 
    openmp/libomptarget/src/omptarget.cpp

Removed: 
    


################################################################################
diff  --git a/openmp/libomptarget/src/omptarget.cpp b/openmp/libomptarget/src/omptarget.cpp
index dfa1686e84e7..447ad73c601b 100644
--- a/openmp/libomptarget/src/omptarget.cpp
+++ b/openmp/libomptarget/src/omptarget.cpp
@@ -1127,11 +1127,13 @@ class PrivateArgumentManagerTy {
   /// 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 @@ class PrivateArgumentManagerTy {
 #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 @@ static int processDataBefore(ident_t *loc, int64_t DeviceId, void *HostPtr,
       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));

diff  --git a/openmp/libomptarget/test/mapping/lambda_by_value.cpp b/openmp/libomptarget/test/mapping/lambda_by_value.cpp
new file mode 100644
index 000000000000..03b044e271c3
--- /dev/null
+++ b/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;
+}
+