Review update: s/Diagonal matrix/blocks array

Co-authored-by: Tobias Ribizel <ribizel@kit.edu>
ginkgo-project · Jul 1, 2021 · 38dfff9 · 38dfff9
1 parent 36337a3
commit 38dfff9
Show file tree

Hide file tree

Showing 12 changed files with 67 additions and 49 deletions.
diff --git a/common/preconditioner/jacobi_kernels.hpp.inc b/common/preconditioner/jacobi_kernels.hpp.inc
@@ -231,9 +231,11 @@ __global__ __launch_bounds__(default_block_size) void scalar_apply(
     const auto col = tidx % num_cols;
 
     if (row < num_rows) {
+        auto diag_val =
+            diag[row] == zero<ValueType>() ? one<ValueType>() : diag[row];
         result_values[row * result_stride + col] =
             result_values[row * result_stride + col] * beta[0] +
-            alpha[0] * source_values[row * source_stride + col] / diag[row];
+            alpha[0] * source_values[row * source_stride + col] / diag_val;
     }
 }
 
@@ -249,8 +251,10 @@ __global__ __launch_bounds__(default_block_size) void simple_scalar_apply(
     const auto col = tidx % num_cols;
 
     if (row < num_rows) {
+        auto diag_val =
+            diag[row] == zero<ValueType>() ? one<ValueType>() : diag[row];
         result_values[row * result_stride + col] =
-            source_values[row * source_stride + col] / diag[row];
+            source_values[row * source_stride + col] / diag_val;
     }
 }
 

diff --git a/core/preconditioner/jacobi.cpp b/core/preconditioner/jacobi.cpp
@@ -78,7 +78,7 @@ void Jacobi<ValueType, IndexType>::apply_impl(const LinOp *b, LinOp *x) const
         [this](auto dense_b, auto dense_x) {
             if (parameters_.max_block_size == 1) {
                 this->get_executor()->run(jacobi::make_simple_scalar_apply(
-                    this->diag_.get(), dense_b, dense_x));
+                    this->blocks_, dense_b, dense_x));
             } else {
                 this->get_executor()->run(jacobi::make_simple_apply(
                     num_blocks_, parameters_.max_block_size, storage_scheme_,
@@ -98,9 +98,8 @@ void Jacobi<ValueType, IndexType>::apply_impl(const LinOp *alpha,
     precision_dispatch_real_complex<ValueType>(
         [this](auto dense_alpha, auto dense_b, auto dense_beta, auto dense_x) {
             if (parameters_.max_block_size == 1) {
-                this->get_executor()->run(
-                    jacobi::make_scalar_apply(this->diag_.get(), dense_alpha,
-                                              dense_b, dense_beta, dense_x));
+                this->get_executor()->run(jacobi::make_scalar_apply(
+                    this->blocks_, dense_alpha, dense_b, dense_beta, dense_x));
             } else {
                 this->get_executor()->run(jacobi::make_apply(
                     num_blocks_, parameters_.max_block_size, storage_scheme_,
@@ -232,8 +231,11 @@ void Jacobi<ValueType, IndexType>::generate(const LinOp *system_matrix,
     using csr_type = matrix::Csr<ValueType, IndexType>;
     const auto exec = this->get_executor();
     if (parameters_.max_block_size == 1) {
-        this->diag_ = as<DiagonalExtractable<ValueType>>(system_matrix)
-                          ->extract_diagonal();
+        auto diag = as<DiagonalExtractable<ValueType>>(system_matrix)
+                        ->extract_diagonal();
+        auto temp = gko::Array<ValueType>::view(
+            exec, system_matrix->get_size()[0], diag->get_values());
+        this->blocks_ = temp;
     } else {
         auto csr_mtx = convert_to_with_sorting<csr_type>(exec, system_matrix,
                                                          skip_sorting);

diff --git a/core/preconditioner/jacobi_kernels.hpp b/core/preconditioner/jacobi_kernels.hpp
@@ -76,7 +76,7 @@ namespace kernels {
 
 #define GKO_DECLARE_JACOBI_SIMPLE_SCALAR_APPLY_KERNEL(ValueType)          \
     void simple_scalar_apply(std::shared_ptr<const DefaultExecutor> exec, \
-                             const matrix::Diagonal<ValueType> *diag,     \
+                             const Array<ValueType> &diag,                \
                              const matrix::Dense<ValueType> *b,           \
                              matrix::Dense<ValueType> *x)
 
@@ -91,13 +91,12 @@ namespace kernels {
         const Array<ValueType> &blocks, const matrix::Dense<ValueType> *b, \
         matrix::Dense<ValueType> *x)
 
-#define GKO_DECLARE_JACOBI_SCALAR_APPLY_KERNEL(ValueType)          \
-    void scalar_apply(std::shared_ptr<const DefaultExecutor> exec, \
-                      const matrix::Diagonal<ValueType> *diag,     \
-                      const matrix::Dense<ValueType> *alpha,       \
-                      const matrix::Dense<ValueType> *b,           \
-                      const matrix::Dense<ValueType> *beta,        \
-                      matrix::Dense<ValueType> *x)
+#define GKO_DECLARE_JACOBI_SCALAR_APPLY_KERNEL(ValueType)                    \
+    void scalar_apply(                                                       \
+        std::shared_ptr<const DefaultExecutor> exec,                         \
+        const Array<ValueType> &diag, const matrix::Dense<ValueType> *alpha, \
+        const matrix::Dense<ValueType> *b,                                   \
+        const matrix::Dense<ValueType> *beta, matrix::Dense<ValueType> *x)
 
 #define GKO_DECLARE_JACOBI_TRANSPOSE_KERNEL(ValueType, IndexType)          \
     void transpose_jacobi(                                                 \

diff --git a/cuda/preconditioner/jacobi_kernels.cu b/cuda/preconditioner/jacobi_kernels.cu
@@ -245,7 +245,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 
 template <typename ValueType>
 void scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                  const matrix::Diagonal<ValueType> *diag,
+                  const Array<ValueType> &diag,
                   const matrix::Dense<ValueType> *alpha,
                   const matrix::Dense<ValueType> *b,
                   const matrix::Dense<ValueType> *beta,
@@ -259,7 +259,7 @@ void scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
     const auto grid_dim = ceildiv(num_rows * num_cols, default_block_size);
 
     const auto b_values = b->get_const_values();
-    const auto diag_values = diag->get_const_values();
+    const auto diag_values = diag.get_const_data();
     auto x_values = x->get_values();
 
     kernel::scalar_apply<<<grid_dim, default_block_size>>>(
@@ -274,7 +274,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_JACOBI_SCALAR_APPLY_KERNEL);
 
 template <typename ValueType>
 void simple_scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                         const matrix::Diagonal<ValueType> *diag,
+                         const Array<ValueType> &diag,
                          const matrix::Dense<ValueType> *b,
                          matrix::Dense<ValueType> *x)
 {
@@ -286,7 +286,7 @@ void simple_scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
     const auto grid_dim = ceildiv(num_rows * num_cols, default_block_size);
 
     const auto b_values = b->get_const_values();
-    const auto diag_values = diag->get_const_values();
+    const auto diag_values = diag.get_const_data();
     auto x_values = x->get_values();
 
     kernel::simple_scalar_apply<<<grid_dim, default_block_size>>>(

diff --git a/cuda/test/preconditioner/jacobi_kernels.cpp b/cuda/test/preconditioner/jacobi_kernels.cpp
@@ -473,10 +473,12 @@ TEST_F(Jacobi, CudaScalarLinearCombinationApplyEquivalentToRef)
     smtx->copy_from(dense_smtx.get());
     auto sb = gko::share(gko::test::generate_random_matrix<Vec>(
         dim, 3, std::uniform_int_distribution<>(1, 1),
-        std::normal_distribution<>(0.0, 1.0), engine, ref));
+        std::normal_distribution<>(0.0, 1.0), engine, ref, gko::dim<2>(dim, 3),
+        4));
     auto sx = gko::share(gko::test::generate_random_matrix<Vec>(
         dim, 3, std::uniform_int_distribution<>(1, 1),
-        std::normal_distribution<>(0.0, 1.0), engine, ref));
+        std::normal_distribution<>(0.0, 1.0), engine, ref, gko::dim<2>(dim, 3),
+        4));
 
     auto d_smtx = gko::share(Mtx::create(cuda));
     auto d_sb = gko::share(Vec::create(cuda));

diff --git a/dpcpp/preconditioner/jacobi_kernels.dp.cpp b/dpcpp/preconditioner/jacobi_kernels.dp.cpp
@@ -282,7 +282,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 
 template <typename ValueType>
 void scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                  const matrix::Diagonal<ValueType> *diag,
+                  const Array<ValueType> &diag,
                   const matrix::Dense<ValueType> *alpha,
                   const matrix::Dense<ValueType> *b,
                   const matrix::Dense<ValueType> *beta,
@@ -293,7 +293,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_JACOBI_SCALAR_APPLY_KERNEL);
 
 template <typename ValueType>
 void simple_scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                         const matrix::Diagonal<ValueType> *diag,
+                         const Array<ValueType> &diag,
                          const matrix::Dense<ValueType> *b,
                          matrix::Dense<ValueType> *x) GKO_NOT_IMPLEMENTED;
 

diff --git a/hip/preconditioner/jacobi_kernels.hip.cpp b/hip/preconditioner/jacobi_kernels.hip.cpp
@@ -259,7 +259,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 
 template <typename ValueType>
 void scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                  const matrix::Diagonal<ValueType> *diag,
+                  const Array<ValueType> &diag,
                   const matrix::Dense<ValueType> *alpha,
                   const matrix::Dense<ValueType> *b,
                   const matrix::Dense<ValueType> *beta,
@@ -273,7 +273,7 @@ void scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
     const auto grid_dim = ceildiv(num_rows * num_cols, default_block_size);
 
     const auto b_values = b->get_const_values();
-    const auto diag_values = diag->get_const_values();
+    const auto diag_values = diag.get_const_data();
     auto x_values = x->get_values();
 
     hipLaunchKernelGGL(
@@ -288,7 +288,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_JACOBI_SCALAR_APPLY_KERNEL);
 
 template <typename ValueType>
 void simple_scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                         const matrix::Diagonal<ValueType> *diag,
+                         const Array<ValueType> &diag,
                          const matrix::Dense<ValueType> *b,
                          matrix::Dense<ValueType> *x)
 {
@@ -300,7 +300,7 @@ void simple_scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
     const auto grid_dim = ceildiv(num_rows * num_cols, default_block_size);
 
     const auto b_values = b->get_const_values();
-    const auto diag_values = diag->get_const_values();
+    const auto diag_values = diag.get_const_data();
     auto x_values = x->get_values();
 
     hipLaunchKernelGGL(kernel::simple_scalar_apply, dim3(grid_dim),

diff --git a/hip/test/preconditioner/jacobi_kernels.cpp b/hip/test/preconditioner/jacobi_kernels.cpp
@@ -504,10 +504,12 @@ TEST_F(Jacobi, HipScalarLinearCombinationApplyEquivalentToRef)
     smtx->copy_from(dense_smtx.get());
     auto sb = gko::share(gko::test::generate_random_matrix<Vec>(
         dim, 3, std::uniform_int_distribution<>(1, 1),
-        std::normal_distribution<>(0.0, 1.0), engine, ref));
+        std::normal_distribution<>(0.0, 1.0), engine, ref, gko::dim<2>(dim, 3),
+        4));
     auto sx = gko::share(gko::test::generate_random_matrix<Vec>(
         dim, 3, std::uniform_int_distribution<>(1, 1),
-        std::normal_distribution<>(0.0, 1.0), engine, ref));
+        std::normal_distribution<>(0.0, 1.0), engine, ref, gko::dim<2>(dim, 3),
+        4));
 
     auto d_smtx = gko::share(Mtx::create(hip));
     auto d_sb = gko::share(Vec::create(hip));

diff --git a/include/ginkgo/core/preconditioner/jacobi.hpp b/include/ginkgo/core/preconditioner/jacobi.hpp
@@ -502,8 +502,7 @@ class Jacobi : public EnableLinOp<Jacobi<ValueType, IndexType>>,
         : EnableLinOp<Jacobi>(exec),
           num_blocks_{},
           blocks_(exec),
-          conditioning_(exec),
-          diag_(matrix::Diagonal<value_type>::create(exec))
+          conditioning_(exec)
     {
         parameters_.block_pointers.set_executor(exec);
         parameters_.storage_optimization.block_wise.set_executor(exec);
@@ -527,8 +526,7 @@ class Jacobi : public EnableLinOp<Jacobi<ValueType, IndexType>>,
           blocks_(factory->get_executor(),
                   storage_scheme_.compute_storage_space(
                       parameters_.block_pointers.get_num_elems() - 1)),
-          conditioning_(factory->get_executor()),
-          diag_(matrix::Diagonal<value_type>::create(factory->get_executor()))
+          conditioning_(factory->get_executor())
     {
         parameters_.block_pointers.set_executor(this->get_executor());
         parameters_.storage_optimization.block_wise.set_executor(
@@ -608,7 +606,6 @@ class Jacobi : public EnableLinOp<Jacobi<ValueType, IndexType>>,
     size_type num_blocks_;
     Array<value_type> blocks_;
     Array<remove_complex<value_type>> conditioning_;
-    std::shared_ptr<matrix::Diagonal<value_type>> diag_;
 };
 
 

diff --git a/omp/preconditioner/jacobi_kernels.cpp b/omp/preconditioner/jacobi_kernels.cpp
@@ -698,7 +698,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 
 template <typename ValueType>
 void scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                  const matrix::Diagonal<ValueType> *diag,
+                  const Array<ValueType> &diag,
                   const matrix::Dense<ValueType> *alpha,
                   const matrix::Dense<ValueType> *b,
                   const matrix::Dense<ValueType> *beta,
@@ -707,9 +707,11 @@ void scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
 #pragma omp parallel for
     for (size_type i = 0; i < x->get_size()[0]; ++i) {
         for (size_type j = 0; j < x->get_size()[1]; ++j) {
-            x->at(i, j) =
-                beta->at(0) * x->at(i, j) +
-                alpha->at(0) * b->at(i, j) / diag->get_const_values()[i];
+            auto diag_val = diag.get_const_data()[i] == zero<ValueType>()
+                                ? one<ValueType>()
+                                : diag.get_const_data()[i];
+            x->at(i, j) = beta->at(0) * x->at(i, j) +
+                          alpha->at(0) * b->at(i, j) / diag_val;
         }
     }
 }
@@ -719,14 +721,17 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_JACOBI_SCALAR_APPLY_KERNEL);
 
 template <typename ValueType>
 void simple_scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                         const matrix::Diagonal<ValueType> *diag,
+                         const Array<ValueType> &diag,
                          const matrix::Dense<ValueType> *b,
                          matrix::Dense<ValueType> *x)
 {
 #pragma omp parallel for
     for (size_type i = 0; i < x->get_size()[0]; ++i) {
         for (size_type j = 0; j < x->get_size()[1]; ++j) {
-            x->at(i, j) = b->at(i, j) / diag->get_const_values()[i];
+            auto diag_val = diag.get_const_data()[i] == zero<ValueType>()
+                                ? one<ValueType>()
+                                : diag.get_const_data()[i];
+            x->at(i, j) = b->at(i, j) / diag_val;
         }
     }
 }

diff --git a/omp/test/preconditioner/jacobi_kernels.cpp b/omp/test/preconditioner/jacobi_kernels.cpp
@@ -472,10 +472,12 @@ TEST_F(Jacobi, OmpScalarLinearCombinationApplyEquivalentToRef)
     smtx->copy_from(dense_smtx.get());
     auto sb = gko::share(gko::test::generate_random_matrix<Vec>(
         dim, 3, std::uniform_int_distribution<>(1, 1),
-        std::normal_distribution<>(0.0, 1.0), engine, ref));
+        std::normal_distribution<>(0.0, 1.0), engine, ref, gko::dim<2>(dim, 3),
+        4));
     auto sx = gko::share(gko::test::generate_random_matrix<Vec>(
         dim, 3, std::uniform_int_distribution<>(1, 1),
-        std::normal_distribution<>(0.0, 1.0), engine, ref));
+        std::normal_distribution<>(0.0, 1.0), engine, ref, gko::dim<2>(dim, 3),
+        4));
 
     auto d_smtx = gko::share(Mtx::create(omp));
     auto d_sb = gko::share(Vec::create(omp));

diff --git a/reference/preconditioner/jacobi_kernels.cpp b/reference/preconditioner/jacobi_kernels.cpp
@@ -562,17 +562,19 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 
 template <typename ValueType>
 void scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                  const matrix::Diagonal<ValueType> *diag,
+                  const Array<ValueType> &diag,
                   const matrix::Dense<ValueType> *alpha,
                   const matrix::Dense<ValueType> *b,
                   const matrix::Dense<ValueType> *beta,
                   matrix::Dense<ValueType> *x)
 {
     for (size_type i = 0; i < x->get_size()[0]; ++i) {
         for (size_type j = 0; j < x->get_size()[1]; ++j) {
-            x->at(i, j) =
-                beta->at(0) * x->at(i, j) +
-                alpha->at(0) * b->at(i, j) / diag->get_const_values()[i];
+            auto diag_val = diag.get_const_data()[i] == zero<ValueType>()
+                                ? one<ValueType>()
+                                : diag.get_const_data()[i];
+            x->at(i, j) = beta->at(0) * x->at(i, j) +
+                          alpha->at(0) * b->at(i, j) / diag_val;
         }
     }
 }
@@ -582,13 +584,16 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_JACOBI_SCALAR_APPLY_KERNEL);
 
 template <typename ValueType>
 void simple_scalar_apply(std::shared_ptr<const DefaultExecutor> exec,
-                         const matrix::Diagonal<ValueType> *diag,
+                         const Array<ValueType> &diag,
                          const matrix::Dense<ValueType> *b,
                          matrix::Dense<ValueType> *x)
 {
     for (size_type i = 0; i < x->get_size()[0]; ++i) {
         for (size_type j = 0; j < x->get_size()[1]; ++j) {
-            x->at(i, j) = b->at(i, j) / diag->get_const_values()[i];
+            auto diag_val = diag.get_const_data()[i] == zero<ValueType>()
+                                ? one<ValueType>()
+                                : diag.get_const_data()[i];
+            x->at(i, j) = b->at(i, j) / diag_val;
         }
     }
 }