[OpenMP][libomptarget] Fix data sharing and globalization infrastructure to work in SPMD mode

[gtbercea]

Summary

This patch fixes the data sharing infrastructure to work for the SPMD and non-SPMD cases.

[gtbercea]

Reset StackP correctly.

[ABataev]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

47

Better to check !isSPMDMode() at first

337

Seems to me the code is not clang-formatted

467

You can just compare pointers as is, no need to cast

[grokos]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

416

Better use SafeMalloc.

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

295–296

Leftover comment from the previous implementation? In data_sharing.cu we said that for uniformity the master thread uses slots of the same size as the workers.

314

Is master_rootS still needed? From what I understand, all threads (master and workers) now use the same slot type, that's why worker_rootS was increased in size from WARPSIZE-1 to WARPSIZE.

[gtbercea]

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

314

It is for the ibm-devel data sharing. We can take this out once we deprecate/remove the "old" data sharing scheme.

[gtbercea]

Address comments and fix formatting.

[gtbercea]

Clean-up.

[ABataev]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

478

USe SafeFree

[grokos]

I don't have any further comments.

@ABataev: When you are also happy with formatting etc. please accept the patch.

[gtbercea]

SafeFree.

[ABataev]

LG

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

__device__ static bool IsWarpMasterActiveThread() {
   unsigned long long Mask = getActiveThreadsMask();
   unsigned long long ShNum = WARPSIZE - (getThreadId() % WARPSIZE);
   unsigned long long Sh = Mask << ShNum;
   // Truncate Sh to the 32 lower bits
   return (unsigned)Sh == 0;
 }
 // Return true if this is the master thread.
 __device__ static bool IsMasterThread() {
-  return getMasterThreadId() == getThreadId();
+  return !isSPMDMode() && getMasterThreadId() == getThreadId();
 }
 
 /// Return the provided size aligned to the size of a pointer.
 __device__ static size_t AlignVal(size_t Val) {
   const size_t Align = (size_t)sizeof(void *);
   if (Val & (Align - 1)) {
     Val += Align;
     Val &= ~(Align - 1);
__kmpc_get_data_sharing_environment_frame(int32_t SourceThreadID,
   DSPRINT0(DSFLAG, "Exiting __kmpc_get_data_sharing_environment_frame\n");
   return P;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Runtime functions for trunk data sharing scheme.
 ////////////////////////////////////////////////////////////////////////////////
 
-// Initialize data sharing data structure. This function needs to be called
-// once at the beginning of a data sharing context (coincides with the kernel
-// 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.
-
-  // Initialize the data sharing structures. This section should only be
-  // executed by the warp active master threads.
-  if (IsWarpMasterActiveThread()) {
-    unsigned WID = getWarpId();
-    omptarget_nvptx_TeamDescr *teamDescr =
-        &omptarget_nvptx_threadPrivateContext->TeamContext();
-    __kmpc_data_sharing_slot *RootS = teamDescr->RootS(WID, IsMasterThread());
+INLINE void data_sharing_init_stack_common() {
 
-    // If a valid address has been returned then proceed with the initalization.
-    // Otherwise the initialization of the slot has already happened in a
-    // previous call to this function.
-    if (RootS) {
-      DataSharingState.SlotPtr[WID] = RootS;
-      DataSharingState.TailPtr[WID] = RootS;
-      DataSharingState.StackPtr[WID] = (void *)&RootS->Data[0];
-    }
-  }
+  for (int WID = 0; WID < WARPSIZE; WID++) {
+    __kmpc_data_sharing_slot *RootS = teamDescr->GetPreallocatedSlotAddr(WID);
 
-  // Currently we only support the sharing of variables between master and
-  // workers. The list of references to shared variables exists only for
-  // the master thread.
-  if (IsMasterThread()) {
-    // Initialize the list of references to arguments.
-    omptarget_nvptx_globalArgs.Init();
-  }
+// Initialize data sharing data structure. This function needs to be called
+// once at the beginning of a data sharing context (coincides with the kernel
+// initialization). This function is called only by the MASTER thread of each
+// team in non-SPMD mode.
+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.
+  data_sharing_init_stack_common();
+
+// Initialize data sharing data structure. This function needs to be called
+// once at the beginning of a data sharing context (coincides with the kernel
+// initialization). This function is called in SPMD mode only.
+EXTERN void __kmpc_data_sharing_init_stack_spmd() {
+  // 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.
+  if (threadIdx.x == 0)
+    data_sharing_init_stack_common();
+
+  __threadfence_block();
 }
 
 // Called at the time of the kernel initialization. This is used to initilize
 // the list of references to shared variables and to pre-allocate global storage
 // for holding the globalized variables.
 //
 // By default the globalized variables are stored in global memory. If the
 // UseSharedMemory is set to true, the runtime will attempt to use shared memory
 // as long as the size requested fits the pre-allocated size.
-//
-// Called by: master, TODO: call by workers
 EXTERN void* __kmpc_data_sharing_push_stack(size_t DataSize,
     int16_t UseSharedMemory) {
   // Frame pointer must be visible to all workers in the same warp.
   unsigned WID = getWarpId();
   void *&FrameP = DataSharingState.FramePtr[WID];
 
   // Only warp active master threads manage the stack.
   if (IsWarpMasterActiveThread()) {
     // SlotP will point to either the shared memory slot or an existing
     // global memory slot.
     __kmpc_data_sharing_slot *&SlotP = DataSharingState.SlotPtr[WID];
-    __kmpc_data_sharing_slot *&TailSlotP = DataSharingState.TailPtr[WID];
     void *&StackP = DataSharingState.StackPtr[WID];
 
     // Compute the total memory footprint of the requested data.
     // The master thread requires a stack only for itself. A worker
     // thread (which at this point is a warp master) will require
     // space for the variables of each thread in the warp,
     // i.e. one DataSize chunk per warp lane.
     // TODO: change WARPSIZE to the number of active threads in the warp.
     size_t PushSize = IsMasterThread() ? DataSize : WARPSIZE * DataSize;
 
     // Check if we have room for the data in the current slot.
     const uintptr_t StartAddress = (uintptr_t)StackP;
     const uintptr_t EndAddress = (uintptr_t)SlotP->DataEnd;
     const uintptr_t RequestedEndAddress = StartAddress + (uintptr_t)PushSize;
 
     // If we requested more data than there is room for in the rest
     // of the slot then we need to either re-use the next slot, if one exists,
     // or create a new slot.
     if (EndAddress < RequestedEndAddress) {
-      size_t NewSize = PushSize;
-
-      // The new or reused slot for holding the data being pushed.
       __kmpc_data_sharing_slot *NewSlot = 0;
+      size_t NewSize = PushSize;
 
-      // Check if there is a next slot.
-      if (__kmpc_data_sharing_slot *ExistingSlot = SlotP->Next) {
-        // Attempt to reuse an existing slot provided the data fits in the slot.
-        // The leftover data space will not be used.
-        ptrdiff_t ExistingSlotSize = (uintptr_t)ExistingSlot->DataEnd -
-                                     (uintptr_t)(&ExistingSlot->Data[0]);
-
-        // Try to add the data in the next available slot. Search for a slot
-        // with enough space.
-        while (ExistingSlotSize < NewSize) {
-          SlotP->Next = ExistingSlot->Next;
-          SlotP->Next->Prev = ExistingSlot->Prev;
-          free(ExistingSlot);
-          ExistingSlot = SlotP->Next;
-          if (!ExistingSlot)
-            break;
-          ExistingSlotSize = (uintptr_t)ExistingSlot->DataEnd -
-                             (uintptr_t)(&ExistingSlot->Data[0]);
-        }
-
-        // Check if a slot has been found.
-        if (ExistingSlotSize >= NewSize) {
-          NewSlot = ExistingSlot;
-          NewSlot->PrevSlotStackPtr = StackP;
-        }
-      }
-
-      if (!NewSlot) {
-        // Allocate at least the default size for each type of slot.
-        size_t DefaultSlotSize =
-            IsMasterThread() ? DS_Slot_Size : DS_Worker_Warp_Slot_Size;
-        if (DefaultSlotSize > NewSize)
-          NewSize = DefaultSlotSize;
-        NewSlot = (__kmpc_data_sharing_slot *)malloc(
-            sizeof(__kmpc_data_sharing_slot) + NewSize);
-        NewSlot->Next = 0;
-        NewSlot->Prev = SlotP;
-        NewSlot->PrevSlotStackPtr = StackP;
-        NewSlot->DataEnd = &NewSlot->Data[NewSize];
-
-        // Newly allocated slots are also tail slots.
-        TailSlotP = NewSlot;
+      // Master is a special case and even though there is only one thread,
+      // it can share more things with the workers. For uniformity, it uses
+      // the full size of a worker warp slot.
+      size_t DefaultSlotSize = DS_Worker_Warp_Slot_Size;
+      NewSlot = (__kmpc_data_sharing_slot *) SafeMalloc(
+          sizeof(__kmpc_data_sharing_slot) + NewSize,
+          "Global memory slot allocation.");
+
+      NewSlot->DataEnd = &NewSlot->Data[0] + NewSize;
 
-        // Make previous slot point to the newly allocated slot.
-        SlotP->Next = NewSlot;
-      }
-
       // The current slot becomes the new slot.
       SlotP = NewSlot;
       // The stack pointer always points to the next free stack frame.
-      StackP = &NewSlot->Data[PushSize];
+      StackP = &NewSlot->Data[0] + PushSize;
       // The frame pointer always points to the beginning of the frame.
       FrameP = &NewSlot->Data[0];
     } else {
       // Add the data chunk to the current slot. The frame pointer is set to
       // point to the start of the new frame held in StackP.
       FrameP = StackP;
       // Reset stack pointer to the requested address.
       StackP = (void *)RequestedEndAddress;

 //
 // When the pop operation removes the last global memory slot,
 // reclaim all outstanding global memory slots since it is
 // likely we have reached the end of the kernel.
 EXTERN void __kmpc_data_sharing_pop_stack(void *FrameStart) {
   if (IsWarpMasterActiveThread()) {
     unsigned WID = getWarpId();
 
+    // Current slot
     __kmpc_data_sharing_slot *&SlotP = DataSharingState.SlotPtr[WID];
+
+    // Pointer to next available stack.
     void *&StackP = DataSharingState.StackPtr[WID];
 
-    // Pop current frame from slot.
+    // If the current slot is empty, we need to free the slot after the
+    // pop.
+    bool SlotEmpty = (StackP == &SlotP->Data[0]);
+
+    // Pop the frame.
     StackP = FrameStart;
 
-    // If we try to pop the last frame of the current slot we need to
-    // move to the previous slot if there is one.
-    const uintptr_t StartAddress = (uintptr_t)FrameStart;
-    if (StartAddress == (uintptr_t)&SlotP->Data[0]) {
-      if (SlotP->Prev) {
-        // The new stack pointer is the end of the data field of the
-        // previous slot. This will allow the stack pointer to be
-        // used in the computation of the remaining data space in
-        // the current slot.
-        StackP = SlotP->PrevSlotStackPtr;
-        // Reset SlotP to previous slot.
-        SlotP = SlotP->Prev;
-      }
+    if (SlotEmpty && SlotP->Prev) {
+      // Before removing the slot we need to reset StackP.
 
-      // If this will "pop" the last global memory node then it is likely
-      // that we are at the end of the data sharing region and we can
-      // de-allocate any existing global memory slots.
-      if (!SlotP->Prev) {
-        __kmpc_data_sharing_slot *Tail = DataSharingState.TailPtr[WID];
-        while(Tail->Prev) {
-          Tail = Tail->Prev;
-          free(Tail->Next);
-        }
-        Tail->Next=0;
-      }
+      // Remove the slot.
+      SlotP = SlotP->Prev;
+      free(SlotP->Next);
+      SlotP->Next = 0;
     }
   }
 
   __threadfence_block();
 }
 
 // Begin a data sharing context. Maintain a list of references to shared
 // variables. This list of references to shared variables will be passed

libomptarget/deviceRTLs/nvptx/src/interface.h

 EXTERN void __kmpc_kernel_end_convergent_parallel(void *buffer);
 EXTERN bool __kmpc_kernel_convergent_simd(void *buffer, uint32_t Mask,
                                           bool *IsFinal, int32_t *LaneSource,
                                           int32_t *LaneId, int32_t *NumLanes);
 EXTERN void __kmpc_kernel_end_convergent_simd(void *buffer);
 
 
 EXTERN void __kmpc_data_sharing_init_stack();
+EXTERN void __kmpc_data_sharing_init_stack_spmd();
 EXTERN void *__kmpc_data_sharing_push_stack(size_t size, int16_t UseSharedMemory);
 EXTERN void __kmpc_data_sharing_pop_stack(void *a);
 EXTERN void __kmpc_begin_sharing_variables(void ***GlobalArgs, size_t nArgs);
 EXTERN void __kmpc_end_sharing_variables();
 EXTERN void __kmpc_get_shared_variables(void ***GlobalArgs);
 
 // The slot used for data sharing by the master and worker threads. We use a
 // complete (default size version and an incomplete one so that we allow sizes

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

 };
 
 // Data structure to keep in shared memory that traces the current slot, stack,
 // and frame pointer as well as the active threads that didn't exit the current
 // environment.
 struct DataSharingStateTy {
   __kmpc_data_sharing_slot *SlotPtr[DS_Max_Warp_Number];
   void *StackPtr[DS_Max_Warp_Number];
-  __kmpc_data_sharing_slot *TailPtr[DS_Max_Warp_Number];
   void *FramePtr[DS_Max_Warp_Number];
   int32_t ActiveThreads[DS_Max_Warp_Number];
 };
 // Additional worker slot type which is initialized with the default worker slot
 // size of 4*32 bytes.
 struct __kmpc_data_sharing_worker_slot_static {
   __kmpc_data_sharing_slot *Next;
   __kmpc_data_sharing_slot *Prev;
worker_rootS[wid].DataEnd =
         &worker_rootS[wid].Data[0] + DS_Worker_Warp_Slot_Size;
     // We currently do not have a next slot.
     worker_rootS[wid].Next = 0;
     worker_rootS[wid].Prev = 0;
     worker_rootS[wid].PrevSlotStackPtr = 0;
     return (__kmpc_data_sharing_slot *)&worker_rootS[wid];
   }
 
+  INLINE __kmpc_data_sharing_slot *GetPreallocatedSlotAddr(int wid) {
+    worker_rootS[wid].DataEnd =
+        &worker_rootS[wid].Data[0] + DS_Worker_Warp_Slot_Size;
+    // We currently do not have a next slot.
+    worker_rootS[wid].Next = 0;
+    worker_rootS[wid].Prev = 0;
+    worker_rootS[wid].PrevSlotStackPtr = 0;
+    return (__kmpc_data_sharing_slot *)&worker_rootS[wid];
+  }
+
 private:
   omptarget_nvptx_TaskDescr
       levelZeroTaskDescr; // icv for team master initial thread
   omptarget_nvptx_WorkDescr
       workDescrForActiveParallel; // one, ONLY for the active par
   omp_lock_t criticalLock;
   uint64_t lastprivateIterBuffer;
 
   __align__(16)
-      __kmpc_data_sharing_worker_slot_static worker_rootS[WARPSIZE - 1];
+      __kmpc_data_sharing_worker_slot_static worker_rootS[WARPSIZE];
   __align__(16) __kmpc_data_sharing_master_slot_static master_rootS[1];
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 // thread private data (struct of arrays for better coalescing)
 // tid refers here to the global thread id
 // do not support multiple concurrent kernel a this time
 class omptarget_nvptx_ThreadPrivateContext {
 public: