fix(rocsparse): Fix SpMM batching with COO segmented algorithm when batch count >65536

projects/rocsparse/CHANGELOG.md

@@ -3,6 +3,11 @@
 Documentation for rocSPARSE is available at
 [https://rocm.docs.amd.com/projects/rocSPARSE/en/latest/](https://rocm.docs.amd.com/projects/rocSPARSE/en/latest/).
 
+## (Unreleased) rocSPARSE 5.1.0
+
+### Resolved issues
+* Fixed `rocsparse_spmm` with the segmented COO algorithm, which failed with `hipErrorInvalidConfiguration` for batch counts exceeding 65535 because the batch dimension of the kernel launch grid exceeded the maximum grid dimension.
+
 ## (Unreleased) rocSPARSE 5.0.0
 
 ### Added

projects/rocsparse/CMakeLists.txt

@@ -196,7 +196,7 @@ endif()
 message(STATUS "GPU_TARGETS: ${GPU_TARGETS}")
 
 # Setup version
-set(VERSION_STRING "5.0.0")
+set(VERSION_STRING "5.1.0")
 set(SOVERSION_STRING "1.0")
 set(MIN_ROCPRIM_VERSION "4.0.0" CACHE STRING "Minimum version of rocPRIM.")
 set(MIN_ROCBLAS_VERSION "4.1.0" CACHE STRING "Minimum version of rocBLAS.")

projects/rocsparse/clients/tests/test_spmm_batched_coo.yaml

@@ -494,3 +494,27 @@ Tests:
   rand_gen_max: 100.0
   convert_to_int: true
   filename: [mac_econ_fwd500]
+
+# Large batch count exercising the coo segmented batch grid-stride loop.
+# Covers the nn/nt kernels (segmented only supports transA == none) and both
+# the main (n >= 8) and remainder (n < 8) column tilings via N.
+- name: spmm_batched_coo_file
+  category: stress
+  function: spmm_batched_coo
+  indextype: *i32
+  precision: *single_only_precisions
+  M: 1
+  N: [2, 65]
+  K: 1
+  batch_count_A: [1]
+  batch_count_B: [65537]
+  batch_count_C: [65537]
+  alpha_beta: *alpha_beta_range_stress
+  transA: [rocsparse_operation_none]
+  transB: [rocsparse_operation_none, rocsparse_operation_transpose]
+  baseA: [rocsparse_index_base_zero]
+  matrix: [rocsparse_matrix_file_rocalution]
+  spmm_alg: [rocsparse_spmm_alg_coo_segmented]
+  orderB: [rocsparse_order_column]
+  orderC: [rocsparse_order_column]
+  filename: [nos4]

projects/rocsparse/library/src/common/rocsparse_common.cpp

@@ -126,11 +126,10 @@ namespace rocsparse
     }
 
     template <uint32_t BLOCKSIZE, typename I, typename A, typename T>
-    ROCSPARSE_DEVICE_ILF void scale_2d_device(
-        I m, I n, int64_t ld, int64_t stride, T value, A* __restrict__ array, rocsparse_order order)
+    ROCSPARSE_DEVICE_ILF void
+        scale_2d_device(I m, I n, int64_t ld, T value, A* __restrict__ array, rocsparse_order order)
     {
-        I gid   = hipBlockIdx_x * BLOCKSIZE + hipThreadIdx_x;
-        I batch = hipBlockIdx_y;
+        I gid = hipBlockIdx_x * BLOCKSIZE + hipThreadIdx_x;
 
         if(gid >= m * n)
         {
@@ -142,11 +141,11 @@ namespace rocsparse
 
         if(value == static_cast<T>(0))
         {
-            array[lid + ld * wid + stride * batch] = static_cast<A>(0);
+            array[lid + ld * wid] = static_cast<A>(0);
         }
         else
         {
-            array[lid + ld * wid + stride * batch] *= value;
+            array[lid + ld * wid] *= value;
         }
     }
 
@@ -207,6 +206,7 @@ namespace rocsparse
     void scale_2d_kernel(I       m,
                          I       n,
                          int64_t ld,
+                         int64_t batch_count,
                          int64_t stride,
                          ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, scalar),
                          A* __restrict__ array,
@@ -217,7 +217,11 @@ namespace rocsparse
 
         if(scalar != static_cast<T>(1))
         {
-            rocsparse::scale_2d_device<BLOCKSIZE>(m, n, ld, stride, scalar, array, order);
+            for(int64_t batch = hipBlockIdx_y; batch < batch_count; batch += hipGridDim_y)
+            {
+                rocsparse::scale_2d_device<BLOCKSIZE>(
+                    m, n, ld, scalar, rocsparse::load_pointer(array, batch, stride), order);
+            }
         }
     }
 
@@ -321,13 +325,14 @@ rocsparse_status rocsparse::scale_2d_array(rocsparse_handle handle,
     {
         RETURN_IF_HIPLAUNCHKERNELGGL_ERROR(
             (rocsparse::scale_2d_kernel<256>),
-            dim3((int64_t(m) * n - 1) / 256 + 1, batch_count),
+            dim3((int64_t(m) * n - 1) / 256 + 1, (batch_count > 65535) ? 65535 : batch_count),
             dim3(256),
             0,
             handle->stream,
             m,
             n,
             ld,
+            batch_count,
             stride,
             ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, scalar_device_host),
             array,

projects/rocsparse/library/src/level3/coomm/segmented/kernel_declarations.h

@@ -41,6 +41,7 @@ namespace rocsparse
                                            I       N,
                                            I       K,
                                            int64_t nnz,
+                                           int64_t batch_count,
                                            int64_t batch_stride_A,
                                            ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha),
                                            I* __restrict__ row_block_red,
@@ -75,6 +76,7 @@ namespace rocsparse
                                                 I       N,
                                                 I       K,
                                                 int64_t nnz,
+                                                int64_t batch_count,
                                                 int64_t batch_stride_A,
                                                 ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha),
                                                 I* __restrict__ row_block_red,

projects/rocsparse/library/src/level3/coomm/segmented/kernel_definitions.h

@@ -43,6 +43,7 @@ namespace rocsparse
                                            I       N,
                                            I       K,
                                            int64_t nnz,
+                                           int64_t batch_count,
                                            int64_t batch_stride_A,
                                            ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha),
                                            I* __restrict__ row_block_red,
@@ -64,28 +65,32 @@ namespace rocsparse
 
         if(alpha != static_cast<T>(0))
         {
-            rocsparse::coommnn_segmented_main_device<BLOCKSIZE, WF_SIZE, LOOPS, TRANSB>(
-                conj_A,
-                conj_B,
-                M,
-                N,
-                K,
-                nnz,
-                batch_stride_A,
-                alpha,
-                row_block_red,
-                val_block_red,
-                coo_row_ind,
-                coo_col_ind,
-                coo_val,
-                dense_B,
-                ldb,
-                batch_stride_B,
-                dense_C,
-                ldc,
-                batch_stride_C,
-                order_C,
-                idx_base);
+            for(int64_t batch = hipBlockIdx_z; batch < batch_count; batch += hipGridDim_z)
+            {
+                rocsparse::coommnn_segmented_main_device<BLOCKSIZE, WF_SIZE, LOOPS, TRANSB>(
+                    conj_A,
+                    conj_B,
+                    M,
+                    N,
+                    K,
+                    nnz,
+                    batch_stride_A,
+                    alpha,
+                    row_block_red,
+                    val_block_red,
+                    coo_row_ind,
+                    coo_col_ind,
+                    coo_val,
+                    dense_B,
+                    ldb,
+                    batch_stride_B,
+                    dense_C,
+                    ldc,
+                    batch_stride_C,
+                    order_C,
+                    idx_base,
+                    batch);
+            }
         }
     }
 
@@ -106,6 +111,7 @@ namespace rocsparse
                                                 I       N,
                                                 I       K,
                                                 int64_t nnz,
+                                                int64_t batch_count,
                                                 int64_t batch_stride_A,
                                                 ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha),
                                                 I* __restrict__ row_block_red,
@@ -128,29 +134,33 @@ namespace rocsparse
 
         if(alpha != static_cast<T>(0))
         {
-            rocsparse::coommnn_segmented_remainder_device<BLOCKSIZE, WF_SIZE, LOOPS, TRANSB>(
-                conj_A,
-                conj_B,
-                colB_offset,
-                M,
-                N,
-                K,
-                nnz,
-                batch_stride_A,
-                alpha,
-                row_block_red,
-                val_block_red,
-                coo_row_ind,
-                coo_col_ind,
-                coo_val,
-                dense_B,
-                ldb,
-                batch_stride_B,
-                dense_C,
-                ldc,
-                batch_stride_C,
-                order_C,
-                idx_base);
+            for(int64_t batch = hipBlockIdx_z; batch < batch_count; batch += hipGridDim_z)
+            {
+                rocsparse::coommnn_segmented_remainder_device<BLOCKSIZE, WF_SIZE, LOOPS, TRANSB>(
+                    conj_A,
+                    conj_B,
+                    colB_offset,
+                    M,
+                    N,
+                    K,
+                    nnz,
+                    batch_stride_A,
+                    alpha,
+                    row_block_red,
+                    val_block_red,
+                    coo_row_ind,
+                    coo_col_ind,
+                    coo_val,
+                    dense_B,
+                    ldb,
+                    batch_stride_B,
+                    dense_C,
+                    ldc,
+                    batch_stride_C,
+                    order_C,
+                    idx_base,
+                    batch);
+            }
         }
     }
 }
@@ -164,6 +174,7 @@ namespace rocsparse
             I       N,                                                                 \
             I       K,                                                                 \
             int64_t nnz,                                                               \
+            int64_t batch_count,                                                       \
             int64_t batch_stride_A,                                                    \
             ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha),                             \
             I* __restrict__ row_block_red,                                             \
@@ -191,6 +202,7 @@ namespace rocsparse
             I       N,                                                                      \
             I       K,                                                                      \
             int64_t nnz,                                                                    \
+            int64_t batch_count,                                                            \
             int64_t batch_stride_A,                                                         \
             ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha),                                  \
             I* __restrict__ row_block_red,                                                  \

projects/rocsparse/library/src/level3/coomm_device_segmented.h

@@ -56,15 +56,14 @@ namespace rocsparse
                                                             int64_t              ldc,
                                                             int64_t              batch_stride_C,
                                                             rocsparse_order      order_C,
-                                                            rocsparse_index_base idx_base)
+                                                            rocsparse_index_base idx_base,
+                                                            int64_t              batch)
     {
         const int tid = hipThreadIdx_x;
         const int bid = hipBlockIdx_x;
         const int lid = tid & (WF_SIZE - 1);
         const int wid = tid / WF_SIZE;
 
-        const int batch = hipBlockIdx_z;
-
         __shared__ I shared_row[BLOCKSIZE];
         __shared__ T shared_val_prev[WF_SIZE];
         __shared__ T shared_val[BLOCKSIZE * WF_SIZE];
@@ -259,15 +258,14 @@ namespace rocsparse
                                                                  int64_t         ldc,
                                                                  int64_t         batch_stride_C,
                                                                  rocsparse_order order_C,
-                                                                 rocsparse_index_base idx_base)
+                                                                 rocsparse_index_base idx_base,
+                                                                 int64_t              batch)
     {
         const int tid = hipThreadIdx_x;
         const int bid = hipBlockIdx_x;
         const int lid = tid & (WF_SIZE - 1);
         const int wid = tid / WF_SIZE;
 
-        const int batch = hipBlockIdx_z;
-
         __shared__ I shared_row[BLOCKSIZE];
         __shared__ T shared_val_prev[WF_SIZE];
         __shared__ T shared_val[BLOCKSIZE * WF_SIZE];
@@ -491,47 +489,51 @@ namespace rocsparse
                                       C*              dense_C,
                                       int64_t         ldc,
                                       int64_t         batch_stride_C,
-                                      rocsparse_order order_C)
+                                      rocsparse_order order_C,
+                                      int64_t         batch_count)
     {
-        const int tid   = hipThreadIdx_x;
-        const int batch = hipBlockIdx_z;
+        const int tid = hipThreadIdx_x;
 
         // Shared memory to hold row indices and values for segmented reduction
         __shared__ I shared_row[BLOCKSIZE];
         __shared__ T shared_val[BLOCKSIZE];
 
         const I col = hipBlockIdx_x;
 
-        for(I i = 0; i < nblocks; i += BLOCKSIZE)
+        for(int64_t batch = hipBlockIdx_z; batch < batch_count; batch += hipGridDim_z)
         {
-            // Copy data to reduction buffers
-            shared_row[tid] = (tid + i < nblocks) ? row_block_red[tid + i + nblocks * batch] : -1;
-            shared_val[tid] = (tid + i < nblocks)
-                                  ? val_block_red[tid + i + nblocks * col + nblocks * n * batch]
-                                  : static_cast<T>(0);
+            for(I i = 0; i < nblocks; i += BLOCKSIZE)
+            {
+                // Copy data to reduction buffers
+                shared_row[tid]
+                    = (tid + i < nblocks) ? row_block_red[tid + i + nblocks * batch] : -1;
+                shared_val[tid] = (tid + i < nblocks)
+                                      ? val_block_red[tid + i + nblocks * col + nblocks * n * batch]
+                                      : static_cast<T>(0);
 
-            __syncthreads();
+                __syncthreads();
 
-            // Do segmented block reduction
-            segmented_blockreduce<BLOCKSIZE>(shared_row, shared_val);
+                // Do segmented block reduction
+                segmented_blockreduce<BLOCKSIZE>(shared_row, shared_val);
 
-            // Add reduced sum to C if valid
-            const I row   = shared_row[tid];
-            const I rowp1 = (tid < BLOCKSIZE - 1) ? shared_row[tid + 1] : -1;
+                // Add reduced sum to C if valid
+                const I row   = shared_row[tid];
+                const I rowp1 = (tid < BLOCKSIZE - 1) ? shared_row[tid + 1] : -1;
 
-            if(row != rowp1 && row >= 0)
-            {
-                if(order_C == rocsparse_order_column)
-                {
-                    dense_C[row + ldc * col + batch_stride_C * batch] += shared_val[tid];
-                }
-                else
+                if(row != rowp1 && row >= 0)
                 {
-                    dense_C[col + ldc * row + batch_stride_C * batch] += shared_val[tid];
+                    if(order_C == rocsparse_order_column)
+                    {
+                        dense_C[row + ldc * col + batch_stride_C * batch] += shared_val[tid];
+                    }
+                    else
+                    {
+                        dense_C[col + ldc * row + batch_stride_C * batch] += shared_val[tid];
+                    }
                 }
-            }
 
-            __syncthreads();
+                __syncthreads();
+            }
         }
     }
 }