[OpenMP] Add support for changing stack size in device RTL

[jhuber6]

Summary

This patch adds support for using dynamic shared memory in the stack
using the sized introduced in D110108. This allows dynamic shared memory
to be scaled in the runtime and statically removed by DCE when the size
is zero.

Depends on D110108

[jhuber6]

Making the usage use dynamic memory as well too, now the usage only uses as much
shared memory are there are threads in the block.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

1240

+? Not *?

[jhuber6]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

1240

This is a linear region of memory, we allocate a byte per-thread to save the memory usage per-thread. This allocated enough memory for the stack, and enough memory for all threads that are active.

[jdoerfert]

LG. Two comments.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

932

We should add a TODO here. It's unreasonable that we copy stuff from the device even though the host has the image with the information. I know this is how we do it for other stuff too, in general seems sub-optimal.

1241

Put these things in separate variables with explanation what they mean and how the size is computed. In the current way this is just magic.

openmp/libomptarget/DeviceRTL/include/Interface.h

 uint32_t __kmpc_get_hardware_thread_id_in_block();
 
 /// Kernel
 ///
 ///{
 int8_t __kmpc_is_spmd_exec_mode();
 
 int32_t __kmpc_target_init(IdentTy *Ident, bool IsSPMD,
-                           bool UseGenericStateMachine, bool);
+                           bool UseGenericStateMachine,
+                           bool RequiresFullRuntime, uint32_t StackSize);
 
 void __kmpc_target_deinit(IdentTy *Ident, bool IsSPMD, bool);
 
 ///}
 
 /// Reduction
 ///
 ///{

openmp/libomptarget/DeviceRTL/include/State.h

 #include "Types.h"
 
 #pragma omp declare target
 
 namespace _OMP {
 
 namespace state {
 
-inline constexpr uint32_t SharedScratchpadSize = SHARED_SCRATCHPAD_SIZE;
-
 /// Initialize the state machinery. Must be called by all threads.
-void init(bool IsSPMD);
+void init(bool IsSPMD, uint32_t StackSize);
 
 /// TODO
 enum ValueKind {
   VK_NThreads,
   VK_Level,
   VK_ActiveLevel,
   VK_MaxActiveLevels,
   VK_RunSched,

openmp/libomptarget/DeviceRTL/src/Kernel.cpp

 //===--- Kernel.cpp - OpenMP device kernel interface -------------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the kernel entry points for the device.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Debug.h"
 #include "Interface.h"
 #include "Mapping.h"
 #include "State.h"
 #include "Synchronization.h"
 #include "Types.h"
 
 using namespace _OMP;
 
 #pragma omp declare target
 
-static void inititializeRuntime(bool IsSPMD) {
+static void inititializeRuntime(bool IsSPMD, uint32_t StackSize) {
   // Order is important here.
   synchronize::init(IsSPMD);
   mapping::init(IsSPMD);
-  state::init(IsSPMD);
+  state::init(IsSPMD, StackSize);
 }
 
 /// Simple generic state machine for worker threads.
 static void genericStateMachine(IdentTy *Ident) {
 
   uint32_t TId = mapping::getThreadIdInBlock();
 
   do {

 
 extern "C" {
 
 /// Initialization
 ///
 /// \param Ident               Source location identification, can be NULL.
 ///
 int32_t __kmpc_target_init(IdentTy *Ident, bool IsSPMD,
-                           bool UseGenericStateMachine, bool) {
+                           bool UseGenericStateMachine,
+                           bool RequiresFullRuntime, uint32_t StackSize) {
   if (IsSPMD) {
-    inititializeRuntime(/* IsSPMD */ true);
+    inititializeRuntime(/* IsSPMD */ true, StackSize);
     synchronize::threads();
   } else {
-    inititializeRuntime(/* IsSPMD */ false);
+    inititializeRuntime(/* IsSPMD */ false, StackSize);
     // No need to wait since only the main threads will execute user
     // code and workers will run into a barrier right away.
   }
 
   if (IsSPMD) {
     state::assumeInitialState(IsSPMD);
     return -1;
   }

openmp/libomptarget/DeviceRTL/src/State.cpp

 //===------ State.cpp - OpenMP State & ICV interface ------------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //===----------------------------------------------------------------------===//
 
 #include "State.h"
 #include "Configuration.h"
 #include "Debug.h"
 #include "Interface.h"
 #include "Mapping.h"
 #include "Synchronization.h"
 #include "Types.h"
 #include "Utils.h"
 

 /// In fact, it is a separate stack *per warp*. That means, each warp must push
 /// and pop symmetrically or this breaks, badly. The implementation will (aim
 /// to) detect non-lock-step warps and fallback to malloc/free. The same will
 /// happen if a warp runs out of memory. The master warp in generic memory is
 /// special and is given more memory than the rest.
 ///
 struct SharedMemorySmartStackTy {
   /// Initialize the stack. Must be called by all threads.
-  void init(bool IsSPMD);
+  void init(bool IsSPMD, uint32_t StackSize);
 
   /// Allocate \p Bytes on the stack for the encountering thread. Each thread
   /// can call this function.
   void *push(uint64_t Bytes);
 
   /// Deallocate the last allocation made by the encountering thread and pointed
   /// to by \p Ptr from the stack. Each thread can call this function.
   void pop(void *Ptr, uint32_t Bytes);
 
 private:
   /// Compute the size of the storage space reserved for a thread.
   uint32_t computeThreadStorageTotal() {
     uint32_t NumLanesInBlock = mapping::getNumberOfProcessorElements();
-    return utils::align_down((state::SharedScratchpadSize / NumLanesInBlock),
+    return utils::align_down((SharedScratchpadSize / NumLanesInBlock),
                              Alignment);
   }
 
   /// Return the top address of the warp data stack, that is the first address
   /// this warp will allocate memory at next.
   void *getThreadDataTop(uint32_t TId) {
     return &Data[computeThreadStorageTotal() * TId + Usage[TId]];
   }
 
+  /// The size of the shared scratchpad, initialized by a single thread.
+  uint32_t SharedScratchpadSize;
+
   /// The actual storage, shared among all warps.
-  unsigned char Data[state::SharedScratchpadSize]
-      __attribute__((aligned(Alignment)));
-  unsigned char Usage[mapping::MaxThreadsPerTeam]
-      __attribute__((aligned(Alignment)));
+  unsigned char *Data __attribute__((aligned(Alignment)));
+  unsigned char *Usage __attribute__((aligned(Alignment)));
 };
 
-static_assert(state::SharedScratchpadSize / mapping::MaxThreadsPerTeam <= 256,
-              "Shared scratchpad of this size not supported yet.");
-
 /// The allocation of a single shared memory scratchpad.
 static SharedMemorySmartStackTy SHARED(SharedMemorySmartStack);
 
-void SharedMemorySmartStackTy::init(bool IsSPMD) {
+void SharedMemorySmartStackTy::init(bool IsSPMD, uint32_t StackSize) {
+  if (mapping::getThreadIdInBlock() == 0) {
+    SharedScratchpadSize = StackSize;
+    Data = static_cast<unsigned char *>(memory::getDynamicBuffer()) +
+           utils::align_up(config::getDynamicMemorySize(), Alignment);
+    Usage = Data + utils::align_up(StackSize, Alignment);
+  }
+  synchronize::threads();
   Usage[mapping::getThreadIdInBlock()] = 0;
 }
 
 void *SharedMemorySmartStackTy::push(uint64_t Bytes) {
   // First align the number of requested bytes.
   uint64_t AlignedBytes = utils::align_up(Bytes, Alignment);
 
   uint32_t StorageTotal = computeThreadStorageTotal();
void *SharedMemorySmartStackTy::push(uint64_t Bytes) {
 
   return memory::allocGlobal(AlignedBytes,
                              "Slow path shared memory allocation, insufficient "
                              "shared memory stack memory!");
 }
 
 void SharedMemorySmartStackTy::pop(void *Ptr, uint32_t Bytes) {
   uint64_t AlignedBytes = utils::align_up(Bytes, Alignment);
-  if (Ptr >= &Data[0] && Ptr < &Data[state::SharedScratchpadSize]) {
+  if (Ptr >= &Data[0] && Ptr < &Data[SharedScratchpadSize]) {
     int TId = mapping::getThreadIdInBlock();
     Usage[TId] -= AlignedBytes;
     return;
   }
   memory::freeGlobal(Ptr, "Slow path shared memory deallocation");
 }
 
 } // namespace
void *&state::lookupPtr(ValueKind Kind, bool IsReadonly) {
   case state::VK_ParallelRegionFn:
     return TeamState.ParallelRegionFnVar;
   default:
     break;
   }
   __builtin_unreachable();
 }
 
-void state::init(bool IsSPMD) {
-  SharedMemorySmartStack.init(IsSPMD);
+void state::init(bool IsSPMD, uint32_t StackSize) {
+  if (StackSize > 0)
+    SharedMemorySmartStack.init(IsSPMD, StackSize);
   if (!mapping::getThreadIdInBlock())
     TeamState.init(IsSPMD);
 
   ThreadStates[mapping::getThreadIdInBlock()] = nullptr;
 }
 
 void state::enterDataEnvironment() {
   unsigned TId = mapping::getThreadIdInBlock();

openmp/libomptarget/deviceRTLs/common/include/target.h

 /// Initialization
 ///
 /// Must be called by all threads.
 ///
 /// \param Ident               Source location identification, can be NULL.
 ///
 int32_t __kmpc_target_init(ident_t *Ident, bool IsSPMD,
                            bool UseGenericStateMachine,
-                           bool RequiresFullRuntime);
+                           bool RequiresFullRuntime, uint32_t);
 
 /// De-Initialization
 ///
 /// Must be called by the main thread in generic mode, can be called by all
 /// threads. Must be called by all threads in SPMD mode.
 ///
 /// In non-SPMD, this function releases the workers trapped in a state machine
 /// and also any memory dynamically allocated by the runtime.

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

do {
     __kmpc_barrier_simple_spmd(Ident, TId);
 
   } while (true);
 }
 
 EXTERN
 int32_t __kmpc_target_init(ident_t *Ident, bool IsSPMD,
                            bool UseGenericStateMachine,
-                           bool RequiresFullRuntime) {
+                           bool RequiresFullRuntime, uint32_t) {
   int TId = __kmpc_get_hardware_thread_id_in_block();
   if (IsSPMD)
     __kmpc_spmd_kernel_init(RequiresFullRuntime);
   else
     __kmpc_generic_kernel_init();
 
    if (IsSPMD) {
     __kmpc_barrier_simple_spmd(Ident, TId);

openmp/libomptarget/deviceRTLs/interface.h

 typedef enum omp_proc_bind_t {
   omp_proc_bind_false = 0,
   omp_proc_bind_true = 1,
   omp_proc_bind_master = 2,
   omp_proc_bind_close = 3,
   omp_proc_bind_spread = 4
 } omp_proc_bind_t;
 
 EXTERN double omp_get_wtick(void);
 EXTERN double omp_get_wtime(void);
 
 EXTERN void omp_set_num_threads(int num);
 EXTERN int omp_get_num_threads(void);
 EXTERN int omp_get_max_threads(void);
 EXTERN int omp_get_thread_limit(void);
 EXTERN int omp_get_thread_num(void);
 EXTERN int omp_get_num_procs(void);
 EXTERN int omp_in_parallel(void);
 EXTERN int omp_in_final(void);
 EXTERN void omp_set_dynamic(int flag);
 EXTERN int omp_get_dynamic(void);
 EXTERN void omp_set_nested(int flag);
 EXTERN int omp_get_nested(void);
 EXTERN void omp_set_max_active_levels(int level);
 EXTERN int omp_get_max_active_levels(void);
 EXTERN int omp_get_level(void);
 EXTERN int omp_get_active_level(void);
 EXTERN int omp_get_ancestor_thread_num(int level);
 EXTERN int omp_get_team_size(int level);
 
 EXTERN void omp_init_lock(omp_lock_t *lock);
 EXTERN void omp_init_nest_lock(omp_nest_lock_t *lock);
 EXTERN void omp_destroy_lock(omp_lock_t *lock);
 EXTERN void omp_destroy_nest_lock(omp_nest_lock_t *lock);
 EXTERN void omp_set_lock(omp_lock_t *lock);
 EXTERN void omp_set_nest_lock(omp_nest_lock_t *lock);

 EXTERN int32_t __kmpc_cancellationpoint(kmp_Ident *loc, int32_t global_tid,
                                         int32_t cancelVal);
 EXTERN int32_t __kmpc_cancel(kmp_Ident *loc, int32_t global_tid,
                              int32_t cancelVal);
 
 // non standard
 EXTERN int32_t __kmpc_target_init(ident_t *Ident, bool IsSPMD,
                                   bool UseGenericStateMachine,
-                                  bool RequiresFullRuntime);
+                                  bool RequiresFullRuntime, uint32_t);
 EXTERN void __kmpc_target_deinit(ident_t *Ident, bool IsSPMD,
                                  bool RequiresFullRuntime);
 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_begin_sharing_variables(void ***GlobalArgs, size_t nArgs);

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 //===----RTLs/cuda/src/rtl.cpp - Target RTLs Implementation ------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // RTL for CUDA machine
 //
 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <string>
 #include <unordered_map>
 #include <vector>
 
 #include "Debug.h"
struct KernelTy {
   // 0 - SPMD mode (without master warp)
   // 1 - Generic mode (with master warp)
   // 2 - SPMD mode execution with Generic mode semantics.
   int8_t ExecutionMode;
 
   /// Maximal number of threads per block for this kernel.
   int MaxThreadsPerBlock = 0;
 
-  KernelTy(CUfunction _Func, int8_t _ExecutionMode)
-      : Func(_Func), ExecutionMode(_ExecutionMode) {}
+  /// Amount of dynamic shared memory to allocate for the data sharing stack.
+  uint32_t StackSize = 0;
+
+  KernelTy(CUfunction _Func, int8_t _ExecutionMode, uint32_t _StackSize)
+      : Func(_Func), ExecutionMode(_ExecutionMode), StackSize(_StackSize) {}
 };
 
 /// Device environment data
 /// Manually sync with the deviceRTL side for now, move to a dedicated header
 /// file later.
 struct omptarget_device_environmentTy {
   int32_t debug_level;
   uint32_t num_devices;
for (const __tgt_offload_entry *E = HostBegin; E != HostEnd; ++E) {
           return nullptr;
         }
       } else {
         DP("Loading global exec_mode '%s' - symbol missing, using default "
            "value GENERIC (1)\n",
            ExecModeName);
       }
 
-      KernelsList.emplace_back(Func, ExecModeVal);
+      // Find the amount of dynamic shared memory to allocate for the data
+      // sharing stack in the new device RTL.
+      uint32_t StackSizeVal = 0;
+      std::string StackSizeNameStr(E->name);
+      StackSizeNameStr += "_stack_size";
+      const char *StackSizeName = StackSizeNameStr.c_str();
+
+      CUdeviceptr StackSizePtr;
+      Err = cuModuleGetGlobal(&StackSizePtr, &CUSize, Module, StackSizeName);
+      if (Err == CUDA_SUCCESS) {
+        if (CUSize != sizeof(uint32_t)) {
+          DP("Loading global stack_size '%s' - size mismatch (%zd != %zd)\n",
+             StackSizeName, CUSize, sizeof(uint32_t));
+          return nullptr;
+        }
+
+        Err = cuMemcpyDtoH(&StackSizeVal, StackSizePtr, CUSize);
+        if (Err != CUDA_SUCCESS) {
+          REPORT("Error when copying data from device to host. Pointers: "
+                 "host = " DPxMOD ", device = " DPxMOD ", size = %zd\n",
+                 DPxPTR(&StackSizeVal), DPxPTR(StackSizePtr), CUSize);
+          CUDA_ERR_STRING(Err);
+          return nullptr;
+        }
+      }
+
+      KernelsList.emplace_back(Func, ExecModeVal, StackSizeVal);
 
       __tgt_offload_entry Entry = *E;
       Entry.addr = &KernelsList.back();
       addOffloadEntry(DeviceId, Entry);
     }
 
     // send device environment data to the device
     {
INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
              : "(null)",
          CudaBlocksPerGrid, CudaThreadsPerBlock,
          (KernelInfo->ExecutionMode != SPMD
               ? (KernelInfo->ExecutionMode == GENERIC ? "Generic"
                                                       : "SPMD-Generic")
               : "SPMD"));
 
     CUstream Stream = getStream(DeviceId, AsyncInfo);
+
+    // Round up memory sizes to allow for proper pointer alignment.
+    auto RoundUp = [](uint32_t V, uint32_t Align) {
+      return ((V + uint32_t(Align) - 1) / uint32_t(Align)) * uint32_t(Align);
+    };
+
+    // Amount of shared memory to allocate, must contain enough for the data
+    // sharing stack.
+    uint32_t DynamicSharedMemory = RoundUp(DynamicMemorySize, 8);
+    if (KernelInfo->StackSize > 0)
+      DynamicSharedMemory +=
+          RoundUp(KernelInfo->StackSize, 8) +
+          RoundUp(CudaThreadsPerBlock, DeviceData[DeviceId].WarpSize);
+
     Err = cuLaunchKernel(KernelInfo->Func, CudaBlocksPerGrid, /* gridDimY */ 1,
                          /* gridDimZ */ 1, CudaThreadsPerBlock,
                          /* blockDimY */ 1, /* blockDimZ */ 1,
-                         DynamicMemorySize, Stream, &Args[0], nullptr);
+                         DynamicSharedMemory, Stream, &Args[0], nullptr);
     if (!checkResult(Err, "Error returned from cuLaunchKernel\n"))
       return OFFLOAD_FAIL;
 
     DP("Launch of entry point at " DPxMOD " successful!\n",
        DPxPTR(TgtEntryPtr));
 
     return OFFLOAD_SUCCESS;
   }