[libomptarget] Mapper function spliting

[dtalaashrafi]

Summary

This is an implementation of __tgt_target_data_begin_mapper_issue and __tgt_target_data_begin_mapper_wait
functions, and some other required changes.

[dtalaashrafi]

Ping.

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

__OMP_RTL(__tgt_target_kernel_nowait, false, Int32, IdentPtr, Int64, Int32,
           Int32, VoidPtr, KernelArgsPtr, Int32, VoidPtr, Int32, VoidPtr)
 __OMP_RTL(__tgt_register_requires, false, Void, Int64)
 __OMP_RTL(__tgt_target_data_begin_mapper, false, Void, IdentPtr, Int64, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr)
 __OMP_RTL(__tgt_target_data_begin_nowait_mapper, false, Void, IdentPtr, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr)
 __OMP_RTL(__tgt_target_data_begin_mapper_issue, false, Void, IdentPtr, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr, AsyncInfoPtr)
-__OMP_RTL(__tgt_target_data_begin_mapper_wait, false, Void, Int64, AsyncInfoPtr)
+__OMP_RTL(__tgt_target_data_begin_mapper_wait, false, Void, IdentPtr, Int64, AsyncInfoPtr)
 __OMP_RTL(__tgt_target_data_end_mapper, false, Void, IdentPtr, Int64, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr)
 __OMP_RTL(__tgt_target_data_end_nowait_mapper, false, Void, IdentPtr, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr)
 __OMP_RTL(__tgt_target_data_update_mapper, false, Void, IdentPtr, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr)
 __OMP_RTL(__tgt_target_data_update_nowait_mapper, false, Void, IdentPtr, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr)

openmp/libomptarget/include/omptarget.h

 /// mistakes.
 class AsyncInfoTy {
   /// Locations we used in (potentially) asynchronous calls which should live
   /// as long as this AsyncInfoTy object.
   std::deque<void *> BufferLocations;
 
   __tgt_async_info AsyncInfo;
   DeviceTy &Device;
+  // If the flag is set, synchronize is called in the delete function.
+  bool ShouldSync;
 
 public:
-  AsyncInfoTy(DeviceTy &Device) : Device(Device) {}
-  ~AsyncInfoTy() { synchronize(); }
+  AsyncInfoTy(DeviceTy &Device, bool ShouldSync = true)
+      : Device(Device), ShouldSync(ShouldSync) {}
+
+  ~AsyncInfoTy() {
+    if (ShouldSync)
+      synchronize();
+  }
 
   /// Implicit conversion to the __tgt_async_info which is used in the
   /// plugin interface.
   operator __tgt_async_info *() { return &AsyncInfo; }
 
   /// Synchronize all pending actions.
   ///
   /// \returns OFFLOAD_FAIL or OFFLOAD_SUCCESS appropriately.
void __tgt_target_data_begin_nowait(int64_t DeviceId, int32_t ArgNum,
                                     int32_t DepNum, void *DepList,
                                     int32_t NoAliasDepNum,
                                     void *NoAliasDepList);
 void __tgt_target_data_begin_mapper(ident_t *Loc, int64_t DeviceId,
                                     int32_t ArgNum, void **ArgsBase,
                                     void **Args, int64_t *ArgSizes,
                                     int64_t *ArgTypes, map_var_info_t *ArgNames,
                                     void **ArgMappers);
+// Issues the memory transfer asynchronously and sets a handle.
+void __tgt_target_data_begin_mapper_issue(
+    ident_t *Loc, int64_t DeviceId, int32_t ArgNum, void **ArgsBase,
+    void **Args, int64_t *ArgSizes, int64_t *ArgTypes, map_var_info_t *ArgNames,
+    void **ArgMappers, __tgt_async_info *Handle);
+
+// Waits in the set handle from "__tgt_target_data_begin_mapper_issue"
+// for the memory transfer to finish.
+void __tgt_target_data_begin_mapper_wait(ident_t *Loc, int64_t DeviceId,
+                                         __tgt_async_info *Handle);
+
 void __tgt_target_data_begin_nowait_mapper(
     ident_t *Loc, int64_t DeviceId, int32_t ArgNum, void **ArgsBase,
     void **Args, int64_t *ArgSizes, int64_t *ArgTypes, map_var_info_t *ArgNames,
     void **ArgMappers, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
     void *NoAliasDepList);
 
 // passes data from the target, release target memory and destroys the
 // host-target mapping (top entry from the stack of data maps) created by

openmp/libomptarget/src/exports

global:
     __tgt_target;
     __tgt_target_teams;
     __tgt_target_data_begin_nowait;
     __tgt_target_data_end_nowait;
     __tgt_target_data_update_nowait;
     __tgt_target_nowait;
     __tgt_target_teams_nowait;
     __tgt_target_data_begin_mapper;
+    __tgt_target_data_begin_mapper_issue;
+    __tgt_target_data_begin_mapper_wait;
     __tgt_target_data_end_mapper;
     __tgt_target_data_update_mapper;
     __tgt_target_mapper;
     __tgt_target_teams_mapper;
     __tgt_target_data_begin_nowait_mapper;
     __tgt_target_data_end_nowait_mapper;
     __tgt_target_data_update_nowait_mapper;
     __tgt_target_nowait_mapper;

openmp/libomptarget/src/interface.cpp

 /// libomptarget.so internal structure (an entry in a stack of data maps)
 /// and passes the data to the device.
 EXTERN void __tgt_target_data_begin_mapper(ident_t *Loc, int64_t DeviceId,
                                            int32_t ArgNum, void **ArgsBase,
                                            void **Args, int64_t *ArgSizes,
                                            int64_t *ArgTypes,
                                            map_var_info_t *ArgNames,
                                            void **ArgMappers) {
+  __tgt_target_data_begin_mapper_issue(Loc, DeviceId, ArgNum, ArgsBase, Args,
+                                       ArgSizes, ArgTypes, ArgNames, ArgMappers,
+                                       nullptr);
+}
+
+EXTERN void __tgt_target_data_begin_mapper_issue(
+    ident_t *Loc, int64_t DeviceId, int32_t ArgNum, void **ArgsBase,
+    void **Args, int64_t *ArgSizes, int64_t *ArgTypes, map_var_info_t *ArgNames,
+    void **ArgMappers, __tgt_async_info *Handle) {
   TIMESCOPE_WITH_IDENT(Loc);
   DP("Entering data begin region for device %" PRId64 " with %d mappings\n",
      DeviceId, ArgNum);
   if (checkDeviceAndCtors(DeviceId, Loc)) {
     DP("Not offloading to device %" PRId64 "\n", DeviceId);
     return;
   }
 
   DeviceTy &Device = *PM->Devices[DeviceId];
 
   if (getInfoLevel() & OMP_INFOTYPE_KERNEL_ARGS)
     printKernelArguments(Loc, DeviceId, ArgNum, ArgSizes, ArgTypes, ArgNames,
                          "Entering OpenMP data region");
 #ifdef OMPTARGET_DEBUG
   for (int I = 0; I < ArgNum; ++I) {
     DP("Entry %2d: Base=" DPxMOD ", Begin=" DPxMOD ", Size=%" PRId64
        ", Type=0x%" PRIx64 ", Name=%s\n",
        I, DPxPTR(ArgsBase[I]), DPxPTR(Args[I]), ArgSizes[I], ArgTypes[I],
        (ArgNames) ? getNameFromMapping(ArgNames[I]).c_str() : "unknown");
   }
 #endif
 
-  AsyncInfoTy AsyncInfo(Device);
+  AsyncInfoTy AsyncInfo(Device, !(Handle));
   int Rc = targetDataBegin(Loc, Device, ArgNum, ArgsBase, Args, ArgSizes,
                            ArgTypes, ArgNames, ArgMappers, AsyncInfo);
-  if (Rc == OFFLOAD_SUCCESS)
+
+  if (Rc == OFFLOAD_SUCCESS) {
+    if (Handle)
+      Handle->Queue = ((__tgt_async_info *)AsyncInfo)->Queue;
+    else
       Rc = AsyncInfo.synchronize();
+  }
+
   handleTargetOutcome(Rc == OFFLOAD_SUCCESS, Loc);
 }
 
+EXTERN void __tgt_target_data_begin_mapper_wait(ident_t *Loc, int64_t DeviceId,
+                                                __tgt_async_info *Handle) {
+
+  TIMESCOPE_WITH_IDENT(Loc);
+  assert((Handle && Handle->Queue) && "Incomplete data mapping");
+
+  DP("Entering wait function in data begin region for device %" PRId64 "\n",
+     DeviceId);
+
+  if (checkDevice(DeviceId, Loc)) {
+    DP("Not offloading to device %" PRId64 "\n", DeviceId);
+    return;
+  }
+
+  DeviceTy &Device = *PM->Devices[DeviceId];
+  if (Device.RTL->synchronize)
+    Device.RTL->synchronize(DeviceId, Handle);
+}
+
 EXTERN void __tgt_target_data_begin_nowait_mapper(
     ident_t *Loc, int64_t DeviceId, int32_t ArgNum, void **ArgsBase,
     void **Args, int64_t *ArgSizes, int64_t *ArgTypes, map_var_info_t *ArgNames,
     void **ArgMappers, int32_t DepNum, void *DepList, int32_t NoAliasDepNum,
     void *NoAliasDepList) {
   TIMESCOPE_WITH_IDENT(Loc);
 
   __tgt_target_data_begin_mapper(Loc, DeviceId, ArgNum, ArgsBase, Args,

openmp/libomptarget/src/omptarget.cpp

 }
 
 static bool isOffloadDisabled() {
   if (PM->TargetOffloadPolicy == tgt_default)
     handleDefaultTargetOffload();
   return PM->TargetOffloadPolicy == tgt_disabled;
 }
 
-// If offload is enabled, ensure that device DeviceID has been initialized,
-// global ctors have been executed, and global data has been mapped.
-//
-// The return bool indicates if the offload is to the host device
-// There are three possible results:
-// - Return false if the taregt device is ready for offload
-// - Return true without reporting a runtime error if offload is
-//   disabled, perhaps because the initial device was specified.
-// - Report a runtime error and return true.
-//
-// If DeviceID == OFFLOAD_DEVICE_DEFAULT, set DeviceID to the default device.
-// This step might be skipped if offload is disabled.
-bool checkDeviceAndCtors(int64_t &DeviceID, ident_t *Loc) {
+// If offload is enabled, ensure that device DeviceID has been initialized.
+bool checkDevice(int64_t &DeviceID, ident_t *Loc) {
   if (isOffloadDisabled()) {
     DP("Offload is disabled\n");
     return true;
   }
 
   if (DeviceID == OFFLOAD_DEVICE_DEFAULT) {
     DeviceID = omp_get_default_device();
     DP("Use default device id %" PRId64 "\n", DeviceID);
bool checkDevice(int64_t &DeviceID, ident_t *Loc) {
 
   // Is device ready?
   if (!deviceIsReady(DeviceID)) {
     REPORT("Device %" PRId64 " is not ready.\n", DeviceID);
     handleTargetOutcome(false, Loc);
     return true;
   }
 
+  return false;
+}
+
+// If offload is enabled, ensure that device DeviceID has been initialized,
+// global ctors have been executed, and global data has been mapped.
+//
+// The return bool indicates if the offload is to the host device
+// There are three possible results:
+// - Return false if the taregt device is ready for offload
+// - Return true without reporting a runtime error if offload is
+//   disabled, perhaps because the initial device was specified.
+// - Report a runtime error and return true.
+//
+// If DeviceID == OFFLOAD_DEVICE_DEFAULT, set DeviceID to the default device.
+// This step might be skipped if offload is disabled.
+bool checkDeviceAndCtors(int64_t &DeviceID, ident_t *Loc) {
+
+  if(checkDevice(DeviceID, Loc))
+    return true;
+
   // Get device info.
   DeviceTy &Device = *PM->Devices[DeviceID];
 
   // Check whether global data has been mapped for this device
   bool HasPendingGlobals;
   {
     std::lock_guard<decltype(Device.PendingGlobalsMtx)> LG(
         Device.PendingGlobalsMtx);

openmp/libomptarget/src/private.h

 extern int target(ident_t *Loc, DeviceTy &Device, void *HostPtr, int32_t ArgNum,
                   void **ArgBases, void **Args, int64_t *ArgSizes,
                   int64_t *ArgTypes, map_var_info_t *ArgNames,
                   void **ArgMappers, int32_t TeamNum, int32_t ThreadLimit,
                   uint64_t Tripcount, int IsTeamConstruct,
                   AsyncInfoTy &AsyncInfo);
 
 extern void handleTargetOutcome(bool Success, ident_t *Loc);
+extern bool checkDevice(int64_t &DeviceID, ident_t *Loc);
 extern bool checkDeviceAndCtors(int64_t &DeviceID, ident_t *Loc);
 extern void *targetAllocExplicit(size_t Size, int DeviceNum, int Kind,
                                  const char *Name);
 
 // This structure stores information of a mapped memory region.
 struct MapComponentInfoTy {
   void *Base;
   void *Begin;