From e290f12aec5df41f88f5510052050b8b31d73cd2 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Wed, 19 Aug 2026 16:53:38 -0600 Subject: [PATCH 01/15] Fix unsigned 32-bit grid_size overflow in csrmv LRB The LRB csrmv variant computed its launch grid as a uint32_t product of a bin row count and a per-row workgroup count (nRowsBins[j] * num_wgs_per_row). This product wraps at 2^32, which both undersizes the kernel launch and, via max_required_grid, undersizes the wg_flags synchronization buffer. - Widen the grid products to int64_t before assignment in both the analysis buffer sizing and the long-rows launch; size max_required_grid (and thus wg_flags) from the unwrapped 64-bit value. - Clamp every LRB launch grid (short rows, short rows 2, medium warp-reduce, medium one-block, long rows) to handle->properties.maxGridSize[0]. The long rows cap is rounded down to a whole number of per-row workgroup groups so all cooperating workgroups of a row stay within one grid-stride wave, preserving the spin-loop hand-off. - Add grid-stride loops to the affected LRB device kernels so the clamped grid still covers the full logical work range; pass the total work count to the medium one-block and long-rows kernels. - Also size the adaptive partial_scale_y launch for the scale operation itself, clamp it to maxGridSize, and grid-stride its kernel. Contributes to AC3 of ROCM-28951 (AISPARSE-662). Co-authored-by: Cursor --- .../library/src/level2/csrmv_device.h | 627 ++++++++++-------- .../rocsparse_csrmv_template_adaptive.cpp | 8 +- .../level2/rocsparse_csrmv_template_lrb.cpp | 68 +- 3 files changed, 399 insertions(+), 304 deletions(-) diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index a1c90fe8b9ba..59b3c1c35905 100644 --- a/projects/rocsparse/library/src/level2/csrmv_device.h +++ b/projects/rocsparse/library/src/level2/csrmv_device.h @@ -178,34 +178,35 @@ namespace rocsparse template ROCSPARSE_DEVICE_ILF void partial_scale_y_device(J m, J first_row, J last_row, T beta, Y* y) { - const J gid = hipBlockIdx_x * WG_SIZE + hipThreadIdx_x; const J required_threads = (first_row - 0) + (m - last_row); - if(gid >= required_threads) + // The launch grid is clamped to the device's maximum grid size, so iterate with a + // grid-stride loop to cover every element that needs scaling. + const J stride = (J)WG_SIZE * hipGridDim_x; + for(J gid = (J)hipBlockIdx_x * WG_SIZE + hipThreadIdx_x; gid < required_threads; + gid += stride) { - return; - } - - if(gid < first_row) - { - if(beta == static_cast(0)) - { - y[gid] = static_cast(0); - } - else - { - y[gid] *= beta; - } - } - else - { - if(beta == static_cast(0)) + if(gid < first_row) { - y[last_row + (gid - first_row)] = static_cast(0); + if(beta == static_cast(0)) + { + y[gid] = static_cast(0); + } + else + { + y[gid] *= beta; + } } else { - y[last_row + (gid - first_row)] *= beta; + if(beta == static_cast(0)) + { + y[last_row + (gid - first_row)] = static_cast(0); + } + else + { + y[last_row + (gid - first_row)] *= beta; + } } } } @@ -752,7 +753,6 @@ namespace rocsparse rocsparse_index_base idx_base) { const uint32_t lid = hipThreadIdx_x; - const uint32_t gid = hipBlockIdx_x; // Allocation from the caller is of size [WG_SIZE << bin_id] elements extern __shared__ char shared_memory[]; @@ -761,71 +761,83 @@ namespace rocsparse const J bin_start = n_rows_bins[bin_id]; const J bin_num_rows = n_rows_bins[bin_id + 1] - bin_start; - const uint32_t wg_row_start = gid * BLOCKSIZE; - const uint32_t wg_row_end - = rocsparse::min(wg_row_start + BLOCKSIZE, uint32_t(bin_num_rows)); - const uint32_t wg_num_rows = wg_row_end - wg_row_start; + // Number of workgroups required to cover this bin. The launch grid is clamped to the + // device's maximum grid size, so iterate with a grid-stride loop to cover every block. + const J num_wgs = (bin_num_rows + BLOCKSIZE - 1) / BLOCKSIZE; - // Load a block of row data using all threads in the WG. - for(uint32_t base_idx = 0; base_idx < (BLOCKSIZE << bin_id); base_idx += BLOCKSIZE) + for(J gid = hipBlockIdx_x; gid < num_wgs; gid += (J)hipGridDim_x) { - uint32_t row_idx = wg_row_start + ((base_idx + lid) >> bin_id); - if(row_idx < wg_row_start + wg_num_rows) + const uint32_t wg_row_start = (uint32_t)(gid * BLOCKSIZE); + const uint32_t wg_row_end + = rocsparse::min(wg_row_start + BLOCKSIZE, uint32_t(bin_num_rows)); + const uint32_t wg_num_rows = wg_row_end - wg_row_start; + + // Load a block of row data using all threads in the WG. + for(uint32_t base_idx = 0; base_idx < (BLOCKSIZE << bin_id); base_idx += BLOCKSIZE) { - const J row_id = rows_bins[row_idx + bin_start]; + uint32_t row_idx = wg_row_start + ((base_idx + lid) >> bin_id); + if(row_idx < wg_row_start + wg_num_rows) + { + const J row_id = rows_bins[row_idx + bin_start]; - const I row_start = csr_row_ptr[row_id] - idx_base; - const I row_end = csr_row_ptr[row_id + 1] - idx_base; - const I row_len = row_end - row_start; + const I row_start = csr_row_ptr[row_id] - idx_base; + const I row_end = csr_row_ptr[row_id + 1] - idx_base; + const I row_len = row_end - row_start; - const uint32_t col_idx_in_row = lid & ((1 << bin_id) - 1); - const uint32_t lds_idx = base_idx + lid; + const uint32_t col_idx_in_row = lid & ((1 << bin_id) - 1); + const uint32_t lds_idx = base_idx + lid; - if(col_idx_in_row < row_len) - { - const A val = rocsparse::conj_val(csr_val[row_start + col_idx_in_row], conj); - partialSums[lds_idx] - = alpha * val * x[csr_col_ind[row_start + col_idx_in_row] - idx_base]; - } - else - { - // lds <- 0 - partialSums[lds_idx] = 0; + if(col_idx_in_row < row_len) + { + const A val + = rocsparse::conj_val(csr_val[row_start + col_idx_in_row], conj); + partialSums[lds_idx] + = alpha * val * x[csr_col_ind[row_start + col_idx_in_row] - idx_base]; + } + else + { + // lds <- 0 + partialSums[lds_idx] = 0; + } } } - } - __syncthreads(); - - // For the moment: just have each thread reduce a given row - // TODO: adaptation as per CSR-Adaptive-Stream - if(lid < wg_num_rows) - { - const uint32_t lds_start_idx = (lid << bin_id); - const J row_id = rows_bins[bin_start + wg_row_start + lid]; - T acc = 0; + __syncthreads(); - for(uint32_t idx = 0; idx < (1 << bin_id); idx++) + // For the moment: just have each thread reduce a given row + // TODO: adaptation as per CSR-Adaptive-Stream + if(lid < wg_num_rows) { - acc += partialSums[lds_start_idx + idx]; - } + const uint32_t lds_start_idx = (lid << bin_id); + const J row_id = rows_bins[bin_start + wg_row_start + lid]; + T acc = 0; - if(beta != 0) - { - acc = rocsparse::fma(beta, y[row_id], acc); - } + for(uint32_t idx = 0; idx < (1 << bin_id); idx++) + { + acc += partialSums[lds_start_idx + idx]; + } - // Add contributions from all extra vectors - for(rocsparse_int i = 0; i < num_extra; ++i) - { - T gamma = gamma_device_array[i]; - if(gamma != static_cast(0)) + if(beta != 0) { - acc = rocsparse::fma( - gamma, rocsparse::nontemporal_load(z_arrays[i] + row_id), acc); + acc = rocsparse::fma(beta, y[row_id], acc); } + + // Add contributions from all extra vectors + for(rocsparse_int i = 0; i < num_extra; ++i) + { + T gamma = gamma_device_array[i]; + if(gamma != static_cast(0)) + { + acc = rocsparse::fma( + gamma, rocsparse::nontemporal_load(z_arrays[i] + row_id), acc); + } + } + + y[row_id] = acc; } - y[row_id] = acc; + // Guard against a later grid-stride iteration overwriting partialSums while this + // iteration's reduction is still reading it. + __syncthreads(); } } @@ -859,7 +871,6 @@ namespace rocsparse rocsparse_index_base idx_base) { const uint32_t lid = hipThreadIdx_x; - const uint32_t gid = hipBlockIdx_x; // In LDS-Stream-V2, we have a fixed LDS allocation of CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS elements (example: 1024). // So, each thread will load 1024 / BLOCKSIZE elements, which, depending on the bin, corresponds to a different # rows. @@ -869,73 +880,87 @@ namespace rocsparse const J bin_start = n_rows_bins[bin_id]; const J bin_num_rows = n_rows_bins[bin_id + 1] - bin_start; const uint32_t rows_per_wg = CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS >> bin_id; - const uint32_t wg_row_start = gid * rows_per_wg; - const uint32_t wg_row_end - = rocsparse::min(wg_row_start + rows_per_wg, uint32_t(bin_num_rows)); - const uint32_t wg_num_rows = wg_row_end - wg_row_start; - - // Load a block of row data using all threads in the WG. - for(uint32_t base_idx = 0; base_idx < CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS; - base_idx += BLOCKSIZE) + + // Number of workgroups required to cover this bin. The launch grid is clamped to the + // device's maximum grid size, so iterate with a grid-stride loop to cover every block. + const J num_wgs = (bin_num_rows + rows_per_wg - 1) / rows_per_wg; + + for(J gid = hipBlockIdx_x; gid < num_wgs; gid += (J)hipGridDim_x) { - uint32_t row_idx = wg_row_start + ((base_idx + lid) >> bin_id); - if(row_idx < wg_row_end) + const uint32_t wg_row_start = (uint32_t)(gid * rows_per_wg); + const uint32_t wg_row_end + = rocsparse::min(wg_row_start + rows_per_wg, uint32_t(bin_num_rows)); + const uint32_t wg_num_rows = wg_row_end - wg_row_start; + + // Load a block of row data using all threads in the WG. + for(uint32_t base_idx = 0; base_idx < CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS; + base_idx += BLOCKSIZE) { - const J row_id = rows_bins[row_idx + bin_start]; + uint32_t row_idx = wg_row_start + ((base_idx + lid) >> bin_id); + if(row_idx < wg_row_end) + { + const J row_id = rows_bins[row_idx + bin_start]; - const I row_start = csr_row_ptr[row_id] - idx_base; - const I row_end = csr_row_ptr[row_id + 1] - idx_base; - const I row_len = row_end - row_start; + const I row_start = csr_row_ptr[row_id] - idx_base; + const I row_end = csr_row_ptr[row_id + 1] - idx_base; + const I row_len = row_end - row_start; - const uint32_t col_idx_in_row = lid & ((1 << bin_id) - 1); - const uint32_t lds_idx = base_idx + lid; - if(col_idx_in_row < row_len) - { - const A val = rocsparse::conj_val(csr_val[row_start + col_idx_in_row], conj); - partialSums[lds_idx] - = alpha * val * x[csr_col_ind[row_start + col_idx_in_row] - idx_base]; - } - else - { - // lds <- 0 - partialSums[lds_idx] = 0; + const uint32_t col_idx_in_row = lid & ((1 << bin_id) - 1); + const uint32_t lds_idx = base_idx + lid; + if(col_idx_in_row < row_len) + { + const A val + = rocsparse::conj_val(csr_val[row_start + col_idx_in_row], conj); + partialSums[lds_idx] + = alpha * val * x[csr_col_ind[row_start + col_idx_in_row] - idx_base]; + } + else + { + // lds <- 0 + partialSums[lds_idx] = 0; + } } } - } - __syncthreads(); + __syncthreads(); - // For the moment: just have each thread reduce a given row - // TODO: adaptation as per CSR-Adaptive-Stream - for(uint32_t row_base = 0; row_base < rows_per_wg; row_base += BLOCKSIZE) - { - const uint32_t this_row_offset_in_wg = row_base + lid; - if(this_row_offset_in_wg < wg_num_rows) + // For the moment: just have each thread reduce a given row + // TODO: adaptation as per CSR-Adaptive-Stream + for(uint32_t row_base = 0; row_base < rows_per_wg; row_base += BLOCKSIZE) { - const uint32_t lds_start_idx = (this_row_offset_in_wg << bin_id); - const J row_id = rows_bins[bin_start + wg_row_start + this_row_offset_in_wg]; - - T acc = 0; - for(uint32_t idx = 0; idx < (1 << bin_id); idx++) + const uint32_t this_row_offset_in_wg = row_base + lid; + if(this_row_offset_in_wg < wg_num_rows) { - acc += partialSums[lds_start_idx + idx]; - } + const uint32_t lds_start_idx = (this_row_offset_in_wg << bin_id); + const J row_id + = rows_bins[bin_start + wg_row_start + this_row_offset_in_wg]; - if(beta != 0) - acc = rocsparse::fma(beta, y[row_id], acc); + T acc = 0; + for(uint32_t idx = 0; idx < (1 << bin_id); idx++) + { + acc += partialSums[lds_start_idx + idx]; + } - // Add contributions from all extra vectors - for(rocsparse_int i = 0; i < num_extra; ++i) - { - T gamma = gamma_device_array[i]; - if(gamma != static_cast(0)) + if(beta != 0) + acc = rocsparse::fma(beta, y[row_id], acc); + + // Add contributions from all extra vectors + for(rocsparse_int i = 0; i < num_extra; ++i) { - acc = rocsparse::fma( - gamma, rocsparse::nontemporal_load(z_arrays[i] + row_id), acc); + T gamma = gamma_device_array[i]; + if(gamma != static_cast(0)) + { + acc = rocsparse::fma( + gamma, rocsparse::nontemporal_load(z_arrays[i] + row_id), acc); + } } - } - y[row_id] = acc; + y[row_id] = acc; + } } + + // Guard against a later grid-stride iteration overwriting partialSums while this + // iteration's reduction is still reading it. + __syncthreads(); } } @@ -979,49 +1004,50 @@ namespace rocsparse const int lid = tid & (WF_SIZE - 1); const int wid = tid / WF_SIZE; - const int gid = (BLOCKSIZE / WF_SIZE) * bid + wid; - - if(gid >= count) - { - return; - } - - const J bin_start = n_rows_bins[bin_id]; - const J row = rows_bins[bin_start + gid]; - - T temp_sum = static_cast(0); - const I vecStart = csr_row_ptr[row] - idx_base; - const I vecEnd = csr_row_ptr[row + 1] - idx_base; - - for(I j = vecStart + lid; j < vecEnd; j += WF_SIZE) + // One wavefront reduces one row. The launch grid is clamped to the device's maximum + // grid size, so iterate with a grid-stride loop (in units of wavefronts) to cover every + // row. This kernel uses no block-wide barriers, so divergent iteration counts are safe. + const int64_t warps_per_grid = (int64_t)(BLOCKSIZE / WF_SIZE) * hipGridDim_x; + for(int64_t gid = (int64_t)(BLOCKSIZE / WF_SIZE) * bid + wid; gid < count; + gid += warps_per_grid) { - temp_sum = rocsparse::fma(alpha * rocsparse::conj_val(csr_val[j], conj), - x[csr_col_ind[j] - idx_base], - temp_sum); - } + const J bin_start = n_rows_bins[bin_id]; + const J row = rows_bins[bin_start + gid]; - // Obtain row sum using parallel warp reduction - temp_sum = rocsparse::wfreduce_sum(temp_sum); + T temp_sum = static_cast(0); + const I vecStart = csr_row_ptr[row] - idx_base; + const I vecEnd = csr_row_ptr[row + 1] - idx_base; - if(lid == WF_SIZE - 1) - { - if(beta != static_cast(0)) + for(I j = vecStart + lid; j < vecEnd; j += WF_SIZE) { - temp_sum = rocsparse::fma(beta, y[row], temp_sum); + temp_sum = rocsparse::fma(alpha * rocsparse::conj_val(csr_val[j], conj), + x[csr_col_ind[j] - idx_base], + temp_sum); } - // Add contributions from all extra vectors - for(rocsparse_int i = 0; i < num_extra; ++i) + // Obtain row sum using parallel warp reduction + temp_sum = rocsparse::wfreduce_sum(temp_sum); + + if(lid == WF_SIZE - 1) { - T gamma = gamma_device_array[i]; - if(gamma != static_cast(0)) + if(beta != static_cast(0)) { - temp_sum = rocsparse::fma( - gamma, rocsparse::nontemporal_load(z_arrays[i] + row), temp_sum); + temp_sum = rocsparse::fma(beta, y[row], temp_sum); } - } - y[row] = temp_sum; + // Add contributions from all extra vectors + for(rocsparse_int i = 0; i < num_extra; ++i) + { + T gamma = gamma_device_array[i]; + if(gamma != static_cast(0)) + { + temp_sum = rocsparse::fma( + gamma, rocsparse::nontemporal_load(z_arrays[i] + row), temp_sum); + } + } + + y[row] = temp_sum; + } } } @@ -1036,6 +1062,7 @@ namespace rocsparse typename T> ROCSPARSE_DEVICE_ILF void csrmvn_lrb_medium_rows_device(bool conj, I nnz, + int64_t count, const J* rows_bins, const J* n_rows_bins, uint32_t bin_id, @@ -1052,15 +1079,11 @@ namespace rocsparse rocsparse_index_base idx_base) { const int lid = hipThreadIdx_x; - const int gid = hipBlockIdx_x; // LRB-Vector case currently does exact 1:1 WG:row mapping - not doing multiple rows per WG here; // we'll work on that later if desired. But, that may be neither needed nor wanted, // since we're now launching with LDS and grid size parameters tuned according to row length. - const J bin_start = n_rows_bins[bin_id]; - const J row = rows_bins[bin_start + gid]; - // In CSR-Adaptive-Vector, each WG will allocate BLOCKSIZE LDS, but for rows shorter than BLOCKSIZE, // this wastes space. So, for LRB-Vector, we allocate from the caller according to bin_id (one LDS // element per WG thread). @@ -1068,54 +1091,68 @@ namespace rocsparse // size from the caller; the only requirement is that we have LDS #elems == WG size (max=1024 on gfx90a). __shared__ T partialSums[BLOCKSIZE]; - // In its original form, any CSR-Vector workgroup only calculates, at most, BLOCKSIZE items in its row; - // if there are more items in this row, CSR-LongRows is used instead. - // In its LRB-ized form, we'll calculate as much as we need, with parallelism bounded only by grid size - // (iterating over the row as needed for rows with length > WG size). - // This means that we can more easily process everything with Vector that we would otherwise have done - // with Longrows - which, in turn, means we can guarantee result reproducibility simply by avoiding Longrows - // use (-> no non-deterministic atomics), just set the bin threshold for Longrows to "infinity" (or "32"). - T temp_sum = static_cast(0); - const I vecStart = csr_row_ptr[row] - idx_base; - const I vecEnd = csr_row_ptr[row + 1] - idx_base; - - // Load in a bunch of partial results into your register space, rather than LDS (no contention) - // Then dump the partially reduced answers into the LDS for inter-work-item reduction. - for(I j = vecStart + lid; j < vecEnd; j += BLOCKSIZE) - { - temp_sum = rocsparse::fma(alpha * rocsparse::conj_val(csr_val[j], conj), - x[csr_col_ind[j] - idx_base], - temp_sum); - } + const J bin_start = n_rows_bins[bin_id]; - partialSums[lid] = temp_sum; + // One block reduces one row. The launch grid is clamped to the device's maximum grid + // size, so iterate with a grid-stride loop to cover every row. All threads in a block + // share the same gid, so the block-wide barriers below stay uniform. + for(int64_t gid = hipBlockIdx_x; gid < count; gid += (int64_t)hipGridDim_x) + { + const J row = rows_bins[bin_start + gid]; + + // In its original form, any CSR-Vector workgroup only calculates, at most, BLOCKSIZE items in its row; + // if there are more items in this row, CSR-LongRows is used instead. + // In its LRB-ized form, we'll calculate as much as we need, with parallelism bounded only by grid size + // (iterating over the row as needed for rows with length > WG size). + // This means that we can more easily process everything with Vector that we would otherwise have done + // with Longrows - which, in turn, means we can guarantee result reproducibility simply by avoiding Longrows + // use (-> no non-deterministic atomics), just set the bin threshold for Longrows to "infinity" (or "32"). + T temp_sum = static_cast(0); + const I vecStart = csr_row_ptr[row] - idx_base; + const I vecEnd = csr_row_ptr[row + 1] - idx_base; + + // Load in a bunch of partial results into your register space, rather than LDS (no contention) + // Then dump the partially reduced answers into the LDS for inter-work-item reduction. + for(I j = vecStart + lid; j < vecEnd; j += BLOCKSIZE) + { + temp_sum = rocsparse::fma(alpha * rocsparse::conj_val(csr_val[j], conj), + x[csr_col_ind[j] - idx_base], + temp_sum); + } - __syncthreads(); + partialSums[lid] = temp_sum; - // Reduce partial sums - rocsparse::blockreduce_sum(lid, partialSums); + __syncthreads(); - if(lid == 0) - { - temp_sum = partialSums[0]; + // Reduce partial sums + rocsparse::blockreduce_sum(lid, partialSums); - if(beta != static_cast(0)) + if(lid == 0) { - temp_sum = rocsparse::fma(beta, y[row], temp_sum); - } + temp_sum = partialSums[0]; - // Add contributions from all extra vectors - for(rocsparse_int i = 0; i < num_extra; ++i) - { - T gamma = gamma_device_array[i]; - if(gamma != static_cast(0)) + if(beta != static_cast(0)) { - temp_sum = rocsparse::fma( - gamma, rocsparse::nontemporal_load(z_arrays[i] + row), temp_sum); + temp_sum = rocsparse::fma(beta, y[row], temp_sum); } + + // Add contributions from all extra vectors + for(rocsparse_int i = 0; i < num_extra; ++i) + { + T gamma = gamma_device_array[i]; + if(gamma != static_cast(0)) + { + temp_sum = rocsparse::fma( + gamma, rocsparse::nontemporal_load(z_arrays[i] + row), temp_sum); + } + } + + y[row] = temp_sum; } - y[row] = temp_sum; + // Guard against a later grid-stride iteration overwriting partialSums while this + // iteration's reduction is still reading it. + __syncthreads(); } } @@ -1132,6 +1169,7 @@ namespace rocsparse ROCSPARSE_DEVICE_ILF void csrmvn_lrb_long_rows_device(bool conj, J m, I nnz, + int64_t count, uint32_t* wg_flags, const J* rows_bins, const J* n_rows_bins, @@ -1151,125 +1189,138 @@ namespace rocsparse __shared__ T partialSums[BLOCKSIZE]; const int lid = hipThreadIdx_x; - const int gid = hipBlockIdx_x; const J bin_start = n_rows_bins[bin_id]; const uint32_t bin_max_row_len = (1 << bin_id); const uint32_t num_wgs_per_row = (bin_max_row_len - 1) / (BLOCK_MULTIPLIER * BLOCKSIZE) + 1; - const J row_idx = gid / num_wgs_per_row; - const J wg = gid % num_wgs_per_row; - const J row = rows_bins[bin_start + row_idx]; - // wg_flags contains the flag bits used so that the multiple WGs calculating a long row can - // know when the first workgroup for that row has finished initializing the output - // value. While this bit is the same as the first workgroup's flag bit, this - // workgroup will spin-loop. + // count is the total number of cooperating workgroups (nRowsBins[bin_id] * num_wgs_per_row). + // The launch grid is clamped to the device's maximum grid size and rounded down to a whole + // number of per-row workgroup groups, so iterate with a grid-stride loop. Because the grid + // is a multiple of num_wgs_per_row, every workgroup cooperating on a given row falls within + // the same grid-stride wave, preserving the cross-workgroup spin-loop synchronization + // contract (all workgroups of a row must run concurrently for the flag hand-off to work). + for(int64_t gid = hipBlockIdx_x; gid < count; gid += (int64_t)hipGridDim_x) + { + const J row_idx = gid / num_wgs_per_row; + const J wg = gid % num_wgs_per_row; + const J row = rows_bins[bin_start + row_idx]; - // TODO: Can the wg_flags-based coordination be done instead with cooperative-groups? + // wg_flags contains the flag bits used so that the multiple WGs calculating a long row can + // know when the first workgroup for that row has finished initializing the output + // value. While this bit is the same as the first workgroup's flag bit, this + // workgroup will spin-loop. - // Each workgroup computes exactly one row. - const I vecStart = (I)wg * (I)BLOCK_MULTIPLIER * BLOCKSIZE + csr_row_ptr[row] - idx_base; - const I vecEnd = rocsparse::min(csr_row_ptr[row + 1] - idx_base, - vecStart + I(BLOCK_MULTIPLIER * BLOCKSIZE)); + // TODO: Can the wg_flags-based coordination be done instead with cooperative-groups? - T temp_sum = static_cast(0); + // Each workgroup computes exactly one row. + const I vecStart = (I)wg * (I)BLOCK_MULTIPLIER * BLOCKSIZE + csr_row_ptr[row] - idx_base; + const I vecEnd = rocsparse::min(csr_row_ptr[row + 1] - idx_base, + vecStart + I(BLOCK_MULTIPLIER * BLOCKSIZE)); - // In CSR-LongRows, we have more than one workgroup calculating this row. - // The output values for those types of rows are stored using atomic_add, because - // more than one parallel workgroup's value makes up the final answer. - // Unfortunately, this makes it difficult to do y=Ax, rather than y=Ax+y, because - // the values still left in y will be added in using the atomic_add. - // - // Our solution is to have the first workgroup in one of these long-rows cases - // properly initialize the output vector. All the other workgroups working on this - // row will spin-loop until that workgroup finishes its work. - - // First, figure out which workgroup you are in the row. - // You can use that to find the global ID for the first workgroup calculating - // this long row. - const J first_wg_in_row = gid - wg; - const uint32_t compare_value = wg_flags[gid]; - - // wg_flags[first_wg_in_row] in the first workgroup is the flag that everyone waits on. - if(gid == first_wg_in_row && lid == 0) - { - // The first workgroup handles the output initialization. - const Y out_val = y[row]; - temp_sum = (beta - static_cast(1)) * out_val; - - // All inter thread communication is done using atomics, therefore cache flushes or - // invalidates should not be needed (thus __threadfence() has been removed to regain - // performance). - // Because of atomics being relaxed, however, the compiler is allowed to reorder them - // with respect to ordinary memory accesses (and other relaxed atomic operations). - // In this case, out_val seem to be reordered with the xor and subsequently, accumulation - // ends up being wrong. - // To force the compiler to stick to the order of operations, we need acquire/release fences. - // Workgroup scope is sufficient for this purpose, to only invalidate L1 and avoid L2 - // invalidations. + T temp_sum = static_cast(0); + + // In CSR-LongRows, we have more than one workgroup calculating this row. + // The output values for those types of rows are stored using atomic_add, because + // more than one parallel workgroup's value makes up the final answer. + // Unfortunately, this makes it difficult to do y=Ax, rather than y=Ax+y, because + // the values still left in y will be added in using the atomic_add. + // + // Our solution is to have the first workgroup in one of these long-rows cases + // properly initialize the output vector. All the other workgroups working on this + // row will spin-loop until that workgroup finishes its work. + + // First, figure out which workgroup you are in the row. + // You can use that to find the global ID for the first workgroup calculating + // this long row. + const J first_wg_in_row = gid - wg; + const uint32_t compare_value = wg_flags[gid]; + + // wg_flags[first_wg_in_row] in the first workgroup is the flag that everyone waits on. + if(gid == first_wg_in_row && lid == 0) + { + // The first workgroup handles the output initialization. + const Y out_val = y[row]; + temp_sum = (beta - static_cast(1)) * out_val; + + // All inter thread communication is done using atomics, therefore cache flushes or + // invalidates should not be needed (thus __threadfence() has been removed to regain + // performance). + // Because of atomics being relaxed, however, the compiler is allowed to reorder them + // with respect to ordinary memory accesses (and other relaxed atomic operations). + // In this case, out_val seem to be reordered with the xor and subsequently, accumulation + // ends up being wrong. + // To force the compiler to stick to the order of operations, we need acquire/release fences. + // Workgroup scope is sufficient for this purpose, to only invalidate L1 and avoid L2 + // invalidations. #if defined(__GFX12__) - __threadfence(); + __threadfence(); #else - __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup"); - __builtin_amdgcn_s_waitcnt(0); + __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup"); + __builtin_amdgcn_s_waitcnt(0); #endif - // Release other workgroups - atomicXor(&wg_flags[first_wg_in_row], 1U); - } + // Release other workgroups + atomicXor(&wg_flags[first_wg_in_row], 1U); + } - // Load in a bunch of partial results into your register space, rather than LDS (no - // contention) - // Then dump the partially reduced answers into the LDS for inter-work-item reduction. - for(I j = vecStart + lid; j < vecEnd; j += BLOCKSIZE) - { - temp_sum = rocsparse::fma(alpha * rocsparse::conj_val(csr_val[j], conj), - x[csr_col_ind[j] - idx_base], - temp_sum); - } + // Load in a bunch of partial results into your register space, rather than LDS (no + // contention) + // Then dump the partially reduced answers into the LDS for inter-work-item reduction. + for(I j = vecStart + lid; j < vecEnd; j += BLOCKSIZE) + { + temp_sum = rocsparse::fma(alpha * rocsparse::conj_val(csr_val[j], conj), + x[csr_col_ind[j] - idx_base], + temp_sum); + } - partialSums[lid] = temp_sum; + partialSums[lid] = temp_sum; - __syncthreads(); + __syncthreads(); - // Reduce partial sums - rocsparse::blockreduce_sum(lid, partialSums); + // Reduce partial sums + rocsparse::blockreduce_sum(lid, partialSums); - // For every other workgroup, wg_flags[first_wg_in_row] holds the value they wait on. - // If your flag == first_wg's flag, you spin loop. - // The first workgroup will eventually flip this flag, and you can move forward. - if(lid == 0) - { - if(gid != first_wg_in_row) + // For every other workgroup, wg_flags[first_wg_in_row] holds the value they wait on. + // If your flag == first_wg's flag, you spin loop. + // The first workgroup will eventually flip this flag, and you can move forward. + if(lid == 0) { - while(rocsparse::atomic_max(&wg_flags[first_wg_in_row], 0U) == compare_value) - ; + if(gid != first_wg_in_row) + { + while(rocsparse::atomic_max(&wg_flags[first_wg_in_row], 0U) == compare_value) + ; - // __builtin_amdgcn_s_waitcnt(0); - // __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup"); + // __builtin_amdgcn_s_waitcnt(0); + // __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup"); - // After you've passed the barrier, update your local flag to make sure that - // the next time through, you know what to wait on. - wg_flags[gid] ^= 1U; - } + // After you've passed the barrier, update your local flag to make sure that + // the next time through, you know what to wait on. + wg_flags[gid] ^= 1U; + } - T extra_sum = partialSums[0]; + T extra_sum = partialSums[0]; - // Add contributions from all extra vectors (only for the first workgroup) - if(gid == first_wg_in_row) - { - for(rocsparse_int i = 0; i < num_extra; ++i) + // Add contributions from all extra vectors (only for the first workgroup) + if(gid == first_wg_in_row) { - T gamma = gamma_device_array[i]; - if(gamma != static_cast(0)) + for(rocsparse_int i = 0; i < num_extra; ++i) { - extra_sum = rocsparse::fma( - gamma, rocsparse::nontemporal_load(z_arrays[i] + row), extra_sum); + T gamma = gamma_device_array[i]; + if(gamma != static_cast(0)) + { + extra_sum = rocsparse::fma( + gamma, rocsparse::nontemporal_load(z_arrays[i] + row), extra_sum); + } } } + + rocsparse::atomic_add(y, row, m, extra_sum); } - rocsparse::atomic_add(y, row, m, extra_sum); + // Guard against a later grid-stride iteration overwriting partialSums while this + // iteration's reduction is still reading it. + __syncthreads(); } } } diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp index d309b368920e..66d74ad654b5 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp @@ -767,7 +767,13 @@ rocsparse_status rocsparse::csrmv_adaptive_template_dispatch(rocsparse_handle J last_row = info->adaptive.last_row; J required_threads = (first_row - 0) + (m - last_row); - dim3 scale_blocks((required_threads - 1) / gen_wg + 1); + // Size the grid for the scale operation itself (previously this reused the adaptive + // kernel's grid). Clamp to the device maximum grid size; the kernel grid-strides over + // all required threads. + int64_t scale_wgs + = rocsparse::min(int64_t((required_threads - 1) / gen_wg + 1), + int64_t(handle->properties.maxGridSize[0])); + dim3 scale_blocks(static_cast(scale_wgs)); dim3 scale_threads(gen_wg); #define ROCSPARSE_LAUNCH_PARTIAL_SCALE_Y(GEN_WG) \ RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp index d5ecf87dcc80..33901158d2ae 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp @@ -200,13 +200,17 @@ rocsparse_status rocsparse::csrmv_analysis_lrb_template_dispatch(rocsparse_handl // Longrows synchronization flags: these will be allocated below, during the preprocessing, based // on Longrows bin sizes. We use a different array per kernel launch in order to permit use of >1 stream. - uint32_t max_required_grid = 0; + // The number of cooperating workgroups for a long-row bin is nRowsBins[j] * num_wgs_per_row. + // This product can exceed the range of a 32-bit unsigned integer, so accumulate it in 64-bit + // to avoid wrap-around. This buffer must be sized for the *true* (unclamped) grid because the + // kernel's grid-stride loop indexes wg_flags by the full logical workgroup id. + int64_t max_required_grid = 0; for(int j = LR_THRESHOLD; j < 32; j++) { uint32_t block_size = WG_SIZE; uint32_t bin_max_row_len = (1 << j); uint32_t num_wgs_per_row = (bin_max_row_len - 1) / (BLOCK_MULTIPLIER * block_size) + 1; - uint32_t grid_size = csrmv_info->lrb.nRowsBins[j] * num_wgs_per_row; + int64_t grid_size = int64_t(csrmv_info->lrb.nRowsBins[j]) * num_wgs_per_row; max_required_grid = rocsparse::max(grid_size, max_required_grid); } @@ -394,8 +398,9 @@ namespace rocsparse typename Z, typename T> ROCSPARSE_KERNEL(BLOCKSIZE) - void csrmvn_lrb_medium_rows_kernel(bool conj, - I nnz, + void csrmvn_lrb_medium_rows_kernel(bool conj, + I nnz, + int64_t count, J* __restrict__ rows_bins, J* __restrict__ n_rows_bins, const uint32_t bin_id, @@ -418,6 +423,7 @@ namespace rocsparse { rocsparse::csrmvn_lrb_medium_rows_device(conj, nnz, + count, rows_bins, n_rows_bins, bin_id, @@ -437,9 +443,10 @@ namespace rocsparse template ROCSPARSE_KERNEL(WG_SIZE) - void csrmvn_lrb_long_rows_kernel(bool conj, - J m, - I nnz, + void csrmvn_lrb_long_rows_kernel(bool conj, + J m, + I nnz, + int64_t count, uint32_t* __restrict__ wg_flags, J* __restrict__ rows_bins, J* __restrict__ n_rows_bins, @@ -465,6 +472,7 @@ namespace rocsparse conj, m, nnz, + count, wg_flags, rows_bins, n_rows_bins, @@ -569,8 +577,10 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // Dynamic LDS allocation if(lds_size < CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS * sizeof(T)) { - uint32_t grid_size - = rocsparse::ceil((float)info->lrb.nRowsBins[j] / block_size); + int64_t num_wgs + = (int64_t)rocsparse::ceil((float)info->lrb.nRowsBins[j] / block_size); + uint32_t grid_size = static_cast( + rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( (csrmvn_lrb_short_rows_kernel), @@ -600,8 +610,10 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle else { uint32_t rows_per_wg = CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS >> j; - uint32_t grid_size - = rocsparse::ceil((float)info->lrb.nRowsBins[j] / rows_per_wg); + int64_t num_wgs + = (int64_t)rocsparse::ceil((float)info->lrb.nRowsBins[j] / rows_per_wg); + uint32_t grid_size = static_cast( + rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( (csrmvn_lrb_short_rows_2_kernel), @@ -646,7 +658,9 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // RDNA (wave32): use 128-thread blocks (4 wavefronts) instead of 256. // One wavefront reduces one row, so grid is sized by wavefronts-per-block. constexpr uint32_t mr_block = WG_SIZE_WAVE32; - uint32_t grid_size = (info->lrb.nRowsBins[j] - 1) / (mr_block / 32) + 1; + int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (mr_block / 32) + 1; + uint32_t grid_size = static_cast( + rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( (csrmvn_lrb_medium_rows_warp_reduce_kernel), grid_size, @@ -674,8 +688,10 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle } else { - uint32_t grid_size + int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (256 / handle->wavefront_size) + 1; + uint32_t grid_size = static_cast( + rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( (csrmvn_lrb_medium_rows_warp_reduce_kernel<256, 64>), grid_size, @@ -704,7 +720,9 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle } else // One block per row { - uint32_t grid_size = info->lrb.nRowsBins[j]; // One WG per row + int64_t count = info->lrb.nRowsBins[j]; // One WG per row + uint32_t grid_size = static_cast( + rocsparse::min(count, int64_t(handle->properties.maxGridSize[0]))); if(handle->wavefront_size == 32) { @@ -718,6 +736,7 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle stream, conj, nnz, + count, static_cast(info->lrb.rows_bins), static_cast(info->lrb.n_rows_bins), j, @@ -744,6 +763,7 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle stream, conj, nnz, + count, static_cast(info->lrb.rows_bins), static_cast(info->lrb.n_rows_bins), j, @@ -776,7 +796,24 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle uint32_t bin_max_row_len = (1 << j); uint32_t num_wgs_per_row = (bin_max_row_len - 1) / (BLOCK_MULTIPLIER * block_size) + 1; - uint32_t grid_size = info->lrb.nRowsBins[j] * num_wgs_per_row; + + // Total cooperating workgroups for this bin. Widen to 64-bit before multiplying to + // avoid the 32-bit product wrapping (which would both undersize the launch and, via + // wg_flags sizing, corrupt the cross-workgroup synchronization). + int64_t count = int64_t(info->lrb.nRowsBins[j]) * num_wgs_per_row; + + // Clamp the launch grid to the device maximum, rounded down to a whole number of + // per-row workgroup groups so that all workgroups cooperating on a row stay within + // one grid-stride wave (required by the spin-loop hand-off). The kernel grid-strides + // over the full logical range [0, count). + int64_t max_grid = int64_t(handle->properties.maxGridSize[0]); + int64_t capped = rocsparse::min(count, max_grid); + capped = (capped / num_wgs_per_row) * num_wgs_per_row; + if(capped < num_wgs_per_row) + { + capped = num_wgs_per_row; + } + uint32_t grid_size = static_cast(rocsparse::min(capped, count)); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( (csrmvn_lrb_long_rows_kernel), @@ -787,6 +824,7 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle conj, m, nnz, + count, info->lrb.wg_flags, static_cast(info->lrb.rows_bins), static_cast(info->lrb.n_rows_bins), From f543f2deba4a682692e1ed5fe256c6e892e94569 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Fri, 21 Aug 2026 14:39:31 -0600 Subject: [PATCH 02/15] style(rocsparse): apply clang-format-14 to large-grid changes Co-authored-by: Cursor --- projects/rocsparse/library/src/level2/csrmv_device.h | 8 ++++---- .../src/level2/rocsparse_csrmv_template_adaptive.cpp | 9 ++++----- .../library/src/level2/rocsparse_csrmv_template_lrb.cpp | 2 +- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index 59b3c1c35905..5f92f059b0b3 100644 --- a/projects/rocsparse/library/src/level2/csrmv_device.h +++ b/projects/rocsparse/library/src/level2/csrmv_device.h @@ -931,8 +931,7 @@ namespace rocsparse if(this_row_offset_in_wg < wg_num_rows) { const uint32_t lds_start_idx = (this_row_offset_in_wg << bin_id); - const J row_id - = rows_bins[bin_start + wg_row_start + this_row_offset_in_wg]; + const J row_id = rows_bins[bin_start + wg_row_start + this_row_offset_in_wg]; T acc = 0; for(uint32_t idx = 0; idx < (1 << bin_id); idx++) @@ -1214,8 +1213,9 @@ namespace rocsparse // TODO: Can the wg_flags-based coordination be done instead with cooperative-groups? // Each workgroup computes exactly one row. - const I vecStart = (I)wg * (I)BLOCK_MULTIPLIER * BLOCKSIZE + csr_row_ptr[row] - idx_base; - const I vecEnd = rocsparse::min(csr_row_ptr[row + 1] - idx_base, + const I vecStart + = (I)wg * (I)BLOCK_MULTIPLIER * BLOCKSIZE + csr_row_ptr[row] - idx_base; + const I vecEnd = rocsparse::min(csr_row_ptr[row + 1] - idx_base, vecStart + I(BLOCK_MULTIPLIER * BLOCKSIZE)); T temp_sum = static_cast(0); diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp index 66d74ad654b5..bc33dcb2cd0a 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp @@ -770,11 +770,10 @@ rocsparse_status rocsparse::csrmv_adaptive_template_dispatch(rocsparse_handle // Size the grid for the scale operation itself (previously this reused the adaptive // kernel's grid). Clamp to the device maximum grid size; the kernel grid-strides over // all required threads. - int64_t scale_wgs - = rocsparse::min(int64_t((required_threads - 1) / gen_wg + 1), - int64_t(handle->properties.maxGridSize[0])); - dim3 scale_blocks(static_cast(scale_wgs)); - dim3 scale_threads(gen_wg); + int64_t scale_wgs = rocsparse::min(int64_t((required_threads - 1) / gen_wg + 1), + int64_t(handle->properties.maxGridSize[0])); + dim3 scale_blocks(static_cast(scale_wgs)); + dim3 scale_threads(gen_wg); #define ROCSPARSE_LAUNCH_PARTIAL_SCALE_Y(GEN_WG) \ RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ (rocsparse::partial_scale_y_kernel), \ diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp index 33901158d2ae..3a50e79ec1aa 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp @@ -658,7 +658,7 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // RDNA (wave32): use 128-thread blocks (4 wavefronts) instead of 256. // One wavefront reduces one row, so grid is sized by wavefronts-per-block. constexpr uint32_t mr_block = WG_SIZE_WAVE32; - int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (mr_block / 32) + 1; + int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (mr_block / 32) + 1; uint32_t grid_size = static_cast( rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( From d77abc7b10da58d7bd2aa6631d1c53d1042540c8 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Wed, 2 Sep 2026 13:07:37 -0600 Subject: [PATCH 03/15] refactor(rocsparse): use overflow-safe integer ceiling for csrmv LRB short-rows grids (AISPARSE-662) The two LRB short-rows launches computed their workgroup count with a float-based ceiling (rocsparse::ceil((float)nRowsBins / block)), which loses precision for bins with more than 2^24 rows. Switch to the overflow-safe integer idiom (count - 1)/block + 1 to match the rest of the LRB overflow fix and the sibling-PR review conventions. Review notes for AISPARSE-662: - All ticket sites (analysis buffer sizing, long-rows launch, short-rows, medium warp-reduce, medium one-wg-per-row, and the adaptive scale_y launch) widen the workgroup product to int64 before multiplying, clamp the launch grid via rocsparse::min against maxGridSize[0], and the kernels grid-stride over the full logical range. - max_required_grid (wg_flags sizing) is accumulated in int64 from the unclamped product, so the sync buffer is never undersized. - The long-rows launch rounds the clamped grid down to a whole number of per-row workgroup groups, keeping every cooperating workgroup of a row in the same grid-stride wave and preserving the spin-loop hand-off. - The csrmv_device.h diff is large but is entirely grid-stride wrapping (bodies re-indented into for-loops); ignoring whitespace it is only ~74/23 lines and touches no untouched code, so there is no pure-format churn to restore. Co-authored-by: Cursor --- .../src/level2/rocsparse_csrmv_template_lrb.cpp | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp index 3a50e79ec1aa..00de56777ad3 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp @@ -577,8 +577,9 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // Dynamic LDS allocation if(lds_size < CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS * sizeof(T)) { - int64_t num_wgs - = (int64_t)rocsparse::ceil((float)info->lrb.nRowsBins[j] / block_size); + // Overflow-safe integer ceiling division (float ceil loses precision for + // row counts above 2^24). nRowsBins[j] is guaranteed non-zero here. + int64_t num_wgs = (int64_t(info->lrb.nRowsBins[j]) - 1) / block_size + 1; uint32_t grid_size = static_cast( rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); @@ -610,8 +611,9 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle else { uint32_t rows_per_wg = CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS >> j; - int64_t num_wgs - = (int64_t)rocsparse::ceil((float)info->lrb.nRowsBins[j] / rows_per_wg); + // Overflow-safe integer ceiling division (float ceil loses precision for + // row counts above 2^24). nRowsBins[j] is guaranteed non-zero here. + int64_t num_wgs = (int64_t(info->lrb.nRowsBins[j]) - 1) / rows_per_wg + 1; uint32_t grid_size = static_cast( rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); From 967dc02d95ac970b263ba13487babd041653537d Mon Sep 17 00:00:00 2001 From: kliegeois Date: Tue, 29 Sep 2026 13:27:48 +0000 Subject: [PATCH 04/15] refactor(rocsparse): use rocsparse_grid.hpp helpers (AISPARSE-662) --- .../rocsparse_csrmv_template_adaptive.cpp | 11 +++---- .../level2/rocsparse_csrmv_template_lrb.cpp | 31 +++++++++---------- 2 files changed, 19 insertions(+), 23 deletions(-) diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp index bc33dcb2cd0a..1f0ad0f6692d 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp @@ -27,6 +27,7 @@ #include "rocsparse_control.hpp" #include "rocsparse_csrmv.hpp" #include "rocsparse_csrmv_adaptive_analysis.hpp" +#include "rocsparse_grid.hpp" #include "rocsparse_utility.hpp" #include "internal/generic/rocsparse_v2_spmv.h" @@ -768,12 +769,10 @@ rocsparse_status rocsparse::csrmv_adaptive_template_dispatch(rocsparse_handle J required_threads = (first_row - 0) + (m - last_row); // Size the grid for the scale operation itself (previously this reused the adaptive - // kernel's grid). Clamp to the device maximum grid size; the kernel grid-strides over - // all required threads. - int64_t scale_wgs = rocsparse::min(int64_t((required_threads - 1) / gen_wg + 1), - int64_t(handle->properties.maxGridSize[0])); - dim3 scale_blocks(static_cast(scale_wgs)); - dim3 scale_threads(gen_wg); + // kernel's grid). Clamp the grid; the kernel grid-strides over all required threads. + dim3 scale_blocks(rocsparse::get_grid_size_x( + handle, (int64_t(required_threads) - 1) / gen_wg + 1, gen_wg)); + dim3 scale_threads(gen_wg); #define ROCSPARSE_LAUNCH_PARTIAL_SCALE_Y(GEN_WG) \ RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ (rocsparse::partial_scale_y_kernel), \ diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp index 00de56777ad3..b6b415cb4d05 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp @@ -25,6 +25,7 @@ #include "rocsparse_common.hpp" #include "rocsparse_control.hpp" #include "rocsparse_csrmv.hpp" +#include "rocsparse_grid.hpp" #include "rocsparse_utility.hpp" #include "internal/generic/rocsparse_v2_spmv.h" @@ -580,8 +581,7 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // Overflow-safe integer ceiling division (float ceil loses precision for // row counts above 2^24). nRowsBins[j] is guaranteed non-zero here. int64_t num_wgs = (int64_t(info->lrb.nRowsBins[j]) - 1) / block_size + 1; - uint32_t grid_size = static_cast( - rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); + uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, block_size); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( (csrmvn_lrb_short_rows_kernel), @@ -614,8 +614,7 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // Overflow-safe integer ceiling division (float ceil loses precision for // row counts above 2^24). nRowsBins[j] is guaranteed non-zero here. int64_t num_wgs = (int64_t(info->lrb.nRowsBins[j]) - 1) / rows_per_wg + 1; - uint32_t grid_size = static_cast( - rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); + uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, block_size); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( (csrmvn_lrb_short_rows_2_kernel), @@ -661,8 +660,7 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // One wavefront reduces one row, so grid is sized by wavefronts-per-block. constexpr uint32_t mr_block = WG_SIZE_WAVE32; int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (mr_block / 32) + 1; - uint32_t grid_size = static_cast( - rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); + uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, mr_block); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( (csrmvn_lrb_medium_rows_warp_reduce_kernel), grid_size, @@ -692,8 +690,7 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle { int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (256 / handle->wavefront_size) + 1; - uint32_t grid_size = static_cast( - rocsparse::min(num_wgs, int64_t(handle->properties.maxGridSize[0]))); + uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, 256); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( (csrmvn_lrb_medium_rows_warp_reduce_kernel<256, 64>), grid_size, @@ -722,9 +719,10 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle } else // One block per row { - int64_t count = info->lrb.nRowsBins[j]; // One WG per row - uint32_t grid_size = static_cast( - rocsparse::min(count, int64_t(handle->properties.maxGridSize[0]))); + int64_t count = info->lrb.nRowsBins[j]; // One WG per row + const uint32_t mr_block + = (handle->wavefront_size == 32) ? WG_SIZE_WAVE32 : WG_SIZE; + uint32_t grid_size = rocsparse::get_grid_size_x(handle, count, mr_block); if(handle->wavefront_size == 32) { @@ -804,12 +802,11 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // wg_flags sizing, corrupt the cross-workgroup synchronization). int64_t count = int64_t(info->lrb.nRowsBins[j]) * num_wgs_per_row; - // Clamp the launch grid to the device maximum, rounded down to a whole number of - // per-row workgroup groups so that all workgroups cooperating on a row stay within - // one grid-stride wave (required by the spin-loop hand-off). The kernel grid-strides - // over the full logical range [0, count). - int64_t max_grid = int64_t(handle->properties.maxGridSize[0]); - int64_t capped = rocsparse::min(count, max_grid); + // Clamp the launch grid, rounded down to a whole number of per-row workgroup + // groups so that all workgroups cooperating on a row stay within one grid-stride + // wave (required by the spin-loop hand-off). The kernel grid-strides over the full + // logical range [0, count). + int64_t capped = rocsparse::get_grid_size_x(handle, count, block_size); capped = (capped / num_wgs_per_row) * num_wgs_per_row; if(capped < num_wgs_per_row) { From eee390706dd43e6ed91e8cad4bed2ee9b8123c60 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Tue, 29 Sep 2026 17:32:54 +0000 Subject: [PATCH 05/15] fix(rocsparse): 64-bit grid-stride index in csrmv partial_scale_y (AISPARSE-662) partial_scale_y_device strided gid in J. Once the number of rows to scale exceeds 2^30, gid += stride overflows int32, wraps negative, takes the gid < first_row branch and writes y out of bounds. Use int64_t for gid and stride. --- projects/rocsparse/library/src/level2/csrmv_device.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index 5f92f059b0b3..6c54c487b1c6 100644 --- a/projects/rocsparse/library/src/level2/csrmv_device.h +++ b/projects/rocsparse/library/src/level2/csrmv_device.h @@ -182,8 +182,8 @@ namespace rocsparse // The launch grid is clamped to the device's maximum grid size, so iterate with a // grid-stride loop to cover every element that needs scaling. - const J stride = (J)WG_SIZE * hipGridDim_x; - for(J gid = (J)hipBlockIdx_x * WG_SIZE + hipThreadIdx_x; gid < required_threads; + const int64_t stride = (int64_t)WG_SIZE * hipGridDim_x; + for(int64_t gid = (int64_t)hipBlockIdx_x * WG_SIZE + hipThreadIdx_x; gid < required_threads; gid += stride) { if(gid < first_row) From c3ba9d64100a9812626f47335b6479fb64d7db42 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Tue, 29 Sep 2026 18:13:45 +0000 Subject: [PATCH 06/15] fix(rocsparse): 64-bit row offsets in csrmv LRB short-row kernels (AISPARSE-662) The grid-stride loops iterate in J, but wg_row_start, wg_row_end and row_idx were narrowed to uint32_t, so with J = int64_t a short-row bin of 2^32 or more rows wrapped onto low rows. Keep the bin offset in J and the in-workgroup offset (bounded by BLOCKSIZE / rows_per_wg) in uint32_t. --- .../library/src/level2/csrmv_device.h | 36 ++++++++++--------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index 6c54c487b1c6..b58c075a807a 100644 --- a/projects/rocsparse/library/src/level2/csrmv_device.h +++ b/projects/rocsparse/library/src/level2/csrmv_device.h @@ -767,18 +767,20 @@ namespace rocsparse for(J gid = hipBlockIdx_x; gid < num_wgs; gid += (J)hipGridDim_x) { - const uint32_t wg_row_start = (uint32_t)(gid * BLOCKSIZE); - const uint32_t wg_row_end - = rocsparse::min(wg_row_start + BLOCKSIZE, uint32_t(bin_num_rows)); - const uint32_t wg_num_rows = wg_row_end - wg_row_start; + // Row offsets within the bin need J (the bin can hold 2^32 or more rows when + // J = int64_t); offsets within the workgroup are bounded by BLOCKSIZE. + const J wg_row_start = gid * (J)BLOCKSIZE; + const uint32_t wg_num_rows + = (uint32_t)rocsparse::min(bin_num_rows - wg_row_start, (J)BLOCKSIZE); + const J* wg_rows_bins = rows_bins + bin_start + wg_row_start; // Load a block of row data using all threads in the WG. for(uint32_t base_idx = 0; base_idx < (BLOCKSIZE << bin_id); base_idx += BLOCKSIZE) { - uint32_t row_idx = wg_row_start + ((base_idx + lid) >> bin_id); - if(row_idx < wg_row_start + wg_num_rows) + const uint32_t row_offset = (base_idx + lid) >> bin_id; + if(row_offset < wg_num_rows) { - const J row_id = rows_bins[row_idx + bin_start]; + const J row_id = wg_rows_bins[row_offset]; const I row_start = csr_row_ptr[row_id] - idx_base; const I row_end = csr_row_ptr[row_id + 1] - idx_base; @@ -808,7 +810,7 @@ namespace rocsparse if(lid < wg_num_rows) { const uint32_t lds_start_idx = (lid << bin_id); - const J row_id = rows_bins[bin_start + wg_row_start + lid]; + const J row_id = wg_rows_bins[lid]; T acc = 0; for(uint32_t idx = 0; idx < (1 << bin_id); idx++) @@ -887,19 +889,21 @@ namespace rocsparse for(J gid = hipBlockIdx_x; gid < num_wgs; gid += (J)hipGridDim_x) { - const uint32_t wg_row_start = (uint32_t)(gid * rows_per_wg); - const uint32_t wg_row_end - = rocsparse::min(wg_row_start + rows_per_wg, uint32_t(bin_num_rows)); - const uint32_t wg_num_rows = wg_row_end - wg_row_start; + // Row offsets within the bin need J (the bin can hold 2^32 or more rows when + // J = int64_t); offsets within the workgroup are bounded by rows_per_wg. + const J wg_row_start = gid * (J)rows_per_wg; + const uint32_t wg_num_rows + = (uint32_t)rocsparse::min(bin_num_rows - wg_row_start, (J)rows_per_wg); + const J wg_bin_start = bin_start + wg_row_start; // Load a block of row data using all threads in the WG. for(uint32_t base_idx = 0; base_idx < CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS; base_idx += BLOCKSIZE) { - uint32_t row_idx = wg_row_start + ((base_idx + lid) >> bin_id); - if(row_idx < wg_row_end) + const uint32_t row_offset = (base_idx + lid) >> bin_id; + if(row_offset < wg_num_rows) { - const J row_id = rows_bins[row_idx + bin_start]; + const J row_id = rows_bins[wg_bin_start + row_offset]; const I row_start = csr_row_ptr[row_id] - idx_base; const I row_end = csr_row_ptr[row_id + 1] - idx_base; @@ -931,7 +935,7 @@ namespace rocsparse if(this_row_offset_in_wg < wg_num_rows) { const uint32_t lds_start_idx = (this_row_offset_in_wg << bin_id); - const J row_id = rows_bins[bin_start + wg_row_start + this_row_offset_in_wg]; + const J row_id = rows_bins[wg_bin_start + this_row_offset_in_wg]; T acc = 0; for(uint32_t idx = 0; idx < (1 << bin_id); idx++) From f1b01bb662ef7fc22667e9aa786a00712acf08c9 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Tue, 29 Sep 2026 18:14:17 +0000 Subject: [PATCH 07/15] perf(rocsparse): hoist bin_start out of csrmv LRB warp-reduce loop (AISPARSE-662) n_rows_bins[bin_id] does not depend on the grid-stride index but was re-read every iteration. Load it once before the loop, as the other LRB kernels already do, instead of relying on the compiler to hoist it. --- projects/rocsparse/library/src/level2/csrmv_device.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index b58c075a807a..5a3d95fc9cf0 100644 --- a/projects/rocsparse/library/src/level2/csrmv_device.h +++ b/projects/rocsparse/library/src/level2/csrmv_device.h @@ -1007,6 +1007,8 @@ namespace rocsparse const int lid = tid & (WF_SIZE - 1); const int wid = tid / WF_SIZE; + const J bin_start = n_rows_bins[bin_id]; + // One wavefront reduces one row. The launch grid is clamped to the device's maximum // grid size, so iterate with a grid-stride loop (in units of wavefronts) to cover every // row. This kernel uses no block-wide barriers, so divergent iteration counts are safe. @@ -1014,8 +1016,7 @@ namespace rocsparse for(int64_t gid = (int64_t)(BLOCKSIZE / WF_SIZE) * bid + wid; gid < count; gid += warps_per_grid) { - const J bin_start = n_rows_bins[bin_id]; - const J row = rows_bins[bin_start + gid]; + const J row = rows_bins[bin_start + gid]; T temp_sum = static_cast(0); const I vecStart = csr_row_ptr[row] - idx_base; From 68cda68cb4ecb317a2cad2e9de709cdea542cf62 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Tue, 29 Sep 2026 18:49:53 +0000 Subject: [PATCH 08/15] perf(rocsparse): drop trailing barrier in csrmv LRB medium/long-row loops (AISPARSE-662) blockreduce_sum already ends with a barrier; after it only lid 0 reads partialSums[0], and only lid 0 writes that slot in the next iteration, so the extra barrier guarded no hazard. In the long-rows kernel it also held the other waves while lid 0 spin-waited. The short-row kernels keep theirs (needed). --- projects/rocsparse/library/src/level2/csrmv_device.h | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index 5a3d95fc9cf0..b3c479193bf4 100644 --- a/projects/rocsparse/library/src/level2/csrmv_device.h +++ b/projects/rocsparse/library/src/level2/csrmv_device.h @@ -1154,9 +1154,8 @@ namespace rocsparse y[row] = temp_sum; } - // Guard against a later grid-stride iteration overwriting partialSums while this - // iteration's reduction is still reading it. - __syncthreads(); + // No trailing barrier: blockreduce_sum ends with one, after which only lid 0 + // reads partialSums[0], and only lid 0 writes that slot in the next iteration. } } @@ -1323,9 +1322,8 @@ namespace rocsparse rocsparse::atomic_add(y, row, m, extra_sum); } - // Guard against a later grid-stride iteration overwriting partialSums while this - // iteration's reduction is still reading it. - __syncthreads(); + // No trailing barrier: blockreduce_sum ends with one, after which only lid 0 + // reads partialSums[0], and only lid 0 writes that slot in the next iteration. } } } From 2c8c6a4c610a029945261f3de5e3f31081522f99 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Tue, 29 Sep 2026 19:05:56 +0000 Subject: [PATCH 09/15] style(rocsparse): clang-format csrmv LRB long-rows grid clamp (AISPARSE-662) --- .../library/src/level2/rocsparse_csrmv_template_lrb.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp index b6b415cb4d05..36721bb3f104 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp @@ -806,8 +806,8 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // groups so that all workgroups cooperating on a row stay within one grid-stride // wave (required by the spin-loop hand-off). The kernel grid-strides over the full // logical range [0, count). - int64_t capped = rocsparse::get_grid_size_x(handle, count, block_size); - capped = (capped / num_wgs_per_row) * num_wgs_per_row; + int64_t capped = rocsparse::get_grid_size_x(handle, count, block_size); + capped = (capped / num_wgs_per_row) * num_wgs_per_row; if(capped < num_wgs_per_row) { capped = num_wgs_per_row; From a6022d6fa857148a49cbac6b01ef15f3f075c104 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Tue, 29 Sep 2026 19:46:49 +0000 Subject: [PATCH 10/15] style(rocsparse): use static_cast instead of C-style casts (AISPARSE-662) --- .../library/src/level2/csrmv_device.h | 38 ++++++++++--------- .../rocsparse_csrmv_template_adaptive.cpp | 2 +- .../level2/rocsparse_csrmv_template_lrb.cpp | 10 +++-- 3 files changed, 27 insertions(+), 23 deletions(-) diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index b3c479193bf4..a5abb06626f1 100644 --- a/projects/rocsparse/library/src/level2/csrmv_device.h +++ b/projects/rocsparse/library/src/level2/csrmv_device.h @@ -182,8 +182,9 @@ namespace rocsparse // The launch grid is clamped to the device's maximum grid size, so iterate with a // grid-stride loop to cover every element that needs scaling. - const int64_t stride = (int64_t)WG_SIZE * hipGridDim_x; - for(int64_t gid = (int64_t)hipBlockIdx_x * WG_SIZE + hipThreadIdx_x; gid < required_threads; + const int64_t stride = static_cast(WG_SIZE) * hipGridDim_x; + for(int64_t gid = static_cast(hipBlockIdx_x) * WG_SIZE + hipThreadIdx_x; + gid < required_threads; gid += stride) { if(gid < first_row) @@ -765,13 +766,13 @@ namespace rocsparse // device's maximum grid size, so iterate with a grid-stride loop to cover every block. const J num_wgs = (bin_num_rows + BLOCKSIZE - 1) / BLOCKSIZE; - for(J gid = hipBlockIdx_x; gid < num_wgs; gid += (J)hipGridDim_x) + for(J gid = hipBlockIdx_x; gid < num_wgs; gid += static_cast(hipGridDim_x)) { // Row offsets within the bin need J (the bin can hold 2^32 or more rows when // J = int64_t); offsets within the workgroup are bounded by BLOCKSIZE. - const J wg_row_start = gid * (J)BLOCKSIZE; - const uint32_t wg_num_rows - = (uint32_t)rocsparse::min(bin_num_rows - wg_row_start, (J)BLOCKSIZE); + const J wg_row_start = gid * static_cast(BLOCKSIZE); + const uint32_t wg_num_rows = static_cast( + rocsparse::min(bin_num_rows - wg_row_start, static_cast(BLOCKSIZE))); const J* wg_rows_bins = rows_bins + bin_start + wg_row_start; // Load a block of row data using all threads in the WG. @@ -887,13 +888,13 @@ namespace rocsparse // device's maximum grid size, so iterate with a grid-stride loop to cover every block. const J num_wgs = (bin_num_rows + rows_per_wg - 1) / rows_per_wg; - for(J gid = hipBlockIdx_x; gid < num_wgs; gid += (J)hipGridDim_x) + for(J gid = hipBlockIdx_x; gid < num_wgs; gid += static_cast(hipGridDim_x)) { // Row offsets within the bin need J (the bin can hold 2^32 or more rows when // J = int64_t); offsets within the workgroup are bounded by rows_per_wg. - const J wg_row_start = gid * (J)rows_per_wg; - const uint32_t wg_num_rows - = (uint32_t)rocsparse::min(bin_num_rows - wg_row_start, (J)rows_per_wg); + const J wg_row_start = gid * static_cast(rows_per_wg); + const uint32_t wg_num_rows = static_cast( + rocsparse::min(bin_num_rows - wg_row_start, static_cast(rows_per_wg))); const J wg_bin_start = bin_start + wg_row_start; // Load a block of row data using all threads in the WG. @@ -1012,8 +1013,8 @@ namespace rocsparse // One wavefront reduces one row. The launch grid is clamped to the device's maximum // grid size, so iterate with a grid-stride loop (in units of wavefronts) to cover every // row. This kernel uses no block-wide barriers, so divergent iteration counts are safe. - const int64_t warps_per_grid = (int64_t)(BLOCKSIZE / WF_SIZE) * hipGridDim_x; - for(int64_t gid = (int64_t)(BLOCKSIZE / WF_SIZE) * bid + wid; gid < count; + const int64_t warps_per_grid = static_cast(BLOCKSIZE / WF_SIZE) * hipGridDim_x; + for(int64_t gid = static_cast(BLOCKSIZE / WF_SIZE) * bid + wid; gid < count; gid += warps_per_grid) { const J row = rows_bins[bin_start + gid]; @@ -1100,7 +1101,7 @@ namespace rocsparse // One block reduces one row. The launch grid is clamped to the device's maximum grid // size, so iterate with a grid-stride loop to cover every row. All threads in a block // share the same gid, so the block-wide barriers below stay uniform. - for(int64_t gid = hipBlockIdx_x; gid < count; gid += (int64_t)hipGridDim_x) + for(int64_t gid = hipBlockIdx_x; gid < count; gid += static_cast(hipGridDim_x)) { const J row = rows_bins[bin_start + gid]; @@ -1203,7 +1204,7 @@ namespace rocsparse // is a multiple of num_wgs_per_row, every workgroup cooperating on a given row falls within // the same grid-stride wave, preserving the cross-workgroup spin-loop synchronization // contract (all workgroups of a row must run concurrently for the flag hand-off to work). - for(int64_t gid = hipBlockIdx_x; gid < count; gid += (int64_t)hipGridDim_x) + for(int64_t gid = hipBlockIdx_x; gid < count; gid += static_cast(hipGridDim_x)) { const J row_idx = gid / num_wgs_per_row; const J wg = gid % num_wgs_per_row; @@ -1217,10 +1218,11 @@ namespace rocsparse // TODO: Can the wg_flags-based coordination be done instead with cooperative-groups? // Each workgroup computes exactly one row. - const I vecStart - = (I)wg * (I)BLOCK_MULTIPLIER * BLOCKSIZE + csr_row_ptr[row] - idx_base; - const I vecEnd = rocsparse::min(csr_row_ptr[row + 1] - idx_base, - vecStart + I(BLOCK_MULTIPLIER * BLOCKSIZE)); + const I vecStart = static_cast(wg) * static_cast(BLOCK_MULTIPLIER) * BLOCKSIZE + + csr_row_ptr[row] - idx_base; + const I vecEnd + = rocsparse::min(csr_row_ptr[row + 1] - idx_base, + vecStart + static_cast(BLOCK_MULTIPLIER * BLOCKSIZE)); T temp_sum = static_cast(0); diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp index 1f0ad0f6692d..d275dccf19ed 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp @@ -771,7 +771,7 @@ rocsparse_status rocsparse::csrmv_adaptive_template_dispatch(rocsparse_handle // Size the grid for the scale operation itself (previously this reused the adaptive // kernel's grid). Clamp the grid; the kernel grid-strides over all required threads. dim3 scale_blocks(rocsparse::get_grid_size_x( - handle, (int64_t(required_threads) - 1) / gen_wg + 1, gen_wg)); + handle, (static_cast(required_threads) - 1) / gen_wg + 1, gen_wg)); dim3 scale_threads(gen_wg); #define ROCSPARSE_LAUNCH_PARTIAL_SCALE_Y(GEN_WG) \ RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp index 36721bb3f104..2db6f1d65744 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp @@ -211,7 +211,7 @@ rocsparse_status rocsparse::csrmv_analysis_lrb_template_dispatch(rocsparse_handl uint32_t block_size = WG_SIZE; uint32_t bin_max_row_len = (1 << j); uint32_t num_wgs_per_row = (bin_max_row_len - 1) / (BLOCK_MULTIPLIER * block_size) + 1; - int64_t grid_size = int64_t(csrmv_info->lrb.nRowsBins[j]) * num_wgs_per_row; + int64_t grid_size = static_cast(csrmv_info->lrb.nRowsBins[j]) * num_wgs_per_row; max_required_grid = rocsparse::max(grid_size, max_required_grid); } @@ -580,7 +580,8 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle { // Overflow-safe integer ceiling division (float ceil loses precision for // row counts above 2^24). nRowsBins[j] is guaranteed non-zero here. - int64_t num_wgs = (int64_t(info->lrb.nRowsBins[j]) - 1) / block_size + 1; + int64_t num_wgs + = (static_cast(info->lrb.nRowsBins[j]) - 1) / block_size + 1; uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, block_size); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( @@ -613,7 +614,8 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle uint32_t rows_per_wg = CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS >> j; // Overflow-safe integer ceiling division (float ceil loses precision for // row counts above 2^24). nRowsBins[j] is guaranteed non-zero here. - int64_t num_wgs = (int64_t(info->lrb.nRowsBins[j]) - 1) / rows_per_wg + 1; + int64_t num_wgs + = (static_cast(info->lrb.nRowsBins[j]) - 1) / rows_per_wg + 1; uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, block_size); RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( @@ -800,7 +802,7 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // Total cooperating workgroups for this bin. Widen to 64-bit before multiplying to // avoid the 32-bit product wrapping (which would both undersize the launch and, via // wg_flags sizing, corrupt the cross-workgroup synchronization). - int64_t count = int64_t(info->lrb.nRowsBins[j]) * num_wgs_per_row; + int64_t count = static_cast(info->lrb.nRowsBins[j]) * num_wgs_per_row; // Clamp the launch grid, rounded down to a whole number of per-row workgroup // groups so that all workgroups cooperating on a row stay within one grid-stride From 4b1b55fe9f927adfd8723bcc68c36a2fa015535c Mon Sep 17 00:00:00 2001 From: kliegeois Date: Tue, 29 Sep 2026 20:43:41 +0000 Subject: [PATCH 11/15] test(rocsparse): forced-clamp unit tests for csrmv LRB and adaptive partial_scale_y (AISPARSE-662) Shrink handle->properties.maxGridSize[0] to 1, 2, 3 and 7 blocks for the compute stage so every LRB kernel (short_rows, short_rows_2, medium_rows_warp_reduce, medium_rows, long_rows) and the adaptive partial_scale_y kernel run their grid-stride loops on a small matrix, and compare y against a host reference. float/double, int32/int64 indices, alpha = 2, beta = -3 (and beta = 0 for partial_scale_y), plus unclamped controls. --- .../clients/unittests/CMakeLists.txt | 2 + ...t_test_internal_level2_csrmv_lrb_grids.cpp | 646 ++++++++++++++++++ 2 files changed, 648 insertions(+) create mode 100644 projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp diff --git a/projects/rocsparse/clients/unittests/CMakeLists.txt b/projects/rocsparse/clients/unittests/CMakeLists.txt index 924a650e44a1..58f701597613 100644 --- a/projects/rocsparse/clients/unittests/CMakeLists.txt +++ b/projects/rocsparse/clients/unittests/CMakeLists.txt @@ -191,6 +191,8 @@ set(ROCSPARSE_UNIT_TEST_DEVICE_SOURCES #10158: distinct anchor to avoid stacked-branch merge conflicts unit_test_host_level3.cpp unit_test_internal_level2_extras.cpp + # csrmv LRB / adaptive partial_scale_y forced grid.x clamp [AISPARSE-662] + unit_test_internal_level2_csrmv_lrb_grids.cpp # --- Level 2 (SpMV/gebsrmv) public C-API host tests [AISPARSE-642] -------- # Distinct anchor (end of the level-2 group) so this leaf does not conflict # with sibling unit-test branches that append after rocsparse_unit_test_main. diff --git a/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp b/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp new file mode 100644 index 000000000000..87b08cfa4899 --- /dev/null +++ b/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp @@ -0,0 +1,646 @@ +/*! \file */ +/* ************************************************************************ + * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights Reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + * + * ************************************************************************ */ + +// +// Forced-clamp tests for the grid.x extents of the csrmv LRB kernels and of the +// adaptive partial_scale_y kernel (AISPARSE-662). +// +// Every LRB compute launch clamps grid.x with rocsparse::get_grid_size_x, i.e. +// to min(maxGridSize[0], (2^32 - 1) / blockDim.x) blocks, and the kernels +// grid-stride over the full logical block count. The long-rows grid is further +// rounded down to a whole multiple of num_wgs_per_row, so that all workgroups +// cooperating on one row run in the same grid-stride wave (the spin-wait +// hand-off needs that). At real sizes the clamp needs more than 16.7M blocks. +// These tests shrink handle->properties.maxGridSize[0] to a few blocks for the +// compute stage instead, so a small matrix runs through the clamped, looping +// path, and compare y against a host reference. +// +// BINS. LRB puts row i in bin ceil(log2(nnz_i)) (0 for an empty row) and +// launches one kernel per non-empty bin: +// +// bins 0-1 csrmvn_lrb_short_rows_kernel ceil(rows / 256) blocks +// bins 2-4 csrmvn_lrb_short_rows_2_kernel ceil(rows / (1024 >> bin)) +// bins 5-8 csrmvn_lrb_medium_rows_warp_reduce_kernel, one wavefront per row +// bins 9-10 csrmvn_lrb_medium_rows_kernel one block per row +// bins 11+ csrmvn_lrb_long_rows_kernel num_wgs_per_row blocks per row +// +// The matrix below populates bins 0 to 14 with enough rows that every one of +// those grids is wider than 7 blocks, so each of the clamped limits 1, 2, 3 and +// 7 forces the stride loop. For long rows num_wgs_per_row is 3, 6, 11 and 22 +// for bins 11 to 14, so the limits cover both "one row per wave" (grid raised +// to num_wgs_per_row) and "two rows per wave" (limit 7, bins 11 and 12). +// Results are checked per bin, so a failure names the kernel that dropped rows. +// +// PARTIAL_SCALE_Y. The adaptive analysis trims leading and trailing runs of +// more than 32 * 256 empty rows out of the adaptive kernel; y is then scaled by +// beta over [0, first_row) and [last_row, m) in partial_scale_y_kernel, which +// only runs for beta != 1. The adaptive matrix has 20000 empty rows on each +// side, so that kernel needs about 40000 threads, i.e. well over 7 blocks. +// +// EXACT ARITHMETIC. Values are small integers, so every product and partial +// sum is exactly representable in float and double and the long-rows atomics +// cannot change the result. The comparison still uses a relative tolerance of a +// few ulps. +// +// TARGET: rocsparse-unit-test-device. The tests drive the public rocsparse_spmv +// entry point and need the complete handle type to reach handle->properties. +// +#include "unit_test_utils.hpp" + +#include "rocsparse_handle.hpp" + +#include "rocsparse.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace rocsparse_ut; + +namespace +{ + // Shrinks the grid.x limit that get_grid_size_x clamps against, and restores + // it on scope exit so a failed assertion cannot leak into the next test. + struct ScopedMaxGridSizeX + { + rocsparse_handle handle; + int saved; + + ScopedMaxGridSizeX(rocsparse_handle h, int limit) + : handle(h) + , saved(h->properties.maxGridSize[0]) + { + handle->properties.maxGridSize[0] = limit; + } + + ~ScopedMaxGridSizeX() + { + handle->properties.maxGridSize[0] = saved; + } + + ScopedMaxGridSizeX(const ScopedMaxGridSizeX&) = delete; + + ScopedMaxGridSizeX& operator=(const ScopedMaxGridSizeX&) = delete; + }; + + constexpr int clamped_limits[] = {1, 2, 3, 7}; + + constexpr int max_bin = 32; + + // Host CSR matrix. Converted to the index / value types under test on upload. + struct HostCsr + { + int64_t m = 0; + int64_t n = 0; + int64_t nnz = 0; + std::vector row_ptr; + std::vector col_ind; + std::vector val; + }; + + int lrb_bin(int64_t row_len) + { + int bin = 0; + while((int64_t(1) << bin) < row_len) + { + ++bin; + } + return bin; + } + + const char* lrb_kernel_of_bin(int bin) + { + if(bin < 2) + { + return "csrmvn_lrb_short_rows_kernel"; + } + if(bin < 5) + { + return "csrmvn_lrb_short_rows_2_kernel"; + } + if(bin < 9) + { + return "csrmvn_lrb_medium_rows_warp_reduce_kernel"; + } + if(bin < 11) + { + return "csrmvn_lrb_medium_rows_kernel"; + } + return "csrmvn_lrb_long_rows_kernel"; + } + + // Builds a CSR matrix from a list of row lengths. Each row holds consecutive + // column indices starting at a row-dependent offset, with values in [-2, 2]. + HostCsr make_matrix(const std::vector& row_len, int64_t n) + { + HostCsr a; + a.m = static_cast(row_len.size()); + a.n = n; + a.row_ptr.assign(a.m + 1, 0); + for(int64_t i = 0; i < a.m; ++i) + { + a.row_ptr[i + 1] = a.row_ptr[i] + row_len[i]; + } + a.nnz = a.row_ptr[a.m]; + a.col_ind.reserve(a.nnz); + a.val.reserve(a.nnz); + for(int64_t i = 0; i < a.m; ++i) + { + const int64_t base = (i * 131) % (n - row_len[i] + 1); + for(int64_t k = 0; k < row_len[i]; ++k) + { + a.col_ind.push_back(base + k); + a.val.push_back(static_cast((i + k) % 5) - 2.0); + } + } + return a; + } + + // Rows spread over LRB bins 0 to 14 (and some empty rows), shuffled so that + // the rows of a bin are not contiguous in the matrix. + HostCsr make_lrb_matrix() + { + struct BinRows + { + int64_t min_len; + int64_t max_len; + int64_t rows; + }; + + const BinRows bins[] = { + {0, 0, 64}, // bin 0, empty rows + {1, 1, 3000}, // bin 0, 12 blocks + {2, 2, 3000}, // bin 1, 12 blocks + {3, 4, 4000}, // bin 2, 16 blocks + {5, 8, 2000}, // bin 3, 16 blocks + {9, 16, 2000}, // bin 4, 32 blocks + {17, 32, 200}, // bin 5, 50 blocks of 4 wavefronts + {33, 64, 200}, // bin 6 + {65, 128, 200}, // bin 7 + {129, 256, 200}, // bin 8 + {257, 512, 40}, // bin 9, 40 blocks + {513, 1024, 40}, // bin 10 + {1025, 2048, 12}, // bin 11, 3 blocks per row, 36 blocks + {2049, 4096, 8}, // bin 12, 6 blocks per row, 48 blocks + {4097, 8192, 6}, // bin 13, 11 blocks per row, 66 blocks + {8193, 16384, 4}, // bin 14, 22 blocks per row, 88 blocks + }; + + std::vector row_len; + for(const BinRows& b : bins) + { + const int64_t span = b.max_len - b.min_len + 1; + for(int64_t r = 0; r < b.rows; ++r) + { + // Hit both ends of the bin, including the exact power of two. + row_len.push_back(r == 0 ? b.max_len : b.min_len + (r * 7) % span); + } + } + + std::mt19937 gen(662); + std::shuffle(row_len.begin(), row_len.end(), gen); + + return make_matrix(row_len, 20000); + } + + // 20000 empty rows, 3000 rows of 1 to 64 entries, then 20000 empty rows. + HostCsr make_adaptive_matrix() + { + constexpr int64_t empty = 20000; + constexpr int64_t body = 3000; + + std::vector row_len(empty + body + empty, 0); + for(int64_t i = 0; i < body; ++i) + { + row_len[empty + i] = 1 + (i * 13) % 64; + } + return make_matrix(row_len, 4096); + } + + std::vector make_x(int64_t n) + { + std::vector x(n); + for(int64_t i = 0; i < n; ++i) + { + x[i] = static_cast(i % 7) - 3.0; + } + return x; + } + + std::vector make_y(int64_t m) + { + std::vector y(m); + for(int64_t i = 0; i < m; ++i) + { + y[i] = static_cast(1 + i % 3); + } + return y; + } + + std::vector host_csrmv(const HostCsr& a, + double alpha, + const std::vector& x, + double beta, + std::vector y) + { + for(int64_t i = 0; i < a.m; ++i) + { + double sum = 0.0; + for(int64_t p = a.row_ptr[i]; p < a.row_ptr[i + 1]; ++p) + { + sum += a.val[p] * x[a.col_ind[p]]; + } + y[i] = alpha * sum + beta * y[i]; + } + return y; + } + + template + std::vector convert(const std::vector& v) + { + return std::vector(v.begin(), v.end()); + } + + template + std::vector convert_values(const std::vector& v) + { + std::vector out(v.size()); + std::transform(v.begin(), v.end(), out.begin(), [](double d) { return static_cast(d); }); + return out; + } + + struct SpmvDescrs + { + rocsparse_spmat_descr mat = nullptr; + rocsparse_dnvec_descr x = nullptr; + rocsparse_dnvec_descr y = nullptr; + + SpmvDescrs() = default; + + SpmvDescrs(const SpmvDescrs&) = delete; + + SpmvDescrs& operator=(const SpmvDescrs&) = delete; + + ~SpmvDescrs() + { + if(mat != nullptr) + { + (void)rocsparse_destroy_spmat_descr(mat); + } + if(x != nullptr) + { + (void)rocsparse_destroy_dnvec_descr(x); + } + if(y != nullptr) + { + (void)rocsparse_destroy_dnvec_descr(y); + } + } + }; + + // y = alpha * A * x + beta * y through rocsparse_spmv (buffer size, + // preprocess, compute). A non-zero clamp_limit shrinks maxGridSize[0] for + // the compute stage, which is where the launches under test live. + // + // Returns an empty vector and records a gtest failure on any API error. + template + std::vector device_csrmv(rocsparse_handle handle, + const HostCsr& a, + rocsparse_spmv_alg alg, + double alpha_d, + const std::vector& x, + double beta_d, + const std::vector& y_in, + int clamp_limit) + { + device_vector d_row_ptr(convert(a.row_ptr)); + device_vector d_col_ind(convert(a.col_ind)); + device_vector d_val(convert_values(a.val)); + device_vector d_x(convert_values(x)); + device_vector d_y(convert_values(y_in)); + + if(d_row_ptr.ptr == nullptr || d_col_ind.ptr == nullptr || d_val.ptr == nullptr + || d_x.ptr == nullptr || d_y.ptr == nullptr) + { + ADD_FAILURE() << "device allocation failed"; + return {}; + } + + SpmvDescrs descrs; + if(rocsparse_create_csr_descr(&descrs.mat, + a.m, + a.n, + a.nnz, + d_row_ptr.ptr, + d_col_ind.ptr, + d_val.ptr, + it_of(), + it_of(), + rocsparse_index_base_zero, + dt_of()) + != rocsparse_status_success + || rocsparse_create_dnvec_descr(&descrs.x, a.n, d_x.ptr, dt_of()) + != rocsparse_status_success + || rocsparse_create_dnvec_descr(&descrs.y, a.m, d_y.ptr, dt_of()) + != rocsparse_status_success) + { + ADD_FAILURE() << "descriptor creation failed"; + return {}; + } + + const T alpha = static_cast(alpha_d); + const T beta = static_cast(beta_d); + + auto spmv = [&](rocsparse_spmv_stage stage, size_t* buffer_size, void* buffer) { + return rocsparse_spmv(handle, + rocsparse_operation_none, + &alpha, + descrs.mat, + descrs.x, + &beta, + descrs.y, + dt_of(), + alg, + stage, + buffer_size, + buffer); + }; + + size_t buffer_size = 0; + rocsparse_status status = spmv(rocsparse_spmv_stage_buffer_size, &buffer_size, nullptr); + if(status != rocsparse_status_success) + { + ADD_FAILURE() << "rocsparse_spmv buffer_size returned status " << status; + return {}; + } + + device_vector d_buffer(buffer_size > 0 ? buffer_size : size_t(1)); + if(d_buffer.ptr == nullptr) + { + ADD_FAILURE() << "temp buffer allocation of " << buffer_size << " bytes failed"; + return {}; + } + + status = spmv(rocsparse_spmv_stage_preprocess, &buffer_size, d_buffer.ptr); + if(status != rocsparse_status_success) + { + ADD_FAILURE() << "rocsparse_spmv preprocess returned status " << status; + return {}; + } + + { + std::unique_ptr clamp; + if(clamp_limit > 0) + { + clamp.reset(new ScopedMaxGridSizeX(handle, clamp_limit)); + } + status = spmv(rocsparse_spmv_stage_compute, &buffer_size, d_buffer.ptr); + } + if(status != rocsparse_status_success) + { + ADD_FAILURE() << "rocsparse_spmv compute returned status " << status; + return {}; + } + + if(hipDeviceSynchronize() != hipSuccess) + { + ADD_FAILURE() << "hipDeviceSynchronize failed"; + return {}; + } + + const std::vector y = to_host(d_y); + return std::vector(y.begin(), y.end()); + } + + template + bool near(double got, double want) + { + const double tol = 4.0 * std::numeric_limits::epsilon() * std::max(1.0, std::abs(want)); + return std::abs(got - want) <= tol; + } + + // Compares y per LRB bin, so a failure names the kernel that owns the rows. + template + void expect_lrb_match(const HostCsr& a, + const std::vector& got, + const std::vector& want, + const std::string& what) + { + ASSERT_EQ(got.size(), want.size()); + + std::vector rows(max_bin, 0); + std::vector bad(max_bin, 0); + std::vector first_bad(max_bin, -1); + for(int64_t i = 0; i < a.m; ++i) + { + const int bin = lrb_bin(a.row_ptr[i + 1] - a.row_ptr[i]); + ++rows[bin]; + if(!near(got[i], want[i])) + { + if(bad[bin]++ == 0) + { + first_bad[bin] = i; + } + } + } + + for(int bin = 0; bin < max_bin; ++bin) + { + if(rows[bin] == 0) + { + continue; + } + const int64_t i = first_bad[bin]; + EXPECT_EQ(bad[bin], 0) + << what << ": " << lrb_kernel_of_bin(bin) << " (bin " << bin << ") got " << bad[bin] + << " of " << rows[bin] << " rows wrong; first is row " << i << " with " + << (i >= 0 ? a.row_ptr[i + 1] - a.row_ptr[i] : 0) << " entries, got " + << (i >= 0 ? got[i] : 0.0) << " want " << (i >= 0 ? want[i] : 0.0); + } + } + + template + void expect_adaptive_match(const std::vector& got, + const std::vector& want, + const std::string& what) + { + ASSERT_EQ(got.size(), want.size()); + + int64_t bad = 0; + int64_t first = -1; + for(size_t i = 0; i < got.size(); ++i) + { + if(!near(got[i], want[i])) + { + if(bad++ == 0) + { + first = static_cast(i); + } + } + } + EXPECT_EQ(bad, 0) << what << ": " << bad << " of " << got.size() + << " rows wrong; first is row " << first << ", got " + << (first >= 0 ? got[first] : 0.0) << " want " + << (first >= 0 ? want[first] : 0.0); + } + + std::string describe(const char* alg, int clamp_limit, double beta) + { + std::ostringstream os; + os << alg << ", maxGridSize[0] = "; + if(clamp_limit > 0) + { + os << clamp_limit; + } + else + { + os << "unclamped"; + } + os << ", beta = " << beta; + return os.str(); + } + + constexpr double lrb_alpha = 2.0; + constexpr double lrb_beta = -3.0; + + template + void check_lrb(rocsparse_handle handle, bool clamped) + { + const HostCsr a = make_lrb_matrix(); + const std::vector x = make_x(a.n); + const std::vector y = make_y(a.m); + const std::vector want = host_csrmv(a, lrb_alpha, x, lrb_beta, y); + + const std::vector limits + = clamped ? std::vector(std::begin(clamped_limits), std::end(clamped_limits)) + : std::vector{0}; + for(const int limit : limits) + { + const std::vector got = device_csrmv( + handle, a, rocsparse_spmv_alg_csr_lrb, lrb_alpha, x, lrb_beta, y, limit); + expect_lrb_match(a, got, want, describe("lrb", limit, lrb_beta)); + } + } + + template + void check_adaptive_partial_scale_y(rocsparse_handle handle, bool clamped) + { + const HostCsr a = make_adaptive_matrix(); + const std::vector x = make_x(a.n); + const std::vector y = make_y(a.m); + const double alpha = 2.0; + + const std::vector limits + = clamped ? std::vector(std::begin(clamped_limits), std::end(clamped_limits)) + : std::vector{0}; + + // beta = 0 takes the store-zero branch of partial_scale_y, beta = -3 the + // multiply branch. beta = 1 skips the kernel entirely. + for(const double beta : {-3.0, 0.0}) + { + const std::vector want = host_csrmv(a, alpha, x, beta, y); + for(const int limit : limits) + { + const std::vector got = device_csrmv( + handle, a, rocsparse_spmv_alg_csr_adaptive, alpha, x, beta, y, limit); + expect_adaptive_match(got, want, describe("adaptive", limit, beta)); + } + } + } + + using CsrmvLrbGrids = HandleTest; + using CsrmvAdaptiveGrids = HandleTest; +} + +// --------------------------------------------------------------------------- +// LRB, grid.x clamped to 1, 2, 3 and 7 blocks for every kernel. +// --------------------------------------------------------------------------- + +TEST_F(CsrmvLrbGrids, clamped_grid_stride_f32_i32) +{ + check_lrb(handle, true); +} + +TEST_F(CsrmvLrbGrids, clamped_grid_stride_f64_i32) +{ + check_lrb(handle, true); +} + +TEST_F(CsrmvLrbGrids, clamped_grid_stride_f32_i64) +{ + check_lrb(handle, true); +} + +TEST_F(CsrmvLrbGrids, clamped_grid_stride_f64_i64) +{ + check_lrb(handle, true); +} + +// Control: the same problems with the real grid.x limit. If this fails the +// clamped cases above prove nothing. +TEST_F(CsrmvLrbGrids, unclamped_grid_matches_host) +{ + check_lrb(handle, false); + check_lrb(handle, false); + check_lrb(handle, false); + check_lrb(handle, false); +} + +// --------------------------------------------------------------------------- +// Adaptive partial_scale_y, grid.x clamped to 1, 2, 3 and 7 blocks. +// --------------------------------------------------------------------------- + +TEST_F(CsrmvAdaptiveGrids, partial_scale_y_clamped_f32_i32) +{ + check_adaptive_partial_scale_y(handle, true); +} + +TEST_F(CsrmvAdaptiveGrids, partial_scale_y_clamped_f64_i32) +{ + check_adaptive_partial_scale_y(handle, true); +} + +TEST_F(CsrmvAdaptiveGrids, partial_scale_y_clamped_f32_i64) +{ + check_adaptive_partial_scale_y(handle, true); +} + +TEST_F(CsrmvAdaptiveGrids, partial_scale_y_clamped_f64_i64) +{ + check_adaptive_partial_scale_y(handle, true); +} + +TEST_F(CsrmvAdaptiveGrids, unclamped_grid_matches_host) +{ + check_adaptive_partial_scale_y(handle, false); + check_adaptive_partial_scale_y(handle, false); + check_adaptive_partial_scale_y(handle, false); + check_adaptive_partial_scale_y(handle, false); +} From 97e5ccd0a8680eef385e572a74048c1a2db0fce0 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Wed, 30 Sep 2026 17:10:58 +0000 Subject: [PATCH 12/15] perf(rocsparse): straight-line csrmv LRB kernels when grid.x is not clamped (AISPARSE-662) The grid-stride conversion made every csrmv LRB launch pay for the large-grid case: a 64-bit loop over hipGridDim_x, 64-bit row offsets, and in the short-row kernels an extra block barrier per iteration, even though the loop runs once whenever the grid is not clamped. On gfx942 this raised VGPR and SGPR counts in all five kernels and slowed normal SpMV. Give each LRB kernel a GRID_STRIDE template parameter. The host launches the straight-line variant whenever get_grid_size_x returns the full block count, and the grid-stride variant only when the clamp binds. The straight-line variant is develop's kernel, with 32-bit row offsets; they cannot wrap there because blocks * blockDim.x < 2^32. The grid-stride variant keeps the 64-bit loop, the barrier between short-row iterations, and the long-rows grid rounded to a multiple of num_wgs_per_row. Both variants share one per-block body. Launches go through rocsparse::dispatch_grid_stride_x, added here unchanged from the AISPARSE-677 branch; the long-rows site dispatches on its rounded grid directly. The forced-clamp unit tests run the grid-stride variants and the unclamped control runs the straight-line ones. --- ...t_test_internal_level2_csrmv_lrb_grids.cpp | 10 +- .../library/src/include/rocsparse_grid.hpp | 20 + .../library/src/level2/csrmv_device.h | 199 ++++--- .../level2/rocsparse_csrmv_template_lrb.cpp | 513 ++++++++++-------- 4 files changed, 458 insertions(+), 284 deletions(-) diff --git a/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp b/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp index 87b08cfa4899..1589ae08d809 100644 --- a/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp +++ b/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp @@ -27,8 +27,10 @@ // adaptive partial_scale_y kernel (AISPARSE-662). // // Every LRB compute launch clamps grid.x with rocsparse::get_grid_size_x, i.e. -// to min(maxGridSize[0], (2^32 - 1) / blockDim.x) blocks, and the kernels -// grid-stride over the full logical block count. The long-rows grid is further +// to min(maxGridSize[0], (2^32 - 1) / blockDim.x) blocks. When the clamp binds, +// the host launches the GRID_STRIDE variant of the kernel, which grid-strides +// over the full logical block count; otherwise it launches the straight-line +// variant, one block per block index. The long-rows grid is further // rounded down to a whole multiple of num_wgs_per_row, so that all workgroups // cooperating on one row run in the same grid-stride wave (the spin-wait // hand-off needs that). At real sizes the clamp needs more than 16.7M blocks. @@ -603,8 +605,8 @@ TEST_F(CsrmvLrbGrids, clamped_grid_stride_f64_i64) check_lrb(handle, true); } -// Control: the same problems with the real grid.x limit. If this fails the -// clamped cases above prove nothing. +// Control: the same problems with the real grid.x limit, which run the +// straight-line variants. If this fails the clamped cases above prove nothing. TEST_F(CsrmvLrbGrids, unclamped_grid_matches_host) { check_lrb(handle, false); diff --git a/projects/rocsparse/library/src/include/rocsparse_grid.hpp b/projects/rocsparse/library/src/include/rocsparse_grid.hpp index ba83d82b9f1f..bd92dbc81a8d 100644 --- a/projects/rocsparse/library/src/include/rocsparse_grid.hpp +++ b/projects/rocsparse/library/src/include/rocsparse_grid.hpp @@ -27,6 +27,7 @@ #include "rocsparse_handle.hpp" #include +#include // Grid-extent clamps, one per axis. Each takes the handle and reads the limit // from the device, so a caller never writes a bound itself. Any kernel launched @@ -109,4 +110,23 @@ namespace rocsparse return rocsparse::clamp_grid_extent( count, static_cast(handle->properties.maxGridSize[2])); } + + // Launch with the grid.x extent get_grid_size_x returns, selecting the kernel + // variant at the call site. launch(std::true_type{}, grid) runs when the clamp + // binds and the kernel must grid-stride over count; launch(std::false_type{}, + // grid) runs otherwise, so a kernel can keep a straight-line variant for the + // common case where every row gets its own block. The clamp binds as soon as + // count exceeds dispatch_limit_x(block_size), 16777215 blocks at 256 threads, + // which 32-bit indices reach as readily as 64-bit ones. + template + static __forceinline__ rocsparse_status + dispatch_grid_stride_x(rocsparse_handle handle, J count, int64_t block_size, F&& launch) + { + const uint32_t grid = rocsparse::get_grid_size_x(handle, count, block_size); + if(grid < static_cast(count)) + { + return launch(std::true_type{}, grid); + } + return launch(std::false_type{}, grid); + } } diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index a5abb06626f1..ac059501545d 100644 --- a/projects/rocsparse/library/src/level2/csrmv_device.h +++ b/projects/rocsparse/library/src/level2/csrmv_device.h @@ -727,8 +727,16 @@ namespace rocsparse } } + // The csrmv LRB kernels below come in two variants, selected by GRID_STRIDE. + // The host launches GRID_STRIDE = false only when rocsparse::get_grid_size_x + // did not clamp grid.x, i.e. every block of the bin has its own block index and + // blocks * blockDim.x < 2^32. Every row offset within a bin is then bounded by + // that product, so the straight-line variant keeps 32-bit offsets. When the + // clamp binds, GRID_STRIDE = true iterates over the blocks in 64 bits. + // "Stream" case a la CSR-Adaptive template (hipGridDim_x)) - { - // Row offsets within the bin need J (the bin can hold 2^32 or more rows when - // J = int64_t); offsets within the workgroup are bounded by BLOCKSIZE. - const J wg_row_start = gid * static_cast(BLOCKSIZE); - const uint32_t wg_num_rows = static_cast( - rocsparse::min(bin_num_rows - wg_row_start, static_cast(BLOCKSIZE))); - const J* wg_rows_bins = rows_bins + bin_start + wg_row_start; + // Computes the rows [wg_row_start, wg_row_end) of the bin. + const auto process_block = [&](auto wg_row_start, auto wg_row_end) { + const uint32_t wg_num_rows = static_cast(wg_row_end - wg_row_start); // Load a block of row data using all threads in the WG. for(uint32_t base_idx = 0; base_idx < (BLOCKSIZE << bin_id); base_idx += BLOCKSIZE) { - const uint32_t row_offset = (base_idx + lid) >> bin_id; - if(row_offset < wg_num_rows) + const auto row_idx = wg_row_start + ((base_idx + lid) >> bin_id); + if(row_idx < wg_row_start + wg_num_rows) { - const J row_id = wg_rows_bins[row_offset]; + const J row_id = rows_bins[row_idx + bin_start]; const I row_start = csr_row_ptr[row_id] - idx_base; const I row_end = csr_row_ptr[row_id + 1] - idx_base; @@ -811,7 +810,7 @@ namespace rocsparse if(lid < wg_num_rows) { const uint32_t lds_start_idx = (lid << bin_id); - const J row_id = wg_rows_bins[lid]; + const J row_id = rows_bins[bin_start + wg_row_start + lid]; T acc = 0; for(uint32_t idx = 0; idx < (1 << bin_id); idx++) @@ -837,10 +836,29 @@ namespace rocsparse y[row_id] = acc; } + }; - // Guard against a later grid-stride iteration overwriting partialSums while this - // iteration's reduction is still reading it. - __syncthreads(); + if constexpr(GRID_STRIDE) + { + const int64_t num_wgs = (static_cast(bin_num_rows) - 1) / BLOCKSIZE + 1; + for(int64_t gid = hipBlockIdx_x; gid < num_wgs; gid += hipGridDim_x) + { + const J wg_row_start = static_cast(gid * BLOCKSIZE); + process_block( + wg_row_start, + wg_row_start + + rocsparse::min(bin_num_rows - wg_row_start, static_cast(BLOCKSIZE))); + + // The next iteration overwrites partialSums, which this reduction reads. + __syncthreads(); + } + } + else + { + const uint32_t wg_row_start = hipBlockIdx_x * BLOCKSIZE; + const uint32_t wg_row_end + = rocsparse::min(wg_row_start + BLOCKSIZE, static_cast(bin_num_rows)); + process_block(wg_row_start, wg_row_end); } } @@ -849,6 +867,7 @@ namespace rocsparse // dynamic LDS allocation approach would blow up size requirements beyond reasonable bounds. template > bin_id; - // Number of workgroups required to cover this bin. The launch grid is clamped to the - // device's maximum grid size, so iterate with a grid-stride loop to cover every block. - const J num_wgs = (bin_num_rows + rows_per_wg - 1) / rows_per_wg; - - for(J gid = hipBlockIdx_x; gid < num_wgs; gid += static_cast(hipGridDim_x)) - { - // Row offsets within the bin need J (the bin can hold 2^32 or more rows when - // J = int64_t); offsets within the workgroup are bounded by rows_per_wg. - const J wg_row_start = gid * static_cast(rows_per_wg); - const uint32_t wg_num_rows = static_cast( - rocsparse::min(bin_num_rows - wg_row_start, static_cast(rows_per_wg))); - const J wg_bin_start = bin_start + wg_row_start; + // Computes the rows [wg_row_start, wg_row_end) of the bin. + const auto process_block = [&](auto wg_row_start, auto wg_row_end) { + const uint32_t wg_num_rows = static_cast(wg_row_end - wg_row_start); // Load a block of row data using all threads in the WG. for(uint32_t base_idx = 0; base_idx < CSRMV_LRB_SHORT_ROWS_2_LDS_ELEMS; base_idx += BLOCKSIZE) { - const uint32_t row_offset = (base_idx + lid) >> bin_id; - if(row_offset < wg_num_rows) + const auto row_idx = wg_row_start + ((base_idx + lid) >> bin_id); + if(row_idx < wg_row_start + wg_num_rows) { - const J row_id = rows_bins[wg_bin_start + row_offset]; + const J row_id = rows_bins[row_idx + bin_start]; const I row_start = csr_row_ptr[row_id] - idx_base; const I row_end = csr_row_ptr[row_id + 1] - idx_base; @@ -936,7 +946,7 @@ namespace rocsparse if(this_row_offset_in_wg < wg_num_rows) { const uint32_t lds_start_idx = (this_row_offset_in_wg << bin_id); - const J row_id = rows_bins[wg_bin_start + this_row_offset_in_wg]; + const J row_id = rows_bins[bin_start + wg_row_start + this_row_offset_in_wg]; T acc = 0; for(uint32_t idx = 0; idx < (1 << bin_id); idx++) @@ -961,16 +971,36 @@ namespace rocsparse y[row_id] = acc; } } + }; - // Guard against a later grid-stride iteration overwriting partialSums while this - // iteration's reduction is still reading it. - __syncthreads(); + if constexpr(GRID_STRIDE) + { + const int64_t num_wgs = (static_cast(bin_num_rows) - 1) / rows_per_wg + 1; + for(int64_t gid = hipBlockIdx_x; gid < num_wgs; gid += hipGridDim_x) + { + const J wg_row_start = static_cast(gid * rows_per_wg); + process_block( + wg_row_start, + wg_row_start + + rocsparse::min(bin_num_rows - wg_row_start, static_cast(rows_per_wg))); + + // The next iteration overwrites partialSums, which this reduction reads. + __syncthreads(); + } + } + else + { + const uint32_t wg_row_start = hipBlockIdx_x * rows_per_wg; + const uint32_t wg_row_end + = rocsparse::min(wg_row_start + rows_per_wg, static_cast(bin_num_rows)); + process_block(wg_row_start, wg_row_end); } } // "Vector" case a la CSR-Adaptive using one warp per row template (BLOCKSIZE / WF_SIZE) * hipGridDim_x; - for(int64_t gid = static_cast(BLOCKSIZE / WF_SIZE) * bid + wid; gid < count; - gid += warps_per_grid) - { - const J row = rows_bins[bin_start + gid]; + // One wavefront reduces the gid-th row of the bin. + const auto process_row = [&](auto gid) { + const J bin_start = n_rows_bins[bin_id]; + const J row = rows_bins[bin_start + gid]; T temp_sum = static_cast(0); const I vecStart = csr_row_ptr[row] - idx_base; @@ -1053,11 +1077,34 @@ namespace rocsparse y[row] = temp_sum; } + }; + + if constexpr(GRID_STRIDE) + { + // This kernel uses no block-wide barriers, so divergent iteration counts are safe. + const int64_t warps_per_grid = static_cast(BLOCKSIZE / WF_SIZE) * hipGridDim_x; + for(int64_t gid = static_cast(BLOCKSIZE / WF_SIZE) * bid + wid; gid < count; + gid += warps_per_grid) + { + process_row(gid); + } + } + else + { + const int gid = (BLOCKSIZE / WF_SIZE) * bid + wid; + + if(gid >= count) + { + return; + } + + process_row(gid); } } // "Vector" case a la CSR-Adaptive using one block per row template (hipGridDim_x)) - { + // One block reduces the gid-th row of the bin. + const auto process_row = [&](auto gid) { const J row = rows_bins[bin_start + gid]; // In its original form, any CSR-Vector workgroup only calculates, at most, BLOCKSIZE items in its row; @@ -1154,15 +1198,28 @@ namespace rocsparse y[row] = temp_sum; } + }; - // No trailing barrier: blockreduce_sum ends with one, after which only lid 0 - // reads partialSums[0], and only lid 0 writes that slot in the next iteration. + if constexpr(GRID_STRIDE) + { + // All threads in a block share the same gid, so the block-wide barriers stay uniform. + // No barrier between iterations: blockreduce_sum ends with one, after which only + // lid 0 reads partialSums[0], and only lid 0 writes that slot in the next iteration. + for(int64_t gid = hipBlockIdx_x; gid < count; gid += hipGridDim_x) + { + process_row(gid); + } + } + else + { + process_row(static_cast(hipBlockIdx_x)); } } // "LongRows" aka "VectorL" case a la CSR-Adaptive template (hipGridDim_x)) - { + // Computes the part of a row that the gid-th cooperating workgroup of the bin owns. + const auto process_block = [&](auto gid) { const J row_idx = gid / num_wgs_per_row; const J wg = gid % num_wgs_per_row; const J row = rows_bins[bin_start + row_idx]; @@ -1323,9 +1374,23 @@ namespace rocsparse rocsparse::atomic_add(y, row, m, extra_sum); } + }; - // No trailing barrier: blockreduce_sum ends with one, after which only lid 0 - // reads partialSums[0], and only lid 0 writes that slot in the next iteration. + if constexpr(GRID_STRIDE) + { + // count is nRowsBins[bin_id] * num_wgs_per_row. The host rounds the clamped grid down + // to a multiple of num_wgs_per_row, so all workgroups cooperating on a row run in the + // same grid-stride iteration, which the spin-loop hand-off requires. + // No barrier between iterations: blockreduce_sum ends with one, after which only + // lid 0 reads partialSums[0], and only lid 0 writes that slot in the next iteration. + for(int64_t gid = hipBlockIdx_x; gid < count; gid += hipGridDim_x) + { + process_block(gid); + } + } + else + { + process_block(static_cast(hipBlockIdx_x)); } } } diff --git a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp index 2db6f1d65744..7fdd1e6d2663 100644 --- a/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp +++ b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_lrb.cpp @@ -241,7 +241,14 @@ rocsparse_status rocsparse::csrmv_analysis_lrb_template_dispatch(rocsparse_handl namespace rocsparse { - template + template ROCSPARSE_KERNEL(WG_SIZE) void csrmvn_lrb_short_rows_kernel(bool conj, I nnz, @@ -265,7 +272,7 @@ namespace rocsparse ROCSPARSE_DEVICE_HOST_SCALAR_GET(beta); if(alpha != 0 || beta != 1 || num_extra > 0) { - rocsparse::csrmvn_lrb_short_rows_device( + rocsparse::csrmvn_lrb_short_rows_device( conj, nnz, rows_bins, @@ -285,7 +292,14 @@ namespace rocsparse } } - template + template ROCSPARSE_KERNEL(WG_SIZE) void csrmvn_lrb_short_rows_2_kernel(bool conj, I nnz, @@ -311,6 +325,7 @@ namespace rocsparse { rocsparse::csrmvn_lrb_short_rows_2_device 0) { - rocsparse::csrmvn_lrb_medium_rows_warp_reduce_device( + rocsparse::csrmvn_lrb_medium_rows_warp_reduce_device( conj, nnz, count, @@ -391,6 +407,7 @@ namespace rocsparse } template 0) { - rocsparse::csrmvn_lrb_medium_rows_device(conj, - nnz, - count, - rows_bins, - n_rows_bins, - bin_id, - alpha, - csr_row_ptr, - csr_col_ind, - csr_val, - x, - beta, - y, - num_extra, - gamma_device_array, - z_arrays, - idx_base); + rocsparse::csrmvn_lrb_medium_rows_device(conj, + nnz, + count, + rows_bins, + n_rows_bins, + bin_id, + alpha, + csr_row_ptr, + csr_col_ind, + csr_val, + x, + beta, + y, + num_extra, + gamma_device_array, + z_arrays, + idx_base); } } - template + template ROCSPARSE_KERNEL(WG_SIZE) void csrmvn_lrb_long_rows_kernel(bool conj, J m, @@ -469,26 +493,34 @@ namespace rocsparse ROCSPARSE_DEVICE_HOST_SCALAR_GET(beta); if(alpha != 0 || beta != 1 || num_extra > 0) { - rocsparse::csrmvn_lrb_long_rows_device( - conj, - m, - nnz, - count, - wg_flags, - rows_bins, - n_rows_bins, - bin_id, - alpha, - csr_row_ptr, - csr_col_ind, - csr_val, - x, - beta, - y, - num_extra, - gamma_device_array, - z_arrays, - idx_base); + rocsparse::csrmvn_lrb_long_rows_device(conj, + m, + nnz, + count, + wg_flags, + rows_bins, + n_rows_bins, + bin_id, + alpha, + csr_row_ptr, + csr_col_ind, + csr_val, + x, + beta, + y, + num_extra, + gamma_device_array, + z_arrays, + idx_base); } } } @@ -582,31 +614,37 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // row counts above 2^24). nRowsBins[j] is guaranteed non-zero here. int64_t num_wgs = (static_cast(info->lrb.nRowsBins[j]) - 1) / block_size + 1; - uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, block_size); - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (csrmvn_lrb_short_rows_kernel), - grid_size, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + num_wgs, block_size, - lds_size, - stream, - conj, - nnz, - static_cast(info->lrb.rows_bins), - static_cast(info->lrb.n_rows_bins), - j, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), - csr_row_ptr, - csr_col_ind, - csr_val, - x, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), - y, - num_extra, - gamma_device_array, - z_array, - descr->base, - handle->pointer_mode == rocsparse_pointer_mode_host); + [&](auto grid_stride, uint32_t grid_size) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (csrmvn_lrb_short_rows_kernel), + grid_size, + block_size, + lds_size, + stream, + conj, + nnz, + static_cast(info->lrb.rows_bins), + static_cast(info->lrb.n_rows_bins), + j, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), + csr_row_ptr, + csr_col_ind, + csr_val, + x, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), + y, + num_extra, + gamma_device_array, + z_array, + descr->base, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } // Static LDS allocation, for when dynamic would grow too large else @@ -616,31 +654,37 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // row counts above 2^24). nRowsBins[j] is guaranteed non-zero here. int64_t num_wgs = (static_cast(info->lrb.nRowsBins[j]) - 1) / rows_per_wg + 1; - uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, block_size); - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (csrmvn_lrb_short_rows_2_kernel), - grid_size, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + num_wgs, block_size, - 0, - stream, - conj, - nnz, - static_cast(info->lrb.rows_bins), - static_cast(info->lrb.n_rows_bins), - j, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), - csr_row_ptr, - csr_col_ind, - csr_val, - x, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), - y, - num_extra, - gamma_device_array, - z_array, - descr->base, - handle->pointer_mode == rocsparse_pointer_mode_host); + [&](auto grid_stride, uint32_t grid_size) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (csrmvn_lrb_short_rows_2_kernel), + grid_size, + block_size, + 0, + stream, + conj, + nnz, + static_cast(info->lrb.rows_bins), + static_cast(info->lrb.n_rows_bins), + j, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), + csr_row_ptr, + csr_col_ind, + csr_val, + x, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), + y, + num_extra, + gamma_device_array, + z_array, + descr->base, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } } } @@ -661,126 +705,157 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle // RDNA (wave32): use 128-thread blocks (4 wavefronts) instead of 256. // One wavefront reduces one row, so grid is sized by wavefronts-per-block. constexpr uint32_t mr_block = WG_SIZE_WAVE32; - int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (mr_block / 32) + 1; - uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, mr_block); - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (csrmvn_lrb_medium_rows_warp_reduce_kernel), - grid_size, + int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (mr_block / 32) + 1; + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + num_wgs, mr_block, - 0, - stream, - conj, - nnz, - info->lrb.nRowsBins[j], - static_cast(info->lrb.rows_bins), - static_cast(info->lrb.n_rows_bins), - j, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), - csr_row_ptr, - csr_col_ind, - csr_val, - x, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), - y, - num_extra, - gamma_device_array, - z_array, - descr->base, - handle->pointer_mode == rocsparse_pointer_mode_host); + [&](auto grid_stride, uint32_t grid_size) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (csrmvn_lrb_medium_rows_warp_reduce_kernel< + mr_block, + 32, + decltype(grid_stride)::value>), + grid_size, + mr_block, + 0, + stream, + conj, + nnz, + info->lrb.nRowsBins[j], + static_cast(info->lrb.rows_bins), + static_cast(info->lrb.n_rows_bins), + j, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), + csr_row_ptr, + csr_col_ind, + csr_val, + x, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), + y, + num_extra, + gamma_device_array, + z_array, + descr->base, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } else { int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (256 / handle->wavefront_size) + 1; - uint32_t grid_size = rocsparse::get_grid_size_x(handle, num_wgs, 256); - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (csrmvn_lrb_medium_rows_warp_reduce_kernel<256, 64>), - grid_size, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + num_wgs, 256, - 0, - stream, - conj, - nnz, - info->lrb.nRowsBins[j], - static_cast(info->lrb.rows_bins), - static_cast(info->lrb.n_rows_bins), - j, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), - csr_row_ptr, - csr_col_ind, - csr_val, - x, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), - y, - num_extra, - gamma_device_array, - z_array, - descr->base, - handle->pointer_mode == rocsparse_pointer_mode_host); + [&](auto grid_stride, uint32_t grid_size) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (csrmvn_lrb_medium_rows_warp_reduce_kernel< + 256, + 64, + decltype(grid_stride)::value>), + grid_size, + 256, + 0, + stream, + conj, + nnz, + info->lrb.nRowsBins[j], + static_cast(info->lrb.rows_bins), + static_cast(info->lrb.n_rows_bins), + j, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), + csr_row_ptr, + csr_col_ind, + csr_val, + x, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), + y, + num_extra, + gamma_device_array, + z_array, + descr->base, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } } else // One block per row { - int64_t count = info->lrb.nRowsBins[j]; // One WG per row - const uint32_t mr_block - = (handle->wavefront_size == 32) ? WG_SIZE_WAVE32 : WG_SIZE; - uint32_t grid_size = rocsparse::get_grid_size_x(handle, count, mr_block); + int64_t count = info->lrb.nRowsBins[j]; // One WG per row if(handle->wavefront_size == 32) { // RDNA (wave32): 128-thread block per row (4 wavefronts) improves // occupancy vs the 256-thread (8-wavefront) block for these rows. - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (csrmvn_lrb_medium_rows_kernel), - grid_size, - WG_SIZE_WAVE32, - 0, - stream, - conj, - nnz, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, count, - static_cast(info->lrb.rows_bins), - static_cast(info->lrb.n_rows_bins), - j, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), - csr_row_ptr, - csr_col_ind, - csr_val, - x, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), - y, - num_extra, - gamma_device_array, - z_array, - descr->base, - handle->pointer_mode == rocsparse_pointer_mode_host); + WG_SIZE_WAVE32, + [&](auto grid_stride, uint32_t grid_size) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (csrmvn_lrb_medium_rows_kernel), + grid_size, + WG_SIZE_WAVE32, + 0, + stream, + conj, + nnz, + count, + static_cast(info->lrb.rows_bins), + static_cast(info->lrb.n_rows_bins), + j, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), + csr_row_ptr, + csr_col_ind, + csr_val, + x, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), + y, + num_extra, + gamma_device_array, + z_array, + descr->base, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } else { - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (csrmvn_lrb_medium_rows_kernel), - grid_size, - WG_SIZE, - 0, - stream, - conj, - nnz, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, count, - static_cast(info->lrb.rows_bins), - static_cast(info->lrb.n_rows_bins), - j, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), - csr_row_ptr, - csr_col_ind, - csr_val, - x, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), - y, - num_extra, - gamma_device_array, - z_array, - descr->base, - handle->pointer_mode == rocsparse_pointer_mode_host); + WG_SIZE, + [&](auto grid_stride, uint32_t grid_size) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (csrmvn_lrb_medium_rows_kernel), + grid_size, + WG_SIZE, + 0, + stream, + conj, + nnz, + count, + static_cast(info->lrb.rows_bins), + static_cast(info->lrb.n_rows_bins), + j, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), + csr_row_ptr, + csr_col_ind, + csr_val, + x, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), + y, + num_extra, + gamma_device_array, + z_array, + descr->base, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } } } @@ -816,32 +891,44 @@ rocsparse_status rocsparse::csrmv_lrb_template_dispatch(rocsparse_handle } uint32_t grid_size = static_cast(rocsparse::min(capped, count)); - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (csrmvn_lrb_long_rows_kernel), - grid_size, - block_size, - 0, - stream, - conj, - m, - nnz, - count, - info->lrb.wg_flags, - static_cast(info->lrb.rows_bins), - static_cast(info->lrb.n_rows_bins), - j, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), - csr_row_ptr, - csr_col_ind, - csr_val, - x, - ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), - y, - num_extra, - gamma_device_array, - z_array, - descr->base, - handle->pointer_mode == rocsparse_pointer_mode_host); + const auto launch = [&](auto grid_stride) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (csrmvn_lrb_long_rows_kernel), + grid_size, + block_size, + 0, + stream, + conj, + m, + nnz, + count, + info->lrb.wg_flags, + static_cast(info->lrb.rows_bins), + static_cast(info->lrb.n_rows_bins), + j, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, alpha_device_host), + csr_row_ptr, + csr_col_ind, + csr_val, + x, + ROCSPARSE_DEVICE_HOST_SCALAR_ARGS(handle, beta_device_host), + y, + num_extra, + gamma_device_array, + z_array, + descr->base, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + }; + + if(grid_size < count) + { + RETURN_IF_ROCSPARSE_ERROR(launch(std::true_type{})); + } + else + { + RETURN_IF_ROCSPARSE_ERROR(launch(std::false_type{})); + } } } } From 520dca969697998bd336011e4dcfb87065f1af36 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Thu, 1 Oct 2026 22:20:53 +0000 Subject: [PATCH 13/15] refactor(rocsparse): cache grid-stride start index and stride in csrmv LRB kernels (AISPARSE-662) Hoist the 64-bit start index of the grid-stride loops in partial_scale_y_device and the GRID_STRIDE path of csrmvn_lrb_medium_rows_warp_reduce into a named const, next to the stride, so each loop header reads `for(idx = gid; idx < n; idx += stride)`. No functional change; gfx942 ISA is identical. --- .../library/src/level2/csrmv_device.h | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index ac059501545d..b8df47e31c96 100644 --- a/projects/rocsparse/library/src/level2/csrmv_device.h +++ b/projects/rocsparse/library/src/level2/csrmv_device.h @@ -182,31 +182,30 @@ namespace rocsparse // The launch grid is clamped to the device's maximum grid size, so iterate with a // grid-stride loop to cover every element that needs scaling. - const int64_t stride = static_cast(WG_SIZE) * hipGridDim_x; - for(int64_t gid = static_cast(hipBlockIdx_x) * WG_SIZE + hipThreadIdx_x; - gid < required_threads; - gid += stride) + const int64_t gid = static_cast(hipBlockIdx_x) * WG_SIZE + hipThreadIdx_x; + const int64_t stride = static_cast(hipGridDim_x) * WG_SIZE; + for(int64_t idx = gid; idx < required_threads; idx += stride) { - if(gid < first_row) + if(idx < first_row) { if(beta == static_cast(0)) { - y[gid] = static_cast(0); + y[idx] = static_cast(0); } else { - y[gid] *= beta; + y[idx] *= beta; } } else { if(beta == static_cast(0)) { - y[last_row + (gid - first_row)] = static_cast(0); + y[last_row + (idx - first_row)] = static_cast(0); } else { - y[last_row + (gid - first_row)] *= beta; + y[last_row + (idx - first_row)] *= beta; } } } @@ -1082,9 +1081,9 @@ namespace rocsparse if constexpr(GRID_STRIDE) { // This kernel uses no block-wide barriers, so divergent iteration counts are safe. - const int64_t warps_per_grid = static_cast(BLOCKSIZE / WF_SIZE) * hipGridDim_x; - for(int64_t gid = static_cast(BLOCKSIZE / WF_SIZE) * bid + wid; gid < count; - gid += warps_per_grid) + const int64_t gwid = static_cast(bid) * (BLOCKSIZE / WF_SIZE) + wid; + const int64_t stride = static_cast(hipGridDim_x) * (BLOCKSIZE / WF_SIZE); + for(int64_t gid = gwid; gid < count; gid += stride) { process_row(gid); } From b11c2157ed61eeff8e7511a02df2302420f43f79 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Fri, 2 Oct 2026 14:25:05 +0000 Subject: [PATCH 14/15] test(rocsparse): cover int64 row pointers with int32 columns in csrmv LRB grid tests (AISPARSE-662) --- ...t_test_internal_level2_csrmv_lrb_grids.cpp | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp b/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp index 1589ae08d809..dcb1d3866536 100644 --- a/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp +++ b/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp @@ -605,6 +605,16 @@ TEST_F(CsrmvLrbGrids, clamped_grid_stride_f64_i64) check_lrb(handle, true); } +TEST_F(CsrmvLrbGrids, clamped_grid_stride_f32_i64_i32) +{ + check_lrb(handle, true); +} + +TEST_F(CsrmvLrbGrids, clamped_grid_stride_f64_i64_i32) +{ + check_lrb(handle, true); +} + // Control: the same problems with the real grid.x limit, which run the // straight-line variants. If this fails the clamped cases above prove nothing. TEST_F(CsrmvLrbGrids, unclamped_grid_matches_host) @@ -613,6 +623,8 @@ TEST_F(CsrmvLrbGrids, unclamped_grid_matches_host) check_lrb(handle, false); check_lrb(handle, false); check_lrb(handle, false); + check_lrb(handle, false); + check_lrb(handle, false); } // --------------------------------------------------------------------------- @@ -639,10 +651,22 @@ TEST_F(CsrmvAdaptiveGrids, partial_scale_y_clamped_f64_i64) check_adaptive_partial_scale_y(handle, true); } +TEST_F(CsrmvAdaptiveGrids, partial_scale_y_clamped_f32_i64_i32) +{ + check_adaptive_partial_scale_y(handle, true); +} + +TEST_F(CsrmvAdaptiveGrids, partial_scale_y_clamped_f64_i64_i32) +{ + check_adaptive_partial_scale_y(handle, true); +} + TEST_F(CsrmvAdaptiveGrids, unclamped_grid_matches_host) { check_adaptive_partial_scale_y(handle, false); check_adaptive_partial_scale_y(handle, false); check_adaptive_partial_scale_y(handle, false); check_adaptive_partial_scale_y(handle, false); + check_adaptive_partial_scale_y(handle, false); + check_adaptive_partial_scale_y(handle, false); } From eb15931f19d0430fd526012410e7e33dc34e4a94 Mon Sep 17 00:00:00 2001 From: kliegeois Date: Fri, 2 Oct 2026 14:25:06 +0000 Subject: [PATCH 15/15] test(rocsparse): boundary tests for dispatch_grid_stride_x variant selection (AISPARSE-662) --- .../clients/unittests/CMakeLists.txt | 2 + .../unit_test_internal_grid_dispatch.cpp | 157 ++++++++++++++++++ 2 files changed, 159 insertions(+) create mode 100644 projects/rocsparse/clients/unittests/unit_test_internal_grid_dispatch.cpp diff --git a/projects/rocsparse/clients/unittests/CMakeLists.txt b/projects/rocsparse/clients/unittests/CMakeLists.txt index 58f701597613..9a06018ab196 100644 --- a/projects/rocsparse/clients/unittests/CMakeLists.txt +++ b/projects/rocsparse/clients/unittests/CMakeLists.txt @@ -148,6 +148,8 @@ set(ROCSPARSE_UNIT_TEST_DEVICE_SOURCES unit_test_internal_collectives_readfirstlane.cpp # Per-axis grid-size helpers in rocsparse_grid.hpp (append-only leaf) unit_test_internal_grid.cpp + # dispatch_grid_stride_x variant selection in rocsparse_grid.hpp (append-only leaf) + unit_test_internal_grid_dispatch.cpp unit_test_internal_collectives_boost.cpp unit_test_internal_infostructs.cpp #10156: distinct anchor to avoid stacked-branch merge conflicts diff --git a/projects/rocsparse/clients/unittests/unit_test_internal_grid_dispatch.cpp b/projects/rocsparse/clients/unittests/unit_test_internal_grid_dispatch.cpp new file mode 100644 index 000000000000..5b192326519e --- /dev/null +++ b/projects/rocsparse/clients/unittests/unit_test_internal_grid_dispatch.cpp @@ -0,0 +1,157 @@ +/*! \file */ +/* ************************************************************************ + * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights Reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + * + * ************************************************************************ */ + +// +// Unit tests for rocsparse::dispatch_grid_stride_x in rocsparse_grid.hpp. +// +// dispatch_grid_stride_x picks the straight-line variant (std::false_type) +// while count fits the grid.x cap and the grid-stride variant (std::true_type) +// once it does not. These tests check that choice and the grid it passes at +// the boundary, count == cap and count == cap + 1, against both caps: the +// derived (2^32 - 1) / block_size bound, which binds at block size 32 and up, +// and maxGridSize[0], which binds at block size 1. No kernel is launched. +// +#include "unit_test_utils.hpp" + +#include "rocsparse_grid.hpp" + +#include + +#include +#include +#include + +using namespace rocsparse_ut; + +namespace +{ + constexpr int64_t BLOCK_SIZES[] = {1, 32, 64, 128, 256, 512, 1024}; + + struct Dispatch + { + int calls = 0; + bool grid_stride = false; + uint32_t grid = 0; + }; + + template + Dispatch dispatch(rocsparse_handle handle, J count, int64_t block_size) + { + Dispatch d; + rocsparse_status status = rocsparse::dispatch_grid_stride_x( + handle, count, block_size, [&](auto grid_stride, uint32_t grid) { + ++d.calls; + d.grid_stride = decltype(grid_stride)::value; + d.grid = grid; + return rocsparse_status_success; + }); + EXPECT_EQ(status, rocsparse_status_success); + return d; + } + + int64_t grid_x_cap(rocsparse_handle handle, int64_t block_size) + { + return std::min(static_cast(handle->properties.maxGridSize[0]), + rocsparse::dispatch_limit_x(block_size)); + } +} + +class internal_grid_dispatch : public HandleTest +{ +}; + +// count == cap launches the straight-line variant on the full grid; cap + 1 +// launches the grid-stride variant on the clamped grid. +TEST_F(internal_grid_dispatch, boundary) +{ + for(const int64_t bs : BLOCK_SIZES) + { + const int64_t cap = grid_x_cap(handle, bs); + + const Dispatch at = dispatch(handle, cap, bs); + EXPECT_EQ(at.calls, 1) << "block size " << bs; + EXPECT_FALSE(at.grid_stride) << "block size " << bs << ", count " << cap; + EXPECT_EQ(static_cast(at.grid), cap) << "block size " << bs; + + const Dispatch over = dispatch(handle, cap + 1, bs); + EXPECT_EQ(over.calls, 1) << "block size " << bs; + EXPECT_TRUE(over.grid_stride) << "block size " << bs << ", count " << cap + 1; + EXPECT_EQ(static_cast(over.grid), cap) << "block size " << bs; + } +} + +// The derived bound binds below maxGridSize[0] at 256 threads, and the device +// cap binds at 1 thread. +TEST_F(internal_grid_dispatch, boundary_concrete_caps) +{ + ASSERT_GT(static_cast(handle->properties.maxGridSize[0]), 16777215); + EXPECT_EQ(grid_x_cap(handle, 256), 16777215); + EXPECT_EQ(grid_x_cap(handle, 1), static_cast(handle->properties.maxGridSize[0])); + + EXPECT_FALSE(dispatch(handle, int32_t(16777215), 256).grid_stride); + EXPECT_TRUE(dispatch(handle, int32_t(16777216), 256).grid_stride); +} + +// The choice does not depend on the count type. A 32-bit count crosses the +// bound as readily as a 64-bit one. +TEST_F(internal_grid_dispatch, count_types) +{ + constexpr int64_t bs = 1024; + const int64_t cap = grid_x_cap(handle, bs); + + EXPECT_FALSE(dispatch(handle, static_cast(cap), bs).grid_stride); + EXPECT_TRUE(dispatch(handle, static_cast(cap + 1), bs).grid_stride); + EXPECT_FALSE(dispatch(handle, static_cast(cap), bs).grid_stride); + EXPECT_TRUE(dispatch(handle, static_cast(cap + 1), bs).grid_stride); + EXPECT_FALSE(dispatch(handle, static_cast(cap), bs).grid_stride); + EXPECT_TRUE(dispatch(handle, static_cast(cap + 1), bs).grid_stride); + EXPECT_FALSE(dispatch(handle, static_cast(cap), bs).grid_stride); + EXPECT_TRUE(dispatch(handle, static_cast(cap + 1), bs).grid_stride); +} + +// Small counts and a zero count stay on the straight-line variant with the +// grid equal to count. +TEST_F(internal_grid_dispatch, small_counts) +{ + for(const int64_t bs : BLOCK_SIZES) + { + for(const int64_t count : {int64_t(0), int64_t(1), int64_t(1000)}) + { + const Dispatch d = dispatch(handle, count, bs); + EXPECT_FALSE(d.grid_stride) << "block size " << bs << ", count " << count; + EXPECT_EQ(static_cast(d.grid), count) << "block size " << bs; + } + } +} + +// The launcher's status is returned to the caller. +TEST_F(internal_grid_dispatch, returns_launch_status) +{ + EXPECT_EQ(rocsparse::dispatch_grid_stride_x( + handle, + int64_t(1), + 256, + [](auto, uint32_t) { return rocsparse_status_internal_error; }), + rocsparse_status_internal_error); +}