[OpenMP][libomptarget] Add runtime function for pushing coalesced global records

[gtbercea]

Summary

In the case of coalesced global records, we need to push the exact data size passed in. This patch fixes this by outlining the common functionality of the previous push function and by adding a separate entry point for coalesced pushes. The pop function remains unchanged.

[grokos]

Looks good.

[Hahnfeld]

I guess this will break the case when the DataSize passed to __kmpc_data_sharing_push_stack() needs additional alignment: With this change it is handled in data_sharing_push_stack_common() but __kmpc_data_sharing_push_stack() will determine PushSize without the adjustment and do the final pointer arithmetic.

[ABataev]

In D53141#1261996, @Hahnfeld wrote:

I guess this will break the case when the DataSize passed to __kmpc_data_sharing_push_stack() needs additional alignment: With this change it is handled in data_sharing_push_stack_common() but __kmpc_data_sharing_push_stack() will determine PushSize without the adjustment and do the final pointer arithmetic.

Why DataSize might require an additional alignment? The DataSize must be already aligned.

[gtbercea]

Ensure PushSize is multiple of 8 bytes.

[gtbercea]

Refactor.

[Hahnfeld]

In D53141#1262019, @ABataev wrote:

In D53141#1261996, @Hahnfeld wrote:

I guess this will break the case when the DataSize passed to __kmpc_data_sharing_push_stack() needs additional alignment: With this change it is handled in data_sharing_push_stack_common() but __kmpc_data_sharing_push_stack() will determine PushSize without the adjustment and do the final pointer arithmetic.

Why DataSize might require an additional alignment? The DataSize must be already aligned.

It's required in __kmpc_data_sharing_push_stack() for some cases, please see D52655.
Another problem with this patch is that __kmpc_data_sharing_push_stack() is adding the lane offset to the return value of omptarget_nvptx_SimpleThreadPrivateContext::Allocate(). AFAICS this will break lastprivate which requires the same buffer for all threads.

Anyway, you must do what you think to be correct.

[ABataev]

In D53141#1262378, @Hahnfeld wrote:

In D53141#1262019, @ABataev wrote:

In D53141#1261996, @Hahnfeld wrote:

I guess this will break the case when the DataSize passed to __kmpc_data_sharing_push_stack() needs additional alignment: With this change it is handled in data_sharing_push_stack_common() but __kmpc_data_sharing_push_stack() will determine PushSize without the adjustment and do the final pointer arithmetic.

Why DataSize might require an additional alignment? The DataSize must be already aligned.

It's required in __kmpc_data_sharing_push_stack() for some cases, please see D52655.

It is another problem. The problem is that the StackP pointer is unaligned after all the pointer maths. The DataSize itself should not be aligned.

Another problem with this patch is that __kmpc_data_sharing_push_stack() is adding the lane offset to the return value of omptarget_nvptx_SimpleThreadPrivateContext::Allocate(). AFAICS this will break lastprivate which requires the same buffer for all threads.

Anyway, you must do what you think to be correct.

Yes, there must be an additional check size_t PushSize = (isRuntimeUninitialized() || IsMasterThread()) ? DataSize : WARPSIZE * DataSize;

[gtbercea]

Simply call to common push function.

[gtbercea]

Refactor.

[gtbercea]

Refactor.

[gtbercea]

Refactor.

[ABataev]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

389

This must be removed

437

Also, must be removed

443

It is a very bad idea to have something like this without atomic instructions.
Also, for writing, you need to use atomic instructions (+, possibly, volatile data type). Otherwise, it leads to undefined behavior and problems during optimization.

[gtbercea]

Address comments.

[ABataev]

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

442

Do you realy need the loop here? Seems to me it is better to have it in else branch of the if (IsWarpMaster) statement

[gtbercea]

Move while on else branch.

[ABataev]

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

159

Do you really need to remove this?

libomptarget/deviceRTLs/nvptx/src/supporti.h

191

Same, do you really need to remove this?

[gtbercea]

Reinstate assert.

[gtbercea]

libomptarget/deviceRTLs/nvptx/src/supporti.h

191

Yes, leaving this in will lead to the whole malloc to be optimized out after inlining.

[ABataev]

LG

libomptarget/deviceRTLs/nvptx/src/data_sharing.cu

EXTERN void *__kmpc_data_sharing_environment_begin(
   DSPRINT(DSFLAG, "Data Size %016llx\n", SharingDataSize);
   DSPRINT(DSFLAG, "Default Data Size %016llx\n", SharingDefaultDataSize);
 
   unsigned WID = getWarpId();
   unsigned CurActiveThreads = getActiveThreadsMask();
 
   __kmpc_data_sharing_slot *&SlotP = DataSharingState.SlotPtr[WID];
   void *&StackP = DataSharingState.StackPtr[WID];
-  void *&FrameP = DataSharingState.FramePtr[WID];
+  void * volatile &FrameP = DataSharingState.FramePtr[WID];
   int32_t &ActiveT = DataSharingState.ActiveThreads[WID];
 
   DSPRINT0(DSFLAG, "Save current slot/stack values.\n");
   // Save the current values.
   *SavedSharedSlot = SlotP;
   *SavedSharedStack = StackP;
   *SavedSharedFrame = FrameP;
   *SavedActiveThreads = ActiveT;
if (IsWarpMasterActiveThread()) {
     DSPRINT(DSFLAG, "Active threads: %08x; New mask: %08x\n", CurActive,
             ActiveT);
 
     if (!ActiveT) {
       // No other active threads? Great, lets restore the stack.
 
       __kmpc_data_sharing_slot *&SlotP = DataSharingState.SlotPtr[WID];
       void *&StackP = DataSharingState.StackPtr[WID];
-      void *&FrameP = DataSharingState.FramePtr[WID];
+      void * volatile &FrameP = DataSharingState.FramePtr[WID];
 
       SlotP = *SavedSharedSlot;
       StackP = *SavedSharedStack;
       FrameP = *SavedSharedFrame;
       ActiveT = *SavedActiveThreads;
 
       DSPRINT(DSFLAG, "Restored slot ptr at: %016llx \n", (long long)SlotP);
       DSPRINT(DSFLAG, "Restored stack ptr at: %016llx \n", (long long)StackP);
if (!IsOMPRuntimeInitialized)
     return (void *)&DataSharingState;
 
   // Get the frame used by the requested thread.
 
   unsigned SourceWID = SourceThreadID / WARPSIZE;
 
   DSPRINT(DSFLAG, "Source  warp: %d\n", SourceWID);
 
-  void *P = DataSharingState.FramePtr[SourceWID];
+  void * volatile P = DataSharingState.FramePtr[SourceWID];
   DSPRINT0(DSFLAG, "Exiting __kmpc_get_data_sharing_environment_frame\n");
   return P;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Runtime functions for trunk data sharing scheme.
 ////////////////////////////////////////////////////////////////////////////////
 
EXTERN void __kmpc_data_sharing_init_stack_spmd() {
   // 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.
-EXTERN void* __kmpc_data_sharing_push_stack(size_t DataSize,
-    int16_t UseSharedMemory) {
+INLINE void* data_sharing_push_stack_common(size_t PushSize) {
   if (isRuntimeUninitialized()) {
     ASSERT0(LT_FUSSY, isSPMDMode(),
             "Expected SPMD mode with uninitialized runtime.");
-    return omptarget_nvptx_SimpleThreadPrivateContext::Allocate(DataSize);
+    return omptarget_nvptx_SimpleThreadPrivateContext::Allocate(PushSize);
   }
 
+  // Only warp active master threads manage the stack.
+  bool IsWarpMaster = (getThreadId() % WARPSIZE) == 0;
+
   // Add worst-case padding to DataSize so that future stack allocations are
   // correctly aligned.
   const size_t Alignment = 8;
-  if (DataSize % Alignment != 0) {
-    DataSize += (Alignment - DataSize % Alignment);
-  }
+  PushSize = (PushSize + (Alignment - 1)) / Alignment * Alignment;
 
   // Frame pointer must be visible to all workers in the same warp.
   unsigned WID = getWarpId();
-  void *&FrameP = DataSharingState.FramePtr[WID];
+  void *volatile &FrameP = DataSharingState.FramePtr[WID];
 
-  // Only warp active master threads manage the stack.
-  if (getThreadId() % WARPSIZE == 0) {
+  do {
+    if (IsWarpMaster) {
       // SlotP will point to either the shared memory slot or an existing
       // global memory slot.
       __kmpc_data_sharing_slot *&SlotP = DataSharingState.SlotPtr[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) {
         __kmpc_data_sharing_slot *NewSlot = 0;
         size_t NewSize = PushSize;
 
         // Allocate at least the default size for each type of slot.
         // 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;
         if (DefaultSlotSize > NewSize)
           NewSize = DefaultSlotSize;
         NewSlot = (__kmpc_data_sharing_slot *) SafeMalloc(
             sizeof(__kmpc_data_sharing_slot) + NewSize,
             "Global memory slot allocation.");
 
         NewSlot->Next = 0;
         NewSlot->Prev = SlotP;
         NewSlot->PrevSlotStackPtr = StackP;
         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[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;
       }
     }
+  } while (!FrameP);
 
-  __threadfence_block();
+  return FrameP;
+}
+
+EXTERN void* __kmpc_data_sharing_coalesced_push_stack(size_t DataSize,
+    int16_t UseSharedMemory) {
+  return data_sharing_push_stack_common(DataSize);
+}
+
+// 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.
+EXTERN void* __kmpc_data_sharing_push_stack(size_t DataSize,
+    int16_t UseSharedMemory) {
+  // 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 = (isRuntimeUninitialized() || IsMasterThread()) ?
+      DataSize : WARPSIZE * DataSize;
 
   // Compute the start address of the frame of each thread in the warp.
-  uintptr_t FrameStartAddress = (uintptr_t)FrameP;
+  uintptr_t FrameStartAddress =
+      (uintptr_t) data_sharing_push_stack_common(PushSize);
   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 (isRuntimeUninitialized()) {
     ASSERT0(LT_FUSSY, isSPMDMode(),
             "Expected SPMD mode with uninitialized runtime.");
     return omptarget_nvptx_SimpleThreadPrivateContext::Deallocate(FrameStart);
   }
 
+  __threadfence_block();
+
   if (getThreadId() % WARPSIZE == 0) {
     unsigned WID = getWarpId();
 
     // Current slot
     __kmpc_data_sharing_slot *&SlotP = DataSharingState.SlotPtr[WID];
 
     // Pointer to next available stack.
     void *&StackP = DataSharingState.StackPtr[WID];
if (SlotEmpty && SlotP->Prev) {
       StackP = SlotP->PrevSlotStackPtr;
 
       // Remove the slot.
       SlotP = SlotP->Prev;
       SafeFree(SlotP->Next, "Free slot.");
       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
 // 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) {

libomptarget/deviceRTLs/nvptx/src/interface.h

 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_coalesced_push_stack(size_t size,
+    int16_t UseSharedMemory);
 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];
-  void *FramePtr[DS_Max_Warp_Number];
+  void * volatile 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;

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

INLINE unsigned nsmid() {
   unsigned n;
   asm("mov.u32 %0, %%nsmid;" : "=r"(n));
   return n;
 }
 
 INLINE unsigned smid() {
   unsigned id;
   asm("mov.u32 %0, %%smid;" : "=r"(id));
-  ASSERT0(LT_FUSSY, nsmid() <= MAX_SM,
-          "Expected number of SMs is less than reported.");
   return id;
 }
 
 EXTERN void __kmpc_kernel_init_params(void *Ptr) {
   PRINT(LD_IO, "call to __kmpc_kernel_init_params with version %f\n",
         OMPTARGET_NVPTX_VERSION);
 
   SetTeamsReductionScratchpadPtr(Ptr);
EXTERN void __kmpc_spmd_kernel_init(int ThreadLimit, int16_t RequiresOMPRuntime,
   omptarget_nvptx_TeamDescr &currTeamDescr = getMyTeamDescriptor();
   omptarget_nvptx_WorkDescr &workDescr = getMyWorkDescriptor();
 
   //
   // Initialize task descr for each thread.
   //
   omptarget_nvptx_TaskDescr *newTaskDescr =
       omptarget_nvptx_threadPrivateContext->Level1TaskDescr(threadId);
-  ASSERT0(LT_FUSSY, newTaskDescr, "expected a task descr");
   newTaskDescr->InitLevelOneTaskDescr(ThreadLimit,
                                       currTeamDescr.LevelZeroTaskDescr());
   newTaskDescr->ThreadLimit() = ThreadLimit;
   // install new top descriptor
   omptarget_nvptx_threadPrivateContext->SetTopLevelTaskDescr(threadId,
                                                              newTaskDescr);
 
   // init thread private from init value

libomptarget/deviceRTLs/nvptx/src/supporti.h

ASSERT(LT_FUSSY, (alignment & (alignment - 1)) == 0,
          "alignment %ld is not a power of 2\n", alignment);
   return (~(unsigned long)size + 1) & (alignment - 1);
 }
 
 INLINE void *SafeMalloc(size_t size, const char *msg) // check if success
 {
   void *ptr = malloc(size);
   PRINT(LD_MEM, "malloc data of size %zu for %s: 0x%llx\n", size, msg, P64(ptr));
-  ASSERT(LT_SAFETY, ptr, "failed to allocate %zu bytes for %s\n", size, msg);
   return ptr;
 }
 
 INLINE void *SafeFree(void *ptr, const char *msg) {
   PRINT(LD_MEM, "free data ptr 0x%llx for %s\n", P64(ptr), msg);
   free(ptr);
   return NULL;
 }