[OpenMP] Add support for dynamic shared memory in new RTL

[jhuber6]

Summary

This patch adds support for using dynamic shared memory in the new
device runtime. The new function __kmpc_get_dynamic_shared will return a
pointer to the buffer of dynamic shared memory. Currently the amount of memory
allocated is set by an environment variable.

In the future this amount will be added to the amount used for the smart stack
which will be configured in a similar way.

[tianshilei1992]

LGTM with some nits.

openmp/libomptarget/DeviceRTL/include/Configuration.h

35

getDynamicMemorySize? I just feel since other words like "dynamic" or "size" we all use the full version, we could also use "memory" here?

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

547

LIBOMPTARGET_SHARED_MEMORY_SIZE?

[jhuber6]

Addressing nits and adding some documentation.

[JonChesterfield]

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

102

This struct is also defined in deviceRTLs/common/device_environment.h, which doesn't seem to have the extra field present. Any objection to moving that header up a couple of levels then including it from the various plugins and device runtimes?

uint64_t seems overly ambitious for shared size. The amdgpu hardware is limited to ~ 64k, can cuda go over 32 bits for this?

[jhuber6]

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

102

It would probably be a good idea to keep the device environment stuff consistent. Currently there might be some difference between the new and old runtime so we should make sure that they all share the same size at least.

And you're probably right, the argument in cuLaunchKernel that it's passed to only uses an unsigned int which is most likely 32 bits. Not likely someone is going to allocate a few billion bytes of shared memory.

[JonChesterfield]

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

102

Cool. I've wanted this for a while (rocm and trunk haven't always matched either) so will put up a patch

[JonChesterfield]

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

102

The dynamic_shared_size field is currently unused (and not defined in the devicertl), but the cuda plugin raises an error when the size of the global in the devicertl is different to the size declared in this file. I don't understand how that is working at present.

[JonChesterfield]

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

102

Specifically, this patch missed the deviceRTL definition:
libomptarget/deviceRTLs/common/device_environment.h
struct omptarget_device_environmentTy {

int32_t debug_level;
uint32_t num_devices;
uint32_t device_num;

};

[tianshilei1992]

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

102

This patch only adds the support in the new device runtime, as described in the title.

[JonChesterfield]

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

102

Right, but both old and new device runtime use the same host plugin, and the host plugin raises an error if this struct has the wrong size, and the old device runtime has one without the extra uint64_t. Just put a comment by the size check further down this file

930

size mismatch here with old runtime, as far as I can tell

openmp/docs/design/Runtimes.rst

 variables is defined below.
 
     * ``LIBOMPTARGET_DEBUG=<Num>``
     * ``LIBOMPTARGET_PROFILE=<Filename>``
     * ``LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD=<Num>``
     * ``LIBOMPTARGET_INFO=<Num>``
     * ``LIBOMPTARGET_HEAP_SIZE=<Num>``
     * ``LIBOMPTARGET_STACK_SIZE=<Num>``
+    * ``LIBOMPTARGET_SHARED_MEMORY_SIZE=<Num>``
 
 LIBOMPTARGET_DEBUG
 """"""""""""""""""
 
 ``LIBOMPTARGET_DEBUG`` controls whether or not debugging information will be
 displayed. This feature is only availible if ``libomptarget`` was built with
 ``-DOMPTARGET_DEBUG``. The debugging output provided is intended for use by
 ``libomptarget`` developers. More user-friendly output is presented when using

 LIBOMPTARGET_HEAP_SIZE
 """""""""""""""""""""""
 
 This environment variable sets the amount of memory in bytes that can be
 allocated using ``malloc`` and ``free`` for the CUDA plugin. This is necessary
 for some applications that allocate too much memory either through the user or
 globalization.
 
+LIBOMPTARGET_SHARED_MEMORY_SIZE
+"""""""""""""""""""""""""""""""
+
+This environment variable sets the amount of dynamic shared memory in bytes used 
+by the kernel once it is launched. A pointer to the dynamic memory buffer can 
+currently only be accessed using the ``__kmpc_get_dynamic_shared`` device 
+runtime call.
+
 .. toctree::
    :hidden:
    :maxdepth: 1
 
    Offloading
 
 LLVM/OpenMP Target Host Runtime Plugins (``libomptarget.rtl.XXXX``)
 -------------------------------------------------------------------

openmp/libomptarget/DeviceRTL/include/Configuration.h

 
 /// Return the number of devices in the system, same number as returned on the
 /// host by omp_get_num_devices.
 uint32_t getDeviceNum();
 
 /// Return the user choosen debug level.
 uint32_t getDebugLevel();
 
+/// Return the amount of dynamic shared memory that was allocated at launch.
+uint64_t getDynamicMemorySize();
+
 bool isDebugMode(DebugLevel Level);
 
 } // namespace config
 } // namespace _OMP
 
 #endif

openmp/libomptarget/DeviceRTL/include/Interface.h

 /// called by any thread, allocation happens *per thread*.
 void *__kmpc_alloc_shared(uint64_t Bytes);
 
 /// Deallocate \p Ptr. Needs to be called balanced with __kmpc_alloc_shared like
 /// a stack (push/pop). Can be called by any thread. \p Ptr has to be the
 /// allocated by __kmpc_alloc_shared by the same thread.
 void __kmpc_free_shared(void *Ptr, uint64_t Bytes);
 
+/// Get a pointer to the memory buffer containing dynamically allocated shared
+/// memory configured at launch.
+void *__kmpc_get_dynamic_shared();
+
 /// Allocate sufficient space for \p NumArgs sequential `void*` and store the
 /// allocation address in \p GlobalArgs.
 ///
 /// Called by the main thread prior to a parallel region.
 ///
 /// We also remember it in GlobalArgsPtr to ensure the worker threads and
 /// deallocation function know the allocation address too.
 void __kmpc_begin_sharing_variables(void ***GlobalArgs, uint64_t NumArgs);

openmp/libomptarget/DeviceRTL/include/State.h

 /// Free \p Ptr, alloated via allocShared, for \p Reason.
 ///
 /// Note: See the restrictions on __kmpc_free_shared for proper usage.
 void freeShared(void *Ptr, uint64_t Bytes, const char *Reason);
 
 /// Alloca \p Size bytes in global memory, if possible, for \p Reason.
 void *allocGlobal(uint64_t Size, const char *Reason);
 
+/// Return a pointer to the dynamic shared memory buffer.
+void *getDynamicBuffer();
+
 /// Free \p Ptr, alloated via allocGlobal, for \p Reason.
 void freeGlobal(void *Ptr, const char *Reason);
 
 } // namespace memory
 
 } // namespace _OMP
 
 #pragma omp end declare target
 
 #endif

openmp/libomptarget/DeviceRTL/src/Configuration.cpp

 #include "Types.h"
 
 using namespace _OMP;
 
 struct DeviceEnvironmentTy {
   uint32_t DebugLevel;
   uint32_t NumDevices;
   uint32_t DeviceNum;
+  uint64_t DynamicMemSize;
 };
 
 #pragma omp declare target
 
 extern uint32_t __omp_rtl_debug_kind;
 
 // TOOD: We want to change the name as soon as the old runtime is gone.
 DeviceEnvironmentTy CONSTANT(omptarget_device_environment)
     __attribute__((used));
 
 uint32_t config::getDebugLevel() {
   return __omp_rtl_debug_kind & omptarget_device_environment.DebugLevel;
 }
 
 uint32_t config::getNumDevices() {
   return omptarget_device_environment.NumDevices;
 }
 
 uint32_t config::getDeviceNum() {
   return omptarget_device_environment.DeviceNum;
 }
 
+uint64_t config::getDynamicMemorySize() {
+  return omptarget_device_environment.DynamicMemSize;
+}
+
 bool config::isDebugMode(config::DebugLevel Level) {
   return config::getDebugLevel() > Level;
 }
 
 #pragma omp end declare target

openmp/libomptarget/DeviceRTL/src/State.cpp

 using namespace _OMP;
 
 #pragma omp declare target
 
 /// Memory implementation
 ///
 ///{
 
+/// Add worst-case padding so that future allocations are properly aligned.
+constexpr const uint32_t Alignment = 8;
+
+/// External symbol to access dynamic shared memory.
+extern unsigned char DynamicSharedBuffer[] __attribute__((aligned(Alignment)));
+#pragma omp allocate(DynamicSharedBuffer) allocator(omp_pteam_mem_alloc)
+
 namespace {
 
 /// Fallback implementations are missing to trigger a link time error.
 /// Implementations for new devices, including the host, should go into a
 /// dedicated begin/end declare variant.
 ///
 ///{
 

 }
 
 void free(void *Ptr) {}
 }
 
 #pragma omp end declare variant
 ///}
 
-/// Add worst-case padding so that future allocations are properly aligned.
-constexpr const uint32_t Alignment = 8;
-
 /// A "smart" stack in shared memory.
 ///
 /// The stack exposes a malloc/free interface but works like a stack internally.
 /// 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.
if (Ptr >= &Data[0] && Ptr < &Data[state::SharedScratchpadSize]) {
     Usage[TId] -= AlignedBytes;
     return;
   }
   memory::freeGlobal(Ptr, "Slow path shared memory deallocation");
 }
 
 } // namespace
 
+void *memory::getDynamicBuffer() { return DynamicSharedBuffer; }
+
 void *memory::allocShared(uint64_t Bytes, const char *Reason) {
   return SharedMemorySmartStack.push(Bytes);
 }
 
 void memory::freeShared(void *Ptr, uint64_t Bytes, const char *Reason) {
   SharedMemorySmartStack.pop(Ptr, Bytes);
 }
 

 __attribute__((noinline)) void *__kmpc_alloc_shared(uint64_t Bytes) {
   return memory::allocShared(Bytes, "Frontend alloc shared");
 }
 
 __attribute__((noinline)) void __kmpc_free_shared(void *Ptr, uint64_t Bytes) {
   memory::freeShared(Ptr, Bytes, "Frontend free shared");
 }
 
+__attribute__((noinline)) void *__kmpc_get_dynamic_shared() {
+  return memory::getDynamicBuffer();
+}
+
 /// Allocate storage in shared memory to communicate arguments from the main
 /// thread to the workers in generic mode. If we exceed
 /// NUM_SHARED_VARIABLES_IN_SHARED_MEM we will malloc space for communication.
 constexpr uint64_t NUM_SHARED_VARIABLES_IN_SHARED_MEM = 64;
 
 [[clang::loader_uninitialized]] static void
     *SharedMemVariableSharingSpace[NUM_SHARED_VARIABLES_IN_SHARED_MEM];
 #pragma omp allocate(SharedMemVariableSharingSpace)                            \

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

struct KernelTy {
 
   KernelTy(CUfunction _Func, int8_t _ExecutionMode)
       : Func(_Func), ExecutionMode(_ExecutionMode) {}
 };
 
 /// 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;
   uint32_t device_num;
+  uint64_t dynamic_shared_size;
 };
 
 namespace {
 bool checkResult(CUresult Err, const char *ErrMsg) {
   if (Err == CUDA_SUCCESS)
     return true;
 
   REPORT("%s", ErrMsg);

 class DeviceRTLTy {
   int NumberOfDevices;
   // OpenMP environment properties
   int EnvNumTeams;
   int EnvTeamLimit;
   int EnvTeamThreadLimit;
   // OpenMP requires flags
   int64_t RequiresFlags;
+  // Amount of dynamic shared memory to use at launch.
+  uint64_t DynamicMemorySize;
 
   static constexpr const int HardTeamLimit = 1U << 16U; // 64k
   static constexpr const int HardThreadLimit = 1024;
   static constexpr const int DefaultNumTeams = 128;
   static constexpr const int DefaultNumThreads = 128;
 
   std::unique_ptr<StreamManagerTy> StreamManager;
   std::vector<DeviceDataTy> DeviceData;

 
 public:
   // This class should not be copied
   DeviceRTLTy(const DeviceRTLTy &) = delete;
   DeviceRTLTy(DeviceRTLTy &&) = delete;
 
   DeviceRTLTy()
       : NumberOfDevices(0), EnvNumTeams(-1), EnvTeamLimit(-1),
-        EnvTeamThreadLimit(-1), RequiresFlags(OMP_REQ_UNDEFINED) {
+        EnvTeamThreadLimit(-1), RequiresFlags(OMP_REQ_UNDEFINED),
+        DynamicMemorySize(0) {
 
     DP("Start initializing CUDA\n");
 
     CUresult Err = cuInit(0);
     if (Err == CUDA_ERROR_INVALID_HANDLE) {
       // Can't call cuGetErrorString if dlsym failed
       DP("Failed to load CUDA shared library\n");
       return;
if (const char *EnvStr = getenv("OMP_TEAMS_THREAD_LIMIT")) {
       EnvTeamThreadLimit = std::stoi(EnvStr);
       DP("Parsed OMP_TEAMS_THREAD_LIMIT=%d\n", EnvTeamThreadLimit);
     }
     if (const char *EnvStr = getenv("OMP_NUM_TEAMS")) {
       // OMP_NUM_TEAMS has been set
       EnvNumTeams = std::stoi(EnvStr);
       DP("Parsed OMP_NUM_TEAMS=%d\n", EnvNumTeams);
     }
+    if (const char *EnvStr = getenv("LIBOMPTARGET_SHARED_MEMORY_SIZE")) {
+      // LIBOMPTARGET_SHARED_MEMORY_SIZE has been set
+      DynamicMemorySize = std::stoi(EnvStr);
+      DP("Parsed LIBOMPTARGET_SHARED_MEMORY_SIZE", DynamicMemorySize);
+    }
 
     StreamManager =
         std::make_unique<StreamManagerTy>(NumberOfDevices, DeviceData);
 
     for (int I = 0; I < NumberOfDevices; ++I)
       DeviceAllocators.emplace_back(I, DeviceData);
 
     // Get the size threshold from environment variable
for (const __tgt_offload_entry *E = HostBegin; E != HostEnd; ++E) {
       addOffloadEntry(DeviceId, Entry);
     }
 
     // send device environment data to the device
     {
       // TODO: The device ID used here is not the real device ID used by OpenMP.
       omptarget_device_environmentTy DeviceEnv{
           0, static_cast<uint32_t>(NumberOfDevices),
-          static_cast<uint32_t>(DeviceId)};
+          static_cast<uint32_t>(DeviceId), DynamicMemorySize};
 
 #ifdef OMPTARGET_DEBUG
       if (const char *EnvStr = getenv("LIBOMPTARGET_DEVICE_RTL_DEBUG"))
         DeviceEnv.debug_level = std::stoi(EnvStr);
 #endif
 
       const char *DeviceEnvName = "omptarget_device_environment";
       CUdeviceptr DeviceEnvPtr;
       size_t CUSize;
 
       Err = cuModuleGetGlobal(&DeviceEnvPtr, &CUSize, Module, DeviceEnvName);
       if (Err == CUDA_SUCCESS) {
         if (CUSize != sizeof(DeviceEnv)) {
           REPORT(
               "Global device_environment '%s' - size mismatch (%zu != %zu)\n",
               DeviceEnvName, CUSize, sizeof(int32_t));
           CUDA_ERR_STRING(Err);
           return nullptr;
         }
 
         Err = cuMemcpyHtoD(DeviceEnvPtr, &DeviceEnv, CUSize);
         if (Err != CUDA_SUCCESS) {
INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
               ? (KernelInfo->ExecutionMode == GENERIC ? "Generic"
                                                       : "SPMD-Generic")
               : "SPMD"));
 
     CUstream Stream = getStream(DeviceId, AsyncInfo);
     Err = cuLaunchKernel(KernelInfo->Func, CudaBlocksPerGrid, /* gridDimY */ 1,
                          /* gridDimZ */ 1, CudaThreadsPerBlock,
                          /* blockDimY */ 1, /* blockDimZ */ 1,
-                         /* sharedMemBytes */ 0, Stream, &Args[0], nullptr);
+                         DynamicMemorySize, 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;
   }

openmp/libomptarget/test/api/omp_dynamic_shared_memory.c

+// RUN: %libomptarget-compile-nvptx64-nvidia-cuda -fopenmp-target-new-runtime
+// RUN: env LIBOMPTARGET_SHARED_MEMORY_SIZE=4 \
+// RUN:   %libomptarget-run-nvptx64-nvidia-cuda | %fcheck-nvptx64-nvidia-cuda
+// REQUIRES: nvptx64-nvidia-cuda
+
+#include <omp.h>
+#include <stdio.h>
+
+void *get_dynamic_shared() { return NULL; }
+#pragma omp begin declare variant match(device = {arch(nvptx64)})
+extern void *__kmpc_get_dynamic_shared();
+void *get_dynamic_shared() { return __kmpc_get_dynamic_shared(); }
+#pragma omp end declare variant
+
+int main() {
+  int x;
+#pragma omp target parallel map(from : x)
+  {
+    int *buf = get_dynamic_shared();
+#pragma omp barrier
+    if (omp_get_thread_num() == 0)
+      *buf = 1;
+#pragma omp barrier
+    if (omp_get_thread_num() == 1)
+      x = *buf;
+  }
+
+  // CHECK: PASS
+  if (x == 1)
+    printf("PASS\n");
+}