Refactor parallel communicators, Introduce module_parallel and migrate rhog_io to domain-aware interfaces

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:\
@@ -778,7 +779,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\

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()

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

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

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

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

source/source_base/module_parallel/para_collection.cpp

@@ -0,0 +1,31 @@
+#include "para_collection.h"
+
+namespace Parallel
+{
+
+void ParaCollection::add(std::unique_ptr<ParaWorld> 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

source/source_base/module_parallel/para_collection.h

@@ -0,0 +1,83 @@
+#ifndef PARA_COLLECTION_H
+#define PARA_COLLECTION_H
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#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<ParaWorld> 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 <typename T>
+    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<std::unique_ptr<ParaWorld>> worlds_; ///< owned domains
+};
+
+template <typename T>
+const T* ParaCollection::find_as(const std::string& tag) const
+{
+    for (const auto& world : worlds_)
+    {
+        if (world->tag() == tag)
+        {
+            return dynamic_cast<const T*>(world.get());
+        }
+    }
+    return nullptr;
+}
+
+} // namespace Parallel
+
+#endif // PARA_COLLECTION_H

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

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