From 403bc9bf2a9e75d4b1b847d33e4e71ee179795b8 Mon Sep 17 00:00:00 2001 From: Michael Lykov Date: Tue, 7 Jul 2026 17:13:52 +0200 Subject: [PATCH] Add SELL-P SpMM and advanced SpMM kernels More detailed comments: + spmm and advanced_spmm dispatched from apply() when nrhs > 2. + OMP SpMM implementation (Gustavson + parallel + SIMD). + Reference and CUDA/HIP/DPC++ stubs falling back to SpMV kernels. + OMP test comparing against the reference executor. --- common/cuda_hip/matrix/sellp_kernels.cpp | 27 ++++++ core/device_hooks/common_kernels.inc.cpp | 2 + core/matrix/sellp.cpp | 36 +++++-- core/matrix/sellp_kernels.hpp | 19 ++++ dpcpp/matrix/sellp_kernels.dp.cpp | 27 ++++++ omp/matrix/sellp_kernels.cpp | 117 +++++++++++++++++++++++ omp/test/matrix/CMakeLists.txt | 1 + omp/test/matrix/sellp_spmm_kernels.cpp | 116 ++++++++++++++++++++++ reference/matrix/sellp_kernels.cpp | 27 ++++++ 9 files changed, 364 insertions(+), 8 deletions(-) create mode 100644 omp/test/matrix/sellp_spmm_kernels.cpp diff --git a/common/cuda_hip/matrix/sellp_kernels.cpp b/common/cuda_hip/matrix/sellp_kernels.cpp index 389c7c5cdcc..23fa3ec343c 100644 --- a/common/cuda_hip/matrix/sellp_kernels.cpp +++ b/common/cuda_hip/matrix/sellp_kernels.cpp @@ -133,6 +133,33 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( GKO_DECLARE_SELLP_ADVANCED_SPMV_KERNEL); +template +void spmm(std::shared_ptr exec, + matrix::view::sellp a, + matrix::view::dense b, + matrix::view::dense c) +{ + spmv(exec, a, b, c); +} + +GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_SPMM_KERNEL); + + +template +void advanced_spmm(std::shared_ptr exec, + matrix::view::dense alpha, + matrix::view::sellp a, + matrix::view::dense b, + matrix::view::dense beta, + matrix::view::dense c) +{ + advanced_spmv(exec, alpha, a, b, beta, c); +} + +GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( + GKO_DECLARE_SELLP_ADVANCED_SPMM_KERNEL); + + } // namespace sellp } // namespace GKO_DEVICE_NAMESPACE } // namespace kernels diff --git a/core/device_hooks/common_kernels.inc.cpp b/core/device_hooks/common_kernels.inc.cpp index e574d8917fe..e60b0cf7c16 100644 --- a/core/device_hooks/common_kernels.inc.cpp +++ b/core/device_hooks/common_kernels.inc.cpp @@ -933,6 +933,8 @@ GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_SPMV_KERNEL); GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_FILL_IN_MATRIX_DATA_KERNEL); GKO_STUB_INDEX_TYPE(GKO_DECLARE_SELLP_COMPUTE_SLICE_SETS_KERNEL); GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_ADVANCED_SPMV_KERNEL); +GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_SPMM_KERNEL); +GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_ADVANCED_SPMM_KERNEL); GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_FILL_IN_DENSE_KERNEL); GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_CONVERT_TO_CSR_KERNEL); GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_COUNT_NONZEROS_PER_ROW_KERNEL); diff --git a/core/matrix/sellp.cpp b/core/matrix/sellp.cpp index d1d33756f1f..4b9d58fa20d 100644 --- a/core/matrix/sellp.cpp +++ b/core/matrix/sellp.cpp @@ -30,6 +30,8 @@ namespace { GKO_REGISTER_OPERATION(spmv, sellp::spmv); GKO_REGISTER_OPERATION(advanced_spmv, sellp::advanced_spmv); +GKO_REGISTER_OPERATION(spmm, sellp::spmm); +GKO_REGISTER_OPERATION(advanced_spmm, sellp::advanced_spmm); GKO_REGISTER_OPERATION(convert_idxs_to_ptrs, components::convert_idxs_to_ptrs); GKO_REGISTER_OPERATION(prefix_sum_nonnegative, components::prefix_sum_nonnegative); @@ -182,9 +184,17 @@ void Sellp::apply_impl(const LinOp* b, LinOp* x) const { precision_dispatch_real_complex( [this](auto dense_b, auto dense_x) { - this->get_executor()->run(sellp::make_spmv( - this->get_const_device_view(), dense_b->get_const_device_view(), - dense_x->get_device_view())); + if (dense_b->get_size()[1] <= 2) { + this->get_executor()->run( + sellp::make_spmv(this->get_const_device_view(), + dense_b->get_const_device_view(), + dense_x->get_device_view())); + } else { + this->get_executor()->run( + sellp::make_spmm(this->get_const_device_view(), + dense_b->get_const_device_view(), + dense_x->get_device_view())); + } }, b, x); } @@ -196,11 +206,21 @@ void Sellp::apply_impl(const LinOp* alpha, const LinOp* b, { precision_dispatch_real_complex( [this](auto dense_alpha, auto dense_b, auto dense_beta, auto dense_x) { - this->get_executor()->run(sellp::make_advanced_spmv( - dense_alpha->get_const_device_view(), - this->get_const_device_view(), dense_b->get_const_device_view(), - dense_beta->get_const_device_view(), - dense_x->get_device_view())); + if (dense_b->get_size()[1] <= 2) { + this->get_executor()->run(sellp::make_advanced_spmv( + dense_alpha->get_const_device_view(), + this->get_const_device_view(), + dense_b->get_const_device_view(), + dense_beta->get_const_device_view(), + dense_x->get_device_view())); + } else { + this->get_executor()->run(sellp::make_advanced_spmm( + dense_alpha->get_const_device_view(), + this->get_const_device_view(), + dense_b->get_const_device_view(), + dense_beta->get_const_device_view(), + dense_x->get_device_view())); + } }, alpha, b, beta, x); } diff --git a/core/matrix/sellp_kernels.hpp b/core/matrix/sellp_kernels.hpp index 9100568457f..f650414c3f6 100644 --- a/core/matrix/sellp_kernels.hpp +++ b/core/matrix/sellp_kernels.hpp @@ -33,6 +33,21 @@ namespace kernels { matrix::view::dense beta, \ matrix::view::dense c) +#define GKO_DECLARE_SELLP_SPMM_KERNEL(ValueType, IndexType) \ + void spmm(std::shared_ptr exec, \ + matrix::view::sellp a, \ + matrix::view::dense b, \ + matrix::view::dense c) + +#define GKO_DECLARE_SELLP_ADVANCED_SPMM_KERNEL(ValueType, IndexType) \ + void advanced_spmm( \ + std::shared_ptr exec, \ + matrix::view::dense alpha, \ + matrix::view::sellp a, \ + matrix::view::dense b, \ + matrix::view::dense beta, \ + matrix::view::dense c) + #define GKO_DECLARE_SELLP_FILL_IN_MATRIX_DATA_KERNEL(ValueType, IndexType) \ void fill_in_matrix_data( \ std::shared_ptr exec, \ @@ -76,6 +91,10 @@ namespace kernels { template \ GKO_DECLARE_SELLP_ADVANCED_SPMV_KERNEL(ValueType, IndexType); \ template \ + GKO_DECLARE_SELLP_SPMM_KERNEL(ValueType, IndexType); \ + template \ + GKO_DECLARE_SELLP_ADVANCED_SPMM_KERNEL(ValueType, IndexType); \ + template \ GKO_DECLARE_SELLP_FILL_IN_MATRIX_DATA_KERNEL(ValueType, IndexType); \ template \ GKO_DECLARE_SELLP_COMPUTE_SLICE_SETS_KERNEL(IndexType); \ diff --git a/dpcpp/matrix/sellp_kernels.dp.cpp b/dpcpp/matrix/sellp_kernels.dp.cpp index 1c572d4fad5..de3f7f2e40b 100644 --- a/dpcpp/matrix/sellp_kernels.dp.cpp +++ b/dpcpp/matrix/sellp_kernels.dp.cpp @@ -143,6 +143,33 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( GKO_DECLARE_SELLP_ADVANCED_SPMV_KERNEL); +template +void spmm(std::shared_ptr exec, + matrix::view::sellp a, + matrix::view::dense b, + matrix::view::dense c) +{ + spmv(exec, a, b, c); +} + +GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_SPMM_KERNEL); + + +template +void advanced_spmm(std::shared_ptr exec, + matrix::view::dense alpha, + matrix::view::sellp a, + matrix::view::dense b, + matrix::view::dense beta, + matrix::view::dense c) +{ + advanced_spmv(exec, alpha, a, b, beta, c); +} + +GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( + GKO_DECLARE_SELLP_ADVANCED_SPMM_KERNEL); + + } // namespace sellp } // namespace dpcpp } // namespace kernels diff --git a/omp/matrix/sellp_kernels.cpp b/omp/matrix/sellp_kernels.cpp index 805702c461c..0419d24eaeb 100644 --- a/omp/matrix/sellp_kernels.cpp +++ b/omp/matrix/sellp_kernels.cpp @@ -8,6 +8,7 @@ #include +#include #include @@ -198,6 +199,122 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( GKO_DECLARE_SELLP_ADVANCED_SPMV_KERNEL); +// Gustavson's algorithm: parallel over slices/rows, SIMD over dense columns. +template +void spmm(std::shared_ptr exec, + matrix::view::sellp a, + matrix::view::dense b, + matrix::view::dense c) +{ + const auto slice_lengths = a.slice_lengths; + const auto slice_sets = a.slice_sets; + const auto slice_size = a.slice_size; + const auto num_rows = a.size[0]; + const auto num_cols = c.size[1]; + const auto slice_num = ceildiv(num_rows + slice_size - 1, slice_size); + +#pragma omp parallel + { + array row_acc{exec, num_cols}; + auto* row_acc_vals = row_acc.get_data(); + +#pragma omp for collapse(2) schedule(static) + for (size_type slice = 0; slice < slice_num; ++slice) { + for (size_type row = 0; row < slice_size; ++row) { + const auto global_row = slice * slice_size + row; + if (global_row >= num_rows) { + continue; + } + const auto slice_begin = slice_sets[slice]; + const auto slice_length = slice_lengths[slice]; + std::fill_n(row_acc_vals, num_cols, zero()); + for (size_type idx = 0; idx < slice_length; ++idx) { + const auto val = a.val_at(row, slice_begin, idx); + const auto col = a.col_at(row, slice_begin, idx); + if (col == invalid_index()) { + continue; + } +#pragma omp simd + for (size_type j = 0; j < num_cols; ++j) { + row_acc_vals[j] += val * b(col, j); + } + } +#pragma omp simd + for (size_type j = 0; j < num_cols; ++j) { + c(global_row, j) = row_acc_vals[j]; + } + } + } + } +} + +GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_SPMM_KERNEL); + + +template +void advanced_spmm(std::shared_ptr exec, + matrix::view::dense alpha, + matrix::view::sellp a, + matrix::view::dense b, + matrix::view::dense beta, + matrix::view::dense c) +{ + const auto slice_lengths = a.slice_lengths; + const auto slice_sets = a.slice_sets; + const auto slice_size = a.slice_size; + const auto num_rows = a.size[0]; + const auto num_cols = c.size[1]; + const auto slice_num = ceildiv(num_rows + slice_size - 1, slice_size); + const auto alpha_val = alpha(0, 0); + const auto beta_val = beta(0, 0); + +#pragma omp parallel + { + array row_acc{exec, num_cols}; + auto* row_acc_vals = row_acc.get_data(); + +#pragma omp for collapse(2) schedule(static) + for (size_type slice = 0; slice < slice_num; ++slice) { + for (size_type row = 0; row < slice_size; ++row) { + const auto global_row = slice * slice_size + row; + if (global_row >= num_rows) { + continue; + } + const auto slice_begin = slice_sets[slice]; + const auto slice_length = slice_lengths[slice]; + std::fill_n(row_acc_vals, num_cols, zero()); + for (size_type idx = 0; idx < slice_length; ++idx) { + const auto val = a.val_at(row, slice_begin, idx); + const auto col = a.col_at(row, slice_begin, idx); + if (col == invalid_index()) { + continue; + } +#pragma omp simd + for (size_type j = 0; j < num_cols; ++j) { + row_acc_vals[j] += val * b(col, j); + } + } + if (is_zero(beta_val)) { +#pragma omp simd + for (size_type j = 0; j < num_cols; ++j) { + c(global_row, j) = alpha_val * row_acc_vals[j]; + } + } else { +#pragma omp simd + for (size_type j = 0; j < num_cols; ++j) { + c(global_row, j) = alpha_val * row_acc_vals[j] + + beta_val * c(global_row, j); + } + } + } + } + } +} + +GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( + GKO_DECLARE_SELLP_ADVANCED_SPMM_KERNEL); + + } // namespace sellp } // namespace omp } // namespace kernels diff --git a/omp/test/matrix/CMakeLists.txt b/omp/test/matrix/CMakeLists.txt index 398921ce75a..a220705b240 100644 --- a/omp/test/matrix/CMakeLists.txt +++ b/omp/test/matrix/CMakeLists.txt @@ -1 +1,2 @@ ginkgo_create_omp_test(fbcsr_kernels) +ginkgo_create_omp_test(sellp_spmm_kernels) diff --git a/omp/test/matrix/sellp_spmm_kernels.cpp b/omp/test/matrix/sellp_spmm_kernels.cpp new file mode 100644 index 00000000000..c098d6015b9 --- /dev/null +++ b/omp/test/matrix/sellp_spmm_kernels.cpp @@ -0,0 +1,116 @@ +// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors +// +// SPDX-License-Identifier: BSD-3-Clause + +#include + +#include + +#include +#include +#include + +#include "core/test/utils.hpp" + + +namespace { + + +class SellpSpmm : public ::testing::Test { +protected: + using value_type = double; + using index_type = int; + using Mtx = gko::matrix::Sellp; + using Vec = gko::matrix::Dense; + + static constexpr gko::size_type num_rows = 200; + static constexpr gko::size_type num_cols = 150; + static constexpr gko::size_type num_rhs = 16; + static constexpr value_type tolerance = 1e-12; + + SellpSpmm() : rand_engine(42) {} + + void SetUp() override + { + ref = gko::ReferenceExecutor::create(); + omp = gko::OmpExecutor::create(); + } + + void TearDown() override + { + if (omp != nullptr) { + ASSERT_NO_THROW(omp->synchronize()); + } + } + + template + std::unique_ptr gen_mtx(gko::size_type rows, gko::size_type cols, + int min_nnz_row) + { + return gko::test::generate_random_matrix( + rows, cols, + std::uniform_int_distribution<>(min_nnz_row, + static_cast(cols)), + std::normal_distribution<>(-1.0, 1.0), rand_engine, ref); + } + + void set_up_apply_data() + { + rand_engine.seed(42); + mtx = gen_mtx(num_rows, num_cols, 5); + y = gen_mtx(num_cols, num_rhs, 1); + expected = gen_mtx(num_rows, num_rhs, 1); + alpha = gko::initialize({2.0}, ref); + beta = gko::initialize({-1.5}, ref); + dmtx = Mtx::create(omp); + dmtx->copy_from(mtx); + dy = Vec::create(omp); + dy->copy_from(y); + dresult = Vec::create(omp); + dresult->copy_from(expected); + dalpha = Vec::create(omp); + dalpha->copy_from(alpha); + dbeta = Vec::create(omp); + dbeta->copy_from(beta); + } + + std::shared_ptr ref; + std::shared_ptr omp; + std::default_random_engine rand_engine; + + std::unique_ptr mtx; + std::unique_ptr y; + std::unique_ptr expected; + std::unique_ptr alpha; + std::unique_ptr beta; + std::unique_ptr dmtx; + std::unique_ptr dy; + std::unique_ptr dresult; + std::unique_ptr dalpha; + std::unique_ptr dbeta; +}; + + +TEST_F(SellpSpmm, ApplyIsEquivalentToRef) +{ + set_up_apply_data(); + + mtx->apply(y, expected); + dmtx->apply(dy, dresult); + + GKO_ASSERT_MTX_NEAR(dresult, expected, tolerance); +} + + +TEST_F(SellpSpmm, AdvancedApplyIsEquivalentToRef) +{ + set_up_apply_data(); + + mtx->apply(alpha, y, beta, expected); + dmtx->apply(dalpha, dy, dbeta, dresult); + + GKO_ASSERT_MTX_NEAR(dresult, expected, tolerance); +} + + +} // namespace diff --git a/reference/matrix/sellp_kernels.cpp b/reference/matrix/sellp_kernels.cpp index d62357dd502..b85f247ecfa 100644 --- a/reference/matrix/sellp_kernels.cpp +++ b/reference/matrix/sellp_kernels.cpp @@ -105,6 +105,33 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( GKO_DECLARE_SELLP_ADVANCED_SPMV_KERNEL); +template +void spmm(std::shared_ptr exec, + matrix::view::sellp a, + matrix::view::dense b, + matrix::view::dense c) +{ + spmv(exec, a, b, c); +} + +GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_SPMM_KERNEL); + + +template +void advanced_spmm(std::shared_ptr exec, + matrix::view::dense alpha, + matrix::view::sellp a, + matrix::view::dense b, + matrix::view::dense beta, + matrix::view::dense c) +{ + advanced_spmv(exec, alpha, a, b, beta, c); +} + +GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( + GKO_DECLARE_SELLP_ADVANCED_SPMM_KERNEL); + + template void compute_slice_sets(std::shared_ptr exec, const array& row_ptrs, size_type slice_size,