From 1c0e4e9b4a40a1e25ddb57c78fa119b783dabea4 Mon Sep 17 00:00:00 2001 From: Fritz Goebel Date: Thu, 31 Oct 2024 11:37:57 +0100 Subject: [PATCH] Add row and column scaling --- core/distributed/dd_matrix.cpp | 73 +++++---- include/ginkgo/core/distributed/dd_matrix.hpp | 2 + test/mpi/dd_matrix.cpp | 144 ++++++++++++++++++ 3 files changed, 192 insertions(+), 27 deletions(-) diff --git a/core/distributed/dd_matrix.cpp b/core/distributed/dd_matrix.cpp index cfb2d65db00..91b1d7f603f 100644 --- a/core/distributed/dd_matrix.cpp +++ b/core/distributed/dd_matrix.cpp @@ -309,21 +309,12 @@ void DdMatrix::apply_impl( const LinOp* b, LinOp* x) const { auto exec = this->get_executor(); - auto comm = this->get_communicator(); const auto nrhs = x->get_size()[1]; - if (nrhs != rhs_buffer_->get_size()[1]) { - dim<2> local_buffer_size{rhs_buffer_->get_local_vector()->get_size()[0], - nrhs}; - dim<2> global_buffer_size{rhs_buffer_->get_size()[0], nrhs}; - lhs_buffer_ = Vector::create(exec, comm, global_buffer_size, - local_buffer_size); - rhs_buffer_ = Vector::create(exec, comm, global_buffer_size, - local_buffer_size); - } + check_and_adjust_buffer_size(nrhs); distributed::precision_dispatch_real_complex( [this](const auto dense_b, auto dense_x) { auto exec = this->get_executor(); - this->restriction_->apply(dense_b, lhs_buffer_); + restriction_->apply(dense_b, lhs_buffer_); auto local_b = gko::matrix::Dense::create( exec, lhs_buffer_->get_local_vector()->get_size(), @@ -342,7 +333,7 @@ void DdMatrix::apply_impl( local_mtx_->apply(local_b, local_x); - this->prolongation_->apply(rhs_buffer_, dense_x); + prolongation_->apply(rhs_buffer_, dense_x); }, b, x); } @@ -353,22 +344,13 @@ void DdMatrix::apply_impl( const LinOp* alpha, const LinOp* b, const LinOp* beta, LinOp* x) const { auto exec = this->get_executor(); - auto comm = this->get_communicator(); const auto nrhs = x->get_size()[1]; - if (nrhs != rhs_buffer_->get_size()[1]) { - dim<2> local_buffer_size{rhs_buffer_->get_local_vector()->get_size()[0], - nrhs}; - dim<2> global_buffer_size{rhs_buffer_->get_size()[0], nrhs}; - lhs_buffer_ = Vector::create(exec, comm, global_buffer_size, - local_buffer_size); - rhs_buffer_ = Vector::create(exec, comm, global_buffer_size, - local_buffer_size); - } + check_and_adjust_buffer_size(nrhs); distributed::precision_dispatch_real_complex( [this](const auto local_alpha, const auto dense_b, const auto local_beta, auto dense_x) { auto exec = this->get_executor(); - this->restriction_->apply(dense_b, lhs_buffer_); + restriction_->apply(dense_b, lhs_buffer_); auto local_b = gko::matrix::Dense::create( exec, lhs_buffer_->get_local_vector()->get_size(), @@ -387,23 +369,60 @@ void DdMatrix::apply_impl( local_mtx_->apply(local_b, local_x); - this->prolongation_->apply(local_alpha, rhs_buffer_, local_beta, - dense_x); + prolongation_->apply(local_alpha, rhs_buffer_, local_beta, dense_x); }, alpha, b, beta, x); } +template +void DdMatrix:: + check_and_adjust_buffer_size(const size_type nrhs) const +{ + auto exec = this->get_executor(); + auto comm = this->get_communicator(); + if (nrhs != rhs_buffer_->get_size()[1]) { + dim<2> local_buffer_size{rhs_buffer_->get_local_vector()->get_size()[0], + nrhs}; + dim<2> global_buffer_size{rhs_buffer_->get_size()[0], nrhs}; + lhs_buffer_ = Vector::create(exec, comm, global_buffer_size, + local_buffer_size); + rhs_buffer_ = Vector::create(exec, comm, global_buffer_size, + local_buffer_size); + } +} + + template void DdMatrix::col_scale( ptr_param scaling_factors) -{} +{ + auto exec = this->get_executor(); + check_and_adjust_buffer_size(1u); + size_type n_local_cols = local_mtx_->get_size()[1]; + restriction_->apply(scaling_factors, lhs_buffer_); + const auto scale_diag = gko::matrix::Diagonal::create_const( + exec, n_local_cols, + make_const_array_view(exec, n_local_cols, + lhs_buffer_->get_const_local_values())); + scale_diag->rapply(local_mtx_, local_mtx_); +} template void DdMatrix::row_scale( ptr_param scaling_factors) -{} +{ + auto exec = this->get_executor(); + check_and_adjust_buffer_size(1u); + size_type n_local_cols = local_mtx_->get_size()[1]; + restriction_->apply(scaling_factors, lhs_buffer_); + const auto scale_diag = gko::matrix::Diagonal::create_const( + exec, n_local_cols, + make_const_array_view(exec, n_local_cols, + lhs_buffer_->get_const_local_values())); + scale_diag->apply(local_mtx_, local_mtx_); +} template diff --git a/include/ginkgo/core/distributed/dd_matrix.hpp b/include/ginkgo/core/distributed/dd_matrix.hpp index 361b47c1773..8de92a65550 100644 --- a/include/ginkgo/core/distributed/dd_matrix.hpp +++ b/include/ginkgo/core/distributed/dd_matrix.hpp @@ -436,6 +436,8 @@ class DdMatrix LinOp* x) const override; private: + void check_and_adjust_buffer_size(const size_type nrhs) const; + std::vector send_offsets_; std::vector send_sizes_; std::vector recv_offsets_; diff --git a/test/mpi/dd_matrix.cpp b/test/mpi/dd_matrix.cpp index 1f127ce9156..def8bf9dac2 100644 --- a/test/mpi/dd_matrix.cpp +++ b/test/mpi/dd_matrix.cpp @@ -167,6 +167,7 @@ TYPED_TEST(DdMatrix, ReadsDistributed) non_local_prolongation[rank], 0); } + TYPED_TEST(DdMatrix, CanApplyToSingleVector) { using value_type = typename TestFixture::value_type; @@ -185,6 +186,38 @@ TYPED_TEST(DdMatrix, CanApplyToSingleVector) GKO_ASSERT_MTX_NEAR(this->y->get_local_vector(), result[rank], 0); } + +TYPED_TEST(DdMatrix, CanApplyToMultipleVectors) +{ + using value_type = typename TestFixture::value_type; + using index_type = typename TestFixture::global_index_type; + auto vec_md = + gko::matrix_data{I>{{1, 2}, + {2, 4}, + {3, 6}, + {4, 8}, + {5, 10}, + {6, 12}, + {7, 14}, + {8, 16}, + {9, 18}, + {10, 20}, + {11, 22}, + {12, 24}}}; + I> result[3] = {{{-4, -8}, {-3, -6}, {-2, -4}, {-1, -2}}, + {{0, 0}, {1, 2}, {-1, -2}, {0, 0}}, + {{1, 2}, {2, 4}, {3, 6}, {4, 8}}}; + auto rank = this->comm.rank(); + this->dist_mat->read_distributed(this->dist_input[rank], this->row_part); + this->x->read_distributed(vec_md, this->row_part); + this->y->read_distributed(vec_md, this->row_part); + + this->dist_mat->apply(this->x, this->y); + + GKO_ASSERT_MTX_NEAR(this->y->get_local_vector(), result[rank], 0); +} + + TYPED_TEST(DdMatrix, CanAdvancedApplyToSingleVector) { using value_type = typename TestFixture::value_type; @@ -206,4 +239,115 @@ TYPED_TEST(DdMatrix, CanAdvancedApplyToSingleVector) GKO_ASSERT_MTX_NEAR(this->y->get_local_vector(), result[rank], 0); } + +TYPED_TEST(DdMatrix, CanAdvancedApplyToMultipleVectors) +{ + using value_type = typename TestFixture::value_type; + using index_type = typename TestFixture::global_index_type; + using dense_vec_type = typename TestFixture::dense_vec_type; + auto vec_md = + gko::matrix_data{I>{{1, 2}, + {2, 4}, + {3, 6}, + {4, 8}, + {5, 10}, + {6, 12}, + {7, 14}, + {8, 16}, + {9, 18}, + {10, 20}, + {11, 22}, + {12, 24}}}; + I> result[3] = {{{-3, -6}, {-1, -2}, {1, 2}, {3, 6}}, + {{5, 10}, {7, 14}, {6, 12}, {8, 16}}, + {{10, 20}, {12, 24}, {14, 28}, {16, 32}}}; + auto rank = this->comm.rank(); + this->dist_mat->read_distributed(this->dist_input[rank], this->row_part); + this->x->read_distributed(vec_md, this->row_part); + this->y->read_distributed(vec_md, this->row_part); + auto alpha = gko::initialize({1.0}, this->exec); + auto beta = gko::initialize({1.0}, this->exec); + + this->dist_mat->apply(alpha, this->x, beta, this->y); + + GKO_ASSERT_MTX_NEAR(this->y->get_local_vector(), result[rank], 0); +} + + +TYPED_TEST(DdMatrix, CanColScale) +{ + using value_type = typename TestFixture::value_type; + using index_type = typename TestFixture::global_index_type; + using matrix_data = typename TestFixture::local_matrix_data; + using csr = typename TestFixture::local_matrix_type; + auto local_size = gko::dim<2>{6, 6}; + auto vec_md = gko::matrix_data{I>{ + {1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}, {9}, {10}, {11}, {12}}}; + std::array scaled_result{ + {{local_size, + {{0, 0, 2}, {0, 1, -2}, {0, 3, -4}, {1, 0, -1}, {1, 1, 6}, + {1, 2, -3}, {1, 4, -5}, {2, 1, -2}, {2, 2, 6}, {2, 5, -6}, + {3, 0, -1}, {3, 3, 6}, {3, 4, -2.5}, {4, 1, -2}, {4, 3, -2}, + {4, 4, 10}, {4, 5, -3}, {5, 2, -3}, {5, 4, -2.5}, {5, 5, 9}}}, + {local_size, + {{0, 0, 10}, {0, 1, -3}, {0, 3, -8}, {0, 4, -2}, {1, 0, -2.5}, + {1, 1, 9}, {1, 5, -9}, {2, 2, 10.5}, {2, 3, -4}, {2, 4, -4}, + {3, 0, -5}, {3, 2, -3.5}, {3, 3, 16}, {3, 5, -4.5}, {4, 0, -2.5}, + {4, 2, -7}, {4, 4, 6}, {5, 1, -6}, {5, 3, -4}, {5, 5, 13.5}}}, + {local_size, {{0, 0, 13.5}, {0, 3, -12}, {0, 5, -4}, {1, 1, 20}, + {1, 2, -11}, {1, 4, -7}, {2, 1, -10}, {2, 2, 33}, + {2, 3, -12}, {2, 5, -8}, {3, 0, -9}, {3, 2, -11}, + {3, 3, 24}, {4, 1, -10}, {4, 4, 10.5}, {4, 5, -4}, + {5, 0, -4.5}, {5, 2, -11}, {5, 4, -3.5}, {5, 5, 16}}}}}; + auto rank = this->comm.rank(); + auto res_local = csr::create(this->exec); + res_local->read(scaled_result[rank]); + this->dist_mat->read_distributed(this->dist_input[rank], this->row_part); + this->x->read_distributed(vec_md, this->row_part); + + this->dist_mat->col_scale(this->x); + + GKO_ASSERT_MTX_NEAR(gko::as(this->dist_mat->get_local_matrix()), + res_local, 0); +} + + +TYPED_TEST(DdMatrix, CanRowScale) +{ + using value_type = typename TestFixture::value_type; + using index_type = typename TestFixture::global_index_type; + using matrix_data = typename TestFixture::local_matrix_data; + using csr = typename TestFixture::local_matrix_type; + auto local_size = gko::dim<2>{6, 6}; + auto vec_md = gko::matrix_data{I>{ + {1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}, {9}, {10}, {11}, {12}}}; + std::array scaled_result{ + {{local_size, + {{0, 0, 2}, {0, 1, -1}, {0, 3, -1}, {1, 0, -2}, {1, 1, 6}, + {1, 2, -2}, {1, 4, -2}, {2, 1, -3}, {2, 2, 6}, {2, 5, -3}, + {3, 0, -4}, {3, 3, 6}, {3, 4, -2}, {4, 1, -5}, {4, 3, -2.5}, + {4, 4, 10}, {4, 5, -2.5}, {5, 2, -6}, {5, 4, -3}, {5, 5, 9}}}, + {local_size, + {{0, 0, 10}, {0, 1, -2.5}, {0, 3, -5}, {0, 4, -2.5}, {1, 0, -3}, + {1, 1, 9}, {1, 5, -6}, {2, 2, 10.5}, {2, 3, -3.5}, {2, 4, -7}, + {3, 0, -8}, {3, 2, -4}, {3, 3, 16}, {3, 5, -4}, {4, 0, -2}, + {4, 2, -4}, {4, 4, 6}, {5, 1, -9}, {5, 3, -4.5}, {5, 5, 13.5}}}, + {local_size, + {{0, 0, 13.5}, {0, 3, -9}, {0, 5, -4.5}, {1, 1, 20}, {1, 2, -10}, + {1, 4, -10}, {2, 1, -11}, {2, 2, 33}, {2, 3, -11}, {2, 5, -11}, + {3, 0, -12}, {3, 2, -12}, {3, 3, 24}, {4, 1, -7}, {4, 4, 10.5}, + {4, 5, -3.5}, {5, 0, -4}, {5, 2, -8}, {5, 4, -4}, {5, 5, 16}}}}}; + auto rank = this->comm.rank(); + auto res_local = csr::create(this->exec); + res_local->read(scaled_result[rank]); + this->dist_mat->read_distributed(this->dist_input[rank], this->row_part); + this->x->read_distributed(vec_md, this->row_part); + + this->dist_mat->row_scale(this->x); + + GKO_ASSERT_MTX_NEAR(gko::as(this->dist_mat->get_local_matrix()), + res_local, 0); +} + + #endif \ No newline at end of file