[Libomptarget] Improve device runtime implementation for globalized variables.

[jhuber6]

Summary

Currently the runtime implementation of __kmpc_alloc_shared is extremely slow because it allocated memory for each thread individually. This patch adds a small buffer for the threads to share data and will greatly improve performance for builds where all globalization could not be optimized out. If the shared buffer is full, then memory will not only be allocated per-warp rather than per-thread.

Depends on D97680

[tianshilei1992]

openmp/libomptarget/deviceRTLs/common/src/data_sharing.cu

42

IIRC, in the new deviceRTLs, we only have one stack where the first chunk, which is bigger than the rest, is for the main thread in non-SPMD mode. Why do we want to have two here?

[jdoerfert]

openmp/libomptarget/deviceRTLs/common/src/data_sharing.cu

42

Was easier to write it like this from scratch, either way works, no real difference, this might be actually nicer.

[jdoerfert]

LG

openmp/libomptarget/deviceRTLs/common/omptarget.h

 struct __kmpc_data_sharing_slot {
   __kmpc_data_sharing_slot *Next;
   __kmpc_data_sharing_slot *Prev;
   void *PrevSlotStackPtr;
   void *DataEnd;
   char Data[DS_Worker_Warp_Slot_Size];
 };
 
-// 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 *volatile FramePtr[DS_Max_Warp_Number];
-  __kmpc_impl_lanemask_t ActiveThreads[DS_Max_Warp_Number];
-};
-
-extern DataSharingStateTy EXTERN_SHARED(DataSharingState);
-
 ////////////////////////////////////////////////////////////////////////////////
 // task ICV and (implicit & explicit) task state
 
 class omptarget_nvptx_TaskDescr {
 public:
   // methods for flags
   INLINE omp_sched_t GetRuntimeSched() const;
   INLINE void SetRuntimeSched(omp_sched_t sched);

openmp/libomptarget/deviceRTLs/common/src/data_sharing.cu

 INLINE static bool IsMasterThread(bool isSPMDExecutionMode) {
   return !isSPMDExecutionMode && GetMasterThreadID() == GetThreadIdInBlock();
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Runtime functions for trunk data sharing scheme.
 ////////////////////////////////////////////////////////////////////////////////
 
-INLINE static void data_sharing_init_stack_common() {
-  ASSERT0(LT_FUSSY, isRuntimeInitialized(), "Runtime must be initialized.");
-  omptarget_nvptx_TeamDescr *teamDescr =
-      &omptarget_nvptx_threadPrivateContext->TeamContext();
-
-  for (int WID = 0; WID < DS_Max_Warp_Number; WID++) {
-    __kmpc_data_sharing_slot *RootS = teamDescr->GetPreallocatedSlotAddr(WID);
-    DataSharingState.SlotPtr[WID] = RootS;
-    DataSharingState.StackPtr[WID] = (void *)&RootS->Data[0];
-  }
-}
-
-// 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() {
-  ASSERT0(LT_FUSSY, isRuntimeInitialized(), "Runtime must be initialized.");
-  // 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();
-}
-
-EXTERN void *__kmpc_alloc_shared(size_t DataSize) {
-  return (void *)SafeMalloc(DataSize, "Alloc Shared");
-}
+static constexpr unsigned MinBytes = 8;
 
-EXTERN void __kmpc_free_shared(void *FrameStart) {
-  SafeFree(FrameStart, "Free Shared");
+template <unsigned BytesPerThread, unsigned NThreads = MAX_THREADS_PER_TEAM>
+struct alignas(32) ThreadStackTy {
+  static constexpr unsigned MaxSize = NThreads * BytesPerThread;
+  static constexpr unsigned NumThreads = NThreads;
+  static constexpr unsigned NumWarps = (NThreads + WARPSIZE - 1) / WARPSIZE;
+  static constexpr unsigned MaxSizePerWarp = MaxSize / NumWarps;
+
+  unsigned char Data[MaxSize];
+  char Sizes[MaxSize / MinBytes];
+  char SizeUsage[NumWarps];
+  char Usage[NumWarps];
+};
+
+[[clang::loader_uninitialized]] ThreadStackTy<MinBytes * 8, 1> MainSharedStack;
+#pragma omp allocate(MainSharedStack) allocator(omp_pteam_mem_alloc)
+
+[[clang::loader_uninitialized]] ThreadStackTy<MinBytes * 2,
+                                              MAX_THREADS_PER_TEAM / 8>
+    WorkerSharedStack;
+#pragma omp allocate(WorkerSharedStack) allocator(omp_pteam_mem_alloc)
+
+template <typename AllocTy>
+static void *__kmpc_alloc_for_warp(AllocTy Alloc, unsigned Bytes,
+                                   unsigned WarpBytes) {
+  void *Ptr;
+  __kmpc_impl_lanemask_t CurActive = __kmpc_impl_activemask();
+  unsigned LeaderID = __kmpc_impl_ffs(CurActive) - 1;
+  bool IsWarpLeader = (GetThreadIdInBlock() % WARPSIZE) == LeaderID;
+  if (IsWarpLeader)
+    Ptr = Alloc();
+  // Get address from the first active lane.
+  int *FP = (int *)&Ptr;
+  FP[0] = __kmpc_impl_shfl_sync(CurActive, FP[0], LeaderID);
+  if (sizeof(Ptr) == 8)
+    FP[1] = __kmpc_impl_shfl_sync(CurActive, FP[1], LeaderID);
+  return (void *)&((char *)(Ptr))[(GetLaneId() - LeaderID) * Bytes];
+}
+
+EXTERN void *__kmpc_alloc_shared(size_t Bytes) {
+  Bytes = Bytes + (Bytes % MinBytes);
+  if (IsMasterThread(isSPMDMode())) {
+    // Main thread alone, use shared memory if space is available.
+    if (MainSharedStack.Usage[0] + Bytes <= MainSharedStack.MaxSize) {
+      void *Ptr = &MainSharedStack.Data[MainSharedStack.Usage[0]];
+      MainSharedStack.Usage[0] += Bytes;
+      MainSharedStack.Sizes[MainSharedStack.SizeUsage[0]++] = Bytes;
+      return Ptr;
     }
-
-// 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() {
-  ASSERT0(LT_FUSSY, isRuntimeInitialized(), "Runtime must be initialized.");
-  // 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 (GetThreadIdInBlock() == 0)
-    data_sharing_init_stack_common();
-
-  __kmpc_impl_threadfence_block();
+  } else {
+    int TID = GetThreadIdInBlock();
+    int WID = GetWarpId();
+    unsigned WarpBytes = Bytes * WARPSIZE;
+    auto AllocSharedStack = [&]() {
+      unsigned WarpOffset = WID * WorkerSharedStack.MaxSizePerWarp;
+      void *Ptr =
+          &WorkerSharedStack.Data[WarpOffset + WorkerSharedStack.Usage[WID]];
+      WorkerSharedStack.Usage[WID] += WarpBytes;
+      WorkerSharedStack.Sizes[WorkerSharedStack.SizeUsage[WID]++] = WarpBytes;
+      return Ptr;
+    };
+    if (TID < WorkerSharedStack.NumThreads &&
+        WorkerSharedStack.Usage[WID] + WarpBytes <=
+            WorkerSharedStack.MaxSizePerWarp)
+      return __kmpc_alloc_for_warp(AllocSharedStack, Bytes, WarpBytes);
+  }
+  // Fallback to malloc
+  int TID = GetThreadIdInBlock();
+  unsigned WarpBytes = Bytes * WARPSIZE;
+  auto AllocGlobal = [&] {
+    return SafeMalloc(WarpBytes, "AllocGlobalFallback");
+  };
+  return __kmpc_alloc_for_warp(AllocGlobal, Bytes, WarpBytes);
 }
 
-INLINE static void *data_sharing_push_stack_common(size_t PushSize) {
-  ASSERT0(LT_FUSSY, isRuntimeInitialized(), "Expected initialized runtime.");
-
-  // Only warp active master threads manage the stack.
-  bool IsWarpMaster = (GetThreadIdInBlock() % WARPSIZE) == 0;
-
-  // Add worst-case padding to DataSize so that future stack allocations are
-  // correctly aligned.
-  const size_t Alignment = 8;
-  PushSize = (PushSize + (Alignment - 1)) / Alignment * Alignment;
-
-  // Frame pointer must be visible to all workers in the same warp.
-  const unsigned WID = GetWarpId();
-  void *FrameP = 0;
+EXTERN void __kmpc_free_shared(void *Ptr) {
   __kmpc_impl_lanemask_t CurActive = __kmpc_impl_activemask();
-
-  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];
-
-    // 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 = DataSharingState.FramePtr[WID] = &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 = DataSharingState.FramePtr[WID] = StackP;
-      // Reset stack pointer to the requested address.
-      StackP = (void *)RequestedEndAddress;
+  unsigned LeaderID = __kmpc_impl_ffs(CurActive) - 1;
+  bool IsWarpLeader = (GetThreadIdInBlock() % WARPSIZE) == LeaderID;
+  __kmpc_syncwarp(CurActive);
+  if (IsWarpLeader) {
+    if (Ptr >= &MainSharedStack.Data[0] &&
+        Ptr < &MainSharedStack.Data[MainSharedStack.MaxSize]) {
+      unsigned Bytes = MainSharedStack.Sizes[--MainSharedStack.SizeUsage[0]];
+      MainSharedStack.Usage[0] -= Bytes;
+      return;
     }
+    if (Ptr >= &WorkerSharedStack.Data[0] &&
+        Ptr < &WorkerSharedStack.Data[WorkerSharedStack.MaxSize]) {
+      int WID = GetWarpId();
+      unsigned Bytes =
+          WorkerSharedStack.Sizes[--WorkerSharedStack.SizeUsage[WID]];
+      WorkerSharedStack.Usage[WID] -= Bytes;
+      return;
     }
-  // Get address from lane 0.
-  int *FP = (int *)&FrameP;
-  FP[0] = __kmpc_impl_shfl_sync(CurActive, FP[0], 0);
-  if (sizeof(FrameP) == 8)
-    FP[1] = __kmpc_impl_shfl_sync(CurActive, FP[1], 0);
-
-  return FrameP;
+    SafeFree(Ptr, "FreeGlobalFallback");
   }
-
-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(isSPMDMode()))
-                        ? DataSize
-                        : WARPSIZE * DataSize;
-
-  // Compute the start address of the frame of each thread in the warp.
-  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) {
-  ASSERT0(LT_FUSSY, isRuntimeInitialized(), "Expected initialized runtime.");
-
-  __kmpc_impl_threadfence_block();
-
-  if (GetThreadIdInBlock() % 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];
-
-    // Pop the frame.
-    StackP = FrameStart;
-
-    // If the current slot is empty, we need to free the slot after the
-    // pop.
-    bool SlotEmpty = (StackP == &SlotP->Data[0]);
-
-    if (SlotEmpty && SlotP->Prev) {
-      // Before removing the slot we need to reset StackP.
-      StackP = SlotP->PrevSlotStackPtr;
-
-      // Remove the slot.
-      SlotP = SlotP->Prev;
-      SafeFree(SlotP->Next, "Free slot.");
-      SlotP->Next = 0;
-    }
+EXTERN void __kmpc_data_sharing_init_stack() {
+  for (unsigned i = 0; i < MainSharedStack.NumWarps; ++i) {
+    MainSharedStack.SizeUsage[i] = 0;
+    MainSharedStack.Usage[i] = 0;
+  }
+  for (unsigned i = 0; i < WorkerSharedStack.NumWarps; ++i) {
+    WorkerSharedStack.SizeUsage[i] = 0;
+    WorkerSharedStack.Usage[i] = 0;
   }
 }
 
 // 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.

openmp/libomptarget/deviceRTLs/common/src/omp_data.cu

 volatile omptarget_nvptx_WorkFn SHARED(omptarget_nvptx_workFn);
 
 ////////////////////////////////////////////////////////////////////////////////
 // OpenMP kernel execution parameters
 ////////////////////////////////////////////////////////////////////////////////
 uint32_t SHARED(execution_param);
 
 ////////////////////////////////////////////////////////////////////////////////
-// Data sharing state
-////////////////////////////////////////////////////////////////////////////////
-DataSharingStateTy SHARED(DataSharingState);
-
-////////////////////////////////////////////////////////////////////////////////
 // Scratchpad for teams reduction.
 ////////////////////////////////////////////////////////////////////////////////
 void *SHARED(ReductionScratchpadPtr);
 
 ////////////////////////////////////////////////////////////////////////////////
 // Data sharing related variables.
 ////////////////////////////////////////////////////////////////////////////////
 omptarget_nvptx_SharedArgs SHARED(omptarget_nvptx_globalArgs);
 
 #pragma omp end declare target

openmp/libomptarget/deviceRTLs/common/src/omptarget.cu

EXTERN void __kmpc_kernel_init(int ThreadLimit, int16_t RequiresOMPRuntime) {
   omptarget_nvptx_TaskDescr *currTaskDescr =
       omptarget_nvptx_threadPrivateContext->GetTopLevelTaskDescr(threadId);
   nThreads = GetNumberOfThreadsInBlock();
   threadLimit = ThreadLimit;
 
   if (!isSPMDMode())
     omptarget_nvptx_globalArgs.Init();
 
+  __kmpc_data_sharing_init_stack();
   __kmpc_impl_target_init();
 }
 
 EXTERN void __kmpc_kernel_deinit(int16_t IsOMPRuntimeInitialized) {
   PRINT0(LD_IO, "call to __kmpc_kernel_deinit\n");
   ASSERT0(LT_FUSSY, IsOMPRuntimeInitialized,
           "Generic always requires initialized runtime.");
   // Enqueue omp state object for use by another team.
EXTERN void __kmpc_spmd_kernel_init(int ThreadLimit,
   if (threadId == 0) {
     usedSlotIdx = __kmpc_impl_smid() % MAX_SM;
     parallelLevel[0] =
         1 + (GetNumberOfThreadsInBlock() > 1 ? OMP_ACTIVE_PARALLEL_LEVEL : 0);
   } else if (GetLaneId() == 0) {
     parallelLevel[GetWarpId()] =
         1 + (GetNumberOfThreadsInBlock() > 1 ? OMP_ACTIVE_PARALLEL_LEVEL : 0);
   }
+  __kmpc_data_sharing_init_stack();
   if (!RequiresOMPRuntime) {
     // Runtime is not required - exit.
     __kmpc_impl_syncthreads();
     return;
   }
 
   //
   // Team Context Initialization.

openmp/libomptarget/deviceRTLs/interface.h

 EXTERN void __kmpc_spmd_kernel_init(int ThreadLimit,
                                     int16_t RequiresOMPRuntime);
 EXTERN void __kmpc_spmd_kernel_deinit_v2(int16_t RequiresOMPRuntime);
 EXTERN void __kmpc_kernel_prepare_parallel(void *WorkFn);
 EXTERN bool __kmpc_kernel_parallel(void **WorkFn);
 EXTERN void __kmpc_kernel_end_parallel();
 
 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);
 
 /// Entry point to start a new parallel region.
 ///
 /// \param ident       The source identifier.
 /// \param global_tid  The global thread ID.