[OpenMP] Fixed the issue that target memory deallocation might be called when they're being used

[tianshilei1992]

Summary

This patch fixed the issue that target memory might be deallocated when
they're still being used or before they're used.

[ye-luo]

Thanks for fixing the bug. It should be good for the moment.
When I think about the existence of recursive mapper, we may still have more sync than needed. I think recursion the whole targetDataBegin/targetDataEnd is convenient but sub-optimal choice.
Recursion should only be done on the map/mapper analysis. Just leave my thoughts here. It needs a discussion beyond this patch.

openmp/libomptarget/src/omptarget.cpp

606

If targetDataEnd is not called from target.
AsyncInfo can be nullptr.
AsyncInfo->Queue also be nullptr in "target exit data nowait" if there is no transfer.

At the end of target(), the synchronize needs to check AsyncInfo->Queue.

[tianshilei1992]

openmp/libomptarget/src/omptarget.cpp

606

There is no need to synchronize at the end of target. I forgot to remove it.
For other cases, let me think about them.

[tianshilei1992]

In D84996#2186441, @ye-luo wrote:

Thanks for fixing the bug. It should be good for the moment.
When I think about the existence of recursive mapper, we may still have more sync than needed. I think recursion the whole targetDataBegin/targetDataEnd is convenient but sub-optimal choice.
Recursion should only be done on the map/mapper analysis. Just leave my thoughts here. It needs a discussion beyond this patch.

Recursive mapper does introduce a lot of synchronization overhead as it passes nullptr to those functions, therefore in the plugin, it will first fetch a new stream, enqueue the operation, and synchronize. The whole mapper is using "synchronous" operations.

[tianshilei1992]

Removed the tailing synchronization because it is no use any more

[tianshilei1992]

Fixed issues in comments

[tianshilei1992]

openmp/libomptarget/src/omptarget.cpp

606

I added some check here. Nice catch.

[jdoerfert]

[Drive By]

openmp/libomptarget/src/omptarget.cpp

434

Please use proper names even for argument. Also add doxygen comments to all members.

[tianshilei1992]

Updated the comments and arguments

[tianshilei1992]

"clause" -> "modifier"

[ye-luo]

LGTM

openmp/libomptarget/src/omptarget.cpp

if (arg_types[i] & OMP_TGT_MAPTYPE_PTR_AND_OBJ && !IsHostPtr) {
           HstPtrBase, PointerTgtPtrBegin, TgtPtrBase};
       Device.ShadowMtx.unlock();
     }
   }
 
   return OFFLOAD_SUCCESS;
 }
 
+namespace {
+/// This structure contains information to deallocate a target pointer, aka.
+/// used to call the function \p DeviceTy::deallocTgtPtr.
+struct DeallocTgtPtrInfo {
+  /// Host pointer used to look up into the map table
+  void *HstPtrBegin;
+  /// Size of the data
+  int64_t DataSize;
+  /// Whether it is forced to be removed from the map table
+  bool ForceDelete;
+  /// Whether it has \p close modifier
+  bool HasCloseModifier;
+
+  DeallocTgtPtrInfo(void *HstPtr, int64_t Size, bool ForceDelete,
+                    bool HasCloseModifier)
+      : HstPtrBegin(HstPtr), DataSize(Size), ForceDelete(ForceDelete),
+        HasCloseModifier(HasCloseModifier) {}
+};
+} // namespace
+
 /// Internal function to undo the mapping and retrieve the data from the device.
 int targetDataEnd(DeviceTy &Device, int32_t ArgNum, void **ArgBases,
                   void **Args, int64_t *ArgSizes, int64_t *ArgTypes,
                   void **ArgMappers, __tgt_async_info *AsyncInfo) {
   int Ret;
+  std::vector<DeallocTgtPtrInfo> DeallocTgtPtrs;
   // process each input.
   for (int32_t I = ArgNum - 1; I >= 0; --I) {
     // Ignore private variables and arrays - there is no mapping for them.
     // Also, ignore the use_device_ptr directive, it has no effect here.
     if ((ArgTypes[I] & OMP_TGT_MAPTYPE_LITERAL) ||
         (ArgTypes[I] & OMP_TGT_MAPTYPE_PRIVATE))
       continue;
 
if ((ArgTypes[I] & OMP_TGT_MAPTYPE_FROM) || DelEntry) {
           DP("Removing shadow pointer " DPxMOD "\n", DPxPTR(ShadowHstPtrAddr));
           Itr = Device.ShadowPtrMap.erase(Itr);
         } else {
           ++Itr;
         }
       }
       Device.ShadowMtx.unlock();
 
-      // Deallocate map
-      if (DelEntry) {
-        Ret = Device.deallocTgtPtr(HstPtrBegin, DataSize, ForceDelete,
+      // Add pointer to the buffer for later deallocation
+      if (DelEntry)
+        DeallocTgtPtrs.emplace_back(HstPtrBegin, DataSize, ForceDelete,
                                     HasCloseModifier);
-        if (Ret != OFFLOAD_SUCCESS) {
-          DP("Deallocating data from device failed.\n");
-          return OFFLOAD_FAIL;
-        }
-      }
+    }
+  }
+
+  // We need to synchronize before deallocating data.
+  // If AsyncInfo is nullptr, the previous data transfer (if has) will be
+  // synchronous, so we don't need to synchronize again. If AsyncInfo->Queue is
+  // nullptr, there is no data transfer happened because once there is,
+  // AsyncInfo->Queue will not be nullptr, so again, we don't need to
+  // synchronize.
+  if (AsyncInfo && AsyncInfo->Queue) {
+    Ret = Device.synchronize(AsyncInfo);
+      DP("Failed to synchronize device.\n");
+
+  // Deallocate target pointer
+  for (DeallocTgtPtrInfo &Info : DeallocTgtPtrs) {
+    Ret = Device.deallocTgtPtr(Info.HstPtrBegin, Info.DataSize,
+                               Info.ForceDelete, Info.HasCloseModifier);
+    if (Ret != OFFLOAD_SUCCESS) {
+      DP("Deallocating data from device failed.\n");
+      return OFFLOAD_FAIL;
     }
   }
 
   return OFFLOAD_SUCCESS;
 }
 
 /// Internal function to pass data to/from the target.
 // async_info_ptr is currently unused, added here so target_data_update has the
int target(int64_t DeviceId, void *HostPtr, int32_t ArgNum, void **ArgBases,
   // variables
   Ret = processDataAfter(DeviceId, HostPtr, ArgNum, ArgBases, Args, ArgSizes,
                          ArgTypes, ArgMappers, FPArrays, &AsyncInfo);
   if (Ret != OFFLOAD_SUCCESS) {
     DP("Failed to process data after launching the kernel.\n");
     return OFFLOAD_FAIL;
   }
 
-  return Device.synchronize(&AsyncInfo);
+  return OFFLOAD_SUCCESS;
 }