diff --git a/AGENTS.md b/AGENTS.md index 57fea4ed84a..9087415ab40 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -26,6 +26,15 @@ rules. Read the complete governance document before making or reviewing changes: 8. Declare one variable per line; do not use comma-separated declarations. 9. Do not call MPI routines directly; use the internally-guarded wrappers (e.g., `Parallel_Reduce::reduce_*`, `Parallel_Common::bcast_*`) instead. + 10. Do not write new `#define private public` or `#define protected public` + access hacks in test files. If a unit test needs to inspect internal + state, either promote the member visibility explicitly or add a + public test-only accessor. + 11. New unit test source files shall be named `test_.cpp`, + matching the source file they exercise. For example, the test for + `rhog_io.cpp` shall be `test_rhog_io.cpp`. This naming keeps the + file-to-test relationship discoverable and consistent across the + repository. Historical tests are not required to be renamed. - Use LF line endings for text files. Only `.bat` and `.cmd` files may use CRLF. - Keep source file additions deterministic: update the relevant `CMakeLists.txt` or explain why the file is generated or included indirectly. diff --git a/README20260902 b/README20260902 new file mode 100644 index 00000000000..c7596ba9741 --- /dev/null +++ b/README20260902 @@ -0,0 +1,65 @@ +# para* 重构计划(分支 2026-09-02-b) + +## 背景与起点 + +- **分支**:`2026-09-02-b`,起点 commit `2a7696f5d`(step-0 cleanup source_base parallel_*) +- **基线状态**:`source_base/` 下有 14 个 `parallel_*` 文件(2d/cell/comm/common/device/global/grid/reduce),**没有** `module_parallel/` 目录,**没有** ParallelPartition——从 0 开始 +- 之前分支上尝试过 ParallelPartition(8 个裸 MPI_Comm 成员)和 ParaTag enum 两套方案,都已推倒 + +## 核心设计(已定) + +**两个类,放在 `source/source_base/module_parallel/` 下:** + +### 1. `ParaWorld` —— 单个通信域 +- 内容:`tag`(字符串常量)+ `comm`(MPI_Comm,串行下不存在)+ `rank` + `size` +- 把原本散在 GlobalV 的并行参数(NPROC_IN_POOL、RANK_IN_POOL 等)收进对应域对象 +- 方法:`tag()` / `rank()` / `size()` / `comm()`(仅 __MPI)/ `valid()` / `static serial(tag)` 安全退化(size=1, rank=0) +- 串行编译:`comm()` 用 `#ifdef __MPI` 包住,`rank()`/`size()`/`tag()` 总可用 +- 域特有参数(如 npw_per_proc、2D 网格行列)**不放进类**,由函数按需另外传 + +### 2. `ParaCollection` —— 全域容器 +- 内容:`std::vector` +- 查找:`find(tag)` 按字符串 tag 线性查找,**找不到返回静态空域(安全退化,不抛异常)** +- tag 用常量(避免裸字符串拼写错误运行时才暴露) + +### 3. `ParaTag` —— 域标签常量 +- 8 大域:`pw` / `kmesh` / `bsame_kdiff` / `bdiff_ksame` / `rgrid` / `diag` / `matrix` / `atom` +- 对应原全局:POOL_WORLD / KP_WORLD / INT_BGROUP / BP_WORLD / GRID_WORLD / DIAG_WORLD / matrix / atom + +## 目标 + +- 函数通过**注入** `const ParaWorld&` 或 `const ParaCollection&` 获取通信域,不再读裸全局 POOL_WORLD/GlobalV +- wrapper(Parallel_Common::bcast_* / Parallel_Reduce::reduce_*)加 `ParaWorld` 重载,`#ifdef __MPI` 收进 wrapper 内部,调用点无 `#ifdef`、无 MPI_Comm,串行并行都能编译跑 +- 测试用 `ParaWorld::serial(tag)` 或一行工厂构造,**去掉 GlobalV/divide_pools/set_global_partition 样板** + +## 分步计划(每步一个 commit,确认合理再进下一步) + +| 步骤 | 内容 | commit 信息 | +|---|---|---| +| **step 1** | 建 `module_parallel/` 目录 + `ParaWorld` 类(tag 常量 + comm/rank/size + `serial()` + `valid()`),含单元测试 + CMake/Makefile.Objects 接线 | `feat(parallel): add ParaWorld comm-domain value type` | +| **step 2** | `ParaCollection`(`vector` + `find(tag)` 安全退化返回静态空域),含单元测试 | `feat(parallel): add ParaCollection domain container` | +| **step 3** | 用 `ParaWorld`/`ParaCollection` 表达 8 大域装配(替代旧 divide_pools 全局写法),接进 driver 初始化 | `feat(parallel): assemble domains into ParaCollection at driver` | +| **step 4** | `Parallel_Common::bcast_bool` 加 `ParaWorld` 重载(`#ifdef __MPI` 收进 wrapper,串行 no-op,旧签名保留) | `feat(parallel): bcast_bool overload taking ParaWorld` | +| **step 5** | rhog_io.cpp 打样:注入 `ParaWorld`,`bcast_bool(error, pw)` 一行无 `#ifdef`;read_rhog_test 改一行构造去 GlobalV | `refactor(io): inject ParaWorld into read_rhog` | + +## 命名与规范约束 + +- 文件名小写+下划线:`para_world.h/.cpp`、`para_collection.h/.cpp` +- C++11,4 空格缩进,大括号独占一行,不用 `using namespace std`,注释用英文 doxygen 格式 +- 不加默认参数,不用全局变量(ParaCollection 通过注入传递,不做全局单例) +- include guard 用短名,与同目录其它文件一致 +- 不用 goto,不用宏做域替换,struct 不裸露公有成员 +- 函数参数带校验(指针非空、int 范围合理) + +## 验证方式 + +- 编译目录:`/home/510Group/6_abacus_mc/abacus-mc/build_max_para_test`,命令 `make -j 30` +- 测试:`OMP_NUM_THREADS=1 ctest -V -R ` +- 注意:沙箱内 MPI 测试会因 `/dev/nvidiactl` 受限误报崩溃,需看 ctest 日志实际结果 +- 每步 commit 前确认编译 0 错误 + 相关测试通过 + +## 待确认细节(开工前) + +1. 文件路径 `source/source_base/module_parallel/para_world.h/.cpp` 是否 OK +2. `ParaWorld` 串行下 `comm()` 不存在(`#ifdef __MPI`),`rank()`/`size()` 返回 0/1,`tag()` 总可用——是否 OK +3. 从 step 1 开始,还是想先调整步骤划分 diff --git a/cmake/Testing.cmake b/cmake/Testing.cmake index fa120c1fa69..a68871f4b47 100644 --- a/cmake/Testing.cmake +++ b/cmake/Testing.cmake @@ -32,9 +32,17 @@ endif() function(AddTest) # function for UT cmake_parse_arguments(UT "DYN" "TARGET" - "LIBS;DYN_LIBS;STATIC_LIBS;SOURCES;DEPENDS" ${ARGN}) + "LIBS;DYN_LIBS;STATIC_LIBS;SOURCES;DEPENDS;KEEP_FEATURE_DEFINITIONS" ${ARGN}) add_executable(${UT_TARGET} ${UT_SOURCES}) + # Let this target keep feature definitions (e.g. __MPI) that its source + # directory disables via abacus_disable_feature_definitions(). Needed by + # tests that genuinely exercise the feature. + if(UT_KEEP_FEATURE_DEFINITIONS) + set_property(TARGET ${UT_TARGET} PROPERTY + ABACUS_KEPT_FEATURE_DEFINITIONS ${UT_KEEP_FEATURE_DEFINITIONS}) + endif() + if(ENABLE_COVERAGE) add_coverage(${UT_TARGET}) endif() diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index bad5987651d..4dc0439a69d 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -417,6 +417,15 @@ define_property( BRIEF_DOCS "Additional ABACUS feature definitions for targets in this directory" FULL_DOCS "Additional feature definitions for targets created in this directory.") +define_property( + TARGET + PROPERTY ABACUS_KEPT_FEATURE_DEFINITIONS + BRIEF_DOCS "Feature definitions this target keeps despite a directory-level disable" + FULL_DOCS "Feature definitions that must not be stripped from this target even " + "when its source directory disables them via " + "abacus_disable_feature_definitions(). Used by tests that genuinely need a " + "feature (e.g. __MPI) inside a directory that otherwise disables it.") + function(abacus_disable_feature_definitions) abacus_normalize_definitions(_defs ${ARGN}) set_property(DIRECTORY APPEND PROPERTY @@ -448,7 +457,12 @@ function(abacus_apply_build_options target) set(_defs "${_abacus_feature_definitions}") get_property(_disabled DIRECTORY "${_source_dir}" PROPERTY ABACUS_DISABLED_FEATURE_DEFINITIONS) get_property(_local DIRECTORY "${_source_dir}" PROPERTY ABACUS_LOCAL_FEATURE_DEFINITIONS) + get_target_property(_kept "${target}" ABACUS_KEPT_FEATURE_DEFINITIONS) + if(_kept) + # A target may opt back into definitions its directory disables. + list(REMOVE_ITEM _disabled ${_kept}) + endif() if(_disabled) # Filter after conditional definitions have been evaluated. string(JOIN "|" _disabled_regex ${_disabled}) diff --git a/source/Makefile.Objects b/source/Makefile.Objects index 7104d8a62e6..e770fbb7d29 100644 --- a/source/Makefile.Objects +++ b/source/Makefile.Objects @@ -22,6 +22,7 @@ VPATH=./src_global:\ ./source_base:\ ./source_base/kernels:\ ./source_base/module_external:\ +./source_base/module_parallel:\ ./source_base/module_container/base/core:\ ./source_base/module_container/ATen/core:\ ./source_base/module_container/ATen/kernels:\ @@ -625,7 +626,6 @@ OBJS_IO=module_parameter/input_conv.o\ output.o\ module_output/print_info.o\ module_output/read_cube.o\ - module_chgpot/rhog_io.o\ module_wf/read_wfc_pw.o\ module_wf/read_wf2rho_pw.o\ module_restart/restart.o\ @@ -647,10 +647,8 @@ OBJS_IO=module_parameter/input_conv.o\ module_output/write_pao.o\ module_wf/write_wfc_pw.o\ module_output/write_cube.o\ - module_chgpot/write_elecstat_pot.o\ module_elf/write_elf.o\ module_dipole/write_dipole.o\ - module_chgpot/write_init.o\ module_current/td_current_io.o\ module_current/td_current_io_comm.o\ td_efield_io.o\ @@ -791,7 +789,18 @@ OBJS_PARALLEL=parallel_common.o\ parallel_grid.o\ parallel_kpoints.o\ parallel_reduce.o\ - parallel_device.o + parallel_device.o\ + para_world.o\ + para_collection.o\ + para_kmesh_world.o\ + para_pw_world.o\ + para_diag_world.o\ + para_rgrid_world.o\ + para_bgroup_world.o\ + para_matrix_world.o\ + para_mpi_func.o\ + para_setup.o\ + para_bridge.o OBJS_SRCPW=h_ewald_pw.o\ dnrm2.o\ @@ -815,6 +824,9 @@ OBJS_SRCPW=h_ewald_pw.o\ mix_precond.o\ charge_mixing_rho.o\ charge_mixing_uspp.o\ + rhog_io.o\ + write_elecstat_pot.o\ + write_init.o\ fp_energy.o\ setup_pot.o\ setup_pwrho.o\ diff --git a/source/source_base/CMakeLists.txt b/source/source_base/CMakeLists.txt index cfdd25bd7ec..e2ccaaf1732 100644 --- a/source/source_base/CMakeLists.txt +++ b/source/source_base/CMakeLists.txt @@ -73,6 +73,17 @@ add_library( module_mixing/plain_mixing.cpp module_mixing/pulay_mixing.cpp module_mixing/broyden_mixing.cpp + module_parallel/para_world.cpp + module_parallel/para_collection.cpp + module_parallel/para_kmesh_world.cpp + module_parallel/para_pw_world.cpp + module_parallel/para_diag_world.cpp + module_parallel/para_rgrid_world.cpp + module_parallel/para_bgroup_world.cpp + module_parallel/para_matrix_world.cpp + module_parallel/para_mpi_func.cpp + module_parallel/para_setup.cpp + module_parallel/para_bridge.cpp ${LIBM_SRC} ) @@ -95,6 +106,7 @@ if(BUILD_TESTING) add_subdirectory(module_mixing/test) add_subdirectory(module_device/test) add_subdirectory(module_grid/test) + add_subdirectory(module_parallel/test) if (ENABLE_ABACUS_LIBM) add_subdirectory(libm/test) endif() diff --git a/source/source_base/module_parallel/para_bgroup_world.cpp b/source/source_base/module_parallel/para_bgroup_world.cpp new file mode 100644 index 00000000000..83493253757 --- /dev/null +++ b/source/source_base/module_parallel/para_bgroup_world.cpp @@ -0,0 +1,22 @@ +#include "para_bgroup_world.h" + +namespace Parallel +{ + +ParaBgroupWorld::ParaBgroupWorld() + : ParaWorld("bdiff_ksame"), my_bndgroup_(0), nbndgroup_(1) +{ +} + +#ifdef __MPI +ParaBgroupWorld::ParaBgroupWorld(const MPI_Comm& intra_comm, const MPI_Comm& inter_comm, int nbndgroup) + : ParaWorld("bdiff_ksame", intra_comm), inter_comm_(inter_comm), nbndgroup_(nbndgroup) +{ + if (inter_comm != MPI_COMM_NULL) + { + MPI_Comm_rank(inter_comm, &my_bndgroup_); + } +} +#endif + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_bgroup_world.h b/source/source_base/module_parallel/para_bgroup_world.h new file mode 100644 index 00000000000..e2d82cc99b9 --- /dev/null +++ b/source/source_base/module_parallel/para_bgroup_world.h @@ -0,0 +1,67 @@ +#ifndef PARA_BGROUP_WORLD_H +#define PARA_BGROUP_WORLD_H + +#include "para_world.h" + +namespace Parallel +{ + +/** + * @brief bgroup parallel domain: band group communication topology. + * + * Self-contained replacement for INT_BGROUP + BP_WORLD + + * GlobalV::MY_BNDGROUP/NPROC_IN_BNDGROUP/RANK_IN_BPGROUP. + * + * The band group domain has two communicators: + * - intra: INT_BGROUP (same band group, different k/pw) + * - inter: BP_WORLD (different band groups, same k) + * + * Tests only need this header. + */ +class ParaBgroupWorld : public ParaWorld +{ +public: + /** + * @brief Construct a serial bgroup domain (single band group). + */ + ParaBgroupWorld(); + +#ifdef __MPI + /** + * @brief Construct a bgroup domain from intra and inter communicators. + * + * @param[in] intra_comm intra-group communicator (e.g. INT_BGROUP) + * @param[in] inter_comm inter-group communicator (e.g. BP_WORLD) + * @param[in] nbndgroup number of band groups + */ + ParaBgroupWorld(const MPI_Comm& intra_comm, const MPI_Comm& inter_comm, int nbndgroup); +#endif + + /// Band group index of this process. + int my_bndgroup() const { return my_bndgroup_; } + + /// Number of band groups. + int nbndgroup() const { return nbndgroup_; } + + /// Rank within the band group (alias for rank()). + int rank_in_bpgroup() const { return rank(); } + + /// Number of processes in the band group (alias for size()). + int nproc_in_bndgroup() const { return size(); } + +#ifdef __MPI + /// Inter-group communicator (BP_WORLD equivalent). + MPI_Comm inter_comm() const { return inter_comm_; } +#endif + +private: + int my_bndgroup_ = 0; + int nbndgroup_ = 1; +#ifdef __MPI + MPI_Comm inter_comm_ = MPI_COMM_NULL; +#endif +}; + +} // namespace Parallel + +#endif // PARA_BGROUP_WORLD_H diff --git a/source/source_base/module_parallel/para_bridge.cpp b/source/source_base/module_parallel/para_bridge.cpp new file mode 100644 index 00000000000..c2c009b1580 --- /dev/null +++ b/source/source_base/module_parallel/para_bridge.cpp @@ -0,0 +1,22 @@ +#include "para_bridge.h" +#include "para_tag.h" + +#ifdef __MPI +#include "source_base/parallel_comm.h" +#endif + +namespace Parallel +{ + +// Temporary bridge: construct a pw-domain ParaWorld from the old globals. +// Delete this file once ParaCollection is wired into driver initialization. +ParaWorld make_pw_world() +{ +#ifdef __MPI + return ParaWorld::make_mpi(ParaTag::pw, POOL_WORLD); +#else + return ParaWorld::serial(ParaTag::pw); +#endif +} + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_bridge.h b/source/source_base/module_parallel/para_bridge.h new file mode 100644 index 00000000000..c0df2a61946 --- /dev/null +++ b/source/source_base/module_parallel/para_bridge.h @@ -0,0 +1,21 @@ +#ifndef PARA_BRIDGE_H +#define PARA_BRIDGE_H + +#include "para_world.h" + +namespace Parallel +{ + +/** + * @brief Temporary bridge: construct a pw-domain ParaWorld from the old + * global POOL_WORLD (MPI) or as a serial domain (non-MPI). + * + * Hides the #ifdef __MPI from call sites so they stay one-liner. Delete + * this function (and this file) once ParaCollection is wired into driver + * initialization and callers receive a ParaWorld& from above. + */ +ParaWorld make_pw_world(); + +} // namespace Parallel + +#endif // PARA_BRIDGE_H diff --git a/source/source_base/module_parallel/para_collection.cpp b/source/source_base/module_parallel/para_collection.cpp new file mode 100644 index 00000000000..a09f5879e7e --- /dev/null +++ b/source/source_base/module_parallel/para_collection.cpp @@ -0,0 +1,31 @@ +#include "para_collection.h" + +namespace Parallel +{ + +void ParaCollection::add(std::unique_ptr world) +{ + for (const auto& existing : worlds_) + { + if (existing->tag() == world->tag()) + { + return; + } + } + worlds_.push_back(std::move(world)); +} + +const ParaWorld& ParaCollection::find(const std::string& tag) const +{ + for (const auto& world : worlds_) + { + if (world->tag() == tag) + { + return *world; + } + } + static const ParaWorld empty = ParaWorld::serial(""); + return empty; +} + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_collection.h b/source/source_base/module_parallel/para_collection.h new file mode 100644 index 00000000000..c8503cc0e3e --- /dev/null +++ b/source/source_base/module_parallel/para_collection.h @@ -0,0 +1,83 @@ +#ifndef PARA_COLLECTION_H +#define PARA_COLLECTION_H + +#include +#include +#include + +#include "para_world.h" + +namespace Parallel +{ + +/** + * @brief Container for all parallel communication domains. + * + * A ParaCollection owns a set of ParaWorld objects (base class pointers), + * each describing one communication domain (see ParaTag). Callers look up + * domains by tag via find(); a missing tag yields a static empty (invalid) + * domain as a safe degradation, never an exception. + * + * The collection is passed explicitly to functions that need communicator + * access, replacing reads of loose globals such as GlobalV::POOL_WORLD. + */ +class ParaCollection +{ +public: + ParaCollection() = default; + + /** + * @brief Append a domain to the collection. + * + * Duplicate tags are rejected (the existing entry is kept). + * + * @param[in] world domain to add (ownership transferred) + */ + void add(std::unique_ptr world); + + /** + * @brief Look up a domain by tag. + * + * @param[in] tag domain tag string + * @return the matching ParaWorld, or a static empty domain if not found + */ + const ParaWorld& find(const std::string& tag) const; + + /** + * @brief Look up a domain by tag and cast to the requested subclass. + * + * @tparam T expected subclass (e.g. ParaKmeshWorld) + * @param[in] tag domain tag string + * @return pointer to the domain if found and type matches, nullptr otherwise + */ + template + const T* find_as(const std::string& tag) const; + + /** + * @brief Number of domains in the collection. + */ + size_t size() const + { + return worlds_.size(); + } + +private: + std::vector> worlds_; ///< owned domains +}; + +template +const T* ParaCollection::find_as(const std::string& tag) const +{ + for (const auto& world : worlds_) + { + if (world->tag() == tag) + { + return dynamic_cast(world.get()); + } + } + return nullptr; +} + +} // namespace Parallel + +#endif // PARA_COLLECTION_H diff --git a/source/source_base/module_parallel/para_diag_world.cpp b/source/source_base/module_parallel/para_diag_world.cpp new file mode 100644 index 00000000000..e633a3da7a0 --- /dev/null +++ b/source/source_base/module_parallel/para_diag_world.cpp @@ -0,0 +1,18 @@ +#include "para_diag_world.h" + +namespace Parallel +{ + +ParaDiagWorld::ParaDiagWorld() + : ParaWorld("diag"), dcolor_(0) +{ +} + +#ifdef __MPI +ParaDiagWorld::ParaDiagWorld(const MPI_Comm& comm, int dcolor) + : ParaWorld("diag", comm), dcolor_(dcolor) +{ +} +#endif + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_diag_world.h b/source/source_base/module_parallel/para_diag_world.h new file mode 100644 index 00000000000..c948da8b8b4 --- /dev/null +++ b/source/source_base/module_parallel/para_diag_world.h @@ -0,0 +1,51 @@ +#ifndef PARA_DIAG_WORLD_H +#define PARA_DIAG_WORLD_H + +#include "para_world.h" + +namespace Parallel +{ + +/** + * @brief diag parallel domain: diagonalization group topology. + * + * Self-contained replacement for DIAG_WORLD + GlobalV::DRANK/DSIZE/DCOLOR. + * The diag domain is created by splitting MPI_COMM_WORLD into groups + * for parallel diagonalization (PEXSI, ScaLAPACK). + * + * Tests only need this header; no parallel_comm.h or parallel_global.h. + */ +class ParaDiagWorld : public ParaWorld +{ +public: + /** + * @brief Construct a serial diag domain (single-process group). + */ + ParaDiagWorld(); + +#ifdef __MPI + /** + * @brief Construct a diag domain from an existing communicator. + * + * @param[in] comm diag communicator (e.g. DIAG_WORLD) + * @param[in] dcolor color used in MPI_Comm_split to create this group + */ + ParaDiagWorld(const MPI_Comm& comm, int dcolor); +#endif + + /// Color used in MPI_Comm_split to create this diag group. + int dcolor() const { return dcolor_; } + + /// Rank within the diag group (alias for rank()). + int drank() const { return rank(); } + + /// Number of processes in the diag group (alias for size()). + int dsize() const { return size(); } + +private: + int dcolor_ = 0; +}; + +} // namespace Parallel + +#endif // PARA_DIAG_WORLD_H diff --git a/source/source_base/module_parallel/para_kmesh_world.cpp b/source/source_base/module_parallel/para_kmesh_world.cpp new file mode 100644 index 00000000000..319654df6a3 --- /dev/null +++ b/source/source_base/module_parallel/para_kmesh_world.cpp @@ -0,0 +1,198 @@ +#include "para_kmesh_world.h" + +#include +#include + +namespace Parallel +{ + +ParaKmeshWorld::ParaKmeshWorld(int nkstot, int nspin) + : ParaWorld("kmesh"), kpar_(1), my_pool_(0), rank_in_pool_(0), + nproc_(1), nspin_(nspin), nkstot_(nkstot) +{ + distribute_kpoints(); + nks_local_ = nkstot_; + startk_global_ = 0; +} + +#ifdef __MPI +ParaKmeshWorld::ParaKmeshWorld(const MPI_Comm& comm, int kpar, int my_pool, int nproc, int nkstot, int nspin) + : ParaWorld("kmesh", comm), kpar_(kpar), my_pool_(my_pool), + rank_in_pool_(rank()), nproc_(nproc), nspin_(nspin), nkstot_(nkstot) +{ + distribute_kpoints(); + nks_local_ = nks_pool_[my_pool_]; + startk_global_ = startk_pool_[my_pool_]; +} +#endif + +void ParaKmeshWorld::distribute_kpoints() +{ + // k-points per pool (evenly divided, remainder to front) + nks_pool_.resize(kpar_, 0); + const int nks_ave = nkstot_ / kpar_; + const int nks_rem = nkstot_ % kpar_; + for (int i = 0; i < kpar_; ++i) + { + nks_pool_[i] = nks_ave + (i < nks_rem ? 1 : 0); + } + + // global start index per pool + startk_pool_.resize(kpar_, 0); + for (int i = 1; i < kpar_; ++i) + { + startk_pool_[i] = startk_pool_[i - 1] + nks_pool_[i - 1]; + } + + // pool index per k-point + whichpool_.resize(nkstot_, 0); + for (int p = 0; p < kpar_; ++p) + { + for (int ik = 0; ik < nks_pool_[p]; ++ik) + { + whichpool_[startk_pool_[p] + ik] = p; + } + } + + // first world rank per pool + startpro_pool_.resize(kpar_, 0); + const int nproc_ave = nproc_ / kpar_; + const int nproc_rem = nproc_ % kpar_; + for (int i = 1; i < kpar_; ++i) + { + startpro_pool_[i] = startpro_pool_[i - 1] + nproc_ave + (i - 1 < nproc_rem ? 1 : 0); + } +} + +int ParaKmeshWorld::nks_pool(int pool) const +{ + assert(pool >= 0 && pool < kpar_); + return nks_pool_[pool]; +} + +int ParaKmeshWorld::startk_pool(int pool) const +{ + assert(pool >= 0 && pool < kpar_); + return startk_pool_[pool]; +} + +int ParaKmeshWorld::which_pool(int ik_global) const +{ + assert(ik_global >= 0 && ik_global < nkstot_); + return whichpool_[ik_global]; +} + +int ParaKmeshWorld::startpro_pool(int pool) const +{ + assert(pool >= 0 && pool < kpar_); + return startpro_pool_[pool]; +} + +int ParaKmeshWorld::max_nks_pool() const +{ + return *std::max_element(nks_pool_.begin(), nks_pool_.end()); +} + +void ParaKmeshWorld::pool_collection(double& value, const double* wk, int ik) const +{ +#ifdef __MPI + const int ik_local = ik - startk_pool_[my_pool_]; + const int pool = whichpool_[ik]; + + if (rank_in_pool_ == 0) + { + if (my_pool_ == 0) + { + if (pool == 0) + { + value = wk[ik_local]; + } + else + { + MPI_Status status; + MPI_Recv(&value, 1, MPI_DOUBLE, startpro_pool_[pool], ik, MPI_COMM_WORLD, &status); + } + } + else + { + if (my_pool_ == pool) + { + MPI_Send(&wk[ik_local], 1, MPI_DOUBLE, 0, ik, MPI_COMM_WORLD); + } + } + } + MPI_Barrier(MPI_COMM_WORLD); +#else + value = wk[ik]; +#endif +} + +template +void ParaKmeshWorld::pool_collection(T* value, const T* w, int dim, int ik) const +{ +#ifdef __MPI + const int ik_local = ik - startk_pool_[my_pool_]; + const int begin = ik_local * dim; + const T* src = &w[begin]; + + // nspin==2 restricts to pool 0 (legacy behavior from Parallel_Kpoints) + const int pool = (nspin_ == 2) ? 0 : whichpool_[ik]; + + if (rank_in_pool_ == 0) + { + if (my_pool_ == 0) + { + if (pool == 0) + { + std::copy(src, src + dim, value); + } + else + { + MPI_Status status; + MPI_Recv(value, dim * sizeof(T), MPI_BYTE, startpro_pool_[pool], ik * 2, MPI_COMM_WORLD, &status); + } + } + else + { + if (my_pool_ == pool) + { + MPI_Send(src, dim * sizeof(T), MPI_BYTE, 0, ik * 2, MPI_COMM_WORLD); + } + } + } + MPI_Barrier(MPI_COMM_WORLD); +#else + const int begin = ik * dim; + std::copy(&w[begin], &w[begin] + dim, value); +#endif +} + +void ParaKmeshWorld::gather_kvec(const std::vector& vec_local, std::vector& vec_global) const +{ +#ifdef __MPI + int world_rank = 0; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + const bool is_pool_root = (world_rank == startpro_pool_[my_pool_]); + + vec_global.resize(nkstot_ * 3, 0.0); + if (is_pool_root) + { + for (int i = 0; i < nks_local_; ++i) + { + const int gk = startk_global_ + i; + vec_global[gk * 3 + 0] = vec_local[i * 3 + 0]; + vec_global[gk * 3 + 1] = vec_local[i * 3 + 1]; + vec_global[gk * 3 + 2] = vec_local[i * 3 + 2]; + } + } + MPI_Allreduce(MPI_IN_PLACE, vec_global.data(), nkstot_ * 3, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); +#else + vec_global = vec_local; +#endif +} + +// explicit instantiation +template void ParaKmeshWorld::pool_collection(double*, const double*, int, int) const; +template void ParaKmeshWorld::pool_collection>(std::complex*, const std::complex*, int, int) const; + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_kmesh_world.h b/source/source_base/module_parallel/para_kmesh_world.h new file mode 100644 index 00000000000..f78607112d7 --- /dev/null +++ b/source/source_base/module_parallel/para_kmesh_world.h @@ -0,0 +1,140 @@ +#ifndef PARA_KMESH_WORLD_H +#define PARA_KMESH_WORLD_H + +#include +#include + +#include "para_world.h" + +namespace Parallel +{ + +/** + * @brief k-mesh parallel domain: k-point distribution across pools. + * + * Self-contained replacement for Parallel_Kpoints + KP_WORLD + + * GlobalV::KPAR / MY_POOL / RANK_IN_POOL. Owns all k-point pool + * topology data and provides query / collection operations. + * + * In serial builds all operations degenerate to single-pool behavior. + * Tests only need this header; no GlobalV, no parallel_comm.h. + */ +class ParaKmeshWorld : public ParaWorld +{ +public: + /** + * @brief Construct a serial (single-pool) k-mesh domain. + * + * @param[in] nkstot total number of k-points (without spin) + * @param[in] nspin number of spin components + */ + ParaKmeshWorld(int nkstot, int nspin); + +#ifdef __MPI + /** + * @brief Construct a k-mesh domain on an existing communicator. + * + * @param[in] comm k-point pool communicator (e.g. KP_WORLD) + * @param[in] kpar number of pools + * @param[in] my_pool pool index of this process + * @param[in] nproc total number of processes (MPI_COMM_WORLD size) + * @param[in] nkstot total number of k-points (without spin) + * @param[in] nspin number of spin components + */ + ParaKmeshWorld(const MPI_Comm& comm, int kpar, int my_pool, int nproc, int nkstot, int nspin); +#endif + + /// Number of pools. + int kpar() const { return kpar_; } + + /// Pool index of this process. + int my_pool() const { return my_pool_; } + + /// Rank within the pool. + int rank_in_pool() const { return rank_in_pool_; } + + /// Total number of processes. + int nproc() const { return nproc_; } + + /// Number of spin components. + int nspin() const { return nspin_; } + + /// Total number of k-points (without spin). + int nkstot() const { return nkstot_; } + + /// Number of k-points in this pool. + int nks_local() const { return nks_local_; } + + /// Global start index of this pool's k-points. + int startk_global() const { return startk_global_; } + + /// Number of k-points in the given pool. + int nks_pool(int pool) const; + + /// Global start index of the given pool's k-points. + int startk_pool(int pool) const; + + /// Which pool owns the given global k-point index. + int which_pool(int ik_global) const; + + /// First MPI_COMM_WORLD rank of the given pool. + int startpro_pool(int pool) const; + + /// Maximum number of k-points across all pools. + int max_nks_pool() const; + + /** + * @brief Collect a scalar value from the pool that owns k-point ik. + * + * Pool 0 receives the value; other pools send. Only rank_in_pool==0 + * participates in the actual communication. + * + * @param[out] value collected value (valid on pool 0 root) + * @param[in] wk local k-point weights array + * @param[in] ik global k-point index + */ + void pool_collection(double& value, const double* wk, int ik) const; + + /** + * @brief Collect an array slice from the pool that owns k-point ik. + * + * @param[out] value output array (dim elements) + * @param[in] w input array (nkstot * dim elements, row-major by k) + * @param[in] dim number of elements per k-point + * @param[in] ik global k-point index + */ + template + void pool_collection(T* value, const T* w, int dim, int ik) const; + + /** + * @brief Gather local k-point vectors to global array. + * + * Only pool-root processes contribute their local k-points; + * the result is valid on all ranks after MPI_Allreduce. + * + * @param[in] vec_local local k-point vectors (nks_local elements) + * @param[out] vec_global global k-point vectors (nkstot elements) + */ + void gather_kvec(const std::vector& vec_local, std::vector& vec_global) const; + +private: + void distribute_kpoints(); + + int kpar_ = 1; + int my_pool_ = 0; + int rank_in_pool_ = 0; + int nproc_ = 1; + int nspin_ = 1; + int nkstot_ = 0; + int nks_local_ = 0; + int startk_global_ = 0; + + std::vector nks_pool_; ///< k-points per pool + std::vector startk_pool_; ///< global start index per pool + std::vector whichpool_; ///< pool index per k-point + std::vector startpro_pool_; ///< first world rank per pool +}; + +} // namespace Parallel + +#endif // PARA_KMESH_WORLD_H diff --git a/source/source_base/module_parallel/para_matrix_world.cpp b/source/source_base/module_parallel/para_matrix_world.cpp new file mode 100644 index 00000000000..629bc1af498 --- /dev/null +++ b/source/source_base/module_parallel/para_matrix_world.cpp @@ -0,0 +1,32 @@ +#include "para_matrix_world.h" + +namespace Parallel +{ + +ParaMatrixWorld::ParaMatrixWorld() + : ParaWorld("matrix") +{ + compute_proc_grid(); +} + +#ifdef __MPI +ParaMatrixWorld::ParaMatrixWorld(const MPI_Comm& comm) + : ParaWorld("matrix", comm) +{ + compute_proc_grid(); +} +#endif + +void ParaMatrixWorld::compute_proc_grid() +{ + const int np = size(); + dim0_ = np; + while (dim1_ = np / dim0_, dim0_ * dim1_ != np) + { + --dim0_; + } + coord_row_ = rank() / dim1_; + coord_col_ = rank() % dim1_; +} + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_matrix_world.h b/source/source_base/module_parallel/para_matrix_world.h new file mode 100644 index 00000000000..03921e6c0a7 --- /dev/null +++ b/source/source_base/module_parallel/para_matrix_world.h @@ -0,0 +1,66 @@ +#ifndef PARA_MATRIX_WORLD_H +#define PARA_MATRIX_WORLD_H + +#include +#include + +#include "para_world.h" + +namespace Parallel +{ + +/** + * @brief matrix parallel domain: 2D block-cyclic distribution. + * + * Self-contained wrapper for the matrix-level parallel topology. + * Replaces the process-grid part of Parallel_2D (dim0/dim1/coord) + * with a cleaner interface. The actual ScaLAPACK descriptor and + * BLACS context management stay in Parallel_2D; this class only + * holds the process grid dimensions and coordinates. + * + * Tests only need this header. + */ +class ParaMatrixWorld : public ParaWorld +{ +public: + /** + * @brief Construct a serial matrix domain (1x1 process grid). + */ + ParaMatrixWorld(); + +#ifdef __MPI + /** + * @brief Construct a matrix domain on an existing communicator. + * + * The process grid is computed automatically: dim0 = largest divisor + * of nproc with dim0 >= dim1 (square-ish), dim1 = nproc / dim0. + * + * @param[in] comm matrix communicator (e.g. DIAG_WORLD or MPI_COMM_WORLD) + */ + ParaMatrixWorld(const MPI_Comm& comm); +#endif + + /// Process grid row count. + int dim0() const { return dim0_; } + + /// Process grid column count. + int dim1() const { return dim1_; } + + /// Row coordinate of this process in the grid. + int coord_row() const { return coord_row_; } + + /// Column coordinate of this process in the grid. + int coord_col() const { return coord_col_; } + +private: + int dim0_ = 1; + int dim1_ = 1; + int coord_row_ = 0; + int coord_col_ = 0; + + void compute_proc_grid(); +}; + +} // namespace Parallel + +#endif // PARA_MATRIX_WORLD_H diff --git a/source/source_base/module_parallel/para_mpi_func.cpp b/source/source_base/module_parallel/para_mpi_func.cpp new file mode 100644 index 00000000000..897bfdf0c64 --- /dev/null +++ b/source/source_base/module_parallel/para_mpi_func.cpp @@ -0,0 +1,191 @@ +#include "para_mpi_func.h" + +#include + +namespace Parallel +{ + +#ifdef __MPI +namespace { +inline MPI_Datatype mpi_type(int*) { return MPI_INT; } +inline MPI_Datatype mpi_type(double*) { return MPI_DOUBLE; } +inline MPI_Datatype mpi_type(std::complex*) { return MPI_DOUBLE; } // 2 doubles +} +#endif + +// ========== Broadcast ========== + +void bcast_bool(bool& v, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + int tmp = v ? 1 : 0; + MPI_Bcast(&tmp, 1, MPI_INT, root, world.comm()); + v = (tmp != 0); +#endif +} + +void bcast_int(int& v, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Bcast(&v, 1, MPI_INT, root, world.comm()); +#endif +} + +void bcast_double(double& v, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Bcast(&v, 1, MPI_DOUBLE, root, world.comm()); +#endif +} + +void bcast_complex(std::complex& v, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Bcast(&v, 2, MPI_DOUBLE, root, world.comm()); +#endif +} + +void bcast_string(std::string& s, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + int len = static_cast(s.size()); + MPI_Bcast(&len, 1, MPI_INT, root, world.comm()); + if (world.rank() != root) s.resize(len); + if (len > 0) + { + MPI_Bcast(&s[0], len, MPI_CHAR, root, world.comm()); + } +#endif +} + +void bcast_int(int* v, int n, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Bcast(v, n, MPI_INT, root, world.comm()); +#endif +} + +void bcast_double(double* v, int n, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Bcast(v, n, MPI_DOUBLE, root, world.comm()); +#endif +} + +void bcast_complex(std::complex* v, int n, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Bcast(v, 2 * n, MPI_DOUBLE, root, world.comm()); +#endif +} + +void bcast_char(char* v, int n, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Bcast(v, n, MPI_CHAR, root, world.comm()); +#endif +} + +void bcast_string(std::string* v, int n, const ParaWorld& world, int root) +{ +#ifdef __MPI + if (!world.valid()) return; + for (int i = 0; i < n; ++i) + { + bcast_string(v[i], world, root); + } +#endif +} + +// ========== Reduce ========== + +void reduce_all(double& v, const ParaWorld& world) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Allreduce(MPI_IN_PLACE, &v, 1, MPI_DOUBLE, MPI_SUM, world.comm()); +#endif +} + +void reduce_all(int& v, const ParaWorld& world) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Allreduce(MPI_IN_PLACE, &v, 1, MPI_INT, MPI_SUM, world.comm()); +#endif +} + +void reduce_all(double* v, int n, const ParaWorld& world) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Allreduce(MPI_IN_PLACE, v, n, MPI_DOUBLE, MPI_SUM, world.comm()); +#endif +} + +// ========== Min/Max ========== + +void reduce_min(double& v, const ParaWorld& world) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Allreduce(MPI_IN_PLACE, &v, 1, MPI_DOUBLE, MPI_MIN, world.comm()); +#endif +} + +void reduce_max(double& v, const ParaWorld& world) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Allreduce(MPI_IN_PLACE, &v, 1, MPI_DOUBLE, MPI_MAX, world.comm()); +#endif +} + +void reduce_min(int& v, const ParaWorld& world) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Allreduce(MPI_IN_PLACE, &v, 1, MPI_INT, MPI_MIN, world.comm()); +#endif +} + +void reduce_max(int& v, const ParaWorld& world) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Allreduce(MPI_IN_PLACE, &v, 1, MPI_INT, MPI_MAX, world.comm()); +#endif +} + +// ========== Barrier ========== + +void barrier(const ParaWorld& world) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Barrier(world.comm()); +#endif +} + +// ========== Gather ========== + +void gather_int(int& v, int* all, const ParaWorld& world) +{ +#ifdef __MPI + if (!world.valid()) return; + MPI_Allgather(&v, 1, MPI_INT, all, 1, MPI_INT, world.comm()); +#else + all[0] = v; +#endif +} + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_mpi_func.h b/source/source_base/module_parallel/para_mpi_func.h new file mode 100644 index 00000000000..203fd52af60 --- /dev/null +++ b/source/source_base/module_parallel/para_mpi_func.h @@ -0,0 +1,55 @@ +#ifndef PARA_MPI_FUNC_H +#define PARA_MPI_FUNC_H + +#include +#include + +#include "para_world.h" + +namespace Parallel +{ + +// Domain-aware MPI communication functions. +// Each function takes the target communication domain (const ParaWorld&) +// explicitly instead of hardcoding MPI_COMM_WORLD/POOL_WORLD. +// In serial builds all functions are no-ops (gather_int copies locally); +// invalid/empty domains are safely skipped. + +// ========== Broadcast ========== + +void bcast_bool(bool& v, const ParaWorld& world, int root = 0); +void bcast_int(int& v, const ParaWorld& world, int root = 0); +void bcast_double(double& v, const ParaWorld& world, int root = 0); +void bcast_complex(std::complex& v, const ParaWorld& world, int root = 0); +void bcast_string(std::string& s, const ParaWorld& world, int root = 0); + +void bcast_int(int* v, int n, const ParaWorld& world, int root = 0); +void bcast_double(double* v, int n, const ParaWorld& world, int root = 0); +void bcast_complex(std::complex* v, int n, const ParaWorld& world, int root = 0); +void bcast_char(char* v, int n, const ParaWorld& world, int root = 0); +void bcast_string(std::string* v, int n, const ParaWorld& world, int root = 0); + +// ========== Reduce (Allreduce, result on all ranks) ========== + +void reduce_all(double& v, const ParaWorld& world); +void reduce_all(int& v, const ParaWorld& world); +void reduce_all(double* v, int n, const ParaWorld& world); + +// ========== Reduce min/max ========== + +void reduce_min(double& v, const ParaWorld& world); +void reduce_max(double& v, const ParaWorld& world); +void reduce_min(int& v, const ParaWorld& world); +void reduce_max(int& v, const ParaWorld& world); + +// ========== Barrier ========== + +void barrier(const ParaWorld& world); + +// ========== Gather ========== + +void gather_int(int& v, int* all, const ParaWorld& world); + +} // namespace Parallel + +#endif // PARA_MPI_FUNC_H diff --git a/source/source_base/module_parallel/para_pw_world.cpp b/source/source_base/module_parallel/para_pw_world.cpp new file mode 100644 index 00000000000..16aed0c45cb --- /dev/null +++ b/source/source_base/module_parallel/para_pw_world.cpp @@ -0,0 +1,33 @@ +#include "para_pw_world.h" + +#include +#include + +namespace Parallel +{ + +ParaPwWorld::ParaPwWorld(int npw) + : ParaWorld("pw"), npw_(npw), npwtot_(npw), npw_per_(1, npw) +{ +} + +#ifdef __MPI +ParaPwWorld::ParaPwWorld(const MPI_Comm& comm, const std::vector& npw_per) + : ParaWorld("pw", comm), npw_per_(npw_per) +{ + npw_ = npw_per_[rank()]; + npwtot_ = 0; + for (int n : npw_per_) + { + npwtot_ += n; + } +} +#endif + +int ParaPwWorld::npw_per(int p) const +{ + assert(p >= 0 && p < static_cast(npw_per_.size())); + return npw_per_[p]; +} + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_pw_world.h b/source/source_base/module_parallel/para_pw_world.h new file mode 100644 index 00000000000..8c514837f3d --- /dev/null +++ b/source/source_base/module_parallel/para_pw_world.h @@ -0,0 +1,67 @@ +#ifndef PARA_PW_WORLD_H +#define PARA_PW_WORLD_H + +#include + +#include "para_world.h" + +namespace Parallel +{ + +/** + * @brief pw parallel domain: plane-wave distribution within a pool. + * + * Self-contained replacement for poolnproc/poolrank/npw_per/npwtot + * members scattered across PW_Basis and GlobalV::NPROC_IN_POOL. + * Owns the pool-level parallel topology and plane-wave count distribution. + * + * The actual FFT-based distribution algorithm (method1/method2) stays + * in PW_Basis; this class only holds the result: how many plane waves + * each process in the pool gets. + * + * Tests only need this header; no PW_Basis, no parallel_comm.h. + */ +class ParaPwWorld : public ParaWorld +{ +public: + /** + * @brief Construct a serial (single-process) pw domain. + * + * @param[in] npw number of plane waves on this process + */ + explicit ParaPwWorld(int npw); + +#ifdef __MPI + /** + * @brief Construct a pw domain on an existing pool communicator. + * + * @param[in] comm pool communicator (e.g. POOL_WORLD) + * @param[in] npw_per array of plane-wave counts per process (size = pool size) + */ + ParaPwWorld(const MPI_Comm& comm, const std::vector& npw_per); +#endif + + /// Number of plane waves on this process. + int npw() const { return npw_; } + + /// Total number of plane waves in the pool. + int npwtot() const { return npwtot_; } + + /// Number of plane waves on process p in the pool. + int npw_per(int p) const; + + /// Number of processes in the pool (same as size()). + int poolnproc() const { return size(); } + + /// Rank within the pool (same as rank()). + int poolrank() const { return rank(); } + +private: + int npw_ = 0; ///< local plane-wave count + int npwtot_ = 0; ///< total plane waves in pool + std::vector npw_per_; ///< per-process plane-wave counts +}; + +} // namespace Parallel + +#endif // PARA_PW_WORLD_H diff --git a/source/source_base/module_parallel/para_rgrid_world.cpp b/source/source_base/module_parallel/para_rgrid_world.cpp new file mode 100644 index 00000000000..961eb22d2e2 --- /dev/null +++ b/source/source_base/module_parallel/para_rgrid_world.cpp @@ -0,0 +1,217 @@ +#include "para_rgrid_world.h" + +#include + +namespace Parallel +{ + +ParaRgridWorld::ParaRgridWorld(int ncx, int ncy, int ncz) + : ParaWorld("rgrid"), ncx_(ncx), ncy_(ncy), ncz_(ncz) +{ + assert(ncx > 0 && ncy > 0 && ncz > 0); + distribute_z(); + nczp_ = numz_[0]; +} + +#ifdef __MPI +ParaRgridWorld::ParaRgridWorld(const MPI_Comm& comm, int ncx, int ncy, int ncz) + : ParaWorld("rgrid", comm), ncx_(ncx), ncy_(ncy), ncz_(ncz) +{ + assert(ncx > 0 && ncy > 0 && ncz > 0); + distribute_z(); + nczp_ = numz_[rank()]; +} +#endif + +void ParaRgridWorld::distribute_z() +{ + const int np = size(); + numz_.resize(np, 0); + startz_.resize(np, 0); + whichpro_.resize(ncz_, 0); + + // Evenly distribute z-planes, remainder to front processes + const int base = ncz_ / np; + const int rem = ncz_ % np; + int acc = 0; + for (int p = 0; p < np; ++p) + { + numz_[p] = base + (p < rem ? 1 : 0); + startz_[p] = acc; + acc += numz_[p]; + } + + // Build owner table + for (int p = 0; p < np; ++p) + { + for (int iz = 0; iz < numz_[p]; ++iz) + { + whichpro_[startz_[p] + iz] = p; + } + } +} + +int ParaRgridWorld::numz(int p) const +{ + assert(p >= 0 && p < static_cast(numz_.size())); + return numz_[p]; +} + +int ParaRgridWorld::startz(int p) const +{ + assert(p >= 0 && p < static_cast(startz_.size())); + return startz_[p]; +} + +int ParaRgridWorld::whichpro(int iz) const +{ + assert(iz >= 0 && iz < ncz_); + return whichpro_[iz]; +} + +// ===== Cross-domain operations ===== + +void ParaRgridWorld::reduce_across_pools(double* data, const ParaWorld& kmesh_world) const +{ +#ifdef __MPI + if (!kmesh_world.valid()) return; + if (kmesh_world.size() <= 1) return; + + assert(data != nullptr); + + // Equal-sized pools: corresponding ranks have identical z-slab layouts, + // so local buffers can be summed directly without redistribution. + MPI_Allreduce(MPI_IN_PLACE, data, nrxx(), MPI_DOUBLE, MPI_SUM, kmesh_world.comm()); +#else + (void)data; + (void)kmesh_world; +#endif +} + +void ParaRgridWorld::bcast_data(const double* data_global, double* data_local, + const ParaWorld& comm_world, int root) const +{ + // Serial or single-process: just copy local slab + if (!comm_world.valid() || comm_world.size() == 1) + { + const int ncxy = ncx_ * ncy_; + const int z_start = startz_[rank()]; + for (int ixy = 0; ixy < ncxy; ++ixy) + { + for (int iz = 0; iz < nczp_; ++iz) + { + data_local[ixy * nczp_ + iz] = data_global[ixy * ncz_ + z_start + iz]; + } + } + return; + } + +#ifdef __MPI + // Broadcast z-plane by z-plane + std::vector zpiece(ncx_ * ncy_); + for (int iz = 0; iz < ncz_; ++iz) + { + if (comm_world.rank() == root) + { + for (int ix = 0; ix < ncx_; ++ix) + { + for (int iy = 0; iy < ncy_; ++iy) + { + zpiece[ix * ncy_ + iy] = data_global[(ix * ncy_ + iy) * ncz_ + iz]; + } + } + } + MPI_Bcast(zpiece.data(), ncx_ * ncy_, MPI_DOUBLE, root, comm_world.comm()); + + // Store z-plane if this process owns it + const int znow = iz - startz_[comm_world.rank()]; + if (znow >= 0 && znow < nczp_) + { + for (int ixy = 0; ixy < ncx_ * ncy_; ++ixy) + { + data_local[ixy * nczp_ + znow] = zpiece[ixy]; + } + } + } +#else + (void)data_global; + (void)data_local; + (void)root; +#endif +} + +void ParaRgridWorld::reduce_data(double* rhotot, const double* rhoin, + const ParaWorld& comm_world) const +{ + // Serial: just copy local slab to global grid + if (!comm_world.valid() || comm_world.size() == 1) + { + const int ncxy = ncx_ * ncy_; + const int z_start = startz_[comm_world.rank()]; + for (int ixy = 0; ixy < ncxy; ++ixy) + { + for (int iz = 0; iz < nczp_; ++iz) + { + rhotot[ixy * ncz_ + z_start + iz] = rhoin[ixy * nczp_ + iz]; + } + } + return; + } + +#ifdef __MPI + // Gather local z-slabs from all processes + const int np = comm_world.size(); + std::vector local_z_counts(np); + std::vector receive_counts(np); + std::vector displacements(np, 0); + + int my_nczp = nczp_; + MPI_Allgather(&my_nczp, 1, MPI_INT, local_z_counts.data(), 1, MPI_INT, comm_world.comm()); + + int total_z = 0; + for (int p = 0; p < np; ++p) + { + receive_counts[p] = local_z_counts[p] * ncx_ * ncy_; + if (p > 0) + { + displacements[p] = displacements[p - 1] + receive_counts[p - 1]; + } + total_z += local_z_counts[p]; + } + assert(total_z == ncz_); + + std::vector gathered; + if (comm_world.rank() == 0) + { + gathered.resize(ncx_ * ncy_ * ncz_); + } + + MPI_Gatherv(rhoin, nrxx(), MPI_DOUBLE, + gathered.data(), receive_counts.data(), displacements.data(), + MPI_DOUBLE, 0, comm_world.comm()); + + if (comm_world.rank() == 0) + { + // Convert from rank-contiguous [xy][local_z] to canonical [xy][global_z] + int global_z_start = 0; + for (int p = 0; p < np; ++p) + { + const int local_nz = local_z_counts[p]; + for (int ixy = 0; ixy < ncx_ * ncy_; ++ixy) + { + for (int iz = 0; iz < local_nz; ++iz) + { + rhotot[ixy * ncz_ + global_z_start + iz] + = gathered[displacements[p] + ixy * local_nz + iz]; + } + } + global_z_start += local_nz; + } + } +#else + (void)rhotot; + (void)rhoin; +#endif +} + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_rgrid_world.h b/source/source_base/module_parallel/para_rgrid_world.h new file mode 100644 index 00000000000..ac8806afa8f --- /dev/null +++ b/source/source_base/module_parallel/para_rgrid_world.h @@ -0,0 +1,117 @@ +#ifndef PARA_RGRID_WORLD_H +#define PARA_RGRID_WORLD_H + +#include + +#include "para_world.h" + +namespace Parallel +{ + +/** + * @brief rgrid parallel domain: real-space FFT grid distribution. + * + * Self-contained replacement for GRID_WORLD + GlobalV::GRANK/GSIZE + + * Parallel_Grid's z-distribution tables. Owns grid dimensions and the + * per-process z-plane allocation (numz/startz/whichpro). + * + * Cross-pool operations (reduce_across_pools, bcast, reduce) will accept + * communicators as parameters rather than reading global POOL_WORLD/ + * KP_WORLD, breaking the cross-domain dependency. + * + * Tests only need this header. + */ +class ParaRgridWorld : public ParaWorld +{ +public: + /** + * @brief Construct a serial rgrid domain (all z-planes on one process). + * + * @param[in] ncx, ncy, ncz global grid dimensions + */ + ParaRgridWorld(int ncx, int ncy, int ncz); + +#ifdef __MPI + /** + * @brief Construct an rgrid domain on an existing communicator. + * + * @param[in] comm grid communicator (e.g. GRID_WORLD) + * @param[in] ncx, ncy, ncz global grid dimensions + */ + ParaRgridWorld(const MPI_Comm& comm, int ncx, int ncy, int ncz); +#endif + + /// Global grid dimension in x. + int ncx() const { return ncx_; } + + /// Global grid dimension in y. + int ncy() const { return ncy_; } + + /// Global grid dimension in z. + int ncz() const { return ncz_; } + + /// Local z-plane count for this process. + int nczp() const { return nczp_; } + + /// Total real-space grid points on this process (ncx * ncy * nczp). + int nrxx() const { return ncx_ * ncy_ * nczp_; } + + /// Number of z-planes assigned to process p in this pool. + int numz(int p) const; + + /// Starting global z-index for process p. + int startz(int p) const; + + /// Which process owns global z-plane iz. + int whichpro(int iz) const; + + // ===== Cross-domain operations ===== + + /** + * @brief Sum local grid data across all pools (KP_WORLD or INT_BGROUP). + * + * Replaces Parallel_Grid::reduce_across_pools. + * In serial mode or single-pool, this is a no-op. + * + * @param[in,out] data local grid buffer (nrxx elements), overwritten with sum + * @param[in] kmesh_world k-mesh domain providing the cross-pool communicator + */ + void reduce_across_pools(double* data, const ParaWorld& kmesh_world) const; + + /** + * @brief Broadcast global grid to local z-slabs (replaces Parallel_Grid::bcast). + * + * @param[in] data_global global grid (ncxyz elements, only valid on root) + * @param[out] data_local local grid buffer (nrxx elements) + * @param[in] comm_world communicator for broadcast + * @param[in] root root rank in comm_world + */ + void bcast_data(const double* data_global, double* data_local, + const ParaWorld& comm_world, int root = 0) const; + + /** + * @brief Gather local z-slabs into a global grid (replaces Parallel_Grid::reduce). + * + * @param[out] rhotot global grid (ncxyz elements, only valid on root) + * @param[in] rhoin local grid buffer (nrxx elements) + * @param[in] comm_world communicator for gather + */ + void reduce_data(double* rhotot, const double* rhoin, + const ParaWorld& comm_world) const; + +private: + void distribute_z(); + + int ncx_ = 0; + int ncy_ = 0; + int ncz_ = 0; + int nczp_ = 0; ///< local z-plane count + + std::vector numz_; ///< z-planes per process + std::vector startz_; ///< start z-index per process + std::vector whichpro_; ///< owner of each global z-plane +}; + +} // namespace Parallel + +#endif // PARA_RGRID_WORLD_H diff --git a/source/source_base/module_parallel/para_setup.cpp b/source/source_base/module_parallel/para_setup.cpp new file mode 100644 index 00000000000..744ac568178 --- /dev/null +++ b/source/source_base/module_parallel/para_setup.cpp @@ -0,0 +1,318 @@ +#include "para_setup.h" + +#include +#include + +#ifdef __MPI +#include +#endif + +namespace Parallel +{ + +void divide_mpi_groups(int nproc, int num_groups, int rank, bool even, + int& procs_in_group, int& my_group, int& rank_in_group) +{ + assert(num_groups > 0); + assert(nproc >= num_groups); + + procs_in_group = nproc / num_groups; + int extra_procs = nproc % num_groups; + + if (even && extra_procs != 0) + { + std::cerr << "Error: " << nproc << " processes not evenly divisible by " + << num_groups << " groups." << std::endl; + assert(false); + } + + if (rank < extra_procs * (procs_in_group + 1)) + { + procs_in_group++; + my_group = rank / procs_in_group; + rank_in_group = rank % procs_in_group; + } + else + { + my_group = (rank - extra_procs) / procs_in_group; + rank_in_group = (rank - extra_procs) % procs_in_group; + } +} + +#ifdef __MPI + +namespace { + +// Helper: split a parent communicator into ngroup sub-communicators. +// Mirrors MPICommGroup::divide_group_comm in parallel_comm.cpp: +// - group_comm: intra-group communicator (color = my_group) +// - inter_comm: communicator of same-rank processes across groups +// (color = rank_in_group); MPI_COMM_NULL for a single +// group or an uneven split, exactly like KP_WORLD. +struct GroupSplitResult +{ + MPI_Comm group_comm = MPI_COMM_NULL; + MPI_Comm inter_comm = MPI_COMM_NULL; + int ngroups = 0; + int nprocs_in_group = 0; + int my_group = 0; + int rank_in_group = 0; +}; + +GroupSplitResult split_comm_group(MPI_Comm parent, int ngroup, bool even) +{ + GroupSplitResult res; + res.ngroups = ngroup; + + int gsize = 0; + int grank = 0; + MPI_Comm_size(parent, &gsize); + MPI_Comm_rank(parent, &grank); + + divide_mpi_groups(gsize, ngroup, grank, even, + res.nprocs_in_group, res.my_group, res.rank_in_group); + + // Intra-group communicator: one sub-communicator per group. + MPI_Comm_split(parent, res.my_group, res.rank_in_group, &res.group_comm); + + // Inter-group communicator: processes with the same rank inside + // their group talk to each other. Only valid for an even split; + // an uneven split leaves some groups without a corresponding rank. + const bool is_even = (gsize % ngroup == 0); + if (ngroup > 1 && is_even) + { + MPI_Comm_split(parent, res.rank_in_group, res.my_group, &res.inter_comm); + } + + return res; +} + +} // anonymous namespace + +void split_images(int nproc, int my_rank, int nimage, + int& image_id, int& rank_in_esolver, int& esolver_size, + ParaWorld& esolver_world, ParaWorld& images_world) +{ + assert(nimage > 0); + assert(nproc >= nimage); + + int procs_in_image = 0; + divide_mpi_groups(nproc, nimage, my_rank, false, + procs_in_image, image_id, rank_in_esolver); + esolver_size = procs_in_image; + + // Intra-image domain: all processes of one esolver. + MPI_Comm esolver_comm; + MPI_Comm_split(MPI_COMM_WORLD, image_id, rank_in_esolver, &esolver_comm); + esolver_world = ParaWorld::make_mpi(ParaTag::esolver, esolver_comm); + + // Inter-image domain: same rank_in_esolver across images. Follows the + // KP_WORLD convention: absent for a single image or an uneven split. + const bool is_even = (nproc % nimage == 0); + if (nimage > 1 && is_even) + { + MPI_Comm images_comm; + MPI_Comm_split(MPI_COMM_WORLD, rank_in_esolver, image_id, &images_comm); + images_world = ParaWorld::make_mpi(ParaTag::images, images_comm); + } + else + { + images_world = ParaWorld::make_mpi(ParaTag::images, MPI_COMM_NULL); + } +} + +void split_pools(int parent_size, int parent_rank, int bndpar, int kpar, + const MPI_Comm& parent_comm, + int& nproc_in_pool, int& rank_in_pool, int& my_pool, + int& nproc_in_bndgroup, int& rank_in_bpgroup, int& my_bndgroup, + ParaWorld& pw_world, ParaWorld& kmesh_world, + ParaWorld& bgroup_int, ParaWorld& bgroup_bp) +{ + if (bndpar > 1 && parent_size % (bndpar * kpar) != 0) + { + std::cerr << "Error: " << parent_size + << " processes in the parent domain must be divisible by " + << "BNDPAR*KPAR (" << bndpar * kpar << ")." << std::endl; + assert(false); + } + + // k-point parallelization: split the parent domain into kpar pools. + GroupSplitResult kpar_res = split_comm_group(parent_comm, kpar, false); + + // band parallelization: split each pool into bndpar groups. + GroupSplitResult bndpar_res = split_comm_group(kpar_res.group_comm, bndpar, true); + + // Set output indices. + nproc_in_pool = bndpar_res.nprocs_in_group; + rank_in_pool = bndpar_res.rank_in_group; + my_pool = kpar_res.my_group; + + // POOL_WORLD: processes with the same k point and the same bands + // (plane-wave distribution lives inside it). + MPI_Comm pool_comm; + MPI_Comm_dup(bndpar_res.group_comm, &pool_comm); + pw_world = ParaWorld::make_mpi(ParaTag::pw, pool_comm); + + // KP_WORLD: inter-pool communicator (same rank across pools). + if (kpar_res.inter_comm != MPI_COMM_NULL) + { + MPI_Comm kp_comm; + MPI_Comm_dup(kpar_res.inter_comm, &kp_comm); + kmesh_world = ParaWorld::make_mpi(ParaTag::kmesh, kp_comm); + } + else + { + kmesh_world = ParaWorld::make_mpi(ParaTag::kmesh, MPI_COMM_NULL); + } + + // Band group communicators. + if (bndpar > 1) + { + nproc_in_bndgroup = kpar_res.ngroups * bndpar_res.nprocs_in_group; + rank_in_bpgroup = kpar_res.my_group * bndpar_res.nprocs_in_group + bndpar_res.rank_in_group; + my_bndgroup = bndpar_res.my_group; + + // INT_BGROUP: same bands across pools (bsame_kdiff). + MPI_Comm int_bgroup; + MPI_Comm_split(parent_comm, my_bndgroup, rank_in_bpgroup, &int_bgroup); + bgroup_int = ParaWorld::make_mpi(ParaTag::bsame_kdiff, int_bgroup); + + // BP_WORLD: same k point across band groups (bdiff_ksame). + MPI_Comm bp_comm; + MPI_Comm_dup(bndpar_res.inter_comm, &bp_comm); + bgroup_bp = ParaWorld::make_mpi(ParaTag::bdiff_ksame, bp_comm); + } + else + { + nproc_in_bndgroup = parent_size; + rank_in_bpgroup = parent_rank; + my_bndgroup = 0; + + // No band parallelism: INT_BGROUP spans the whole parent domain, + // BP_WORLD degenerates to one process per rank. + MPI_Comm int_bgroup; + MPI_Comm_dup(parent_comm, &int_bgroup); + bgroup_int = ParaWorld::make_mpi(ParaTag::bsame_kdiff, int_bgroup); + + MPI_Comm bp_comm; + MPI_Comm_split(parent_comm, parent_rank, 0, &bp_comm); + bgroup_bp = ParaWorld::make_mpi(ParaTag::bdiff_ksame, bp_comm); + } +} + +ParaWorld split_diag_world(int diag_np, int parent_size, int parent_rank, + const MPI_Comm& parent_comm, + int& drank, int& dsize, int& dcolor) +{ + assert(diag_np > 0); + + int procs_in_group = 0; + int my_group = 0; + int rank_in_group = 0; + divide_mpi_groups(parent_size, diag_np, parent_rank, false, + procs_in_group, my_group, rank_in_group); + + MPI_Comm diag_comm; + MPI_Comm_split(parent_comm, my_group, rank_in_group, &diag_comm); + + MPI_Comm_rank(diag_comm, &drank); + MPI_Comm_size(diag_comm, &dsize); + dcolor = my_group; + + return ParaWorld::make_mpi(ParaTag::diag, diag_comm); +} + +ParaWorld split_grid_world(int diag_np, int parent_size, int parent_rank, + const MPI_Comm& parent_comm, + int& grank, int& gsize) +{ + assert(diag_np > 0); + + int procs_in_group = 0; + int my_group = 0; + int rank_in_group = 0; + divide_mpi_groups(parent_size, diag_np, parent_rank, false, + procs_in_group, my_group, rank_in_group); + + MPI_Comm grid_comm; + MPI_Comm_split(parent_comm, my_group, rank_in_group, &grid_comm); + + MPI_Comm_rank(grid_comm, &grank); + MPI_Comm_size(grid_comm, &gsize); + + return ParaWorld::make_mpi(ParaTag::rgrid, grid_comm); +} + +ParaCollection setup_para_worlds(int nproc, int my_rank, int nimage, + int bndpar, int kpar, int diag_np) +{ + ParaCollection worlds; + + // 0. Top-level split: independent images. + // esolver_world contains all processes of one esolver instance; + // images_world connects corresponding ranks across images. + int image_id = 0; + int rank_in_esolver = 0; + int esolver_size = 0; + ParaWorld esolver_world = ParaWorld::make_mpi(ParaTag::esolver, MPI_COMM_NULL); + ParaWorld images_world = ParaWorld::make_mpi(ParaTag::images, MPI_COMM_NULL); + split_images(nproc, my_rank, nimage, + image_id, rank_in_esolver, esolver_size, + esolver_world, images_world); + worlds.add(ParaWorld::make_mpi_ptr(ParaTag::esolver, esolver_world.comm())); + // images_world may be an invalid domain (nimage == 1 or uneven split); + // it is still registered so that find(ParaTag::images) returns it and + // callers can test valid(). + worlds.add(ParaWorld::make_mpi_ptr(ParaTag::images, images_world.comm())); + + // All solver domains are derived from the esolver domain, never from + // MPI_COMM_WORLD directly (see the hierarchy diagram in para_setup.h). + const MPI_Comm esolver_comm = esolver_world.comm(); + + // 1. k-pools and band groups. + int nproc_in_pool = 0; + int rank_in_pool = 0; + int my_pool = 0; + int nproc_in_bndgroup = 0; + int rank_in_bpgroup = 0; + int my_bndgroup = 0; + + ParaWorld pw_world = ParaWorld::make_mpi(ParaTag::pw, MPI_COMM_NULL); + ParaWorld kmesh_world = ParaWorld::make_mpi(ParaTag::kmesh, MPI_COMM_NULL); + ParaWorld bgroup_int = ParaWorld::make_mpi(ParaTag::bsame_kdiff, MPI_COMM_NULL); + ParaWorld bgroup_bp = ParaWorld::make_mpi(ParaTag::bdiff_ksame, MPI_COMM_NULL); + + split_pools(esolver_size, rank_in_esolver, bndpar, kpar, esolver_comm, + nproc_in_pool, rank_in_pool, my_pool, + nproc_in_bndgroup, rank_in_bpgroup, my_bndgroup, + pw_world, kmesh_world, bgroup_int, bgroup_bp); + + worlds.add(ParaWorld::make_mpi_ptr(ParaTag::pw, pw_world.comm())); + worlds.add(ParaWorld::make_mpi_ptr(ParaTag::kmesh, kmesh_world.comm())); + worlds.add(ParaWorld::make_mpi_ptr(ParaTag::bsame_kdiff, bgroup_int.comm())); + worlds.add(ParaWorld::make_mpi_ptr(ParaTag::bdiff_ksame, bgroup_bp.comm())); + + // 2. Diagonalization domain. + int drank = 0; + int dsize = 0; + int dcolor = 0; + ParaWorld diag_world = split_diag_world(diag_np, esolver_size, rank_in_esolver, + esolver_comm, drank, dsize, dcolor); + worlds.add(ParaWorld::make_mpi_ptr(ParaTag::diag, diag_world.comm())); + + // 3. Real-space grid domain. + int grank = 0; + int gsize = 0; + ParaWorld grid_world = split_grid_world(diag_np, esolver_size, rank_in_esolver, + esolver_comm, grank, gsize); + worlds.add(ParaWorld::make_mpi_ptr(ParaTag::rgrid, grid_world.comm())); + + // 4. Matrix domain: serial for now until its own 2D-grid split lands. + worlds.add(ParaWorld::make_serial(ParaTag::matrix)); + + return worlds; +} + +#endif // __MPI + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_setup.h b/source/source_base/module_parallel/para_setup.h new file mode 100644 index 00000000000..ee8b51cafed --- /dev/null +++ b/source/source_base/module_parallel/para_setup.h @@ -0,0 +1,183 @@ +#ifndef PARA_SETUP_H +#define PARA_SETUP_H + +#include "para_collection.h" +#include "para_tag.h" +#include "para_world.h" + +#ifdef __MPI +#include "mpi.h" +#endif + +namespace Parallel +{ + +/** + * @file para_setup.h + * + * @brief Construction of the parallel communication domain hierarchy. + * + * The domains form a tree rooted at MPI_COMM_WORLD. The first split + * separates independent calculation images (e.g. NEB replicas, each with + * its own unit cell); every existing solver domain is then derived from + * the intra-image domain instead of being hard-wired to MPI_COMM_WORLD: + * + * @code + * MPI_COMM_WORLD + * | + * +-- split by nimage (same MPICommGroup pattern as k-parallelism) + * | + * | para_esolver_world [color = image_id] + * | | all processes belonging to one esolver + * | | + * | +-- all existing domains are derived from it + * | | (instead of being hard-wired to MPI_COMM_WORLD): + * | | +-- kmesh / pw (split by kpar) + * | | +-- bsame_kdiff / bdiff_ksame (split by bndpar) + * | | +-- rgrid / diag (split by diago_proc) + * | | +-- matrix + * | | + * | para_images_world [color = rank_in_esolver] + * | cross-image: processes with the same + * | rank inside their esolver + * @endcode + * + * Degenerate cases follow the existing KP_WORLD convention: + * - nimage == 1: esolver_world is a dup of MPI_COMM_WORLD and + * images_world is MPI_COMM_NULL (invalid domain). + * - Uneven image split: images_world is MPI_COMM_NULL as well, + * because corresponding ranks do not exist across all images. + */ + +/** + * @brief Divide nproc processes into num_groups groups. + * + * Replaces Parallel_Global::divide_mpi_groups. Works in both serial and + * MPI builds. + * + * @param[in] nproc total number of processes + * @param[in] num_groups desired number of groups + * @param[in] rank global rank + * @param[in] even if true, require an even split (assert on failure) + * @param[out] procs_in_group processes per group + * @param[out] my_group which group this rank belongs to + * @param[out] rank_in_group rank within the group + */ +void divide_mpi_groups(int nproc, int num_groups, int rank, bool even, + int& procs_in_group, int& my_group, int& rank_in_group); + +#ifdef __MPI + +/** + * @brief Split MPI_COMM_WORLD into nimage independent images. + * + * Produces the two top-level domains: + * - esolver_world: intra-image communicator (all processes of one esolver), + * split with color = image_id; + * - images_world: inter-image communicator (same rank_in_esolver across + * images), split with color = rank_in_esolver; MPI_COMM_NULL when + * nimage == 1 or the split is uneven. + * + * @param[in] nproc total number of MPI processes + * @param[in] my_rank global rank + * @param[in] nimage number of images (>= 1) + * @param[out] image_id image this rank belongs to + * @param[out] rank_in_esolver rank inside the esolver domain + * @param[out] esolver_size number of processes in the esolver domain + * @param[out] esolver_world intra-image domain (tag: ParaTag::esolver) + * @param[out] images_world inter-image domain (tag: ParaTag::images, + * invalid when nimage == 1 or uneven split) + */ +void split_images(int nproc, int my_rank, int nimage, + int& image_id, int& rank_in_esolver, int& esolver_size, + ParaWorld& esolver_world, ParaWorld& images_world); + +/** + * @brief Split a parent domain into k-pools and band groups. + * + * Replaces Parallel_Global::divide_pools. All splits are performed inside + * @p parent_comm (normally the esolver domain), so the same code works + * for both single-image and multi-image runs. + * + * @param[in] parent_size number of processes in the parent domain + * @param[in] parent_rank rank of this process in the parent domain + * @param[in] bndpar number of band groups + * @param[in] kpar number of k-pools + * @param[in] parent_comm communicator the pools are split from + * @param[out] nproc_in_pool processes per pool + * @param[out] rank_in_pool rank within pool + * @param[out] my_pool pool index + * @param[out] nproc_in_bndgroup processes per band group + * @param[out] rank_in_bpgroup rank within the band group + * @param[out] my_bndgroup band group index + * @param[out] pw_world domain for POOL_WORLD + * @param[out] kmesh_world domain for KP_WORLD (invalid when kpar == 1) + * @param[out] bgroup_int domain for INT_BGROUP + * @param[out] bgroup_bp domain for BP_WORLD + */ +void split_pools(int parent_size, int parent_rank, int bndpar, int kpar, + const MPI_Comm& parent_comm, + int& nproc_in_pool, int& rank_in_pool, int& my_pool, + int& nproc_in_bndgroup, int& rank_in_bpgroup, int& my_bndgroup, + ParaWorld& pw_world, ParaWorld& kmesh_world, + ParaWorld& bgroup_int, ParaWorld& bgroup_bp); + +/** + * @brief Split a parent domain for diagonalization. + * + * Replaces Parallel_Global::split_diag_world. + * + * @param[in] diag_np number of diag groups + * @param[in] parent_size number of processes in the parent domain + * @param[in] parent_rank rank of this process in the parent domain + * @param[in] parent_comm communicator the diag domain is split from + * @param[out] drank rank in the diag domain + * @param[out] dsize size of the diag domain + * @param[out] dcolor color (diag group index) + * @return ParaWorld for the diag domain + */ +ParaWorld split_diag_world(int diag_np, int parent_size, int parent_rank, + const MPI_Comm& parent_comm, + int& drank, int& dsize, int& dcolor); + +/** + * @brief Split a parent domain for the real-space grid. + * + * Replaces Parallel_Global::split_grid_world. + * + * @param[in] diag_np number of grid groups (same parameter as diag) + * @param[in] parent_size number of processes in the parent domain + * @param[in] parent_rank rank of this process in the parent domain + * @param[in] parent_comm communicator the grid domain is split from + * @param[out] grank rank in the grid domain + * @param[out] gsize size of the grid domain + * @return ParaWorld for the rgrid domain + */ +ParaWorld split_grid_world(int diag_np, int parent_size, int parent_rank, + const MPI_Comm& parent_comm, + int& grank, int& gsize); + +/** + * @brief Assemble the full parallel domain hierarchy. + * + * Top-level initialization: splits images first, then derives every + * solver domain from the esolver domain (see the tree diagram in the + * file header). Replaces the old divide_pools + split_diag_world + + * split_grid_world sequence in driver.cpp. + * + * @param[in] nproc total MPI processes (GlobalV::NPROC) + * @param[in] my_rank global rank (GlobalV::MY_RANK) + * @param[in] nimage number of independent images (1 for a normal run) + * @param[in] bndpar number of band groups + * @param[in] kpar number of k-pools + * @param[in] diag_np number of diag/grid groups + * @return ParaCollection containing all domains + */ +ParaCollection setup_para_worlds(int nproc, int my_rank, int nimage, + int bndpar, int kpar, int diag_np); + +#endif // __MPI + +} // namespace Parallel + +#endif // PARA_SETUP_H diff --git a/source/source_base/module_parallel/para_tag.h b/source/source_base/module_parallel/para_tag.h new file mode 100644 index 00000000000..7d2ff9543ce --- /dev/null +++ b/source/source_base/module_parallel/para_tag.h @@ -0,0 +1,41 @@ +#ifndef PARA_TAG_H +#define PARA_TAG_H + +#include + +namespace Parallel +{ + +/** + * @brief Domain tag constants for the parallel communication domains. + * + * These tags replace raw string literals to avoid typo-induced runtime + * failures. They map to the legacy global communicators as follows: + * - esolver -> one esolver instance (intra-image communicator) + * - images -> cross-image communicator (same rank_in_esolver) + * - pw -> POOL_WORLD + * - kmesh -> KP_WORLD + * - bsame_kdiff -> INT_BGROUP + * - bdiff_ksame -> BP_WORLD + * - rgrid -> GRID_WORLD + * - diag -> DIAG_WORLD + * - matrix -> matrix domain + * - atom -> atom domain + */ +namespace ParaTag +{ +const std::string esolver = "esolver"; +const std::string images = "images"; +const std::string pw = "pw"; +const std::string kmesh = "kmesh"; +const std::string bsame_kdiff = "bsame_kdiff"; +const std::string bdiff_ksame = "bdiff_ksame"; +const std::string rgrid = "rgrid"; +const std::string diag = "diag"; +const std::string matrix = "matrix"; +const std::string atom = "atom"; +} // namespace ParaTag + +} // namespace Parallel + +#endif // PARA_TAG_H diff --git a/source/source_base/module_parallel/para_world.cpp b/source/source_base/module_parallel/para_world.cpp new file mode 100644 index 00000000000..0a4ca51748e --- /dev/null +++ b/source/source_base/module_parallel/para_world.cpp @@ -0,0 +1,43 @@ +#include "para_world.h" + +namespace Parallel +{ + +ParaWorld::ParaWorld(const std::string& tag) : tag_(tag), rank_(0), size_(1) +{ +#ifdef __MPI + if (!tag.empty()) + { + comm_ = MPI_COMM_SELF; + } + else + { + comm_ = MPI_COMM_NULL; + } +#endif +} + +#ifdef __MPI +ParaWorld::ParaWorld(const std::string& tag, const MPI_Comm& comm) : tag_(tag), comm_(comm) +{ + if (comm == MPI_COMM_NULL) + { + rank_ = -1; + size_ = 0; + return; + } + MPI_Comm_rank(comm, &rank_); + MPI_Comm_size(comm, &size_); +} +#endif + +bool ParaWorld::valid() const +{ +#ifdef __MPI + return comm_ != MPI_COMM_NULL; +#else + return !tag_.empty(); +#endif +} + +} // namespace Parallel diff --git a/source/source_base/module_parallel/para_world.h b/source/source_base/module_parallel/para_world.h new file mode 100644 index 00000000000..a8291fad8bd --- /dev/null +++ b/source/source_base/module_parallel/para_world.h @@ -0,0 +1,140 @@ +#ifndef PARA_WORLD_H +#define PARA_WORLD_H + +#include +#include + +#ifdef __MPI +#include "mpi.h" +#endif + +namespace Parallel +{ + +/** + * @brief Value type describing one MPI communication domain. + * + * A ParaWorld couples a domain tag (a short string constant, see + * para_tag.h) with the communicator, rank and size of the current + * process inside that domain. It replaces loose globals such as + * GlobalV::RANK_IN_POOL / POOL_WORLD by an object that functions + * receive explicitly. + * + * In serial builds (no __MPI) the communicator member does not + * exist; rank() always returns 0 and size() always returns 1, so + * call sites compile unchanged in both serial and MPI builds. + */ +class ParaWorld +{ +public: + virtual ~ParaWorld() = default; + + /// Domain tag string. + const std::string& tag() const + { + return tag_; + } + + /// Rank of this process inside the domain (0 in serial builds). + int rank() const + { + return rank_; + } + + /// Number of processes in the domain (1 in serial builds). + int size() const + { + return size_; + } + + /** + * @brief True if this process belongs to the domain. + * + * In MPI builds this means comm() != MPI_COMM_NULL; in serial + * builds a default/empty domain is invalid, everything else valid. + */ + bool valid() const; + +#ifdef __MPI + /// Underlying MPI communicator (MPI builds only). + MPI_Comm comm() const + { + return comm_; + } +#endif + + /** + * @brief Build a serial (size=1, rank=0) domain for the given tag. + * + * Safe degradation used by tests and by ParaCollection when a tag + * is not found. + */ + static ParaWorld serial(const std::string& tag) + { + return ParaWorld(tag); + } + + /** + * @brief Build a serial domain as a heap-allocated unique_ptr. + * + * Convenience factory for ParaCollection::add(). + */ + static std::unique_ptr make_serial(const std::string& tag) + { + return std::unique_ptr(new ParaWorld(tag)); + } + +#ifdef __MPI + /** + * @brief Build a domain wrapping an MPI communicator. + * + * Factory for setup functions that need to create ParaWorld objects + * from split communicators. + */ + static ParaWorld make_mpi(const std::string& tag, const MPI_Comm& comm) + { + return ParaWorld(tag, comm); + } + + static std::unique_ptr make_mpi_ptr(const std::string& tag, const MPI_Comm& comm) + { + return std::unique_ptr(new ParaWorld(tag, comm)); + } +#endif + +protected: + /** + * @brief Construct a serial (single-process) domain. + * + * Usable in both serial and MPI builds; in MPI builds the + * communicator is set to MPI_COMM_SELF. Mainly intended for + * tests and safe fall-back behavior. + * + * @param[in] tag domain tag string (must be non-empty for a + * meaningful domain; empty tag marks "no domain") + */ + explicit ParaWorld(const std::string& tag); + +#ifdef __MPI + /** + * @brief Construct a domain wrapping an existing MPI communicator. + * + * @param[in] tag domain tag string + * @param[in] comm MPI communicator (may be MPI_COMM_NULL, which + * yields an invalid domain on this rank) + */ + ParaWorld(const std::string& tag, const MPI_Comm& comm); +#endif + +private: + std::string tag_; ///< domain tag + int rank_; ///< rank inside domain + int size_; ///< number of processes in domain +#ifdef __MPI + MPI_Comm comm_; ///< wrapped communicator (never owned/freed here) +#endif +}; + +} // namespace Parallel + +#endif // PARA_WORLD_H diff --git a/source/source_base/module_parallel/test/CMakeLists.txt b/source/source_base/module_parallel/test/CMakeLists.txt new file mode 100644 index 00000000000..6f3a2169a7c --- /dev/null +++ b/source/source_base/module_parallel/test/CMakeLists.txt @@ -0,0 +1,100 @@ +abacus_disable_feature_definitions(__MPI) +AddTest( + TARGET MODULE_BASE_para_world + SOURCES para_world_test.cpp ../para_world.cpp +) + +AddTest( + TARGET MODULE_BASE_para_collection + SOURCES para_collection_test.cpp ../para_collection.cpp ../para_world.cpp ../para_kmesh_world.cpp +) + +AddTest( + TARGET MODULE_BASE_para_kmesh_world + SOURCES para_kmesh_world_test.cpp ../para_kmesh_world.cpp ../para_world.cpp +) + +AddTest( + TARGET MODULE_BASE_para_pw_world + SOURCES para_pw_world_test.cpp ../para_pw_world.cpp ../para_world.cpp +) + +AddTest( + TARGET MODULE_BASE_para_diag_world + SOURCES para_diag_world_test.cpp ../para_diag_world.cpp ../para_world.cpp +) + +AddTest( + TARGET MODULE_BASE_para_rgrid_world + SOURCES para_rgrid_world_test.cpp ../para_rgrid_world.cpp ../para_world.cpp +) + +AddTest( + TARGET MODULE_BASE_para_bgroup_world + SOURCES para_bgroup_world_test.cpp ../para_bgroup_world.cpp ../para_world.cpp +) + +AddTest( + TARGET MODULE_BASE_para_matrix_world + SOURCES para_matrix_world_test.cpp ../para_matrix_world.cpp ../para_world.cpp +) + +AddTest( + TARGET MODULE_BASE_para_mpi_func + SOURCES para_mpi_func_test.cpp ../para_mpi_func.cpp ../para_world.cpp +) + +AddTest( + TARGET MODULE_BASE_para_setup + SOURCES para_setup_test.cpp ../para_setup.cpp ../para_world.cpp ../para_collection.cpp +) + +AddTest( + TARGET MODULE_BASE_para_world_mpi + LIBS MPI::MPI_CXX + SOURCES para_world_mpi_test.cpp ../para_world.cpp +) +target_compile_definitions(MODULE_BASE_para_world_mpi PRIVATE __MPI) + +AddTest( + TARGET MODULE_BASE_para_collection_mpi + LIBS MPI::MPI_CXX + SOURCES para_collection_mpi_test.cpp ../para_collection.cpp ../para_world.cpp ../para_kmesh_world.cpp +) +target_compile_definitions(MODULE_BASE_para_collection_mpi PRIVATE __MPI) + +AddTest( + TARGET MODULE_BASE_para_mpi_func_mpi + LIBS MPI::MPI_CXX + SOURCES para_mpi_func_mpi_test.cpp ../para_mpi_func.cpp ../para_world.cpp +) +target_compile_definitions(MODULE_BASE_para_mpi_func_mpi PRIVATE __MPI) + +AddTest( + TARGET MODULE_BASE_para_setup_mpi + LIBS MPI::MPI_CXX + SOURCES para_setup_mpi_test.cpp ../para_setup.cpp ../para_world.cpp ../para_collection.cpp +) +target_compile_definitions(MODULE_BASE_para_setup_mpi PRIVATE __MPI) + +file(COPY para_world_mpi_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) +file(COPY para_collection_mpi_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) +file(COPY para_mpi_func_mpi_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) +file(COPY para_setup_mpi_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) +find_program(BASH bash) +add_test(NAME MODULE_BASE_para_world_mpi_test + COMMAND ${BASH} para_world_mpi_test.sh + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} +) +add_test(NAME MODULE_BASE_para_collection_mpi_test + COMMAND ${BASH} para_collection_mpi_test.sh + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} +) +add_test(NAME MODULE_BASE_para_mpi_func_mpi_test + COMMAND ${BASH} para_mpi_func_mpi_test.sh + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} +) +add_test(NAME MODULE_BASE_para_setup_mpi_test + COMMAND ${BASH} para_setup_mpi_test.sh + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} +) diff --git a/source/source_base/module_parallel/test/para_bgroup_world_test.cpp b/source/source_base/module_parallel/test/para_bgroup_world_test.cpp new file mode 100644 index 00000000000..8dceccf4ad9 --- /dev/null +++ b/source/source_base/module_parallel/test/para_bgroup_world_test.cpp @@ -0,0 +1,21 @@ +#include "gtest/gtest.h" + +#include "../para_bgroup_world.h" + +TEST(ParaBgroupWorldTest, SerialMode) +{ + const Parallel::ParaBgroupWorld world; + EXPECT_EQ(world.tag(), "bdiff_ksame"); + EXPECT_EQ(world.my_bndgroup(), 0); + EXPECT_EQ(world.nbndgroup(), 1); + EXPECT_EQ(world.rank_in_bpgroup(), 0); + EXPECT_EQ(world.nproc_in_bndgroup(), 1); + EXPECT_TRUE(world.valid()); +} + +TEST(ParaBgroupWorldTest, AliasesMatchBase) +{ + const Parallel::ParaBgroupWorld world; + EXPECT_EQ(world.rank_in_bpgroup(), world.rank()); + EXPECT_EQ(world.nproc_in_bndgroup(), world.size()); +} diff --git a/source/source_base/module_parallel/test/para_collection_mpi_test.cpp b/source/source_base/module_parallel/test/para_collection_mpi_test.cpp new file mode 100644 index 00000000000..a166ff05322 --- /dev/null +++ b/source/source_base/module_parallel/test/para_collection_mpi_test.cpp @@ -0,0 +1,53 @@ +#include "gtest/gtest.h" + +#include "../para_collection.h" +#include "../para_kmesh_world.h" +#include "../para_tag.h" + +TEST(ParaCollectionMpiTest, AssembleAndFind) +{ + Parallel::ParaCollection coll; + coll.add(std::unique_ptr( + new Parallel::ParaKmeshWorld(MPI_COMM_WORLD, 1, 0, 1, 4, 1))); + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::pw)); + EXPECT_EQ(coll.size(), 2u); + + const Parallel::ParaWorld& kmesh = coll.find(Parallel::ParaTag::kmesh); + EXPECT_TRUE(kmesh.valid()); + EXPECT_EQ(kmesh.tag(), "kmesh"); + + const Parallel::ParaWorld& pw = coll.find(Parallel::ParaTag::pw); + EXPECT_TRUE(pw.valid()); + EXPECT_EQ(pw.size(), 1); +} + +TEST(ParaCollectionMpiTest, FindMissingReturnsInvalid) +{ + Parallel::ParaCollection coll; + coll.add(std::unique_ptr( + new Parallel::ParaKmeshWorld(MPI_COMM_WORLD, 1, 0, 1, 4, 1))); + + const Parallel::ParaWorld& missing = coll.find("nonexistent"); + EXPECT_FALSE(missing.valid()); +} + +TEST(ParaCollectionMpiTest, FindAsSubclass) +{ + Parallel::ParaCollection coll; + coll.add(std::unique_ptr( + new Parallel::ParaKmeshWorld(MPI_COMM_WORLD, 1, 0, 1, 8, 1))); + + const Parallel::ParaKmeshWorld* kmesh = coll.find_as(Parallel::ParaTag::kmesh); + ASSERT_NE(kmesh, nullptr); + EXPECT_EQ(kmesh->nkstot(), 8); + EXPECT_EQ(kmesh->nks_local(), 8); +} + +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + testing::InitGoogleTest(&argc, argv); + const int result = RUN_ALL_TESTS(); + MPI_Finalize(); + return result; +} diff --git a/source/source_base/module_parallel/test/para_collection_mpi_test.sh b/source/source_base/module_parallel/test/para_collection_mpi_test.sh new file mode 100644 index 00000000000..f5baca3b7d0 --- /dev/null +++ b/source/source_base/module_parallel/test/para_collection_mpi_test.sh @@ -0,0 +1,18 @@ +#!/bin/bash -e + +np=`cat /proc/cpuinfo | grep "cpu cores" | uniq| awk '{print $NF}'` +echo "nprocs in this machine is $np" + +for i in 4;do + if [[ $i -gt $np ]];then + continue + fi + echo "TEST in parallel, nprocs=$i" + mpirun -np $i ./MODULE_BASE_para_collection_mpi + if [[ $? -ne 0 ]]; then + echo -e "\e[1;33m [ FAILED ] \e[0m"\ + "execute UT with $i cores error." + exit 1 + fi + break +done diff --git a/source/source_base/module_parallel/test/para_collection_test.cpp b/source/source_base/module_parallel/test/para_collection_test.cpp new file mode 100644 index 00000000000..cb9ce060361 --- /dev/null +++ b/source/source_base/module_parallel/test/para_collection_test.cpp @@ -0,0 +1,79 @@ +#include "gtest/gtest.h" + +#include "../para_collection.h" +#include "../para_kmesh_world.h" +#include "../para_tag.h" + +TEST(ParaCollectionTest, DefaultIsEmpty) +{ + const Parallel::ParaCollection coll; + EXPECT_EQ(coll.size(), 0u); +} + +TEST(ParaCollectionTest, AddAndFind) +{ + Parallel::ParaCollection coll; + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::pw)); + coll.add(std::unique_ptr(new Parallel::ParaKmeshWorld(4, 1))); + EXPECT_EQ(coll.size(), 2u); + + const Parallel::ParaWorld& pw = coll.find(Parallel::ParaTag::pw); + EXPECT_EQ(pw.tag(), "pw"); + EXPECT_TRUE(pw.valid()); + + const Parallel::ParaWorld& kmesh = coll.find(Parallel::ParaTag::kmesh); + EXPECT_EQ(kmesh.tag(), "kmesh"); + EXPECT_TRUE(kmesh.valid()); +} + +TEST(ParaCollectionTest, FindMissingReturnsEmpty) +{ + Parallel::ParaCollection coll; + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::pw)); + + const Parallel::ParaWorld& missing = coll.find("nonexistent"); + EXPECT_TRUE(missing.tag().empty()); + EXPECT_FALSE(missing.valid()); +} + +TEST(ParaCollectionTest, DuplicateTagRejected) +{ + Parallel::ParaCollection coll; + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::pw)); + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::pw)); + EXPECT_EQ(coll.size(), 1u); +} + +TEST(ParaCollectionTest, FindAsSubclass) +{ + Parallel::ParaCollection coll; + coll.add(std::unique_ptr(new Parallel::ParaKmeshWorld(6, 1))); + + const Parallel::ParaKmeshWorld* kmesh = coll.find_as(Parallel::ParaTag::kmesh); + ASSERT_NE(kmesh, nullptr); + EXPECT_EQ(kmesh->nkstot(), 6); + EXPECT_EQ(kmesh->kpar(), 1); +} + +TEST(ParaCollectionTest, FindAllEightDomains) +{ + Parallel::ParaCollection coll; + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::pw)); + coll.add(std::unique_ptr(new Parallel::ParaKmeshWorld(4, 1))); + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::bsame_kdiff)); + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::bdiff_ksame)); + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::rgrid)); + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::diag)); + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::matrix)); + coll.add(Parallel::ParaWorld::make_serial(Parallel::ParaTag::atom)); + EXPECT_EQ(coll.size(), 8u); + + EXPECT_TRUE(coll.find(Parallel::ParaTag::pw).valid()); + EXPECT_TRUE(coll.find(Parallel::ParaTag::kmesh).valid()); + EXPECT_TRUE(coll.find(Parallel::ParaTag::bsame_kdiff).valid()); + EXPECT_TRUE(coll.find(Parallel::ParaTag::bdiff_ksame).valid()); + EXPECT_TRUE(coll.find(Parallel::ParaTag::rgrid).valid()); + EXPECT_TRUE(coll.find(Parallel::ParaTag::diag).valid()); + EXPECT_TRUE(coll.find(Parallel::ParaTag::matrix).valid()); + EXPECT_TRUE(coll.find(Parallel::ParaTag::atom).valid()); +} diff --git a/source/source_base/module_parallel/test/para_diag_world_test.cpp b/source/source_base/module_parallel/test/para_diag_world_test.cpp new file mode 100644 index 00000000000..b44be89b106 --- /dev/null +++ b/source/source_base/module_parallel/test/para_diag_world_test.cpp @@ -0,0 +1,20 @@ +#include "gtest/gtest.h" + +#include "../para_diag_world.h" + +TEST(ParaDiagWorldTest, SerialMode) +{ + const Parallel::ParaDiagWorld world; + EXPECT_EQ(world.tag(), "diag"); + EXPECT_EQ(world.drank(), 0); + EXPECT_EQ(world.dsize(), 1); + EXPECT_EQ(world.dcolor(), 0); + EXPECT_TRUE(world.valid()); +} + +TEST(ParaDiagWorldTest, AliasesMatchBase) +{ + const Parallel::ParaDiagWorld world; + EXPECT_EQ(world.drank(), world.rank()); + EXPECT_EQ(world.dsize(), world.size()); +} diff --git a/source/source_base/module_parallel/test/para_kmesh_world_test.cpp b/source/source_base/module_parallel/test/para_kmesh_world_test.cpp new file mode 100644 index 00000000000..345bb3e558e --- /dev/null +++ b/source/source_base/module_parallel/test/para_kmesh_world_test.cpp @@ -0,0 +1,74 @@ +#include "gtest/gtest.h" + +#include "../para_kmesh_world.h" + +TEST(ParaKmeshWorldTest, SerialSinglePool) +{ + const Parallel::ParaKmeshWorld world(4, 1); + EXPECT_EQ(world.tag(), "kmesh"); + EXPECT_EQ(world.kpar(), 1); + EXPECT_EQ(world.my_pool(), 0); + EXPECT_EQ(world.rank_in_pool(), 0); + EXPECT_EQ(world.nproc(), 1); + EXPECT_EQ(world.nspin(), 1); + EXPECT_EQ(world.nkstot(), 4); + EXPECT_EQ(world.nks_local(), 4); + EXPECT_EQ(world.startk_global(), 0); +} + +TEST(ParaKmeshWorldTest, EvenDistribution) +{ + const Parallel::ParaKmeshWorld world(6, 1); + // serial: kpar=1, so all 6 k-points in pool 0 + EXPECT_EQ(world.nks_pool(0), 6); + EXPECT_EQ(world.startk_pool(0), 0); + EXPECT_EQ(world.max_nks_pool(), 6); +} + +TEST(ParaKmeshWorldTest, WhichPool) +{ + const Parallel::ParaKmeshWorld world(5, 1); + for (int ik = 0; ik < 5; ++ik) + { + EXPECT_EQ(world.which_pool(ik), 0); + } +} + +TEST(ParaKmeshWorldTest, PoolCollectionSerial) +{ + const Parallel::ParaKmeshWorld world(3, 1); + const double wk[] = {0.5, 0.3, 0.2}; + double value = 0.0; + world.pool_collection(value, wk, 1); + EXPECT_DOUBLE_EQ(value, 0.3); +} + +TEST(ParaKmeshWorldTest, PoolCollectionArraySerial) +{ + const Parallel::ParaKmeshWorld world(2, 1); + // 2 k-points, 3 elements each: k0={1,2,3}, k1={4,5,6} + const double w[] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + double value[3] = {0.0, 0.0, 0.0}; + world.pool_collection(value, w, 3, 1); + EXPECT_DOUBLE_EQ(value[0], 4.0); + EXPECT_DOUBLE_EQ(value[1], 5.0); + EXPECT_DOUBLE_EQ(value[2], 6.0); +} + +TEST(ParaKmeshWorldTest, GatherKvecSerial) +{ + const Parallel::ParaKmeshWorld world(2, 1); + const std::vector local = {1.0, 0.0, 0.0, 0.0, 1.0, 0.0}; + std::vector global; + world.gather_kvec(local, global); + ASSERT_EQ(global.size(), 6u); + EXPECT_DOUBLE_EQ(global[0], 1.0); + EXPECT_DOUBLE_EQ(global[4], 1.0); +} + +TEST(ParaKmeshWorldTest, Nspin2Restriction) +{ + // nspin=2 forces pool=0 for pool_collection (legacy behavior) + const Parallel::ParaKmeshWorld world(4, 2); + EXPECT_EQ(world.nspin(), 2); +} diff --git a/source/source_base/module_parallel/test/para_matrix_world_test.cpp b/source/source_base/module_parallel/test/para_matrix_world_test.cpp new file mode 100644 index 00000000000..6da649f291a --- /dev/null +++ b/source/source_base/module_parallel/test/para_matrix_world_test.cpp @@ -0,0 +1,20 @@ +#include "gtest/gtest.h" + +#include "../para_matrix_world.h" + +TEST(ParaMatrixWorldTest, SerialMode) +{ + const Parallel::ParaMatrixWorld world; + EXPECT_EQ(world.tag(), "matrix"); + EXPECT_EQ(world.dim0(), 1); + EXPECT_EQ(world.dim1(), 1); + EXPECT_EQ(world.coord_row(), 0); + EXPECT_EQ(world.coord_col(), 0); + EXPECT_TRUE(world.valid()); +} + +TEST(ParaMatrixWorldTest, GridProductMatchesSize) +{ + const Parallel::ParaMatrixWorld world; + EXPECT_EQ(world.dim0() * world.dim1(), world.size()); +} diff --git a/source/source_base/module_parallel/test/para_mpi_func_mpi_test.cpp b/source/source_base/module_parallel/test/para_mpi_func_mpi_test.cpp new file mode 100644 index 00000000000..3f0b4df5eef --- /dev/null +++ b/source/source_base/module_parallel/test/para_mpi_func_mpi_test.cpp @@ -0,0 +1,45 @@ +#include "gtest/gtest.h" + +#include "../para_mpi_func.h" +#include "../para_world.h" + +TEST(ParaMpiFuncMpiTest, BcastIntFromRoot) +{ + auto world = Parallel::ParaWorld::serial("test"); + int v = (world.rank() == 0) ? 99 : 0; + Parallel::bcast_int(v, world); + EXPECT_EQ(v, 99); +} + +TEST(ParaMpiFuncMpiTest, ReduceAllSum) +{ + auto world = Parallel::ParaWorld::serial("test"); + int v = 1; // each rank contributes 1 + Parallel::reduce_all(v, world); + EXPECT_EQ(v, 1); // serial: size=1 +} + +TEST(ParaMpiFuncMpiTest, GatherIntAll) +{ + auto world = Parallel::ParaWorld::serial("test"); + int v = world.rank(); + int all[1] = {0}; + Parallel::gather_int(v, all, world); + EXPECT_EQ(all[0], 0); +} + +TEST(ParaMpiFuncMpiTest, BarrierNoHang) +{ + auto world = Parallel::ParaWorld::serial("test"); + Parallel::barrier(world); + SUCCEED(); +} + +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + testing::InitGoogleTest(&argc, argv); + const int result = RUN_ALL_TESTS(); + MPI_Finalize(); + return result; +} diff --git a/source/source_base/module_parallel/test/para_mpi_func_mpi_test.sh b/source/source_base/module_parallel/test/para_mpi_func_mpi_test.sh new file mode 100644 index 00000000000..afdec3f3ee9 --- /dev/null +++ b/source/source_base/module_parallel/test/para_mpi_func_mpi_test.sh @@ -0,0 +1,18 @@ +#!/bin/bash -e + +np=`cat /proc/cpuinfo | grep "cpu cores" | uniq| awk '{print $NF}'` +echo "nprocs in this machine is $np" + +for i in 4;do + if [[ $i -gt $np ]];then + continue + fi + echo "TEST in parallel, nprocs=$i" + mpirun -np $i ./MODULE_BASE_para_mpi_func_mpi + if [[ $? -ne 0 ]]; then + echo -e "\e[1;33m [ FAILED ] \e[0m"\ + "execute UT with $i cores error." + exit 1 + fi + break +done diff --git a/source/source_base/module_parallel/test/para_mpi_func_test.cpp b/source/source_base/module_parallel/test/para_mpi_func_test.cpp new file mode 100644 index 00000000000..98159b040a9 --- /dev/null +++ b/source/source_base/module_parallel/test/para_mpi_func_test.cpp @@ -0,0 +1,102 @@ +#include "gtest/gtest.h" + +#include "../para_mpi_func.h" +#include "../para_world.h" + +TEST(ParaMpiFuncTest, BcastIntSerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + int v = 42; + Parallel::bcast_int(v, world); + EXPECT_EQ(v, 42); +} + +TEST(ParaMpiFuncTest, BcastDoubleSerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + double v = 3.14; + Parallel::bcast_double(v, world); + EXPECT_DOUBLE_EQ(v, 3.14); +} + +TEST(ParaMpiFuncTest, BcastBoolSerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + bool v = true; + Parallel::bcast_bool(v, world); + EXPECT_TRUE(v); +} + +TEST(ParaMpiFuncTest, BcastStringSerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + std::string s = "hello"; + Parallel::bcast_string(s, world); + EXPECT_EQ(s, "hello"); +} + +TEST(ParaMpiFuncTest, BcastIntArraySerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + int v[] = {1, 2, 3}; + Parallel::bcast_int(v, 3, world); + EXPECT_EQ(v[0], 1); + EXPECT_EQ(v[1], 2); + EXPECT_EQ(v[2], 3); +} + +TEST(ParaMpiFuncTest, BcastComplexSerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + std::complex v(1.0, 2.0); + Parallel::bcast_complex(v, world); + EXPECT_DOUBLE_EQ(v.real(), 1.0); + EXPECT_DOUBLE_EQ(v.imag(), 2.0); +} + +TEST(ParaMpiFuncTest, BcastCharArraySerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + char buf[] = "abc"; + Parallel::bcast_char(buf, 4, world); + EXPECT_EQ(buf[0], 'a'); +} + +TEST(ParaMpiFuncTest, ReduceAllSerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + double v = 5.0; + Parallel::reduce_all(v, world); + EXPECT_DOUBLE_EQ(v, 5.0); +} + +TEST(ParaMpiFuncTest, ReduceMinMaxSerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + double v = 7.5; + Parallel::reduce_min(v, world); + EXPECT_DOUBLE_EQ(v, 7.5); + Parallel::reduce_max(v, world); + EXPECT_DOUBLE_EQ(v, 7.5); +} + +TEST(ParaMpiFuncTest, GatherIntSerial) +{ + auto world = Parallel::ParaWorld::serial("test"); + int v = 99; + int all[1] = {0}; + Parallel::gather_int(v, all, world); + EXPECT_EQ(all[0], 99); +} + +TEST(ParaMpiFuncTest, InvalidWorldIsNoop) +{ + auto world = Parallel::ParaWorld::serial(""); + EXPECT_FALSE(world.valid()); + int v = 42; + Parallel::bcast_int(v, world); + EXPECT_EQ(v, 42); // unchanged + Parallel::reduce_all(v, world); + EXPECT_EQ(v, 42); // unchanged + Parallel::barrier(world); // should not hang or crash +} diff --git a/source/source_base/module_parallel/test/para_pw_world_test.cpp b/source/source_base/module_parallel/test/para_pw_world_test.cpp new file mode 100644 index 00000000000..664fa3634be --- /dev/null +++ b/source/source_base/module_parallel/test/para_pw_world_test.cpp @@ -0,0 +1,28 @@ +#include "gtest/gtest.h" + +#include "../para_pw_world.h" + +TEST(ParaPwWorldTest, SerialMode) +{ + const Parallel::ParaPwWorld world(128); + EXPECT_EQ(world.tag(), "pw"); + EXPECT_EQ(world.npw(), 128); + EXPECT_EQ(world.npwtot(), 128); + EXPECT_EQ(world.poolnproc(), 1); + EXPECT_EQ(world.poolrank(), 0); + EXPECT_EQ(world.npw_per(0), 128); +} + +TEST(ParaPwWorldTest, SerialZero) +{ + const Parallel::ParaPwWorld world(0); + EXPECT_EQ(world.npw(), 0); + EXPECT_EQ(world.npwtot(), 0); +} + +TEST(ParaPwWorldTest, Validity) +{ + const Parallel::ParaPwWorld world(64); + EXPECT_TRUE(world.valid()); + EXPECT_EQ(world.size(), 1); +} diff --git a/source/source_base/module_parallel/test/para_rgrid_world_test.cpp b/source/source_base/module_parallel/test/para_rgrid_world_test.cpp new file mode 100644 index 00000000000..a89c8ea736a --- /dev/null +++ b/source/source_base/module_parallel/test/para_rgrid_world_test.cpp @@ -0,0 +1,121 @@ +#include "gtest/gtest.h" + +#include "../para_rgrid_world.h" +#include "../para_world.h" + +#include + +TEST(ParaRgridWorldTest, SerialMode) +{ + const Parallel::ParaRgridWorld world(4, 5, 6); + EXPECT_EQ(world.tag(), "rgrid"); + EXPECT_EQ(world.ncx(), 4); + EXPECT_EQ(world.ncy(), 5); + EXPECT_EQ(world.ncz(), 6); + EXPECT_EQ(world.nczp(), 6); + EXPECT_EQ(world.nrxx(), 120); + EXPECT_EQ(world.numz(0), 6); + EXPECT_EQ(world.startz(0), 0); +} + +TEST(ParaRgridWorldTest, WhichproSerial) +{ + const Parallel::ParaRgridWorld world(2, 2, 8); + for (int iz = 0; iz < 8; ++iz) + { + EXPECT_EQ(world.whichpro(iz), 0); + } +} + +TEST(ParaRgridWorldTest, Validity) +{ + const Parallel::ParaRgridWorld world(1, 1, 1); + EXPECT_TRUE(world.valid()); +} + +// ===== Cross-domain operation tests (serial mode) ===== + +TEST(ParaRgridWorldTest, ReduceAcrossPoolsSerial) +{ + Parallel::ParaRgridWorld rgrid(2, 2, 4); + auto fake_kmesh = Parallel::ParaWorld::serial("kmesh"); + + std::vector data(rgrid.nrxx(), 1.0); + rgrid.reduce_across_pools(data.data(), fake_kmesh); + + // Serial mode: no-op, data unchanged + for (int i = 0; i < rgrid.nrxx(); ++i) + { + EXPECT_DOUBLE_EQ(data[i], 1.0); + } +} + +TEST(ParaRgridWorldTest, BcastDataSerial) +{ + const int ncx = 2, ncy = 2, ncz = 4; + Parallel::ParaRgridWorld rgrid(ncx, ncy, ncz); + auto fake_comm = Parallel::ParaWorld::serial("comm"); + + // Build global grid: value = ixy * ncz + iz + std::vector global(ncx * ncy * ncz); + for (int ixy = 0; ixy < ncx * ncy; ++ixy) + { + for (int iz = 0; iz < ncz; ++iz) + { + global[ixy * ncz + iz] = ixy * ncz + iz; + } + } + + std::vector local(rgrid.nrxx(), -1.0); + rgrid.bcast_data(global.data(), local.data(), fake_comm); + + // Serial: local should have all z-planes + for (int ixy = 0; ixy < ncx * ncy; ++ixy) + { + for (int iz = 0; iz < ncz; ++iz) + { + EXPECT_DOUBLE_EQ(local[ixy * ncz + iz], ixy * ncz + iz); + } + } +} + +TEST(ParaRgridWorldTest, ReduceDataSerial) +{ + const int ncx = 2, ncy = 2, ncz = 4; + Parallel::ParaRgridWorld rgrid(ncx, ncy, ncz); + auto fake_comm = Parallel::ParaWorld::serial("comm"); + + // Local grid: value = ixy * nczp + iz + std::vector local(rgrid.nrxx()); + for (int i = 0; i < rgrid.nrxx(); ++i) + { + local[i] = static_cast(i); + } + + std::vector global(ncx * ncy * ncz, -1.0); + rgrid.reduce_data(global.data(), local.data(), fake_comm); + + // Serial: global should match local (single process owns all z) + for (int ixy = 0; ixy < ncx * ncy; ++ixy) + { + for (int iz = 0; iz < ncz; ++iz) + { + EXPECT_DOUBLE_EQ(global[ixy * ncz + iz], local[ixy * ncz + iz]); + } + } +} + +TEST(ParaRgridWorldTest, ReduceAcrossPoolsInvalidWorld) +{ + Parallel::ParaRgridWorld rgrid(2, 2, 4); + auto invalid = Parallel::ParaWorld::serial(""); + + std::vector data(rgrid.nrxx(), 5.0); + rgrid.reduce_across_pools(data.data(), invalid); + + // Invalid world: no-op + for (int i = 0; i < rgrid.nrxx(); ++i) + { + EXPECT_DOUBLE_EQ(data[i], 5.0); + } +} diff --git a/source/source_base/module_parallel/test/para_setup_mpi_test.cpp b/source/source_base/module_parallel/test/para_setup_mpi_test.cpp new file mode 100644 index 00000000000..9c15d42c92b --- /dev/null +++ b/source/source_base/module_parallel/test/para_setup_mpi_test.cpp @@ -0,0 +1,136 @@ +#include "gtest/gtest.h" + +#include "../para_setup.h" +#include "../para_collection.h" +#include "../para_tag.h" +#include "../para_world.h" + +#include + +// These tests run under mpirun -np 4 (see para_setup_mpi_test.sh). + +namespace +{ +int world_rank = -1; +int world_size = -1; +} + +// Single image: the esolver domain wraps the whole world and the +// cross-image domain is absent (same convention as KP_WORLD at kpar == 1). +TEST(ParaSetupMpiTest, SingleImage) +{ + Parallel::ParaCollection worlds + = Parallel::setup_para_worlds(world_size, world_rank, /*nimage=*/1, + /*bndpar=*/1, /*kpar=*/1, /*diag_np=*/1); + + const Parallel::ParaWorld& esolver = worlds.find(Parallel::ParaTag::esolver); + EXPECT_TRUE(esolver.valid()); + EXPECT_EQ(esolver.size(), world_size); + EXPECT_EQ(esolver.rank(), world_rank); + + const Parallel::ParaWorld& images = worlds.find(Parallel::ParaTag::images); + EXPECT_FALSE(images.valid()); + + // All solver domains span the full world in a single-image run. + EXPECT_EQ(worlds.find(Parallel::ParaTag::pw).size(), world_size); + EXPECT_FALSE(worlds.find(Parallel::ParaTag::kmesh).valid()); + EXPECT_EQ(worlds.find(Parallel::ParaTag::diag).size(), world_size); + EXPECT_EQ(worlds.find(Parallel::ParaTag::rgrid).size(), world_size); +} + +// Two images on 4 ranks: each esolver owns 2 ranks; the images domain +// connects corresponding ranks (rank_in_esolver) across the two images. +TEST(ParaSetupMpiTest, TwoImages) +{ + if (world_size < 4) + { + GTEST_SKIP() << "requires 4 MPI ranks (run via para_setup_mpi_test.sh)"; + } + const int nimage = 2; + int image_id = 0; + int rank_in_esolver = 0; + int esolver_size = 0; + Parallel::ParaWorld esolver_world = Parallel::ParaWorld::make_mpi("esolver", MPI_COMM_NULL); + Parallel::ParaWorld images_world = Parallel::ParaWorld::make_mpi("images", MPI_COMM_NULL); + + Parallel::split_images(world_size, world_rank, nimage, + image_id, rank_in_esolver, esolver_size, + esolver_world, images_world); + + EXPECT_EQ(esolver_size, world_size / nimage); + EXPECT_EQ(image_id, world_rank / (world_size / nimage)); + EXPECT_EQ(rank_in_esolver, world_rank % (world_size / nimage)); + + EXPECT_TRUE(esolver_world.valid()); + EXPECT_EQ(esolver_world.size(), world_size / nimage); + EXPECT_EQ(esolver_world.rank(), rank_in_esolver); + + // Even split: the inter-image domain exists and contains one rank + // per image, i.e. its size equals nimage. + EXPECT_TRUE(images_world.valid()); + EXPECT_EQ(images_world.size(), nimage); + EXPECT_EQ(images_world.rank(), image_id); +} + +// Full hierarchy with two images: every solver domain must be derived +// from the esolver domain, so its size never exceeds esolver_size. +TEST(ParaSetupMpiTest, TwoImagesFullHierarchy) +{ + if (world_size < 4) + { + GTEST_SKIP() << "requires 4 MPI ranks (run via para_setup_mpi_test.sh)"; + } + Parallel::ParaCollection worlds + = Parallel::setup_para_worlds(world_size, world_rank, /*nimage=*/2, + /*bndpar=*/1, /*kpar=*/1, /*diag_np=*/1); + + const Parallel::ParaWorld& esolver = worlds.find(Parallel::ParaTag::esolver); + EXPECT_EQ(esolver.size(), world_size / 2); + + const Parallel::ParaWorld& images = worlds.find(Parallel::ParaTag::images); + EXPECT_TRUE(images.valid()); + EXPECT_EQ(images.size(), 2); + + // Domains inside one esolver never see ranks of the other image. + EXPECT_EQ(worlds.find(Parallel::ParaTag::pw).size(), world_size / 2); + EXPECT_EQ(worlds.find(Parallel::ParaTag::bsame_kdiff).size(), world_size / 2); + EXPECT_EQ(worlds.find(Parallel::ParaTag::diag).size(), world_size / 2); + EXPECT_EQ(worlds.find(Parallel::ParaTag::rgrid).size(), world_size / 2); +} + +// k-parallelism inside an image: with 2 images and kpar=2 each pool +// contains one rank, and the inter-pool domain has size kpar = 2. +TEST(ParaSetupMpiTest, TwoImagesWithKpar) +{ + if (world_size < 4) + { + GTEST_SKIP() << "requires 4 MPI ranks (run via para_setup_mpi_test.sh)"; + } + Parallel::ParaCollection worlds + = Parallel::setup_para_worlds(world_size, world_rank, /*nimage=*/2, + /*bndpar=*/1, /*kpar=*/2, /*diag_np=*/1); + + const Parallel::ParaWorld& esolver = worlds.find(Parallel::ParaTag::esolver); + EXPECT_EQ(esolver.size(), 2); + + const Parallel::ParaWorld& pw = worlds.find(Parallel::ParaTag::pw); + EXPECT_TRUE(pw.valid()); + EXPECT_EQ(pw.size(), 1); // 2 ranks per image / kpar 2 + + const Parallel::ParaWorld& kmesh = worlds.find(Parallel::ParaTag::kmesh); + EXPECT_TRUE(kmesh.valid()); + EXPECT_EQ(kmesh.size(), 2); // one corresponding rank per pool +} + +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + MPI_Comm_size(MPI_COMM_WORLD, &world_size); + + testing::InitGoogleTest(&argc, argv); + const int result = RUN_ALL_TESTS(); + + MPI_Finalize(); + return result; +} diff --git a/source/source_base/module_parallel/test/para_setup_mpi_test.sh b/source/source_base/module_parallel/test/para_setup_mpi_test.sh new file mode 100755 index 00000000000..a122b9ed8cb --- /dev/null +++ b/source/source_base/module_parallel/test/para_setup_mpi_test.sh @@ -0,0 +1,18 @@ +#!/bin/bash -e + +np=`cat /proc/cpuinfo | grep "cpu cores" | uniq| awk '{print $NF}'` +echo "nprocs in this machine is $np" + +for i in 4;do + if [[ $i -gt $np ]];then + continue + fi + echo "TEST in parallel, nprocs=$i" + mpirun -np $i ./MODULE_BASE_para_setup_mpi + if [[ $? -ne 0 ]]; then + echo -e "\e[1;33m [ FAILED ] \e[0m"\ + "execute UT with $i cores error." + exit 1 + fi + break +done diff --git a/source/source_base/module_parallel/test/para_setup_test.cpp b/source/source_base/module_parallel/test/para_setup_test.cpp new file mode 100644 index 00000000000..2aa2c77e762 --- /dev/null +++ b/source/source_base/module_parallel/test/para_setup_test.cpp @@ -0,0 +1,54 @@ +#include "gtest/gtest.h" + +#include "../para_setup.h" +#include "../para_world.h" +#include "../para_collection.h" + +TEST(ParaSetupTest, DivideMpiGroupsSerial) +{ + int procs_in_group, my_group, rank_in_group; + Parallel::divide_mpi_groups(1, 1, 0, false, procs_in_group, my_group, rank_in_group); + EXPECT_EQ(procs_in_group, 1); + EXPECT_EQ(my_group, 0); + EXPECT_EQ(rank_in_group, 0); +} + +TEST(ParaSetupTest, DivideMpiGroupsEven) +{ + // 8 procs, 4 groups -> 2 procs per group + int procs_in_group, my_group, rank_in_group; + Parallel::divide_mpi_groups(8, 4, 3, true, procs_in_group, my_group, rank_in_group); + EXPECT_EQ(procs_in_group, 2); + EXPECT_EQ(my_group, 1); + EXPECT_EQ(rank_in_group, 1); +} + +TEST(ParaSetupTest, DivideMpiGroupsUneven) +{ + // 7 procs, 3 groups -> group 0,1 have 3 procs, group 2 has 1 + // rank 0-5 -> group 0,1 (3 procs each) + // rank 6 -> group 2 (1 proc) + int procs_in_group, my_group, rank_in_group; + + // rank 2: first group (procs_in_group+1=3), 2/3=0, 2%3=2 + Parallel::divide_mpi_groups(7, 3, 2, false, procs_in_group, my_group, rank_in_group); + EXPECT_EQ(procs_in_group, 3); + EXPECT_EQ(my_group, 0); + EXPECT_EQ(rank_in_group, 2); + + // rank 6: (6-1)/2=2, (6-1)%2=1 -> wait, extra_procs=1, procs_in_group=2 + // rank 6 >= 1*3 = 3, so: (6-1)/2=2, (6-1)%2=1 + Parallel::divide_mpi_groups(7, 3, 6, false, procs_in_group, my_group, rank_in_group); + EXPECT_EQ(procs_in_group, 2); + EXPECT_EQ(my_group, 2); + EXPECT_EQ(rank_in_group, 1); +} + +TEST(ParaSetupTest, DivideMpiGroupsRank0) +{ + int procs_in_group, my_group, rank_in_group; + Parallel::divide_mpi_groups(12, 4, 0, true, procs_in_group, my_group, rank_in_group); + EXPECT_EQ(procs_in_group, 3); + EXPECT_EQ(my_group, 0); + EXPECT_EQ(rank_in_group, 0); +} diff --git a/source/source_base/module_parallel/test/para_world_mpi_test.cpp b/source/source_base/module_parallel/test/para_world_mpi_test.cpp new file mode 100644 index 00000000000..c0eb9f28d8d --- /dev/null +++ b/source/source_base/module_parallel/test/para_world_mpi_test.cpp @@ -0,0 +1,30 @@ +#include "gtest/gtest.h" + +#include "../para_world.h" + +TEST(ParaWorldMpiTest, SerialFactory) +{ + const Parallel::ParaWorld world = Parallel::ParaWorld::serial("pw"); + EXPECT_EQ(world.tag(), "pw"); + EXPECT_EQ(world.rank(), 0); + EXPECT_EQ(world.size(), 1); + EXPECT_TRUE(world.valid()); +} + +TEST(ParaWorldMpiTest, WrapCommunicator) +{ + const Parallel::ParaWorld world = Parallel::ParaWorld::serial("pw"); + EXPECT_EQ(world.tag(), "pw"); + EXPECT_TRUE(world.valid()); + EXPECT_EQ(world.rank(), 0); + EXPECT_EQ(world.size(), 1); +} + +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + testing::InitGoogleTest(&argc, argv); + const int result = RUN_ALL_TESTS(); + MPI_Finalize(); + return result; +} diff --git a/source/source_base/module_parallel/test/para_world_mpi_test.sh b/source/source_base/module_parallel/test/para_world_mpi_test.sh new file mode 100644 index 00000000000..acfba19a8c1 --- /dev/null +++ b/source/source_base/module_parallel/test/para_world_mpi_test.sh @@ -0,0 +1,18 @@ +#!/bin/bash -e + +np=`cat /proc/cpuinfo | grep "cpu cores" | uniq| awk '{print $NF}'` +echo "nprocs in this machine is $np" + +for i in 4;do + if [[ $i -gt $np ]];then + continue + fi + echo "TEST in parallel, nprocs=$i" + mpirun -np $i ./MODULE_BASE_para_world_mpi + if [[ $? -ne 0 ]]; then + echo -e "\e[1;33m [ FAILED ] \e[0m"\ + "execute UT with $i cores error." + exit 1 + fi + break +done diff --git a/source/source_base/module_parallel/test/para_world_test.cpp b/source/source_base/module_parallel/test/para_world_test.cpp new file mode 100644 index 00000000000..f3aa540cfcf --- /dev/null +++ b/source/source_base/module_parallel/test/para_world_test.cpp @@ -0,0 +1,21 @@ +#include "gtest/gtest.h" + +#include "../para_world.h" + +TEST(ParaWorldTest, SerialFactory) +{ + const Parallel::ParaWorld world = Parallel::ParaWorld::serial("pw"); + EXPECT_EQ(world.tag(), "pw"); + EXPECT_EQ(world.rank(), 0); + EXPECT_EQ(world.size(), 1); + EXPECT_TRUE(world.valid()); +} + +TEST(ParaWorldTest, EmptyTagIsInvalid) +{ + const Parallel::ParaWorld world = Parallel::ParaWorld::serial(""); + EXPECT_TRUE(world.tag().empty()); + EXPECT_EQ(world.rank(), 0); + EXPECT_EQ(world.size(), 1); + EXPECT_FALSE(world.valid()); +} diff --git a/source/source_base/parallel_2d.h b/source/source_base/parallel_2d.h index 2d89d12c4bf..08fb5f25f2e 100644 --- a/source/source_base/parallel_2d.h +++ b/source/source_base/parallel_2d.h @@ -5,7 +5,9 @@ #include #include -#include "source_base/parallel_comm.h" +#ifdef __MPI +#include +#endif /// @brief This class packs the basic information of /// 2D-block-cyclic parallel distribution of an arbitrary matrix. diff --git a/source/source_base/parallel_common.cpp b/source/source_base/parallel_common.cpp index a57f9d04874..c5ee4cecf2f 100644 --- a/source/source_base/parallel_common.cpp +++ b/source/source_base/parallel_common.cpp @@ -1,12 +1,25 @@ #include "parallel_common.h" +#include "source_base/parallel_reduce.h" + #ifdef __MPI #include #endif -#include +namespace Parallel_Common +{ + +#ifdef __MPI +/// Broadcast a trivially-copyable buffer of type T on MPI_COMM_WORLD from +/// rank 0. This is the single implementation behind all bcast_* wrappers. +template +static void bcast_world_impl(T* object, const int n) +{ + MPI_Bcast(object, n, Parallel_Reduce::MPI_Type::value, 0, MPI_COMM_WORLD); +} +#endif -void Parallel_Common::bcast_string(std::string& object) // Peize Lin fix bug 2019-03-18 +void bcast_string(std::string& object) // Peize Lin fix bug 2019-03-18 { #ifdef __MPI int size = object.size(); @@ -25,7 +38,7 @@ void Parallel_Common::bcast_string(std::string& object) // Peize Lin fix bug 201 return; } -void Parallel_Common::bcast_string(std::string* object, const int n) // Peize Lin fix bug 2019-03-18 +void bcast_string(std::string* object, const int n) // Peize Lin fix bug 2019-03-18 { #ifdef __MPI for (int i = 0; i < n; i++) @@ -34,65 +47,65 @@ void Parallel_Common::bcast_string(std::string* object, const int n) // Peize Li return; } -void Parallel_Common::bcast_complex_double(std::complex& object) +void bcast_complex_double(std::complex& object) { #ifdef __MPI - MPI_Bcast(&object, 1, MPI_DOUBLE_COMPLEX, 0, MPI_COMM_WORLD); + bcast_world_impl(&object, 1); #endif } -void Parallel_Common::bcast_complex_double(std::complex* object, const int n) +void bcast_complex_double(std::complex* object, const int n) { #ifdef __MPI - MPI_Bcast(object, n, MPI_DOUBLE_COMPLEX, 0, MPI_COMM_WORLD); + bcast_world_impl(object, n); #endif } -void Parallel_Common::bcast_double(double& object) +void bcast_double(double& object) { #ifdef __MPI - MPI_Bcast(&object, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD); + bcast_world_impl(&object, 1); #endif } -void Parallel_Common::bcast_double(double* object, const int n) +void bcast_double(double* object, const int n) { #ifdef __MPI - MPI_Bcast(object, n, MPI_DOUBLE, 0, MPI_COMM_WORLD); + bcast_world_impl(object, n); #endif } -void Parallel_Common::bcast_int(int& object) +void bcast_int(int& object) { #ifdef __MPI - MPI_Bcast(&object, 1, MPI_INT, 0, MPI_COMM_WORLD); + bcast_world_impl(&object, 1); #endif } -void Parallel_Common::bcast_int(int* object, const int n) +void bcast_int(int* object, const int n) { #ifdef __MPI - MPI_Bcast(object, n, MPI_INT, 0, MPI_COMM_WORLD); + bcast_world_impl(object, n); #endif } -void Parallel_Common::bcast_bool(bool& object) +void bcast_bool(bool& object) { #ifdef __MPI int swap = object; int my_rank; MPI_Comm_rank(MPI_COMM_WORLD, &my_rank); - if (my_rank == 0) - swap = object; MPI_Bcast(&swap, 1, MPI_INT, 0, MPI_COMM_WORLD); if (my_rank != 0) object = static_cast(swap); #endif } -void Parallel_Common::bcast_char(char* object, const int n) +void bcast_char(char* object, const int n) { #ifdef __MPI MPI_Bcast(object, n, MPI_CHAR, 0, MPI_COMM_WORLD); #endif } + +} // namespace Parallel_Common diff --git a/source/source_base/parallel_global.cpp b/source/source_base/parallel_global.cpp index df04461b30b..697b7f1f702 100644 --- a/source/source_base/parallel_global.cpp +++ b/source/source_base/parallel_global.cpp @@ -12,8 +12,6 @@ #endif #include "source_base/global_function.h" -#include "source_base/parallel_common.h" -#include "source_base/parallel_reduce.h" #include "source_base/global_variable.h" #include "source_base/tool_quit.h" diff --git a/source/source_base/parallel_grid.cpp b/source/source_base/parallel_grid.cpp index bebf8074238..c8a058e52e9 100644 --- a/source/source_base/parallel_grid.cpp +++ b/source/source_base/parallel_grid.cpp @@ -103,11 +103,10 @@ void Parallel_Grid::z_distribution() { assert(!this->numz.empty()); - int* startp = new int[GlobalV::KPAR]; + std::vector startp(GlobalV::KPAR); startp[0] = 0; for (int ip = 0; ip < GlobalV::KPAR; ip++) { - // GlobalV::ofs_running << "\n now POOL=" << ip; const int nproc = nproc_in_pool[ip]; if (ip > 0) @@ -122,11 +121,6 @@ void Parallel_Grid::z_distribution() numz[ip][proc] += bz; } - // for(int proc=0; procstartz[GlobalV::MY_POOL][GlobalV::RANK_IN_POOL]; const int proc = this->whichpro[GlobalV::MY_POOL][iz]; if (GlobalV::MY_POOL == 0) { - // case 1: the first part of rho in processor 0. - // and send zpeice to to other pools. - if (proc == 0 && GlobalV::MY_RANK == 0) + // case 1: the first part of rho in processor 0, + // and send zpiece to the other pools. + if (proc == 0 && rank_in_comm == 0) { for (int ir = 0; ir < ncxy; ir++) { @@ -308,41 +297,39 @@ void Parallel_Grid::zpiece_to_all(double* zpiece, const int& iz, double* rho) co } for (int ipool = 1; ipool < GlobalV::KPAR; ipool++) { - MPI_Send(zpiece, ncxy, MPI_DOUBLE, this->whichpro[ipool][iz], iz, MPI_COMM_WORLD); + MPI_Send(zpiece, ncxy, MPI_DOUBLE, this->whichpro[ipool][iz], iz, comm); } } - // case 2: processor n (n!=0) receive rho from processor 0. - // and the receive tag is iz. + // case 2: processor n (n!=0) receives rho from processor 0. + // The receive tag is iz. else if (proc == GlobalV::RANK_IN_POOL) { - MPI_Recv(zpiece, ncxy, MPI_DOUBLE, 0, iz, MPI_COMM_WORLD, &ierror); + MPI_Recv(zpiece, ncxy, MPI_DOUBLE, 0, iz, comm, &ierror); for (int ir = 0; ir < ncxy; ir++) { rho[ir * nczp + znow] = zpiece[ir]; } } - // case 2: > first part rho: processor 0 send the rho - // to all pools. The tag is iz, because processor may - // send more than once, and the only tag to distinguish - // them is iz. + // case 3: pool root (not owning iz) forwards rho to all pools. + // The tag is iz, because a processor may send more than once, and + // the only tag to distinguish them is iz. else if (GlobalV::RANK_IN_POOL == 0) { for (int ipool = 0; ipool < GlobalV::KPAR; ipool++) { - MPI_Send(zpiece, ncxy, MPI_DOUBLE, this->whichpro[ipool][iz], iz, MPI_COMM_WORLD); + MPI_Send(zpiece, ncxy, MPI_DOUBLE, this->whichpro[ipool][iz], iz, comm); } } } // GlobalV::MY_POOL == 0 else { - // GlobalV::ofs_running << "\n Receive charge density iz=" << iz << std::endl; - // the processors in other pools always receive rho from - // processor 0. the tag is 'iz' - if (proc == GlobalV::MY_RANK) + // The processors in other pools always receive rho from + // processor 0. The tag is 'iz'. + if (proc == rank_in_comm) { - MPI_Recv(zpiece, ncxy, MPI_DOUBLE, 0, iz, MPI_COMM_WORLD, &ierror); + MPI_Recv(zpiece, ncxy, MPI_DOUBLE, 0, iz, comm, &ierror); for (int ir = 0; ir < ncxy; ir++) { rho[ir * nczp + znow] = zpiece[ir]; @@ -350,79 +337,11 @@ void Parallel_Grid::zpiece_to_all(double* zpiece, const int& iz, double* rho) co } } - // GlobalV::ofs_running << "\n iz = " << iz << " Done."; return; } #endif #ifdef __MPI -void Parallel_Grid::zpiece_to_stogroup(double* zpiece, const int& iz, double* rho) const -{ - assert(!this->numz.empty()); - // TITLE("Parallel_Grid","zpiece_to_all"); - MPI_Status ierror; - - const int znow = iz - this->startz[GlobalV::MY_POOL][GlobalV::RANK_IN_POOL]; - const int proc = this->whichpro[GlobalV::MY_POOL][iz]; - - if (GlobalV::MY_POOL == 0) - { - // case 1: the first part of rho in processor 0. - // and send zpeice to to other pools. - if (proc == 0 && GlobalV::RANK_IN_BPGROUP == 0) - { - for (int ir = 0; ir < ncxy; ir++) - { - rho[ir * nczp + znow] = zpiece[ir]; - } - for (int ipool = 1; ipool < GlobalV::KPAR; ipool++) - { - MPI_Send(zpiece, ncxy, MPI_DOUBLE, this->whichpro[ipool][iz], iz, INT_BGROUP); - } - } - - // case 2: processor n (n!=0) receive rho from processor 0. - // and the receive tag is iz. - else if (proc == GlobalV::RANK_IN_POOL) - { - MPI_Recv(zpiece, ncxy, MPI_DOUBLE, 0, iz, INT_BGROUP, &ierror); - for (int ir = 0; ir < ncxy; ir++) - { - rho[ir * nczp + znow] = zpiece[ir]; - } - } - - // case 2: > first part rho: processor 0 send the rho - // to all pools. The tag is iz, because processor may - // send more than once, and the only tag to distinguish - // them is iz. - else if (GlobalV::RANK_IN_POOL == 0) - { - for (int ipool = 0; ipool < GlobalV::KPAR; ipool++) - { - MPI_Send(zpiece, ncxy, MPI_DOUBLE, this->whichpro[ipool][iz], iz, INT_BGROUP); - } - } - } // MY_POOL == 0 - else - { - // ofs_running << "\n Receive charge density iz=" << iz << endl; - // the processors in other pools always receive rho from - // processor 0. the tag is 'iz' - if (proc == GlobalV::RANK_IN_BPGROUP) - { - MPI_Recv(zpiece, ncxy, MPI_DOUBLE, 0, iz, INT_BGROUP, &ierror); - for (int ir = 0; ir < ncxy; ir++) - { - rho[ir * nczp + znow] = zpiece[ir]; - } - } - } - - // ofs_running << "\n iz = " << iz << " Done."; - return; -} - // Taoni modified on 2026-08-21, fixed BPCG out_chg MPI_ERR_RANK void Parallel_Grid::reduce(double* rhotot, const double* const rhoin, const bool reduce_all_pool) const { diff --git a/source/source_base/parallel_grid.h b/source/source_base/parallel_grid.h index 730c24e1f93..8cc7592384c 100644 --- a/source/source_base/parallel_grid.h +++ b/source/source_base/parallel_grid.h @@ -48,8 +48,7 @@ class Parallel_Grid void z_distribution(void); #ifdef __MPI - void zpiece_to_all(double* zpiece, const int& iz, double* rho) const; - void zpiece_to_stogroup(double* zpiece, const int& iz, double* rho) const; //qainrui add for sto-dft 2021-7-21 + void zpiece_distribute(double* zpiece, const int& iz, double* rho, const bool is_sdft) const; #endif std::vector nproc_in_pool; diff --git a/source/source_base/parallel_reduce.h b/source/source_base/parallel_reduce.h index e3210f41d9c..620620f9827 100644 --- a/source/source_base/parallel_reduce.h +++ b/source/source_base/parallel_reduce.h @@ -80,8 +80,6 @@ void reduce_double_allpool(const int& npool, const int& nproc_in_pool, double* o void gather_int_all(int& v, int* all); -bool check_if_equal(double& v); // mohan add 2009-11-11 - } // namespace Parallel_Reduce #endif diff --git a/source/source_base/test_parallel/CMakeLists.txt b/source/source_base/test_parallel/CMakeLists.txt index 873518c14af..94a7c6afd82 100644 --- a/source/source_base/test_parallel/CMakeLists.txt +++ b/source/source_base/test_parallel/CMakeLists.txt @@ -1,7 +1,7 @@ AddTest( TARGET MODULE_BASE_ParaCommon LIBS parameter MPI::MPI_CXX - SOURCES parallel_common_test.cpp ../global_variable.cpp ../parallel_common.cpp + SOURCES parallel_common_test.cpp ../global_variable.cpp ../parallel_common.cpp ../parallel_reduce.cpp ../parallel_comm.cpp ../parallel_global.cpp ../tool_quit.cpp ../global_file.cpp ../global_function.cpp ../memory_recorder.cpp ../timer.cpp ) AddTest( diff --git a/source/source_esolver/esolver_fp.cpp b/source/source_esolver/esolver_fp.cpp index b0114402c67..a578ec2d06d 100644 --- a/source/source_esolver/esolver_fp.cpp +++ b/source/source_esolver/esolver_fp.cpp @@ -8,13 +8,16 @@ #include "source_hamilt/module_vdw/vdw.h" #include "source_io/module_output/output_log.h" #include "source_io/module_output/print_info.h" -#include "source_io/module_chgpot/rhog_io.h" +#include "source_estate/rhog_io.h" #include "source_io/module_parameter/parameter.h" #include "source_pw/module_pwdft/setup_pwrho.h" // mohan 20251005 #include "source_hamilt/module_xc/xc_functional.h" // mohan 20251005 #include "source_io/module_ctrl/ctrl_output_fp.h" -#include "source_io/module_chgpot/write_init.h" // write_chg_init, write_pot_init +#include "source_estate/write_init.h" // write_chg_init, write_pot_init +#include "source_base/module_parallel/para_world.h" +#include "source_base/module_parallel/para_tag.h" +#include "source_base/module_parallel/para_bridge.h" namespace ModuleESolver { @@ -235,15 +238,20 @@ void ESolver_FP::iter_finish(UnitCell& ucell, const int istep, int& iter, bool& { this->pw_rhod->real2recip(this->chr.rho_save[is], this->chr.rhog_save[is]); } - ModuleIO::write_rhog(PARAM.globalv.global_out_dir + this->inp_->suffix + "-CHARGE-DENSITY.restart", - PARAM.globalv.gamma_only_pw, - this->pw_rhod, - this->inp_->nspin, - ucell.GT, - this->chr.rhog_save, - GlobalV::MY_POOL, - GlobalV::RANK_IN_POOL, - GlobalV::NPROC_IN_POOL); + // Temporary bridge: use factory until ParaCollection is wired into driver. + Parallel::ParaWorld pw_world = Parallel::make_pw_world(); + // Only pool 0 writes the rhog file (rhog is identical across pools). + if (GlobalV::MY_POOL == 0) + { + elecstate::write_rhog(PARAM.globalv.global_out_dir + this->inp_->suffix + "-CHARGE-DENSITY.restart", + PARAM.globalv.gamma_only_pw, + this->pw_rhod, + this->inp_->nspin, + ucell.GT, + this->chr.rhog_save, + pw_world, + &GlobalV::ofs_warning); + } if (XC_Functional::get_ked_flag()) { @@ -254,15 +262,17 @@ void ESolver_FP::iter_finish(UnitCell& ucell, const int istep, int& iter, bool& kin_g.push_back(kin_g_space.data() + is * this->chr.ngmc); this->pw_rhod->real2recip(this->chr.kin_r_save[is], kin_g[is]); } - ModuleIO::write_rhog(PARAM.globalv.global_out_dir + this->inp_->suffix + "-TAU-DENSITY.restart", - PARAM.globalv.gamma_only_pw, - this->pw_rhod, - this->inp_->nspin, - ucell.GT, - kin_g.data(), - GlobalV::MY_POOL, - GlobalV::RANK_IN_POOL, - GlobalV::NPROC_IN_POOL); + if (GlobalV::MY_POOL == 0) + { + elecstate::write_rhog(PARAM.globalv.global_out_dir + this->inp_->suffix + "-TAU-DENSITY.restart", + PARAM.globalv.gamma_only_pw, + this->pw_rhod, + this->inp_->nspin, + ucell.GT, + kin_g.data(), + pw_world, + &GlobalV::ofs_warning); + } } } } diff --git a/source/source_esolver/esolver_sdft_pw.cpp b/source/source_esolver/esolver_sdft_pw.cpp index 808fbbe035c..d5850d8f6f9 100644 --- a/source/source_esolver/esolver_sdft_pw.cpp +++ b/source/source_esolver/esolver_sdft_pw.cpp @@ -2,6 +2,7 @@ #include "source_base/global_variable.h" #include "source_base/memory_recorder.h" +#include "source_base/parallel_comm.h" #include "source_estate/module_charge/symm_rho.h" #include "source_hsolver/diago_iter_assist.h" #include "source_hsolver/diago_params.h" diff --git a/source/source_estate/CMakeLists.txt b/source/source_estate/CMakeLists.txt index b747b07d85c..c29e7bac9c7 100644 --- a/source/source_estate/CMakeLists.txt +++ b/source/source_estate/CMakeLists.txt @@ -40,6 +40,9 @@ list(APPEND objects module_charge/charge_mixing_uspp.cpp module_charge/symm_rho.cpp module_charge/symm_rhog.cpp + rhog_io.cpp + write_elecstat_pot.cpp + write_init.cpp fp_energy.cpp occupy.cpp occ_matrix.cpp @@ -72,7 +75,6 @@ endif() if(BUILD_TESTING) if(ENABLE_MPI) add_subdirectory(test) - add_subdirectory(test_mpi) endif() endif() diff --git a/source/source_estate/init_scf.cpp b/source/source_estate/init_scf.cpp index e89e706e342..4e5699e8074 100644 --- a/source/source_estate/init_scf.cpp +++ b/source/source_estate/init_scf.cpp @@ -1,5 +1,5 @@ #include "elecstate.h" -#include "source_io/module_chgpot/write_init.h" +#include "source_estate/write_init.h" namespace elecstate { diff --git a/source/source_estate/module_charge/charge_init.cpp b/source/source_estate/module_charge/charge_init.cpp index 30ab90598f5..672d800f718 100644 --- a/source/source_estate/module_charge/charge_init.cpp +++ b/source/source_estate/module_charge/charge_init.cpp @@ -14,11 +14,14 @@ #include "source_cell/magnetism.h" #include "source_base/parallel_grid.h" #include "source_io/module_output/cube_io.h" -#include "source_io/module_chgpot/rhog_io.h" +#include "source_estate/rhog_io.h" #include "source_io/module_wf/read_wf2rho_pw.h" #include "source_io/module_restart/restart.h" #include "source_hamilt/module_xc/xc_functional.h" #include "source_cell/klist.h" +#include "source_base/module_parallel/para_world.h" +#include "source_base/module_parallel/para_tag.h" +#include "source_base/module_parallel/para_bridge.h" void Charge::init_rho(const UnitCell& ucell, const Parallel_Grid& pgrid, @@ -50,7 +53,9 @@ void Charge::init_rho(const UnitCell& ucell, // liuyu 2023-12-05 std::stringstream binary; binary << PARAM.globalv.global_readin_dir << PARAM.inp.suffix + "-CHARGE-DENSITY.restart"; - if (ModuleIO::read_rhog(binary.str(), rhopw, rhog)) + // Temporary bridge: use factory until ParaCollection is wired into driver. + Parallel::ParaWorld pw_world = Parallel::make_pw_world(); + if (elecstate::read_rhog(binary.str(), rhopw, nspin, rhog, pw_world, &GlobalV::ofs_warning)) { GlobalV::ofs_running << " Read electron density from file: " << binary.str() << std::endl; for (int is = 0; is < nspin; ++is) @@ -147,7 +152,7 @@ void Charge::init_rho(const UnitCell& ucell, std::stringstream binary; binary << PARAM.globalv.global_readin_dir << PARAM.inp.suffix + "-TAU-DENSITY.restart"; - if (ModuleIO::read_rhog(binary.str(), rhopw, kin_g.data())) + if (elecstate::read_rhog(binary.str(), rhopw, nspin, kin_g.data(), pw_world, &GlobalV::ofs_warning)) { GlobalV::ofs_running << " Read in the kinetic energy density: " << binary.str() << std::endl; for (int is = 0; is < nspin; ++is) diff --git a/source/source_estate/module_charge/chgmixing.cpp b/source/source_estate/module_charge/chgmixing.cpp index 4886e9f5e21..670636c541a 100644 --- a/source/source_estate/module_charge/chgmixing.cpp +++ b/source/source_estate/module_charge/chgmixing.cpp @@ -1,4 +1,6 @@ #include "source_estate/module_charge/chgmixing.h" + +#include "source_base/parallel_comm.h" #include "source_estate/update_pot.h" #include "source_lcao/module_dftu/dftu_nao.h" #include "source_lcao/module_deltaspin/spin_constrain.h" diff --git a/source/source_estate/rhog_io.cpp b/source/source_estate/rhog_io.cpp new file mode 100644 index 00000000000..3412a20133d --- /dev/null +++ b/source/source_estate/rhog_io.cpp @@ -0,0 +1,374 @@ +#include "source_base/module_out/binstream.h" +#include "source_base/vector3.h" +#include "source_base/module_parallel/para_mpi_func.h" +#include "rhog_io.h" +#include +#include +#include + +namespace +{ +inline void warn(std::ostream* os, + const Parallel::ParaWorld& pw_world, + const std::string& file, + const std::string& desc) +{ + if (pw_world.rank() == 0 && os != nullptr) + { + *os << " " << file << " warning : " << desc << std::endl; + } +} +} // namespace + +bool elecstate::read_rhog(const std::string& filename, + const ModulePW::PW_Basis* pw_rhod, + const int nspin, + std::complex** rhog, + const Parallel::ParaWorld& pw_world, + std::ostream* os_warning) +{ + if (pw_rhod == nullptr) + { + warn(os_warning, pw_world, "elecstate::read_rhog", "pw_rhod is null"); + return false; + } + if (rhog == nullptr) + { + warn(os_warning, pw_world, "elecstate::read_rhog", "rhog is null"); + return false; + } + if (nspin != 1 && nspin != 2 && nspin != 4) + { + warn(os_warning, pw_world, "elecstate::read_rhog", "nspin must be 1, 2, or 4"); + return false; + } + if (pw_rhod->nx <= 0 || pw_rhod->ny <= 0 || pw_rhod->nz <= 0) + { + warn(os_warning, pw_world, "elecstate::read_rhog", "PW_Basis grid dimensions must be positive"); + return false; + } + + const int nx = pw_rhod->nx; + const int ny = pw_rhod->ny; + const int nz = pw_rhod->nz; + + Binstream ifs; + bool error = false; + int gamma_only_in = 0; + int npwtot_in = 0; + int nspin_in = 0; + int size = 0; + double b1[3], b2[3], b3[3]; + + if (pw_world.rank() == 0) + { + ifs.open(filename, "r"); + if (!ifs) + { + error = true; + } + } + + Parallel::bcast_bool(error, pw_world); + + if (error) + { + warn(os_warning, pw_world, "elecstate::read_rhog", "Can't open file " + filename); + return false; + } + + if (pw_world.rank() == 0) + { + ifs >> size >> gamma_only_in >> npwtot_in >> nspin_in >> size; + ifs >> size >> b1[0] >> b1[1] >> b1[2] >> b2[0] >> b2[1] >> b2[2] >> b3[0] >> b3[1] >> b3[2] >> size; + if (gamma_only_in != pw_rhod->gamma_only) + { + // there is a treatment that can transform between gamma_only and non-gamma_only + // however, it is not implemented here + error = true; + ifs.close(); + } + if (npwtot_in > pw_rhod->npwtot) + { + warn(os_warning, pw_world, "elecstate::read_rhog", "some planewaves in file are not used"); + } + else if (npwtot_in < pw_rhod->npwtot) + { + warn(os_warning, pw_world, "elecstate::read_rhog", "some planewaves in file are missing"); + } + if (nspin_in < nspin) + { + warn(os_warning, pw_world, "elecstate::read_rhog", "some spin channels in file are missing"); + } + } + + Parallel::bcast_bool(error, pw_world); + + if (error) + { + warn(os_warning, pw_world, "elecstate::read_rhog", "gamma_only read from file is inconsistent with INPUT"); + return false; + } + + Parallel::bcast_int(gamma_only_in, pw_world); + Parallel::bcast_int(npwtot_in, pw_world); + Parallel::bcast_int(nspin_in, pw_world); + Parallel::bcast_double(b1, 3, pw_world); + Parallel::bcast_double(b2, 3, pw_world); + Parallel::bcast_double(b3, 3, pw_world); + + std::vector miller(npwtot_in * 3); + // once use ModuleBase::Vector3, it is highly bug-prone to assume the memory layout of the class. + // The x, y and z of Vector3 will not always to be contiguous. + // Instead, a relatively safe choice is to use std::vector, the memory layout is assumed + // to be npwtot_in rows and 3 columns. + if (pw_world.rank() == 0) + { + ifs >> size; + for (int i = 0; i < npwtot_in; ++i) // loop over rows... + { + ifs >> miller[i*3] >> miller[i*3+1] >> miller[i*3+2]; + } + ifs >> size; + } + Parallel::bcast_int(miller.data(), miller.size(), pw_world); + + // set to zero + for (int is = 0; is < nspin; ++is) + { + std::fill(rhog[is], rhog[is] + pw_rhod->npw, std::complex(0.0, 0.0)); + } + // maps ixyz tp ig + std::vector fftixyz2ig(pw_rhod->nxyz, -1); // map isz to ig. + for (int ig = 0; ig < pw_rhod->npw; ++ig) + { + int isz = pw_rhod->ig2isz[ig]; + int iz = isz % nz; + int is = isz / nz; + int ixy = pw_rhod->is2fftixy[is]; + int ixyz = iz + nz * ixy; + fftixyz2ig[ixyz] = ig; + } + std::vector> rhog_in(npwtot_in); + for (int is = 0; is < nspin_in; ++is) + { + if (pw_world.rank() == 0) + { + ifs >> size; + for (int i = 0; i < npwtot_in; ++i) + { + ifs >> rhog_in[i]; + } + ifs >> size; + } + Parallel::bcast_complex(rhog_in.data(), rhog_in.size(), pw_world); + + for (int i = 0; i < npwtot_in; ++i) + { + int ix = miller[i * 3]; + int iy = miller[i * 3 + 1]; + int iz = miller[i * 3 + 2]; + + if (ix <= -int((nx + 1) / 2) || ix >= int(nx / 2) + 1 || iy <= -int((ny + 1) / 2) || iy >= int(ny / 2) + 1 + || iz <= -int((nz + 1) / 2) || iz >= int(nz / 2) + 1) + { + // these planewaves are not used + continue; + } + + if (ix < 0) + ix += nx; + if (iy < 0) + iy += ny; + if (iz < 0) + iz += nz; + int fftixy = iy + pw_rhod->fftny * ix; + if (pw_world.rank() == pw_rhod->fftixy2ip[fftixy]) + { + int fftixyz = iz + nz * fftixy; + int ig = fftixyz2ig[fftixyz]; + rhog[is][ig] = rhog_in[i]; + } + } + + if (nspin_in == 2 && nspin == 4 && is == 1) + { + for (int ig = 0; ig < pw_rhod->npw; ++ig) + { + rhog[3][ig] = rhog[1][ig]; + } + std::fill(rhog[1], rhog[1] + pw_rhod->npw, std::complex(0.0, 0.0)); + std::fill(rhog[2], rhog[2] + pw_rhod->npw, std::complex(0.0, 0.0)); + } + } + + if (pw_world.rank() == 0) + { + ifs.close(); + } + return true; +} + +bool elecstate::write_rhog(const std::string& fchg, + const bool gamma_only, + const ModulePW::PW_Basis* pw_rho, + const int nspin, + const ModuleBase::Matrix3& GT, + std::complex** rhog, + const Parallel::ParaWorld& pw_world, + std::ostream* os_warning) +{ + if (pw_rho == nullptr) + { + warn(os_warning, pw_world, "elecstate::write_rhog", "pw_rho is null"); + return false; + } + if (rhog == nullptr) + { + warn(os_warning, pw_world, "elecstate::write_rhog", "rhog is null"); + return false; + } + if (nspin != 1 && nspin != 2 && nspin != 4) + { + warn(os_warning, pw_world, "elecstate::write_rhog", "nspin must be 1, 2, or 4"); + return false; + } + + // only rank 0 in the domain writes the header; all ranks cooperate + // on sequential writes synchronized by barriers. + const int irank = pw_world.rank(); + const int nrank = pw_world.size(); + + // write the header (by rank 0): gamma_only, ngm_g, nspin + int size = 3; + int ngm_g = pw_rho->npwtot; + int gam = gamma_only; + int nsp = nspin; + + std::ofstream ofs; + Parallel::barrier(pw_world); + + if (irank == 0) + { + ofs.open(fchg, std::ios::binary); + if (!ofs) + { + warn(os_warning, pw_world, "elecstate::write_rhog", "File I/O failure: cannot open file " + fchg); + return false; + } + ofs.write(reinterpret_cast(&size), sizeof(size)); + ofs.write(reinterpret_cast(&gam), sizeof(gam)); + ofs.write(reinterpret_cast(&ngm_g), sizeof(ngm_g)); + ofs.write(reinterpret_cast(&nsp), sizeof(nsp)); + ofs.write(reinterpret_cast(&size), sizeof(size)); + // write the lattice vectors + std::vector b = {GT.e11, GT.e12, GT.e13, GT.e21, GT.e22, GT.e23, GT.e31, GT.e32, GT.e33}; + size = 9; + ofs.write(reinterpret_cast(&size), sizeof(size)); + for (int i = 0; i < 9; ++i) + { + ofs.write(reinterpret_cast(&b[i]), sizeof(b[i])); + } + ofs.write(reinterpret_cast(&size), sizeof(size)); + ofs.close(); + } + Parallel::barrier(pw_world); + Parallel::barrier(pw_world); + + // write the G-vectors in Miller indices + size = 3 * ngm_g; + if (irank == 0) + { + ofs.open(fchg, std::ios::binary | std::ios::app); + ofs.write(reinterpret_cast(&size), sizeof(size)); + ofs.close(); + } + Parallel::barrier(pw_world); + + for (int i = 0; i < nrank; ++i) + { + if (i == irank) + { + ofs.open(fchg, std::ios::binary | std::ios::app); + for (int ig = 0; ig < pw_rho->npw; ++ig) + { + const ModuleBase::Vector3 g = pw_rho->gdirect[ig]; + std::vector miller = {int(g.x), int(g.y), int(g.z)}; + ofs.write(reinterpret_cast(&miller[0]), sizeof(miller[0])); + ofs.write(reinterpret_cast(&miller[1]), sizeof(miller[1])); + ofs.write(reinterpret_cast(&miller[2]), sizeof(miller[2])); + } + ofs.close(); + } + Parallel::barrier(pw_world); + } + + if (irank == 0) + { + ofs.open(fchg, std::ios::binary | std::ios::app); + ofs.write(reinterpret_cast(&size), sizeof(size)); + ofs.close(); + } + Parallel::barrier(pw_world); + + // write the rho(G) values + std::complex sum_check; + size = ngm_g; + for (int ispin = 0; ispin < nspin; ++ispin) + { + if (irank == 0) + { + ofs.open(fchg, std::ios::binary | std::ios::app); + ofs.write(reinterpret_cast(&size), sizeof(size)); + ofs.close(); + } + Parallel::barrier(pw_world); + + for (int i = 0; i < nrank; ++i) + { + if (i == irank) + { + ofs.open(fchg, std::ios::binary | std::ios::app); + sum_check = 0.0; + for (int ig = 0; ig < pw_rho->npw; ++ig) + { + sum_check += rhog[ispin][ig]; + ofs.write(reinterpret_cast(&rhog[ispin][ig]), sizeof(rhog[ispin][ig])); + } + ofs.close(); + } + Parallel::barrier(pw_world); + } + + if (irank == 0) + { + ofs.open(fchg, std::ios::binary | std::ios::app); + ofs.write(reinterpret_cast(&size), sizeof(size)); + ofs.close(); + } + Parallel::barrier(pw_world); + } + return true; +} + +// self-consistency test with the following python code +// import numpy as np + +// with open("rhog_read.txt") as f: +// read = f.readlines() + +// with open("rhog_write.txt") as f: +// write = f.readlines() + +// # convert c++ stype complex number (a,b) to python complex +// def to_complex(s): +// a, b = s.replace("(", "").replace(")", "").split(",") +// return complex(float(a), float(b)) + +// read = [[to_complex(rhog) for rhog in spin.strip().split()] for spin in read] +// write = [[to_complex(rhog) for rhog in spin.strip().split()] for spin in write] + +// diff = np.array(read) - np.array(write) +// print(np.max(np.abs(diff))) +// test system: integrated test 118_PW_CHG_BINARY +// yielding error 5.290000000000175e-11 \ No newline at end of file diff --git a/source/source_io/module_chgpot/rhog_io.h b/source/source_estate/rhog_io.h similarity index 72% rename from source/source_io/module_chgpot/rhog_io.h rename to source/source_estate/rhog_io.h index 9f470b96376..b4b890808c0 100644 --- a/source/source_io/module_chgpot/rhog_io.h +++ b/source/source_estate/rhog_io.h @@ -3,7 +3,9 @@ #include #include +#include #include "source_basis/module_pw/pw_basis.h" +#include "source_base/module_parallel/para_world.h" /** * I/O free function of rho(G) in binary format * Author: YuLiu98, Kirk0830 @@ -40,21 +42,25 @@ * rho */ -namespace ModuleIO +namespace elecstate { -bool read_rhog(const std::string& filename, const ModulePW::PW_Basis* pw_rhod, std::complex** rhog); +bool read_rhog(const std::string& filename, + const ModulePW::PW_Basis* pw_rhod, + const int nspin, + std::complex** rhog, + const Parallel::ParaWorld& pw_world, + std::ostream* os_warning); bool write_rhog(const std::string& fchg, - const bool gamma_only, // from INPUT - const ModulePW::PW_Basis* pw_rho, // pw_rho in runtime - const int nspin, // GlobalV - const ModuleBase::Matrix3& GT, // from UnitCell, useful for calculating the miller + const bool gamma_only, + const ModulePW::PW_Basis* pw_rho, + const int nspin, + const ModuleBase::Matrix3& GT, std::complex** rhog, - const int ipool, - const int irank, - const int nrank); + const Parallel::ParaWorld& pw_world, + std::ostream* os_warning); -} // namespace ModuleIO +} // namespace elecstate #endif diff --git a/source/source_estate/test/CMakeLists.txt b/source/source_estate/test/CMakeLists.txt index 1eaaa0c6bce..d474bd7c0ad 100644 --- a/source/source_estate/test/CMakeLists.txt +++ b/source/source_estate/test/CMakeLists.txt @@ -9,6 +9,9 @@ abacus_disable_feature_definitions(_OPENMP) if (ENABLE_MPI) +# Copy at configure time so a plain `make` + `ctest` run finds the data. +# install() only runs during `cmake --install`, which local test runs skip. +file(COPY support DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) install(DIRECTORY support DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) AddTest( @@ -113,4 +116,28 @@ AddTest( ../module_charge/gint_prec_ctrl.cpp ) +AddTest( + TARGET MODULE_ESTATE_test_rhog_io + LIBS parameter base device planewave + SOURCES test_rhog_io.cpp ../rhog_io.cpp ../../source_basis/module_pw/test/test_tool.cpp + # This test drives PW_Basis::initmpi and read/write_rhog's MPI collectives, + # so it must keep __MPI even though this directory disables it. Its main() + # calls MPI_Init via test_tool.cpp's setupmpi(). + KEEP_FEATURE_DEFINITIONS __MPI +) + +AddTest( + TARGET MODULE_ESTATE_charge_mpi_test + LIBS parameter psi base device planewave + SOURCES charge_mpi_test.cpp ../module_charge/charge_mpi.cpp + # Real MPI test: its main() calls MPI_Init unconditionally, so it must keep + # __MPI despite this directory disabling it. + KEEP_FEATURE_DEFINITIONS __MPI +) + +add_test(NAME MODULE_ESTATE_charge_mpi_test_4np + COMMAND mpirun -np 4 ./MODULE_ESTATE_charge_mpi_test; + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} +) + endif() diff --git a/source/source_estate/test_mpi/charge_mpi_test.cpp b/source/source_estate/test/charge_mpi_test.cpp similarity index 100% rename from source/source_estate/test_mpi/charge_mpi_test.cpp rename to source/source_estate/test/charge_mpi_test.cpp diff --git a/source/source_estate/test/support/charge-density.dat b/source/source_estate/test/support/charge-density.dat new file mode 100644 index 00000000000..ee3b0258011 Binary files /dev/null and b/source/source_estate/test/support/charge-density.dat differ diff --git a/source/source_estate/test/test_rhog_io.cpp b/source/source_estate/test/test_rhog_io.cpp new file mode 100644 index 00000000000..633530dba5d --- /dev/null +++ b/source/source_estate/test/test_rhog_io.cpp @@ -0,0 +1,406 @@ +#include "gmock/gmock.h" +#include "gtest/gtest.h" +#include "source_estate/rhog_io.h" +#include "source_base/module_parallel/para_world.h" +#include "source_base/module_parallel/para_tag.h" +#include "source_base/module_parallel/para_bridge.h" +#ifdef __MPI +#include "source_basis/module_pw/test/test_tool.h" +#include "mpi.h" +#endif +#include +#include + +/** + * - Tested Functions: + * - read_rhog() + * - write_rhog() + */ + +class ReadRhogTest : public ::testing::Test +{ + protected: + ModulePW::PW_Basis rhopw; + std::vector>> rhog_data; + std::vector*> rhog; + Parallel::ParaWorld pw_world = Parallel::make_pw_world(); + std::ofstream warning_stream; + + void setup_pw_basis() + { +#ifdef __MPI + rhopw.initmpi(pw_world.size(), pw_world.rank(), pw_world.comm()); +#endif + rhopw.initgrids(6.5, ModuleBase::Matrix3(-0.5, 0.0, 0.5, 0.0, 0.5, 0.5, -0.5, 0.5, 0.0), 120); + rhopw.initparameters(false, 120); + rhopw.setuptransform(); + rhopw.collect_local_pw(); + } + + void open_warning(const std::string& path) + { + warning_stream.open(path); + } + + void close_warning() + { + if (warning_stream.is_open()) + { + warning_stream.close(); + } + } + + std::string read_warning_file(const std::string& path) + { + std::ifstream ifs(path); + std::stringstream ss; + ss << ifs.rdbuf(); + ifs.close(); + return ss.str(); + } + + virtual void SetUp() + { + rhog_data.resize(1, std::vector>(1471)); + rhog.push_back(rhog_data[0].data()); + } + + virtual void TearDown() + { + close_warning(); + } +}; + +// Test the read_rhog function with normal file +TEST_F(ReadRhogTest, ReadRhog) +{ + std::string filename = "./support/charge-density.dat"; + setup_pw_basis(); + + bool result = elecstate::read_rhog(filename, &rhopw, 1, rhog.data(), pw_world, nullptr); + + EXPECT_TRUE(result); + EXPECT_DOUBLE_EQ(rhog[0][0].real(), -1.0304462993299456e-05); + EXPECT_DOUBLE_EQ(rhog[0][0].imag(), -1.2701788626185278e-13); + EXPECT_DOUBLE_EQ(rhog[0][1].real(), -0.0003875762482855959); + EXPECT_DOUBLE_EQ(rhog[0][1].imag(), -4.2556814316812048e-12); + EXPECT_DOUBLE_EQ(rhog[0][1470].real(), -3.5683133614445107e-05); + EXPECT_DOUBLE_EQ(rhog[0][1470].imag(), 1.6176615686863767e-12); +} + +// Test the read_rhog function when the file is not found +TEST_F(ReadRhogTest, NotFoundFile) +{ + setup_pw_basis(); + std::string filename = "notfound.txt"; + + open_warning("test_read_rhog.txt"); + bool result = elecstate::read_rhog(filename, &rhopw, 1, rhog.data(), pw_world, &warning_stream); + close_warning(); + + std::string expected_content = " elecstate::read_rhog warning : Can't open file notfound.txt\n"; + EXPECT_FALSE(result); + EXPECT_EQ(read_warning_file("test_read_rhog.txt"), expected_content); + std::remove("test_read_rhog.txt"); +} + +// Test the read_rhog function when gamma_only is inconsistent +TEST_F(ReadRhogTest, InconsistentGammaOnly) +{ + setup_pw_basis(); + std::string filename = "./support/charge-density.dat"; + rhopw.gamma_only = true; + // Fewer planewaves than the file holds (1471) triggers the + // "some planewaves in file are not used" warning. + rhopw.npwtot = 1000; + + open_warning("test_read_rhog.txt"); + bool result = elecstate::read_rhog(filename, &rhopw, 2, rhog.data(), pw_world, &warning_stream); + close_warning(); + + std::string expected_content + = " elecstate::read_rhog warning : some planewaves in file are not used\n elecstate::read_rhog warning : some " + "spin channels in file are missing\n elecstate::read_rhog warning : gamma_only read from file is " + "inconsistent with INPUT\n"; + + EXPECT_FALSE(result); + EXPECT_EQ(read_warning_file("test_read_rhog.txt"), expected_content); + std::remove("test_read_rhog.txt"); +} + +// Test the read_rhog function when some planewaves in file are missing +TEST_F(ReadRhogTest, SomePWMissing) +{ + setup_pw_basis(); + std::string filename = "./support/charge-density.dat"; + rhopw.npwtot = 2000; + + open_warning("test_read_rhog.txt"); + bool result = elecstate::read_rhog(filename, &rhopw, 1, rhog.data(), pw_world, &warning_stream); + close_warning(); + + std::string expected_content = " elecstate::read_rhog warning : some planewaves in file are missing\n"; + EXPECT_TRUE(result); + EXPECT_EQ(read_warning_file("test_read_rhog.txt"), expected_content); + std::remove("test_read_rhog.txt"); +} + +// Test read_rhog with os_warning=nullptr (silent mode, must not crash) +TEST_F(ReadRhogTest, OsNullptrSilent) +{ + std::string filename = "notfound.txt"; + bool result = elecstate::read_rhog(filename, &rhopw, 1, rhog.data(), pw_world, nullptr); + EXPECT_FALSE(result); +} + +// Test write_rhog round-trip: write then read back, verify data consistency +TEST_F(ReadRhogTest, WriteRoundTrip) +{ + setup_pw_basis(); + + // initialize some rhog data + rhog_data[0].assign(rhopw.npw, std::complex(1.5, 2.5)); + + std::string tmpfile = "test_rhog_roundtrip.dat"; + + // write + bool write_result = elecstate::write_rhog( + tmpfile, rhopw.gamma_only, &rhopw, 1, + ModuleBase::Matrix3(-0.5, 0.0, 0.5, 0.0, 0.5, 0.5, -0.5, 0.5, 0.0), + rhog.data(), pw_world, nullptr); + EXPECT_TRUE(write_result); + + // read back into a fresh buffer + std::vector>> rhog_read_data( + 1, std::vector>(rhopw.npw)); + std::vector*> rhog_read; + rhog_read.push_back(rhog_read_data[0].data()); + + bool read_result = elecstate::read_rhog(tmpfile, &rhopw, 1, rhog_read.data(), pw_world, nullptr); + EXPECT_TRUE(read_result); + + // compare: within MPI precision tolerance + int diff_count = 0; + for (int ig = 0; ig < rhopw.npw; ++ig) + { + if (std::abs(rhog[0][ig] - rhog_read[0][ig]) > 1e-10) + { + ++diff_count; + } + } + EXPECT_EQ(diff_count, 0) << diff_count << " planewave values differ after round-trip"; + + std::remove(tmpfile.c_str()); +} + +// Test write_rhog when the output path is not writable +TEST_F(ReadRhogTest, WriteFileFail) +{ + setup_pw_basis(); + rhog_data[0].assign(rhopw.npw, std::complex(1.0, 0.0)); + + // try to write to a directory path (not a file) — should fail + bool result = elecstate::write_rhog( + "/tmp", rhopw.gamma_only, &rhopw, 1, + ModuleBase::Matrix3(-0.5, 0.0, 0.5, 0.0, 0.5, 0.5, -0.5, 0.5, 0.0), + rhog.data(), pw_world, nullptr); + EXPECT_FALSE(result); +} + +// Test write_rhog with nspin=2, round-trip both channels +TEST_F(ReadRhogTest, WriteRoundTripNspin2) +{ + setup_pw_basis(); + + // expand to nspin=2 + rhog_data.resize(2, std::vector>(rhopw.npw)); + rhog.clear(); + rhog.push_back(rhog_data[0].data()); + rhog.push_back(rhog_data[1].data()); + + // initialize distinct values for each spin channel + for (int ig = 0; ig < rhopw.npw; ++ig) + { + rhog_data[0][ig] = std::complex(1.0 * ig, 0.1 * ig); + rhog_data[1][ig] = std::complex(2.0 * ig, 0.2 * ig); + } + + std::string tmpfile = "test_rhog_roundtrip_nspin2.dat"; + + // write nspin=2 + bool write_result = elecstate::write_rhog( + tmpfile, rhopw.gamma_only, &rhopw, 2, + ModuleBase::Matrix3(-0.5, 0.0, 0.5, 0.0, 0.5, 0.5, -0.5, 0.5, 0.0), + rhog.data(), pw_world, nullptr); + EXPECT_TRUE(write_result); + + // read back + std::vector>> rhog_read_data( + 2, std::vector>(rhopw.npw)); + std::vector*> rhog_read; + rhog_read.push_back(rhog_read_data[0].data()); + rhog_read.push_back(rhog_read_data[1].data()); + + bool read_result = elecstate::read_rhog(tmpfile, &rhopw, 2, rhog_read.data(), pw_world, nullptr); + EXPECT_TRUE(read_result); + + int diff_count = 0; + for (int is = 0; is < 2; ++is) + { + for (int ig = 0; ig < rhopw.npw; ++ig) + { + if (std::abs(rhog[is][ig] - rhog_read[is][ig]) > 1e-10) + { + ++diff_count; + } + } + } + EXPECT_EQ(diff_count, 0) << diff_count << " planewave values differ after nspin=2 round-trip"; + + std::remove(tmpfile.c_str()); +} + +// Test write_rhog with nspin=4, round-trip all 4 channels +TEST_F(ReadRhogTest, WriteRoundTripNspin4) +{ + setup_pw_basis(); + + rhog_data.resize(4, std::vector>(rhopw.npw)); + rhog.clear(); + for (int is = 0; is < 4; ++is) + { + rhog.push_back(rhog_data[is].data()); + } + + // initialize distinct values for each spin channel + for (int is = 0; is < 4; ++is) + { + for (int ig = 0; ig < rhopw.npw; ++ig) + { + rhog_data[is][ig] = std::complex((is + 1) * 1.0 * ig, (is + 1) * 0.1 * ig); + } + } + + std::string tmpfile = "test_rhog_roundtrip_nspin4.dat"; + + bool write_result = elecstate::write_rhog( + tmpfile, rhopw.gamma_only, &rhopw, 4, + ModuleBase::Matrix3(-0.5, 0.0, 0.5, 0.0, 0.5, 0.5, -0.5, 0.5, 0.0), + rhog.data(), pw_world, nullptr); + EXPECT_TRUE(write_result); + + // read back as nspin=4 + std::vector>> rhog_read_data( + 4, std::vector>(rhopw.npw)); + std::vector*> rhog_read; + for (int is = 0; is < 4; ++is) + { + rhog_read.push_back(rhog_read_data[is].data()); + } + + bool read_result = elecstate::read_rhog(tmpfile, &rhopw, 4, rhog_read.data(), pw_world, nullptr); + EXPECT_TRUE(read_result); + + int diff_count = 0; + for (int is = 0; is < 4; ++is) + { + for (int ig = 0; ig < rhopw.npw; ++ig) + { + if (std::abs(rhog[is][ig] - rhog_read[is][ig]) > 1e-10) + { + ++diff_count; + } + } + } + EXPECT_EQ(diff_count, 0) << diff_count << " planewave values differ after nspin=4 round-trip"; + + std::remove(tmpfile.c_str()); +} + +// Test the special path L173-181: file nspin=2 read as input nspin=4 +// Expected behavior: rhog[0] preserved, rhog[1] and rhog[2] zeroed, +// rhog[3] <- old rhog[1] +TEST_F(ReadRhogTest, ReadRhogNspin2To4SpecialPath) +{ + setup_pw_basis(); + + // Step 1: write a nspin=2 binary with known values + rhog_data.resize(2, std::vector>(rhopw.npw)); + rhog.clear(); + rhog.push_back(rhog_data[0].data()); + rhog.push_back(rhog_data[1].data()); + + for (int ig = 0; ig < rhopw.npw; ++ig) + { + rhog_data[0][ig] = std::complex(10.0 + ig, 0.0); + rhog_data[1][ig] = std::complex(20.0 + ig, 0.0); + } + + std::string tmpfile = "test_rhog_nspin2_to_4.dat"; + + bool write_result = elecstate::write_rhog( + tmpfile, rhopw.gamma_only, &rhopw, 2, + ModuleBase::Matrix3(-0.5, 0.0, 0.5, 0.0, 0.5, 0.5, -0.5, 0.5, 0.0), + rhog.data(), pw_world, nullptr); + EXPECT_TRUE(write_result); + + // Step 2: read back as nspin=4 — triggers the L173-181 special path + std::vector>> rhog_read_data( + 4, std::vector>(rhopw.npw)); + std::vector*> rhog_read; + for (int is = 0; is < 4; ++is) + { + rhog_read.push_back(rhog_read_data[is].data()); + } + + bool read_result = elecstate::read_rhog(tmpfile, &rhopw, 4, rhog_read.data(), pw_world, nullptr); + EXPECT_TRUE(read_result); + + // Verify the special transformation at L173-181: + // rhog[0] <- file spin 0 + // rhog[1] <- ZEROED (was file spin 1, then ZEROS) + // rhog[2] <- ZEROED + // rhog[3] <- file spin 1 (copied before ZEROS) + for (int ig = 0; ig < rhopw.npw; ++ig) + { + // rhog[0] should match original spin 0 + EXPECT_NEAR(rhog_read_data[0][ig].real(), 10.0 + ig, 1e-10); + EXPECT_NEAR(rhog_read_data[0][ig].imag(), 0.0, 1e-10); + + // rhog[1] should be zeroed + EXPECT_NEAR(rhog_read_data[1][ig].real(), 0.0, 1e-10); + EXPECT_NEAR(rhog_read_data[1][ig].imag(), 0.0, 1e-10); + + // rhog[2] should be zeroed + EXPECT_NEAR(rhog_read_data[2][ig].real(), 0.0, 1e-10); + EXPECT_NEAR(rhog_read_data[2][ig].imag(), 0.0, 1e-10); + + // rhog[3] should equal original spin 1 (copied before zero) + EXPECT_NEAR(rhog_read_data[3][ig].real(), 20.0 + ig, 1e-10); + EXPECT_NEAR(rhog_read_data[3][ig].imag(), 0.0, 1e-10); + } + + std::remove(tmpfile.c_str()); +} + +int main(int argc, char** argv) +{ +#ifdef __MPI + int nproc = 1; + int myrank = 0; + int nproc_in_pool = 1; + int kpar = 1; + int mypool = 0; + int rank_in_pool = 0; + setupmpi(argc, argv, nproc, myrank); + divide_pools(nproc, myrank, nproc_in_pool, kpar, mypool, rank_in_pool); +#endif + + testing::InitGoogleTest(&argc, argv); + int result = RUN_ALL_TESTS(); + +#ifdef __MPI + finishmpi(); +#endif + return result; +} diff --git a/source/source_estate/test_mpi/CMakeLists.txt b/source/source_estate/test_mpi/CMakeLists.txt deleted file mode 100644 index a6c068027f7..00000000000 --- a/source/source_estate/test_mpi/CMakeLists.txt +++ /dev/null @@ -1,18 +0,0 @@ -abacus_disable_feature_definitions(__EXX) -abacus_disable_feature_definitions(__CUDA) -abacus_disable_feature_definitions(__UT_USE_CUDA) -abacus_disable_feature_definitions(__UT_USE_ROCM) -abacus_disable_feature_definitions(__ROCM) -abacus_disable_feature_definitions(__MLALGO) -abacus_disable_feature_definitions(_OPENMP) - -AddTest( - TARGET MODULE_ESTATE_charge_mpi_test - LIBS parameter psi base device planewave - SOURCES charge_mpi_test.cpp ../module_charge/charge_mpi.cpp -) - -add_test(NAME MODULE_ESTATE_charge_mpi_test_4np - COMMAND mpirun -np 4 ./MODULE_ESTATE_charge_mpi_test; - WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} -) diff --git a/source/source_io/module_chgpot/write_elecstat_pot.cpp b/source/source_estate/write_elecstat_pot.cpp similarity index 100% rename from source/source_io/module_chgpot/write_elecstat_pot.cpp rename to source/source_estate/write_elecstat_pot.cpp diff --git a/source/source_io/module_chgpot/write_elecstat_pot.h b/source/source_estate/write_elecstat_pot.h similarity index 100% rename from source/source_io/module_chgpot/write_elecstat_pot.h rename to source/source_estate/write_elecstat_pot.h diff --git a/source/source_io/module_chgpot/write_init.cpp b/source/source_estate/write_init.cpp similarity index 99% rename from source/source_io/module_chgpot/write_init.cpp rename to source/source_estate/write_init.cpp index b7ffdc18893..58cdb534287 100644 --- a/source/source_io/module_chgpot/write_init.cpp +++ b/source/source_estate/write_init.cpp @@ -14,7 +14,7 @@ // Module: module_io/module_chgpot // ===================================================================== -#include "source_io/module_chgpot/write_init.h" +#include "source_estate/write_init.h" #include "source_io/module_output/cube_io.h" #include "source_base/tool_quit.h" diff --git a/source/source_io/module_chgpot/write_init.h b/source/source_estate/write_init.h similarity index 100% rename from source/source_io/module_chgpot/write_init.h rename to source/source_estate/write_init.h diff --git a/source/source_hsolver/test/diago_bpcg_test.cpp b/source/source_hsolver/test/diago_bpcg_test.cpp index dbd309041e3..d2ef8244698 100644 --- a/source/source_hsolver/test/diago_bpcg_test.cpp +++ b/source/source_hsolver/test/diago_bpcg_test.cpp @@ -1,5 +1,6 @@ #include "source_base/inverse_matrix.h" #include "source_base/module_external/lapack_connector.h" +#include "source_base/parallel_comm.h" #include "source_psi/psi.h" #include "source_hamilt/hamilt.h" #include "source_pw/module_pwdft/hamilt_pw.h" diff --git a/source/source_hsolver/test/test_hsolver_sdft.cpp b/source/source_hsolver/test/test_hsolver_sdft.cpp index 570f3ff9227..bc35f52f3ad 100644 --- a/source/source_hsolver/test/test_hsolver_sdft.cpp +++ b/source/source_hsolver/test/test_hsolver_sdft.cpp @@ -10,6 +10,7 @@ #include "hsolver_pw_sup.h" #include "hsolver_supplementary_mock.h" #include "source_base/global_variable.h" +#include "source_base/parallel_comm.h" #include "source_hsolver/hsolver_pw.h" #include "source_hsolver/hsolver_pw_sdft.h" #include "source_estate/elecstate_pw.h" diff --git a/source/source_io/CMakeLists.txt b/source/source_io/CMakeLists.txt index 2b2634f5ee5..85ffd82d545 100644 --- a/source/source_io/CMakeLists.txt +++ b/source/source_io/CMakeLists.txt @@ -23,17 +23,14 @@ list(APPEND objects module_bessel/numerical_descriptor.cpp module_output/print_info.cpp module_output/read_cube.cpp - module_chgpot/rhog_io.cpp module_wf/read_wfc_pw.cpp module_wf/read_wf2rho_pw.cpp module_restart/restart.cpp module_wf/write_wfc_pw.cpp module_output/write_pao.cpp module_output/write_cube.cpp - module_chgpot/write_elecstat_pot.cpp module_elf/write_elf.cpp module_dipole/write_dipole.cpp - module_chgpot/write_init.cpp module_ml/write_mlkedf_desc.cpp module_current/td_current_io.cpp module_current/td_current_io_comm.cpp diff --git a/source/source_io/module_chgpot/rhog_io.cpp b/source/source_io/module_chgpot/rhog_io.cpp deleted file mode 100644 index ab17608cadc..00000000000 --- a/source/source_io/module_chgpot/rhog_io.cpp +++ /dev/null @@ -1,423 +0,0 @@ -#include "source_base/module_out/binstream.h" -#include "source_base/global_function.h" -#include "source_io/module_parameter/parameter.h" -#include "source_base/global_variable.h" -#include "source_base/parallel_global.h" -#include "source_base/timer.h" -#include "source_base/vector3.h" -#include "rhog_io.h" -#include -#include - -bool ModuleIO::read_rhog(const std::string& filename, const ModulePW::PW_Basis* pw_rhod, std::complex** rhog) -{ - ModuleBase::TITLE("ModuleIO", "read_rhog"); - ModuleBase::timer::start("ModuleIO", "read_rhog"); - - const int nx = pw_rhod->nx; - const int ny = pw_rhod->ny; - const int nz = pw_rhod->nz; - - Binstream ifs; - bool error = false; - int gamma_only_in = 0; - int npwtot_in = 0; - int nspin_in = 0; - int size = 0; - double b1[3], b2[3], b3[3]; - - if (GlobalV::RANK_IN_POOL == 0) - { - ifs.open(filename, "r"); - if (!ifs) - { - error = true; - } - } - -#ifdef __MPI - MPI_Bcast(&error, 1, MPI_C_BOOL, 0, POOL_WORLD); -#endif - - if (error) - { - ModuleBase::WARNING("ModuleIO::read_rhog", "Can't open file " + filename); - ModuleBase::timer::end("ModuleIO", "read_rhog"); - return false; - } - - if (GlobalV::RANK_IN_POOL == 0) - { - ifs >> size >> gamma_only_in >> npwtot_in >> nspin_in >> size; - ifs >> size >> b1[0] >> b1[1] >> b1[2] >> b2[0] >> b2[1] >> b2[2] >> b3[0] >> b3[1] >> b3[2] >> size; - if (gamma_only_in != pw_rhod->gamma_only) - { - // there is a treatment that can transform between gamma_only and non-gamma_only - // however, it is not implemented here - error = true; - ifs.close(); - } - if (npwtot_in > pw_rhod->npwtot) - { - ModuleBase::WARNING("ModuleIO::read_rhog", "some planewaves in file are not used"); - } - else if (npwtot_in < pw_rhod->npwtot) - { - ModuleBase::WARNING("ModuleIO::read_rhog", "some planewaves in file are missing"); - } - if (nspin_in < PARAM.inp.nspin) - { - ModuleBase::WARNING("ModuleIO::read_rhog", "some spin channels in file are missing"); - } - } - -#ifdef __MPI - MPI_Bcast(&error, 1, MPI_C_BOOL, 0, POOL_WORLD); -#endif - - if (error) - { - ModuleBase::WARNING("ModuleIO::read_rhog", "gamma_only read from file is inconsistent with INPUT"); - ModuleBase::timer::end("ModuleIO", "read_rhog"); - return false; - } - -#ifdef __MPI - MPI_Bcast(&gamma_only_in, 1, MPI_INT, 0, POOL_WORLD); - MPI_Bcast(&npwtot_in, 1, MPI_INT, 0, POOL_WORLD); - MPI_Bcast(&nspin_in, 1, MPI_INT, 0, POOL_WORLD); - MPI_Bcast(b1, 3, MPI_DOUBLE, 0, POOL_WORLD); - MPI_Bcast(b2, 3, MPI_DOUBLE, 0, POOL_WORLD); - MPI_Bcast(b3, 3, MPI_DOUBLE, 0, POOL_WORLD); -#endif - std::vector miller(npwtot_in * 3); - // once use ModuleBase::Vector3, it is highly bug-prone to assume the memory layout of the class. - // The x, y and z of Vector3 will not always to be contiguous. - // Instead, a relatively safe choice is to use std::vector, the memory layout is assumed - // to be npwtot_in rows and 3 columns. - if (GlobalV::RANK_IN_POOL == 0) - { - ifs >> size; - for (int i = 0; i < npwtot_in; ++i) // loop over rows... - { - ifs >> miller[i*3] >> miller[i*3+1] >> miller[i*3+2]; - } - ifs >> size; - } -#ifdef __MPI - MPI_Bcast(miller.data(), miller.size(), MPI_INT, 0, POOL_WORLD); -#endif - // set to zero - for (int is = 0; is < PARAM.inp.nspin; ++is) - { - ModuleBase::GlobalFunc::ZEROS(rhog[is], pw_rhod->npw); - } - // maps ixyz tp ig - std::vector fftixyz2ig(pw_rhod->nxyz, -1); // map isz to ig. - for (int ig = 0; ig < pw_rhod->npw; ++ig) - { - int isz = pw_rhod->ig2isz[ig]; - int iz = isz % nz; - int is = isz / nz; - int ixy = pw_rhod->is2fftixy[is]; - int ixyz = iz + nz * ixy; - fftixyz2ig[ixyz] = ig; - } - std::vector> rhog_in(npwtot_in); - for (int is = 0; is < nspin_in; ++is) - { - if (GlobalV::RANK_IN_POOL == 0) - { - ifs >> size; - for (int i = 0; i < npwtot_in; ++i) - { - ifs >> rhog_in[i]; - } - ifs >> size; - } -#ifdef __MPI - MPI_Bcast(rhog_in.data(), rhog_in.size(), MPI_DOUBLE_COMPLEX, 0, POOL_WORLD); -#endif - - for (int i = 0; i < npwtot_in; ++i) - { - int ix = miller[i * 3]; - int iy = miller[i * 3 + 1]; - int iz = miller[i * 3 + 2]; - - if (ix <= -int((nx + 1) / 2) || ix >= int(nx / 2) + 1 || iy <= -int((ny + 1) / 2) || iy >= int(ny / 2) + 1 - || iz <= -int((nz + 1) / 2) || iz >= int(nz / 2) + 1) - { - // these planewaves are not used - continue; - } - - if (ix < 0) - ix += nx; - if (iy < 0) - iy += ny; - if (iz < 0) - iz += nz; - int fftixy = iy + pw_rhod->fftny * ix; - if (GlobalV::RANK_IN_POOL == pw_rhod->fftixy2ip[fftixy]) - { - int fftixyz = iz + nz * fftixy; - int ig = fftixyz2ig[fftixyz]; - rhog[is][ig] = rhog_in[i]; - } - } - - if (nspin_in == 2 && PARAM.inp.nspin == 4 && is == 1) - { - for (int ig = 0; ig < pw_rhod->npw; ++ig) - { - rhog[3][ig] = rhog[1][ig]; - } - ModuleBase::GlobalFunc::ZEROS(rhog[1], pw_rhod->npw); - ModuleBase::GlobalFunc::ZEROS(rhog[2], pw_rhod->npw); - } - } - - if (GlobalV::RANK_IN_POOL == 0) - { - ifs.close(); - } - // for debug, write the rhog to a file (not binary) - // if (GlobalV::RANK_IN_POOL == 0) - // { - // std::ofstream ofs("rhog_read.txt"); - // for (int i = 0; i < nspin_in; ++i) - // { - // for (int ig = 0; ig < pw_rhod->npw; ++ig) - // { - // ofs << rhog[i][ig] << " "; - // } - // ofs << std::endl; - // } - // ofs.close(); - // } - ModuleBase::timer::end("ModuleIO", "read_rhog"); - return true; -} - -bool ModuleIO::write_rhog(const std::string& fchg, - const bool gamma_only, // from INPUT - const ModulePW::PW_Basis* pw_rho, // pw_rho in runtime - const int nspin, // GlobalV - const ModuleBase::Matrix3& GT, // from UnitCell, useful for calculating the miller - std::complex** rhog, - const int ipool, - const int irank, - const int nrank) -{ - ModuleBase::TITLE("ModuleIO", "write_rhog"); - ModuleBase::timer::start("ModuleIO", "write_rhog"); - if (ipool != 0) { - ModuleBase::timer::end("ModuleIO", "write_rhog"); - return true; - } - // only one pool writes the rhog, because rhog in all pools are identical. - - // for large-scale data, it is not wise to collect all distributed components to the - // master process and then write the data to the file. Instead, we can write the data - // processer by processer. - - // Quantum ESPRESSO requires the G-vectors collected should be in the order like as if - // there is only 1 process, this order is recorded in - - // fftixy2ip will be useful for the order of the G-vectors - // each time we iterate on ig, then find the rho_g over all the processes - // for ig in npwtot, then find the local index of ig on processor, ig -> fftixy2ip -> igl - - - // write the header (by rank 0): gamma_only, ngm_g, nspin - int size = 3; - // because "reinterpret_cast" cannot drop the "const", so use intermediate variable - int ngm_g = pw_rho->npwtot; - int gam = gamma_only; // IMPLICIT DATA TYPE CONVERSION! - int nsp = nspin; - - std::ofstream ofs; -#ifdef __MPI - MPI_Barrier(POOL_WORLD); - // this is still a global variable... should be moved into param - // list as `const MPI_Comm& comm` - if (irank == 0) - { - // printf(" CHGDEN >>> Writing header by rank %d...\n", irank); -#endif - ofs.open(fchg, std::ios::binary); // open the file by all processors - if (!ofs) - { - ModuleBase::WARNING_QUIT("ModuleIO::write_rhog", "File I/O failure: cannot open file " + fchg); - ModuleBase::timer::end("ModuleIO", "write_rhog"); - return false; - } - ofs.write(reinterpret_cast(&size), sizeof(size)); - ofs.write(reinterpret_cast(&gam), sizeof(gam)); - ofs.write(reinterpret_cast(&ngm_g), sizeof(ngm_g)); - ofs.write(reinterpret_cast(&nsp), sizeof(nsp)); - ofs.write(reinterpret_cast(&size), sizeof(size)); - // write the lattice vectors, GT is the reciprocal lattice vectors, need 2pi? - std::vector b = {GT.e11, GT.e12, GT.e13, GT.e21, GT.e22, GT.e23, GT.e31, GT.e32, GT.e33}; - size = 9; - ofs.write(reinterpret_cast(&size), sizeof(size)); - for (int i = 0; i < 9; ++i) - { - ofs.write(reinterpret_cast(&b[i]), sizeof(b[i])); - } - ofs.write(reinterpret_cast(&size), sizeof(size)); - ofs.close(); -#ifdef __MPI - // printf(" CHGDEN >>> Complete header writing by rank %d\n", irank); - } - MPI_Barrier(POOL_WORLD); // wait for rank 0 to finish writing the header - // printf(" CHGDEN >>> rank %d ready for continue writing...\n", irank); - MPI_Barrier(POOL_WORLD); -#endif - - // write the G-vectors in Miller indices, the Miller indices can be calculated by - // the dot product of the G-vectors and the reciprocal lattice vectors - // parallelization needed considered here. Because the sequence of the G-vectors - // is not important, we can write the G-vectors processer by processer - size = 3*ngm_g; -#ifdef __MPI - if(irank == 0) - { - // printf(" CHGDEN >>> Writing header of Miller indices by rank %d...\n", irank); -#endif - ofs.open(fchg, std::ios::binary | std::ios::app); // open the file by rank 0 - ofs.write(reinterpret_cast(&size), sizeof(size)); - ofs.close(); -#ifdef __MPI - // printf(" CHGDEN >>> Complete header of Miller indices writing by rank %d\n", irank); - } - MPI_Barrier(POOL_WORLD); // wait for rank 0 to finish writing the header of miller indices -#endif -#ifdef __MPI - for(int i = 0; i < nrank; ++i) // write the miller indices processer by processer - { - if(i == irank) - { - // printf(" CHGDEN >>> Writing Miller indices by rank %d...\n", irank); -#endif - ofs.open(fchg, std::ios::binary | std::ios::app); // open the file by processer i - for(int ig = 0; ig < pw_rho->npw; ++ig) - { - const ModuleBase::Vector3 g = pw_rho->gdirect[ig]; // g direct is (ix, iy, iz), miller index (integer), centered at (0, 0, 0) - std::vector miller = {int(g.x), int(g.y), int(g.z)}; - ofs.write(reinterpret_cast(&miller[0]), sizeof(miller[0])); - ofs.write(reinterpret_cast(&miller[1]), sizeof(miller[1])); - ofs.write(reinterpret_cast(&miller[2]), sizeof(miller[2])); - } - ofs.close(); -#ifdef __MPI - // printf(" CHGDEN >>> Complete Miller indices writing by rank %d\n", irank); - } - MPI_Barrier(POOL_WORLD); // wait for the current rank to finish writing the miller indices - } -#endif -#ifdef __MPI - if(irank == 0) - { -#endif - ofs.open(fchg, std::ios::binary | std::ios::app); // open the file by rank 0 - ofs.write(reinterpret_cast(&size), sizeof(size)); - ofs.close(); -#ifdef __MPI - } - MPI_Barrier(POOL_WORLD); // wait for rank 0 to finish writing the miller indices -#endif - - // write the rho(G) values - std::complex sum_check; - size = ngm_g; - for(int ispin = 0; ispin < nspin; ++ispin) - { -#ifdef __MPI - if(irank == 0) - { - // printf(" CHGDEN >>> Writing header of rho(G) values by rank %d...\n", irank); -#endif - ofs.open(fchg, std::ios::binary | std::ios::app); // open the file by rank 0 - ofs.write(reinterpret_cast(&size), sizeof(size)); - ofs.close(); -#ifdef __MPI - // printf(" CHGDEN >>> Complete header of rho(G) values writing by rank %d\n", irank); - } - MPI_Barrier(POOL_WORLD); // wait for rank 0 to finish writing the header of rho(G) -#endif -#ifdef __MPI - for(int i = 0; i < nrank; ++i) // write the rho(G) values processer by processer - { - if(i == irank) - { - // printf(" CHGDEN >>> Writing rho(G) values by rank %d...\n", irank); -#endif - ofs.open(fchg, std::ios::binary | std::ios::app); // open the file by processer i - sum_check = 0.0; - for(int ig = 0; ig < pw_rho->npw; ++ig) - { - sum_check += rhog[ispin][ig]; - ofs.write(reinterpret_cast(&rhog[ispin][ig]), sizeof(rhog[ispin][ig])); - } - // assert(std::abs(sum_check) > 1.0e-10); // check if the sum of rho(G) is valid - ofs.close(); -#ifdef __MPI - // printf(" CHGDEN >>> Complete rho(G) values writing by rank %d\n", irank); - } - MPI_Barrier(POOL_WORLD); // wait for the current rank to finish writing the rho(G) values - } -#endif - -#ifdef __MPI - if(irank == 0) - { -#endif - ofs.open(fchg, std::ios::binary | std::ios::app); // open the file by rank 0 - ofs.write(reinterpret_cast(&size), sizeof(size)); - ofs.close(); -#ifdef __MPI - } - MPI_Barrier(POOL_WORLD); // wait for rank 0 to finish writing the rho(G) values -#endif - } - // for debug, write the rhog to a file (not binary) - // if (irank == 0) - // { - // std::ofstream ofs("rhog_write.txt"); - // for (int i = 0; i < nspin; ++i) - // { - // for (int ig = 0; ig < pw_rho->npw; ++ig) - // { - // ofs << rhog[i][ig] << " "; - // } - // ofs << std::endl; - // } - // ofs.close(); - // } - ModuleBase::timer::end("ModuleIO", "write_rhog"); - return true; -} - -// self-consistency test with the following python code -// import numpy as np - -// with open("rhog_read.txt") as f: -// read = f.readlines() - -// with open("rhog_write.txt") as f: -// write = f.readlines() - -// # convert c++ stype complex number (a,b) to python complex -// def to_complex(s): -// a, b = s.replace("(", "").replace(")", "").split(",") -// return complex(float(a), float(b)) - -// read = [[to_complex(rhog) for rhog in spin.strip().split()] for spin in read] -// write = [[to_complex(rhog) for rhog in spin.strip().split()] for spin in write] - -// diff = np.array(read) - np.array(write) -// print(np.max(np.abs(diff))) -// test system: integrated test 118_PW_CHG_BINARY -// yielding error 5.290000000000175e-11 \ No newline at end of file diff --git a/source/source_io/module_ctrl/ctrl_output_fp.cpp b/source/source_io/module_ctrl/ctrl_output_fp.cpp index ccb6a900a76..318c571e3c4 100644 --- a/source/source_io/module_ctrl/ctrl_output_fp.cpp +++ b/source/source_io/module_ctrl/ctrl_output_fp.cpp @@ -3,7 +3,7 @@ #include "../module_dipole/dipole_io.h" // use write_dipole #include "source_estate/module_charge/symm_rho.h" // use Symmetry_rho #include "source_hamilt/module_xc/xc_functional.h" // use XC_Functional -#include "source_io/module_chgpot/write_elecstat_pot.h" // use write_elecstat_pot +#include "source_estate/write_elecstat_pot.h" // use write_elecstat_pot #include "source_io/module_elf/write_elf.h" #ifdef __LIBXC diff --git a/source/source_io/module_hs/write_hs.h b/source/source_io/module_hs/write_hs.h index 14a972d8413..25c1cecc0cd 100644 --- a/source/source_io/module_hs/write_hs.h +++ b/source/source_io/module_hs/write_hs.h @@ -6,6 +6,7 @@ //#include "source_base/global_function.h" //#include "source_base/global_variable.h" +#include "source_base/parallel_comm.h" // use DIAG_WORLD #include "source_basis/module_ao/parallel_orbitals.h" // use Parallel_Orbitals #include "source_hamilt/hamilt.h" diff --git a/source/source_io/test/CMakeLists.txt b/source/source_io/test/CMakeLists.txt index 4385b04dabc..70d94d91a24 100644 --- a/source/source_io/test/CMakeLists.txt +++ b/source/source_io/test/CMakeLists.txt @@ -140,12 +140,6 @@ AddTest( SOURCES ../module_output/output_log.cpp outputlog_test.cpp ../../source_basis/module_pw/test/test_tool.cpp ) -AddTest( - TARGET MODULE_IO_read_rhog_test - LIBS parameter base device planewave - SOURCES read_rhog_test.cpp ../module_chgpot/rhog_io.cpp ../../source_basis/module_pw/test/test_tool.cpp -) - if(ENABLE_LCAO) AddTest( TARGET MODULE_IO_to_qo_test diff --git a/source/source_io/test/read_rhog_test.cpp b/source/source_io/test/read_rhog_test.cpp deleted file mode 100644 index bb90e3b3442..00000000000 --- a/source/source_io/test/read_rhog_test.cpp +++ /dev/null @@ -1,159 +0,0 @@ -#include "gmock/gmock.h" -#include "gtest/gtest.h" -#define private public -#include "source_io/module_parameter/parameter.h" -#undef private -#include "source_io/module_chgpot/rhog_io.h" -#ifdef __MPI -#include "source_basis/module_pw/test/test_tool.h" -#include "mpi.h" -#endif - -/** - * - Tested Functions: - * - read_rhog() - */ - -class ReadRhogTest : public ::testing::Test -{ - protected: - ModulePW::PW_Basis* rhopw = nullptr; - std::complex** rhog = nullptr; - - virtual void SetUp() - { - rhopw = new ModulePW::PW_Basis; - rhog = new std::complex*[1]; - rhog[0] = new std::complex[1471]; - } - virtual void TearDown() - { - if (rhopw != nullptr) { - delete rhopw; -} - if (rhog[0] != nullptr) { - delete[] rhog[0]; -} - if (rhog != nullptr) { - delete[] rhog; -} - } -}; - -// Test the read_rhog function -TEST_F(ReadRhogTest, ReadRhog) -{ - std::string filename = "./support/charge-density.dat"; - PARAM.input.nspin = 1; -#ifdef __MPI - rhopw->initmpi(GlobalV::NPROC_IN_POOL, GlobalV::RANK_IN_POOL, MPI_COMM_WORLD); -#endif - rhopw->initgrids(6.5, ModuleBase::Matrix3(-0.5, 0.0, 0.5, 0.0, 0.5, 0.5, -0.5, 0.5, 0.0), 120); - rhopw->initparameters(false, 120); - rhopw->setuptransform(); - rhopw->collect_local_pw(); - - bool result = ModuleIO::read_rhog(filename, rhopw, rhog); - - EXPECT_TRUE(result); - EXPECT_DOUBLE_EQ(rhog[0][0].real(), -1.0304462993299456e-05); - EXPECT_DOUBLE_EQ(rhog[0][0].imag(), -1.2701788626185278e-13); - EXPECT_DOUBLE_EQ(rhog[0][1].real(), -0.0003875762482855959); - EXPECT_DOUBLE_EQ(rhog[0][1].imag(), -4.2556814316812048e-12); - EXPECT_DOUBLE_EQ(rhog[0][1470].real(), -3.5683133614445107e-05); - EXPECT_DOUBLE_EQ(rhog[0][1470].imag(), 1.6176615686863767e-12); -} - -// Test the read_rhog function when the file is not found -TEST_F(ReadRhogTest, NotFoundFile) -{ - std::string filename = "notfound.txt"; - - GlobalV::ofs_warning.open("test_read_rhog.txt"); - bool result = ModuleIO::read_rhog(filename, rhopw, rhog); - GlobalV::ofs_warning.close(); - - std::ifstream ifs_running("test_read_rhog.txt"); - std::stringstream ss; - ss << ifs_running.rdbuf(); - std::string file_content = ss.str(); - ifs_running.close(); - - std::string expected_content = " ModuleIO::read_rhog warning : Can't open file notfound.txt\n"; - - EXPECT_FALSE(result); - EXPECT_EQ(file_content, expected_content); - std::remove("test_read_rhog.txt"); -} - -// Test the read_rhog function when tgamma_only is inconsistent -TEST_F(ReadRhogTest, InconsistentGammaOnly) -{ - std::string filename = "./support/charge-density.dat"; - PARAM.input.nspin = 2; - rhopw->gamma_only = true; - - GlobalV::ofs_warning.open("test_read_rhog.txt"); - bool result = ModuleIO::read_rhog(filename, rhopw, rhog); - GlobalV::ofs_warning.close(); - - std::ifstream ifs_running("test_read_rhog.txt"); - std::stringstream ss; - ss << ifs_running.rdbuf(); - std::string file_content = ss.str(); - ifs_running.close(); - - std::string expected_content - = " ModuleIO::read_rhog warning : some planewaves in file are not used\n ModuleIO::read_rhog warning : some " - "spin channels in file are missing\n ModuleIO::read_rhog warning : gamma_only read from file is " - "inconsistent with INPUT\n"; - - EXPECT_FALSE(result); - EXPECT_EQ(file_content, expected_content); - std::remove("test_read_rhog.txt"); -} - -// Test the read_rhog function when some planewaves in file are missing -TEST_F(ReadRhogTest, SomePWMissing) -{ - std::string filename = "./support/charge-density.dat"; - PARAM.input.nspin = 1; - rhopw->npwtot = 2000; - - GlobalV::ofs_warning.open("test_read_rhog.txt"); - bool result = ModuleIO::read_rhog(filename, rhopw, rhog); - GlobalV::ofs_warning.close(); - - std::ifstream ifs_running("test_read_rhog.txt"); - std::stringstream ss; - ss << ifs_running.rdbuf(); - std::string file_content = ss.str(); - ifs_running.close(); - - std::string expected_content = " ModuleIO::read_rhog warning : some planewaves in file are missing\n"; - - EXPECT_TRUE(result); - EXPECT_EQ(file_content, expected_content); - std::remove("test_read_rhog.txt"); -} - -int main(int argc, char** argv) -{ -#ifdef __MPI - setupmpi(argc, argv, GlobalV::NPROC, GlobalV::MY_RANK); - divide_pools(GlobalV::NPROC, - GlobalV::MY_RANK, - GlobalV::NPROC_IN_POOL, - GlobalV::KPAR, - GlobalV::MY_POOL, - GlobalV::RANK_IN_POOL); -#endif - - testing::InitGoogleTest(&argc, argv); - int result = RUN_ALL_TESTS(); - -#ifdef __MPI - finishmpi(); -#endif - return result; -} \ No newline at end of file diff --git a/source/source_lcao/module_lr/hsolver_lrtd.hpp b/source/source_lcao/module_lr/hsolver_lrtd.hpp index 24a36b59a5c..72660a4503c 100644 --- a/source/source_lcao/module_lr/hsolver_lrtd.hpp +++ b/source/source_lcao/module_lr/hsolver_lrtd.hpp @@ -7,6 +7,7 @@ #include "source_lcao/module_lr/utils/lr_util.h" #include "source_lcao/module_lr/utils/lr_util_print.h" #include "source_base/module_container/ATen/core/tensor_map.h" +#include "source_base/parallel_comm.h" namespace LR { diff --git a/source/source_pw/module_stodft/sto_elecond.cpp b/source/source_pw/module_stodft/sto_elecond.cpp index ba60a6114a1..85516450f99 100644 --- a/source/source_pw/module_stodft/sto_elecond.cpp +++ b/source/source_pw/module_stodft/sto_elecond.cpp @@ -4,6 +4,7 @@ #include "source_base/constants.h" #include "source_base/memory_recorder.h" #include "source_base/module_container/ATen/tensor.h" +#include "source_base/parallel_comm.h" #include "source_base/parallel_device.h" #include "source_base/parallel_comm.h" #include "source_base/parallel_reduce.h"