From e06b44c55a4bb751ea7de029dc891bcb8d31c814 Mon Sep 17 00:00:00 2001 From: Audrey-777 <3043005175chy@gmail.com> Date: Fri, 26 Jun 2026 20:43:57 +0800 Subject: [PATCH 1/2] perf: parallelize neighbor list construction --- .../module_neighlist/bin_manager.cpp | 105 ++++++++++++++---- .../module_neighlist/neighbor_search.cpp | 9 +- 2 files changed, 92 insertions(+), 22 deletions(-) diff --git a/source/source_cell/module_neighlist/bin_manager.cpp b/source/source_cell/module_neighlist/bin_manager.cpp index 1077b91dd75..5217e944f0e 100644 --- a/source/source_cell/module_neighlist/bin_manager.cpp +++ b/source/source_cell/module_neighlist/bin_manager.cpp @@ -3,6 +3,16 @@ #include #include #include "bin_manager.h" +#include "source_base/timer.h" + +#ifdef _OPENMP +#include +#endif + +namespace +{ +constexpr int neighbor_build_openmp_threshold = 256; +} // ========== Bin class implementation ========== @@ -69,6 +79,7 @@ void BinManager::init_bins( const std::vector& ghost_atoms ) { + ModuleBase::timer::start("BinManager", "init_bins"); sradius_ = sr; if(inside_atoms.empty() && ghost_atoms.empty()) { @@ -77,6 +88,7 @@ void BinManager::init_bins( nbinx_ = nbiny_ = nbinz_ = 1; bins_.clear(); bins_.resize(1); + ModuleBase::timer::end("BinManager", "init_bins"); return; } @@ -129,6 +141,7 @@ void BinManager::init_bins( } } } + ModuleBase::timer::end("BinManager", "init_bins"); } void BinManager::do_binning( @@ -136,6 +149,7 @@ void BinManager::do_binning( const std::vector& ghost_atoms ) { + ModuleBase::timer::start("BinManager", "do_binning"); auto bin_atom = [&](const NeighborAtom& atom) { int ix = std::min( @@ -160,6 +174,7 @@ void BinManager::do_binning( for (const auto& atom : inside_atoms) bin_atom(atom); for (const auto& atom : ghost_atoms) bin_atom(atom); + ModuleBase::timer::end("BinManager", "do_binning"); } int BinManager::bin_index(int ix, int iy, int iz) const { @@ -171,29 +186,25 @@ void BinManager::build_atom_neighbors( std::vector& atoms ) { + ModuleBase::timer::start("BinManager", "build_atom_neighbors"); assert(atoms.size() == static_cast(neighbor_list.get_nlocal())); - double sradius2 = sradius_ * sradius_; - - neighbor_list.reset(); - - std::vector neigh_tmp; + const int nlocal = static_cast(atoms.size()); + const double sradius2 = sradius_ * sradius_; - for (int i = 0; i < atoms.size(); i++) + auto visit_neighbors = [&](const int i, const auto& emit) { - neigh_tmp.clear(); - - int ix = std::min( + const int ix = std::min( std::max(int((atoms[i].position_x - x_min_) / bin_sizex_), 0), nbinx_ - 1 ); - int iy = std::min( + const int iy = std::min( std::max(int((atoms[i].position_y - y_min_) / bin_sizey_), 0), nbiny_ - 1 ); - int iz = std::min( + const int iz = std::min( std::max(int((atoms[i].position_z - z_min_) / bin_sizez_), 0), nbinz_ - 1 ); @@ -204,33 +215,84 @@ void BinManager::build_atom_neighbors( { for (int dz = -1; dz <= 1; dz++) { - int jx = ix + dx; - int jy = iy + dy; - int jz = iz + 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; + } - int nidx = bin_index(jx, jy, jz); + const int nidx = bin_index(jx, jy, jz); for (const NeighborAtom& natom : bins_[nidx].get_atoms()) { - double dx = atoms[i].position_x - natom.position_x; - double dy = atoms[i].position_y - natom.position_y; - double dz = atoms[i].position_z - natom.position_z; + const double dx = atoms[i].position_x - natom.position_x; + const double dy = atoms[i].position_y - natom.position_y; + const double dz = atoms[i].position_z - natom.position_z; - double dist2 = dx * dx + dy * dy + dz * dz; + const double dist2 = dx * dx + dy * dy + dz * dz; if (dist2 <= sradius2 && dist2 != 0) { - neigh_tmp.push_back(natom.atom_id); + emit(natom.atom_id); } } } } } + }; + + neighbor_list.reset(); + +#ifdef _OPENMP + const bool use_parallel = nlocal >= neighbor_build_openmp_threshold && omp_get_max_threads() > 1; + if (use_parallel) + { + std::vector neighbor_counts(nlocal, 0); + +#pragma omp parallel for schedule(static) + for (int i = 0; i < nlocal; i++) + { + int count = 0; + visit_neighbors(i, [&](const int) { ++count; }); + neighbor_counts[i] = count; + } + + for (int i = 0; i < nlocal; i++) + { + const int n = neighbor_counts[i]; + neighbor_list.firstneigh_[i] = neighbor_list.allocator_.allocate(n); + neighbor_list.numneigh_[i] = n; + } + +#pragma omp parallel for schedule(static) + for (int i = 0; i < nlocal; i++) + { + int* ptr = neighbor_list.firstneigh_[i]; + int k = 0; + visit_neighbors(i, [&](const int atom_id) + { + assert(ptr != nullptr); + ptr[k++] = atom_id; + }); + assert(k == neighbor_counts[i]); + } + + ModuleBase::timer::end("BinManager", "build_atom_neighbors"); + return; + } +#endif + + std::vector neigh_tmp; + + for (int i = 0; i < nlocal; i++) + { + neigh_tmp.clear(); + visit_neighbors(i, [&](const int atom_id) { neigh_tmp.push_back(atom_id); }); int n = neigh_tmp.size(); @@ -245,6 +307,7 @@ void BinManager::build_atom_neighbors( neighbor_list.firstneigh_[i] = ptr; neighbor_list.numneigh_[i] = n; } + ModuleBase::timer::end("BinManager", "build_atom_neighbors"); } void BinManager::clear() @@ -255,4 +318,4 @@ void BinManager::clear() } bins_.clear(); -} \ No newline at end of file +} diff --git a/source/source_cell/module_neighlist/neighbor_search.cpp b/source/source_cell/module_neighlist/neighbor_search.cpp index 912515bf9d5..63333cf3907 100644 --- a/source/source_cell/module_neighlist/neighbor_search.cpp +++ b/source/source_cell/module_neighlist/neighbor_search.cpp @@ -3,6 +3,7 @@ #include #include #include +#include "source_base/timer.h" // ========== Getter methods ========== @@ -171,6 +172,7 @@ void NeighborSearch::check_expand_condition(const AtomProvider& ucell) void NeighborSearch::set_member_variables(const AtomProvider& ucell) { + ModuleBase::timer::start("NeighborSearch", "set_member_variables"); all_atoms_.clear(); ModuleBase::Vector3 vec1(ucell.get_latvec().e11, ucell.get_latvec().e12, ucell.get_latvec().e13); @@ -209,12 +211,14 @@ void NeighborSearch::set_member_variables(const AtomProvider& ucell) } } } + ModuleBase::timer::end("NeighborSearch", "set_member_variables"); } // ========== Main public interface ========== void NeighborSearch::init(const AtomProvider& ucell, double sr, int mpi_rank) { + ModuleBase::timer::start("NeighborSearch", "init"); // clear possible residual data from previous runs inside_atoms_.clear(); ghost_atoms_.clear(); @@ -322,13 +326,16 @@ void NeighborSearch::init(const AtomProvider& ucell, double sr, int mpi_rank) } neighbor_list_.initialize(inside_atoms_.size(), all_atoms_.size() * neighbor_reserve_factor); + ModuleBase::timer::end("NeighborSearch", "init"); } void NeighborSearch::build_neighbors() { + ModuleBase::timer::start("NeighborSearch", "build_neighbors"); bin_manager_.init_bins(search_radius_, inside_atoms_, ghost_atoms_); bin_manager_.do_binning(inside_atoms_, ghost_atoms_); bin_manager_.build_atom_neighbors(neighbor_list_, inside_atoms_); + ModuleBase::timer::end("NeighborSearch", "build_neighbors"); } // ========== Utility methods ========== @@ -374,4 +381,4 @@ void NeighborSearch::decompose(int mpi_size, int &nx, int &ny, int &nz) break; } } -} \ No newline at end of file +} From 6dba4951c5cee9492c8f89eacaaddafa174b24f6 Mon Sep 17 00:00:00 2001 From: Audrey-777 <3043005175chy@gmail.com> Date: Fri, 26 Jun 2026 20:52:18 +0800 Subject: [PATCH 2/2] fix: keep neighbor list OpenMP code C++11 compatible --- .../module_neighlist/bin_manager.cpp | 114 +++++++++--------- .../module_neighlist/bin_manager.h | 10 +- 2 files changed, 69 insertions(+), 55 deletions(-) diff --git a/source/source_cell/module_neighlist/bin_manager.cpp b/source/source_cell/module_neighlist/bin_manager.cpp index 5217e944f0e..545ff8da21b 100644 --- a/source/source_cell/module_neighlist/bin_manager.cpp +++ b/source/source_cell/module_neighlist/bin_manager.cpp @@ -181,70 +181,76 @@ int BinManager::bin_index(int ix, int iy, int iz) const { return ix * nbiny_ * nbinz_ + iy * nbinz_ + iz; } -void BinManager::build_atom_neighbors( - NeighborList& neighbor_list, - std::vector& atoms -) +template +void BinManager::visit_neighbors( + int i, + const std::vector& atoms, + double sradius2, + const Emit& emit +) const { - ModuleBase::timer::start("BinManager", "build_atom_neighbors"); - assert(atoms.size() == static_cast(neighbor_list.get_nlocal())); - - const int nlocal = static_cast(atoms.size()); - const double sradius2 = sradius_ * sradius_; - - auto visit_neighbors = [&](const int i, const auto& emit) + const int ix = std::min( + std::max(int((atoms[i].position_x - x_min_) / bin_sizex_), 0), + nbinx_ - 1 + ); + + const int iy = std::min( + std::max(int((atoms[i].position_y - y_min_) / bin_sizey_), 0), + nbiny_ - 1 + ); + + const int iz = std::min( + std::max(int((atoms[i].position_z - z_min_) / bin_sizez_), 0), + nbinz_ - 1 + ); + + for (int dx = -1; dx <= 1; dx++) { - const int ix = std::min( - std::max(int((atoms[i].position_x - x_min_) / bin_sizex_), 0), - nbinx_ - 1 - ); - - const int iy = std::min( - std::max(int((atoms[i].position_y - y_min_) / bin_sizey_), 0), - nbiny_ - 1 - ); - - const int iz = std::min( - std::max(int((atoms[i].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 dy = -1; dy <= 1; dy++) + for (int dz = -1; dz <= 1; dz++) { - for (int dz = -1; dz <= 1; dz++) - { - const int jx = ix + dx; - const int jy = iy + dy; - const int jz = iz + 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; - } + if (jx < 0 || jx >= nbinx_ || + jy < 0 || jy >= nbiny_ || + jz < 0 || jz >= nbinz_) + { + continue; + } - const int nidx = bin_index(jx, jy, jz); + const int nidx = bin_index(jx, jy, jz); - for (const NeighborAtom& natom : bins_[nidx].get_atoms()) - { - const double dx = atoms[i].position_x - natom.position_x; - const double dy = atoms[i].position_y - natom.position_y; - const double dz = atoms[i].position_z - natom.position_z; + for (const NeighborAtom& natom : bins_[nidx].get_atoms()) + { + const double dx = atoms[i].position_x - natom.position_x; + const double dy = atoms[i].position_y - natom.position_y; + const double dz = atoms[i].position_z - natom.position_z; - const double dist2 = dx * dx + dy * dy + dz * dz; + const double dist2 = dx * dx + dy * dy + dz * dz; - if (dist2 <= sradius2 && dist2 != 0) - { - emit(natom.atom_id); - } + if (dist2 <= sradius2 && dist2 != 0) + { + emit(natom.atom_id); } } } } - }; + } +} + +void BinManager::build_atom_neighbors( + NeighborList& neighbor_list, + std::vector& atoms +) +{ + ModuleBase::timer::start("BinManager", "build_atom_neighbors"); + assert(atoms.size() == static_cast(neighbor_list.get_nlocal())); + + const int nlocal = static_cast(atoms.size()); + const double sradius2 = sradius_ * sradius_; neighbor_list.reset(); @@ -258,7 +264,7 @@ void BinManager::build_atom_neighbors( for (int i = 0; i < nlocal; i++) { int count = 0; - visit_neighbors(i, [&](const int) { ++count; }); + visit_neighbors(i, atoms, sradius2, [&](const int) { ++count; }); neighbor_counts[i] = count; } @@ -274,7 +280,7 @@ void BinManager::build_atom_neighbors( { int* ptr = neighbor_list.firstneigh_[i]; int k = 0; - visit_neighbors(i, [&](const int atom_id) + visit_neighbors(i, atoms, sradius2, [&](const int atom_id) { assert(ptr != nullptr); ptr[k++] = atom_id; @@ -292,7 +298,7 @@ void BinManager::build_atom_neighbors( for (int i = 0; i < nlocal; i++) { neigh_tmp.clear(); - visit_neighbors(i, [&](const int atom_id) { neigh_tmp.push_back(atom_id); }); + visit_neighbors(i, atoms, sradius2, [&](const int atom_id) { neigh_tmp.push_back(atom_id); }); int n = neigh_tmp.size(); diff --git a/source/source_cell/module_neighlist/bin_manager.h b/source/source_cell/module_neighlist/bin_manager.h index 22b94d394a0..e069ade809c 100644 --- a/source/source_cell/module_neighlist/bin_manager.h +++ b/source/source_cell/module_neighlist/bin_manager.h @@ -218,6 +218,14 @@ class BinManager * @return Flat index in the bins_ array. */ int bin_index(int ix, int iy, int iz) const; + + template + void visit_neighbors( + int i, + const std::vector& atoms, + double sradius2, + const Emit& emit + ) const; }; -#endif // BIN_MANAGER_H \ No newline at end of file +#endif // BIN_MANAGER_H