Refactor parallel communicators

source/Makefile.Objects

@@ -774,6 +774,7 @@ OBJS_MODULE_RI=conv_coulomb_pot_k.o\
 OBJS_PARALLEL=parallel_common.o\
     parallel_cell.o\
     parallel_global.o\
+    parallel_partition.o\
     parallel_comm.o\
     parallel_grid.o\
     parallel_kpoints.o\

source/source_base/CMakeLists.txt

@@ -57,6 +57,7 @@ add_library(
     parallel_reduce.cpp
     parallel_device.cpp
     parallel_grid.cpp
+    parallel_partition.cpp
     sph_bessel_tf.cpp
     cubic_spline.cpp
     parallel_2d.cpp

source/source_base/parallel_2d.h

@@ -5,7 +5,9 @@
 #include <cstdint>
 #include <vector>
 
-#include "source_base/parallel_comm.h"
+#ifdef __MPI
+#include <mpi.h>
+#endif
 
 /// @brief  This class packs the basic information of
 /// 2D-block-cyclic parallel distribution of an arbitrary matrix.

source/source_base/parallel_common.cpp

@@ -1,12 +1,25 @@
 #include "parallel_common.h"
 
+#include "source_base/parallel_reduce.h"
+
 #ifdef __MPI
 #include <mpi.h>
 #endif
 
-#include <cstring>
+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 <typename T>
+static void bcast_world_impl(T* object, const int n)
+{
+    MPI_Bcast(object, n, Parallel_Reduce::MPI_Type<T>::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<double>& object)
+void bcast_complex_double(std::complex<double>& 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<double>* object, const int n)
+void bcast_complex_double(std::complex<double>* 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<bool>(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

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"
 
@@ -284,6 +282,99 @@ void Parallel_Global::divide_pools(const int& NPROC,
     }
     return;
 }
+#endif
+
+ParallelPartition Parallel_Global::create_partition(int world_nproc,
+                                                    int my_rank,
+                                                    int kpar,
+                                                    int bndpar,
+                                                    int diag_np,
+                                                    int /*grid_np*/)
+{
+#ifdef __MPI
+    // ---- Build the nproc_in_pool vector (arithmetic, no MPI). ----
+    //
+    // divide_mpi_groups(..., num_groups=kpar, even=false) assigns
+    // the first `extra_procs` groups (extra_procs = world_nproc %
+    // kpar) to size (base+1) and the remaining (kpar-extra_procs)
+    // groups to size base, where base = world_nproc / kpar. No MPI;
+    // build the vector directly.
+    const int base = world_nproc / kpar;
+    const int extra_procs = world_nproc % kpar;
+    std::vector<int> nproc_in_pool(kpar, base);
+    for (int g = 0; g < extra_procs; ++g) { nproc_in_pool[g] = base + 1; }
+
+    // ---- Validation: sum of nproc_in_pool must equal world_nproc. ----
+    {
+        int total = 0;
+        for (int s : nproc_in_pool) { total += s; }
+        if (total != world_nproc)
+        {
+            ModuleBase::WARNING_QUIT("Parallel_Global::create_partition",
+                                     "internal: nproc_in_pool sum differs from world_nproc.");
+        }
+    }
+
+    // ---- Reproduce the exact legacy driver call order: ----
+    //   1. split_diag_world  (writes DIAG_WORLD)
+    //   2. split_grid_world (writes GRID_WORLD)
+    //   3. divide_pools     (writes POOL_WORLD, KP_WORLD, INT_BGROUP, BP_WORLD)
+    // Do NOT reorder these calls without a separate validation PR.
+    //
+    // diag_np == 0 is not meaningful; fall back to 1 so the
+    // even-partition guard in divide_mpi_groups (called inside
+    // split_diag_world / split_grid_world) does not fire on
+    // "num_groups == 0".
+    const int effective_diag_np = (diag_np > 0) ? diag_np : 1;
+    int drank = -1, dsize = -1, dcolor = -1;
+    Parallel_Global::split_diag_world(effective_diag_np, world_nproc, my_rank, drank, dsize, dcolor);
+    int grank = -1, gsize = -1;
+    Parallel_Global::split_grid_world(effective_diag_np, world_nproc, my_rank, grank, gsize);
+
+    // ---- Invoke the legacy divide_pools flow: it fills the 6 legacy
+    //      global communicators AND writes the GlobalV scalar fields
+    //      (NPROC_IN_POOL, RANK_IN_POOL, MY_POOL, etc.). ----
+    //
+    // NOTE(mohan): order matters. divide_pools internally calls
+    //   kpar_group .divide_group_comm(KPAR, false);
+    //   bndpar_group.divide_group_comm(BNDPAR, true);
+    //   -> sets POOL_WORLD, KP_WORLD, INT_BGROUP, BP_WORLD.
+    Parallel_Global::divide_pools(world_nproc, my_rank, bndpar, kpar,
+                                  GlobalV::NPROC_IN_BNDGROUP,
+                                  GlobalV::RANK_IN_BPGROUP,
+                                  GlobalV::MY_BNDGROUP,
+                                  GlobalV::NPROC_IN_POOL,
+                                  GlobalV::RANK_IN_POOL,
+                                  GlobalV::MY_POOL);
+
+    return ParallelPartition(world_nproc,
+                             my_rank,
+                             kpar,
+                             GlobalV::MY_POOL,
+                             GlobalV::RANK_IN_POOL,
+                             nproc_in_pool,
+                             bndpar,
+                             GlobalV::MY_BNDGROUP,
+                             GlobalV::RANK_IN_BPGROUP,
+                             GlobalV::NPROC_IN_BNDGROUP,
+                             POOL_WORLD,      // -> pw_world_comm        (legacy duped handle)
+                             KP_WORLD,        // -> kmesh_world_comm     (KP_WORLD alias back)
+                             INT_BGROUP,      // -> bsame_kdiff_world_comm
+                             BP_WORLD,        // -> bdiff_ksame_world_comm
+                             GRID_WORLD,      // -> rgrid_world_comm
+                             DIAG_WORLD,      // -> diag_world_comm
+                             MPI_COMM_NULL,   // -> matrix_world_comm (caller-filled later)
+                             MPI_COMM_NULL);  // -> atom_world_comm   (caller-filled later)
+#else
+    // Serial / non-MPI fallback: a single-process trivial partition.
+    (void)world_nproc;
+    (void)my_rank;
+    (void)kpar;
+    (void)bndpar;
+    (void)diag_np;
+    return ParallelPartition();
+#endif
+}
 
 void Parallel_Global::divide_mpi_groups(const int& procs,
                                         const int& num_groups,
@@ -317,7 +408,10 @@ void Parallel_Global::divide_mpi_groups(const int& procs,
     {
         std::cout << "Error: Number of processes (" << procs << ") must be evenly divisible by the number of groups ("
                   << num_groups << " in the even partition case)." << std::endl;
-        exit(1);
+        ModuleBase::WARNING_QUIT(
+            "Parallel_Global::divide_mpi_groups",
+            "Even partition requested but procs is not divisible by num_groups."
+        );
     }
 
     if(rank < extra_procs * (procs_in_group + 1))
@@ -334,5 +428,3 @@ void Parallel_Global::divide_mpi_groups(const int& procs,
         rank_in_group = (rank - extra_procs) % procs_in_group;
     }
 }
-
-#endif

source/source_base/parallel_global.h

@@ -7,6 +7,7 @@
 #define PARALLEL_GLOBAL_H
 
 #include "parallel_comm.h"
+#include "parallel_partition.h"
 
 namespace Parallel_Global
 {
@@ -19,6 +20,46 @@ extern int omp_number;
 // changed from read_mpi_parameters in 2024-1018
 void read_pal_param(int argc, char** argv, int& NPROC, int& NTHREAD_PER_PROC, int& MY_RANK);
 
+/**
+ * @brief Build a ParallelPartition snapshot for the given parallel parameters.
+ *
+ * This is the single factory that knows how to:
+ *   * split MPI_COMM_WORLD into KPAR k-pools via divide_group_comm(even=false);
+ *   * split each pool into BNDPAR band groups via divide_group_comm(even=true);
+ *   * derive INT_BGROUP (bsame_kdiff_world) / BP_WORLD (bdiff_ksame_world);
+ *   * split diag_np-based DIAG_WORLD and diag_np-grouped GRID_WORLD.
+ *
+ * The returned ParallelPartition::pw_world_comm (previously POOL_WORLD) is the
+ * smallest PW tile: the intersection of one k-pool and one band-group.
+ *
+ * matrix_world_comm and atom_world_comm are left as MPI_COMM_NULL in the
+ * returned value; callers that know which distributed view is required for
+ * a given step (Parallel_2D / Parallel_Orbitals / DomainDecomposition) are
+ * expected to fill them in from the appropriate view before passing the
+ * partition down.
+ *
+ * Note: Under __MPI the factory splits the real communicators as described
+ * below. The non-MPI build path returns the default ParallelPartition
+ * object, which already represents the single-process trivial partition.
+ *
+ * @param[in] world_nproc Size of MPI_COMM_WORLD
+ * @param[in] my_rank    Rank in MPI_COMM_WORLD
+ * @param[in] kpar       KPAR   from INPUT (k-point parallelism)
+ * @param[in] bndpar     BNDPAR from INPUT (band  parallelism)
+ * @param[in] diag_np    Number of diag worlds (also serves as group count
+ *                       for the real-space grid world: GRID_WORLD groups are
+ *                       the contiguous blocks produced by split_grid_world).
+ * @param[in] grid_np    Reserved; currently the real-space grid world is
+ *                       tied to diag_np via split_grid_world(diag_np, ...).
+ */
+ParallelPartition create_partition(int world_nproc,
+                                   int my_rank,
+                                   int kpar,
+                                   int bndpar,
+                                   int diag_np,
+                                   int grid_np);
+
+
 /**-------------------------------------------
  * call to split the "diago world"
  * the unit of first proc of each grid group