Merge branch 'main' into enh-launch-kernel-extended

Author: achirkin

cpp/include/raft/linalg/detail/strided_reduction.cuh

@@ -17,6 +17,16 @@ namespace raft {
 namespace linalg {
 namespace detail {
 
+/**
+ * Upper bound on the number of blocks covering the reduced dimension, i.e. on @c gridDim.y .
+ *
+ * Both kernels below stride over the reduced dimension, so covering it with fewer blocks than
+ * elements is always correct. The bound must stay at or below the CUDA limit of 65535 on the y
+ * dimension of the grid, which the reduced dimension can otherwise easily exceed: reducing a
+ * matrix of 10M rows along the rows is a common case.
+ */
+constexpr int kMaxBlocksDimY = 8192;
+
 // Kernel to perform summation along the strided dimension
 // of the matrix, i.e. reduce along columns for row major or reduce along rows
 // for column major layout
@@ -150,10 +160,10 @@ void stridedReduction(OutType* dots,
     constexpr dim3 Block(ColsPerBlk, TPB / ColsPerBlk);
     constexpr int MinRowsPerThread = 16;
     constexpr int MinRowsPerBlk    = Block.y * MinRowsPerThread;
-    constexpr int MaxBlocksDimY    = 8192;
 
-    const dim3 grid(raft::ceildiv(D, (IdxType)ColsPerBlk),
-                    raft::min((IdxType)MaxBlocksDimY, raft::ceildiv(N, (IdxType)MinRowsPerBlk)));
+    const dim3 grid(
+      raft::div_rounding_up_safe(D, (IdxType)ColsPerBlk),
+      raft::min((IdxType)kMaxBlocksDimY, raft::div_rounding_up_safe(N, (IdxType)MinRowsPerBlk)));
     const size_t shmemSize = sizeof(OutType) * Block.x * 2;
 
     raft::launch_kernel({stream, shmemSize},
@@ -169,10 +179,11 @@ void stridedReduction(OutType* dots,
   } else {
     // Arbitrary numbers for now, probably need to tune
     const dim3 thrds(32, 16);
-    IdxType elemsPerThread = raft::ceildiv(N, (IdxType)thrds.y);
+    IdxType elemsPerThread = raft::div_rounding_up_safe(N, (IdxType)thrds.y);
     elemsPerThread         = (elemsPerThread > 8) ? 8 : elemsPerThread;
-    const dim3 nblks(raft::ceildiv(D, (IdxType)thrds.x),
-                     raft::ceildiv(N, (IdxType)thrds.y * elemsPerThread));
+    const dim3 nblks(raft::div_rounding_up_safe(D, (IdxType)thrds.x),
+                     raft::min((IdxType)kMaxBlocksDimY,
+                               raft::div_rounding_up_safe(N, (IdxType)thrds.y * elemsPerThread)));
     const size_t shmemSize = sizeof(OutType) * thrds.x * thrds.y;
 
     raft::launch_kernel({stream, shmemSize},

cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh

@@ -6,6 +6,7 @@
 #pragma once
 
 #include <raft/core/detail/macros.hpp>
+#include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resource/device_properties.hpp>
 #include <raft/core/resource/dry_run_flag.hpp>
 #include <raft/core/resource/thrust_policy.hpp>
@@ -174,6 +175,7 @@ Graph_COO<vertex_t, edge_t, weight_t> MST_solver<vertex_t, edge_t, weight_t, alt
 
     // copy this iteration's results and store
     prev_mst_edge_count.set_value_async(curr_mst_edge_count, stream);
+    resource::sync_stream(handle, stream);
   }
 
   // result packaging

cpp/tests/linalg/strided_reduction.cu

@@ -6,9 +6,11 @@
 #include "../test_utils.cuh"
 #include "reduce.cuh"
 
+#include <raft/core/device_mdarray.hpp>
 #include <raft/core/operators.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/linalg/strided_reduction.cuh>
+#include <raft/matrix/init.cuh>
 #include <raft/random/rng.cuh>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
@@ -122,5 +124,44 @@ INSTANTIATE_TEST_CASE_P(stridedReductionTests, stridedReductionTestF, ::testing:
 
 INSTANTIATE_TEST_CASE_P(stridedReductionTests, stridedReductionTestD, ::testing::ValuesIn(inputsd));
 
+/*
+ * The reduced dimension is mapped onto the y dimension of the grid, which CUDA limits to 65535
+ * blocks. That is far fewer blocks than the number of rows a caller may reduce along, so the grid
+ * must be bounded and the kernels must stride over the remaining rows.
+ */
+TEST(stridedReductionTest, LargeReducedDimension)
+{
+  raft::resources handle;
+  auto stream = resource::get_cuda_stream(handle);
+  // More rows than the largest grid the kernels are launched with can cover one row per thread.
+  constexpr int64_t kRows = 8'400'000;
+  constexpr int64_t kCols = 2;
+
+  auto data = raft::make_device_matrix<float, int64_t>(handle, kRows, kCols);
+  raft::matrix::fill(handle, data.view(), 1.0f);
+  auto data_view = raft::make_const_mdspan(data.view());
+
+  // Summing into the input type takes the compensated-summation path.
+  auto sums_same_type = raft::make_device_vector<float, int64_t>(handle, kCols);
+  strided_reduction(handle, data_view, sums_same_type.view(), 0.0f);
+
+  // Summing into a wider type takes the generic path, which handles an arbitrary reduce operation.
+  auto sums_wider_type = raft::make_device_vector<double, int64_t>(handle, kCols);
+  strided_reduction(handle, data_view, sums_wider_type.view(), 0.0, false, raft::cast_op<double>{});
+
+  resource::sync_stream(handle, stream);
+
+  ASSERT_TRUE(devArrMatch(static_cast<float>(kRows),
+                          sums_same_type.data_handle(),
+                          kCols,
+                          raft::CompareApprox<float>(1e-6f),
+                          stream));
+  ASSERT_TRUE(devArrMatch(static_cast<double>(kRows),
+                          sums_wider_type.data_handle(),
+                          kCols,
+                          raft::CompareApprox<double>(1e-12),
+                          stream));
+}
+
 }  // end namespace linalg
 }  // end namespace raft