[OpenMP][libomptarget] Enable globalization for workers

[gtbercea]

Summary

This patch allows worker to have a global memory stack managed by the runtime. This patch is needed for completeness and consistency with the globalization policy: if a worker-side variable escapes the current context it then needs to be globalized.
Until now, only the master thread was allowed to have such a stack. These global values can now potentially be shared amongst workers if the semantics of the OpenMP program require it.

[grokos]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

42–44

Remove this change, that was a bug which has been fixed already. I have pushed it upstream.

[gtbercea]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

42–44

Was there a bug fix for this published on fabricator prior to me posting it in this patch?

[grokos]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

42–44

There was no revision on the Phabricator for this bug. I had the fix ready from clang-ykt but I hadn't pushed it onto the trunk. We don't need to call __popc(), just compare against 0.

[gtbercea]

Rebase patch.

[gtbercea]

Rebase patch.

[gtbercea]

Choose correct default size depending on slot type: worker or master.

[gtbercea]

Share FrameP with all threads in the same warp.

[gtbercea]

Always pop current frame from slot.

[gtbercea]

Fix comment.

[grokos]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

429

Can you just use and if here? Since NewSize does not change value if DefaultSlotSize <= NewSize, it's clearer to say

if (DefaultSlotSize > NewSize)
  NewSize = DefaultSlotSize;

503–509

I think you can speed things up here. Since you plan to remove the next slot (free(Tail->Next)), there is no point in setting new values for the fields of that slot which will be removed right away. You only need to do this for the very final slot (the one which is statically allocated). The two "//Extra" assignments look redundant in this sense. Also, the loop condition can be simplified. Since there is the statically allocated first slot, Tail will always point to some slot, it can never be NULL, and all you need to check for is while(Tail->Prev). Finally, Tail->Next = 0 can also be deferred until the end of the loop because the slot Tail points to may also be removed in the next iteration. So I think an equivalent (in terms of functionality) version would be:

while(Tail->Prev) {
  Tail = Tail->Prev;
  free(Tail->Next);
}
Tail->Next=0;

[gtbercea]

Address comments.

[grokos]

LGTM.

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

 #include <stdio.h>
 
 // Number of threads in the CUDA block.
 __device__ static unsigned getNumThreads() { return blockDim.x; }
 // Thread ID in the CUDA block
 __device__ static unsigned getThreadId() { return threadIdx.x; }
 // Warp ID in the CUDA block
 __device__ static unsigned getWarpId() { return threadIdx.x / WARPSIZE; }
+// Lane ID in the CUDA warp.
+__device__ static unsigned getLaneId() { return threadIdx.x % WARPSIZE; }
 
 // The CUDA thread ID of the master thread.
 __device__ static unsigned getMasterThreadId() {
   unsigned Mask = WARPSIZE - 1;
   return (getNumThreads() - 1) & (~Mask);
 }
 
 // Find the active threads in the warp - return a mask whose n-th bit is set if
 // the n-th thread in the warp is active.
 __device__ static unsigned getActiveThreadsMask() {
   return __BALLOT_SYNC(0xFFFFFFFF, true);
 }
 
 // Return true if this is the first active thread in the warp.
 __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 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 *);

 //
 // 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) {
-  if (IsMasterThread()) {
-    unsigned WID = getWarpId();
+  // Frame pointer must be visible to all workers in the same warp.
+  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];
-    void *FrameP = 0;
+
+    // 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)DataSize;
+    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 = DataSize;
+      size_t NewSize = PushSize;
 
       // The new or reused slot for holding the data being pushed.
       __kmpc_data_sharing_slot *NewSlot = 0;
 
       // 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.
if (EndAddress < RequestedEndAddress) {
         // Check if a slot has been found.
         if (ExistingSlotSize >= NewSize) {
           NewSlot = ExistingSlot;
           NewSlot->PrevSlotStackPtr = StackP;
         }
       }
 
       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;
+        // Allocate at least the default size for each type of slot.
+        size_t DefaultSlotSize =
+            IsMasterThread() ? DS_Slot_Size : DS_Worker_Warp_Slot_Size;
+        NewSize = DefaultSlotSize > NewSize ? DefaultSlotSize : NewSize;
         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;
 
         // 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];
+      StackP = &NewSlot->Data[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;
     }
-
-    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.
+  __threadfence_block();
 
-  // TODO: implement sharing across workers.
-  return 0;
+  // Compute the start address of the frame of each thread in the warp.
+  uintptr_t FrameStartAddress = (uintptr_t)FrameP;
+  FrameStartAddress += (uintptr_t) (getLaneId() * DataSize);
+  return (void *)FrameStartAddress;
 }
 
 // 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 *FrameStart) {
-  if (IsMasterThread()) {
+  if (IsWarpMasterActiveThread()) {
     unsigned WID = getWarpId();
 
     __kmpc_data_sharing_slot *&SlotP = DataSharingState.SlotPtr[WID];
     void *&StackP = DataSharingState.StackPtr[WID];
 
+    // Pop current frame from slot.
+    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 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 && Tail->Prev) {
+          // Extra
+          Tail->PrevSlotStackPtr = 0;
           Tail = Tail->Prev;
+          // Extra
+          Tail->Next->Prev = 0;
           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.
-
-  // TODO: implement sharing across workers.
+  __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
 // to one or more threads.
 // In L0 data sharing this is called by master thread.
 // In L1 data sharing this is called by active warp master thread.
 EXTERN void __kmpc_begin_sharing_variables(void ***GlobalArgs, size_t nArgs) {