[OpenMP] Move synchronization into `__tgt_async_info`

[jdoerfert]

Summary

The AsyncInfo should be passed everywhere and it should offer a way to
ensure synchronization, given a libomptarget Device.

This replaces D96431.

[tianshilei1992]

openmp/libomptarget/include/omptarget.h

128

I suggest we wrap this info into another class, and bind the device object such that we will make sure that the async info object will not be synced by a wrong device.

class AsyncInfo {
  __tgt_async_info TheInfo;
  DeviceTy &Device;
public:
  int synchronize();
};

Then __tgt_async_info is only used when we're communicating with plugin. AsyncInfo is used across the whole libomptarget.

[jdoerfert]

Introduced wrapper

[tianshilei1992]

openmp/libomptarget/include/omptarget.h

145

Not sure the wrapper has a bonus to fix issues mentioned in D96431.

openmp/libomptarget/src/omptarget.cpp

1143

The last argument of DeviceTy::submitData is __tgt_async_info *. I'm thinking we can pass AsyncInfoTy & here and all conversion is done in DeviceTy such that all interfaces in libomptarget use same thing.

[jdoerfert]

openmp/libomptarget/include/omptarget.h

128

That's what is happening after this patch and a second to replace the remaining uses of __tgt_async_info in the libomptarget layer. Already with this patch synchronize will work on the right device.

145

I don't understand.

[tianshilei1992]

openmp/libomptarget/include/omptarget.h

145

If we have event, AsyncInfoTy's lifetime can still be out of its scope.

[jdoerfert]

openmp/libomptarget/include/omptarget.h

145

Let's cross that bridge when we get events? I was assuming we copy the __tgt_async_info into the appropriate locations of the runtime and null it in the AsyncInfoTy object. Or something along those lines.

[tianshilei1992]

openmp/libomptarget/include/omptarget.h

145

That should be fine. Actually we cannot use AsyncInfoTy in libomp because it's a C library. If it is a C++ library, we also don't need to worry about that because we can use smart pointer then,

[tianshilei1992]

openmp/libomptarget/src/omptarget.cpp

646

This will cause synchronize to be called twice: one here, and another when AsyncInfo is destroyed. For now __tgt_rtl_synchronize is not designed to be called on a same object twice because it will assert(async_info_ptr->Queue), and every time we call synchronize, Queue will be set to nullptr.

[jdoerfert]

openmp/libomptarget/src/omptarget.cpp

646

AsyncInfoTy::synchronize() checks the Queue and only forwards the call if it is not null. Calling it twice is OK.

[tianshilei1992]

LG

openmp/libomptarget/include/omptarget.h

 
 /// This struct contains the offload entries identified by the target runtime
 struct __tgt_target_table {
   __tgt_offload_entry *EntriesBegin; // Begin of the table with all the entries
   __tgt_offload_entry
       *EntriesEnd; // End of the table with all the entries (non inclusive)
 };
 
+// clang-format on
+
+struct DeviceTy;
+
 /// This struct contains information exchanged between different asynchronous
 /// operations for device-dependent optimization and potential synchronization
 struct __tgt_async_info {
-  // A pointer to a queue-like structure where offloading operations are issued.
-  // We assume to use this structure to do synchronization. In CUDA backend, it
-  // is CUstream.
+  /// Synchronize the queue with \p Device, if necessary, and resets it to
+  /// indicate no pending actions.
+  ///
+  /// \returns OFFLOAD_FAIL or OFFLOAD_SUCCESS appropriately.
+  int synchronize(DeviceTy &Device);
+
+  /// A pointer to a queue-like structure where offloading operations are
+  /// issued. We assume to use this structure to do synchronization. In CUDA
+  /// backend, it is CUstream. If it is a nullptr there are no pending actions
+  /// that are not synchronized yet.
   void *Queue = nullptr;
 };
 
 /// This struct is a record of non-contiguous information
 struct __tgt_target_non_contig {
   uint64_t Offset;
   uint64_t Count;
   uint64_t Stride;
 };
 
-// clang-format on
-
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 int omp_get_num_devices(void);
 int omp_get_initial_device(void);
 void *omp_target_alloc(size_t size, int device_num);
 void omp_target_free(void *device_ptr, int device_num);

openmp/libomptarget/src/omptarget.cpp

 #include "omptarget.h"
 #include "device.h"
 #include "private.h"
 #include "rtl.h"
 
 #include <cassert>
 #include <vector>
 
+int __tgt_async_info::synchronize(DeviceTy &Device) {
+  int Result = OFFLOAD_SUCCESS;
+  if (Queue) {
+    // If we have a queue we need to synchronize it now. As it is synchronized
+    // afterwards we can get rid of the queue indicating no pending actions.
+    Result = Device.synchronize(this);
+    Queue = nullptr;
+  }
+  return Result;
+}
+
 /* All begin addresses for partially mapped structs must be 8-aligned in order
  * to ensure proper alignment of members. E.g.
  *
  * struct S {
  *   int a;   // 4-aligned
  *   int b;   // 4-aligned
  *   int *p;  // 8-aligned
  * } s1;
struct DeallocTgtPtrInfo {
   /// 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) {}
 };
-
-/// Synchronize device
-static int syncDevice(DeviceTy &Device, __tgt_async_info *AsyncInfo) {
-  assert(AsyncInfo && AsyncInfo->Queue && "Invalid AsyncInfo");
-  if (Device.synchronize(AsyncInfo) != OFFLOAD_SUCCESS) {
-    REPORT("Failed to synchronize device.\n");
-    return OFFLOAD_FAIL;
-  }
-
-  return OFFLOAD_SUCCESS;
-}
 } // namespace
 
 /// Internal function to undo the mapping and retrieve the data from the device.
 int targetDataEnd(ident_t *loc, DeviceTy &Device, int32_t ArgNum,
                   void **ArgBases, void **Args, int64_t *ArgSizes,
                   int64_t *ArgTypes, map_var_info_t *ArgNames,
                   void **ArgMappers, __tgt_async_info *AsyncInfo) {
   int Ret;
int targetDataEnd(ident_t *loc, DeviceTy &Device, int32_t ArgNum,
   }
 
   // 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 = syncDevice(Device, AsyncInfo);
+  if (AsyncInfo) {
+    Ret = AsyncInfo->synchronize(Device);
     if (Ret != OFFLOAD_SUCCESS)
       return OFFLOAD_FAIL;
   }
 
   // Deallocate target pointer
   for (DeallocTgtPtrInfo &Info : DeallocTgtPtrs) {
     Ret = Device.deallocTgtPtr(Info.HstPtrBegin, Info.DataSize,
                                Info.ForceDelete, Info.HasCloseModifier);
if (!(ArgTypes[I] & OMP_TGT_MAPTYPE_TARGET_PARAM)) {
           DP("Unified memory is active, no need to map lambda captured"
              "variable (" DPxMOD ")\n",
              DPxPTR(HstPtrVal));
           continue;
         }
         DP("Update lambda reference (" DPxMOD ") -> [" DPxMOD "]\n",
            DPxPTR(PointerTgtPtrBegin), DPxPTR(TgtPtrBegin));
         Ret = Device.submitData(TgtPtrBegin, &PointerTgtPtrBegin,
                                 sizeof(void *), AsyncInfo);
         if (Ret != OFFLOAD_SUCCESS) {
           REPORT("Copying data to device failed.\n");
           return OFFLOAD_FAIL;
         }
       }
       continue;
     }
     void *HstPtrBegin = Args[I];
if (ArgNum) {
     // variables
     Ret = processDataAfter(loc, DeviceId, HostPtr, ArgNum, ArgBases, Args,
                            ArgSizes, ArgTypes, ArgNames, ArgMappers,
                            PrivateArgumentManager, AsyncInfo);
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Failed to process data after launching the kernel.\n");
       return OFFLOAD_FAIL;
     }
-  } else if (AsyncInfo->Queue) {
+  } else {
+    // TODO: We should not synchronize here but on the outer level once we pass
+    // in a reference AsyncInfo object.
     // If ArgNum is zero, but AsyncInfo.Queue is valid, then the kernel doesn't
     // hava any argument, and the device supports async operations, so we need a
     // sync at this point.
-    return syncDevice(Device, AsyncInfo);
+    return AsyncInfo->synchronize(Device);
   }
 
   return OFFLOAD_SUCCESS;
 }