[Libomptarget] Use new tripcount argument in the runtime.

[jhuber6]

Summary

The previous patch added an argument to the __tgt_target_kernel
runtime function which includes the tripcount used for the loop clause.
This was originally passed in via the __kmpc_push_target_tripcount
function. Now we move this logic to the kernel launch itself and remove
the need for the push function.

Depends on D128816

[jhuber6]

Rebasing and making special tripcuont -1 for the legacy API to use. If we have this tripcount we don't set it in target and instead expect it to be set by __kmpc_push_target_tripcount

[jhuber6]

Adding a static check on the struct size.

[jdoerfert]

openmp/libomptarget/src/interface.cpp

388

Doesn't target already do this? And we should get rid of this in a follow up.

[jhuber6]

openmp/libomptarget/src/interface.cpp

388

Oh yeah, I did it here and forgot to delete it after I decided to move it. I'll get rid of it. And why do we get rid of it in a follow-up?

[jdoerfert]

openmp/libomptarget/src/interface.cpp

388

we don't need this mutex and stuff anymore. Now the information is passed right into the target launch and not separately. See below.

openmp/libomptarget/src/omptarget.cpp

1557

And then delete everything that has to do with Device.LoopTripCnt

[jdoerfert]

OK, keep it till we branched 15 but don't use it when we come from the new interface.

[jhuber6]

Updating

[jdoerfert]

LG, one typo, see below.

openmp/libomptarget/src/omptarget.cpp

1515

openmp/libomptarget/include/omptarget.h

 /// This struct contains all of the arguments to a target kernel region launch.
 struct __tgt_kernel_arguments {
   int32_t NumArgs;    // Number of arguments in each input pointer.
   void **ArgBasePtrs; // Base pointer of each argument (e.g. a struct).
   void **ArgPtrs;     // Pointer to the argument data.
   int64_t *ArgSizes;  // Size of the argument data in bytes.
   int64_t *ArgTypes;  // Type of the data (e.g. to / from).
   void **ArgNames;    // Name of the data for debugging, possibly null.
-  void **ArgMappers;  // User-defined mappers, possible null.
+  void **ArgMappers;  // User-defined mappers, possibly null.
+  int64_t Tripcount;  // Tripcount for the teams / distribute loop, 0 otherwise.
 };
 
 /// This struct is a record of an entry point or global. For a function
 /// entry point the size is expected to be zero
 struct __tgt_offload_entry {
   void *addr;   // Pointer to the offload entry info (function or global)
   char *name;   // Name of the function or global
   size_t size;  // Size of the entry info (0 if it is a function)

openmp/libomptarget/src/interface.cpp

 }
 
 EXTERN int __tgt_target_mapper(ident_t *loc, int64_t device_id, void *host_ptr,
                                int32_t arg_num, void **args_base, void **args,
                                int64_t *arg_sizes, int64_t *arg_types,
                                map_var_info_t *arg_names, void **arg_mappers) {
   TIMESCOPE_WITH_IDENT(loc);
   __tgt_kernel_arguments Args{arg_num,   args_base, args,       arg_sizes,
-                               arg_types, arg_names, arg_mappers};
   return __tgt_target_kernel(loc, device_id, -1, -1, host_ptr, &Args);
 }
 
 EXTERN int __tgt_target_nowait_mapper(
     ident_t *loc, int64_t device_id, void *host_ptr, int32_t arg_num,
     void **args_base, void **args, int64_t *arg_sizes, int64_t *arg_types,
     map_var_info_t *arg_names, void **arg_mappers, int32_t depNum,
     void *depList, int32_t noAliasDepNum, void *noAliasDepList) {
EXTERN int __tgt_target_teams_mapper(ident_t *loc, int64_t device_id,
                                      void **args_base, void **args,
                                      int64_t *arg_sizes, int64_t *arg_types,
                                      map_var_info_t *arg_names,
                                      void **arg_mappers, int32_t team_num,
                                      int32_t thread_limit) {
   TIMESCOPE_WITH_IDENT(loc);
 
   __tgt_kernel_arguments Args{arg_num,   args_base, args,       arg_sizes,
-                               arg_types, arg_names, arg_mappers};
   return __tgt_target_kernel(loc, device_id, team_num, thread_limit, host_ptr,
                              &Args);
 }
 
 EXTERN int __tgt_target_teams_nowait_mapper(
     ident_t *loc, int64_t device_id, void *host_ptr, int32_t arg_num,
     void **args_base, void **args, int64_t *arg_sizes, int64_t *arg_types,
     map_var_info_t *arg_names, void **arg_mappers, int32_t team_num,

 #endif
 
   bool IsTeams = NumTeams != -1;
   if (!IsTeams) {
     NumTeams = 0;
     ThreadLimit = 0;
   }
 
+  if (Args->Tripcount != 0) {
+    PM->TblMapMtx.lock();
+    PM->Devices[DeviceId]->LoopTripCnt.emplace(__kmpc_global_thread_num(Loc),
+                                               Args->Tripcount);
+    PM->TblMapMtx.unlock();
+  }
+
   DeviceTy &Device = *PM->Devices[DeviceId];
   AsyncInfoTy AsyncInfo(Device);
   int rc = target(Loc, Device, HostPtr, Args->NumArgs, Args->ArgBasePtrs,
                   Args->ArgPtrs, Args->ArgSizes, Args->ArgTypes, Args->ArgNames,
-                  Args->ArgMappers, NumTeams, ThreadLimit, IsTeams, AsyncInfo);
+                  Args->ArgMappers, NumTeams, ThreadLimit, Args->Tripcount,
+                  IsTeams, AsyncInfo);
   if (rc == OFFLOAD_SUCCESS)
     rc = AsyncInfo.synchronize();
   handleTargetOutcome(rc == OFFLOAD_SUCCESS, Loc);
   assert(rc == OFFLOAD_SUCCESS && "__tgt_target_kernel unexpected failure!");
   return OMP_TGT_SUCCESS;
 }
 
 EXTERN int __tgt_target_kernel_nowait(

openmp/libomptarget/src/omptarget.cpp

static int InitLibrary(DeviceTy &Device) {
   if (!Device.PendingCtorsDtors.empty()) {
     AsyncInfoTy AsyncInfo(Device);
     // Call all ctors for all libraries registered so far
     for (auto &lib : Device.PendingCtorsDtors) {
       if (!lib.second.PendingCtors.empty()) {
         DP("Has pending ctors... call now\n");
         for (auto &entry : lib.second.PendingCtors) {
           void *ctor = entry;
-          int rc =
-              target(nullptr, Device, ctor, 0, nullptr, nullptr, nullptr,
-                     nullptr, nullptr, nullptr, 1, 1, true /*team*/, AsyncInfo);
+          int rc = target(nullptr, Device, ctor, 0, nullptr, nullptr, nullptr,
+                          nullptr, nullptr, nullptr, 1, 1, 0, true /*team*/,
+                          AsyncInfo);
           if (rc != OFFLOAD_SUCCESS) {
             REPORT("Running ctor " DPxMOD " failed.\n", DPxPTR(ctor));
             return OFFLOAD_FAIL;
           }
         }
         // Clear the list to indicate that this device has been used
         lib.second.PendingCtors.clear();
         DP("Done with pending ctors for lib " DPxMOD "\n", DPxPTR(lib.first));

 } // namespace
 
 /// performs the same actions as data_begin in case arg_num is
 /// non-zero and initiates run of the offloaded region on the target platform;
 /// if arg_num is non-zero after the region execution is done it also
 /// performs the same action as data_update and data_end above. This function
 /// returns 0 if it was able to transfer the execution to a target and an
 /// integer different from zero otherwise.
-int target(ident_t *loc, DeviceTy &Device, void *HostPtr, int32_t ArgNum,
+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, int IsTeamConstruct, AsyncInfoTy &AsyncInfo) {
+           int32_t ThreadLimit, uint64_t Tripcount, int IsTeamConstruct,
+           AsyncInfoTy &AsyncInfo) {
   int32_t DeviceId = Device.DeviceID;
 
   TableMap *TM = getTableMap(HostPtr);
   // No map for this host pointer found!
   if (!TM) {
     REPORT("Host ptr " DPxMOD " does not have a matching target pointer.\n",
            DPxPTR(HostPtr));
     return OFFLOAD_FAIL;
   }
 
   // get target table.
   __tgt_target_table *TargetTable = nullptr;
   {
     std::lock_guard<std::mutex> TrlTblLock(PM->TrlTblMtx);
     assert(TM->Table->TargetsTable.size() > (size_t)DeviceId &&
            "Not expecting a device ID outside the table's bounds!");
     TargetTable = TM->Table->TargetsTable[DeviceId];
   }
   assert(TargetTable && "Global data has not been mapped\n");
 
+  // Set tripcount.
+  {
+    std::lock_guard<std::mutex> TblMapMtx(PM->TrlTblMtx);
+    Device.LoopTripCnt.emplace(__kmpc_global_thread_num(Loc), Tripcount);
+  }
+
   // We need to keep bases and offsets separate. Sometimes (e.g. in OpenCL) we
   // need to manifest base pointers prior to launching a kernel. Even if we have
   // mapped an object only partially, e.g. A[N:M], although the kernel is
   // expected to access elements starting at address &A[N] and beyond, we still
   // need to manifest the base of the array &A[0]. In other cases, e.g. the COI
   // API, we need the begin address itself, i.e. &A[N], as the API operates on
   // begin addresses, not bases. That's why we pass args and offsets as two
   // separate entities so that each plugin can do what it needs. This behavior
   // was introdued via https://reviews.llvm.org/D33028 and commit 1546d319244c.
   std::vector<void *> TgtArgs;
   std::vector<ptrdiff_t> TgtOffsets;
 
   PrivateArgumentManagerTy PrivateArgumentManager(Device, AsyncInfo);
 
   int Ret;
   if (ArgNum) {
     // Process data, such as data mapping, before launching the kernel
-    Ret = processDataBefore(loc, DeviceId, HostPtr, ArgNum, ArgBases, Args,
+    Ret = processDataBefore(Loc, DeviceId, HostPtr, ArgNum, ArgBases, Args,
                             ArgSizes, ArgTypes, ArgNames, ArgMappers, TgtArgs,
                             TgtOffsets, PrivateArgumentManager, AsyncInfo);
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Failed to process data before launching the kernel.\n");
       return OFFLOAD_FAIL;
     }
   }
 
   // Launch device execution.
   void *TgtEntryPtr = TargetTable->EntriesBegin[TM->Index].addr;
   DP("Launching target execution %s with pointer " DPxMOD " (index=%d).\n",
      TargetTable->EntriesBegin[TM->Index].name, DPxPTR(TgtEntryPtr), TM->Index);
 
   {
     TIMESCOPE_WITH_NAME_AND_IDENT(
-        IsTeamConstruct ? "runTargetTeamRegion" : "runTargetRegion", loc);
+        IsTeamConstruct ? "runTargetTeamRegion" : "runTargetRegion", Loc);
     if (IsTeamConstruct)
       Ret = Device.runTeamRegion(TgtEntryPtr, &TgtArgs[0], &TgtOffsets[0],
                                  TgtArgs.size(), TeamNum, ThreadLimit,
                                  getLoopTripCount(DeviceId), AsyncInfo);
     else
       Ret = Device.runRegion(TgtEntryPtr, &TgtArgs[0], &TgtOffsets[0],
                              TgtArgs.size(), AsyncInfo);
   }
 
   if (Ret != OFFLOAD_SUCCESS) {
     REPORT("Executing target region abort target.\n");
     return OFFLOAD_FAIL;
   }
 
   if (ArgNum) {
     // Transfer data back and deallocate target memory for (first-)private
     // variables
-    Ret = processDataAfter(loc, DeviceId, HostPtr, ArgNum, ArgBases, Args,
+    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;
     }
   }
 
   return OFFLOAD_SUCCESS;
 }

openmp/libomptarget/src/private.h

extern int targetDataEnd(ident_t *loc, DeviceTy &Device, int32_t ArgNum,
                          bool FromMapper = false);
 
 extern int targetDataUpdate(ident_t *loc, DeviceTy &Device, int32_t arg_num,
                             void **args_base, void **args, int64_t *arg_sizes,
                             int64_t *arg_types, map_var_info_t *arg_names,
                             void **arg_mappers, AsyncInfoTy &AsyncInfo,
                             bool FromMapper = false);
 
-extern int target(ident_t *loc, DeviceTy &Device, void *HostPtr, int32_t ArgNum,
+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 *arg_names,
                   void **ArgMappers, int32_t TeamNum, int32_t ThreadLimit,
-                  int IsTeamConstruct, AsyncInfoTy &AsyncInfo);
+                  uint64_t Tripcount, int IsTeamConstruct,
+                  AsyncInfoTy &AsyncInfo);
 
 extern void handleTargetOutcome(bool Success, ident_t *Loc);
 extern bool checkDeviceAndCtors(int64_t &DeviceID, ident_t *Loc);
 extern void *targetAllocExplicit(size_t size, int device_num, int kind,
                                  const char *name);
 
 // This structure stores information of a mapped memory region.
 struct MapComponentInfoTy {

openmp/libomptarget/src/rtl.cpp

for (auto *R : UsedRTLs) {
       // if its PendingCtors list has been emptied.
       for (int32_t i = 0; i < FoundRTL->NumberOfDevices; ++i) {
         DeviceTy &Device = *PM->Devices[FoundRTL->Idx + i];
         Device.PendingGlobalsMtx.lock();
         if (Device.PendingCtorsDtors[desc].PendingCtors.empty()) {
           AsyncInfoTy AsyncInfo(Device);
           for (auto &dtor : Device.PendingCtorsDtors[desc].PendingDtors) {
             int rc = target(nullptr, Device, dtor, 0, nullptr, nullptr, nullptr,
-                            nullptr, nullptr, nullptr, 1, 1, true /*team*/,
+                            nullptr, nullptr, nullptr, 1, 1, 0, true /*team*/,
                             AsyncInfo);
             if (rc != OFFLOAD_SUCCESS) {
               DP("Running destructor " DPxMOD " failed.\n", DPxPTR(dtor));
             }
           }
           // Remove this library's entry from PendingCtorsDtors
           Device.PendingCtorsDtors.erase(desc);
           // All constructors have been issued, wait for them now.