[OpenMP][libomptarget] Fix master warp check

[gtbercea]

Summary

The check for the master warp must take into consideration the actual number of warps: the master warp is equal to the last active warp not necessarily WARPSIZE - 1.

[grokos]

Looks good.

[gtbercea]

Insert guards in init function.

[gtbercea]

Adjust comments.

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

__kmpc_initialize_data_sharing_environment(__kmpc_data_sharing_slot *rootS,
   DSPRINT0(DSFLAG_INIT,
            "Entering __kmpc_initialize_data_sharing_environment\n");
 
   unsigned WID = getWarpId();
   DSPRINT(DSFLAG_INIT, "Warp ID: %d\n", WID);
 
   omptarget_nvptx_TeamDescr *teamDescr =
       &omptarget_nvptx_threadPrivateContext->TeamContext();
-  __kmpc_data_sharing_slot *RootS = teamDescr->RootS(WID);
+  __kmpc_data_sharing_slot *RootS = teamDescr->RootS(WID, IsMasterThread());
 
   DataSharingState.SlotPtr[WID] = RootS;
   DataSharingState.StackPtr[WID] = (void *)&RootS->Data[0];
 
   // We don't need to initialize the frame and active threads.
 
   DSPRINT(DSFLAG_INIT, "Initial data size: %08x \n", InitialDataSize);
   DSPRINT(DSFLAG_INIT, "Root slot at: %016llx \n", (long long)RootS);

 // initialization).
 EXTERN void __kmpc_data_sharing_init_stack() {
   // This function initializes the stack pointer with the pointer to the
   // statically allocated shared memory slots. The size of a shared memory
   // slot is pre-determined to be 256 bytes.
   unsigned WID = getWarpId();
   omptarget_nvptx_TeamDescr *teamDescr =
       &omptarget_nvptx_threadPrivateContext->TeamContext();
-  __kmpc_data_sharing_slot *RootS = teamDescr->RootS(WID);
+  __kmpc_data_sharing_slot *RootS = teamDescr->RootS(WID, IsMasterThread());
 
   DataSharingState.SlotPtr[WID] = RootS;
   DataSharingState.TailPtr[WID] = RootS;
   DataSharingState.StackPtr[WID] = (void *)&RootS->Data[0];
 
   // We initialize the list of references to arguments here.
   omptarget_nvptx_globalArgs.Init();
 }

libomptarget/deviceRTLs/nvptx/src/omptarget-nvptx.h

INLINE omptarget_nvptx_WorkDescr &WorkDescr() {
     return workDescrForActiveParallel;
   }
   INLINE omp_lock_t *CriticalLock() { return &criticalLock; }
   INLINE uint64_t *getLastprivateIterBuffer() { return &lastprivateIterBuffer; }
 
   // init
   INLINE void InitTeamDescr();
 
-  INLINE __kmpc_data_sharing_slot *RootS(int wid) {
+  INLINE __kmpc_data_sharing_slot *RootS(int wid, bool IsMasterThread) {
     // If this is invoked by the master thread of the master warp then intialize
     // it with a smaller slot.
-    if (wid == WARPSIZE - 1) {
+    if (IsMasterThread) {
       // Initialize the pointer to the end of the slot given the size of the
       // data section. DataEnd is non-inclusive.
       master_rootS[0].DataEnd = &master_rootS[0].Data[0] + DS_Slot_Size;
       // We currently do not have a next slot.
       master_rootS[0].Next = 0;
       master_rootS[0].Prev = 0;
       master_rootS[0].PrevSlotStackPtr = 0;
       return (__kmpc_data_sharing_slot *)&master_rootS[0];

libomptarget/deviceRTLs/nvptx/src/omptarget-nvptx.cu

EXTERN void __kmpc_spmd_kernel_init(int ThreadLimit, int16_t RequiresOMPRuntime,
   PRINT(LD_PAR,
         "thread will execute parallel region with id %d in a team of "
         "%d threads\n",
         newTaskDescr->ThreadId(), newTaskDescr->ThreadsInTeam());
 
   if (RequiresDataSharing && threadId % WARPSIZE == 0) {
     // Warp master innitializes data sharing environment.
     unsigned WID = threadId / WARPSIZE;
-    __kmpc_data_sharing_slot *RootS = currTeamDescr.RootS(WID);
+    __kmpc_data_sharing_slot *RootS = currTeamDescr.RootS(
+        WID, WID == WARPSIZE - 1);
     DataSharingState.SlotPtr[WID] = RootS;
     DataSharingState.StackPtr[WID] = (void *)&RootS->Data[0];
   }
 }
 
 EXTERN void __kmpc_spmd_kernel_deinit() {
   // We're not going to pop the task descr stack of each thread since
   // there are no more parallel regions in SPMD mode.