[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

385

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

[jhuber6]

openmp/libomptarget/src/interface.cpp

385

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

385

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

1556

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

1514

openmp/libomptarget/include/omptarget.h

 struct __tgt_kernel_arguments {
   int32_t Version;    // Version of this struct for ABI compatibility.
   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.
 };
+static_assert(sizeof(__tgt_kernel_arguments) == 64 && "Invalid struct size");
 
 /// 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)
   int32_t flags; // Flags associated with the entry, e.g. 'link'.

openmp/libomptarget/src/interface.cpp

DP("Entry %2d: Base=" DPxMOD ", Begin=" DPxMOD ", Size=%" PRId64
                         : "unknown");
   }
 #endif
 
   bool IsTeams = NumTeams != -1;
   if (!IsTeams)
     NumTeams = 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));

 /// 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,
            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
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);
     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;

openmp/libomptarget/src/private.h

extern int targetDataUpdate(ident_t *Loc, DeviceTy &Device, int32_t ArgNum,
                             int64_t *ArgTypes, map_var_info_t *ArgNames,
                             void **ArgMappers, AsyncInfoTy &AsyncInfo,
                             bool FromMapper = false);
 
 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,
-                  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 DeviceNum, 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.