feat(module_parallel): add esolver/images domains and parent-comm splits

Add two top-level parallel domains for multi-image calculations
(e.g. NEB replicas with independent unit cells):

- para_esolver_world [color = image_id]: all processes of one
  esolver instance; every existing solver domain is now derived
  from it instead of being hard-wired to MPI_COMM_WORLD.
- para_images_world [color = rank_in_esolver]: cross-image
  communicator connecting corresponding ranks; MPI_COMM_NULL for
  nimage == 1 or uneven splits, following the KP_WORLD convention.

split_pools/split_diag_world/split_grid_world now take a parent
MPI_Comm, and setup_para_worlds takes nimage and builds the full
hierarchy: WORLD -> images split -> esolver domain -> kmesh/pw/
bgroup/diag/rgrid/matrix. The hierarchy tree is documented in
para_setup.h.

New MPI test (mpirun -np 4) verifies single-image degradation,
2-image split ranks/sizes, derived-domain sizes, and kpar=2 inside
images; multi-image tests skip under single-rank direct execution.

Author: abacus_fixer

source/source_base/module_parallel/para_setup.cpp

@@ -43,34 +43,43 @@ void divide_mpi_groups(int nproc, int num_groups, int rank, bool even,
 
 namespace {
 
-// Helper: split a communicator into ngroup sub-communicators
+// 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;       // intra-group communicator
-    MPI_Comm inter_comm;       // inter-group communicator (MPI_COMM_NULL if single group)
-    int ngroups;
-    int nprocs_in_group;
-    int my_group;
-    int rank_in_group;
+    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.group_comm = MPI_COMM_NULL;
-    res.inter_comm = MPI_COMM_NULL;
     res.ngroups = ngroup;
 
-    int gsize = 0, grank = 0;
+    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);
 
-    if (ngroup > 1)
+    // 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);
     }
@@ -80,37 +89,71 @@ GroupSplitResult split_comm_group(MPI_Comm parent, int ngroup, bool even)
 
 } // anonymous namespace
 
-void split_pools(int nproc, int my_rank, int bndpar, int kpar,
+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 && nproc % (bndpar * kpar) != 0)
+    if (bndpar > 1 && parent_size % (bndpar * kpar) != 0)
     {
-        std::cerr << "Error: NPROC (" << nproc
-                  << ") must be divisible by BNDPAR*KPAR ("
-                  << bndpar * kpar << ")." << std::endl;
+        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 WORLD into kpar pools
-    GroupSplitResult kpar_res = split_comm_group(MPI_COMM_WORLD, kpar, 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
+    // band parallelization: split each pool into bndpar groups.
     GroupSplitResult bndpar_res = split_comm_group(kpar_res.group_comm, bndpar, true);
 
-    // Set output indices
+    // 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
+    // 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)
+    // KP_WORLD: inter-pool communicator (same rank across pools).
     if (kpar_res.inter_comm != MPI_COMM_NULL)
     {
         MPI_Comm kp_comm;
@@ -122,48 +165,55 @@ void split_pools(int nproc, int my_rank, int bndpar, int kpar,
         kmesh_world = ParaWorld::make_mpi(ParaTag::kmesh, MPI_COMM_NULL);
     }
 
-    // Band group communicators
+    // 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(MPI_COMM_WORLD, my_bndgroup, rank_in_bpgroup, &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 = nproc;
-        rank_in_bpgroup = my_rank;
+        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(MPI_COMM_WORLD, &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(MPI_COMM_WORLD, my_rank, 0, &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 nproc, int my_rank,
+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, my_group = 0, rank_in_group = 0;
-    divide_mpi_groups(nproc, diag_np, my_rank, false,
+    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(MPI_COMM_WORLD, my_group, rank_in_group, &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);
@@ -172,38 +222,67 @@ ParaWorld split_diag_world(int diag_np, int nproc, int my_rank,
     return ParaWorld::make_mpi(ParaTag::diag, diag_comm);
 }
 
-ParaWorld split_grid_world(int diag_np, int nproc, int my_rank,
+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, my_group = 0, rank_in_group = 0;
-    divide_mpi_groups(nproc, diag_np, my_rank, false,
+    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(MPI_COMM_WORLD, my_group, rank_in_group, &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 bndpar, int kpar, int diag_np)
+ParaCollection setup_para_worlds(int nproc, int my_rank, int nimage,
+                                 int bndpar, int kpar, int diag_np)
 {
     ParaCollection worlds;
 
-    // 1. POOL_WORLD + KP_WORLD + band group comms
-    int nproc_in_pool, rank_in_pool, my_pool;
-    int nproc_in_bndgroup, rank_in_bpgroup, my_bndgroup;
+    // 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(nproc, my_rank, bndpar, kpar,
+    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);
@@ -213,17 +292,22 @@ ParaCollection setup_para_worlds(int nproc, int my_rank, int bndpar, int kpar, i
     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. DIAG_WORLD
-    int drank, dsize, dcolor;
-    ParaWorld diag_world = split_diag_world(diag_np, nproc, my_rank, drank, dsize, dcolor);
+    // 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. GRID_WORLD
-    int grank, gsize;
-    ParaWorld grid_world = split_grid_world(diag_np, nproc, my_rank, grank, gsize);
+    // 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, will get its own split later)
+    // 4. Matrix domain: serial for now until its own 2D-grid split lands.
     worlds.add(ParaWorld::make_serial(ParaTag::matrix));
 
     return worlds;