Skip to content

Add CSR SpMM and advanced SpMM kernels #2044

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 CSR SpMM and advanced SpMM kernels #2044

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
99 changes: 99 additions & 0 deletions common/cuda_hip/matrix/csr_kernels.instantiate.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,105 @@ GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT7(
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT8(
GKO_DECLARE_CSR_ADVANCED_SPMV_KERNEL, int64);

// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT1(GKO_DECLARE_CSR_SPMM_KERNEL,
int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT2(GKO_DECLARE_CSR_SPMM_KERNEL,
int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT3(GKO_DECLARE_CSR_SPMM_KERNEL,
int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT4(GKO_DECLARE_CSR_SPMM_KERNEL,
int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT5(GKO_DECLARE_CSR_SPMM_KERNEL,
int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT6(GKO_DECLARE_CSR_SPMM_KERNEL,
int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT7(GKO_DECLARE_CSR_SPMM_KERNEL,
int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT8(GKO_DECLARE_CSR_SPMM_KERNEL,
int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT1(GKO_DECLARE_CSR_SPMM_KERNEL,
int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT2(GKO_DECLARE_CSR_SPMM_KERNEL,
int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT3(GKO_DECLARE_CSR_SPMM_KERNEL,
int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT4(GKO_DECLARE_CSR_SPMM_KERNEL,
int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT5(GKO_DECLARE_CSR_SPMM_KERNEL,
int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT6(GKO_DECLARE_CSR_SPMM_KERNEL,
int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT7(GKO_DECLARE_CSR_SPMM_KERNEL,
int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT8(GKO_DECLARE_CSR_SPMM_KERNEL,
int64);


// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT1(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT2(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT3(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT4(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT5(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT6(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT7(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT8(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int32);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT1(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT2(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT3(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT4(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT5(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT6(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT7(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int64);
// split
GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_TYPE_SPLIT8(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL, int64);


// split
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_TRANSPOSE_KERNEL);
Expand Down
26 changes: 26 additions & 0 deletions common/cuda_hip/matrix/csr_kernels.template.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2486,6 +2486,32 @@ void advanced_spmv(std::shared_ptr<const DefaultExecutor> exec,
}


// Stub implementation.
template <typename MatrixValueType, typename InputValueType,
typename OutputValueType, typename IndexType>
void spmm(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Csr<MatrixValueType, IndexType>* a,
matrix::view::dense<const InputValueType> b,
matrix::view::dense<OutputValueType> c)
{
spmv(exec, a, b, c);
}


// Stub implementation.
template <typename MatrixValueType, typename InputValueType,
typename OutputValueType, typename IndexType>
void advanced_spmm(std::shared_ptr<const DefaultExecutor> exec,
matrix::view::dense<const MatrixValueType> alpha,
const matrix::Csr<MatrixValueType, IndexType>* a,
matrix::view::dense<const InputValueType> b,
matrix::view::dense<const OutputValueType> beta,
matrix::view::dense<OutputValueType> c)
{
advanced_spmv(exec, alpha, a, b, beta, c);
}


template <typename ValueType, typename IndexType>
void spgemm(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Csr<ValueType, IndexType>* a,
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 @@ -769,6 +769,8 @@ namespace csr {

GKO_STUB_MIXED_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_SPMV_KERNEL);
GKO_STUB_MIXED_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_ADVANCED_SPMV_KERNEL);
GKO_STUB_MIXED_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_SPMM_KERNEL);
GKO_STUB_MIXED_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL);
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_SPGEMM_KERNEL);
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_ADVANCED_SPGEMM_KERNEL);
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_SPGEMM_REUSE_KERNEL);
Expand Down
32 changes: 24 additions & 8 deletions core/matrix/csr.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@ namespace {

GKO_REGISTER_OPERATION(spmv, csr::spmv);
GKO_REGISTER_OPERATION(advanced_spmv, csr::advanced_spmv);
GKO_REGISTER_OPERATION(spmm, csr::spmm);
GKO_REGISTER_OPERATION(advanced_spmm, csr::advanced_spmm);
GKO_REGISTER_OPERATION(spgemm, csr::spgemm);
GKO_REGISTER_OPERATION(advanced_spgemm, csr::advanced_spgemm);
GKO_REGISTER_OPERATION(spgemm_reuse, csr::spgemm_reuse);
Expand Down Expand Up @@ -266,9 +268,15 @@ void Csr<ValueType, IndexType>::apply_impl(const LinOp* b, LinOp* x) const
} else {
mixed_precision_dispatch_real_complex<ValueType>(
[this](auto dense_b, auto dense_x) {
this->get_executor()->run(
csr::make_spmv(this, dense_b->get_const_device_view(),
dense_x->get_device_view()));
if (dense_b->get_size()[1] <= 2) {
this->get_executor()->run(
csr::make_spmv(this, dense_b->get_const_device_view(),
dense_x->get_device_view()));
} else {
this->get_executor()->run(
csr::make_spmm(this, dense_b->get_const_device_view(),
dense_x->get_device_view()));
}
},
b, x);
}
Expand Down Expand Up @@ -305,11 +313,19 @@ void Csr<ValueType, IndexType>::apply_impl(const LinOp* alpha, const LinOp* b,
auto dense_beta = make_temporary_conversion<
typename std::decay_t<decltype(*dense_x)>::value_type>(
beta);
this->get_executor()->run(csr::make_advanced_spmv(
dense_alpha->get_const_device_view(), this,
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(csr::make_advanced_spmv(
dense_alpha->get_const_device_view(), this,
dense_b->get_const_device_view(),
dense_beta->get_const_device_view(),
dense_x->get_device_view()));
} else {
this->get_executor()->run(csr::make_advanced_spmm(
dense_alpha->get_const_device_view(), this,
dense_b->get_const_device_view(),
dense_beta->get_const_device_view(),
dense_x->get_device_view()));
}
},
b, x);
}
Expand Down
24 changes: 24 additions & 0 deletions core/matrix/csr_kernels.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,22 @@ namespace kernels {
matrix::view::dense<const OutputValueType> beta, \
matrix::view::dense<OutputValueType> c)

#define GKO_DECLARE_CSR_SPMM_KERNEL(MatrixValueType, InputValueType, \
OutputValueType, IndexType) \
void spmm(std::shared_ptr<const DefaultExecutor> exec, \
const matrix::Csr<MatrixValueType, IndexType>* a, \
matrix::view::dense<const InputValueType> b, \
matrix::view::dense<OutputValueType> c)

#define GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL(MatrixValueType, InputValueType, \
OutputValueType, IndexType) \
void advanced_spmm(std::shared_ptr<const DefaultExecutor> exec, \
matrix::view::dense<const MatrixValueType> alpha, \
const matrix::Csr<MatrixValueType, IndexType>* a, \
matrix::view::dense<const InputValueType> b, \
matrix::view::dense<const OutputValueType> beta, \
matrix::view::dense<OutputValueType> c)

#define GKO_DECLARE_CSR_SPGEMM_KERNEL(ValueType, IndexType) \
void spgemm(std::shared_ptr<const DefaultExecutor> exec, \
const matrix::Csr<ValueType, IndexType>* a, \
Expand Down Expand Up @@ -296,6 +312,14 @@ namespace kernels {
typename OutputValueType, typename IndexType> \
GKO_DECLARE_CSR_ADVANCED_SPMV_KERNEL(MatrixValueType, InputValueType, \
OutputValueType, IndexType); \
template <typename MatrixValueType, typename InputValueType, \
typename OutputValueType, typename IndexType> \
GKO_DECLARE_CSR_SPMM_KERNEL(MatrixValueType, InputValueType, \
OutputValueType, IndexType); \
template <typename MatrixValueType, typename InputValueType, \
typename OutputValueType, typename IndexType> \
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL(MatrixValueType, InputValueType, \
OutputValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_CSR_SPGEMM_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
Expand Down
32 changes: 32 additions & 0 deletions dpcpp/matrix/csr_kernels.dp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1663,6 +1663,38 @@ GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_CSR_ADVANCED_SPMV_KERNEL);


// Stub implementation.
template <typename MatrixValueType, typename InputValueType,
typename OutputValueType, typename IndexType>
void spmm(std::shared_ptr<const DpcppExecutor> exec,
const matrix::Csr<MatrixValueType, IndexType>* a,
matrix::view::dense<const InputValueType> b,
matrix::view::dense<OutputValueType> c)
{
spmv(exec, a, b, c);
}

GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_CSR_SPMM_KERNEL);


// Stub implementation.
template <typename MatrixValueType, typename InputValueType,
typename OutputValueType, typename IndexType>
void advanced_spmm(std::shared_ptr<const DpcppExecutor> exec,
matrix::view::dense<const MatrixValueType> alpha,
const matrix::Csr<MatrixValueType, IndexType>* a,
matrix::view::dense<const InputValueType> b,
matrix::view::dense<const OutputValueType> beta,
matrix::view::dense<OutputValueType> c)
{
advanced_spmv(exec, alpha, a, b, beta, c);
}

GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL);


namespace kernel {


Expand Down
111 changes: 111 additions & 0 deletions omp/matrix/csr_kernels.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,117 @@ GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_CSR_ADVANCED_SPMV_KERNEL);


// Gustavson's algorithm with row-level parallelism and SIMD over dense columns
template <typename MatrixValueType, typename InputValueType,
typename OutputValueType, typename IndexType>
void spmm(std::shared_ptr<const OmpExecutor> exec,
const matrix::Csr<MatrixValueType, IndexType>* a,
matrix::view::dense<const InputValueType> b,
matrix::view::dense<OutputValueType> c)
{
using arithmetic_type =
highest_precision<MatrixValueType, InputValueType, OutputValueType>;

const auto row_ptrs = a->get_const_row_ptrs();
const auto col_idxs = a->get_const_col_idxs();

const auto a_vals =
acc::helper::build_const_rrm_accessor<arithmetic_type>(a);
const auto b_vals =
acc::helper::build_const_rrm_accessor<arithmetic_type>(b);
auto c_vals = acc::helper::build_rrm_accessor<arithmetic_type>(c);

const auto num_rows = a->get_size()[0];
const auto num_cols = c.size[1];

#pragma omp parallel
{
array<arithmetic_type> row_acc{exec, num_cols};
auto* row_acc_vals = row_acc.get_data();
const auto row_acc_size = row_acc.get_size();

#pragma omp for schedule(static)
for (size_type row = 0; row < num_rows; ++row) {
std::fill_n(row_acc_vals, row_acc_size, zero<arithmetic_type>());
for (size_type nz = row_ptrs[row];
nz < static_cast<size_type>(row_ptrs[row + 1]); ++nz) {
const arithmetic_type a_ik = a_vals(nz);
const auto k_row_b = static_cast<size_type>(col_idxs[nz]);
#pragma omp simd
for (size_type j = 0; j < num_cols; ++j) {
row_acc_vals[j] += a_ik * b_vals(k_row_b, j);
}
}
#pragma omp simd
for (size_type j = 0; j < num_cols; ++j) {
c_vals(row, j) = row_acc_vals[j];
}
}
}
}

GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_CSR_SPMM_KERNEL);


template <typename MatrixValueType, typename InputValueType,
typename OutputValueType, typename IndexType>
void advanced_spmm(std::shared_ptr<const OmpExecutor> exec,
matrix::view::dense<const MatrixValueType> alpha,
const matrix::Csr<MatrixValueType, IndexType>* a,
matrix::view::dense<const InputValueType> b,
matrix::view::dense<const OutputValueType> beta,
matrix::view::dense<OutputValueType> c)
{
using arithmetic_type =
highest_precision<MatrixValueType, InputValueType, OutputValueType>;

const auto row_ptrs = a->get_const_row_ptrs();
const auto col_idxs = a->get_const_col_idxs();

const auto a_vals =
acc::helper::build_const_rrm_accessor<arithmetic_type>(a);
const auto b_vals =
acc::helper::build_const_rrm_accessor<arithmetic_type>(b);
auto c_vals = acc::helper::build_rrm_accessor<arithmetic_type>(c);

const auto num_rows = a->get_size()[0];
const auto num_cols = c.size[1];
const auto valpha = static_cast<arithmetic_type>(alpha(0, 0));
const auto vbeta = static_cast<arithmetic_type>(beta(0, 0));

#pragma omp parallel
{
array<arithmetic_type> row_acc{exec, num_cols};
auto* row_acc_vals = row_acc.get_data();
const auto row_acc_size = row_acc.get_size();

#pragma omp for schedule(static)
for (size_type row = 0; row < num_rows; ++row) {
std::fill_n(row_acc_vals, row_acc_size, zero<arithmetic_type>());
for (size_type nz = row_ptrs[row];
nz < static_cast<size_type>(row_ptrs[row + 1]); ++nz) {
const arithmetic_type a_ik = a_vals(nz);
const auto k_row_b = static_cast<size_type>(col_idxs[nz]);
#pragma omp simd
for (size_type j = 0; j < num_cols; ++j) {
row_acc_vals[j] += a_ik * b_vals(k_row_b, j);
}
}
#pragma omp simd
for (size_type j = 0; j < num_cols; ++j) {
const auto scaled_c =
is_zero(vbeta) ? zero(vbeta) : c_vals(row, j) * vbeta;
c_vals(row, j) = valpha * row_acc_vals[j] + scaled_c;
}
}
}
}

GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_CSR_ADVANCED_SPMM_KERNEL);


namespace {


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(csr_spmm_kernels)
ginkgo_create_omp_test(fbcsr_kernels)
Loading