diff --git a/projects/rocsparse/clients/unittests/CMakeLists.txt b/projects/rocsparse/clients/unittests/CMakeLists.txt index 40cd19fed61d..52b65be7ea22 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 @@ -203,6 +205,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_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); +} 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..dcb1d3866536 --- /dev/null +++ b/projects/rocsparse/clients/unittests/unit_test_internal_level2_csrmv_lrb_grids.cpp @@ -0,0 +1,672 @@ +/*! \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. 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. +// 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); +} + +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) +{ + check_lrb(handle, false); + check_lrb(handle, false); + 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, 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); +} diff --git a/projects/rocsparse/library/src/level2/csrmv_device.h b/projects/rocsparse/library/src/level2/csrmv_device.h index a1c90fe8b9ba..b8df47e31c96 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 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) { - return; - } - - if(gid < first_row) - { - if(beta == static_cast(0)) + if(idx < first_row) { - y[gid] = static_cast(0); - } - else - { - y[gid] *= beta; - } - } - else - { - if(beta == static_cast(0)) - { - y[last_row + (gid - first_row)] = static_cast(0); + if(beta == static_cast(0)) + { + y[idx] = static_cast(0); + } + else + { + y[idx] *= beta; + } } else { - y[last_row + (gid - first_row)] *= beta; + if(beta == static_cast(0)) + { + y[last_row + (idx - first_row)] = static_cast(0); + } + else + { + y[last_row + (idx - first_row)] *= beta; + } } } } @@ -725,8 +726,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 (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) - { - uint32_t row_idx = wg_row_start + ((base_idx + lid) >> bin_id); - if(row_idx < wg_row_start + wg_num_rows) + // 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]; + 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[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; + } + } + } + __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; - if(col_idx_in_row < row_len) + for(uint32_t idx = 0; idx < (1 << bin_id); idx++) { - 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]; + acc += partialSums[lds_start_idx + idx]; } - else + + if(beta != 0) { - // lds <- 0 - partialSums[lds_idx] = 0; + acc = rocsparse::fma(beta, y[row_id], acc); } - } - } - __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; + // 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); + } + } - for(uint32_t idx = 0; idx < (1 << bin_id); idx++) - { - acc += partialSums[lds_start_idx + idx]; + y[row_id] = acc; } + }; - if(beta != 0) + 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) { - acc = rocsparse::fma(beta, y[row_id], acc); - } + 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))); - // 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); - } + // The next iteration overwrites partialSums, which this reduction reads. + __syncthreads(); } - - y[row_id] = acc; + } + 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); } } @@ -834,6 +866,7 @@ namespace rocsparse // dynamic LDS allocation approach would blow up size requirements beyond reasonable bounds. template > 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) - { - 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; + // 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); - 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 + // 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 auto row_idx = wg_row_start + ((base_idx + lid) >> bin_id); + if(row_idx < wg_row_start + wg_num_rows) { - // lds <- 0 - partialSums[lds_idx] = 0; + 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 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; + 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 = count) - { - return; - } + 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; - const J bin_start = n_rows_bins[bin_id]; - const J row = rows_bins[bin_start + gid]; + for(I j = vecStart + lid; j < vecEnd; j += WF_SIZE) + { + temp_sum = rocsparse::fma(alpha * rocsparse::conj_val(csr_val[j], conj), + x[csr_col_ind[j] - idx_base], + 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; + // Obtain row sum using parallel warp reduction + temp_sum = rocsparse::wfreduce_sum(temp_sum); - for(I j = vecStart + lid; j < vecEnd; j += WF_SIZE) - { - temp_sum = rocsparse::fma(alpha * rocsparse::conj_val(csr_val[j], conj), - x[csr_col_ind[j] - idx_base], - temp_sum); - } + if(lid == WF_SIZE - 1) + { + if(beta != static_cast(0)) + { + temp_sum = rocsparse::fma(beta, y[row], temp_sum); + } - // Obtain row sum using parallel warp reduction - temp_sum = rocsparse::wfreduce_sum(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; + } + }; - if(lid == WF_SIZE - 1) + if constexpr(GRID_STRIDE) { - if(beta != static_cast(0)) + // This kernel uses no block-wide barriers, so divergent iteration counts are safe. + 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) { - temp_sum = rocsparse::fma(beta, y[row], temp_sum); + process_row(gid); } + } + else + { + const int gid = (BLOCKSIZE / WF_SIZE) * bid + wid; - // Add contributions from all extra vectors - for(rocsparse_int i = 0; i < num_extra; ++i) + if(gid >= count) { - 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); - } + return; } - y[row] = temp_sum; + process_row(gid); } } // "Vector" case a la CSR-Adaptive using one block per row template 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,14 +1130,12 @@ 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 @@ -1068,60 +1144,81 @@ 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); - } - - partialSums[lid] = temp_sum; + // 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; + // 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; + 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 (0); + // Each workgroup computes exactly one row. + 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)); - // 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); + 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_adaptive.cpp b/projects/rocsparse/library/src/level2/rocsparse_csrmv_template_adaptive.cpp index d309b368920e..d275dccf19ed 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" @@ -767,7 +768,10 @@ 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 the grid; the kernel grid-strides over all required threads. + dim3 scale_blocks(rocsparse::get_grid_size_x( + 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 d5ecf87dcc80..7fdd1e6d2663 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" @@ -200,13 +201,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 = static_cast(csrmv_info->lrb.nRowsBins[j]) * num_wgs_per_row; max_required_grid = rocsparse::max(grid_size, max_required_grid); } @@ -236,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, @@ -260,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, @@ -280,7 +292,14 @@ namespace rocsparse } } - template + template ROCSPARSE_KERNEL(WG_SIZE) void csrmvn_lrb_short_rows_2_kernel(bool conj, I nnz, @@ -306,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, @@ -386,6 +407,7 @@ namespace rocsparse } template 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, @@ -416,30 +439,39 @@ namespace rocsparse ROCSPARSE_DEVICE_HOST_SCALAR_GET(beta); if(alpha != 0 || beta != 1 || num_extra > 0) { - rocsparse::csrmvn_lrb_medium_rows_device(conj, - nnz, - 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, - 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, @@ -461,25 +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, - 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); } } } @@ -569,63 +610,81 @@ 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); + // 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 + = (static_cast(info->lrb.nRowsBins[j]) - 1) / block_size + 1; - 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 { 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); + // 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 + = (static_cast(info->lrb.nRowsBins[j]) - 1) / rows_per_wg + 1; - 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; + })); } } } @@ -646,119 +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; - uint32_t grid_size = (info->lrb.nRowsBins[j] - 1) / (mr_block / 32) + 1; - 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 { - uint32_t grid_size + int64_t num_wgs = (info->lrb.nRowsBins[j] - 1) / (256 / handle->wavefront_size) + 1; - 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 { - uint32_t grid_size = info->lrb.nRowsBins[j]; // One WG per row + 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, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + count, WG_SIZE_WAVE32, - 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_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, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + count, WG_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_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; + })); } } } @@ -776,33 +873,62 @@ 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; - - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (csrmvn_lrb_long_rows_kernel), - grid_size, - block_size, - 0, - stream, - conj, - m, - nnz, - 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); + + // 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 = 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 + // 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) + { + capped = num_wgs_per_row; + } + uint32_t grid_size = static_cast(rocsparse::min(capped, count)); + + 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{})); + } } } }