From f5ea63504617f69bdda79f4e74ae306fcd749336 Mon Sep 17 00:00:00 2001 From: linpz Date: Wed, 5 Aug 2026 11:40:20 +0800 Subject: [PATCH] 1. add OpenMP in XC_Functional_Libxc::v_xc_libxc() 2. add OpenMP in RI_2D_Comm::split_m2D_ktoR() --- .../module_xc/xc_functional_libxc_vxc.cpp | 44 ++- source/module_ri/RI_2D_Comm.h | 32 ++- source/module_ri/RI_2D_Comm.hpp | 267 +++++++++++++----- 3 files changed, 262 insertions(+), 81 deletions(-) diff --git a/source/module_hamilt_general/module_xc/xc_functional_libxc_vxc.cpp b/source/module_hamilt_general/module_xc/xc_functional_libxc_vxc.cpp index 2456abe5ba..f69efc1b94 100644 --- a/source/module_hamilt_general/module_xc/xc_functional_libxc_vxc.cpp +++ b/source/module_hamilt_general/module_xc/xc_functional_libxc_vxc.cpp @@ -10,7 +10,7 @@ #include "module_base/tool_title.h" #include - +#include #include std::tuple XC_Functional_Libxc::v_xc_libxc( // Peize Lin update for nspin==4 at 2023.01.14 @@ -94,16 +94,46 @@ std::tuple XC_Functional_Libxc::v_xc_libxc( / switch( func.info->family ) { case XC_FAMILY_LDA: - // call Libxc function: xc_lda_exc_vxc - xc_lda_exc_vxc( &func, nrxx, rho.data(), - exc.data(), vrho.data() ); + { + constexpr int nr_batch_size = 1024; + #ifdef _OPENMP + #pragma omp parallel for schedule(static, nr_batch_size) + #endif + for( int ir_start = 0; ir_start < nrxx; ir_start += nr_batch_size ) + { + const int ir_end = std::min(ir_start + nr_batch_size, nrxx); + const int nrxx_thread = ir_end - ir_start; + xc_lda_exc_vxc( + &func, + nrxx_thread, + rho.data() + ir_start * nspin, + exc.data() + ir_start, + vrho.data() + ir_start * nspin ); + } break; + } case XC_FAMILY_GGA: case XC_FAMILY_HYB_GGA: - // call Libxc function: xc_gga_exc_vxc - xc_gga_exc_vxc( &func, nrxx, rho.data(), sigma.data(), - exc.data(), vrho.data(), vsigma.data() ); + { + constexpr int nr_batch_size = 1024; + #ifdef _OPENMP + #pragma omp parallel for schedule(static, nr_batch_size) + #endif + for( int ir_start = 0; ir_start < nrxx; ir_start += nr_batch_size ) + { + const int ir_end = std::min(ir_start + nr_batch_size, nrxx); + const int nrxx_thread = ir_end - ir_start; + xc_gga_exc_vxc( + &func, + nrxx_thread, + rho.data() + ir_start * nspin, + sigma.data() + ir_start * ((1==nspin)?1:3), + exc.data() + ir_start, + vrho.data() + ir_start * nspin, + vsigma.data() + ir_start * ((1==nspin)?1:3) ); + } break; + } default: throw std::domain_error("func.info->family ="+std::to_string(func.info->family) +" unfinished in "+std::string(__FILE__)+" line "+std::to_string(__LINE__)); diff --git a/source/module_ri/RI_2D_Comm.h b/source/module_ri/RI_2D_Comm.h index c7d023e735..ca7e210870 100644 --- a/source/module_ri/RI_2D_Comm.h +++ b/source/module_ri/RI_2D_Comm.h @@ -29,14 +29,30 @@ namespace RI_2D_Comm using TAC = std::pair; //public: - template - extern std::vector>>> - split_m2D_ktoR(const UnitCell& ucell, - const K_Vectors& kv, - const std::vector& mks_2D, - const Parallel_2D& pv, - const int nspin, - const bool spgsym = false); + template + extern std::vector>>> split_m2D_ktoR( + const UnitCell& ucell, + const K_Vectors& kv, + const std::vector& mks_2D, + const Parallel_2D& pv, + const int nspin, + const bool spgsym = false); + + template + extern std::vector>>> split_m2D_ktoR_gamma( + const UnitCell& ucell, + const std::vector& mks_2D, + const Parallel_2D& pv, + const int nspin); + + template + extern std::vector>>> split_m2D_ktoR_k( + const UnitCell& ucell, + const K_Vectors& kv, + const std::vector& mks_2D, + const Parallel_2D& pv, + const int nspin, + const bool spgsym = false); // judge[is] = {s0, s1} extern std::vector, std::set>> diff --git a/source/module_ri/RI_2D_Comm.hpp b/source/module_ri/RI_2D_Comm.hpp index f40e6f59d4..056d6aff6e 100644 --- a/source/module_ri/RI_2D_Comm.hpp +++ b/source/module_ri/RI_2D_Comm.hpp @@ -30,93 +30,228 @@ inline RI::Tensor> tensor_conj(const RI::Tensor auto RI_2D_Comm::split_m2D_ktoR(const UnitCell& ucell, - const K_Vectors & kv, - const std::vector&mks_2D, - const Parallel_2D & pv, - const int nspin, + const K_Vectors & kv, + const std::vector&mks_2D, + const Parallel_2D & pv, + const int nspin, const bool spgsym) -> std::vector>>> { ModuleBase::TITLE("RI_2D_Comm","split_m2D_ktoR"); ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR"); - const TC period = RI_Util::get_Born_vonKarmen_period(kv); + std::vector>>> mRs_a2D + = (period == TC{1, 1, 1}) + ? RI_2D_Comm::split_m2D_ktoR_gamma(ucell, mks_2D, pv, nspin) + : RI_2D_Comm::split_m2D_ktoR_k(ucell, kv, mks_2D, pv, nspin, spgsym); + ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR"); + return mRs_a2D; +} + +template +auto RI_2D_Comm::split_m2D_ktoR_gamma(const UnitCell& ucell, + const std::vector& mks_2D, + const Parallel_2D& pv, + const int nspin) +-> std::vector>>> +{ + ModuleBase::TITLE("RI_2D_Comm","split_m2D_ktoR_gamma"); + ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR_gamma"); + const std::map nspin_k = {{1,1}, {2,2}, {4,1}}; const double SPIN_multiple = std::map{ {1,0.5}, {2,1}, {4,1} }.at(nspin); // why? + const TC cell = {0, 0, 0}; std::vector>>> mRs_a2D(nspin); + + #ifdef _OPENMP + // pre-init all outer maps mRs_a2D[is_b][iat] to avoid concurrent std::map rebalancing + for (int is_b = 0; is_b < nspin; ++is_b) + for (int iat0 = 0; iat0 < ucell.nat; ++iat0) + mRs_a2D[is_b][iat0]; + + std::vector locks(ucell.nat); + for (auto& l : locks) + omp_init_lock(&l); + #endif + for (int is_k = 0; is_k < nspin_k.at(nspin); ++is_k) - { - const std::vector ik_list = RI_2D_Comm::get_ik_list(kv, is_k); - for(const TC &cell : RI_Util::get_Born_von_Karmen_cells(period)) - { - RI::Tensor mR_2D; - int ik_full = 0; - for (const int ik : ik_list) + { + using Tdata_m = typename Tmatrix::value_type; + RI::Tensor mk_2D + = RI_Util::Vector_to_Tensor(*mks_2D[is_k], pv.get_col_size(), pv.get_row_size()); + const Tdata_m frac = RI::Global_Func::convert(SPIN_multiple); + RI::Tensor mR_2D = RI::Global_Func::convert(mk_2D * frac); + + #ifdef _OPENMP + #pragma omp parallel for schedule(dynamic) + #endif + for (int iwt0_2D = 0; iwt0_2D != mR_2D.shape[0]; ++iwt0_2D) + { + const int iwt0 = ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver) + ? pv.local2global_col(iwt0_2D) + : pv.local2global_row(iwt0_2D); + int iat0, iw0_b, is0_b; + std::tie(iat0, iw0_b, is0_b) = RI_2D_Comm::get_iat_iw_is_block(ucell, iwt0); + const int it0 = ucell.iat2it[iat0]; + for (int iwt1_2D = 0; iwt1_2D != mR_2D.shape[1]; ++iwt1_2D) { - auto set_mR_2D = [&mR_2D](auto&& mk_frac) { - if (mR_2D.empty()) { - mR_2D = RI::Global_Func::convert(mk_frac); - } else { - mR_2D - = mR_2D + RI::Global_Func::convert(mk_frac); - } - }; - using Tdata_m = typename Tmatrix::value_type; - if (!spgsym) + const int iwt1 = ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver) + ? pv.local2global_row(iwt1_2D) + : pv.local2global_col(iwt1_2D); + int iat1, iw1_b, is1_b; + std::tie(iat1, iw1_b, is1_b) = RI_2D_Comm::get_iat_iw_is_block(ucell, iwt1); + const int it1 = ucell.iat2it[iat1]; + + const int is_b = RI_2D_Comm::get_is_block(is_k, is0_b, is1_b); + #ifdef _OPENMP + omp_set_lock(&locks[iat0]); + #endif + RI::Tensor& mR_a2D = mRs_a2D[is_b][iat0][{iat1, cell}]; + if (mR_a2D.empty()) { - RI::Tensor mk_2D = RI_Util::Vector_to_Tensor(*mks_2D[ik], pv.get_col_size(), pv.get_row_size()); - const Tdata_m frac = SPIN_multiple - * RI::Global_Func::convert(std::exp( - -ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * (kv.kvec_c[ik] * (RI_Util::array3_to_Vector3(cell) * ucell.latvec)))); - if (static_cast(std::round(SPIN_multiple * kv.wk[ik] * kv.get_nkstot_full())) == 2) - { set_mR_2D(mk_2D * (frac * 0.5) + tensor_conj(mk_2D * (frac * 0.5))); } - else { set_mR_2D(mk_2D * frac); } + mR_a2D = RI::Tensor( + {static_cast(ucell.atoms[it0].nw), + static_cast(ucell.atoms[it1].nw)}); } + mR_a2D(iw0_b, iw1_b) = mR_2D(iwt0_2D, iwt1_2D); + #ifdef _OPENMP + omp_unset_lock(&locks[iat0]); + #endif + } + } + } + + #ifdef _OPENMP + for (auto& l : locks) + omp_destroy_lock(&l); + + // prune empty inner maps created by pre-init + for (int is_b = 0; is_b < nspin; ++is_b) + for (auto it = mRs_a2D[is_b].begin(); it != mRs_a2D[is_b].end();) + { + if (it->second.empty()) + it = mRs_a2D[is_b].erase(it); else - { // traverse kstar, ik means ik_ibz - for (auto& isym_kvd : kv.kstars[ik % ik_list.size()]) + ++it; + } + #endif + + ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR_gamma"); + return mRs_a2D; +} + +template +auto RI_2D_Comm::split_m2D_ktoR_k(const UnitCell& ucell, + const K_Vectors& kv, + const std::vector& mks_2D, + const Parallel_2D& pv, + const int nspin, + const bool spgsym) +-> std::vector>>> +{ + ModuleBase::TITLE("RI_2D_Comm","split_m2D_ktoR_k"); + ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR_k"); + + const TC period = RI_Util::get_Born_vonKarmen_period(kv); + const std::map nspin_k = {{1,1}, {2,2}, {4,1}}; + const double SPIN_multiple = std::map{ {1,0.5}, {2,1}, {4,1} }.at(nspin); // why? + + std::vector>>> mRs_a2D(nspin); + #ifdef _OPENMP + #pragma omp parallel + #endif + { + std::vector>>> mRs_a2D_thread(nspin); + for (int is_k = 0; is_k < nspin_k.at(nspin); ++is_k) + { + const std::vector ik_list = RI_2D_Comm::get_ik_list(kv, is_k); + const auto cells = RI_Util::get_Born_von_Karmen_cells(period); + #pragma omp for schedule(dynamic) + for (size_t icell = 0; icell < cells.size(); ++icell) + { + const TC& cell = cells[icell]; + RI::Tensor mR_2D; + int ik_full = 0; + for (const int ik : ik_list) + { + using Tdata_m = typename Tmatrix::value_type; + auto set_mR_2D = [&mR_2D](RI::Tensor&& mk_frac) + { + if (mR_2D.empty()) + { mR_2D = RI::Global_Func::convert(mk_frac); } + else + { mR_2D = mR_2D + RI::Global_Func::convert(mk_frac); } + }; + if (!spgsym) { - RI::Tensor mk_2D = RI_Util::Vector_to_Tensor(*mks_2D[ik_full + is_k * kv.get_nkstot_full()], pv.get_col_size(), pv.get_row_size()); + RI::Tensor mk_2D = RI_Util::Vector_to_Tensor(*mks_2D[ik], pv.get_col_size(), pv.get_row_size()); const Tdata_m frac = SPIN_multiple * RI::Global_Func::convert(std::exp( - -ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * ((isym_kvd.second * ucell.G) * (RI_Util::array3_to_Vector3(cell) * ucell.latvec)))); - set_mR_2D(mk_2D * frac); - ++ik_full; + -ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * (kv.kvec_c[ik] * (RI_Util::array3_to_Vector3(cell) * ucell.latvec)))); + if (static_cast(std::round(SPIN_multiple * kv.wk[ik] * kv.get_nkstot_full())) == 2) + { set_mR_2D(mk_2D * (frac * 0.5) + tensor_conj(mk_2D * (frac * 0.5))); } + else + { set_mR_2D(mk_2D * frac); } } - } - } - for(int iwt0_2D=0; iwt0_2D!=mR_2D.shape[0]; ++iwt0_2D) - { - const int iwt0 =ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver) - ? pv.local2global_col(iwt0_2D) - : pv.local2global_row(iwt0_2D); - int iat0, iw0_b, is0_b; - std::tie(iat0,iw0_b,is0_b) = RI_2D_Comm::get_iat_iw_is_block(ucell,iwt0); - const int it0 = ucell.iat2it[iat0]; - for(int iwt1_2D=0; iwt1_2D!=mR_2D.shape[1]; ++iwt1_2D) - { - const int iwt1 =ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver) - ? pv.local2global_row(iwt1_2D) - : pv.local2global_col(iwt1_2D); - int iat1, iw1_b, is1_b; - std::tie(iat1,iw1_b,is1_b) = RI_2D_Comm::get_iat_iw_is_block(ucell,iwt1); - const int it1 = ucell.iat2it[iat1]; - - const int is_b = RI_2D_Comm::get_is_block(is_k, is0_b, is1_b); - RI::Tensor &mR_a2D = mRs_a2D[is_b][iat0][{iat1,cell}]; - if (mR_a2D.empty()) { - mR_a2D = RI::Tensor( - {static_cast(ucell.atoms[it0].nw), - static_cast( - ucell.atoms[it1].nw)}); + else + { // traverse kstar, ik means ik_ibz + for (auto& isym_kvd : kv.kstars[ik % ik_list.size()]) + { + RI::Tensor mk_2D = RI_Util::Vector_to_Tensor(*mks_2D[ik_full + is_k * kv.get_nkstot_full()], pv.get_col_size(), pv.get_row_size()); + const Tdata_m frac = SPIN_multiple + * RI::Global_Func::convert(std::exp( + -ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * ((isym_kvd.second * ucell.G) * (RI_Util::array3_to_Vector3(cell) * ucell.latvec)))); + set_mR_2D(mk_2D * frac); + ++ik_full; + } + } + } // end for ik + for(int iwt0_2D=0; iwt0_2D!=mR_2D.shape[0]; ++iwt0_2D) + { + const int iwt0 =ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver) + ? pv.local2global_col(iwt0_2D) + : pv.local2global_row(iwt0_2D); + int iat0, iw0_b, is0_b; + std::tie(iat0,iw0_b,is0_b) = RI_2D_Comm::get_iat_iw_is_block(ucell,iwt0); + const int it0 = ucell.iat2it[iat0]; + for(int iwt1_2D=0; iwt1_2D!=mR_2D.shape[1]; ++iwt1_2D) + { + const int iwt1 =ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver) + ? pv.local2global_row(iwt1_2D) + : pv.local2global_col(iwt1_2D); + int iat1, iw1_b, is1_b; + std::tie(iat1,iw1_b,is1_b) = RI_2D_Comm::get_iat_iw_is_block(ucell,iwt1); + const int it1 = ucell.iat2it[iat1]; + + const int is_b = RI_2D_Comm::get_is_block(is_k, is0_b, is1_b); + RI::Tensor& mR_a2D = mRs_a2D_thread[is_b][iat0][{iat1, cell}]; + if (mR_a2D.empty()) + { + mR_a2D = RI::Tensor( + {static_cast(ucell.atoms[it0].nw), + static_cast(ucell.atoms[it1].nw)}); + } + mR_a2D(iw0_b, iw1_b) = mR_2D(iwt0_2D, iwt1_2D); + } // for iwt1_2D + } // end for iwt0_2D + } // end for icell + } // end for is_k + + #ifdef _OPENMP + #pragma omp critical + #endif + { + for(int is=0; is