Merge remote-tracking branch 'upstream/develop' into 2026-08-10-b

Author: abacus_fixer

source/source_cell/module_neighlist/bin_manager.cpp

@@ -5,6 +5,15 @@
 #include <stdexcept>
 #include "bin_manager.h"
 
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
+namespace
+{
+constexpr int neighbor_build_openmp_threshold = 256;
+}
+
 // ========== Bin class implementation ==========
 
 const std::vector<ModuleNeighList::LocalAtomIndex>& Bin::get_atom_indices() const {
@@ -176,6 +185,63 @@ int BinManager::bin_index(int ix, int iy, int iz) const {
     return ix * nbiny_ * nbinz_ + iy * nbinz_ + iz;
 }
 
+template <typename Emit>
+void BinManager::visit_neighbors(const NeighborAtom& atom,
+                                 const std::vector<NeighborAtom>& binned_atoms,
+                                 double sradius2,
+                                 const Emit& emit) const
+{
+    const int ix = std::min(
+        std::max(int((atom.position_x - x_min_) / bin_sizex_), 0),
+        nbinx_ - 1
+    );
+
+    const int iy = std::min(
+        std::max(int((atom.position_y - y_min_) / bin_sizey_), 0),
+        nbiny_ - 1
+    );
+
+    const int iz = std::min(
+        std::max(int((atom.position_z - z_min_) / bin_sizez_), 0),
+        nbinz_ - 1
+    );
+
+    for (int dx = -1; dx <= 1; dx++)
+    {
+        for (int dy = -1; dy <= 1; dy++)
+        {
+            for (int dz = -1; dz <= 1; dz++)
+            {
+                const int jx = ix + dx;
+                const int jy = iy + dy;
+                const int jz = iz + dz;
+
+                if (jx < 0 || jx >= nbinx_ ||
+                    jy < 0 || jy >= nbiny_ ||
+                    jz < 0 || jz >= nbinz_)
+                {
+                    continue;
+                }
+
+                const int nidx = bin_index(jx, jy, jz);
+                for (const ModuleNeighList::LocalAtomIndex binned_atom_index : bins_[nidx].get_atom_indices())
+                {
+                    const NeighborAtom& natom = binned_atoms[static_cast<std::size_t>(binned_atom_index)];
+                    const double delta_x = atom.position_x - natom.position_x;
+                    const double delta_y = atom.position_y - natom.position_y;
+                    const double delta_z = atom.position_z - natom.position_z;
+                    const double dist2 = delta_x * delta_x + delta_y * delta_y + delta_z * delta_z;
+
+                    if (natom.atom_id != atom.atom_id && dist2 <= sradius2)
+                    {
+                        emit(natom.atom_id);
+                    }
+                }
+            }
+        }
+    }
+}
+
 void BinManager::build_atom_neighbors(
     NeighborList& neighbor_list,
     const std::vector<NeighborAtom>& atoms,
@@ -184,71 +250,63 @@ void BinManager::build_atom_neighbors(
 {
     assert(atoms.size() == static_cast<size_t>(neighbor_list.get_nlocal()));
 
-    double sradius2 = sradius_ * sradius_;
+    const double sradius2 = sradius_ * sradius_;
 
     neighbor_list.reset();
 
-    std::vector<int> neigh_tmp;
-
     const int nlocal = neighbor_list.get_nlocal();
-    for (int i = 0; i < nlocal; i++)
-    {
-        neigh_tmp.clear();
-        const NeighborAtom& atom = atoms[i];
-
-        int ix = std::min(
-            std::max(int((atom.position_x - x_min_) / bin_sizex_), 0),
-            nbinx_ - 1
-        );
 
-        int iy = std::min(
-            std::max(int((atom.position_y - y_min_) / bin_sizey_), 0),
-            nbiny_ - 1
-        );
-
-        int iz = std::min(
-            std::max(int((atom.position_z - z_min_) / bin_sizez_), 0),
-            nbinz_ - 1
-        );
+#ifdef _OPENMP
+    const bool use_parallel = nlocal >= neighbor_build_openmp_threshold && omp_get_max_threads() > 1;
+    if (use_parallel)
+    {
+        std::vector<std::size_t> neighbor_counts(static_cast<std::size_t>(nlocal), 0);
 
-        for (int dx = -1; dx <= 1; dx++)
+#pragma omp parallel for schedule(static)
+        for (int i = 0; i < nlocal; i++)
         {
-            for (int dy = -1; dy <= 1; dy++)
-            {
-                for (int dz = -1; dz <= 1; dz++)
-                {
-                    int jx = ix + dx;
-                    int jy = iy + dy;
-                    int jz = iz + dz;
-
-                    if (jx < 0 || jx >= nbinx_ ||
-                        jy < 0 || jy >= nbiny_ ||
-                        jz < 0 || jz >= nbinz_)
-                        continue;
-
-                    int nidx = bin_index(jx, jy, jz);
+            std::size_t count = 0;
+            visit_neighbors(atoms[i], binned_atoms, sradius2,
+                            [&count](ModuleNeighList::LocalAtomIndex) { ++count; });
+            neighbor_counts[static_cast<std::size_t>(i)] = count;
+        }
 
-                    for (const ModuleNeighList::LocalAtomIndex binned_atom_index : bins_[nidx].get_atom_indices())
-                    {
-                        const NeighborAtom& natom = binned_atoms[static_cast<std::size_t>(binned_atom_index)];
-                        double dx = atom.position_x - natom.position_x;
-                        double dy = atom.position_y - natom.position_y;
-                        double dz = atom.position_z - natom.position_z;
+        for (int i = 0; i < nlocal; i++)
+        {
+            const int n = ModuleNeighList::checked_int_size(
+                neighbor_counts[static_cast<std::size_t>(i)],
+                "BinManager neighbor count"
+            );
+            neighbor_list.firstneigh_[i] = neighbor_list.allocator_.allocate(n);
+            neighbor_list.numneigh_[i] = n;
+        }
 
-                        double dist2 = dx * dx + dy * dy + dz * dz;
+#pragma omp parallel for schedule(static)
+        for (int i = 0; i < nlocal; i++)
+        {
+            int* ptr = neighbor_list.firstneigh_[i];
+            int k = 0;
+            visit_neighbors(atoms[i], binned_atoms, sradius2,
+                            [&](ModuleNeighList::LocalAtomIndex atom_id)
+                            {
+                                assert(ptr != nullptr);
+                                ptr[k++] = atom_id;
+                            });
+            assert(k == neighbor_list.numneigh_[i]);
+        }
+        return;
+    }
+#endif
 
-                        if (natom.atom_id == atom.atom_id)
-                        {
-                            continue;
-                        }
-                        if (dist2 <= sradius2)
+    std::vector<int> neigh_tmp;
+    for (int i = 0; i < nlocal; i++)
+    {
+        neigh_tmp.clear();
+        visit_neighbors(atoms[i], binned_atoms, sradius2,
+                        [&neigh_tmp](ModuleNeighList::LocalAtomIndex atom_id)
                         {
-                            neigh_tmp.push_back(natom.atom_id);
-                        }
-                    }
-                }
-            }
-        }
+                            neigh_tmp.push_back(atom_id);
+                        });
 
         const int n = ModuleNeighList::checked_int_size(neigh_tmp.size(), "BinManager neighbor count");
 

source/source_cell/module_neighlist/bin_manager.h

@@ -198,6 +198,21 @@ class BinManager
      * @return Flat index in the bins_ array.
      */
     int bin_index(int ix, int iy, int iz) const;
+
+    /**
+     * @brief Visit neighbors of one atom in the existing deterministic bin order.
+     *
+     * @tparam Emit Callable accepting a rank-local neighbor atom ID.
+     * @param atom Atom used as the neighbor-list center.
+     * @param binned_atoms All atoms assigned to bins by do_binning().
+     * @param sradius2 Squared search radius.
+     * @param emit Callback invoked once for every accepted neighbor.
+     */
+    template <typename Emit>
+    void visit_neighbors(const NeighborAtom& atom,
+                         const std::vector<NeighborAtom>& binned_atoms,
+                         double sradius2,
+                         const Emit& emit) const;
 };
 
 #endif // BIN_MANAGER_H

source/source_cell/module_neighlist/neighbor_search.cpp

@@ -1,6 +1,7 @@
 #include "source_cell/module_neighlist/neighbor_search.h"
 #include <cmath>
 #include <algorithm>
+#include <cstdint>
 #include <limits>
 #include <cassert>
 #include <array>
@@ -112,13 +113,14 @@ void NeighborSearch::init(const AtomProvider& ucell, double sr)
     {
         for (int j = 0; j < ucell.get_na(i); j++)
         {
+            const ModuleBase::Vector3<double> position = ucell.get_tau(i, j);
             const ModuleNeighList::LocalAtomIndex atom_count
                 = ModuleNeighList::checked_local_atom_index(all_atoms_.size(),
                                                             "NeighborSearch atom id");
             NeighborAtom atom(
-                ucell.get_tau(i,j).x,
-                ucell.get_tau(i,j).y,
-                ucell.get_tau(i,j).z,
+                position.x,
+                position.y,
+                position.z,
                 i,
                 j,
                 atom_count
@@ -189,37 +191,111 @@ void NeighborSearch::check_expand_condition(const AtomProvider& ucell, int& glay
 
 void NeighborSearch::set_member_variables(const AtomProvider& ucell, int glayerX_minus, int glayerX, int glayerY_minus, int glayerY, int glayerZ_minus, int glayerZ)
 {
-    ModuleBase::Vector3<double> vec1(ucell.get_latvec().e11, ucell.get_latvec().e12, ucell.get_latvec().e13);
-    ModuleBase::Vector3<double> vec2(ucell.get_latvec().e21, ucell.get_latvec().e22, ucell.get_latvec().e23);
-    ModuleBase::Vector3<double> vec3(ucell.get_latvec().e31, ucell.get_latvec().e32, ucell.get_latvec().e33);
+    if (glayerX_minus < 0 || glayerX < 0 || glayerY_minus < 0 || glayerY < 0
+        || glayerZ_minus < 0 || glayerZ < 0)
+    {
+        throw std::invalid_argument("NeighborSearch periodic image layers must be non-negative.");
+    }
 
-    for (int ix = -glayerX_minus; ix < glayerX; ix++)
+    const ModuleBase::Matrix3& lattice = ucell.get_latvec();
+    const ModuleBase::Vector3<double> vec1(lattice.e11, lattice.e12, lattice.e13);
+    const ModuleBase::Vector3<double> vec2(lattice.e21, lattice.e22, lattice.e23);
+    const ModuleBase::Vector3<double> vec3(lattice.e31, lattice.e32, lattice.e33);
+
+    const std::size_t image_count_x
+        = ModuleNeighList::checked_size_sum(static_cast<std::size_t>(glayerX_minus),
+                                            static_cast<std::size_t>(glayerX),
+                                            "NeighborSearch x image count");
+    const std::size_t image_count_y
+        = ModuleNeighList::checked_size_sum(static_cast<std::size_t>(glayerY_minus),
+                                            static_cast<std::size_t>(glayerY),
+                                            "NeighborSearch y image count");
+    const std::size_t image_count_z
+        = ModuleNeighList::checked_size_sum(static_cast<std::size_t>(glayerZ_minus),
+                                            static_cast<std::size_t>(glayerZ),
+                                            "NeighborSearch z image count");
+    const std::size_t image_count_yz
+        = ModuleNeighList::checked_size_product(image_count_y,
+                                                image_count_z,
+                                                "NeighborSearch yz image count");
+    const std::size_t image_count
+        = ModuleNeighList::checked_size_product(image_count_x,
+                                                image_count_yz,
+                                                "NeighborSearch periodic image count");
+    if (image_count == 0)
     {
-        for (int iy = -glayerY_minus; iy < glayerY; iy++)
-        {
-            for (int iz = -glayerZ_minus; iz < glayerZ; iz++)
-            {
-                if(ix==0 && iy==0 && iz==0)
-                {
-                    continue;
-                }
-                for (int i = 0; i < ucell.get_ntype(); i++)
-                {
-                    for (int j = 0; j < ucell.get_na(i); j++)
-                    {
-                        double atom_x = ucell.get_tau(i,j).x + vec1[0] * ix + vec2[0] * iy + vec3[0] * iz;
-                        double atom_y = ucell.get_tau(i,j).y + vec1[1] * ix + vec2[1] * iy + vec3[1] * iz;
-                        double atom_z = ucell.get_tau(i,j).z + vec1[2] * ix + vec2[2] * iy + vec3[2] * iz;
-
-                        const ModuleNeighList::LocalAtomIndex atom_count
-                            = ModuleNeighList::checked_local_atom_index(all_atoms_.size(),
-                                                                        "NeighborSearch atom id");
-                        NeighborAtom atom(atom_x, atom_y, atom_z, i, j, atom_count);
-                        ghost_atoms_.push_back(atom);
-                        all_atoms_.push_back(atom);
-                    }
-                }
-            }
-        }
+        throw std::invalid_argument("NeighborSearch periodic image range must contain the primary cell.");
+    }
+
+    const std::size_t origin_xy
+        = ModuleNeighList::checked_size_sum(
+            ModuleNeighList::checked_size_product(static_cast<std::size_t>(glayerX_minus),
+                                                  image_count_y,
+                                                  "NeighborSearch primary image index"),
+            static_cast<std::size_t>(glayerY_minus),
+            "NeighborSearch primary image index");
+    const std::size_t origin_image
+        = ModuleNeighList::checked_size_sum(
+            ModuleNeighList::checked_size_product(origin_xy,
+                                                  image_count_z,
+                                                  "NeighborSearch primary image index"),
+            static_cast<std::size_t>(glayerZ_minus),
+            "NeighborSearch primary image index");
+
+    const std::size_t base_atom_count = inside_atoms_.size();
+    const std::size_t ghost_image_count = image_count - 1;
+    const std::size_t generated_atom_count
+        = ModuleNeighList::checked_size_product(ghost_image_count,
+                                                base_atom_count,
+                                                "NeighborSearch generated atom count");
+    const std::size_t final_atom_count
+        = ModuleNeighList::checked_size_sum(base_atom_count,
+                                            generated_atom_count,
+                                            "NeighborSearch total atom count");
+    if (final_atom_count > static_cast<std::size_t>(std::numeric_limits<ModuleNeighList::LocalAtomIndex>::max()))
+    {
+        throw std::overflow_error("NeighborSearch total atom count exceeds local atom index range.");
+    }
+    if (generated_atom_count == 0)
+    {
+        return;
+    }
+
+    const NeighborAtom placeholder(0.0, 0.0, 0.0, 0, 0, 0);
+    all_atoms_.insert(all_atoms_.end(), generated_atom_count, placeholder);
+    ghost_atoms_.assign(generated_atom_count, placeholder);
+
+    // Thread creation costs more than it saves for small unit cells.
+    const std::size_t parallel_threshold = 100000;
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static) if(generated_atom_count >= parallel_threshold)
+#endif
+    for (std::int64_t generated_index = 0;
+         generated_index < static_cast<std::int64_t>(generated_atom_count);
+         ++generated_index)
+    {
+        const std::size_t output_index = static_cast<std::size_t>(generated_index);
+        const std::size_t ghost_image = output_index / base_atom_count;
+        const std::size_t base_atom = output_index % base_atom_count;
+        const std::size_t image = ghost_image < origin_image ? ghost_image : ghost_image + 1;
+        const std::size_t image_x = image / image_count_yz;
+        const std::size_t image_yz = image % image_count_yz;
+        const std::size_t image_y = image_yz / image_count_z;
+        const std::size_t image_z = image_yz % image_count_z;
+        const double ix = static_cast<double>(image_x) - static_cast<double>(glayerX_minus);
+        const double iy = static_cast<double>(image_y) - static_cast<double>(glayerY_minus);
+        const double iz = static_cast<double>(image_z) - static_cast<double>(glayerZ_minus);
+
+        const NeighborAtom& source = inside_atoms_[base_atom];
+        const ModuleNeighList::LocalAtomIndex atom_id
+            = static_cast<ModuleNeighList::LocalAtomIndex>(base_atom_count + output_index);
+        const NeighborAtom atom(source.position_x + vec1[0] * ix + vec2[0] * iy + vec3[0] * iz,
+                                source.position_y + vec1[1] * ix + vec2[1] * iy + vec3[1] * iz,
+                                source.position_z + vec1[2] * ix + vec2[2] * iy + vec3[2] * iz,
+                                source.atom_type,
+                                source.atom_index,
+                                atom_id);
+        ghost_atoms_[output_index] = atom;
+        all_atoms_[base_atom_count + output_index] = atom;
     }
 }