Add ref and omp count diag kernel, tests for rem diag kernel.

ginkgo-project · Sep 30, 2019 · 27078b3 · 27078b3
1 parent 8f7d48f
commit 27078b3
Show file tree

Hide file tree

Showing 5 changed files with 116 additions and 5 deletions.
diff --git a/include/ginkgo/core/matrix/sparsity.hpp b/include/ginkgo/core/matrix/sparsity.hpp
@@ -83,7 +83,6 @@ class Sparsity : public EnableLinOp<Sparsity<ValueType, IndexType>>,
 
  using value_type = ValueType;
  using index_type = IndexType;
- using value_type = ValueType;
 
  using mat_data = matrix_data<ValueType, IndexType>;
 

diff --git a/omp/matrix/sparsity_kernels.cpp b/omp/matrix/sparsity_kernels.cpp
@@ -125,7 +125,21 @@ template <typename ValueType, typename IndexType>
 void count_num_diagonal_elements(
  std::shared_ptr<const OmpExecutor> exec,
  const matrix::Sparsity<ValueType, IndexType> *matrix,
- size_type &num_diagonal_elements) GKO_NOT_IMPLEMENTED;
+ size_type &num_diagonal_elements)
+{
+ auto num_rows = matrix->get_size()[0];
+ auto row_ptrs = matrix->get_const_row_ptrs();
+ auto col_idxs = matrix->get_const_col_idxs();
+ size_type num_diag = 0;
+ for (auto i = 0; i < num_rows; ++i) {
+ for (auto j = row_ptrs[i]; j < row_ptrs[i + 1]; ++j) {
+ if (col_idxs[j] == i) {
+ num_diag++;
+ }
+ }
+ }
+ num_diagonal_elements = num_diag;
+}
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
  GKO_DECLARE_SPARSITY_COUNT_NUM_DIAGONAL_ELEMENTS_KERNEL);
@@ -151,6 +165,7 @@ void remove_diagonal_elements(std::shared_ptr<const OmpExecutor> exec,
  }
  }
  }
+#pragma omp parallel for
  for (auto i = 0; i < temp_idxs.size(); ++i) {
  adj_idxs[i] = temp_idxs[i];
  }

diff --git a/omp/test/matrix/sparsity_kernels.cpp b/omp/test/matrix/sparsity_kernels.cpp
@@ -37,6 +37,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <random>
 #include <utility>
 
+
 #include <gtest/gtest.h>
 
 
@@ -80,8 +81,8 @@ class Sparsity : public ::testing::Test {
  {
  return gko::test::generate_random_sparsity_matrix<MtxType>(
  num_rows, num_cols,
- std::uniform_int_distribution<>(min_nnz_row, num_cols), rand_engine,
- ref);
+ std::uniform_int_distribution<>(min_nnz_row, num_cols), 1.0,
+ rand_engine, ref);
  }
 
  void set_up_apply_data(int num_vectors = 1)
@@ -217,6 +218,41 @@ TEST_F(Sparsity, TransposeIsEquivalentToRef)
 }
 
 
+TEST_F(Sparsity, CountsNumberOfDiagElementsIsEqualToRef)
+{
+ set_up_apply_data();
+ gko::size_type num_diags = 0;
+ gko::size_type d_num_diags = 0;
+
+ gko::kernels::reference::sparsity::count_num_diagonal_elements(
+ ref, mtx.get(), num_diags);
+ gko::kernels::omp::sparsity::count_num_diagonal_elements(omp, dmtx.get(),
+ d_num_diags);
+
+ ASSERT_EQ(d_num_diags, num_diags);
+}
+
+
+TEST_F(Sparsity, RemovesDiagElementsKernelIsEquivalentToRef)
+{
+ set_up_apply_data();
+ gko::size_type num_diags = 0;
+ gko::kernels::reference::sparsity::count_num_diagonal_elements(
+ ref, mtx.get(), num_diags);
+ auto tmp =
+ Mtx::create(ref, mtx->get_size(), mtx->get_num_nonzeros() - num_diags);
+ auto d_tmp = Mtx::create(omp, dmtx->get_size(),
+ dmtx->get_num_nonzeros() - num_diags);
+ gko::kernels::reference::sparsity::remove_diagonal_elements(
+ ref, tmp.get(), mtx->get_const_row_ptrs(), mtx->get_const_col_idxs());
+ gko::kernels::omp::sparsity::remove_diagonal_elements(
+ omp, d_tmp.get(), dmtx->get_const_row_ptrs(),
+ dmtx->get_const_col_idxs());
+
+ GKO_ASSERT_MTX_NEAR(tmp.get(), d_tmp.get(), 0.0);
+}
+
+
 TEST_F(Sparsity, RecognizeSortedMatrixIsEquivalentToRef)
 {
  set_up_apply_data();

diff --git a/reference/matrix/sparsity_kernels.cpp b/reference/matrix/sparsity_kernels.cpp
@@ -120,7 +120,21 @@ template <typename ValueType, typename IndexType>
 void count_num_diagonal_elements(
  std::shared_ptr<const ReferenceExecutor> exec,
  const matrix::Sparsity<ValueType, IndexType> *matrix,
- size_type &num_diagonal_elements) GKO_NOT_IMPLEMENTED;
+ size_type &num_diagonal_elements)
+{
+ auto num_rows = matrix->get_size()[0];
+ auto row_ptrs = matrix->get_const_row_ptrs();
+ auto col_idxs = matrix->get_const_col_idxs();
+ size_type num_diag = 0;
+ for (auto i = 0; i < num_rows; ++i) {
+ for (auto j = row_ptrs[i]; j < row_ptrs[i + 1]; ++j) {
+ if (col_idxs[j] == i) {
+ num_diag++;
+ }
+ }
+ }
+ num_diagonal_elements = num_diag;
+}
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
  GKO_DECLARE_SPARSITY_COUNT_NUM_DIAGONAL_ELEMENTS_KERNEL);

diff --git a/reference/test/matrix/sparsity_kernels.cpp b/reference/test/matrix/sparsity_kernels.cpp
@@ -272,6 +272,53 @@ TEST_F(Sparsity, NonSquareMtxIsTransposable)
 }
 
 
+TEST_F(Sparsity, CountsCorrectNumberOfDiagonalElements)
+{
+ // clang-format off
+ auto mtx2 = gko::initialize<gko::matrix::Sparsity<>>(
+ {{1.0, 1.0, 1.0},
+ {0.0, 1.0, 0.0},
+ {0.0, 1.0, 1.0}}, exec);
+ auto mtx_s = gko::initialize<gko::matrix::Sparsity<>>(
+ {{1.0, 1.0, 1.0},
+ {0.0, 0.0, 0.0},
+ {0.0, 1.0, 1.0}}, exec);
+ // clang-format on
+ gko::size_type m1_num_diags = 0;
+ gko::size_type ms_num_diags = 0;
+ gko::kernels::reference::sparsity::count_num_diagonal_elements(
+ exec, mtx2.get(), m1_num_diags);
+ gko::kernels::reference::sparsity::count_num_diagonal_elements(
+ exec, mtx_s.get(), ms_num_diags);
+
+ ASSERT_EQ(m1_num_diags, 3);
+ ASSERT_EQ(ms_num_diags, 2);
+}
+
+
+TEST_F(Sparsity, RemovesDiagonalElements)
+{
+ // clang-format off
+ auto mtx2 = gko::initialize<gko::matrix::Sparsity<>>(
+ {{1.0, 1.0, 1.0},
+ {0.0, 1.0, 0.0},
+ {0.0, 1.0, 1.0}}, exec);
+ auto mtx_s = gko::initialize<gko::matrix::Sparsity<>>(
+ {{0.0, 1.0, 1.0},
+ {0.0, 0.0, 0.0},
+ {0.0, 1.0, 0.0}}, exec);
+ // clang-format on
+ auto tmp_mtx = gko::matrix::Sparsity<>::create(exec, mtx_s->get_size(),
+ mtx_s->get_num_nonzeros());
+ tmp_mtx->copy_from(mtx2.get());
+ gko::kernels::reference::sparsity::remove_diagonal_elements(
+ exec, tmp_mtx.get(), mtx2->get_const_row_ptrs(),
+ mtx2->get_const_col_idxs());
+
+ GKO_ASSERT_MTX_NEAR(tmp_mtx.get(), mtx_s.get(), 0.0);
+}
+
+
 TEST_F(Sparsity, SquareMtxIsConvertibleToAdjacencyMatrix)
 {
  // clang-format off