[Openmp-commits] [openmp] 3817ba1 - [OpenMP] Add environment variables to change stack / heap size in the CUDA plugin

[via Openmp-commits]

Author: Joseph Huber
Date: 2021-07-22T21:40:02-04:00
New Revision: 3817ba13aea3fdf189ee4048e4a5bf1eba2a3044

URL: https://github.com/llvm/llvm-project/commit/3817ba13aea3fdf189ee4048e4a5bf1eba2a3044
DIFF: https://github.com/llvm/llvm-project/commit/3817ba13aea3fdf189ee4048e4a5bf1eba2a3044.diff

LOG: [OpenMP] Add environment variables to change stack / heap size in the CUDA plugin

This patch adds support for two environment variables to configure the device.
``LIBOMPTARGET_STACK_SIZE`` sets the amount of memory in bytes that each thread
has for its stack. ``LIBOMPTARGET_HEAP_SIZE`` sets the amount of heap memory
that can be allocated using malloc / free on the device.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D106627

Added: 
    

Modified: 
    openmp/docs/design/Runtimes.rst
    openmp/libomptarget/plugins/cuda/src/rtl.cpp

Removed: 
    


################################################################################
diff  --git a/openmp/docs/design/Runtimes.rst b/openmp/docs/design/Runtimes.rst
index 08bdb450b3ac..b8b8da843b4f 100644
--- a/openmp/docs/design/Runtimes.rst
+++ b/openmp/docs/design/Runtimes.rst
@@ -30,6 +30,8 @@ variables is defined below.
     * ``LIBOMPTARGET_PROFILE=<Filename>``
     * ``LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD=<Num>``
     * ``LIBOMPTARGET_INFO=<Num>``
+    * ``LIBOMPTARGET_HEAP_SIZE=<Num>``
+    * ``LIBOMPTARGET_STACK_SIZE=<Num>``
 
 LIBOMPTARGET_DEBUG
 """"""""""""""""""
@@ -321,6 +323,21 @@ default. The solution is to add an explicit map clause in the target region.
       return sum;
     }
 
+LIBOMPTARGET_STACK_SIZE
+"""""""""""""""""""""""
+
+This environment variable sets the stack size in bytes for the CUDA plugin. This
+can be used to increase or decrease the standard amount of memory reserved for
+each thread's stack.
+
+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.
+
 .. toctree::
    :hidden:
    :maxdepth: 1

diff  --git a/openmp/libomptarget/plugins/cuda/src/rtl.cpp b/openmp/libomptarget/plugins/cuda/src/rtl.cpp
index 60b97dc1a462..da38c68a7142 100644
--- a/openmp/libomptarget/plugins/cuda/src/rtl.cpp
+++ b/openmp/libomptarget/plugins/cuda/src/rtl.cpp
@@ -642,11 +642,34 @@ class DeviceRTLTy {
       DeviceData[DeviceId].BlocksPerGrid = EnvTeamLimit;
     }
 
+    size_t StackLimit;
+    size_t HeapLimit;
+    if (const char *EnvStr = getenv("LIBOMPTARGET_STACK_SIZE")) {
+      StackLimit = std::stol(EnvStr);
+      if (cuCtxSetLimit(CU_LIMIT_STACK_SIZE, StackLimit) != CUDA_SUCCESS)
+        return OFFLOAD_FAIL;
+    } else {
+      if (cuCtxGetLimit(&StackLimit, CU_LIMIT_STACK_SIZE) != CUDA_SUCCESS)
+        return OFFLOAD_FAIL;
+    }
+    if (const char *EnvStr = getenv("LIBOMPTARGET_HEAP_SIZE")) {
+      HeapLimit = std::stol(EnvStr);
+      if (cuCtxSetLimit(CU_LIMIT_MALLOC_HEAP_SIZE, HeapLimit) != CUDA_SUCCESS)
+        return OFFLOAD_FAIL;
+    } else {
+      if (cuCtxGetLimit(&HeapLimit, CU_LIMIT_MALLOC_HEAP_SIZE) != CUDA_SUCCESS)
+        return OFFLOAD_FAIL;
+    }
+
     INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
          "Device supports up to %d CUDA blocks and %d threads with a "
          "warp size of %d\n",
          DeviceData[DeviceId].BlocksPerGrid,
          DeviceData[DeviceId].ThreadsPerBlock, DeviceData[DeviceId].WarpSize);
+    INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
+         "Device heap size is %d Bytes, device stack size is %d Bytes per "
+         "thread\n",
+         (int)HeapLimit, (int)StackLimit);
 
     // Set default number of teams
     if (EnvNumTeams > 0) {