Optimize MDCell migration output

Author: Fei Yang

docs/advanced/input_files/input-main.md

@@ -3518,7 +3518,7 @@
 ### md_dumpfreq
 
 - **Type**: Integer
-- **Description**: The output frequency of OUT.${suffix}/MD_dump in molecular dynamics calculations, which includes lattice and atomic information. Set to 0 to disable MD_dump output.
+- **Description**: The output frequency of OUT.${suffix}/MD_dump in molecular dynamics calculations, which includes lattice and atomic information. Set to 0 to disable MD_dump output; scalar MD progress remains printed to the terminal and running_md.log every step.
 - **Default**: 1
 
 ### md_out_force

docs/parameters.yaml

@@ -1465,7 +1465,7 @@ parameters:
     category: Molecular dynamics
     type: Integer
     description: |
-      The output frequency of OUT.${suffix}/MD_dump in molecular dynamics calculations, which includes lattice and atomic information. Set to 0 to disable MD_dump output.
+      The output frequency of OUT.${suffix}/MD_dump in molecular dynamics calculations, which includes lattice and atomic information. Set to 0 to disable MD_dump output; scalar MD progress remains printed to the terminal and running_md.log every step.
     default_value: "1"
     unit: ""
     availability: ""

source/source_cell/module_neighlist/domain_decomposition.cpp

@@ -11,6 +11,7 @@
 #include <limits>
 #include <map>
 #include <stdexcept>
+#include <utility>
 
 DomainDecomposition::DomainDecomposition()
     : comm_(MPI_COMM_NULL),
@@ -675,66 +676,114 @@ void DomainDecomposition::accumulate_ghost_forces(std::vector<LocalAtom>& owned_
 
 void DomainDecomposition::migrate_owned_atoms(std::vector<LocalAtom>& owned_atoms) const
 {
-    std::vector<std::vector<PackedAtom> > send_atoms(static_cast<std::size_t>(size_));
-    for (std::size_t i = 0; i < owned_atoms.size(); ++i)
+    const int direction_count = 6;
+    const int axis[direction_count] = {0, 0, 1, 1, 2, 2};
+    const int step[direction_count] = {-1, 1, -1, 1, -1, 1};
+    std::array<int, direction_count> neighbors;
+    for (int idir = 0; idir < direction_count; ++idir)
     {
-        LocalAtom atom = owned_atoms[i];
-        atom.frac = wrapped_frac_from_cart(atom.cart);
-        atom.cart = atom.frac * latvec_;
-        atom.owner_rank = owner_rank_from_frac(atom.frac);
-        const std::array<int, 3> no_shift = {{0, 0, 0}};
-        send_atoms[static_cast<std::size_t>(atom.owner_rank)].push_back(pack_atom(atom, no_shift));
+        std::array<int, 3> neighbor_coords = coords_;
+        neighbor_coords[axis[idir]] = positive_mod(neighbor_coords[axis[idir]] + step[idir], dims_[axis[idir]]);
+        neighbors[idir] = rank_from_coords(neighbor_coords);
     }
 
-    std::vector<int> send_counts(static_cast<std::size_t>(size_), 0);
-    std::vector<int> recv_counts(static_cast<std::size_t>(size_), 0);
-    for (int irank = 0; irank < size_; ++irank)
-    {
-        send_counts[static_cast<std::size_t>(irank)]
-            = static_cast<int>(send_atoms[static_cast<std::size_t>(irank)].size() * sizeof(PackedAtom));
-    }
-    MPI_Alltoall(&send_counts[0], 1, MPI_INT, &recv_counts[0], 1, MPI_INT, comm_);
+    std::vector<LocalAtom> pending_atoms;
+    pending_atoms.swap(owned_atoms);
+    std::vector<LocalAtom> retained_atoms;
+    retained_atoms.reserve(pending_atoms.size());
+    const std::array<int, 3> no_shift = {{0, 0, 0}};
 
-    std::vector<int> send_displs(static_cast<std::size_t>(size_), 0);
-    std::vector<int> recv_displs(static_cast<std::size_t>(size_), 0);
-    int total_send_bytes = 0;
-    int total_recv_bytes = 0;
-    for (int irank = 0; irank < size_; ++irank)
+    long long global_outgoing = 0;
+    do
     {
-        send_displs[static_cast<std::size_t>(irank)] = total_send_bytes;
-        recv_displs[static_cast<std::size_t>(irank)] = total_recv_bytes;
-        total_send_bytes += send_counts[static_cast<std::size_t>(irank)];
-        total_recv_bytes += recv_counts[static_cast<std::size_t>(irank)];
-    }
+        std::array<std::vector<PackedAtom>, direction_count> send_atoms;
+        for (std::size_t i = 0; i < pending_atoms.size(); ++i)
+        {
+            LocalAtom atom = std::move(pending_atoms[i]);
+            atom.frac = wrapped_frac_from_cart(atom.cart);
+            atom.cart = atom.frac * latvec_;
 
-    std::vector<PackedAtom> send_buffer(static_cast<std::size_t>(total_send_bytes / static_cast<int>(sizeof(PackedAtom))));
-    int send_index = 0;
-    for (int irank = 0; irank < size_; ++irank)
-    {
-        const std::vector<PackedAtom>& atoms = send_atoms[static_cast<std::size_t>(irank)];
-        for (std::size_t i = 0; i < atoms.size(); ++i)
+            std::array<int, 3> owner_coords;
+            const double frac[3] = {atom.frac.x, atom.frac.y, atom.frac.z};
+            for (int idim = 0; idim < 3; ++idim)
+            {
+                owner_coords[idim] = std::min(static_cast<int>(std::floor(frac[idim] * dims_[idim])), dims_[idim] - 1);
+            }
+            atom.owner_rank = rank_from_coords(owner_coords);
+            if (atom.owner_rank == rank_)
+            {
+                retained_atoms.push_back(std::move(atom));
+                continue;
+            }
+
+            int direction = -1;
+            for (int idim = 0; idim < 3 && direction < 0; ++idim)
+            {
+                int delta = owner_coords[idim] - coords_[idim];
+                if (delta > dims_[idim] / 2) delta -= dims_[idim];
+                if (delta < -dims_[idim] / 2) delta += dims_[idim];
+                if (delta != 0) direction = 2 * idim + (delta > 0 ? 1 : 0);
+            }
+            assert(direction >= 0);
+            send_atoms[direction].push_back(pack_atom(atom, no_shift));
+        }
+        pending_atoms.clear();
+
+        std::array<int, direction_count> send_counts;
+        std::array<int, direction_count> recv_counts;
+        long long local_outgoing = 0;
+        for (int idir = 0; idir < direction_count; ++idir)
         {
-            send_buffer[static_cast<std::size_t>(send_index++)] = atoms[i];
+            const std::size_t bytes = send_atoms[idir].size() * sizeof(PackedAtom);
+            if (bytes > static_cast<std::size_t>(std::numeric_limits<int>::max()))
+            {
+                throw std::overflow_error("DomainDecomposition migration send count exceeds int range.");
+            }
+            send_counts[idir] = static_cast<int>(bytes);
+            local_outgoing += static_cast<long long>(send_atoms[idir].size());
         }
-    }
+        MPI_Allreduce(&local_outgoing, &global_outgoing, 1, MPI_LONG_LONG, MPI_SUM, comm_);
+        if (global_outgoing == 0) break;
 
-    std::vector<PackedAtom> recv_buffer(static_cast<std::size_t>(total_recv_bytes / static_cast<int>(sizeof(PackedAtom))));
-    MPI_Alltoallv(total_send_bytes > 0 ? reinterpret_cast<const char*>(&send_buffer[0]) : 0,
-                  &send_counts[0],
-                  &send_displs[0],
-                  MPI_BYTE,
-                  total_recv_bytes > 0 ? reinterpret_cast<char*>(&recv_buffer[0]) : 0,
-                  &recv_counts[0],
-                  &recv_displs[0],
-                  MPI_BYTE,
-                  comm_);
+        std::array<MPI_Request, 2 * direction_count> requests;
+        for (int idir = 0; idir < direction_count; ++idir)
+        {
+            const int opposite = idir ^ 1;
+            MPI_Irecv(&recv_counts[idir], 1, MPI_INT, neighbors[idir], 100 + opposite, comm_, &requests[idir]);
+            MPI_Isend(&send_counts[idir], 1, MPI_INT, neighbors[idir], 100 + idir, comm_, &requests[direction_count + idir]);
+        }
+        MPI_Waitall(2 * direction_count, &requests[0], MPI_STATUSES_IGNORE);
 
-    owned_atoms.clear();
-    owned_atoms.reserve(recv_buffer.size());
-    for (std::size_t i = 0; i < recv_buffer.size(); ++i)
-    {
-        owned_atoms.push_back(unpack_owned_atom(recv_buffer[i]));
-    }
+        std::array<std::vector<PackedAtom>, direction_count> recv_atoms;
+        for (int idir = 0; idir < direction_count; ++idir)
+        {
+            if (recv_counts[idir] < 0 || recv_counts[idir] % static_cast<int>(sizeof(PackedAtom)) != 0)
+            {
+                throw std::runtime_error("Invalid DomainDecomposition migration receive count.");
+            }
+            recv_atoms[idir].resize(static_cast<std::size_t>(recv_counts[idir] / static_cast<int>(sizeof(PackedAtom))));
+        }
+
+        for (int idir = 0; idir < direction_count; ++idir)
+        {
+            const int opposite = idir ^ 1;
+            MPI_Irecv(recv_atoms[idir].empty() ? NULL : reinterpret_cast<char*>(&recv_atoms[idir][0]),
+                      recv_counts[idir], MPI_BYTE, neighbors[idir], 200 + opposite, comm_, &requests[idir]);
+            MPI_Isend(send_atoms[idir].empty() ? NULL : reinterpret_cast<const char*>(&send_atoms[idir][0]),
+                      send_counts[idir], MPI_BYTE, neighbors[idir], 200 + idir, comm_, &requests[direction_count + idir]);
+        }
+        MPI_Waitall(2 * direction_count, &requests[0], MPI_STATUSES_IGNORE);
+
+        for (int idir = 0; idir < direction_count; ++idir)
+        {
+            for (std::size_t i = 0; i < recv_atoms[idir].size(); ++i)
+            {
+                pending_atoms.push_back(unpack_owned_atom(recv_atoms[idir][i]));
+            }
+        }
+    } while (global_outgoing > 0);
+
+    owned_atoms.swap(retained_atoms);
 }
 
 #endif // __MPI

source/source_cell/module_neighlist/test/md_cell_migrate_mpi_test.cpp

@@ -6,6 +6,7 @@
 #include <mpi.h>
 
 #include <cstdint>
+#include <mpi.h>
 #include <string>
 #include <vector>
 
@@ -71,6 +72,15 @@ TEST(MdCellMigrateMpiTest, AtomCrossingDomainMigratesToNewOwner)
     ASSERT_EQ(mdcell.mpi_size(), size);
     if (size == 2)
     {
+        ASSERT_EQ(mdcell.nlocal(), 1);
+        mdcell.mutable_owned_atoms()[0].vel.x = static_cast<double>(rank + 1);
+        mdcell.mutable_owned_atoms()[0].force.y = static_cast<double>(rank + 3);
+        mdcell.migrate_owned_atoms();
+        ASSERT_EQ(mdcell.nlocal(), 1);
+        EXPECT_EQ(mdcell.owned_atoms()[0].owner_rank, rank);
+        EXPECT_EQ(mdcell.owned_atoms()[0].vel.x, static_cast<double>(rank + 1));
+        EXPECT_EQ(mdcell.owned_atoms()[0].force.y, static_cast<double>(rank + 3));
+
         if (rank == 0 && mdcell.nlocal() == 1)
         {
             mdcell.mutable_owned_atoms()[0].cart.x = 0.8;

source/source_io/module_output/output_log.cpp

@@ -293,8 +293,10 @@ void print_force(std::ofstream& ofs, const MDCell& cell, const std::string& name
     };
 
 #ifdef __MPI
-    const int rank = cell.mpi_rank();
-    const int size = cell.mpi_size();
+    int rank = 0;
+    int size = 1;
+    MPI_Comm_rank(cell.communicator(), &rank);
+    MPI_Comm_size(cell.communicator(), &size);
     if (rank != 0)
     {
         const int nlocal = cell.nlocal();

source/source_md/run_md.cpp

@@ -86,11 +86,10 @@ void md_line(MDCell& mdcell,
                                                      mdrun->frozen_freedom_);
         }
 
+        mdrun->print_md(GlobalV::ofs_running, PARAM.inp.cal_stress);
         if (param_in.mdp.md_dumpfreq > 0
             && (mdrun->step_ + mdrun->step_rst_) % param_in.mdp.md_dumpfreq == 0)
         {
-            mdrun->print_md(GlobalV::ofs_running, PARAM.inp.cal_stress);
-
             MD_func::dump_info(mdrun->step_ + mdrun->step_rst_,
                                PARAM.globalv.global_out_dir,
                                mdcell,

tests/04_FF/01_LJ_Anderson/result.ref

@@ -1,5 +1,5 @@
-etotref -2.234489982055187
-etotperatomref -0.0698278119
-totalforceref 2.434994
-totalstressref 28.475724
-totaltimeref 0.03
+etotref -2.239520155656396
+etotperatomref -0.0699850049
+totalforceref 2.331355
+totalstressref 28.360820
+totaltimeref 0.02

tests/04_FF/04_LJ_Langevin/result.ref

@@ -1,5 +1,5 @@
-etotref -2.242881714733034
-etotperatomref -0.0700900536
-totalforceref 2.351158
-totalstressref 28.243185
-totaltimeref 0.04
+etotref -2.24671151561839
+etotperatomref -0.0702097349
+totalforceref 2.191421
+totalstressref 28.146284
+totaltimeref 0.02