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

[grokos]

Summary

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

[tianshilei1992]

LGTM

openmp/libomptarget/src/omptarget.cpp

 public:
   /// Constructor
   PrivateArgumentManagerTy(DeviceTy &Dev, AsyncInfoTy &AsyncInfo)
       : Device(Dev), AsyncInfo(AsyncInfo) {}
 
   /// 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",
            (IsFirstPrivate ? "first-" : ""), DPxPTR(HstPtr));
         return OFFLOAD_FAIL;
       }
 #ifdef OMPTARGET_DEBUG
       void *TgtPtrBase = (void *)((intptr_t)TgtPtr + ArgOffset);
       DP("Allocated %" PRId64 " bytes of target memory at " DPxMOD
          " for %sprivate array " DPxMOD " - pushing target argument " DPxMOD
          "\n",
          ArgSize, DPxPTR(TgtPtr), (IsFirstPrivate ? "first-" : ""),
          DPxPTR(HstPtr), DPxPTR(TgtPtrBase));
 #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");
           return OFFLOAD_FAIL;
         }
       }
       TgtPtrs.push_back(TgtPtr);
     } else {
for (int32_t I = 0; I < ArgNum; ++I) {
     bool IsLast, IsHostPtr; // unused.
     if (ArgTypes[I] & OMP_TGT_MAPTYPE_LITERAL) {
       DP("Forwarding first-private value " DPxMOD " to the target construct\n",
          DPxPTR(HstPtrBase));
       TgtPtrBegin = HstPtrBase;
       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));
         return OFFLOAD_FAIL;
       }
     } else {
       if (ArgTypes[I] & OMP_TGT_MAPTYPE_PTR_AND_OBJ)
         HstPtrBase = *reinterpret_cast<void **>(HstPtrBase);

openmp/libomptarget/test/mapping/lambda_by_value.cpp

+// 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;
+}
+