[SYCL][CUDA] Add env variable to specify max local mem size (#6173)

Previously the max local mem allocation in DPC++ for CUDA backend was 48KB for most devices.

However as https://docs.nvidia.com/cuda/ampere-tuning-guide/index.html notes, for the A100 the max local mem dynamic allocation is in fact 164KB.

This PR introduces an environment variable SYCL_PI_CUDA_MAX_LOCAL_MEM_SZ which allows you to manually specify the max local memory in bytes allowed to be allocated per kernel for a given application. If an invalid value is specified (one that exceeds the device's capabilities/is negative) then a runtime error will be thrown.

Using:
SYCL_PI_CUDA_MAX_LOCAL_MEM_SZ=166912 ./a.out

Allows the application to use up to 163KB of local memory, if the device supports it.

Author: hdelan

sycl/doc/EnvironmentVariables.md

@@ -69,6 +69,12 @@ Note that all device selectors will throw an exception if the filtered list of d
 
 `(*) Note: Any means this environment variable is effective when set to any non-null value.`
 
+## Controlling DPC++ CUDA Plugin
+
+| Environment variable | Values | Description |
+| -------------------- | ------ | ----------- |
+| `SYCL_PI_CUDA_MAX_LOCAL_MEM_SIZE` | Integer | Specifies the maximum size of a local memory allocation in bytes. If the value exceeds the device's capabilities then a `sycl::runtime_error` is thrown. In order for the full error message to be printed, `SYCL_RT_WARNING_LEVEL=2` must be set. The default value for `SYCL_PI_CUDA_MAX_LOCAL_MEM_SIZE` is determined by the hardware. |
+
 ## Tools variables
 
 | Environment variable | Values | Description |

sycl/include/CL/sycl/detail/common.hpp

@@ -131,19 +131,21 @@ static inline std::string codeToString(cl_int code) {
 #ifndef SYCL_SUPPRESS_EXCEPTIONS
 #include <CL/sycl/exception.hpp>
 // SYCL 1.2.1 exceptions
-#define __SYCL_REPORT_OCL_ERR_TO_EXC(expr, exc)                                \
+#define __SYCL_REPORT_OCL_ERR_TO_EXC(expr, exc, str)                           \
   {                                                                            \
     auto code = expr;                                                          \
     if (code != CL_SUCCESS) {                                                  \
+      std::string err_str =                                                    \
+          str ? "\n" + std::string(str) + "\n" : std::string{};                \
       throw exc(__SYCL_OCL_ERROR_REPORT +                                      \
-                    cl::sycl::detail::codeToString(code),                      \
+                    cl::sycl::detail::codeToString(code) + err_str,            \
                 code);                                                         \
     }                                                                          \
   }
-#define __SYCL_REPORT_OCL_ERR_TO_EXC_THROW(code, exc)                          \
-  __SYCL_REPORT_OCL_ERR_TO_EXC(code, exc)
+#define __SYCL_REPORT_OCL_ERR_TO_EXC_THROW(code, exc, str)                     \
+  __SYCL_REPORT_OCL_ERR_TO_EXC(code, exc, str)
 #define __SYCL_REPORT_OCL_ERR_TO_EXC_BASE(code)                                \
-  __SYCL_REPORT_OCL_ERR_TO_EXC(code, cl::sycl::runtime_error)
+  __SYCL_REPORT_OCL_ERR_TO_EXC(code, cl::sycl::runtime_error, nullptr)
 #else
 #define __SYCL_REPORT_OCL_ERR_TO_EXC_BASE(code)                                \
   __SYCL_REPORT_OCL_ERR_TO_STREAM(code)
@@ -164,15 +166,19 @@ static inline std::string codeToString(cl_int code) {
 #ifdef __SYCL_SUPPRESS_OCL_ERROR_REPORT
 // SYCL 1.2.1 exceptions
 #define __SYCL_CHECK_OCL_CODE(X) (void)(X)
-#define __SYCL_CHECK_OCL_CODE_THROW(X, EXC) (void)(X)
+#define __SYCL_CHECK_OCL_CODE_THROW(X, EXC, STR)                               \
+  {                                                                            \
+    (void)(X);                                                                 \
+    (void)(STR);                                                               \
+  }
 #define __SYCL_CHECK_OCL_CODE_NO_EXC(X) (void)(X)
 // SYCL 2020 exceptions
 #define __SYCL_CHECK_CODE_THROW_VIA_ERRC(X, ERRC) (void)(X)
 #else
 // SYCL 1.2.1 exceptions
 #define __SYCL_CHECK_OCL_CODE(X) __SYCL_REPORT_OCL_ERR_TO_EXC_BASE(X)
-#define __SYCL_CHECK_OCL_CODE_THROW(X, EXC)                                    \
-  __SYCL_REPORT_OCL_ERR_TO_EXC_THROW(X, EXC)
+#define __SYCL_CHECK_OCL_CODE_THROW(X, EXC, STR)                               \
+  __SYCL_REPORT_OCL_ERR_TO_EXC_THROW(X, EXC, STR)
 #define __SYCL_CHECK_OCL_CODE_NO_EXC(X) __SYCL_REPORT_OCL_ERR_TO_STREAM(X)
 // SYCL 2020 exceptions
 #define __SYCL_CHECK_CODE_THROW_VIA_ERRC(X, ERRC)                              \

sycl/plugins/cuda/pi_cuda.cpp

@@ -2867,6 +2867,28 @@ pi_result cuda_piEnqueueKernelLaunch(
       retImplEv->start();
     }
 
+    // Set local mem max size if env var is present
+    static const char *local_mem_sz_ptr =
+        std::getenv("SYCL_PI_CUDA_MAX_LOCAL_MEM_SIZE");
+
+    if (local_mem_sz_ptr) {
+      int device_max_local_mem = 0;
+      cuDeviceGetAttribute(
+          &device_max_local_mem,
+          CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
+          command_queue->get_device()->get());
+
+      static const int env_val = std::atoi(local_mem_sz_ptr);
+      if (env_val <= 0 || env_val > device_max_local_mem) {
+        setErrorMessage("Invalid value specified for "
+                        "SYCL_PI_CUDA_MAX_LOCAL_MEM_SIZE",
+                        PI_PLUGIN_SPECIFIC_ERROR);
+        return PI_PLUGIN_SPECIFIC_ERROR;
+      }
+      PI_CHECK_ERROR(cuFuncSetAttribute(
+          cuFunc, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, env_val));
+    }
+
     retError = PI_CHECK_ERROR(cuLaunchKernel(
         cuFunc, blocksPerGrid[0], blocksPerGrid[1], blocksPerGrid[2],
         threadsPerBlock[0], threadsPerBlock[1], threadsPerBlock[2], local_size,

sycl/plugins/cuda/pi_cuda.hpp

@@ -442,6 +442,8 @@ struct _pi_queue {
 
   _pi_context *get_context() const { return context_; };
 
+  _pi_device *get_device() const { return device_; };
+
   pi_uint32 increment_reference_count() noexcept { return ++refCount_; }
 
   pi_uint32 decrement_reference_count() noexcept { return --refCount_; }

sycl/source/detail/error_handling/enqueue_kernel.cpp

@@ -22,7 +22,7 @@ namespace detail {
 
 namespace enqueue_kernel_launch {
 
-bool handleInvalidWorkGroupSize(const device_impl &DeviceImpl, pi_kernel Kernel,
+void handleInvalidWorkGroupSize(const device_impl &DeviceImpl, pi_kernel Kernel,
                                 const NDRDescT &NDRDesc) {
   const bool HasLocalSize = (NDRDesc.LocalSize[0] != 0);
 
@@ -246,7 +246,7 @@ bool handleInvalidWorkGroupSize(const device_impl &DeviceImpl, pi_kernel Kernel,
       "PI backend failed. PI backend returns: " + codeToString(Error), Error);
 }
 
-bool handleInvalidWorkItemSize(const device_impl &DeviceImpl,
+void handleInvalidWorkItemSize(const device_impl &DeviceImpl,
                                const NDRDescT &NDRDesc) {
 
   const plugin &Plugin = DeviceImpl.getPlugin();
@@ -265,10 +265,9 @@ bool handleInvalidWorkItemSize(const device_impl &DeviceImpl,
               " > " + std::to_string(MaxWISize[I]),
           PI_INVALID_WORK_ITEM_SIZE);
   }
-  return 0;
 }
 
-bool handleInvalidValue(const device_impl &DeviceImpl,
+void handleInvalidValue(const device_impl &DeviceImpl,
                         const NDRDescT &NDRDesc) {
   const plugin &Plugin = DeviceImpl.getPlugin();
   RT::PiDevice Device = DeviceImpl.getHandleRef();
@@ -293,8 +292,8 @@ bool handleInvalidValue(const device_impl &DeviceImpl,
       "Native API failed. Native API returns: " + codeToString(Error), Error);
 }
 
-bool handleError(pi_result Error, const device_impl &DeviceImpl,
-                 pi_kernel Kernel, const NDRDescT &NDRDesc) {
+void handleErrorOrWarning(pi_result Error, const device_impl &DeviceImpl,
+                          pi_kernel Kernel, const NDRDescT &NDRDesc) {
   assert(Error != PI_SUCCESS &&
          "Success is expected to be handled on caller side");
   switch (Error) {
@@ -343,6 +342,14 @@ bool handleError(pi_result Error, const device_impl &DeviceImpl,
   case PI_INVALID_VALUE:
     return handleInvalidValue(DeviceImpl, NDRDesc);
 
+  case PI_PLUGIN_SPECIFIC_ERROR:
+    // checkPiResult does all the necessary handling for
+    // PI_PLUGIN_SPECIFIC_ERROR, making sure an error is thrown or not,
+    // depending on whether PI_PLUGIN_SPECIFIC_ERROR contains an error or a
+    // warning. It also ensures that the contents of the error message buffer
+    // (used only by PI_PLUGIN_SPECIFIC_ERROR) get handled correctly.
+    return DeviceImpl.getPlugin().checkPiResult(Error);
+
     // TODO: Handle other error codes
 
   default:

sycl/source/detail/error_handling/error_handling.hpp

@@ -25,7 +25,8 @@ namespace enqueue_kernel_launch {
 ///
 /// This function actually never returns and always throws an exception with
 /// error description.
-bool handleError(pi_result, const device_impl &, pi_kernel, const NDRDescT &);
+void handleErrorOrWarning(pi_result, const device_impl &, pi_kernel,
+                          const NDRDescT &);
 } // namespace enqueue_kernel_launch
 
 } // namespace detail

sycl/source/detail/plugin.hpp

@@ -114,8 +114,8 @@ class plugin {
   /// \throw Exception if pi_result is not a PI_SUCCESS.
   template <typename Exception = cl::sycl::runtime_error>
   void checkPiResult(RT::PiResult pi_result) const {
+    char *message = nullptr;
     if (pi_result == PI_PLUGIN_SPECIFIC_ERROR) {
-      char *message = nullptr;
       pi_result = call_nocheck<PiApiKind::piPluginGetLastError>(&message);
 
       // If the warning level is greater then 2 emit the message
@@ -126,7 +126,7 @@ class plugin {
       if (pi_result == PI_SUCCESS)
         return;
     }
-    __SYCL_CHECK_OCL_CODE_THROW(pi_result, Exception);
+    __SYCL_CHECK_OCL_CODE_THROW(pi_result, Exception, message);
   }
 
   /// \throw SYCL 2020 exception(errc) if pi_result is not PI_SUCCESS

sycl/source/detail/scheduler/commands.cpp

@@ -2185,8 +2185,8 @@ cl_int enqueueImpKernel(
     // If we have got non-success error code, let's analyze it to emit nice
     // exception explaining what was wrong
     const device_impl &DeviceImpl = *(Queue->getDeviceImplPtr());
-    return detail::enqueue_kernel_launch::handleError(Error, DeviceImpl, Kernel,
-                                                      NDRDesc);
+    detail::enqueue_kernel_launch::handleErrorOrWarning(Error, DeviceImpl,
+                                                        Kernel, NDRDesc);
   }
 
   return PI_SUCCESS;