Skip to content

Commit

Permalink
Add dpcpp kernels + tests.
Browse files Browse the repository at this point in the history
  • Loading branch information
@pratikvn
pratikvn committed Nov 9, 2021
1 parent 16a5165 commit c3c0351
Show file tree
Hide file tree
Showing 2 changed files with 168 additions and 2 deletions.
88 changes: 86 additions & 2 deletions dpcpp/matrix/csr_kernels.dp.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1625,11 +1625,79 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_CSR_ADVANCED_SPMV_KERNEL);


namespace kernel {


template <typename IndexType>
void calc_nnz_in_span(const span row_span, const span col_span,
const IndexType* __restrict__ row_ptrs,
const IndexType* __restrict__ col_idxs,
IndexType* __restrict__ nnz_per_row,
sycl::nd_item<3> item_ct1)
{
const auto tidx = thread::get_thread_id_flat(item_ct1);
if (tidx < row_span.end) {
nnz_per_row[tidx - row_span.begin] = zero<IndexType>();
for (size_type col = row_ptrs[tidx]; col < row_ptrs[tidx + 1]; ++col) {
if (col_idxs[col] >= col_span.begin &&
col_idxs[col] < col_span.end) {
nnz_per_row[tidx - row_span.begin]++;
}
}
}
}

GKO_ENABLE_DEFAULT_HOST(calc_nnz_in_span, calc_nnz_in_span);


template <typename ValueType, typename IndexType>
void compute_submat(size_type num_rows, size_type num_cols, size_type num_nnz,
size_type row_offset, size_type col_offset,
const IndexType* __restrict__ src_row_ptrs,
const IndexType* __restrict__ src_col_idxs,
const ValueType* __restrict__ src_values,
const IndexType* __restrict__ res_row_ptrs,
IndexType* __restrict__ res_col_idxs,
ValueType* __restrict__ res_values,
sycl::nd_item<3> item_ct1)
{
const auto tidx = thread::get_thread_id_flat(item_ct1);
if (tidx < num_rows) {
size_type res_nnz = res_row_ptrs[tidx];
for (size_type nnz = src_row_ptrs[row_offset + tidx];
nnz < src_row_ptrs[row_offset + tidx + 1]; ++nnz) {
if ((src_col_idxs[nnz] < (col_offset + num_cols) &&
src_col_idxs[nnz] >= col_offset)) {
res_col_idxs[res_nnz] = src_col_idxs[nnz] - col_offset;
res_values[res_nnz] = src_values[nnz];
res_nnz++;
}
}
}
}

GKO_ENABLE_DEFAULT_HOST(compute_submat, compute_submat);


} // namespace kernel


template <typename ValueType, typename IndexType>
void calculate_nonzeros_per_row_in_span(
std::shared_ptr<const DefaultExecutor> exec,
const matrix::Csr<ValueType, IndexType>* source, const span& row_span,
const span& col_span, Array<IndexType>* row_nnz) GKO_NOT_IMPLEMENTED;
const span& col_span, Array<IndexType>* row_nnz)
{
const auto num_rows = source->get_size()[0];
auto row_ptrs = source->get_const_row_ptrs();
auto col_idxs = source->get_const_col_idxs();
auto grid_dim = ceildiv(row_span.length(), default_block_size);
auto block_dim = default_block_size;

kernel::calc_nnz_in_span(grid_dim, block_dim, 0, exec->get_queue(),
row_span, col_span, row_ptrs, col_idxs,
row_nnz->get_data());
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_CSR_CALC_NNZ_PER_ROW_IN_SPAN_KERNEL);
Expand All @@ -1640,7 +1708,23 @@ void compute_submatrix(std::shared_ptr<const DefaultExecutor> exec,
const matrix::Csr<ValueType, IndexType>* source,
gko::span row_span, gko::span col_span,
matrix::Csr<ValueType, IndexType>* result)
GKO_NOT_IMPLEMENTED;
{
const auto row_offset = row_span.begin;
const auto col_offset = col_span.begin;
const auto num_rows = result->get_size()[0];
const auto num_cols = result->get_size()[1];
const auto row_ptrs = source->get_const_row_ptrs();

const auto num_nnz = source->get_num_stored_elements();
auto grid_dim = ceildiv(num_rows, default_block_size);
auto block_dim = default_block_size;
kernel::compute_submat(
grid_dim, block_dim, 0, exec->get_queue(), num_rows, num_cols, num_nnz,
row_offset, col_offset, source->get_const_row_ptrs(),
source->get_const_col_idxs(), source->get_const_values(),
result->get_const_row_ptrs(), result->get_col_idxs(),
result->get_values());
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
GKO_DECLARE_CSR_COMPUTE_SUB_MATRIX_KERNEL);
Expand Down
82 changes: 82 additions & 0 deletions dpcpp/test/matrix/csr_kernels.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#include <ginkgo/core/matrix/sparsity_csr.hpp>


#include "core/components/prefix_sum.hpp"
#include "core/matrix/csr_kernels.hpp"
#include "core/test/utils.hpp"
#include "core/test/utils/unsort_matrix.hpp"
Expand Down Expand Up @@ -102,6 +103,14 @@ class Csr : public ::testing::Test {
std::normal_distribution<vtype>(-1.0, 1.0), rand_engine, ref);
}

void set_up_apply_data()
{
mtx2 = Mtx::create(ref);
mtx2->copy_from(gen_mtx<Mtx>(mtx_size[0], mtx_size[1], 5));
dmtx2 = Mtx::create(dpcpp);
dmtx2->copy_from(mtx2.get());
}

void set_up_apply_data(std::shared_ptr<Mtx::strategy_type> strategy,
int num_vectors = 1)
{
Expand Down Expand Up @@ -156,6 +165,7 @@ class Csr : public ::testing::Test {
std::ranlux48 rand_engine;

std::unique_ptr<Mtx> mtx;
std::unique_ptr<Mtx> mtx2;
std::unique_ptr<ComplexMtx> complex_mtx;
std::unique_ptr<Mtx> square_mtx;
std::unique_ptr<Vec> expected;
Expand All @@ -164,6 +174,7 @@ class Csr : public ::testing::Test {
std::unique_ptr<Vec> beta;

std::unique_ptr<Mtx> dmtx;
std::unique_ptr<Mtx> dmtx2;
std::unique_ptr<ComplexMtx> complex_dmtx;
std::unique_ptr<Mtx> square_dmtx;
std::unique_ptr<Vec> dresult;
Expand Down Expand Up @@ -930,4 +941,75 @@ TEST_F(Csr, OutplaceAbsoluteComplexMatrixIsEquivalentToRef)
}


TEST_F(Csr, CalculateNnzPerRowInSpanIsEquivalentToRef)
{
using Mtx = gko::matrix::Csr<vtype, int>;
set_up_apply_data();
gko::span rspan{7, 51};
gko::span cspan{22, 88};
auto size = this->mtx2->get_size();
auto row_nnz = gko::Array<int>(this->ref, rspan.length() + 1);
auto drow_nnz = gko::Array<int>(this->dpcpp, row_nnz);

gko::kernels::reference::csr::calculate_nonzeros_per_row_in_span(
this->ref, this->mtx2.get(), rspan, cspan, &row_nnz);
gko::kernels::dpcpp::csr::calculate_nonzeros_per_row_in_span(
this->dpcpp, this->dmtx2.get(), rspan, cspan, &drow_nnz);

GKO_ASSERT_ARRAY_EQ(row_nnz, drow_nnz);
}


TEST_F(Csr, ComputeSubmatrixIsEquivalentToRef)
{
using Mtx = gko::matrix::Csr<vtype, int>;
using IndexType = int;
using ValueType = vtype;
set_up_apply_data();
gko::span rspan{7, 51};
gko::span cspan{22, 88};
auto size = this->mtx2->get_size();
auto row_nnz = gko::Array<int>(this->ref, rspan.length() + 1);
row_nnz.fill(gko::zero<int>());
gko::kernels::reference::csr::calculate_nonzeros_per_row_in_span(
this->ref, this->mtx2.get(), rspan, cspan, &row_nnz);
gko::kernels::reference::components::prefix_sum(
this->ref, row_nnz.get_data(), row_nnz.get_num_elems());
auto num_nnz = row_nnz.get_data()[rspan.length()];
auto drow_nnz = gko::Array<int>(this->dpcpp, row_nnz);
auto smat1 =
Mtx::create(this->ref, gko::dim<2>(rspan.length(), cspan.length()),
std::move(gko::Array<ValueType>(this->ref, num_nnz)),
std::move(gko::Array<IndexType>(this->ref, num_nnz)),
std::move(row_nnz));
auto sdmat1 =
Mtx::create(this->dpcpp, gko::dim<2>(rspan.length(), cspan.length()),
std::move(gko::Array<ValueType>(this->dpcpp, num_nnz)),
std::move(gko::Array<IndexType>(this->dpcpp, num_nnz)),
std::move(drow_nnz));


gko::kernels::reference::csr::compute_submatrix(this->ref, this->mtx2.get(),
rspan, cspan, smat1.get());
gko::kernels::dpcpp::csr::compute_submatrix(this->dpcpp, this->dmtx2.get(),
rspan, cspan, sdmat1.get());

GKO_ASSERT_MTX_NEAR(sdmat1, smat1, 0.0);
}


TEST_F(Csr, CreateSubMatrixIsEquivalentToRef)
{
using Mtx = gko::matrix::Csr<vtype>;
set_up_apply_data();
gko::span rspan{47, 81};
gko::span cspan{2, 31};

auto smat1 = this->mtx2->create_submatrix(rspan, cspan);
auto sdmat1 = this->dmtx2->create_submatrix(rspan, cspan);

GKO_ASSERT_MTX_NEAR(sdmat1, smat1, 0.0);
}


} // namespace

0 comments on commit c3c0351

Please sign in to comment.