diff --git a/projects/rocsparse/clients/unittests/CMakeLists.txt b/projects/rocsparse/clients/unittests/CMakeLists.txt index 924a650e44a1..46d086a4877e 100644 --- a/projects/rocsparse/clients/unittests/CMakeLists.txt +++ b/projects/rocsparse/clients/unittests/CMakeLists.txt @@ -160,6 +160,12 @@ set(ROCSPARSE_UNIT_TEST_DEVICE_SOURCES unit_test_internal_hostblocks_itilu0.cpp #10153: distinct anchor to avoid stacked-branch merge conflicts unit_test_host_extra_bsrgemm.cpp + # --- csrgemm / bsrgemm calc grid.x clamps [AISPARSE-677] ----------------- + # Distinct anchor (after the bsrgemm host tests) so this leaf does not conflict + # with the sibling grid-clamp branches that append at the end of the list. + # Drives the public rocsparse_spgemm entry point with + # handle->properties.maxGridSize[0] shrunk; no compile-in TUs. + unit_test_internal_extra_gemm_calc_grids.cpp # Residual in-scope internal-routine coverage: device __global__ wrappers for # collective/atomic/hash device helpers, plus host-path primitive helpers. # unit_test_internal_primitives_host_extras.cpp additionally needs the diff --git a/projects/rocsparse/clients/unittests/unit_test_grid_clamp.hpp b/projects/rocsparse/clients/unittests/unit_test_grid_clamp.hpp new file mode 100644 index 000000000000..0ce9d85534f8 --- /dev/null +++ b/projects/rocsparse/clients/unittests/unit_test_grid_clamp.hpp @@ -0,0 +1,72 @@ +/*! \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. + * + * ************************************************************************ */ + +#pragma once + +// +// Test-only helper shared by the large-grid unit tests (AISPARSE-699/700). +// +// The launches hardened by those tickets clamp grid.x with +// rocsparse::get_grid_size_x, to min(maxGridSize[0], (2^32 - 1) / blockDim.x) +// blocks, and grid-stride over whatever the clamp drops. For a 256-thread block +// that clamp is 16,777,215 blocks. The tests shrink maxGridSize[0] for speed +// rather than allocate a problem that large, and run a small problem through +// the same code path. This is the AISPARSE-702 idiom. +// +// A new header rather than an addition to unit_test_utils.hpp: that file is +// shared with the sibling unit-test branches, and a new leaf file cannot +// conflict with them. +// + +// Internal handle definition: the complete _rocsparse_handle type is required to +// reach handle->properties. +#include "rocsparse_handle.hpp" + +namespace rocsparse_ut +{ + // Temporarily shrink the grid.x limit that the launches clamp against, and + // restore it on scope exit so a failed assertion cannot leak the override + // into the next test sharing the fixture's handle. + 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; + }; +} diff --git a/projects/rocsparse/clients/unittests/unit_test_internal_extra_gemm_calc_grids.cpp b/projects/rocsparse/clients/unittests/unit_test_internal_extra_gemm_calc_grids.cpp new file mode 100644 index 000000000000..f7f5e86da602 --- /dev/null +++ b/projects/rocsparse/clients/unittests/unit_test_internal_extra_gemm_calc_grids.cpp @@ -0,0 +1,844 @@ +/*! \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 csrgemm and bsrgemm calc grids (AISPARSE-677). +// +// FOCUS. csrgemm_intermediate_products, the csrgemm nnz / fill / symbolic / +// numeric bucket kernels (wf-per-row groups 0 and 1, block-per-row groups 2 to +// 9, the group 10 multipass fallback) and every bsrgemm fill bucket clamp grid.x +// with rocsparse::get_grid_size_x and grid-stride over the rows of their bucket. +// rocsparse::dispatch_grid_stride_x selects the GRID_STRIDE instantiation only +// when the clamp binds, so a normally sized problem never runs the looping +// variant at all. +// +// WHY NOT TEST THE REAL THRESHOLD. The clamp binds at (2^32 - 1) / 256 = +// 16,777,215 blocks, i.e. about 16.7M rows in one block-per-row bucket. These +// tests shrink handle->properties.maxGridSize[0] with ScopedMaxGridSizeX to 1, +// 3, 7 and 64 blocks instead, so a problem of a few thousand rows takes the +// clamped, looping path. This is the AISPARSE-702 idiom. +// +// WHAT MAKES EACH CASE LOAD-BEARING. The operands are built from row classes, +// each aimed at one bucket (see csr_classes and bsr_classes). Every non-empty +// bucket holds at least 8 rows, so limits 1, 3 and 7 clamp every bucketed +// launch. The wf-per-row buckets, csrgemm_intermediate_products and the +// block_dim > 8 bsrgemm launch are more than 64 blocks wide, so limit 64 clamps +// those too. Each test checks the bucket populations on the host reference +// before it trusts a pass. +// +// EXACT ARITHMETIC. All values, alpha and beta are small integers, so every +// entry of C is exactly representable in float and does not depend on the +// order in which the hash tables accumulate. That allows an exact comparison +// against the host reference and against the unclamped run. +// +// TARGET: rocsparse-unit-test-device. The tests drive the public rocsparse_spgemm +// entry point and need the complete handle type to shrink the grid limit. +// +#include "unit_test_utils.hpp" + +// ScopedMaxGridSizeX: shrinks handle->properties.maxGridSize[0], the limit +// get_grid_size_x clamps grid.x against. +#include "unit_test_grid_clamp.hpp" + +#include "rocsparse.h" + +#include +#include +#include +#include +#include +#include +#include + +using namespace rocsparse_ut; + +namespace +{ + // grid.x limits the tests shrink maxGridSize[0] to. 0 means the device limit. + constexpr int clamped_limits[] = {1, 3, 7, 64}; + + constexpr double alpha_value = 2.0; + constexpr double beta_value = 3.0; + + // Host sparse matrix. For BSR, m, n, ptr and ind are at block level and val + // holds block_dim * block_dim entries per block in the block direction. + struct HostSparse + { + int64_t m = 0; + int64_t n = 0; + int64_t block_dim = 1; + std::vector ptr; + std::vector ind; + std::vector val; + + int64_t nnz() const + { + return ptr.empty() ? 0 : ptr.back(); + } + }; + + // Rows of A that each pick a_per_row distinct rows of B from a private pool of + // B rows with b_row_len random columns. Row nnz of C is then about the union + // of a_per_row random column sets of length b_row_len. + struct RowClass + { + int64_t rows; + int64_t a_per_row; + int64_t b_row_len; + }; + + // n = 50000 columns. The comment on each class is the csrgemm fill bucket + // (row nnz of C, D included) it lands in; the nnz stage buckets by + // intermediate products and sees classes 0 to 9 in its lower buckets and + // class 10 (5 * 15000 = 75000 products) in its group 10. + constexpr int64_t csr_n = 50000; + constexpr RowClass csr_classes[] = { + {2400, 2, 4}, // <= 16, wf-per-row group 0, 75 blocks of 32 rows + {1200, 3, 7}, // 17 - 32, wf-per-row group 1, 75 blocks of 16 rows + {72, 4, 25}, // 33 - 256, block-per-row group 2 + {72, 4, 90}, // 257 - 512, group 3 + {72, 4, 190}, // 513 - 1024, group 4 + {72, 4, 380}, // 1025 - 2048, group 5 + {72, 4, 770}, // 2049 - 4096, group 6 + {12, 4, 1600}, // 4097 - 8192, group 7 (group 10 if shared memory is short) + {12, 4, 3300}, // 8193 - 16384, group 8 (likewise) + {12, 4, 7000}, // 16385 - 32768, group 9 (likewise) + {8, 5, 15000}, // > 32768, group 10 multipass + }; + + // Rows of C with at most 16 entries, so csrgemm skips the grouping and runs + // one wf-per-row launch over all rows: 20000 rows are 625 blocks. + constexpr int64_t short_n = 20000; + constexpr RowClass short_classes[] = {{20000, 4, 4}}; + + // nb = 2000 block columns. The comment is the bsrgemm bucket (nnzb per row + // of C, D included). block_dim > 8 skips the grouping and runs all 2620 block + // rows through one block-per-row launch. + constexpr int64_t bsr_nb = 2000; + constexpr RowClass bsr_classes[] = { + {1100, 2, 3}, // <= 8, group 0 + {1100, 3, 4}, // 9 - 16, group 1 + {100, 3, 8}, // 17 - 32, group 2 + {100, 4, 12}, // 33 - 64, group 3 + {100, 4, 25}, // 65 - 128, group 4 + {100, 4, 50}, // 129 - 256, group 5 + {12, 4, 100}, // 257 - 512, group 6 + {8, 4, 200}, // > 512, group 7 + }; + + // Entries of D per row, at block level for BSR. + constexpr int64_t d_per_row = 2; + + // Distinct sorted random columns in [0, n). + std::vector random_columns(std::mt19937& rng, int64_t n, int64_t count) + { + std::vector cols; + std::vector used(n, 0); + std::uniform_int_distribution col(0, n - 1); + while(static_cast(cols.size()) < count) + { + const int64_t c = col(rng); + if(!used[c]) + { + used[c] = 1; + cols.push_back(c); + } + } + std::sort(cols.begin(), cols.end()); + return cols; + } + + // Small nonzero integers so every product and sum is exact. + double random_value(std::mt19937& rng) + { + static const double values[] = {-3.0, -2.0, -1.0, 1.0, 2.0, 3.0}; + return values[std::uniform_int_distribution(0, 5)(rng)]; + } + + void append_row(HostSparse& a, const std::vector& cols, std::mt19937& rng) + { + const int64_t block_size = a.block_dim * a.block_dim; + for(const int64_t c : cols) + { + a.ind.push_back(c); + for(int64_t e = 0; e < block_size; ++e) + { + a.val.push_back(random_value(rng)); + } + } + a.ptr.push_back(static_cast(a.ind.size())); + } + + struct Operands + { + HostSparse a; + HostSparse b; + HostSparse d; + }; + + // Build A (m x k), B (k x n) and D (m x n) from the row classes. The rows of + // A are shuffled so every bucket's rows are scattered and the permutation + // that groups them is far from the identity. + template + Operands make_operands(const RowClass (&classes)[N], int64_t n, int64_t block_dim, bool with_d) + { + std::mt19937 rng(677); + + Operands op; + op.b.n = n; + op.b.block_dim = block_dim; + op.b.ptr.push_back(0); + + // One pool of B rows per class. + std::vector pool_begin(N); + std::vector pool_size(N); + for(size_t c = 0; c < N; ++c) + { + pool_begin[c] = op.b.m; + pool_size[c] = std::max(8, 4 * classes[c].a_per_row); + for(int64_t r = 0; r < pool_size[c]; ++r) + { + append_row(op.b, random_columns(rng, n, classes[c].b_row_len), rng); + ++op.b.m; + } + } + + std::vector row_class; + for(size_t c = 0; c < N; ++c) + { + row_class.insert(row_class.end(), classes[c].rows, c); + } + std::shuffle(row_class.begin(), row_class.end(), rng); + + op.a.m = static_cast(row_class.size()); + op.a.n = op.b.m; + op.a.block_dim = block_dim; + op.a.ptr.push_back(0); + + op.d.m = op.a.m; + op.d.n = n; + op.d.block_dim = block_dim; + op.d.ptr.push_back(0); + + for(const size_t c : row_class) + { + std::vector picks = random_columns(rng, pool_size[c], classes[c].a_per_row); + for(int64_t& p : picks) + { + p += pool_begin[c]; + } + append_row(op.a, picks, rng); + append_row( + op.d, with_d ? random_columns(rng, n, d_per_row) : std::vector(), rng); + } + return op; + } + + int64_t block_entry(rocsparse_direction dir, int64_t block_dim, int64_t r, int64_t c) + { + return (dir == rocsparse_direction_row) ? r * block_dim + c : c * block_dim + r; + } + + // C = alpha * A * B + beta * D on the host, with the columns of each row + // sorted. The pattern is the structural union, as rocsparse computes it. + HostSparse host_gemm(const Operands& op, bool with_d, rocsparse_direction dir) + { + const HostSparse& a = op.a; + const HostSparse& b = op.b; + const HostSparse& d = op.d; + const int64_t bd = a.block_dim; + const int64_t bs = bd * bd; + + HostSparse c; + c.m = a.m; + c.n = b.n; + c.block_dim = bd; + c.ptr.push_back(0); + + std::vector slot(b.n, -1); + std::vector cols; + std::vector acc; + + for(int64_t i = 0; i < a.m; ++i) + { + cols.clear(); + acc.clear(); + + auto slot_of = [&](int64_t col) { + if(slot[col] < 0) + { + slot[col] = static_cast(cols.size()); + cols.push_back(col); + acc.resize(acc.size() + bs, 0.0); + } + return slot[col]; + }; + + for(int64_t pa = a.ptr[i]; pa < a.ptr[i + 1]; ++pa) + { + const int64_t k = a.ind[pa]; + const double* a_blk = &a.val[pa * bs]; + for(int64_t pb = b.ptr[k]; pb < b.ptr[k + 1]; ++pb) + { + const double* b_blk = &b.val[pb * bs]; + double* c_blk = &acc[slot_of(b.ind[pb]) * bs]; + for(int64_t r = 0; r < bd; ++r) + { + for(int64_t cc = 0; cc < bd; ++cc) + { + double sum = 0.0; + for(int64_t t = 0; t < bd; ++t) + { + sum += a_blk[block_entry(dir, bd, r, t)] + * b_blk[block_entry(dir, bd, t, cc)]; + } + c_blk[block_entry(dir, bd, r, cc)] += alpha_value * sum; + } + } + } + } + + if(with_d) + { + for(int64_t pd = d.ptr[i]; pd < d.ptr[i + 1]; ++pd) + { + double* c_blk = &acc[slot_of(d.ind[pd]) * bs]; + for(int64_t e = 0; e < bs; ++e) + { + c_blk[e] += beta_value * d.val[pd * bs + e]; + } + } + } + + std::vector sorted = cols; + std::sort(sorted.begin(), sorted.end()); + for(const int64_t col : sorted) + { + c.ind.push_back(col); + c.val.insert(c.val.end(), &acc[slot[col] * bs], &acc[slot[col] * bs] + bs); + } + for(const int64_t col : cols) + { + slot[col] = -1; + } + c.ptr.push_back(static_cast(c.ind.size())); + } + return c; + } + + // Rows of m whose entry count lies in [lo, hi]. + int64_t rows_with_nnz(const HostSparse& m, int64_t lo, int64_t hi) + { + int64_t count = 0; + for(int64_t i = 0; i < m.m; ++i) + { + const int64_t nnz = m.ptr[i + 1] - m.ptr[i]; + count += (nnz >= lo && nnz <= hi) ? 1 : 0; + } + return count; + } + + // Empty if equal, otherwise where the first difference is. + std::string first_difference(const HostSparse& got, const HostSparse& want) + { + std::ostringstream msg; + if(got.ptr != want.ptr) + { + size_t i = 0; + while(i < got.ptr.size() && i < want.ptr.size() && got.ptr[i] == want.ptr[i]) + { + ++i; + } + msg << "row pointer differs at " << i; + } + else if(got.ind != want.ind) + { + size_t i = 0; + while(got.ind[i] == want.ind[i]) + { + ++i; + } + msg << "column index " << i << " is " << got.ind[i] << ", want " << want.ind[i]; + } + else if(got.val != want.val) + { + size_t i = 0; + while(got.val[i] == want.val[i]) + { + ++i; + } + msg << "value " << i << " is " << got.val[i] << ", want " << want.val[i]; + } + return msg.str(); + } + + template + struct DeviceSparse + { + std::unique_ptr> ptr; + std::unique_ptr> ind; + std::unique_ptr> val; + + explicit DeviceSparse(const HostSparse& h) + { + ptr.reset(new device_vector(std::vector(h.ptr.begin(), h.ptr.end()))); + ind.reset(new device_vector(std::vector(h.ind.begin(), h.ind.end()))); + val.reset(new device_vector(std::vector(h.val.begin(), h.val.end()))); + } + }; + + // Owns the spgemm descriptors so each one created is destroyed on every + // return path. + struct SpgemmDescrs + { + rocsparse_spmat_descr a = nullptr; + rocsparse_spmat_descr b = nullptr; + rocsparse_spmat_descr c = nullptr; + rocsparse_spmat_descr d = nullptr; + + SpgemmDescrs() = default; + + SpgemmDescrs(const SpgemmDescrs&) = delete; + + SpgemmDescrs& operator=(const SpgemmDescrs&) = delete; + + ~SpgemmDescrs() + { + for(rocsparse_spmat_descr m : {a, b, c, d}) + { + if(m != nullptr) + { + (void)rocsparse_destroy_spmat_descr(m); + } + } + } + }; + + enum class Flow + { + // buffer_size, nnz, compute: csrgemm_nnz_calc and csrgemm_calc (bsrgemm_calc). + nnz_compute, + // buffer_size, nnz, symbolic, numeric: csrgemm_symbolic_calc and + // csrgemm_numeric_calc. CSR only. + symbolic_numeric + }; + + template + rocsparse_status create_descr(rocsparse_spmat_descr* descr, + const HostSparse& h, + int64_t nnz, + void* ptr, + void* ind, + void* val, + bool bsr, + rocsparse_direction dir) + { + if(bsr) + { + return rocsparse_create_bsr_descr(descr, + h.m, + h.n, + nnz, + dir, + h.block_dim, + ptr, + ind, + val, + it_of(), + it_of(), + rocsparse_index_base_zero, + dt_of()); + } + return rocsparse_create_csr_descr(descr, + h.m, + h.n, + nnz, + ptr, + ind, + val, + it_of(), + it_of(), + rocsparse_index_base_zero, + dt_of()); + } + + // Run C = alpha * A * B (+ beta * D) through rocsparse_spgemm with grid.x + // limited to `limit` blocks (0 keeps the device limit) for every stage, and + // return C. Records a gtest failure and returns an empty matrix on error. + template + HostSparse device_gemm(rocsparse_handle handle, + const Operands& op, + bool with_d, + bool bsr, + rocsparse_direction dir, + Flow flow, + int limit) + { + const T alpha = static_cast(alpha_value); + const T beta = static_cast(beta_value); + + DeviceSparse da(op.a); + DeviceSparse db(op.b); + DeviceSparse dd(op.d); + device_vector c_ptr(static_cast(op.a.m + 1)); + if(da.ptr->ptr == nullptr || da.ind->ptr == nullptr || da.val->ptr == nullptr + || db.ptr->ptr == nullptr || db.ind->ptr == nullptr || db.val->ptr == nullptr + || dd.ptr->ptr == nullptr || c_ptr.ptr == nullptr) + { + ADD_FAILURE() << "device allocation failed"; + return {}; + } + + HostSparse c_shape; + c_shape.m = op.a.m; + c_shape.n = op.b.n; + c_shape.block_dim = op.a.block_dim; + + SpgemmDescrs descrs; + if(create_descr( + &descrs.a, op.a, op.a.nnz(), da.ptr->ptr, da.ind->ptr, da.val->ptr, bsr, dir) + != rocsparse_status_success + || create_descr( + &descrs.b, op.b, op.b.nnz(), db.ptr->ptr, db.ind->ptr, db.val->ptr, bsr, dir) + != rocsparse_status_success + || create_descr( + &descrs.d, op.d, op.d.nnz(), dd.ptr->ptr, dd.ind->ptr, dd.val->ptr, bsr, dir) + != rocsparse_status_success + || create_descr(&descrs.c, c_shape, 0, c_ptr.ptr, nullptr, nullptr, bsr, dir) + != rocsparse_status_success) + { + ADD_FAILURE() << "descriptor creation failed"; + return {}; + } + + std::unique_ptr clamp; + if(limit > 0) + { + clamp.reset(new ScopedMaxGridSizeX(handle, limit)); + } + + std::unique_ptr> c_ind; + std::unique_ptr> c_val; + size_t buffer_size = 0; + void* buffer = nullptr; + + auto spgemm = [&](rocsparse_spgemm_stage stage) { + return rocsparse_spgemm(handle, + rocsparse_operation_none, + rocsparse_operation_none, + &alpha, + descrs.a, + descrs.b, + with_d ? &beta : nullptr, + descrs.d, + descrs.c, + dt_of(), + rocsparse_spgemm_alg_default, + stage, + &buffer_size, + buffer); + }; + + rocsparse_status status = spgemm(rocsparse_spgemm_stage_buffer_size); + if(status == rocsparse_status_success + && hipMalloc(&buffer, std::max(buffer_size, 1)) != hipSuccess) + { + buffer = nullptr; + status = rocsparse_status_memory_error; + } + if(status == rocsparse_status_success) + { + status = spgemm(rocsparse_spgemm_stage_nnz); + } + + int64_t rows = 0; + int64_t cols = 0; + int64_t nnz = 0; + if(status == rocsparse_status_success) + { + status = rocsparse_spmat_get_size(descrs.c, &rows, &cols, &nnz); + } + if(status == rocsparse_status_success) + { + const int64_t bs = op.a.block_dim * op.a.block_dim; + c_ind.reset(new device_vector(static_cast(std::max(nnz, 1)))); + c_val.reset(new device_vector(static_cast(std::max(nnz * bs, 1)))); + if(c_ind->ptr == nullptr || c_val->ptr == nullptr) + { + status = rocsparse_status_memory_error; + } + else if(bsr) + { + status = rocsparse_bsr_set_pointers(descrs.c, c_ptr.ptr, c_ind->ptr, c_val->ptr); + } + else + { + status = rocsparse_csr_set_pointers(descrs.c, c_ptr.ptr, c_ind->ptr, c_val->ptr); + } + } + if(status == rocsparse_status_success) + { + if(flow == Flow::nnz_compute) + { + status = spgemm(rocsparse_spgemm_stage_compute); + } + else + { + status = spgemm(rocsparse_spgemm_stage_symbolic); + if(status == rocsparse_status_success) + { + status = spgemm(rocsparse_spgemm_stage_numeric); + } + } + } + if(status == rocsparse_status_success && hipDeviceSynchronize() != hipSuccess) + { + status = rocsparse_status_internal_error; + } + if(buffer != nullptr) + { + (void)hipFree(buffer); + } + if(status != rocsparse_status_success) + { + ADD_FAILURE() << "rocsparse_spgemm failed with status " << status << " at limit " + << limit; + return {}; + } + + const std::vector h_ptr = to_host(c_ptr); + const std::vector h_ind = to_host(c_ind->ptr, static_cast(nnz)); + const std::vector h_val + = to_host(c_val->ptr, static_cast(nnz * op.a.block_dim * op.a.block_dim)); + + c_shape.ptr.assign(h_ptr.begin(), h_ptr.end()); + c_shape.ind.assign(h_ind.begin(), h_ind.end()); + c_shape.val.assign(h_val.begin(), h_val.end()); + return c_shape; + } + + // Unclamped control against the host reference, then every clamped limit + // against both the host reference and the unclamped result. + template + void check_all_limits(rocsparse_handle handle, + const Operands& op, + const HostSparse& want, + bool with_d, + bool bsr, + rocsparse_direction dir, + Flow flow) + { + const HostSparse unclamped = device_gemm(handle, op, with_d, bsr, dir, flow, 0); + const std::string control = first_difference(unclamped, want); + ASSERT_TRUE(control.empty()) << "unclamped control: " << control; + + for(const int limit : clamped_limits) + { + const HostSparse got = device_gemm(handle, op, with_d, bsr, dir, flow, limit); + const std::string vs_host = first_difference(got, want); + EXPECT_TRUE(vs_host.empty()) + << "grid.x clamped to " << limit << " blocks, against the host: " << vs_host; + const std::string vs_unclamped = first_difference(got, unclamped); + EXPECT_TRUE(vs_unclamped.empty()) + << "grid.x clamped to " << limit << " blocks, against unclamped: " << vs_unclamped; + } + } + + // The csrgemm fill buckets of csr_classes must each hold enough rows that + // the clamp binds; see the file header. + void expect_csr_buckets(const HostSparse& c) + { + EXPECT_GT(rows_with_nnz(c, 0, 16), 64 * 32) << "wf-per-row group 0"; + EXPECT_GT(rows_with_nnz(c, 17, 32), 64 * 16) << "wf-per-row group 1"; + const int64_t bounds[] = {32, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768}; + for(size_t g = 0; g + 1 < sizeof(bounds) / sizeof(bounds[0]); ++g) + { + EXPECT_GE(rows_with_nnz(c, bounds[g] + 1, bounds[g + 1]), 8) + << "block-per-row bucket " << bounds[g] + 1 << " - " << bounds[g + 1]; + } + EXPECT_GE(rows_with_nnz(c, 32769, c.n), 8) << "group 10"; + } + + void expect_bsr_buckets(const HostSparse& c) + { + const int64_t bounds[] = {0, 8, 16, 32, 64, 128, 256, 512}; + for(size_t g = 0; g + 1 < sizeof(bounds) / sizeof(bounds[0]); ++g) + { + EXPECT_GE(rows_with_nnz(c, bounds[g] + 1, bounds[g + 1]), 8) + << "bucket " << bounds[g] + 1 << " - " << bounds[g + 1]; + } + EXPECT_GE(rows_with_nnz(c, 513, c.n), 8) << "group 7"; + } + + template + void check_csrgemm_buckets(rocsparse_handle handle, bool with_d, Flow flow) + { + const Operands op = make_operands(csr_classes, csr_n, 1, with_d); + const HostSparse want = host_gemm(op, with_d, rocsparse_direction_row); + expect_csr_buckets(want); + check_all_limits(handle, op, want, with_d, false, rocsparse_direction_row, flow); + } + + template + void check_bsrgemm(rocsparse_handle handle, int64_t block_dim, rocsparse_direction dir) + { + const Operands op = make_operands(bsr_classes, bsr_nb, block_dim, true); + const HostSparse want = host_gemm(op, true, dir); + expect_bsr_buckets(want); + check_all_limits(handle, op, want, true, true, dir, Flow::nnz_compute); + } + + using GemmCalcGrids = HandleTest; +} + +// --------------------------------------------------------------------------- +// csrgemm, C = alpha * A * B + beta * D, nnz + compute stages. +// --------------------------------------------------------------------------- + +TEST_F(GemmCalcGrids, csrgemm_multadd_compute_f32_i32) +{ + check_csrgemm_buckets(handle, true, Flow::nnz_compute); +} + +TEST_F(GemmCalcGrids, csrgemm_multadd_compute_f64_i32) +{ + check_csrgemm_buckets(handle, true, Flow::nnz_compute); +} + +TEST_F(GemmCalcGrids, csrgemm_multadd_compute_f32_i64) +{ + check_csrgemm_buckets(handle, true, Flow::nnz_compute); +} + +TEST_F(GemmCalcGrids, csrgemm_multadd_compute_f64_i64) +{ + check_csrgemm_buckets(handle, true, Flow::nnz_compute); +} + +// 64-bit row pointers with 32-bit column indices. +TEST_F(GemmCalcGrids, csrgemm_multadd_compute_f64_i64_i32) +{ + check_csrgemm_buckets(handle, true, Flow::nnz_compute); +} + +// --------------------------------------------------------------------------- +// csrgemm, symbolic + numeric stages (csrgemm_symbolic_calc, csrgemm_numeric_calc). +// --------------------------------------------------------------------------- + +TEST_F(GemmCalcGrids, csrgemm_multadd_symbolic_numeric_f32_i32) +{ + check_csrgemm_buckets(handle, true, Flow::symbolic_numeric); +} + +TEST_F(GemmCalcGrids, csrgemm_multadd_symbolic_numeric_f64_i64) +{ + check_csrgemm_buckets(handle, true, Flow::symbolic_numeric); +} + +// --------------------------------------------------------------------------- +// csrgemm, C = alpha * A * B (mul only, beta == nullptr). +// --------------------------------------------------------------------------- + +TEST_F(GemmCalcGrids, csrgemm_mult_compute_f64_i32) +{ + check_csrgemm_buckets(handle, false, Flow::nnz_compute); +} + +TEST_F(GemmCalcGrids, csrgemm_mult_compute_f32_i64) +{ + check_csrgemm_buckets(handle, false, Flow::nnz_compute); +} + +TEST_F(GemmCalcGrids, csrgemm_mult_symbolic_numeric_f64_i32) +{ + check_csrgemm_buckets(handle, false, Flow::symbolic_numeric); +} + +// Every row of C has at most 16 entries, so csrgemm skips the grouping and runs +// a single wf-per-row launch of 625 blocks over all rows, with no permutation. +TEST_F(GemmCalcGrids, csrgemm_short_rows_single_group) +{ + const Operands op = make_operands(short_classes, short_n, 1, false); + const HostSparse want = host_gemm(op, false, rocsparse_direction_row); + ASSERT_EQ(rows_with_nnz(want, 0, 16), want.m); + check_all_limits( + handle, op, want, false, false, rocsparse_direction_row, Flow::nnz_compute); + check_all_limits( + handle, op, want, false, false, rocsparse_direction_row, Flow::symbolic_numeric); +} + +// --------------------------------------------------------------------------- +// bsrgemm, C = alpha * A * B + beta * D. block_dim 2 takes the 2x2 fill kernels, +// 4 the 3-4 path, 8 the 5-8 path, and 16 the ungrouped block-per-row multipass. +// --------------------------------------------------------------------------- + +TEST_F(GemmCalcGrids, bsrgemm_bd2_row_f32_i32) +{ + check_bsrgemm(handle, 2, rocsparse_direction_row); +} + +TEST_F(GemmCalcGrids, bsrgemm_bd2_column_f32_i32) +{ + check_bsrgemm(handle, 2, rocsparse_direction_column); +} + +TEST_F(GemmCalcGrids, bsrgemm_bd2_row_f64_i64) +{ + check_bsrgemm(handle, 2, rocsparse_direction_row); +} + +TEST_F(GemmCalcGrids, bsrgemm_bd2_column_f64_i64) +{ + check_bsrgemm(handle, 2, rocsparse_direction_column); +} + +TEST_F(GemmCalcGrids, bsrgemm_bd4_row_f64_i32) +{ + check_bsrgemm(handle, 4, rocsparse_direction_row); +} + +TEST_F(GemmCalcGrids, bsrgemm_bd4_column_f32_i64) +{ + check_bsrgemm(handle, 4, rocsparse_direction_column); +} + +TEST_F(GemmCalcGrids, bsrgemm_bd8_row_f32_i64) +{ + check_bsrgemm(handle, 8, rocsparse_direction_row); +} + +TEST_F(GemmCalcGrids, bsrgemm_bd8_column_f64_i32) +{ + check_bsrgemm(handle, 8, rocsparse_direction_column); +} + +TEST_F(GemmCalcGrids, bsrgemm_bd16_row_f64_i64) +{ + check_bsrgemm(handle, 16, rocsparse_direction_row); +} + +TEST_F(GemmCalcGrids, bsrgemm_bd16_column_f32_i32) +{ + check_bsrgemm(handle, 16, rocsparse_direction_column); +} diff --git a/projects/rocsparse/library/src/extra/bsrgemm_device.h b/projects/rocsparse/library/src/extra/bsrgemm_device.h index a1f3289589c1..a565403b60c9 100644 --- a/projects/rocsparse/library/src/extra/bsrgemm_device.h +++ b/projects/rocsparse/library/src/extra/bsrgemm_device.h @@ -230,7 +230,8 @@ namespace rocsparse typename J, typename T> ROCSPARSE_DEVICE_ILF void - bsrgemm_fill_wf_per_row_2x2_device(rocsparse_direction dir, + bsrgemm_fill_wf_per_row_2x2_device(J block_offset, + rocsparse_direction dir, J mb, J nkb, const J* __restrict__ offset, @@ -266,8 +267,9 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each (sub)wavefront processes a row - J row = hipBlockIdx_x * (BLOCKSIZE / WFSIZE) + wid; + // Each (sub)wavefront processes a row (block_offset supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = block_offset + wid; // Hash table in shared memory __shared__ J stable[(BLOCKSIZE / WFSIZE) * HASHSIZE]; @@ -503,7 +505,8 @@ namespace rocsparse typename J, typename T> ROCSPARSE_DEVICE_ILF void - bsrgemm_fill_block_per_row_2x2_device(rocsparse_direction dir, + bsrgemm_fill_block_per_row_2x2_device(J block_id, + rocsparse_direction dir, J mb, J nkb, const J* __restrict__ offset, @@ -557,8 +560,9 @@ namespace rocsparse __syncthreads(); - // Each block processes a row (apply permutation) - J row = perm ? perm[hipBlockIdx_x + *offset] : hipBlockIdx_x; + // Each block processes a row (block_id supplied by the grid-stride loop in the + // kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm ? perm[block_id + *offset] : block_id; // alpha * A * B part if(mul == true) @@ -762,7 +766,8 @@ namespace rocsparse typename I, typename J, typename T> - ROCSPARSE_DEVICE_ILF void bsrgemm_fill_wf_per_row_device(rocsparse_direction dir, + ROCSPARSE_DEVICE_ILF void bsrgemm_fill_wf_per_row_device(J block_offset, + rocsparse_direction dir, J mb, J nkb, J block_dim, @@ -804,8 +809,9 @@ namespace rocsparse int c = slid & (BLOCKDIM - 1); int r = slid / BLOCKDIM; - // Each (sub)wavefront processes a row - J row = hipBlockIdx_x * (BLOCKSIZE / WFSIZE) + wid; + // Each (sub)wavefront processes a row (block_offset supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = block_offset + wid; // Hash table in shared memory __shared__ J stable[(BLOCKSIZE / WFSIZE) * HASHSIZE]; @@ -993,7 +999,8 @@ namespace rocsparse typename I, typename J, typename T> - ROCSPARSE_DEVICE_ILF void bsrgemm_fill_block_per_row_device(rocsparse_direction dir, + ROCSPARSE_DEVICE_ILF void bsrgemm_fill_block_per_row_device(J block_id, + rocsparse_direction dir, J mb, J nkb, J block_dim, @@ -1030,8 +1037,9 @@ namespace rocsparse int c = lid & (BLOCKDIM - 1); int r = lid / BLOCKDIM; - // Each block processes a row (apply permutation) - J row = perm ? perm[hipBlockIdx_x + *offset] : hipBlockIdx_x; + // Each block processes a row (block_id supplied by the grid-stride loop in the + // kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm ? perm[block_id + *offset] : block_id; // Hash table in shared memory __shared__ J table[HASHSIZE]; @@ -1207,7 +1215,8 @@ namespace rocsparse typename J, typename T> ROCSPARSE_DEVICE_ILF void - bsrgemm_block_per_row_atomic_multipass_device(rocsparse_direction dir, + bsrgemm_block_per_row_atomic_multipass_device(J block_id, + rocsparse_direction dir, J nb, J block_dim, const J* __restrict__ offset_, @@ -1244,8 +1253,9 @@ namespace rocsparse int c = lid & (BLOCKDIM - 1); int r = lid / BLOCKDIM; - // Each block processes a row (apply permutation) - J row = perm ? perm[hipBlockIdx_x + *offset_] : hipBlockIdx_x; + // Each block processes a row (block_id supplied by the grid-stride loop in the + // kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm ? perm[block_id + *offset_] : block_id; // Row entry marker and value accumulator __shared__ int table[CHUNKSIZE]; @@ -1538,7 +1548,8 @@ namespace rocsparse typename J, typename T> ROCSPARSE_DEVICE_ILF void - bsrgemm_block_per_row_multipass_device(rocsparse_direction dir, + bsrgemm_block_per_row_multipass_device(J block_id, + rocsparse_direction dir, J nb, J block_dim, const J* __restrict__ offset_, @@ -1572,8 +1583,9 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / (BLOCKSIZE / BLOCKDIM); - // Each block processes a row (apply permutation) - J row = perm ? perm[hipBlockIdx_x + *offset_] : hipBlockIdx_x; + // Each block processes a row (block_id supplied by the grid-stride loop in the + // kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm ? perm[block_id + *offset_] : block_id; // Row entry marker and value accumulator __shared__ bool table[CHUNKSIZE]; diff --git a/projects/rocsparse/library/src/extra/csrgemm_device.h b/projects/rocsparse/library/src/extra/csrgemm_device.h index 8e601f336e43..dc4a368049ac 100644 --- a/projects/rocsparse/library/src/extra/csrgemm_device.h +++ b/projects/rocsparse/library/src/extra/csrgemm_device.h @@ -70,7 +70,7 @@ namespace rocsparse } // Compute number of intermediate products of each row - template + template ROCSPARSE_KERNEL(BLOCKSIZE) void csrgemm_intermediate_products(J m, const I* __restrict__ csr_row_ptr_A, @@ -88,50 +88,53 @@ namespace rocsparse // Lane id int lid = hipThreadIdx_x & (WFSIZE - 1); - // Each (sub)wavefront processes a row - J row = (hipBlockIdx_x * BLOCKSIZE + hipThreadIdx_x) / WFSIZE; - - // Bounds check - if(row >= m) + // Each (sub)wavefront processes a row, grid-strided so a grid clamped by + // get_grid_size_x still covers every row. The first row fits in 32 bits because + // the dispatch caps grid.x * BLOCKSIZE below 2^32. + for(int64_t row = (hipBlockIdx_x * BLOCKSIZE + hipThreadIdx_x) / WFSIZE; row < m; + row += static_cast(hipGridDim_x) * (BLOCKSIZE / WFSIZE)) { - return; - } + // Initialize intermediate product counter of current row + I nprod = 0; - // Initialize intermediate product counter of current row - I nprod = 0; + // alpha * A * B part + if(mul == true) + { + // Row begin and row end of A matrix + I row_begin_A = csr_row_ptr_A[row] - idx_base_A; + I row_end_A = csr_row_ptr_A[row + 1] - idx_base_A; - // alpha * A * B part - if(mul == true) - { - // Row begin and row end of A matrix - I row_begin_A = csr_row_ptr_A[row] - idx_base_A; - I row_end_A = csr_row_ptr_A[row + 1] - idx_base_A; + // Loop over columns of A in current row + for(I j = row_begin_A + lid; j < row_end_A; j += WFSIZE) + { + // Current column of A + J col_A = csr_col_ind_A[j] - idx_base_A; - // Loop over columns of A in current row - for(I j = row_begin_A + lid; j < row_end_A; j += WFSIZE) - { - // Current column of A - J col_A = csr_col_ind_A[j] - idx_base_A; + // Accumulate non zero entries of B in row col_A + nprod += (csr_row_ptr_B[col_A + 1] - csr_row_ptr_B[col_A]); + } - // Accumulate non zero entries of B in row col_A - nprod += (csr_row_ptr_B[col_A + 1] - csr_row_ptr_B[col_A]); + // Gather nprod + nprod = rocsparse::wfreduce_sum(nprod); } - // Gather nprod - nprod = rocsparse::wfreduce_sum(nprod); - } - - // Last lane writes result - if(lid == WFSIZE - 1) - { - // beta * D part - if(add == true) + // Last lane writes result + if(lid == WFSIZE - 1) { - nprod += (csr_row_ptr_D[row + 1] - csr_row_ptr_D[row]); + // beta * D part + if(add == true) + { + nprod += (csr_row_ptr_D[row + 1] - csr_row_ptr_D[row]); + } + + // Write number of intermediate products of the current row + int_prod[row] = nprod; } - // Write number of intermediate products of the current row - int_prod[row] = nprod; + if constexpr(!GRID_STRIDE) + { + break; + } } } @@ -456,6 +459,7 @@ namespace rocsparse uint32_t WFSIZE, uint32_t HASHSIZE, uint32_t HASHVAL, + bool GRID_STRIDE, typename I, typename J> ROCSPARSE_KERNEL(BLOCKSIZE) @@ -485,84 +489,87 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each (sub)wavefront processes a row - J row = hipBlockIdx_x * BLOCKSIZE / WFSIZE + wid; - // Hash table in shared memory __shared__ J stable[BLOCKSIZE / WFSIZE * HASHSIZE]; // Local hash table J* table = &stable[wid * HASHSIZE]; - // Initialize hash table - for(uint32_t i = lid; i < HASHSIZE; i += WFSIZE) + // Grid-stride over the (sub)wavefront rows so a grid clamped by get_grid_size_x + // covers all rows. The first row fits in 32 bits because the dispatch caps + // grid.x * BLOCKSIZE below 2^32. + for(int64_t idx = (hipBlockIdx_x * BLOCKSIZE + hipThreadIdx_x) / WFSIZE; idx < m; + idx += static_cast(hipGridDim_x) * (BLOCKSIZE / WFSIZE)) { - table[i] = -1; - } + // Initialize hash table + for(uint32_t i = lid; i < HASHSIZE; i += WFSIZE) + { + table[i] = -1; + } - __threadfence_block(); + __threadfence_block(); - // Bounds check - if(row >= m) - { - return; - } + // Apply permutation, if available + J row = perm ? perm[idx + *offset] : static_cast(idx); - // Apply permutation, if available - row = perm ? perm[row + *offset] : row; + // Initialize row nnz + J nnz = 0; - // Initialize row nnz - J nnz = 0; + // alpha * A * B part + if(mul == true) + { + // Get row boundaries of the current row in A + I row_begin_A = csr_row_ptr_A[row] - idx_base_A; + I row_end_A = csr_row_ptr_A[row + 1] - idx_base_A; - // alpha * A * B part - if(mul == true) - { - // Get row boundaries of the current row in A - I row_begin_A = csr_row_ptr_A[row] - idx_base_A; - I row_end_A = csr_row_ptr_A[row + 1] - idx_base_A; + // Loop over columns of A in current row + for(I j = row_begin_A + lid; j < row_end_A; j += WFSIZE) + { + // Column of A in current row + J col_A = csr_col_ind_A[j] - idx_base_A; - // Loop over columns of A in current row - for(I j = row_begin_A + lid; j < row_end_A; j += WFSIZE) - { - // Column of A in current row - J col_A = csr_col_ind_A[j] - idx_base_A; + // Loop over columns of B in row col_A + I row_begin_B = csr_row_ptr_B[col_A] - idx_base_B; + I row_end_B = csr_row_ptr_B[col_A + 1] - idx_base_B; - // Loop over columns of B in row col_A - I row_begin_B = csr_row_ptr_B[col_A] - idx_base_B; - I row_end_B = csr_row_ptr_B[col_A + 1] - idx_base_B; + // Insert all columns of B into hash table + for(I k = row_begin_B; k < row_end_B; ++k) + { + // Count the actual insertions to obtain row nnz of C + nnz += insert_key(csr_col_ind_B[k] - idx_base_B, table); + } + } + } - // Insert all columns of B into hash table - for(I k = row_begin_B; k < row_end_B; ++k) + // beta * D part + if(add == true) + { + // Get row boundaries of the current row in D + I row_begin_D = csr_row_ptr_D[row] - idx_base_D; + I row_end_D = csr_row_ptr_D[row + 1] - idx_base_D; + + // Loop over columns of D in current row and insert all columns of D into hash table + for(I j = row_begin_D + lid; j < row_end_D; j += WFSIZE) { // Count the actual insertions to obtain row nnz of C - nnz += insert_key(csr_col_ind_B[k] - idx_base_B, table); + nnz += insert_key(csr_col_ind_D[j] - idx_base_D, table); } } - } - // beta * D part - if(add == true) - { - // Get row boundaries of the current row in D - I row_begin_D = csr_row_ptr_D[row] - idx_base_D; - I row_end_D = csr_row_ptr_D[row + 1] - idx_base_D; + // Accumulate all row nnz within each (sub)wavefront to obtain the total row nnz + // of the current row + nnz = rocsparse::wfreduce_sum(nnz); - // Loop over columns of D in current row and insert all columns of D into hash table - for(I j = row_begin_D + lid; j < row_end_D; j += WFSIZE) + // Write result to global memory + if(lid == WFSIZE - 1) { - // Count the actual insertions to obtain row nnz of C - nnz += insert_key(csr_col_ind_D[j] - idx_base_D, table); + row_nnz[row] = nnz; } - } - // Accumulate all row nnz within each (sub)wavefront to obtain the total row nnz - // of the current row - nnz = rocsparse::wfreduce_sum(nnz); - - // Write result to global memory - if(lid == WFSIZE - 1) - { - row_nnz[row] = nnz; + if constexpr(!GRID_STRIDE) + { + break; + } } } @@ -571,10 +578,12 @@ namespace rocsparse uint32_t WFSIZE, uint32_t HASHSIZE, uint32_t HASHVAL, + bool GRID_STRIDE, typename I, typename J> ROCSPARSE_KERNEL(BLOCKSIZE) - void csrgemm_nnz_block_per_row(const J* __restrict__ offset, + void csrgemm_nnz_block_per_row(J size, + const J* __restrict__ offset, const J* __restrict__ perm, const I* __restrict__ csr_row_ptr_A, const J* __restrict__ csr_col_ind_A, @@ -599,91 +608,105 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each block processes a row (apply permutation) - J row = perm[hipBlockIdx_x + *offset]; - // Hash table in shared memory extern __shared__ char shared_memory[]; J* table = (J*)shared_memory; - // Initialize hash table - for(uint32_t i = hipThreadIdx_x; i < HASHSIZE; i += BLOCKSIZE) + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) { - table[i] = -1; - } + // Each block processes a row (apply permutation) + J row = perm[block_id + *offset]; - // Wait for all threads to finish initialization - __syncthreads(); + // Initialize hash table + for(uint32_t i = hipThreadIdx_x; i < HASHSIZE; i += BLOCKSIZE) + { + table[i] = -1; + } - // Initialize row nnz - J nnz = 0; + // Wait for all threads to finish initialization + __syncthreads(); - // alpha * A * B part - if(mul == true) - { - // Get row boundaries of the current row in A - I row_begin_A = csr_row_ptr_A[row] - idx_base_A; - I row_end_A = csr_row_ptr_A[row + 1] - idx_base_A; + // Initialize row nnz + J nnz = 0; - // Loop over columns of A in current row - for(I j = row_begin_A + wid; j < row_end_A; j += BLOCKSIZE / WFSIZE) + // alpha * A * B part + if(mul == true) { - // Column of A in current row - J col_A = csr_col_ind_A[j] - idx_base_A; + // Get row boundaries of the current row in A + I row_begin_A = csr_row_ptr_A[row] - idx_base_A; + I row_end_A = csr_row_ptr_A[row + 1] - idx_base_A; - // Loop over columns of B in row col_A - I row_begin_B = csr_row_ptr_B[col_A] - idx_base_B; - I row_end_B = csr_row_ptr_B[col_A + 1] - idx_base_B; + // Loop over columns of A in current row + for(I j = row_begin_A + wid; j < row_end_A; j += BLOCKSIZE / WFSIZE) + { + // Column of A in current row + J col_A = csr_col_ind_A[j] - idx_base_A; - for(I k = row_begin_B + lid; k < row_end_B; k += WFSIZE) + // Loop over columns of B in row col_A + I row_begin_B = csr_row_ptr_B[col_A] - idx_base_B; + I row_end_B = csr_row_ptr_B[col_A + 1] - idx_base_B; + + for(I k = row_begin_B + lid; k < row_end_B; k += WFSIZE) + { + // Count the actual insertions to obtain row nnz of C + nnz += insert_key(csr_col_ind_B[k] - idx_base_B, table); + } + } + } + + // beta * D part + if(add == true) + { + // Get row boundaries of the current row in D + I row_begin_D = csr_row_ptr_D[row] - idx_base_D; + I row_end_D = csr_row_ptr_D[row + 1] - idx_base_D; + + // Loop over columns of D in current row and insert all columns of D into hash table + for(I j = row_begin_D + wid; j < row_end_D; j += BLOCKSIZE / WFSIZE) { // Count the actual insertions to obtain row nnz of C - nnz += insert_key(csr_col_ind_B[k] - idx_base_B, table); + nnz += insert_key(csr_col_ind_D[j] - idx_base_D, table); } } - } - // beta * D part - if(add == true) - { - // Get row boundaries of the current row in D - I row_begin_D = csr_row_ptr_D[row] - idx_base_D; - I row_end_D = csr_row_ptr_D[row + 1] - idx_base_D; + // Wait for all threads to finish hash operation + __syncthreads(); - // Loop over columns of D in current row and insert all columns of D into hash table - for(I j = row_begin_D + wid; j < row_end_D; j += BLOCKSIZE / WFSIZE) + // Accumulate all row nnz within each (sub)wavefront to obtain the total row nnz + // of the current row + nnz = rocsparse::wfreduce_sum(nnz); + + // Write result to shared memory for final reduction by first wavefront + if(lid == WFSIZE - 1) { - // Count the actual insertions to obtain row nnz of C - nnz += insert_key(csr_col_ind_D[j] - idx_base_D, table); + table[wid] = nnz; } - } - - // Wait for all threads to finish hash operation - __syncthreads(); - - // Accumulate all row nnz within each (sub)wavefront to obtain the total row nnz - // of the current row - nnz = rocsparse::wfreduce_sum(nnz); - // Write result to shared memory for final reduction by first wavefront - if(lid == WFSIZE - 1) - { - table[wid] = nnz; - } + // Wait for all threads to finish reduction + __syncthreads(); - // Wait for all threads to finish reduction - __syncthreads(); + // Gather row nnz for the whole block + nnz = (hipThreadIdx_x < BLOCKSIZE / WFSIZE) ? table[hipThreadIdx_x] : 0; - // Gather row nnz for the whole block - nnz = (hipThreadIdx_x < BLOCKSIZE / WFSIZE) ? table[hipThreadIdx_x] : 0; + // First wavefront computes final sum + nnz = rocsparse::wfreduce_sum(nnz); - // First wavefront computes final sum - nnz = rocsparse::wfreduce_sum(nnz); + // Write result to global memory + if(hipThreadIdx_x == BLOCKSIZE / WFSIZE - 1) + { + row_nnz[row] = nnz; + } - // Write result to global memory - if(hipThreadIdx_x == BLOCKSIZE / WFSIZE - 1) - { - row_nnz[row] = nnz; + if constexpr(GRID_STRIDE) + { + // The next row re-initialises the shared hash table read above + __syncthreads(); + } + else + { + break; + } } } @@ -691,9 +714,15 @@ namespace rocsparse // Splitting row into several chunks such that we can use shared memory to store whether // a column index is populated or not. // Each row has at least 8193 intermediate products to compute. - template + template ROCSPARSE_KERNEL(BLOCKSIZE) - void csrgemm_nnz_block_per_row_multipass(J n, + void csrgemm_nnz_block_per_row_multipass(J size, + J n, const J* __restrict__ offset, const J* __restrict__ perm, const I* __restrict__ csr_row_ptr_A, @@ -718,9 +747,6 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each block processes a row (apply permutation) - J row = perm[hipBlockIdx_x + *offset]; - // Row nnz marker __shared__ bool table[CHUNKSIZE]; @@ -731,169 +757,181 @@ namespace rocsparse // current chunk __shared__ J next_chunk; - // Begin of the current row chunk (this is the column index of the current row) - J chunk_begin = 0; - J chunk_end = CHUNKSIZE; - - // Initialize row nnz for the full row - if(hipThreadIdx_x == 0) + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) { - nnz = 0; - } - - // Get row boundaries of the current row in A - I row_begin_A = (mul == true) ? csr_row_ptr_A[row] - idx_base_A : 0; - I row_end_A = (mul == true) ? csr_row_ptr_A[row + 1] - idx_base_A : 0; + // Each block processes a row (apply permutation) + J row = perm[block_id + *offset]; - // Loop over the row chunks until the end of the row has been reached (which is - // the number of total columns) - while(chunk_begin < n) - { - // Initialize row nnz table - for(uint32_t i = hipThreadIdx_x; i < CHUNKSIZE; i += BLOCKSIZE) - { - table[i] = false; - } + // Begin of the current row chunk (this is the column index of the current row) + J chunk_begin = 0; + J chunk_end = CHUNKSIZE; - // Initialize next chunk column index + // Initialize row nnz for the full row if(hipThreadIdx_x == 0) { - next_chunk = n; + nnz = 0; } - // Wait for all threads to finish initialization - __syncthreads(); - - // Initialize the beginning of the next chunk - J min_col = n; + // Get row boundaries of the current row in A + I row_begin_A = (mul == true) ? csr_row_ptr_A[row] - idx_base_A : 0; + I row_end_A = (mul == true) ? csr_row_ptr_A[row + 1] - idx_base_A : 0; - // alpha * A * B part - if(mul == true) + // Loop over the row chunks until the end of the row has been reached (which is + // the number of total columns) + while(chunk_begin < n) { - // Loop over columns of A in current row - for(I j = row_begin_A + wid; j < row_end_A; j += BLOCKSIZE / WFSIZE) + // Initialize row nnz table + for(uint32_t i = hipThreadIdx_x; i < CHUNKSIZE; i += BLOCKSIZE) { - // Column of A in current row - J col_A = csr_col_ind_A[j] - idx_base_A; + table[i] = false; + } - // Loop over columns of B in row col_A - I row_begin_B - = (chunk_begin == 0) ? csr_row_ptr_B[col_A] - idx_base_B : workspace_B[j]; - I row_end_B = csr_row_ptr_B[col_A + 1] - idx_base_B; + // Initialize next chunk column index + if(hipThreadIdx_x == 0) + { + next_chunk = n; + } - // Keep track of the first k where the column index of B is exceeding - // the current chunks end point - I next_k = row_begin_B + lid; + // Wait for all threads to finish initialization + __syncthreads(); - for(I k = next_k; k < row_end_B; k += WFSIZE) + // Initialize the beginning of the next chunk + J min_col = n; + + // alpha * A * B part + if(mul == true) + { + // Loop over columns of A in current row + for(I j = row_begin_A + wid; j < row_end_A; j += BLOCKSIZE / WFSIZE) { - // Column of B in row col_A - J col_B = csr_col_ind_B[k] - idx_base_B; + // Column of A in current row + J col_A = csr_col_ind_A[j] - idx_base_A; - if(col_B >= chunk_begin && col_B < chunk_end) + // Loop over columns of B in row col_A + I row_begin_B = (chunk_begin == 0) ? csr_row_ptr_B[col_A] - idx_base_B + : workspace_B[j]; + I row_end_B = csr_row_ptr_B[col_A + 1] - idx_base_B; + + // Keep track of the first k where the column index of B is exceeding + // the current chunks end point + I next_k = row_begin_B + lid; + + for(I k = next_k; k < row_end_B; k += WFSIZE) { - // Mark nnz table if entry at col_B - table[col_B - chunk_begin] = true; + // Column of B in row col_A + J col_B = csr_col_ind_B[k] - idx_base_B; + + if(col_B >= chunk_begin && col_B < chunk_end) + { + // Mark nnz table if entry at col_B + table[col_B - chunk_begin] = true; + } + else if(col_B >= chunk_end) + { + // If column index exceeds chunks end point, store k as starting + // point of the columns of B for the next pass + next_k = k; + + // Store the first column index of B that exceeds the current chunk + min_col = rocsparse::min(min_col, col_B); + break; + } } - else if(col_B >= chunk_end) - { - // If column index exceeds chunks end point, store k as starting - // point of the columns of B for the next pass - next_k = k; - // Store the first column index of B that exceeds the current chunk - min_col = rocsparse::min(min_col, col_B); - break; + // Obtain the minimum of all k that exceed the current chunks end point + rocsparse::wfreduce_min(&next_k); + + // Store the minimum globally for the next chunk + if(lid == WFSIZE - 1) + { + workspace_B[j] = next_k; } } + } - // Obtain the minimum of all k that exceed the current chunks end point - rocsparse::wfreduce_min(&next_k); + // beta * D part + if(add == true) + { + // Get row boundaries of the current row in D + I row_begin_D = csr_row_ptr_D[row] - idx_base_D; + I row_end_D = csr_row_ptr_D[row + 1] - idx_base_D; - // Store the minimum globally for the next chunk - if(lid == WFSIZE - 1) + // Loop over columns of D in current row and insert all columns of D into hash table + for(I j = row_begin_D + hipThreadIdx_x; j < row_end_D; j += BLOCKSIZE) { - workspace_B[j] = next_k; + // Column of D in current row + J col_D = csr_col_ind_D[j] - idx_base_D; + + if(col_D >= chunk_begin && col_D < chunk_end) + { + // Mark nnz table if entry at col_D + table[col_D - chunk_begin] = true; + } + else if(col_D >= chunk_end) + { + // Store the first column index of D that exceeds the current chunk + min_col = rocsparse::min(min_col, col_D); + break; + } + + // Performance can potentially improved by adding another temporary + // workspace of dimension sizeof(J) * nnz, which is significant! } } - } - // beta * D part - if(add == true) - { - // Get row boundaries of the current row in D - I row_begin_D = csr_row_ptr_D[row] - idx_base_D; - I row_end_D = csr_row_ptr_D[row + 1] - idx_base_D; + // Gather wavefront-wide minimum for the next chunks starting column index + rocsparse::wfreduce_min(&min_col); - // Loop over columns of D in current row and insert all columns of D into hash table - for(I j = row_begin_D + hipThreadIdx_x; j < row_end_D; j += BLOCKSIZE) + // Last thread in each wavefront finds block-wide minimum atomically + if(lid == WFSIZE - 1) { - // Column of D in current row - J col_D = csr_col_ind_D[j] - idx_base_D; + // Atomically determine the new chunks beginning (minimum column index of B + // that is larger than the current chunks end point) + rocsparse::atomic_min(&next_chunk, min_col); + } - if(col_D >= chunk_begin && col_D < chunk_end) - { - // Mark nnz table if entry at col_D - table[col_D - chunk_begin] = true; - } - else if(col_D >= chunk_end) - { - // Store the first column index of D that exceeds the current chunk - min_col = rocsparse::min(min_col, col_D); - break; - } + // Wait for all threads to finish row nnz operation + __syncthreads(); - // Performance can potentially improved by adding another temporary - // workspace of dimension sizeof(J) * nnz, which is significant! + // Each thread loads its entry for the current chunk + J chunk_nnz = 0; + for(uint32_t i = hipThreadIdx_x; i < CHUNKSIZE; i += BLOCKSIZE) + { + chunk_nnz += (table[i] == true) ? 1 : 0; } - } - // Gather wavefront-wide minimum for the next chunks starting column index - rocsparse::wfreduce_min(&min_col); + // Gather wavefront-wide nnz for the current chunk + chunk_nnz = rocsparse::wfreduce_sum(chunk_nnz); - // Last thread in each wavefront finds block-wide minimum atomically - if(lid == WFSIZE - 1) - { - // Atomically determine the new chunks beginning (minimum column index of B - // that is larger than the current chunks end point) - rocsparse::atomic_min(&next_chunk, min_col); - } + // Last thread in each wavefront accumulates block-wide nnz atomically + if(lid == WFSIZE - 1) + { + // Atomically add this chunks nnz to the total row nnz + rocsparse::atomic_add(&nnz, chunk_nnz); + } - // Wait for all threads to finish row nnz operation - __syncthreads(); + // Wait for atomics to be processed + __syncthreads(); - // Each thread loads its entry for the current chunk - J chunk_nnz = 0; - for(uint32_t i = hipThreadIdx_x; i < CHUNKSIZE; i += BLOCKSIZE) - { - chunk_nnz += (table[i] == true) ? 1 : 0; - } + // Each thread loads the new chunk beginning and end point + chunk_begin = next_chunk; + chunk_end = chunk_begin + CHUNKSIZE; - // Gather wavefront-wide nnz for the current chunk - chunk_nnz = rocsparse::wfreduce_sum(chunk_nnz); + // Wait for all threads to finish load from shared memory + __syncthreads(); + } - // Last thread in each wavefront accumulates block-wide nnz atomically - if(lid == WFSIZE - 1) + // Write accumulated total row nnz to global memory + if(hipThreadIdx_x == 0) { - // Atomically add this chunks nnz to the total row nnz - rocsparse::atomic_add(&nnz, chunk_nnz); + row_nnz[row] = nnz; } - // Wait for atomics to be processed - __syncthreads(); - - // Each thread loads the new chunk beginning and end point - chunk_begin = next_chunk; - chunk_end = chunk_begin + CHUNKSIZE; - - // Wait for all threads to finish load from shared memory - __syncthreads(); - } - - // Write accumulated total row nnz to global memory - if(hipThreadIdx_x == 0) - { - row_nnz[row] = nnz; + if constexpr(!GRID_STRIDE) + { + break; + } } } @@ -905,7 +943,8 @@ namespace rocsparse typename I, typename J, typename T> - ROCSPARSE_DEVICE_ILF void csrgemm_fill_wf_per_row_device(J m, + ROCSPARSE_DEVICE_ILF void csrgemm_fill_wf_per_row_device(J block_offset, + J m, J nk, const J* __restrict__ offset, const J* __restrict__ perm, @@ -940,8 +979,9 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each (sub)wavefront processes a row - J row = hipBlockIdx_x * BLOCKSIZE / WFSIZE + wid; + // Each (sub)wavefront processes a row (block_offset supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = block_offset + wid; // Hash table in shared memory __shared__ J stable[BLOCKSIZE / WFSIZE * HASHSIZE]; @@ -1072,7 +1112,8 @@ namespace rocsparse typename I, typename J, typename T> - ROCSPARSE_DEVICE_ILF void csrgemm_fill_block_per_row_device(J nk, + ROCSPARSE_DEVICE_ILF void csrgemm_fill_block_per_row_device(J block_id, + J nk, const J* __restrict__ offset_, const J* __restrict__ perm, T alpha, @@ -1121,8 +1162,9 @@ namespace rocsparse // Wait for all threads to finish initialization __syncthreads(); - // Each block processes a row (apply permutation) - J row = perm[hipBlockIdx_x + *offset_]; + // Each block processes a row (apply permutation; block_id supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm[block_id + *offset_]; // alpha * A * B part if(mul == true) @@ -1292,7 +1334,8 @@ namespace rocsparse typename J, typename T> ROCSPARSE_DEVICE_ILF void - csrgemm_fill_block_per_row_multipass_device(J n, + csrgemm_fill_block_per_row_multipass_device(J block_id, + J n, const J* __restrict__ offset_, const J* __restrict__ perm, T alpha, @@ -1325,8 +1368,9 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each block processes a row (apply permutation) - J row = perm[hipBlockIdx_x + *offset_]; + // Each block processes a row (apply permutation; block_id supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm[block_id + *offset_]; // Row entry marker and value accumulator __shared__ bool table[CHUNKSIZE]; diff --git a/projects/rocsparse/library/src/extra/csrgemm_numeric_device.h b/projects/rocsparse/library/src/extra/csrgemm_numeric_device.h index a85f6de1b78b..cdad6745392b 100644 --- a/projects/rocsparse/library/src/extra/csrgemm_numeric_device.h +++ b/projects/rocsparse/library/src/extra/csrgemm_numeric_device.h @@ -155,7 +155,8 @@ namespace rocsparse typename J, typename T> ROCSPARSE_DEVICE_ILF void - csrgemm_numeric_fill_wf_per_row_device(J m, + csrgemm_numeric_fill_wf_per_row_device(J block_offset, + J m, J nk, const J* __restrict__ offset, const J* __restrict__ perm, @@ -190,8 +191,9 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each (sub)wavefront processes a row - J row = hipBlockIdx_x * BLOCKSIZE / WFSIZE + wid; + // Each (sub)wavefront processes a row (block_offset supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = block_offset + wid; // Hash table in shared memory __shared__ J stable[BLOCKSIZE / WFSIZE * HASHSIZE]; @@ -319,7 +321,8 @@ namespace rocsparse typename J, typename T> ROCSPARSE_DEVICE_ILF void - csrgemm_numeric_fill_block_per_row_device(J nk, + csrgemm_numeric_fill_block_per_row_device(J block_id, + J nk, const J* __restrict__ offset_, const J* __restrict__ perm, T alpha, @@ -368,8 +371,9 @@ namespace rocsparse // Wait for all threads to finish initialization __syncthreads(); - // Each block processes a row (apply permutation) - J row = perm[hipBlockIdx_x + *offset_]; + // Each block processes a row (apply permutation; block_id supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm[block_id + *offset_]; // alpha * A * B part if(mul) @@ -536,7 +540,8 @@ namespace rocsparse typename J, typename T> ROCSPARSE_DEVICE_ILF void - csrgemm_numeric_fill_block_per_row_multipass_device(J n, + csrgemm_numeric_fill_block_per_row_multipass_device(J block_id, + J n, const J* __restrict__ offset_, const J* __restrict__ perm, T alpha, @@ -569,8 +574,9 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each block processes a row (apply permutation) - J row = perm[hipBlockIdx_x + *offset_]; + // Each block processes a row (apply permutation; block_id supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm[block_id + *offset_]; // Row entry marker and value accumulator __shared__ bool table[CHUNKSIZE]; diff --git a/projects/rocsparse/library/src/extra/csrgemm_symbolic_device.h b/projects/rocsparse/library/src/extra/csrgemm_symbolic_device.h index 30cb157bfbf6..0ba8ecfa7e59 100644 --- a/projects/rocsparse/library/src/extra/csrgemm_symbolic_device.h +++ b/projects/rocsparse/library/src/extra/csrgemm_symbolic_device.h @@ -39,7 +39,8 @@ namespace rocsparse typename I, typename J> ROCSPARSE_DEVICE_ILF void - csrgemm_symbolic_fill_block_per_row_multipass_device(J n, + csrgemm_symbolic_fill_block_per_row_multipass_device(J block_id, + J n, const J* __restrict__ offset_, const J* __restrict__ perm, const I* __restrict__ csr_row_ptr_A, @@ -66,8 +67,9 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each block processes a row (apply permutation) - J row = perm[hipBlockIdx_x + *offset_]; + // Each block processes a row (apply permutation; block_id supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm[block_id + *offset_]; // Row entry marker and value accumulator __shared__ bool table[CHUNKSIZE]; @@ -571,7 +573,8 @@ namespace rocsparse typename I, typename J> ROCSPARSE_DEVICE_ILF void - csrgemm_symbolic_fill_wf_per_row_device(J m, + csrgemm_symbolic_fill_wf_per_row_device(J block_offset, + J m, J nk, const J* __restrict__ offset, const J* __restrict__ perm, @@ -600,8 +603,9 @@ namespace rocsparse // Wavefront id int wid = hipThreadIdx_x / WFSIZE; - // Each (sub)wavefront processes a row - J row = hipBlockIdx_x * BLOCKSIZE / WFSIZE + wid; + // Each (sub)wavefront processes a row (block_offset supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = block_offset + wid; // Hash table in shared memory __shared__ J stable[BLOCKSIZE / WFSIZE * HASHSIZE]; @@ -732,7 +736,8 @@ namespace rocsparse typename I, typename J> ROCSPARSE_DEVICE_ILF void - csrgemm_symbolic_fill_block_per_row_device(J nk, + csrgemm_symbolic_fill_block_per_row_device(J block_id, + J nk, const J* __restrict__ offset_, const J* __restrict__ perm, const I* __restrict__ csr_row_ptr_A, @@ -773,8 +778,9 @@ namespace rocsparse // Wait for all threads to finish initialization __syncthreads(); - // Each block processes a row (apply permutation) - J row = perm[hipBlockIdx_x + *offset_]; + // Each block processes a row (apply permutation; block_id supplied by the grid-stride + // loop in the kernel wrapper so a grid clamped by get_grid_size_x still covers all rows) + J row = perm[block_id + *offset_]; // alpha * A * B part if(mul == true) diff --git a/projects/rocsparse/library/src/extra/rocsparse_bsrgemm_calc.cpp b/projects/rocsparse/library/src/extra/rocsparse_bsrgemm_calc.cpp index b06db8e3f74a..1c71c1171c1c 100644 --- a/projects/rocsparse/library/src/extra/rocsparse_bsrgemm_calc.cpp +++ b/projects/rocsparse/library/src/extra/rocsparse_bsrgemm_calc.cpp @@ -30,6 +30,7 @@ #include "rocsparse_bsrgemm.hpp" #include "rocsparse_control.hpp" #include "rocsparse_csrgemm.hpp" +#include "rocsparse_grid.hpp" #include "rocsparse_utility.hpp" #include "rocsparse_primitives.hpp" @@ -44,6 +45,7 @@ namespace rocsparse uint32_t WF_SIZE, uint32_t HASHSIZE, uint32_t HASHVAL, + bool GRID_STRIDE, typename I, typename J, typename T> @@ -77,32 +79,45 @@ namespace rocsparse { ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::bsrgemm_fill_wf_per_row_2x2_device( - dir, - mb, - nkb, - offset, - perm, - alpha, - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - beta, - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the (sub)wavefront rows so a grid clamped by get_grid_size_x + // covers all rows + for(int64_t block_offset = static_cast(hipBlockIdx_x) * (BLOCKSIZE / WF_SIZE); + block_offset < mb; + block_offset += static_cast(hipGridDim_x) * (BLOCKSIZE / WF_SIZE)) + { + rocsparse::bsrgemm_fill_wf_per_row_2x2_device( + static_cast(block_offset), + dir, + mb, + nkb, + offset, + perm, + alpha, + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + beta, + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(!GRID_STRIDE) + { + break; + } + } } template @@ -115,6 +130,7 @@ namespace rocsparse uint32_t WFSIZE, uint32_t HASHSIZE, uint32_t HASHVAL, + bool GRID_STRIDE, typename I, typename J, typename T> @@ -148,32 +164,47 @@ namespace rocsparse { ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::bsrgemm_fill_block_per_row_2x2_device( - dir, - mb, - nkb, - offset, - perm, - alpha, - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - beta, - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < mb; block_id += hipGridDim_x) + { + rocsparse::bsrgemm_fill_block_per_row_2x2_device( + static_cast(block_id), + dir, + mb, + nkb, + offset, + perm, + alpha, + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + beta, + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(GRID_STRIDE) + { + // The next row re-initialises the shared hash table read above + __syncthreads(); + } + else + { + break; + } + } } template @@ -215,39 +247,64 @@ namespace rocsparse { ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::bsrgemm_fill_wf_per_row_device( - dir, - mb, - nkb, - block_dim, - offset, - perm, - alpha, - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - beta, - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the (sub)wavefront rows so a grid clamped by get_grid_size_x + // covers all rows + for(int64_t block_offset = static_cast(hipBlockIdx_x) * (BLOCKSIZE / WF_SIZE); + block_offset < mb; + block_offset += static_cast(hipGridDim_x) * (BLOCKSIZE / WF_SIZE)) + { + rocsparse:: + bsrgemm_fill_wf_per_row_device( + static_cast(block_offset), + dir, + mb, + nkb, + block_dim, + offset, + perm, + alpha, + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + beta, + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(GRID_STRIDE) + { + // The next row re-initialises the shared hash table read above, which a + // 64-lane (sub)wavefront shares across hardware wavefronts on wave32. Skipped + // on the last row, where out-of-range wavefronts have already returned. + if(block_offset + static_cast(hipGridDim_x) * (BLOCKSIZE / WF_SIZE) < mb) + { + __syncthreads(); + } + } + else + { + break; + } + } } template @@ -282,38 +339,54 @@ namespace rocsparse { ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::bsrgemm_fill_block_per_row_device( - dir, - mb, - nkb, - block_dim, - offset, - perm, - alpha, - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - beta, - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < mb; block_id += hipGridDim_x) + { + rocsparse::bsrgemm_fill_block_per_row_device( + static_cast(block_id), + dir, + mb, + nkb, + block_dim, + offset, + perm, + alpha, + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + beta, + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(GRID_STRIDE) + { + // The next row re-initialises the shared hash table read above + __syncthreads(); + } + else + { + break; + } + } } template @@ -321,6 +394,7 @@ namespace rocsparse void bsrgemm_block_per_row_atomic_multipass(rocsparse_direction dir, J nb, J block_dim, + J size, const J* __restrict__ offset, const J* __restrict__ perm, ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha), @@ -348,38 +422,50 @@ namespace rocsparse { ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::bsrgemm_block_per_row_atomic_multipass_device( - dir, - nb, - block_dim, - offset, - perm, - alpha, - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - beta, - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace_B, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) + { + rocsparse:: + bsrgemm_block_per_row_atomic_multipass_device( + static_cast(block_id), + dir, + nb, + block_dim, + offset, + perm, + alpha, + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + beta, + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace_B, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(!GRID_STRIDE) + { + break; + } + } } template @@ -387,6 +473,7 @@ namespace rocsparse void bsrgemm_block_per_row_multipass(rocsparse_direction dir, J nb, J block_dim, + J size, const J* __restrict__ offset, const J* __restrict__ perm, ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha), @@ -414,33 +501,43 @@ namespace rocsparse { ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::bsrgemm_block_per_row_multipass_device( - dir, - nb, - block_dim, - offset, - perm, - alpha, - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - beta, - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace_B, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) + { + rocsparse::bsrgemm_block_per_row_multipass_device( + static_cast(block_id), + dir, + nb, + block_dim, + offset, + perm, + alpha, + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + beta, + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace_B, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(!GRID_STRIDE) + { + break; + } + } } template , - hipFuncAttributeMaxDynamicSharedMemorySize, - bsrgemm_fill_block_per_row_2x2_shared_memory_size())); - - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_block_per_row_2x2), - dim3(group_size), - dim3(BSRGEMM_BLOCKSIZE), - (bsrgemm_fill_block_per_row_2x2_shared_memory_size()), - handle->stream, - dir, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, group_size, - rocsparse::max(kb, nb), - group_offset, - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - base_C, - base_D, - mul, - add, - handle->pointer_mode == rocsparse_pointer_mode_host); + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIP_ERROR(hipFuncSetAttribute( + (const void*) + rocsparse::bsrgemm_fill_block_per_row_2x2, + hipFuncAttributeMaxDynamicSharedMemorySize, + bsrgemm_fill_block_per_row_2x2_shared_memory_size())); + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_block_per_row_2x2), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + (bsrgemm_fill_block_per_row_2x2_shared_memory_size()), + handle->stream, + dir, + group_size, + rocsparse::max(kb, nb), + group_offset, + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + base_C, + base_D, + mul, + add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_HASHSIZE @@ -632,41 +738,49 @@ namespace rocsparse #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_WFSIZE 16 #define BSRGEMM_HASHSIZE 8 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_wf_per_row_2x2), - dim3((BSRGEMM_WFSIZE * (int64_t)group_size[0] - 1) / 256 + 1), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - group_size[0], - rocsparse::max(kb, nb), - &group_offset[0], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (BSRGEMM_WFSIZE * static_cast(group_size[0]) - 1) / 256 + 1, + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_wf_per_row_2x2), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + group_size[0], + rocsparse::max(kb, nb), + &group_offset[0], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_WFSIZE #undef BSRGEMM_HASHSIZE @@ -678,41 +792,49 @@ namespace rocsparse #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_WFSIZE 16 #define BSRGEMM_HASHSIZE 16 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_wf_per_row_2x2), - dim3((BSRGEMM_WFSIZE * (int64_t)group_size[1] - 1) / 256 + 1), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - group_size[1], - rocsparse::max(kb, nb), - &group_offset[1], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (BSRGEMM_WFSIZE * static_cast(group_size[1]) - 1) / 256 + 1, + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_wf_per_row_2x2), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + group_size[1], + rocsparse::max(kb, nb), + &group_offset[1], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_WFSIZE #undef BSRGEMM_HASHSIZE @@ -724,41 +846,49 @@ namespace rocsparse #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_WFSIZE 16 #define BSRGEMM_HASHSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_wf_per_row_2x2), - dim3((BSRGEMM_WFSIZE * (int64_t)group_size[2] - 1) / 256 + 1), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - group_size[2], - rocsparse::max(kb, nb), - &group_offset[2], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (BSRGEMM_WFSIZE * static_cast(group_size[2]) - 1) / 256 + 1, + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_wf_per_row_2x2), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + group_size[2], + rocsparse::max(kb, nb), + &group_offset[2], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_WFSIZE #undef BSRGEMM_HASHSIZE @@ -769,41 +899,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_HASHSIZE 64 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_block_per_row_2x2), - dim3(group_size[3]), - dim3(BSRGEMM_BLOCKSIZE), - (bsrgemm_fill_block_per_row_2x2_shared_memory_size()), - stream, - dir, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, group_size[3], - rocsparse::max(kb, nb), - &group_offset[3], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_block_per_row_2x2), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + (bsrgemm_fill_block_per_row_2x2_shared_memory_size()), + stream, + dir, + group_size[3], + rocsparse::max(kb, nb), + &group_offset[3], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_HASHSIZE } @@ -813,41 +953,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_HASHSIZE 128 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_block_per_row_2x2), - dim3(group_size[4]), - dim3(BSRGEMM_BLOCKSIZE), - (bsrgemm_fill_block_per_row_2x2_shared_memory_size()), - stream, - dir, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, group_size[4], - rocsparse::max(kb, nb), - &group_offset[4], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_block_per_row_2x2), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + (bsrgemm_fill_block_per_row_2x2_shared_memory_size()), + stream, + dir, + group_size[4], + rocsparse::max(kb, nb), + &group_offset[4], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_HASHSIZE } @@ -857,41 +1007,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_HASHSIZE 256 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_block_per_row_2x2), - dim3(group_size[5]), - dim3(BSRGEMM_BLOCKSIZE), - (bsrgemm_fill_block_per_row_2x2_shared_memory_size()), - stream, - dir, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, group_size[5], - rocsparse::max(kb, nb), - &group_offset[5], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_block_per_row_2x2), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + (bsrgemm_fill_block_per_row_2x2_shared_memory_size()), + stream, + dir, + group_size[5], + rocsparse::max(kb, nb), + &group_offset[5], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_HASHSIZE } @@ -935,41 +1095,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 256 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[7]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[7], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[7], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 2, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[7], + &group_offset[7], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1032,43 +1202,51 @@ namespace rocsparse #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_WFSIZE 64 #define BSRGEMM_HASHSIZE 8 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_wf_per_row), - dim3((BSRGEMM_WFSIZE * (int64_t)group_size[0] - 1) / BSRGEMM_BLOCKSIZE + 1), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - group_size[0], - rocsparse::max(kb, nb), - block_dim, - &group_offset[0], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (BSRGEMM_WFSIZE * static_cast(group_size[0]) - 1) / BSRGEMM_BLOCKSIZE + 1, + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_wf_per_row), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + group_size[0], + rocsparse::max(kb, nb), + block_dim, + &group_offset[0], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_WFSIZE #undef BSRGEMM_HASHSIZE @@ -1080,43 +1258,51 @@ namespace rocsparse #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_WFSIZE 64 #define BSRGEMM_HASHSIZE 16 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_wf_per_row), - dim3((BSRGEMM_WFSIZE * (int64_t)group_size[1] - 1) / BSRGEMM_BLOCKSIZE + 1), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - group_size[1], - rocsparse::max(kb, nb), - block_dim, - &group_offset[1], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (BSRGEMM_WFSIZE * static_cast(group_size[1]) - 1) / BSRGEMM_BLOCKSIZE + 1, + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_wf_per_row), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + group_size[1], + rocsparse::max(kb, nb), + block_dim, + &group_offset[1], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_WFSIZE #undef BSRGEMM_HASHSIZE @@ -1127,41 +1313,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[2]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[2], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[2], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 4, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[2], + &group_offset[2], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1171,41 +1367,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 64 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[3]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[3], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[3], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 4, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[3], + &group_offset[3], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1215,41 +1421,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 128 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[4]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[4], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[4], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 4, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[4], + &group_offset[4], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1259,41 +1475,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 128 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[5]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[5], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[5], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 4, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[5], + &group_offset[5], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1303,41 +1529,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 128 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[6]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[6], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[6], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 4, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[6], + &group_offset[6], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1347,41 +1583,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 128 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[7]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[7], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[7], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 4, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[7], + &group_offset[7], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1444,43 +1690,51 @@ namespace rocsparse #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_WFSIZE 64 #define BSRGEMM_HASHSIZE 8 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_fill_wf_per_row), - dim3((BSRGEMM_WFSIZE * (int64_t)group_size[0] - 1) / BSRGEMM_BLOCKSIZE + 1), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - group_size[0], - rocsparse::max(kb, nb), - block_dim, - &group_offset[0], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (BSRGEMM_WFSIZE * static_cast(group_size[0]) - 1) / BSRGEMM_BLOCKSIZE + 1, + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_fill_wf_per_row), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + group_size[0], + rocsparse::max(kb, nb), + block_dim, + &group_offset[0], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_WFSIZE #undef BSRGEMM_HASHSIZE @@ -1491,41 +1745,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 16 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[1]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[1], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[1], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 8, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[1], + &group_offset[1], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1535,41 +1799,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[2]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[2], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[2], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 8, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[2], + &group_offset[2], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1579,41 +1853,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[3]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[3], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[3], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 8, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[3], + &group_offset[3], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1623,41 +1907,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[4]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[4], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[4], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 8, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[4], + &group_offset[4], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1667,41 +1961,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[5]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[5], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[5], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 8, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[5], + &group_offset[5], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1711,41 +2015,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[6]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[6], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[6], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 8, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[6], + &group_offset[6], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1755,41 +2069,51 @@ namespace rocsparse { #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_atomic_multipass), - dim3(group_size[7]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[7], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[7], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_atomic_multipass< + BSRGEMM_BLOCKSIZE, + BSRGEMM_CHUNKSIZE, + 8, + decltype(grid_stride)::value>), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[7], + &group_offset[7], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE } @@ -1848,39 +2172,50 @@ namespace rocsparse #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 8 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_multipass), - dim3(group_size[0]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[0], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[0], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_multipass), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[0], + &group_offset[0], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE @@ -1938,39 +2273,50 @@ namespace rocsparse #define BSRGEMM_BLOCKSIZE 256 #define BSRGEMM_CHUNKSIZE 2 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::bsrgemm_block_per_row_multipass), - dim3(group_size[0]), - dim3(BSRGEMM_BLOCKSIZE), - 0, - stream, - dir, - nb, - block_dim, - &group_offset[0], - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - bsr_row_ptr_A, - bsr_col_ind_A, - bsr_val_A, - bsr_row_ptr_B, - bsr_col_ind_B, - bsr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - bsr_row_ptr_D, - bsr_col_ind_D, - bsr_val_D, - bsr_row_ptr_C, - bsr_col_ind_C, - bsr_val_C, - workspace, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size[0], + BSRGEMM_BLOCKSIZE, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::bsrgemm_block_per_row_multipass), + dim3(grid), + dim3(BSRGEMM_BLOCKSIZE), + 0, + stream, + dir, + nb, + block_dim, + group_size[0], + &group_offset[0], + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + bsr_row_ptr_A, + bsr_col_ind_A, + bsr_val_A, + bsr_row_ptr_B, + bsr_col_ind_B, + bsr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + bsr_row_ptr_D, + bsr_col_ind_D, + bsr_val_D, + bsr_row_ptr_C, + bsr_col_ind_C, + bsr_val_C, + workspace, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef BSRGEMM_BLOCKSIZE #undef BSRGEMM_CHUNKSIZE diff --git a/projects/rocsparse/library/src/extra/rocsparse_csrgemm_calc.cpp b/projects/rocsparse/library/src/extra/rocsparse_csrgemm_calc.cpp index 7ed6d08ecc1d..81e80cc6b897 100644 --- a/projects/rocsparse/library/src/extra/rocsparse_csrgemm_calc.cpp +++ b/projects/rocsparse/library/src/extra/rocsparse_csrgemm_calc.cpp @@ -28,6 +28,7 @@ #include "internal/extra/rocsparse_csrgemm.h" #include "rocsparse_control.hpp" #include "rocsparse_csrgemm.hpp" +#include "rocsparse_grid.hpp" #include "rocsparse_utility.hpp" #include "rocsparse_primitives.hpp" @@ -38,6 +39,7 @@ namespace rocsparse uint32_t WFSIZE, uint32_t HASHSIZE, uint32_t HASHVAL, + bool GRID_STRIDE, typename I, typename J, typename T> @@ -71,31 +73,44 @@ namespace rocsparse ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::csrgemm_fill_wf_per_row_device( - m, - nk, - offset, - perm, - alpha, - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - beta, - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the (sub)wavefront rows so a grid clamped by get_grid_size_x + // covers all rows + for(int64_t block_offset = static_cast(hipBlockIdx_x) * (BLOCKSIZE / WFSIZE); + block_offset < m; + block_offset += static_cast(hipGridDim_x) * (BLOCKSIZE / WFSIZE)) + { + rocsparse::csrgemm_fill_wf_per_row_device( + static_cast(block_offset), + m, + nk, + offset, + perm, + alpha, + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + beta, + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(!GRID_STRIDE) + { + break; + } + } } template @@ -109,11 +124,13 @@ namespace rocsparse uint32_t HASHSIZE, uint32_t HASHVAL, uint32_t WARPSIZE, + bool GRID_STRIDE, typename I, typename J, typename T> ROCSPARSE_KERNEL(BLOCKSIZE) - void csrgemm_fill_block_per_row(J nk, + void csrgemm_fill_block_per_row(J size, + J nk, const J* __restrict__ offset, const J* __restrict__ perm, ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha), @@ -141,42 +158,59 @@ namespace rocsparse ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse:: - csrgemm_fill_block_per_row_device( - nk, - offset, - perm, - alpha, - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - beta, - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) + { + rocsparse:: + csrgemm_fill_block_per_row_device( + static_cast(block_id), + nk, + offset, + perm, + alpha, + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + beta, + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(GRID_STRIDE) + { + // The next row re-initialises the shared hash table read above + __syncthreads(); + } + else + { + break; + } + } } template ROCSPARSE_KERNEL(BLOCKSIZE) - void csrgemm_fill_block_per_row_multipass(J n, + void csrgemm_fill_block_per_row_multipass(J size, + J n, const J* __restrict__ offset, const J* __restrict__ perm, ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha), @@ -204,32 +238,42 @@ namespace rocsparse { ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse:: - csrgemm_fill_block_per_row_multipass_device( - n, - offset, - perm, - alpha, - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - beta, - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - workspace_B, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) + { + rocsparse:: + csrgemm_fill_block_per_row_multipass_device( + static_cast(block_id), + n, + offset, + perm, + alpha, + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + beta, + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + workspace_B, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(!GRID_STRIDE) + { + break; + } + } } } @@ -424,132 +468,121 @@ rocsparse_status rocsparse::csrgemm_calc_template(rocsparse_handle hand } // Compute columns and accumulate values for each group +#define CSRGEMM_FILL_WF_PER_ROW_IMPL(GROUP_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, GS) \ + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ + (rocsparse::csrgemm_fill_wf_per_row), \ + dim3(rocsparse::get_grid_size_x( \ + handle, (h_group_size[GROUP_ID] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1, CSRGEMM_DIM)), \ + dim3(CSRGEMM_DIM), \ + 0, \ + stream, \ + h_group_size[GROUP_ID], \ + rocsparse::max(k, n), \ + &d_group_offset[GROUP_ID], \ + d_perm, \ + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), \ + csr_row_ptr_A, \ + csr_col_ind_A, \ + csr_val_A, \ + csr_row_ptr_B, \ + csr_col_ind_B, \ + csr_val_B, \ + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), \ + csr_row_ptr_D, \ + csr_col_ind_D, \ + csr_val_D, \ + csr_row_ptr_C, \ + csr_col_ind_C, \ + csr_val_C, \ + base_A, \ + base_B, \ + descr_C->base, \ + base_D, \ + info_C->csrgemm_info->mul, \ + info_C->csrgemm_info->add, \ + handle->pointer_mode == rocsparse_pointer_mode_host) + +#define CSRGEMM_FILL_WF_PER_ROW(GROUP_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE) \ + if(rocsparse::get_grid_size_x( \ + handle, (h_group_size[GROUP_ID] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1, CSRGEMM_DIM) \ + < static_cast((h_group_size[GROUP_ID] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1)) \ + { \ + CSRGEMM_FILL_WF_PER_ROW_IMPL(GROUP_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, true); \ + } \ + else \ + { \ + CSRGEMM_FILL_WF_PER_ROW_IMPL(GROUP_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, false); \ + } + // Group 0: 0 - 16 non-zeros per row if(h_group_size[0] > 0) { -#define CSRGEMM_DIM 256 -#define CSRGEMM_SUB 8 -#define CSRGEMM_HASHSIZE 16 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_fill_wf_per_row), - dim3((h_group_size[0] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1), - dim3(CSRGEMM_DIM), - 0, - stream, - h_group_size[0], - rocsparse::max(k, n), - &d_group_offset[0], - d_perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); -#undef CSRGEMM_HASHSIZE -#undef CSRGEMM_SUB -#undef CSRGEMM_DIM + CSRGEMM_FILL_WF_PER_ROW(0, 256, 8, 16); } // Group 1: 17 - 32 non-zeros per row if(h_group_size[1] > 0) { -#define CSRGEMM_DIM 256 -#define CSRGEMM_SUB 16 -#define CSRGEMM_HASHSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_fill_wf_per_row), - dim3((h_group_size[1] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1), - dim3(CSRGEMM_DIM), - 0, - stream, - h_group_size[1], - rocsparse::max(k, n), - &d_group_offset[1], - d_perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); -#undef CSRGEMM_HASHSIZE -#undef CSRGEMM_SUB -#undef CSRGEMM_DIM + CSRGEMM_FILL_WF_PER_ROW(1, 256, 16, 32); } -#define CSRGEMM_FILL_BLOCK_PER_ROW( \ - GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ - (rocsparse::csrgemm_fill_block_per_row), \ - dim3(h_group_size[GROUP_SIZE_ID]), \ - dim3(CSRGEMM_DIM), \ - (csrgemm_fill_block_per_row_shared_memory_size()), \ - stream, \ - rocsparse::max(k, n), \ - &d_group_offset[GROUP_SIZE_ID], \ - d_perm, \ - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), \ - csr_row_ptr_A, \ - csr_col_ind_A, \ - csr_val_A, \ - csr_row_ptr_B, \ - csr_col_ind_B, \ - csr_val_B, \ - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), \ - csr_row_ptr_D, \ - csr_col_ind_D, \ - csr_val_D, \ - csr_row_ptr_C, \ - csr_col_ind_C, \ - csr_val_C, \ - base_A, \ - base_B, \ - descr_C->base, \ - base_D, \ - info_C->csrgemm_info->mul, \ - info_C->csrgemm_info->add, \ - handle->pointer_mode == rocsparse_pointer_mode_host); +#define CSRGEMM_FILL_BLOCK_PER_ROW_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, GS) \ + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ + (rocsparse::csrgemm_fill_block_per_row), \ + dim3(rocsparse::get_grid_size_x(handle, h_group_size[GROUP_SIZE_ID], CSRGEMM_DIM)), \ + dim3(CSRGEMM_DIM), \ + (csrgemm_fill_block_per_row_shared_memory_size()), \ + stream, \ + h_group_size[GROUP_SIZE_ID], \ + rocsparse::max(k, n), \ + &d_group_offset[GROUP_SIZE_ID], \ + d_perm, \ + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), \ + csr_row_ptr_A, \ + csr_col_ind_A, \ + csr_val_A, \ + csr_row_ptr_B, \ + csr_col_ind_B, \ + csr_val_B, \ + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), \ + csr_row_ptr_D, \ + csr_col_ind_D, \ + csr_val_D, \ + csr_row_ptr_C, \ + csr_col_ind_C, \ + csr_val_C, \ + base_A, \ + base_B, \ + descr_C->base, \ + base_D, \ + info_C->csrgemm_info->mul, \ + info_C->csrgemm_info->add, \ + handle->pointer_mode == rocsparse_pointer_mode_host) + +// Dispatch the straight-line variant unless the clamp binds (AISPARSE-677). +#define CSRGEMM_FILL_BLOCK_PER_ROW( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ + if(rocsparse::get_grid_size_x(handle, h_group_size[GROUP_SIZE_ID], CSRGEMM_DIM) \ + < static_cast(h_group_size[GROUP_SIZE_ID])) \ + { \ + CSRGEMM_FILL_BLOCK_PER_ROW_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, true); \ + } \ + else \ + { \ + CSRGEMM_FILL_BLOCK_PER_ROW_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, false); \ + } #define CSRGEMM_FILL_BLOCK_PER_ROW_2( \ GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ @@ -559,6 +592,19 @@ rocsparse_status rocsparse::csrgemm_calc_template(rocsparse_handle hand CSRGEMM_HASHSIZE, \ CSRGEMM_FLL_HASH, \ CSRGEMM_WARPSIZE, \ + true, \ + I, \ + J, \ + T>, \ + hipFuncAttributeMaxDynamicSharedMemorySize, \ + csrgemm_fill_block_per_row_shared_memory_size())); \ + RETURN_IF_HIP_ERROR(hipFuncSetAttribute( \ + (const void*)rocsparse::csrgemm_fill_block_per_row, \ @@ -568,42 +614,58 @@ rocsparse_status rocsparse::csrgemm_calc_template(rocsparse_handle hand CSRGEMM_FILL_BLOCK_PER_ROW( \ GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) -#define CSRGEMM_FILL_BLOCK_PER_ROW_MULTIPASS( \ - GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ - (rocsparse::csrgemm_fill_block_per_row_multipass), \ - dim3(h_group_size[10]), \ - dim3(CSRGEMM_DIM), \ - 0, \ - stream, \ - n, \ - &d_group_offset[10], \ - d_perm, \ - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), \ - csr_row_ptr_A, \ - csr_col_ind_A, \ - csr_val_A, \ - csr_row_ptr_B, \ - csr_col_ind_B, \ - csr_val_B, \ - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), \ - csr_row_ptr_D, \ - csr_col_ind_D, \ - csr_val_D, \ - csr_row_ptr_C, \ - csr_col_ind_C, \ - csr_val_C, \ - workspace_B, \ - base_A, \ - base_B, \ - descr_C->base, \ - base_D, \ - info_C->csrgemm_info->mul, \ - info_C->csrgemm_info->add, \ - handle->pointer_mode == rocsparse_pointer_mode_host); +#define CSRGEMM_FILL_BLOCK_PER_ROW_MULTIPASS_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, GS) \ + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ + (rocsparse::csrgemm_fill_block_per_row_multipass), \ + dim3(rocsparse::get_grid_size_x(handle, h_group_size[10], CSRGEMM_DIM)), \ + dim3(CSRGEMM_DIM), \ + 0, \ + stream, \ + h_group_size[10], \ + n, \ + &d_group_offset[10], \ + d_perm, \ + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), \ + csr_row_ptr_A, \ + csr_col_ind_A, \ + csr_val_A, \ + csr_row_ptr_B, \ + csr_col_ind_B, \ + csr_val_B, \ + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), \ + csr_row_ptr_D, \ + csr_col_ind_D, \ + csr_val_D, \ + csr_row_ptr_C, \ + csr_col_ind_C, \ + csr_val_C, \ + workspace_B, \ + base_A, \ + base_B, \ + descr_C->base, \ + base_D, \ + info_C->csrgemm_info->mul, \ + info_C->csrgemm_info->add, \ + handle->pointer_mode == rocsparse_pointer_mode_host) + +#define CSRGEMM_FILL_BLOCK_PER_ROW_MULTIPASS( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ + if(rocsparse::get_grid_size_x(handle, h_group_size[10], CSRGEMM_DIM) \ + < static_cast(h_group_size[10])) \ + { \ + CSRGEMM_FILL_BLOCK_PER_ROW_MULTIPASS_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, true); \ + } \ + else \ + { \ + CSRGEMM_FILL_BLOCK_PER_ROW_MULTIPASS_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, false); \ + } // Group 2: 33 - 256 non-zeros per row if(h_group_size[2] > 0) diff --git a/projects/rocsparse/library/src/extra/rocsparse_csrgemm_nnz_calc.cpp b/projects/rocsparse/library/src/extra/rocsparse_csrgemm_nnz_calc.cpp index 2eb8d3148e42..a66093fbaccf 100644 --- a/projects/rocsparse/library/src/extra/rocsparse_csrgemm_nnz_calc.cpp +++ b/projects/rocsparse/library/src/extra/rocsparse_csrgemm_nnz_calc.cpp @@ -27,6 +27,7 @@ #include "csrgemm_device.h" #include "rocsparse_control.hpp" #include "rocsparse_csrgemm.hpp" +#include "rocsparse_grid.hpp" #include "rocsparse_utility.hpp" #include "rocsparse_primitives.hpp" @@ -92,21 +93,30 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, // Compute number of intermediate products for each row #define CSRGEMM_DIM 256 #define CSRGEMM_SUB 8 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_intermediate_products), - dim3((m - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1), - dim3(CSRGEMM_DIM), - 0, - stream, - m, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_row_ptr_D, - csr_row_ptr_C, - base_A, - mul, - add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (m - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1, + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_intermediate_products), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + stream, + m, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_row_ptr_D, + csr_row_ptr_C, + base_A, + mul, + add); + return rocsparse_status_success; + })); #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -234,30 +244,38 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 128 #define CSRGEMM_SUB 4 #define CSRGEMM_HASHSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_wf_per_row), - dim3((h_group_size[0] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1), - dim3(CSRGEMM_DIM), - 0, - stream, - h_group_size[0], - &d_group_offset[0], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - mul, - add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (h_group_size[0] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1, + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_wf_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + stream, + h_group_size[0], + &d_group_offset[0], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + mul, + add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -269,30 +287,38 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 256 #define CSRGEMM_SUB 8 #define CSRGEMM_HASHSIZE 64 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_wf_per_row), - dim3((h_group_size[1] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1), - dim3(CSRGEMM_DIM), - 0, - stream, - h_group_size[1], - &d_group_offset[1], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - mul, - add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (h_group_size[1] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1, + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_wf_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + stream, + h_group_size[1], + &d_group_offset[1], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + mul, + add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -304,29 +330,38 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 128 #define CSRGEMM_SUB 8 #define CSRGEMM_HASHSIZE 512 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_block_per_row), - dim3(h_group_size[2]), - dim3(CSRGEMM_DIM), - (csrgemm_nnz_block_per_row_shared_memory_size()), - stream, - &d_group_offset[2], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[2], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_block_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + (csrgemm_nnz_block_per_row_shared_memory_size()), + stream, + h_group_size[2], + &d_group_offset[2], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -338,29 +373,38 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 128 #define CSRGEMM_SUB 8 #define CSRGEMM_HASHSIZE 1024 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_block_per_row), - dim3(h_group_size[3]), - dim3(CSRGEMM_DIM), - (csrgemm_nnz_block_per_row_shared_memory_size()), - stream, - &d_group_offset[3], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[3], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_block_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + (csrgemm_nnz_block_per_row_shared_memory_size()), + stream, + h_group_size[3], + &d_group_offset[3], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -372,29 +416,38 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 256 #define CSRGEMM_SUB 16 #define CSRGEMM_HASHSIZE 2048 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_block_per_row), - dim3(h_group_size[4]), - dim3(CSRGEMM_DIM), - (csrgemm_nnz_block_per_row_shared_memory_size()), - stream, - &d_group_offset[4], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[4], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_block_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + (csrgemm_nnz_block_per_row_shared_memory_size()), + stream, + h_group_size[4], + &d_group_offset[4], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -406,29 +459,38 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 512 #define CSRGEMM_SUB 16 #define CSRGEMM_HASHSIZE 4096 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_block_per_row), - dim3(h_group_size[5]), - dim3(CSRGEMM_DIM), - (csrgemm_nnz_block_per_row_shared_memory_size()), - stream, - &d_group_offset[5], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[5], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_block_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + (csrgemm_nnz_block_per_row_shared_memory_size()), + stream, + h_group_size[5], + &d_group_offset[5], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -440,29 +502,38 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 1024 #define CSRGEMM_SUB 32 #define CSRGEMM_HASHSIZE 8192 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_block_per_row), - dim3(h_group_size[6]), - dim3(CSRGEMM_DIM), - (csrgemm_nnz_block_per_row_shared_memory_size()), - stream, - &d_group_offset[6], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[6], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_block_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + (csrgemm_nnz_block_per_row_shared_memory_size()), + stream, + h_group_size[6], + &d_group_offset[6], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -474,39 +545,49 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 1024 #define CSRGEMM_SUB 32 #define CSRGEMM_HASHSIZE 16384 - RETURN_IF_HIP_ERROR(hipFuncSetAttribute( - (const void*)rocsparse::csrgemm_nnz_block_per_row, - hipFuncAttributeMaxDynamicSharedMemorySize, - csrgemm_nnz_block_per_row_shared_memory_size())); - - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_block_per_row), - dim3(h_group_size[7]), - dim3(CSRGEMM_DIM), - (csrgemm_nnz_block_per_row_shared_memory_size()), - handle->stream, - &d_group_offset[7], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[7], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIP_ERROR(hipFuncSetAttribute( + (const void*)rocsparse::csrgemm_nnz_block_per_row, + hipFuncAttributeMaxDynamicSharedMemorySize, + csrgemm_nnz_block_per_row_shared_memory_size())); + + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_block_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + (csrgemm_nnz_block_per_row_shared_memory_size()), + handle->stream, + h_group_size[7], + &d_group_offset[7], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -518,39 +599,49 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 1024 #define CSRGEMM_SUB 32 #define CSRGEMM_HASHSIZE 32768 - RETURN_IF_HIP_ERROR(hipFuncSetAttribute( - (const void*)rocsparse::csrgemm_nnz_block_per_row, - hipFuncAttributeMaxDynamicSharedMemorySize, - csrgemm_nnz_block_per_row_shared_memory_size())); - - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_block_per_row), - dim3(h_group_size[8]), - dim3(CSRGEMM_DIM), - (csrgemm_nnz_block_per_row_shared_memory_size()), - handle->stream, - &d_group_offset[8], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[8], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIP_ERROR(hipFuncSetAttribute( + (const void*)rocsparse::csrgemm_nnz_block_per_row, + hipFuncAttributeMaxDynamicSharedMemorySize, + csrgemm_nnz_block_per_row_shared_memory_size())); + + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_block_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + (csrgemm_nnz_block_per_row_shared_memory_size()), + handle->stream, + h_group_size[8], + &d_group_offset[8], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -562,39 +653,49 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, #define CSRGEMM_DIM 1024 #define CSRGEMM_SUB 32 #define CSRGEMM_HASHSIZE 65536 - RETURN_IF_HIP_ERROR(hipFuncSetAttribute( - (const void*)rocsparse::csrgemm_nnz_block_per_row, - hipFuncAttributeMaxDynamicSharedMemorySize, - csrgemm_nnz_block_per_row_shared_memory_size())); - - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_nnz_block_per_row), - dim3(h_group_size[9]), - dim3(CSRGEMM_DIM), - (csrgemm_nnz_block_per_row_shared_memory_size()), - handle->stream, - &d_group_offset[9], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - base_A, - base_B, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[9], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIP_ERROR(hipFuncSetAttribute( + (const void*)rocsparse::csrgemm_nnz_block_per_row, + hipFuncAttributeMaxDynamicSharedMemorySize, + csrgemm_nnz_block_per_row_shared_memory_size())); + + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_block_per_row), + dim3(grid), + dim3(CSRGEMM_DIM), + (csrgemm_nnz_block_per_row_shared_memory_size()), + handle->stream, + h_group_size[9], + &d_group_offset[9], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + base_A, + base_B, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -623,29 +724,39 @@ rocsparse_status rocsparse::csrgemm_nnz_calc(rocsparse_handle handle, rocsparse_hipMallocAsync(&workspace_B, sizeof(I) * nnz_A, handle->stream)); } - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse:: - csrgemm_nnz_block_per_row_multipass), - dim3(h_group_size[10]), - dim3(CSRGEMM_DIM), - 0, - stream, - n, - &d_group_offset[10], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - workspace_B, - base_A, - base_B, - base_D, - mul, - add); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[10], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_nnz_block_per_row_multipass), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + stream, + h_group_size[10], + n, + &d_group_offset[10], + d_perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + workspace_B, + base_A, + base_B, + base_D, + mul, + add); + return rocsparse_status_success; + })); if(info_C->csrgemm_info->mul == true) { diff --git a/projects/rocsparse/library/src/extra/rocsparse_csrgemm_numeric_calc.cpp b/projects/rocsparse/library/src/extra/rocsparse_csrgemm_numeric_calc.cpp index 2bb23b2bc212..954f0ac0145d 100644 --- a/projects/rocsparse/library/src/extra/rocsparse_csrgemm_numeric_calc.cpp +++ b/projects/rocsparse/library/src/extra/rocsparse_csrgemm_numeric_calc.cpp @@ -30,6 +30,7 @@ #include "csrgemm_numeric_device.h" #include "rocsparse_control.hpp" +#include "rocsparse_grid.hpp" #include "rocsparse_utility.hpp" namespace rocsparse @@ -38,6 +39,7 @@ namespace rocsparse uint32_t WFSIZE, uint32_t HASHSIZE, uint32_t HASHVAL, + bool GRID_STRIDE, typename I, typename J, typename T> @@ -70,31 +72,44 @@ namespace rocsparse { ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::csrgemm_numeric_fill_wf_per_row_device( - m, - nk, - offset, - perm, - alpha, - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - beta, - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the (sub)wavefront rows so a grid clamped by get_grid_size_x + // covers all rows + for(int64_t block_offset = static_cast(hipBlockIdx_x) * (BLOCKSIZE / WFSIZE); + block_offset < m; + block_offset += static_cast(hipGridDim_x) * (BLOCKSIZE / WFSIZE)) + { + rocsparse::csrgemm_numeric_fill_wf_per_row_device( + static_cast(block_offset), + m, + nk, + offset, + perm, + alpha, + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + beta, + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(!GRID_STRIDE) + { + break; + } + } } template @@ -108,11 +123,13 @@ namespace rocsparse uint32_t HASHSIZE, uint32_t HASHVAL, uint32_t WARPSIZE, + bool GRID_STRIDE, typename I, typename J, typename T> ROCSPARSE_KERNEL(BLOCKSIZE) - void csrgemm_numeric_fill_block_per_row_kernel(J nk, + void csrgemm_numeric_fill_block_per_row_kernel(J size, + J nk, const J* __restrict__ offset, const J* __restrict__ perm, ROCSPARSE_DEVICE_HOST_SCALAR_PARAMS(T, alpha), @@ -139,44 +156,61 @@ namespace rocsparse { ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::csrgemm_numeric_fill_block_per_row_device(nk, - offset, - perm, - alpha, - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - beta, - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) + { + rocsparse::csrgemm_numeric_fill_block_per_row_device(static_cast(block_id), + nk, + offset, + perm, + alpha, + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + beta, + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(GRID_STRIDE) + { + // The next row re-initialises the shared hash table read above + __syncthreads(); + } + else + { + break; + } + } } template ROCSPARSE_KERNEL(BLOCKSIZE) void csrgemm_numeric_fill_block_per_row_multipass_kernel( + J size, J n, const J* __restrict__ offset, const J* __restrict__ perm, @@ -206,33 +240,44 @@ namespace rocsparse ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(mul, alpha); ROCSPARSE_DEVICE_HOST_SCALAR_GET_IF(add, beta); - rocsparse::csrgemm_numeric_fill_block_per_row_multipass_device(n, - offset, - perm, - alpha, - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - beta, - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - workspace_B, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) + { + rocsparse::csrgemm_numeric_fill_block_per_row_multipass_device( + static_cast(block_id), + n, + offset, + perm, + alpha, + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + beta, + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + workspace_B, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(!GRID_STRIDE) + { + break; + } + } } template , + hipFuncAttributeMaxDynamicSharedMemorySize, + csrgemm_numeric_fill_block_per_row_kernel_shared_memory_size())); + RETURN_IF_HIP_ERROR(hipFuncSetAttribute( + (const void*)rocsparse::csrgemm_numeric_fill_block_per_row_kernel, @@ -292,42 +352,53 @@ namespace rocsparse J, T>())); - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_numeric_fill_block_per_row_kernel), - dim3(group_size), - dim3(CSRGEMM_DIM), - (csrgemm_numeric_fill_block_per_row_kernel_shared_memory_size()), - handle->stream, - rocsparse::max(k, n), - group_offset, - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - base_A, - base_B, - base_C, - base_D, - mul, - add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size, + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_numeric_fill_block_per_row_kernel< + CSRGEMM_DIM, + CSRGEMM_SUB, + CSRGEMM_HASHSIZE, + CSRGEMM_FLL_HASH, + CSRGEMM_WARPSIZE, + decltype(grid_stride)::value>), + dim3(grid), + dim3(CSRGEMM_DIM), + (csrgemm_numeric_fill_block_per_row_kernel_shared_memory_size< + CSRGEMM_HASHSIZE, + J, + T>()), + handle->stream, + group_size, + rocsparse::max(k, n), + group_offset, + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + base_A, + base_B, + base_C, + base_D, + mul, + add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef CSRGEMM_SUB #undef CSRGEMM_DIM } @@ -354,79 +425,97 @@ namespace rocsparse if(handle->wavefront_size == 32) { - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_numeric_fill_block_per_row_multipass_kernel< - CSRGEMM_DIM, - CSRGEMM_SUB, - CSRGEMM_CHUNKSIZE, - 32>), - dim3(group_size), - dim3(CSRGEMM_DIM), - 0, - handle->stream, - n, - group_offset, - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - workspace_B, - base_A, - base_B, - base_C, - base_D, - mul, - add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size, + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_numeric_fill_block_per_row_multipass_kernel< + CSRGEMM_DIM, + CSRGEMM_SUB, + CSRGEMM_CHUNKSIZE, + 32, + decltype(grid_stride)::value>), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + handle->stream, + group_size, + n, + group_offset, + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + workspace_B, + base_A, + base_B, + base_C, + base_D, + mul, + add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } else { - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_numeric_fill_block_per_row_multipass_kernel< - CSRGEMM_DIM, - CSRGEMM_SUB, - CSRGEMM_CHUNKSIZE, - 64>), - dim3(group_size), - dim3(CSRGEMM_DIM), - 0, - handle->stream, - n, - group_offset, - perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - workspace_B, - base_A, - base_B, - base_C, - base_D, - mul, - add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + group_size, + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_numeric_fill_block_per_row_multipass_kernel< + CSRGEMM_DIM, + CSRGEMM_SUB, + CSRGEMM_CHUNKSIZE, + 64, + decltype(grid_stride)::value>), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + handle->stream, + group_size, + n, + group_offset, + perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + workspace_B, + base_A, + base_B, + base_C, + base_D, + mul, + add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } if(mul) { @@ -515,18 +604,20 @@ rocsparse_status rocsparse::csrgemm_numeric_calc_template(rocsparse_handle ha rocsparse_index_base base_D = info_C->csrgemm_info->add ? descr_D->base : rocsparse_index_base_zero; -#define CSRGEMM_NUMERIC_FILL_BLOCK_PER_ROW( \ - GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ +#define CSRGEMM_NUMERIC_FILL_BLOCK_PER_ROW_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, GS) \ RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ (rocsparse::csrgemm_numeric_fill_block_per_row_kernel), \ - dim3(h_group_size[GROUP_SIZE_ID]), \ + CSRGEMM_WARPSIZE, \ + GS>), \ + dim3(rocsparse::get_grid_size_x(handle, h_group_size[GROUP_SIZE_ID], CSRGEMM_DIM)), \ dim3(CSRGEMM_DIM), \ (csrgemm_numeric_fill_block_per_row_kernel_shared_memory_size()), \ stream, \ + h_group_size[GROUP_SIZE_ID], \ rocsparse::max(k, n), \ &d_group_offset[GROUP_SIZE_ID], \ d_perm, \ @@ -550,7 +641,22 @@ rocsparse_status rocsparse::csrgemm_numeric_calc_template(rocsparse_handle ha base_D, \ info_C->csrgemm_info->mul, \ info_C->csrgemm_info->add, \ - handle->pointer_mode == rocsparse_pointer_mode_host); + handle->pointer_mode == rocsparse_pointer_mode_host) + +// Dispatch the straight-line variant unless the clamp binds (AISPARSE-677). +#define CSRGEMM_NUMERIC_FILL_BLOCK_PER_ROW( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ + if(rocsparse::get_grid_size_x(handle, h_group_size[GROUP_SIZE_ID], CSRGEMM_DIM) \ + < static_cast(h_group_size[GROUP_SIZE_ID])) \ + { \ + CSRGEMM_NUMERIC_FILL_BLOCK_PER_ROW_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, true); \ + } \ + else \ + { \ + CSRGEMM_NUMERIC_FILL_BLOCK_PER_ROW_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, false); \ + } #define CSRGEMM_NUMERIC_LAUNCHER(GROUP_SIZE_ID, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ RETURN_IF_ROCSPARSE_ERROR( \ @@ -592,40 +698,49 @@ rocsparse_status rocsparse::csrgemm_numeric_calc_template(rocsparse_handle ha #define CSRGEMM_DIM 256 #define CSRGEMM_SUB 8 #define CSRGEMM_HASHSIZE 16 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_numeric_fill_wf_per_row_kernel), - dim3((h_group_size[0] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1), - dim3(CSRGEMM_DIM), - 0, - stream, - h_group_size[0], - rocsparse::max(k, n), - &d_group_offset[0], - d_perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (h_group_size[0] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1, + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_numeric_fill_wf_per_row_kernel< + CSRGEMM_DIM, + CSRGEMM_SUB, + CSRGEMM_HASHSIZE, + CSRGEMM_FLL_HASH, + decltype(grid_stride)::value>), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + stream, + h_group_size[0], + rocsparse::max(k, n), + &d_group_offset[0], + d_perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -637,40 +752,49 @@ rocsparse_status rocsparse::csrgemm_numeric_calc_template(rocsparse_handle ha #define CSRGEMM_DIM 256 #define CSRGEMM_SUB 16 #define CSRGEMM_HASHSIZE 32 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_numeric_fill_wf_per_row_kernel), - dim3((h_group_size[1] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1), - dim3(CSRGEMM_DIM), - 0, - stream, - h_group_size[1], - rocsparse::max(k, n), - &d_group_offset[1], - d_perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + (h_group_size[1] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1, + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_numeric_fill_wf_per_row_kernel< + CSRGEMM_DIM, + CSRGEMM_SUB, + CSRGEMM_HASHSIZE, + CSRGEMM_FLL_HASH, + decltype(grid_stride)::value>), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + stream, + h_group_size[1], + rocsparse::max(k, n), + &d_group_offset[1], + d_perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); #undef CSRGEMM_HASHSIZE #undef CSRGEMM_SUB #undef CSRGEMM_DIM @@ -829,77 +953,97 @@ rocsparse_status rocsparse::csrgemm_numeric_calc_template(rocsparse_handle ha if(handle->wavefront_size == 32) { - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_numeric_fill_block_per_row_multipass_kernel), - dim3(h_group_size[10]), - dim3(CSRGEMM_DIM), - 0, - stream, - n, - &d_group_offset[10], - d_perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - workspace_B, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[10], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_numeric_fill_block_per_row_multipass_kernel< + CSRGEMM_DIM, + CSRGEMM_SUB, + CSRGEMM_CHUNKSIZE, + 32, + decltype(grid_stride)::value>), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + stream, + h_group_size[10], + n, + &d_group_offset[10], + d_perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + workspace_B, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } else { - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_numeric_fill_block_per_row_multipass_kernel), - dim3(h_group_size[10]), - dim3(CSRGEMM_DIM), - 0, - stream, - n, - &d_group_offset[10], - d_perm, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), - csr_row_ptr_A, - csr_col_ind_A, - csr_val_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_val_B, - ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), - csr_row_ptr_D, - csr_col_ind_D, - csr_val_D, - csr_row_ptr_C, - csr_col_ind_C, - csr_val_C, - workspace_B, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add, - handle->pointer_mode == rocsparse_pointer_mode_host); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::dispatch_grid_stride_x( + handle, + h_group_size[10], + CSRGEMM_DIM, + [&](auto grid_stride, uint32_t grid) -> rocsparse_status { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( + (rocsparse::csrgemm_numeric_fill_block_per_row_multipass_kernel< + CSRGEMM_DIM, + CSRGEMM_SUB, + CSRGEMM_CHUNKSIZE, + 64, + decltype(grid_stride)::value>), + dim3(grid), + dim3(CSRGEMM_DIM), + 0, + stream, + h_group_size[10], + n, + &d_group_offset[10], + d_perm, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, alpha_device_host), + csr_row_ptr_A, + csr_col_ind_A, + csr_val_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_val_B, + ROCSPARSE_DEVICE_HOST_SCALAR_PERMISSIVE_ARGS(handle, beta_device_host), + csr_row_ptr_D, + csr_col_ind_D, + csr_val_D, + csr_row_ptr_C, + csr_col_ind_C, + csr_val_C, + workspace_B, + base_A, + base_B, + descr_C->base, + base_D, + info_C->csrgemm_info->mul, + info_C->csrgemm_info->add, + handle->pointer_mode == rocsparse_pointer_mode_host); + return rocsparse_status_success; + })); } if(info_C->csrgemm_info->mul) diff --git a/projects/rocsparse/library/src/extra/rocsparse_csrgemm_symbolic_calc.cpp b/projects/rocsparse/library/src/extra/rocsparse_csrgemm_symbolic_calc.cpp index 17a51bb03627..05f9d196a9dd 100644 --- a/projects/rocsparse/library/src/extra/rocsparse_csrgemm_symbolic_calc.cpp +++ b/projects/rocsparse/library/src/extra/rocsparse_csrgemm_symbolic_calc.cpp @@ -29,6 +29,7 @@ #include "rocsparse_csrgemm.hpp" #include "rocsparse_control.hpp" +#include "rocsparse_grid.hpp" #include "rocsparse_utility.hpp" #include "csrgemm_symbolic_device.h" @@ -40,10 +41,12 @@ namespace rocsparse uint32_t WFSIZE, uint32_t CHUNKSIZE, uint32_t WARPSIZE, + bool GRID_STRIDE, typename I, typename J> ROCSPARSE_KERNEL(BLOCKSIZE) - void csrgemm_symbolic_fill_block_per_row_multipass(J n, + void csrgemm_symbolic_fill_block_per_row_multipass(J size, + J n, const J* __restrict__ offset, const J* __restrict__ perm, const I* __restrict__ csr_row_ptr_A, @@ -62,33 +65,45 @@ namespace rocsparse bool mul, bool add) { - rocsparse::csrgemm_symbolic_fill_block_per_row_multipass_device(n, - offset, - perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - csr_col_ind_C, - workspace_B, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) + { + rocsparse::csrgemm_symbolic_fill_block_per_row_multipass_device( + static_cast(block_id), + n, + offset, + perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + csr_col_ind_C, + workspace_B, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(!GRID_STRIDE) + { + break; + } + } } template ROCSPARSE_KERNEL(BLOCKSIZE) @@ -116,25 +131,39 @@ namespace rocsparse bool mul, bool add) { - rocsparse::csrgemm_symbolic_fill_wf_per_row_device( - m, - nk, - offset, - perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - csr_col_ind_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the (sub)wavefront rows so a grid clamped by get_grid_size_x + // covers all rows + for(int64_t block_offset = static_cast(hipBlockIdx_x) * (BLOCKSIZE / WFSIZE); + block_offset < m; + block_offset += static_cast(hipGridDim_x) * (BLOCKSIZE / WFSIZE)) + { + rocsparse:: + csrgemm_symbolic_fill_wf_per_row_device( + static_cast(block_offset), + m, + nk, + offset, + perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + csr_col_ind_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(!GRID_STRIDE) + { + break; + } + } } template @@ -148,10 +177,12 @@ namespace rocsparse uint32_t HASHSIZE, uint32_t HASHVAL, uint32_t WARPSIZE, + bool GRID_STRIDE, typename I, typename J> ROCSPARSE_KERNEL(BLOCKSIZE) - void csrgemm_symbolic_fill_block_per_row(J nk, + void csrgemm_symbolic_fill_block_per_row(J size, + J nk, const J* __restrict__ offset, const J* __restrict__ perm, const I* __restrict__ csr_row_ptr_A, @@ -169,27 +200,43 @@ namespace rocsparse bool mul, bool add) { - rocsparse::csrgemm_symbolic_fill_block_per_row_device(nk, - offset, - perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - csr_col_ind_C, - idx_base_A, - idx_base_B, - idx_base_C, - idx_base_D, - mul, - add); + // Grid-stride over the block rows so a grid clamped by get_grid_size_x covers all rows + for(int64_t block_id = hipBlockIdx_x; block_id < size; block_id += hipGridDim_x) + { + rocsparse::csrgemm_symbolic_fill_block_per_row_device( + static_cast(block_id), + nk, + offset, + perm, + csr_row_ptr_A, + csr_col_ind_A, + csr_row_ptr_B, + csr_col_ind_B, + csr_row_ptr_D, + csr_col_ind_D, + csr_row_ptr_C, + csr_col_ind_C, + idx_base_A, + idx_base_B, + idx_base_C, + idx_base_D, + mul, + add); + + if constexpr(GRID_STRIDE) + { + // The next row re-initialises the shared hash table read above + __syncthreads(); + } + else + { + break; + } + } } } @@ -405,111 +452,109 @@ rocsparse_status rocsparse::csrgemm_symbolic_calc_template(rocsparse_handle rocsparse_index_base base_D = info_C->csrgemm_info->add ? descr_D->base : rocsparse_index_base_zero; +#define CSRGEMM_SYMBOLIC_FILL_WF_PER_ROW_IMPL( \ + GROUP_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, GS) \ + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ + (rocsparse::csrgemm_symbolic_fill_wf_per_row), \ + dim3(rocsparse::get_grid_size_x( \ + handle, (h_group_size[GROUP_ID] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1, CSRGEMM_DIM)), \ + dim3(CSRGEMM_DIM), \ + 0, \ + stream, \ + h_group_size[GROUP_ID], \ + rocsparse::max(k, n), \ + &d_group_offset[GROUP_ID], \ + d_perm, \ + csr_row_ptr_A, \ + csr_col_ind_A, \ + csr_row_ptr_B, \ + csr_col_ind_B, \ + csr_row_ptr_D, \ + csr_col_ind_D, \ + csr_row_ptr_C, \ + csr_col_ind_C, \ + base_A, \ + base_B, \ + descr_C->base, \ + base_D, \ + info_C->csrgemm_info->mul, \ + info_C->csrgemm_info->add) + +#define CSRGEMM_SYMBOLIC_FILL_WF_PER_ROW(GROUP_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE) \ + if(rocsparse::get_grid_size_x( \ + handle, (h_group_size[GROUP_ID] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1, CSRGEMM_DIM) \ + < static_cast((h_group_size[GROUP_ID] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1)) \ + { \ + CSRGEMM_SYMBOLIC_FILL_WF_PER_ROW_IMPL( \ + GROUP_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, true); \ + } \ + else \ + { \ + CSRGEMM_SYMBOLIC_FILL_WF_PER_ROW_IMPL( \ + GROUP_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, false); \ + } + // Group 0: 0 - 16 non-zeros per row if(h_group_size[0] > 0) { -#define CSRGEMM_DIM 256 -#define CSRGEMM_SUB 8 -#define CSRGEMM_HASHSIZE 16 - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_symbolic_fill_wf_per_row), - dim3((h_group_size[0] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1), - dim3(CSRGEMM_DIM), - 0, - stream, - h_group_size[0], - rocsparse::max(k, n), - &d_group_offset[0], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - csr_col_ind_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); -#undef CSRGEMM_HASHSIZE -#undef CSRGEMM_SUB -#undef CSRGEMM_DIM + CSRGEMM_SYMBOLIC_FILL_WF_PER_ROW(0, 256, 8, 16); } // Group 1: 17 - 32 non-zeros per row if(h_group_size[1] > 0) { -#define CSRGEMM_DIM 256 -#define CSRGEMM_SUB 16 -#define CSRGEMM_HASHSIZE 32 - - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( - (rocsparse::csrgemm_symbolic_fill_wf_per_row), - dim3((h_group_size[1] - 1) / (CSRGEMM_DIM / CSRGEMM_SUB) + 1), - dim3(CSRGEMM_DIM), - 0, - stream, - h_group_size[1], - rocsparse::max(k, n), - &d_group_offset[1], - d_perm, - csr_row_ptr_A, - csr_col_ind_A, - csr_row_ptr_B, - csr_col_ind_B, - csr_row_ptr_D, - csr_col_ind_D, - csr_row_ptr_C, - csr_col_ind_C, - base_A, - base_B, - descr_C->base, - base_D, - info_C->csrgemm_info->mul, - info_C->csrgemm_info->add); -#undef CSRGEMM_HASHSIZE -#undef CSRGEMM_SUB -#undef CSRGEMM_DIM + CSRGEMM_SYMBOLIC_FILL_WF_PER_ROW(1, 256, 16, 32); + } +#define CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, GS) \ + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ + (rocsparse::csrgemm_symbolic_fill_block_per_row), \ + dim3(rocsparse::get_grid_size_x(handle, h_group_size[GROUP_SIZE_ID], CSRGEMM_DIM)), \ + dim3(CSRGEMM_DIM), \ + (csrgemm_symbolic_fill_block_per_row_shared_memory_size()), \ + handle->stream, \ + h_group_size[GROUP_SIZE_ID], \ + rocsparse::max(k, n), \ + &d_group_offset[GROUP_SIZE_ID], \ + d_perm, \ + csr_row_ptr_A, \ + csr_col_ind_A, \ + csr_row_ptr_B, \ + csr_col_ind_B, \ + csr_row_ptr_D, \ + csr_col_ind_D, \ + csr_row_ptr_C, \ + csr_col_ind_C, \ + base_A, \ + base_B, \ + descr_C->base, \ + base_D, \ + info_C->csrgemm_info->mul, \ + info_C->csrgemm_info->add) + +// Dispatch the straight-line variant unless the clamp binds (AISPARSE-677). +#define CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ + if(rocsparse::get_grid_size_x(handle, h_group_size[GROUP_SIZE_ID], CSRGEMM_DIM) \ + < static_cast(h_group_size[GROUP_SIZE_ID])) \ + { \ + CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, true); \ + } \ + else \ + { \ + CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, false); \ } -#define CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW( \ - GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ - (rocsparse::csrgemm_symbolic_fill_block_per_row), \ - dim3(h_group_size[GROUP_SIZE_ID]), \ - dim3(CSRGEMM_DIM), \ - (csrgemm_symbolic_fill_block_per_row_shared_memory_size()), \ - handle->stream, \ - rocsparse::max(k, n), \ - &d_group_offset[GROUP_SIZE_ID], \ - d_perm, \ - csr_row_ptr_A, \ - csr_col_ind_A, \ - csr_row_ptr_B, \ - csr_col_ind_B, \ - csr_row_ptr_D, \ - csr_col_ind_D, \ - csr_row_ptr_C, \ - csr_col_ind_C, \ - base_A, \ - base_B, \ - descr_C->base, \ - base_D, \ - info_C->csrgemm_info->mul, \ - info_C->csrgemm_info->add); #define CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW_2( \ GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ RETURN_IF_HIP_ERROR(hipFuncSetAttribute( \ @@ -518,6 +563,18 @@ rocsparse_status rocsparse::csrgemm_symbolic_calc_template(rocsparse_handle CSRGEMM_HASHSIZE, \ CSRGEMM_FLL_HASH, \ CSRGEMM_WARPSIZE, \ + true, \ + I, \ + J>, \ + hipFuncAttributeMaxDynamicSharedMemorySize, \ + csrgemm_symbolic_fill_block_per_row_shared_memory_size())); \ + RETURN_IF_HIP_ERROR(hipFuncSetAttribute( \ + (const void*)rocsparse::csrgemm_symbolic_fill_block_per_row, \ hipFuncAttributeMaxDynamicSharedMemorySize, \ @@ -525,47 +582,63 @@ rocsparse_status rocsparse::csrgemm_symbolic_calc_template(rocsparse_handle CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW( \ GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) -#define CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW_MULTIPASS( \ - GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ - I* workspace_B = nullptr; \ - \ - if(info_C->csrgemm_info->mul == true) \ - { \ - RETURN_IF_HIP_ERROR( \ - rocsparse_hipMallocAsync(&workspace_B, sizeof(I) * nnz_A, handle->stream)); \ - } \ - \ - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ - (rocsparse::csrgemm_symbolic_fill_block_per_row_multipass), \ - dim3(h_group_size[GROUP_SIZE_ID]), \ - dim3(CSRGEMM_DIM), \ - 0, \ - stream, \ - n, \ - &d_group_offset[GROUP_SIZE_ID], \ - d_perm, \ - csr_row_ptr_A, \ - csr_col_ind_A, \ - csr_row_ptr_B, \ - csr_col_ind_B, \ - csr_row_ptr_D, \ - csr_col_ind_D, \ - csr_row_ptr_C, \ - csr_col_ind_C, \ - workspace_B, \ - base_A, \ - base_B, \ - descr_C->base, \ - base_D, \ - info_C->csrgemm_info->mul, \ - info_C->csrgemm_info->add); \ - \ - if(info_C->csrgemm_info->mul == true) \ - { \ - RETURN_IF_HIP_ERROR(rocsparse_hipFreeAsync(workspace_B, handle->stream)); \ +#define CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW_MULTIPASS_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, GS) \ + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR( \ + (rocsparse::csrgemm_symbolic_fill_block_per_row_multipass), \ + dim3(rocsparse::get_grid_size_x(handle, h_group_size[GROUP_SIZE_ID], CSRGEMM_DIM)), \ + dim3(CSRGEMM_DIM), \ + 0, \ + stream, \ + h_group_size[GROUP_SIZE_ID], \ + n, \ + &d_group_offset[GROUP_SIZE_ID], \ + d_perm, \ + csr_row_ptr_A, \ + csr_col_ind_A, \ + csr_row_ptr_B, \ + csr_col_ind_B, \ + csr_row_ptr_D, \ + csr_col_ind_D, \ + csr_row_ptr_C, \ + csr_col_ind_C, \ + workspace_B, \ + base_A, \ + base_B, \ + descr_C->base, \ + base_D, \ + info_C->csrgemm_info->mul, \ + info_C->csrgemm_info->add) + +#define CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW_MULTIPASS( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE) \ + I* workspace_B = nullptr; \ + \ + if(info_C->csrgemm_info->mul == true) \ + { \ + RETURN_IF_HIP_ERROR( \ + rocsparse_hipMallocAsync(&workspace_B, sizeof(I) * nnz_A, handle->stream)); \ + } \ + \ + if(rocsparse::get_grid_size_x(handle, h_group_size[GROUP_SIZE_ID], CSRGEMM_DIM) \ + < static_cast(h_group_size[GROUP_SIZE_ID])) \ + { \ + CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW_MULTIPASS_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, true); \ + } \ + else \ + { \ + CSRGEMM_SYMBOLIC_FILL_BLOCK_PER_ROW_MULTIPASS_IMPL( \ + GROUP_SIZE_ID, CSRGEMM_DIM, CSRGEMM_SUB, CSRGEMM_HASHSIZE, CSRGEMM_WARPSIZE, false); \ + } \ + \ + if(info_C->csrgemm_info->mul == true) \ + { \ + RETURN_IF_HIP_ERROR(rocsparse_hipFreeAsync(workspace_B, handle->stream)); \ } // Group 2: 33 - 256 non-zeros per row 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); + } }