diff --git a/common/preconditioner/jacobi_kernels.cpp b/common/preconditioner/jacobi_kernels.cpp index 85d4bce3ff7..442a1b47809 100644 --- a/common/preconditioner/jacobi_kernels.cpp +++ b/common/preconditioner/jacobi_kernels.cpp @@ -59,7 +59,7 @@ void scalar_conj(std::shared_ptr exec, [] GKO_KERNEL(auto elem, auto diag, auto conj_diag) { conj_diag[elem] = conj(diag[elem]); }, - diag.get_num_elems(), diag.get_const_data(), conj_diag.get_data()); + diag.get_num_elems(), diag, conj_diag); } GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_JACOBI_SCALAR_CONJ_KERNEL); @@ -69,14 +69,12 @@ template void invert_diagonal(std::shared_ptr exec, const Array &diag, Array &inv_diag) { - auto one_val = one(); run_kernel( exec, - [] GKO_KERNEL(auto elem, auto diag, auto inv_diag, auto one_val) { - inv_diag[elem] = one_val / diag[elem]; + [] GKO_KERNEL(auto elem, auto diag, auto inv_diag) { + inv_diag[elem] = safe_divide(one(diag[elem]), diag[elem]); }, - diag.get_num_elems(), diag.get_const_data(), inv_diag.get_data(), - one_val); + diag.get_num_elems(), diag, inv_diag); } GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_JACOBI_INVERT_DIAGONAL_KERNEL); @@ -98,7 +96,7 @@ void scalar_apply(std::shared_ptr exec, x(row, col) = beta[col] * x(row, col) + alpha[col] * b(row, col) * diag[row]; }, - x->get_size(), diag.get_const_data(), alpha->get_const_values(), b, + x->get_size(), diag, alpha->get_const_values(), b, beta->get_const_values(), x); } else { run_kernel( @@ -108,7 +106,7 @@ void scalar_apply(std::shared_ptr exec, x(row, col) = beta[0] * x(row, col) + alpha[0] * b(row, col) * diag[row]; }, - x->get_size(), diag.get_const_data(), alpha->get_const_values(), b, + x->get_size(), diag, alpha->get_const_values(), b, beta->get_const_values(), x); } } @@ -127,7 +125,7 @@ void simple_scalar_apply(std::shared_ptr exec, [] GKO_KERNEL(auto row, auto col, auto diag, auto b, auto x) { x(row, col) = b(row, col) * diag[row]; }, - x->get_size(), diag.get_const_data(), b, x); + x->get_size(), diag, b, x); } GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE( @@ -137,20 +135,17 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE( template void scalar_convert_to_dense(std::shared_ptr exec, const Array &blocks, - ValueType *result_values, - const gko::dim<2> &matrix_size, - size_type result_stride) + matrix::Dense *result) { run_kernel( exec, - [] GKO_KERNEL(auto row, auto col, auto stride, auto diag, - auto result_values) { - result_values[row * stride + col] = zero(); + [] GKO_KERNEL(auto row, auto col, auto diag, auto result) { + result(row, col) = zero(diag[row]); if (row == col) { - result_values[row * stride + col] = diag[row]; + result(row, col) = diag[row]; } }, - matrix_size, result_stride, blocks.get_const_data(), result_values); + result->get_size(), blocks, result); } GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE( diff --git a/core/preconditioner/jacobi.cpp b/core/preconditioner/jacobi.cpp index d52d2959f6d..d62d97ec0b5 100644 --- a/core/preconditioner/jacobi.cpp +++ b/core/preconditioner/jacobi.cpp @@ -124,8 +124,7 @@ void Jacobi::convert_to( auto exec = this->get_executor(); auto tmp = matrix::Dense::create(exec, this->get_size()); if (parameters_.max_block_size == 1) { - exec->run(jacobi::make_scalar_convert_to_dense( - blocks_, tmp->get_values(), tmp->get_size(), tmp->get_stride())); + exec->run(jacobi::make_scalar_convert_to_dense(blocks_, tmp.get())); } else { exec->run(jacobi::make_convert_to_dense( num_blocks_, parameters_.storage_optimization.block_wise, diff --git a/core/preconditioner/jacobi_kernels.hpp b/core/preconditioner/jacobi_kernels.hpp index 2960e175de4..7745438abda 100644 --- a/core/preconditioner/jacobi_kernels.hpp +++ b/core/preconditioner/jacobi_kernels.hpp @@ -130,11 +130,10 @@ namespace kernels { &storage_scheme, \ Array &out_blocks) -#define GKO_DECLARE_JACOBI_SCALAR_CONVERT_TO_DENSE_KERNEL(ValueType) \ - void scalar_convert_to_dense( \ - std::shared_ptr exec, \ - const Array &blocks, ValueType *result_values, \ - const gko::dim<2> &mat_size, size_type result_stride) +#define GKO_DECLARE_JACOBI_SCALAR_CONVERT_TO_DENSE_KERNEL(ValueType) \ + void scalar_convert_to_dense(std::shared_ptr exec, \ + const Array &blocks, \ + matrix::Dense *result) #define GKO_DECLARE_JACOBI_CONVERT_TO_DENSE_KERNEL(ValueType, IndexType) \ void convert_to_dense( \ diff --git a/reference/preconditioner/jacobi_kernels.cpp b/reference/preconditioner/jacobi_kernels.cpp index 9799329f773..e8abd8b8958 100644 --- a/reference/preconditioner/jacobi_kernels.cpp +++ b/reference/preconditioner/jacobi_kernels.cpp @@ -613,8 +613,10 @@ void invert_diagonal(std::shared_ptr exec, const Array &diag, Array &inv_diag) { for (size_type i = 0; i < diag.get_num_elems(); ++i) { - inv_diag.get_data()[i] = - static_cast(1.0) / diag.get_const_data()[i]; + auto diag_val = diag.get_const_data()[i] == zero() + ? one() + : diag.get_const_data()[i]; + inv_diag.get_data()[i] = one() / diag_val; } } @@ -696,16 +698,14 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( template void scalar_convert_to_dense(std::shared_ptr exec, const Array &blocks, - ValueType *result_values, - const gko::dim<2> &matrix_size, - size_type result_stride) + matrix::Dense *result) { + auto matrix_size = result->get_size(); for (size_type i = 0; i < matrix_size[0]; ++i) { for (size_type j = 0; j < matrix_size[1]; ++j) { - result_values[i * result_stride + j] = zero(); + result->at(i, j) = zero(); if (i == j) { - result_values[i * result_stride + j] = - blocks.get_const_data()[i]; + result->at(i, j) = blocks.get_const_data()[i]; } } } diff --git a/reference/test/preconditioner/jacobi.cpp b/reference/test/preconditioner/jacobi.cpp index 390a575290e..330195cee6a 100644 --- a/reference/test/preconditioner/jacobi.cpp +++ b/reference/test/preconditioner/jacobi.cpp @@ -107,9 +107,7 @@ class Jacobi : public ::testing::Test { template void init_array(T *arr, std::initializer_list vals) { - for (auto elem : vals) { - *(arr++) = elem; - } + std::copy(std::begin(vals), std::end(vals), arr); } template @@ -325,11 +323,15 @@ TYPED_TEST(Jacobi, ScalarJacobiConvertsToDense) dense_j->copy_from(scalar_j.get()); auto j_val = scalar_j->get_blocks(); - EXPECT_EQ(dense_j->at(0, 0), j_val[0]); - EXPECT_EQ(dense_j->at(1, 1), j_val[1]); - EXPECT_EQ(dense_j->at(2, 2), j_val[2]); - EXPECT_EQ(dense_j->at(3, 3), j_val[3]); - EXPECT_EQ(dense_j->at(4, 4), j_val[4]); + for (auto i = 0; i < dense_j->get_size()[0]; ++i) { + for (auto j = 0; j < dense_j->get_size()[1]; ++j) { + if (i == j) { + EXPECT_EQ(dense_j->at(i, j), j_val[j]); + } else { + EXPECT_EQ(dense_j->at(i, j), value_type{0.0}); + } + } + } } @@ -349,6 +351,12 @@ TYPED_TEST(Jacobi, ScalarJacobiCanBeTransposed) auto trans_j = gko::as(t_j.get())->get_blocks(); auto scal_j = scalar_j->get_blocks(); + ASSERT_EQ(scalar_j->get_size(), gko::dim<2>(5, 5)); + ASSERT_EQ(scalar_j->get_num_stored_elements(), 5); + ASSERT_EQ(scalar_j->get_parameters().max_block_size, 1); + ASSERT_EQ(scalar_j->get_num_blocks(), 5); + ASSERT_EQ(scalar_j->get_parameters().block_pointers.get_const_data(), + nullptr); EXPECT_EQ(trans_j[0], scal_j[0]); EXPECT_EQ(trans_j[1], scal_j[1]); EXPECT_EQ(trans_j[2], scal_j[2]); @@ -360,9 +368,7 @@ TYPED_TEST(Jacobi, ScalarJacobiCanBeTransposed) template void init_array(T *arr, std::initializer_list vals) { - for (auto elem : vals) { - *(arr++) = elem; - } + std::copy(std::begin(vals), std::end(vals), arr); } @@ -390,6 +396,12 @@ TEST(Jacobi, ScalarJacobiCanBeConjTransposed) auto trans_j = gko::as(t_j.get())->get_blocks(); auto scal_j = scalar_j->get_blocks(); + ASSERT_EQ(scalar_j->get_size(), gko::dim<2>(5, 5)); + ASSERT_EQ(scalar_j->get_num_stored_elements(), 5); + ASSERT_EQ(scalar_j->get_parameters().max_block_size, 1); + ASSERT_EQ(scalar_j->get_num_blocks(), 5); + ASSERT_EQ(scalar_j->get_parameters().block_pointers.get_const_data(), + nullptr); EXPECT_EQ(trans_j[0], gko::conj(scal_j[0])); EXPECT_EQ(trans_j[1], gko::conj(scal_j[1])); EXPECT_EQ(trans_j[2], gko::conj(scal_j[2]));