[OpenMP][libomptarget] Enable multiple frames per global memory slot

[gtbercea]

Summary

To save on calls to malloc, this patch enables the re-use of pre-allocated global memory slots.

[gtbercea]

Add init of tail pointer.

[grokos]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

400

So here we are only inspecting the very next slot hoping that it will be large enough to accommodate our request. In case the next slot is not big enough, could there be slots after the next which are suitable? If this scenario is possible, then why are we only inspecting the very next slot and delete everything thereafter if it's not big enough?

502–505

I think this loop will delete all slots apart from the very first (the last iteration will be when Tail points to the first slot and we just deallocate Tail->next). Don't we want to delete the first slot as well?

[gtbercea]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

400

It's definitely do-able but chances of that logic being applied are small. We expect the vast majority of requests to fit the default size of the the slot. For data that is larger than the default case it is vastly more likely that there's no next slot and we just create an entirely new slot.
There is also the issue of space: keeping around a list that occasionally has completely empty slots will be wasteful (you could argue that since it's a linked list we can just delete those nodes and redo the links). What I can do is as I delete the "next" node I can check if the data fits, if it does I can just place it there.

502–505

That's intentional. The head of the list is a statically allocated shared memory node so that one we don't need to call free on.

[grokos]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

400

Yes, that's what I had in mind. We are traversing the list anyway, so checking before deleting comes for free - it doesn't raise the complexity.

[gtbercea]

New algorithm for adding data into an existing slot.

[gtbercea]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

400

I made the change (for now).

[gtbercea]

Rebase patch.

[grokos]

Looks good.

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

EXTERN void __kmpc_data_sharing_init_stack() {
   // 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);
 
   DataSharingState.SlotPtr[WID] = RootS;
-  DataSharingState.StackPtr[WID] = (void *)&RootS->Data[0];
+  DataSharingState.TailPtr[WID] = RootS;
+
+  // Initialize the stack pointer to be equal to the end of
+  // the shared memory slot. This way we ensure that the global
+  // version of the stack will be used.
+  // TODO: remove this:
+  DataSharingState.StackPtr[WID] = RootS->DataEnd;
+
+  // TODO: When the use of shared memory is enabled we will have to
+  // initialize this with the start of the Data region like so:
+  // DataSharingState.StackPtr[WID] = (void *)&RootS->Data[0];
 
   // We initialize the list of references to arguments here.
   omptarget_nvptx_globalArgs.Init();
 }
 
 // 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.
 //
-// TODO: allow more than one push per slot to save on calls to malloc.
-// Currently there is only one slot for each push so the data size in the slot
-// is the same size as the size being requested.
-//
 // Called by: master, TODO: call by workers
-EXTERN void* __kmpc_data_sharing_push_stack(size_t size,
+EXTERN void* __kmpc_data_sharing_push_stack(size_t DataSize,
     int16_t UseSharedMemory) {
   // TODO: Add shared memory support. For now, use global memory only for
   // storing the data sharing slots so ignore the pre-allocated
   // shared memory slot.
 
   // Use global memory for storing the stack.
   if (IsMasterThread()) {
     unsigned WID = getWarpId();
 
     // 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];
+    void *FrameP = 0;
 
-    // The slot for holding the data we are pushing.
+    // 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)DataSize;
+
+    // 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 = DataSize;
+
+      // The new or reused slot for holding the data being pushed.
       __kmpc_data_sharing_slot *NewSlot = 0;
-    size_t NewSize = size;
 
       // Check if there is a next slot.
       if (__kmpc_data_sharing_slot *ExistingSlot = SlotP->Next) {
-      // Attempt to re-use an existing slot provided the data fits in the slot.
+        // 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]);
-      if (ExistingSlotSize >= NewSize)
+        if (ExistingSlotSize >= NewSize) {
           NewSlot = ExistingSlot;
-      else
+          NewSlot->PrevSlotStackPtr = StackP;
+        } else {
+          // Free all slots following the existing slot if any.
+          while(TailSlotP && TailSlotP != ExistingSlot) {
+            TailSlotP = TailSlotP->Prev;
+            free(TailSlotP->Next);
+            TailSlotP->Next=0;
+          }
+
+          // Free next existing slot.
           free(ExistingSlot);
+
+          // Reset next and tail slot pointers correctly.
+          SlotP->Next = 0;
+          TailSlotP = SlotP;
+        }
       }
 
       if (!NewSlot) {
+        // Allocate at least the default size.
+        // TODO: generalize this for workers which need a larger data slot
+        // i.e. using DS_Worker_Warp_Slot_Size.
+        if (DS_Slot_Size > DataSize)
+          NewSize = DS_Slot_Size;
         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];
 
-      // This is the last slot, save it.
+        // Newly allocated slots are also tail slots.
         TailSlotP = NewSlot;
-    }
-
-    NewSlot->DataEnd = &NewSlot->Data[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[DataSize];
+      // 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;
+    }
 
-    return (void*)&SlotP->Data[0];
+    return FrameP;
   }
 
   // TODO: add memory fence here when this function can be called by
   // worker threads also. For now, this function is only called by the
   // master thread of each team.
 
   // TODO: implement sharing across workers.
   return 0;
 }
 
 // Pop the stack and free any memory which can be reclaimed.
 //
 // 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 *a) {
+EXTERN void __kmpc_data_sharing_pop_stack(void *FrameStart) {
   if (IsMasterThread()) {
     unsigned WID = getWarpId();
 
-    __kmpc_data_sharing_slot *S = DataSharingState.SlotPtr[WID];
+    __kmpc_data_sharing_slot *&SlotP = DataSharingState.SlotPtr[WID];
+    void *&StackP = DataSharingState.StackPtr[WID];
 
-    if (S->Prev)
-      S = S->Prev;
+    // 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 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 (!S->Prev) {
+      if (!SlotP->Prev) {
         __kmpc_data_sharing_slot *Tail = DataSharingState.TailPtr[WID];
 
         while(Tail && Tail->Prev) {
           Tail = Tail->Prev;
           free(Tail->Next);
           Tail->Next=0;
         }
       }
+    } else {
+      // This is not the last frame popped from this slot.
+      // Reset StackP
+      StackP = FrameStart;
+    }
 
     return;
   }
 
   // TODO: add memory fence here when this function can be called by
   // worker threads also. For now, this function is only called by the
   // master thread of each team.
 

libomptarget/deviceRTLs/nvptx/src/interface.h

 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
 // greater than the default).
 struct __kmpc_data_sharing_slot {
   __kmpc_data_sharing_slot *Next;
   __kmpc_data_sharing_slot *Prev;
+  void *PrevSlotStackPtr;
   void *DataEnd;
   char Data[];
 };
 EXTERN void
 __kmpc_initialize_data_sharing_environment(__kmpc_data_sharing_slot *RootS,
                                            size_t InitialDataSize);
 EXTERN void *__kmpc_data_sharing_environment_begin(
     __kmpc_data_sharing_slot **SavedSharedSlot, void **SavedSharedStack,

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

struct DataSharingStateTy {
   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;
+  void *PrevSlotStackPtr;
   void *DataEnd;
   char Data[DS_Worker_Warp_Slot_Size];
 };
 // Additional master slot type which is initialized with the default master slot
 // size of 4 bytes.
 struct __kmpc_data_sharing_master_slot_static {
   __kmpc_data_sharing_slot *Next;
   __kmpc_data_sharing_slot *Prev;
+  void *PrevSlotStackPtr;
   void *DataEnd;
   char Data[DS_Slot_Size];
 };
 extern __device__ __shared__ DataSharingStateTy DataSharingState;
 
 ////////////////////////////////////////////////////////////////////////////////
 // task ICV and (implicit & explicit) task state
 
INLINE __kmpc_data_sharing_slot *RootS(int wid) {
     // it with a smaller slot.
     if (wid == WARPSIZE - 1) {
       // 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];
     }
     // Initialize the pointer to the end of the slot given the size of the data
     // section. DataEnd is non-inclusive.
     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