Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jsandham merged 17 commits into from
Aug 7, 2026
Merged

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

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
6829c89
Fix SpMM COO segmented algorithm when using a batch count larger than…
jsandham Jul 24, 2026
46b9d2e
Clang formatting
jsandham Jul 27, 2026
bab2119
PR fixes
jsandham Jul 30, 2026
af239ab
Clang formatting
jsandham Jul 30, 2026
96c4c50
PR fixes
jsandham Jul 30, 2026
9b08ed8
Merge branch 'develop' of https://github.com/ROCm/rocm-libraries into…
jsandham Aug 4, 2026
3c00c8a
Merge branch 'develop' of https://github.com/ROCm/rocm-libraries into…
jsandham Aug 4, 2026
6938773
Merge branch 'develop' into spmm_batched_coo_seg2
jsandham Aug 4, 2026
fd14201
Merge branch 'develop' of https://github.com/ROCm/rocm-libraries into…
jsandham Aug 5, 2026
d2e1f98
Merge branch 'spmm_batched_coo_seg2' of https://github.com/jsandham/r…
jsandham Aug 5, 2026
4a25f36
Merge branch 'develop' of https://github.com/ROCm/rocm-libraries into…
jsandham Aug 5, 2026
e6aa85c
Edit changelog entry
jsandham Aug 5, 2026
4c31361
Merge branch 'develop' of https://github.com/ROCm/rocm-libraries into…
jsandham Aug 5, 2026
7b623bb
PR fixes
jsandham Aug 5, 2026
312bd65
Merge branch 'develop' into spmm_batched_coo_seg2
jsandham Aug 5, 2026
aa0371f
Merge branch 'develop' into spmm_batched_coo_seg2
jsandham Aug 6, 2026
6f08c2b
Merge branch 'develop' into spmm_batched_coo_seg2
jsandham Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions projects/rocsparse/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ Documentation for rocSPARSE is available at
* Added batched support to the SpMM algorithm `rocsparse_spmm_alg_csr_nnz_split`.

### 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.
Comment thread
jsandham marked this conversation as resolved.
Resolved
* Fixed `rocsparse_spmm` with the segmented atomic 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.
* Fixed `rocsparse_spmm` with the row-split CSR algorithm, which failed with `hipErrorInvalidConfiguration` for batch counts exceeding 65535 because the batch dimension of the kernel launch grid exceeded the maximum grid dimension.

Expand Down
2 changes: 1 addition & 1 deletion projects/rocsparse/clients/tests/test_spmm_batched_coo.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -511,7 +511,7 @@ Tests:
transB: [rocsparse_operation_none, rocsparse_operation_transpose]
baseA: [rocsparse_index_base_zero]
matrix: [rocsparse_matrix_file_rocalution]
spmm_alg: [rocsparse_spmm_alg_coo_segmented_atomic]
spmm_alg: [rocsparse_spmm_alg_coo_segmented, rocsparse_spmm_alg_coo_segmented_atomic]
orderB: [rocsparse_order_column]
orderC: [rocsparse_order_column]
filename: [nos4]
2 changes: 2 additions & 0 deletions projects/rocsparse/library/src/level3/coomm/segmented/kernel_declarations.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
102 changes: 57 additions & 45 deletions projects/rocsparse/library/src/level3/coomm/segmented/kernel_definitions.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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);
}
}
}

Expand All @@ -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,
Expand All @@ -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);
}
}
}
}
Expand All @@ -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, \
Expand Down Expand Up @@ -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, \
Expand Down
64 changes: 33 additions & 31 deletions projects/rocsparse/library/src/level3/coomm_device_segmented.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down Expand Up @@ -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];
Expand Down Expand Up @@ -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();
}
}
}
}
17 changes: 13 additions & 4 deletions projects/rocsparse/library/src/level3/rocsparse_coomm_template_segmented.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,12 @@

namespace rocsparse
{
template <typename J>
static uint16_t get_y_grid_size(J batch_count)
{
return (batch_count > 65535) ? 32768 : batch_count;
Comment thread
jsandham marked this conversation as resolved.
Resolved
}

template <typename T, typename I, typename A>
rocsparse_status coomm_buffer_size_template_segmented(rocsparse_handle handle,
rocsparse_operation trans_A,
Expand Down Expand Up @@ -62,7 +68,7 @@ namespace rocsparse
#define LAUNCH_COOMMNN_SEGMENTED_MAIN_KERNEL(COOMMNN_DIM, WF_SIZE, LOOPS, TRANSB) \
RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \
(rocsparse::coommnn_segmented_main_kernel<COOMMNN_DIM, WF_SIZE, LOOPS, TRANSB>), \
dim3(nblocks, (main - 1) / WF_SIZE + 1, batch_count_C), \
dim3(nblocks, (main - 1) / WF_SIZE + 1, get_y_grid_size<I>(batch_count_C)), \
dim3(COOMMNN_DIM), \
0, \
stream, \
Expand All @@ -72,6 +78,7 @@ namespace rocsparse
n, \
k, \
nnz, \
batch_count_C, \
batch_stride_A, \
ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), \
row_block_red, \
Expand All @@ -92,7 +99,7 @@ namespace rocsparse
#define LAUNCH_COOMMNN_SEGMENTED_REMAINDER_KERNEL(COOMMNN_DIM, WF_SIZE, LOOPS, TRANSB) \
RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \
(rocsparse::coommnn_segmented_remainder_kernel<COOMMNN_DIM, WF_SIZE, LOOPS, TRANSB>), \
dim3(nblocks, 1, batch_count_C), \
dim3(nblocks, 1, get_y_grid_size<I>(batch_count_C)), \
dim3(COOMMNN_DIM), \
0, \
stream, \
Expand All @@ -103,6 +110,7 @@ namespace rocsparse
n, \
k, \
nnz, \
batch_count_C, \
batch_stride_A, \
ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), \
row_block_red, \
Expand Down Expand Up @@ -295,7 +303,7 @@ namespace rocsparse
#undef LOOPS

RETURN_IF_HIPLAUNCHKERNELGGL_ERROR((rocsparse::coommnn_general_block_reduce<1024>),
dim3(n, 1, batch_count_C),
dim3(n, 1, get_y_grid_size<I>(batch_count_C)),
1024,
0,
stream,
Expand All @@ -306,7 +314,8 @@ namespace rocsparse
dense_C,
ldc,
batch_stride_C,
order_C);
order_C,
batch_count_C);
}
else
{
Expand Down
Loading