Skip to content

Add FBCSR SpMM and advanced SpMM kernels #2046

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 FBCSR SpMM and advanced SpMM kernels #2046

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
3 changes: 3 additions & 0 deletions common/cuda_hip/matrix/fbcsr_kernels.instantiate.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,9 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_SPMV_KERNEL);
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_FBCSR_ADVANCED_SPMV_KERNEL);
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_SPMM_KERNEL);
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_FBCSR_ADVANCED_SPMM_KERNEL);
// split
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_FBCSR_TRANSPOSE_KERNEL);
Expand Down
22 changes: 22 additions & 0 deletions common/cuda_hip/matrix/fbcsr_kernels.template.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -551,6 +551,28 @@ void advanced_spmv(std::shared_ptr<const DefaultExecutor> exec,
}


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


template <typename ValueType, typename IndexType>
void advanced_spmm(std::shared_ptr<const DefaultExecutor> exec,
matrix::view::dense<const ValueType> alpha,
const matrix::Fbcsr<ValueType, 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);
}


namespace {


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 @@ -827,6 +827,8 @@ namespace fbcsr {

GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_SPMV_KERNEL);
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_ADVANCED_SPMV_KERNEL);
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_SPMM_KERNEL);
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_ADVANCED_SPMM_KERNEL);
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_FILL_IN_MATRIX_DATA_KERNEL);
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_FILL_IN_DENSE_KERNEL);
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_CONVERT_TO_CSR_KERNEL);
Expand Down
32 changes: 24 additions & 8 deletions core/matrix/fbcsr.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ namespace {

GKO_REGISTER_OPERATION(spmv, fbcsr::spmv);
GKO_REGISTER_OPERATION(advanced_spmv, fbcsr::advanced_spmv);
GKO_REGISTER_OPERATION(spmm, fbcsr::spmm);
GKO_REGISTER_OPERATION(advanced_spmm, fbcsr::advanced_spmm);
GKO_REGISTER_OPERATION(fill_in_matrix_data, fbcsr::fill_in_matrix_data);
GKO_REGISTER_OPERATION(convert_to_csr, fbcsr::convert_to_csr);
GKO_REGISTER_OPERATION(fill_in_dense, fbcsr::fill_in_dense);
Expand Down Expand Up @@ -112,9 +114,15 @@ void Fbcsr<ValueType, IndexType>::apply_impl(const LinOp* b, LinOp* x) const
// otherwise we assume that b is dense and compute a SpMV/SpMM
precision_dispatch_real_complex<ValueType>(
[this](auto dense_b, auto dense_x) {
this->get_executor()->run(
fbcsr::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(
fbcsr::make_spmv(this, dense_b->get_const_device_view(),
dense_x->get_device_view()));
} else {
this->get_executor()->run(
fbcsr::make_spmm(this, dense_b->get_const_device_view(),
dense_x->get_device_view()));
}
},
b, x);
}
Expand All @@ -136,11 +144,19 @@ void Fbcsr<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(fbcsr::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(fbcsr::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(fbcsr::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()));
}
},
alpha, b, beta, x);
}
Expand Down
18 changes: 18 additions & 0 deletions core/matrix/fbcsr_kernels.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,20 @@ namespace kernels {
matrix::view::dense<const ValueType> beta, \
matrix::view::dense<ValueType> c)

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

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

#define GKO_DECLARE_FBCSR_FILL_IN_MATRIX_DATA_KERNEL(ValueType, IndexType) \
void fill_in_matrix_data(std::shared_ptr<const DefaultExecutor> exec, \
device_matrix_data<ValueType, IndexType>& data, \
Expand Down Expand Up @@ -82,6 +96,10 @@ namespace kernels {
template <typename ValueType, typename IndexType> \
GKO_DECLARE_FBCSR_ADVANCED_SPMV_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_FBCSR_SPMM_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_FBCSR_ADVANCED_SPMM_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_FBCSR_FILL_IN_MATRIX_DATA_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_FBCSR_FILL_IN_DENSE_KERNEL(ValueType, IndexType); \
Expand Down
27 changes: 27 additions & 0 deletions dpcpp/matrix/fbcsr_kernels.dp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,33 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_FBCSR_ADVANCED_SPMV_KERNEL);


template <typename ValueType, typename IndexType>
void spmm(std::shared_ptr<const DpcppExecutor> exec,
const matrix::Fbcsr<ValueType, 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_FBCSR_SPMM_KERNEL);


template <typename ValueType, typename IndexType>
void advanced_spmm(std::shared_ptr<const DpcppExecutor> exec,
matrix::view::dense<const ValueType> alpha,
const matrix::Fbcsr<ValueType, 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_FBCSR_ADVANCED_SPMM_KERNEL);


template <typename ValueType, typename IndexType>
void fill_in_matrix_data(std::shared_ptr<const DefaultExecutor> exec,
device_matrix_data<ValueType, IndexType>& data,
Expand Down
213 changes: 213 additions & 0 deletions omp/matrix/fbcsr_kernels.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,219 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_FBCSR_ADVANCED_SPMV_KERNEL);


namespace {


template <typename ValueType, typename IndexType>
void spmm_tiled(std::shared_ptr<const OmpExecutor> exec,
const matrix::Fbcsr<ValueType, IndexType>* a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<ValueType> c)
{
constexpr int R = 2;
GKO_ASSERT(a->get_block_size() == R);
const auto K = static_cast<IndexType>(b.size[1]);

const IndexType nbrows = a->get_num_block_rows();
const size_type nbnz = a->get_num_stored_blocks();
const auto row_ptrs = a->get_const_row_ptrs();
const auto col_idxs = a->get_const_col_idxs();
const acc::range<acc::block_col_major<const ValueType, 3>> avalues{
to_std_array<acc::size_type>(nbnz, R, R), a->get_const_values()};
const auto acc_size = static_cast<size_type>(R) * static_cast<size_type>(K);

#pragma omp parallel
{
array<ValueType> cacc{exec, acc_size};
array<ValueType> brow{exec, acc_size};
auto* cacc_vals = cacc.get_data();
auto* brow_vals = brow.get_data();

#pragma omp for schedule(static)
for (IndexType ibrow = 0; ibrow < nbrows; ++ibrow) {
std::fill_n(cacc_vals, acc_size, zero<ValueType>());
for (IndexType inz = row_ptrs[ibrow]; inz < row_ptrs[ibrow + 1];
++inz) {
const IndexType bcol = col_idxs[inz] * R;

ValueType aval[R][R];
for (int ib = 0; ib < R; ++ib) {
for (int jb = 0; jb < R; ++jb) {
aval[ib][jb] = avalues(inz, ib, jb);
}
}

for (int jb = 0; jb < R; ++jb) {
#pragma omp simd
for (IndexType j = 0; j < K; ++j) {
brow_vals[static_cast<size_type>(jb) *
static_cast<size_type>(K) +
j] = b(bcol + jb, j);
}
}

for (int ib = 0; ib < R; ++ib) {
for (int jb = 0; jb < R; ++jb) {
#pragma omp simd
for (IndexType j = 0; j < K; ++j) {
cacc_vals[static_cast<size_type>(ib) *
static_cast<size_type>(K) +
j] +=
aval[ib][jb] *
brow_vals[static_cast<size_type>(jb) *
static_cast<size_type>(K) +
j];
}
}
}
}

for (int ib = 0; ib < R; ++ib) {
#pragma omp simd
for (IndexType j = 0; j < K; ++j) {
c(ibrow * R + ib, j) =
cacc_vals[static_cast<size_type>(ib) *
static_cast<size_type>(K) +
j];
}
}
}
}
}


template <typename ValueType, typename IndexType>
void advanced_spmm_tiled(std::shared_ptr<const OmpExecutor> exec,
ValueType valpha, ValueType vbeta,
const matrix::Fbcsr<ValueType, IndexType>* a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<ValueType> c)
{
constexpr int R = 2;
GKO_ASSERT(a->get_block_size() == R);
const auto K = static_cast<IndexType>(b.size[1]);

const IndexType nbrows = a->get_num_block_rows();
const size_type nbnz = a->get_num_stored_blocks();
const auto row_ptrs = a->get_const_row_ptrs();
const auto col_idxs = a->get_const_col_idxs();
const acc::range<acc::block_col_major<const ValueType, 3>> avalues{
to_std_array<acc::size_type>(nbnz, R, R), a->get_const_values()};
const auto acc_size = static_cast<size_type>(R) * static_cast<size_type>(K);

#pragma omp parallel
{
array<ValueType> cacc{exec, acc_size};
array<ValueType> brow{exec, acc_size};
auto* cacc_vals = cacc.get_data();
auto* brow_vals = brow.get_data();

#pragma omp for schedule(static)
for (IndexType ibrow = 0; ibrow < nbrows; ++ibrow) {
std::fill_n(cacc_vals, acc_size, zero<ValueType>());
for (IndexType inz = row_ptrs[ibrow]; inz < row_ptrs[ibrow + 1];
++inz) {
const IndexType bcol = col_idxs[inz] * R;

ValueType aval[R][R];
for (int ib = 0; ib < R; ++ib) {
for (int jb = 0; jb < R; ++jb) {
aval[ib][jb] = avalues(inz, ib, jb);
}
}

for (int jb = 0; jb < R; ++jb) {
#pragma omp simd
for (IndexType j = 0; j < K; ++j) {
brow_vals[static_cast<size_type>(jb) *
static_cast<size_type>(K) +
j] = b(bcol + jb, j);
}
}

for (int ib = 0; ib < R; ++ib) {
for (int jb = 0; jb < R; ++jb) {
#pragma omp simd
for (IndexType j = 0; j < K; ++j) {
cacc_vals[static_cast<size_type>(ib) *
static_cast<size_type>(K) +
j] +=
aval[ib][jb] *
brow_vals[static_cast<size_type>(jb) *
static_cast<size_type>(K) +
j];
}
}
}
}

if (is_zero(vbeta)) {
for (int ib = 0; ib < R; ++ib) {
#pragma omp simd
for (IndexType j = 0; j < K; ++j) {
c(ibrow * R + ib, j) =
valpha * cacc_vals[static_cast<size_type>(ib) *
static_cast<size_type>(K) +
j];
}
}
} else {
for (int ib = 0; ib < R; ++ib) {
#pragma omp simd
for (IndexType j = 0; j < K; ++j) {
c(ibrow * R + ib, j) =
valpha * cacc_vals[static_cast<size_type>(ib) *
static_cast<size_type>(K) +
j] +
vbeta * c(ibrow * R + ib, j);
}
}
}
}
}
}


} // namespace


template <typename ValueType, typename IndexType>
void spmm(std::shared_ptr<const OmpExecutor> exec,
const matrix::Fbcsr<ValueType, IndexType>* a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<ValueType> c)
{
if (a->get_block_size() == 2) {
spmm_tiled(exec, a, b, c);
} else {
spmv(exec, a, b, c);
}
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_FBCSR_SPMM_KERNEL);


template <typename ValueType, typename IndexType>
void advanced_spmm(std::shared_ptr<const OmpExecutor> exec,
matrix::view::dense<const ValueType> alpha,
const matrix::Fbcsr<ValueType, IndexType>* a,
matrix::view::dense<const ValueType> b,
matrix::view::dense<const ValueType> beta,
matrix::view::dense<ValueType> c)
{
const auto valpha = alpha(0, 0);
const auto vbeta = beta(0, 0);
if (a->get_block_size() == 2) {
advanced_spmm_tiled(exec, valpha, vbeta, a, b, c);
} else {
advanced_spmv(exec, alpha, a, b, beta, c);
}
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_FBCSR_ADVANCED_SPMM_KERNEL);


template <typename ValueType, typename IndexType>
void fill_in_matrix_data(std::shared_ptr<const DefaultExecutor> exec,
device_matrix_data<ValueType, IndexType>& data,
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(fbcsr_spmm_kernels)
Loading