Skip to content

Add SELL-P SpMM and advanced SpMM kernels #2045

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
mlykov wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Add SELL-P SpMM and advanced SpMM kernels #2045

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions common/cuda_hip/matrix/sellp_kernels.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,33 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_SELLP_ADVANCED_SPMV_KERNEL);


template <typename ValueType, typename IndexType>
void spmm(std::shared_ptr<const DefaultExecutor> exec,
matrix::view::sellp<const ValueType, const IndexType> a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<ValueType> c)
{
spmv(exec, a, b, c);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_SPMM_KERNEL);


template <typename ValueType, typename IndexType>
void advanced_spmm(std::shared_ptr<const DefaultExecutor> exec,
matrix::view::dense<const ValueType> alpha,
matrix::view::sellp<const ValueType, const IndexType> a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<const ValueType> beta,
matrix::view::dense<ValueType> 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
Expand Down
2 changes: 2 additions & 0 deletions core/device_hooks/common_kernels.inc.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
36 changes: 28 additions & 8 deletions core/matrix/sellp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -182,9 +184,17 @@ void Sellp<ValueType, IndexType>::apply_impl(const LinOp* b, LinOp* x) const
{
precision_dispatch_real_complex<ValueType>(
[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);
}
Expand All @@ -196,11 +206,21 @@ void Sellp<ValueType, IndexType>::apply_impl(const LinOp* alpha, const LinOp* b,
{
precision_dispatch_real_complex<ValueType>(
[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);
}
Expand Down
19 changes: 19 additions & 0 deletions core/matrix/sellp_kernels.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,21 @@ namespace kernels {
matrix::view::dense<const ValueType> beta, \
matrix::view::dense<ValueType> c)

#define GKO_DECLARE_SELLP_SPMM_KERNEL(ValueType, IndexType) \
void spmm(std::shared_ptr<const DefaultExecutor> exec, \
matrix::view::sellp<const ValueType, const IndexType> a, \
matrix::view::dense<const ValueType> b, \
matrix::view::dense<ValueType> c)

#define GKO_DECLARE_SELLP_ADVANCED_SPMM_KERNEL(ValueType, IndexType) \
void advanced_spmm( \
std::shared_ptr<const DefaultExecutor> exec, \
matrix::view::dense<const ValueType> alpha, \
matrix::view::sellp<const ValueType, const IndexType> a, \
matrix::view::dense<const ValueType> b, \
matrix::view::dense<const ValueType> beta, \
matrix::view::dense<ValueType> c)

#define GKO_DECLARE_SELLP_FILL_IN_MATRIX_DATA_KERNEL(ValueType, IndexType) \
void fill_in_matrix_data( \
std::shared_ptr<const DefaultExecutor> exec, \
Expand Down Expand Up @@ -76,6 +91,10 @@ namespace kernels {
template <typename ValueType, typename IndexType> \
GKO_DECLARE_SELLP_ADVANCED_SPMV_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_SELLP_SPMM_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_SELLP_ADVANCED_SPMM_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_SELLP_FILL_IN_MATRIX_DATA_KERNEL(ValueType, IndexType); \
template <typename IndexType> \
GKO_DECLARE_SELLP_COMPUTE_SLICE_SETS_KERNEL(IndexType); \
Expand Down
27 changes: 27 additions & 0 deletions dpcpp/matrix/sellp_kernels.dp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,33 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_SELLP_ADVANCED_SPMV_KERNEL);


template <typename ValueType, typename IndexType>
void spmm(std::shared_ptr<const DpcppExecutor> exec,
matrix::view::sellp<const ValueType, const IndexType> a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<ValueType> c)
{
spmv(exec, a, b, c);
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_SPMM_KERNEL);


template <typename ValueType, typename IndexType>
void advanced_spmm(std::shared_ptr<const DpcppExecutor> exec,
matrix::view::dense<const ValueType> alpha,
matrix::view::sellp<const ValueType, const IndexType> a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<const ValueType> beta,
matrix::view::dense<ValueType> 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
Expand Down
117 changes: 117 additions & 0 deletions omp/matrix/sellp_kernels.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

#include <omp.h>

#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>


Expand Down Expand Up @@ -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 <typename ValueType, typename IndexType>
void spmm(std::shared_ptr<const OmpExecutor> exec,
matrix::view::sellp<const ValueType, const IndexType> a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<ValueType> 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<ValueType> 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<ValueType>());
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<IndexType>()) {
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 <typename ValueType, typename IndexType>
void advanced_spmm(std::shared_ptr<const OmpExecutor> exec,
matrix::view::dense<const ValueType> alpha,
matrix::view::sellp<const ValueType, const IndexType> a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<const ValueType> beta,
matrix::view::dense<ValueType> 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<ValueType> 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<ValueType>());
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<IndexType>()) {
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
Expand Down
1 change: 1 addition & 0 deletions omp/test/matrix/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1,2 @@
ginkgo_create_omp_test(fbcsr_kernels)
ginkgo_create_omp_test(sellp_spmm_kernels)
Loading