[Openmp-commits] [PATCH] D106627: [OpenMP] Add environment variables to change stack / heap size in the CUDA plugin

[Joseph Huber via Phabricator via Openmp-commits]

jhuber6 created this revision.
jhuber6 added a reviewer: jdoerfert.
Herald added subscribers: guansong, yaxunl.
jhuber6 requested review of this revision.
Herald added subscribers: openmp-commits, sstefan1.
Herald added a project: OpenMP.

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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D106627

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


Index: openmp/libomptarget/plugins/cuda/src/rtl.cpp
===================================================================
--- openmp/libomptarget/plugins/cuda/src/rtl.cpp
+++ openmp/libomptarget/plugins/cuda/src/rtl.cpp
@@ -642,11 +642,34 @@
       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) {
Index: openmp/docs/design/Runtimes.rst
===================================================================
--- openmp/docs/design/Runtimes.rst
+++ openmp/docs/design/Runtimes.rst
@@ -30,6 +30,8 @@
     * ``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 @@
       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


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D106627.361058.patch
Type: text/x-patch
Size: 2661 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/openmp-commits/attachments/20210723/071e7290/attachment.bin>