From e6a5c3579b3a6ae09ece383c597817de257b435a Mon Sep 17 00:00:00 2001 From: Peng Xingliang <91927439+pxlxingliang@users.noreply.github.com> Date: Sat, 22 Jun 2024 09:15:13 +0800 Subject: [PATCH] add function get_global_row/col_size() in parallel_2d (#4440) * add function get_global_row/col_size() in parallel_2d * [pre-commit.ci lite] apply automatic fixes * add const --------- Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com> --- source/module_basis/module_ao/parallel_2d.cpp | 54 ++++---- source/module_basis/module_ao/parallel_2d.h | 78 ++++++++---- .../module_ao/test/parallel_2d_test.cpp | 120 ++++++++++-------- 3 files changed, 146 insertions(+), 106 deletions(-) diff --git a/source/module_basis/module_ao/parallel_2d.cpp b/source/module_basis/module_ao/parallel_2d.cpp index 4ac2bfefe5..3a225eec86 100644 --- a/source/module_basis/module_ao/parallel_2d.cpp +++ b/source/module_basis/module_ao/parallel_2d.cpp @@ -1,16 +1,34 @@ #include "parallel_2d.h" -#include -#include - #include "module_base/blacs_connector.h" #include "module_base/scalapack_connector.h" +#include +#include + bool Parallel_2D::in_this_processor(const int iw1_all, const int iw2_all) const { return global2local_row(iw1_all) != -1 && global2local_col(iw2_all) != -1; } +int Parallel_2D::get_global_row_size() const +{ +#ifdef __MPI + return desc[2]; +#else + return nrow; +#endif +} + +int Parallel_2D::get_global_col_size() const +{ +#ifdef __MPI + return desc[3]; +#else + return ncol; +#endif +} + #ifdef __MPI void Parallel_2D::_init_proc_grid(const MPI_Comm comm, const bool mode) { @@ -19,9 +37,15 @@ void Parallel_2D::_init_proc_grid(const MPI_Comm comm, const bool mode) int num_proc = 0; MPI_Comm_size(comm, &num_proc); dim0 = static_cast(std::sqrt(num_proc + 0.5)); - while (dim1 = num_proc / dim0, dim0 * dim1 != num_proc) { --dim0; } + while (dim1 = num_proc / dim0, dim0 * dim1 != num_proc) + { + --dim0; + } - if (mode) { std::swap(dim0, dim1); } + if (mode) + { + std::swap(dim0, dim1); + } // create a 2D Cartesian MPI communicator (row-major by default) int period[2] = {1, 1}; @@ -42,7 +66,6 @@ void Parallel_2D::_init_proc_grid(const MPI_Comm comm, const bool mode) // ***This needs to be verified*** } - void Parallel_2D::_set_dist_info(const int mg, const int ng, const int nb) { this->nb = nb; @@ -75,28 +98,14 @@ void Parallel_2D::_set_dist_info(const int mg, const int ng, const int nb) } } - -int Parallel_2D::init( - const int mg, - const int ng, - const int nb, - const MPI_Comm comm, - const bool mode -) +int Parallel_2D::init(const int mg, const int ng, const int nb, const MPI_Comm comm, const bool mode) { _init_proc_grid(comm, mode); _set_dist_info(mg, ng, nb); return nrow == 0 || ncol == 0; } - -int Parallel_2D::set( - const int mg, - const int ng, - const int nb, - const MPI_Comm comm_2D, - const int blacs_ctxt -) +int Parallel_2D::set(const int mg, const int ng, const int nb, const MPI_Comm comm_2D, const int blacs_ctxt) { this->comm_2D = comm_2D; this->blacs_ctxt = blacs_ctxt; @@ -127,4 +136,3 @@ void Parallel_2D::set_serial(const int mg, const int ng) blacs_ctxt = -1; #endif } - diff --git a/source/module_basis/module_ao/parallel_2d.h b/source/module_basis/module_ao/parallel_2d.h index 9905973698..04ddd3a69c 100644 --- a/source/module_basis/module_ao/parallel_2d.h +++ b/source/module_basis/module_ao/parallel_2d.h @@ -1,8 +1,8 @@ #ifndef _PARALLEL_2D_H_ #define _PARALLEL_2D_H_ -#include #include +#include #ifdef __MPI #include @@ -12,36 +12,66 @@ /// 2D-block-cyclic parallel distribution of an arbitrary matrix. class Parallel_2D { -public: + public: Parallel_2D() = default; ~Parallel_2D() = default; /// number of local rows - int get_row_size() const { return nrow; }; + int get_row_size() const + { + return nrow; + }; /// number of local columns - int get_col_size() const { return ncol; }; + int get_col_size() const + { + return ncol; + }; + + /// number of global rows + int get_global_row_size() const; + + /// number of global columns + int get_global_col_size() const; /// number of local matrix elements - int64_t get_local_size() const { return nloc; }; + int64_t get_local_size() const + { + return nloc; + }; /// get the local index of a global index (row) - int global2local_row(const int igr) const { return global2local_row_[igr]; } + int global2local_row(const int igr) const + { + return global2local_row_[igr]; + } /// get the local index of a global index (col) - int global2local_col(const int igc) const { return global2local_col_[igc]; } + int global2local_col(const int igc) const + { + return global2local_col_[igc]; + } /// get the global index of a local index (row) - int local2global_row(const int ilr) const { return local2global_row_[ilr]; } + int local2global_row(const int ilr) const + { + return local2global_row_[ilr]; + } /// get the global index of a local index (col) - int local2global_col(const int ilc) const { return local2global_col_[ilc]; } + int local2global_col(const int ilc) const + { + return local2global_col_[ilc]; + } /// check whether a global index is in this process bool in_this_processor(const int iw1_all, const int iw2_all) const; /// side length of 2d square block - int get_block_size() const { return nb; }; + int get_block_size() const + { + return nb; + }; #ifdef __MPI /** @@ -49,26 +79,22 @@ class Parallel_2D * and set up the info of a block-cyclic distribution. * */ - int init( - const int mg, - const int ng, - const int nb, // square block is assumed - const MPI_Comm comm, - bool mode = 0 - ); + int init(const int mg, + const int ng, + const int nb, // square block is assumed + const MPI_Comm comm, + bool mode = 0); /** * @brief Set up the info of a block-cyclic distribution using given * MPI communicator and BLACS context. * */ - int set( - const int mg, - const int ng, - const int nb, // square block is assumed - const MPI_Comm comm_2D, - const int blacs_ctxt - ); + int set(const int mg, + const int ng, + const int nb, // square block is assumed + const MPI_Comm comm_2D, + const int blacs_ctxt); /// BLACS context int blacs_ctxt = -1; @@ -104,9 +130,7 @@ class Parallel_2D /// test parameter int testpb = 0; - -protected: - + protected: /// map from global index to local index std::vector global2local_row_; std::vector global2local_col_; diff --git a/source/module_basis/module_ao/test/parallel_2d_test.cpp b/source/module_basis/module_ao/test/parallel_2d_test.cpp index 8e73713e38..206231398a 100644 --- a/source/module_basis/module_ao/test/parallel_2d_test.cpp +++ b/source/module_basis/module_ao/test/parallel_2d_test.cpp @@ -1,34 +1,35 @@ -#include "gtest/gtest.h" #include "../parallel_2d.h" + +#include "gtest/gtest.h" /*********************************************************** * unit test of class "Parallel_2D" ***********************************************************/ - /* Tested functions (in order): - * - * - init - * initialize the MPI & BLACS 2d Cartesian grid and set up - * the info of block-cyclic distribution. - * - * - set_serial (serial) - * set the local(=global) sizes. - * - * - some getters: - * - get_row_size, get_col_size, get_local_size, get_block_size - * - in_this_processor - * - * Result check: - * - local sizes - * - index maps - * - desc[9] - ***********************************************************/ +/* Tested functions (in order): + * + * - init + * initialize the MPI & BLACS 2d Cartesian grid and set up + * the info of block-cyclic distribution. + * + * - set_serial (serial) + * set the local(=global) sizes. + * + * - some getters: + * - get_row_size, get_col_size, get_local_size, get_block_size + * - in_this_processor + * + * Result check: + * - local sizes + * - index maps + * - desc[9] + ***********************************************************/ class test_para2d : public testing::Test { -protected: + protected: int dsize; int my_rank = 0; - std::vector> sizes{ {30, 35}, {49, 94}, {57, 57} }; - std::vector nbs{ 1,2,3 }; + std::vector> sizes{{30, 35}, {49, 94}, {57, 57}}; + std::vector nbs{1, 2, 3}; #ifdef __MPI void SetUp() override { @@ -44,42 +45,43 @@ class test_para2d : public testing::Test #ifdef __MPI TEST_F(test_para2d, Divide2D) { - for (auto& size : sizes) + for (auto& size: sizes) { int gr = size.first; int gc = size.second; - for (auto nb : nbs) + for (auto nb: nbs) { Parallel_2D p2d; - for (auto mode : { 0,1 }) + for (auto mode: {0, 1}) { p2d.init(gr, gc, nb, MPI_COMM_WORLD, mode); EXPECT_EQ(p2d.get_block_size(), nb); - //1. dim0 and dim1 + // 1. dim0 and dim1 EXPECT_EQ(p2d.dim0 * p2d.dim1, dsize); - if (mode)EXPECT_LE(p2d.dim1, p2d.dim0); - else EXPECT_LE(p2d.dim0, p2d.dim1); + if (mode) + EXPECT_LE(p2d.dim1, p2d.dim0); + else + EXPECT_LE(p2d.dim0, p2d.dim1); - //2. MPI 2d communicator + // 2. MPI 2d communicator EXPECT_NE(p2d.comm_2D, MPI_COMM_NULL); - //3. local2global and local sizes + // 3. local2global and local sizes int lr = p2d.get_row_size(); int lc = p2d.get_col_size(); EXPECT_EQ(lr * lc, p2d.get_local_size()); - auto cal_lsize = [](const int& gsize, const int& nb, const int& np, const int& pcoord) -> int - { - int nblock = gsize / nb; - return nblock / np * nb + static_cast(nblock % np > pcoord) * nb //full blocks' contribution - + static_cast(nblock % np == pcoord) * (gsize % nb); // the last block's contribution - }; + auto cal_lsize = [](const int& gsize, const int& nb, const int& np, const int& pcoord) -> int { + int nblock = gsize / nb; + return nblock / np * nb + static_cast(nblock % np > pcoord) * nb // full blocks' contribution + + static_cast(nblock % np == pcoord) * (gsize % nb); // the last block's contribution + }; EXPECT_EQ(lr, cal_lsize(gr, nb, p2d.dim0, p2d.coord[0])); EXPECT_EQ(lc, cal_lsize(gc, nb, p2d.dim1, p2d.coord[1])); - //4. ScaLAPACK descriptor + // 4. ScaLAPACK descriptor EXPECT_EQ(p2d.desc[0], 1); EXPECT_EQ(p2d.desc[1], p2d.blacs_ctxt); EXPECT_EQ(p2d.desc[2], gr); @@ -90,22 +92,25 @@ TEST_F(test_para2d, Divide2D) EXPECT_EQ(p2d.desc[7], 0); EXPECT_EQ(p2d.desc[8], lr); - //5. global2local - auto sum_array = [&p2d](const int& gr, const int& gc) -> std::pair - { - int sum_row = 0; int sum_col = 0; - for (int i = 0; i < gr; ++i) - sum_row += p2d.global2local_row(i); - for (int i = 0; i < gc; ++i) - sum_col += p2d.global2local_col(i); - return { sum_row, sum_col }; - }; + // 5. global2local + auto sum_array = [&p2d](const int& gr, const int& gc) -> std::pair { + int sum_row = 0; + int sum_col = 0; + for (int i = 0; i < gr; ++i) + sum_row += p2d.global2local_row(i); + for (int i = 0; i < gc; ++i) + sum_col += p2d.global2local_col(i); + return {sum_row, sum_col}; + }; std::pair sumrc = sum_array(gr, gc); EXPECT_EQ(std::get<0>(sumrc), lr * (lr - 1) / 2 - (gr - lr)); EXPECT_EQ(std::get<1>(sumrc), lc * (lc - 1) / 2 - (gc - lc)); - for (int i = 0;i < lr;++i) - for (int j = 0;j < lc;++j) + for (int i = 0; i < lr; ++i) + for (int j = 0; j < lc; ++j) EXPECT_TRUE(p2d.in_this_processor(p2d.local2global_row(i), p2d.local2global_col(j))); + + EXPECT_EQ(p2d.get_global_row_size(), gr); + EXPECT_EQ(p2d.get_global_col_size(), gc); } } } @@ -113,12 +118,12 @@ TEST_F(test_para2d, Divide2D) TEST_F(test_para2d, DescReuseCtxt) { - for (auto nb : nbs) + for (auto nb: nbs) { Parallel_2D p1; p1.init(sizes[0].first, sizes[0].second, nb, MPI_COMM_WORLD); - Parallel_2D p2; // use 2 different sizes, but they can share the same ctxt + Parallel_2D p2; // use 2 different sizes, but they can share the same ctxt p2.set(sizes[1].first, sizes[1].second, nb, p1.comm_2D, p1.blacs_ctxt); EXPECT_EQ(p1.desc[1], p2.desc[1]); @@ -131,7 +136,7 @@ TEST_F(test_para2d, DescReuseCtxt) #else TEST_F(test_para2d, Serial) { - for (auto& size : sizes) + for (auto& size: sizes) { int gr = size.first; int gc = size.second; @@ -152,11 +157,15 @@ TEST_F(test_para2d, Serial) for (int i = 0; i < gc; ++i) EXPECT_EQ(p2d.local2global_col(i), i); - //3. global2local - for (int i = 0;i < gr;++i) + // 3. global2local + for (int i = 0; i < gr; ++i) EXPECT_EQ(p2d.global2local_row(i), i); - for (int i = 0;i < gc;++i) + for (int i = 0; i < gc; ++i) EXPECT_EQ(p2d.global2local_col(i), i); + + // 4. get_global_row_size, get_global_col_size + EXPECT_EQ(p2d.get_global_row_size(), gr); + EXPECT_EQ(p2d.get_global_col_size(), gc); } } #endif @@ -173,4 +182,3 @@ int main(int argc, char** argv) #endif return result; } -