fix(cuda_helpers): fix data race in set_shmem_of_kernel and add docstring

Author: np96

cpp/src/utilities/cuda_helpers.cuh

@@ -12,6 +12,7 @@
 #include <thrust/host_vector.h>
 #include <thrust/tuple.h>
 #include <mutex>
+#include <shared_mutex>
 #include <raft/core/device_span.hpp>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
@@ -175,29 +176,49 @@ HDI To bit_cast(const From& src)
   return *(To*)(&src);
 }
 
+/**
+ * @brief Raises the dynamic shared-memory limit for a CUDA kernel, with caching.
+ *
+ * Calls cudaFuncSetAttribute(cudaFuncAttributeMaxDynamicSharedMemorySize) only when
+ * @p dynamic_request_size exceeds the previously set limit for @p function.  The
+ * per-kernel high-water mark is stored in a process-wide cache so that repeated
+ * calls with the same or smaller sizes are cheap shared-lock reads.
+ *
+ * Thread safety: safe to call concurrently from multiple host threads.
+ *
+ * @param function             Host pointer to the __global__ kernel function.
+ * @param dynamic_request_size Requested dynamic shared memory in bytes.
+ *                             A value of 0 is a no-op and always returns true.
+ * @return true  if the attribute was successfully set (or was already sufficient).
+ * @return false if cudaFuncSetAttribute failed (e.g. size exceeds device limit);
+ *               the sticky CUDA error is consumed so it cannot surface later.
+ */
 template <typename Function>
 inline bool set_shmem_of_kernel(Function* function, size_t dynamic_request_size)
 {
-  static std::mutex mtx;
+  static std::shared_mutex mtx;
   static std::unordered_map<Function*, size_t> shmem_sizes;
 
   if (dynamic_request_size != 0) {
     dynamic_request_size = raft::alignTo(dynamic_request_size, size_t(1024));
-    size_t current_size  = shmem_sizes[function];
+
+    {
+      std::shared_lock<std::shared_mutex> rlock(mtx);
+      auto it = shmem_sizes.find(function);
+      if (it != shmem_sizes.end() && dynamic_request_size <= it->second) { return true; }
+    }
+
+    std::unique_lock<std::shared_mutex> wlock(mtx);
+    size_t current_size = shmem_sizes.count(function) ? shmem_sizes[function] : 0;
     if (dynamic_request_size > current_size) {
-      std::lock_guard<std::mutex> lock(mtx);
-      current_size = shmem_sizes[function];
-
-      if (dynamic_request_size > current_size) {
-        auto err = cudaFuncSetAttribute(
-          function, cudaFuncAttributeMaxDynamicSharedMemorySize, dynamic_request_size);
-        if (err == cudaSuccess) {
-          shmem_sizes[function] = dynamic_request_size;
-          return true;
-        } else {
-          cudaGetLastError();  // clear sticky error so later RAFT_CHECK_CUDA doesn't catch it
-          return false;
-        }
+      auto err = cudaFuncSetAttribute(
+        function, cudaFuncAttributeMaxDynamicSharedMemorySize, dynamic_request_size);
+      if (err == cudaSuccess) {
+        shmem_sizes[function] = dynamic_request_size;
+        return true;
+      } else {
+        cudaGetLastError();  // clear sticky error so later RAFT_CHECK_CUDA doesn't catch it
+        return false;
       }
     }
   }

cpp/tests/routing/unit_tests/set_shmem_of_kernel.cu

@@ -14,58 +14,69 @@
 namespace cuopt {
 namespace test {
 
+/// @brief Dummy kernel used to test a zero-byte shared-memory request.
 __global__ void kernel_zero() {}
+/// @brief Dummy kernel used to test a normal (within-limit) shared-memory request.
 __global__ void kernel_normal() {}
+/// @brief Dummy kernel used to test a too-large shared-memory request (first call).
 __global__ void kernel_too_large_a() {}
+/// @brief Dummy kernel used to test a too-large shared-memory request (repeated call).
 __global__ void kernel_too_large_b() {}
+/// @brief Dummy kernel used to verify that a failed request leaves no sticky CUDA error.
 __global__ void kernel_sticky_error() {}
 
-// Zero request is a no-op and must return true.
+/// @brief Zero request is a no-op and must return true.
 TEST(set_shmem_of_kernel, zero_request)
 {
   EXPECT_TRUE(set_shmem_of_kernel(kernel_zero, 0));
   EXPECT_EQ(cudaSuccess, cudaGetLastError());
 }
 
-// A modest request well within device limits must succeed.
+/// @brief A modest request well within device limits must succeed.
 TEST(set_shmem_of_kernel, normal_request)
 {
   EXPECT_TRUE(set_shmem_of_kernel(kernel_normal, 4096));
   EXPECT_EQ(cudaSuccess, cudaGetLastError());
 }
 
-// Requesting more shared memory than the device supports must return false.
+/// @brief Requesting more shared memory than the device supports must return false.
 TEST(set_shmem_of_kernel, too_large_returns_false)
 {
   int shmem_max{};
-  cudaDeviceGetAttribute(&shmem_max, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0);
+  ASSERT_EQ(cudaSuccess,
+            cudaDeviceGetAttribute(&shmem_max, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0))
+    << "cudaDeviceGetAttribute(cudaDevAttrMaxSharedMemoryPerBlockOptin) failed";
   size_t too_large = static_cast<size_t>(shmem_max) + 1024;
 
   EXPECT_FALSE(set_shmem_of_kernel(kernel_too_large_a, too_large));
   EXPECT_EQ(cudaSuccess, cudaGetLastError());
 }
 
-// A second call with the same too-large size must still return false
+/// @brief A second call with the same too-large size must still return false.
 TEST(set_shmem_of_kernel, cache_not_poisoned_on_failure)
 {
   int shmem_max{};
-  cudaDeviceGetAttribute(&shmem_max, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0);
+  ASSERT_EQ(cudaSuccess,
+            cudaDeviceGetAttribute(&shmem_max, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0))
+    << "cudaDeviceGetAttribute(cudaDevAttrMaxSharedMemoryPerBlockOptin) failed";
   size_t too_large = static_cast<size_t>(shmem_max) + 1024;
 
   EXPECT_FALSE(set_shmem_of_kernel(kernel_too_large_b, too_large));
   EXPECT_FALSE(set_shmem_of_kernel(kernel_too_large_b, too_large));  // must not return true
   EXPECT_EQ(cudaSuccess, cudaGetLastError());
 }
 
-// A failed call must not leave a sticky CUDA error that would be caught
-// later by an unrelated RAFT_CHECK_CUDA.
+/// @brief A failed call must not leave a sticky CUDA error that would be caught
+/// later by an unrelated RAFT_CHECK_CUDA.
 TEST(set_shmem_of_kernel, no_sticky_error_after_failure)
 {
   int shmem_max{};
-  cudaDeviceGetAttribute(&shmem_max, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0);
+  ASSERT_EQ(cudaSuccess,
+            cudaDeviceGetAttribute(&shmem_max, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0))
+    << "cudaDeviceGetAttribute(cudaDevAttrMaxSharedMemoryPerBlockOptin) failed";
   size_t too_large = static_cast<size_t>(shmem_max) + 1024;
 
-  set_shmem_of_kernel(kernel_sticky_error, too_large);
+  EXPECT_FALSE(set_shmem_of_kernel(kernel_sticky_error, too_large));  // confirm failure branch taken
   EXPECT_EQ(cudaSuccess, cudaGetLastError());
 }