Skip to content

Commit ee763df

Browse files
linpeizePeizeLin
linpeize
and
PeizeLin
authored
1. add OpenMP in XC_Functional_Libxc::v_xc_libxc() (deepmodeling#7777)
2. add OpenMP in RI_2D_Comm::split_m2D_ktoR() Co-authored-by: linpz <linpz@mail.ustc.edu.cn>
1 parent e148bc4 commit ee763df

3 files changed

Lines changed: 262 additions & 81 deletions

File tree

‎source/module_hamilt_general/module_xc/xc_functional_libxc_vxc.cpp‎

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
#include "module_base/tool_title.h"
1111

1212
#include <xc.h>
13-
13+
#include <omp.h>
1414
#include <vector>
1515

1616
std::tuple<double,double,ModuleBase::matrix> XC_Functional_Libxc::v_xc_libxc( // Peize Lin update for nspin==4 at 2023.01.14
@@ -94,16 +94,46 @@ std::tuple<double,double,ModuleBase::matrix> XC_Functional_Libxc::v_xc_libxc( /
9494
switch( func.info->family )
9595
{
9696
case XC_FAMILY_LDA:
97-
// call Libxc function: xc_lda_exc_vxc
98-
xc_lda_exc_vxc( &func, nrxx, rho.data(),
99-
exc.data(), vrho.data() );
97+
{
98+
constexpr int nr_batch_size = 1024;
99+
#ifdef _OPENMP
100+
#pragma omp parallel for schedule(static, nr_batch_size)
101+
#endif
102+
for( int ir_start = 0; ir_start < nrxx; ir_start += nr_batch_size )
103+
{
104+
const int ir_end = std::min(ir_start + nr_batch_size, nrxx);
105+
const int nrxx_thread = ir_end - ir_start;
106+
xc_lda_exc_vxc(
107+
&func,
108+
nrxx_thread,
109+
rho.data() + ir_start * nspin,
110+
exc.data() + ir_start,
111+
vrho.data() + ir_start * nspin );
112+
}
100113
break;
114+
}
101115
case XC_FAMILY_GGA:
102116
case XC_FAMILY_HYB_GGA:
103-
// call Libxc function: xc_gga_exc_vxc
104-
xc_gga_exc_vxc( &func, nrxx, rho.data(), sigma.data(),
105-
exc.data(), vrho.data(), vsigma.data() );
117+
{
118+
constexpr int nr_batch_size = 1024;
119+
#ifdef _OPENMP
120+
#pragma omp parallel for schedule(static, nr_batch_size)
121+
#endif
122+
for( int ir_start = 0; ir_start < nrxx; ir_start += nr_batch_size )
123+
{
124+
const int ir_end = std::min(ir_start + nr_batch_size, nrxx);
125+
const int nrxx_thread = ir_end - ir_start;
126+
xc_gga_exc_vxc(
127+
&func,
128+
nrxx_thread,
129+
rho.data() + ir_start * nspin,
130+
sigma.data() + ir_start * ((1==nspin)?1:3),
131+
exc.data() + ir_start,
132+
vrho.data() + ir_start * nspin,
133+
vsigma.data() + ir_start * ((1==nspin)?1:3) );
134+
}
106135
break;
136+
}
107137
default:
108138
throw std::domain_error("func.info->family ="+std::to_string(func.info->family)
109139
+" unfinished in "+std::string(__FILE__)+" line "+std::to_string(__LINE__));

‎source/module_ri/RI_2D_Comm.h‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,30 @@ namespace RI_2D_Comm
2929
using TAC = std::pair<TA,TC>;
3030

3131
//public:
32-
template<typename Tdata, typename Tmatrix>
33-
extern std::vector<std::map<TA,std::map<TAC,RI::Tensor<Tdata>>>>
34-
split_m2D_ktoR(const UnitCell& ucell,
35-
const K_Vectors& kv,
36-
const std::vector<const Tmatrix*>& mks_2D,
37-
const Parallel_2D& pv,
38-
const int nspin,
39-
const bool spgsym = false);
32+
template <typename Tdata, typename Tmatrix>
33+
extern std::vector<std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>> split_m2D_ktoR(
34+
const UnitCell& ucell,
35+
const K_Vectors& kv,
36+
const std::vector<const Tmatrix*>& mks_2D,
37+
const Parallel_2D& pv,
38+
const int nspin,
39+
const bool spgsym = false);
40+
41+
template <typename Tdata, typename Tmatrix>
42+
extern std::vector<std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>> split_m2D_ktoR_gamma(
43+
const UnitCell& ucell,
44+
const std::vector<const Tmatrix*>& mks_2D,
45+
const Parallel_2D& pv,
46+
const int nspin);
47+
48+
template <typename Tdata, typename Tmatrix>
49+
extern std::vector<std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>> split_m2D_ktoR_k(
50+
const UnitCell& ucell,
51+
const K_Vectors& kv,
52+
const std::vector<const Tmatrix*>& mks_2D,
53+
const Parallel_2D& pv,
54+
const int nspin,
55+
const bool spgsym = false);
4056

4157
// judge[is] = {s0, s1}
4258
extern std::vector<std::tuple<std::set<TA>, std::set<TA>>>

‎source/module_ri/RI_2D_Comm.hpp‎

Lines changed: 201 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -30,93 +30,228 @@ inline RI::Tensor<std::complex<double>> tensor_conj(const RI::Tensor<std::comple
3030
}
3131
template<typename Tdata, typename Tmatrix>
3232
auto RI_2D_Comm::split_m2D_ktoR(const UnitCell& ucell,
33-
const K_Vectors & kv,
34-
const std::vector<const Tmatrix*>&mks_2D,
35-
const Parallel_2D & pv,
36-
const int nspin,
33+
const K_Vectors & kv,
34+
const std::vector<const Tmatrix*>&mks_2D,
35+
const Parallel_2D & pv,
36+
const int nspin,
3737
const bool spgsym)
3838
-> std::vector<std::map<TA,std::map<TAC,RI::Tensor<Tdata>>>>
3939
{
4040
ModuleBase::TITLE("RI_2D_Comm","split_m2D_ktoR");
4141
ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR");
42-
4342
const TC period = RI_Util::get_Born_vonKarmen_period(kv);
43+
std::vector<std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>> mRs_a2D
44+
= (period == TC{1, 1, 1})
45+
? RI_2D_Comm::split_m2D_ktoR_gamma<Tdata, Tmatrix>(ucell, mks_2D, pv, nspin)
46+
: RI_2D_Comm::split_m2D_ktoR_k<Tdata, Tmatrix>(ucell, kv, mks_2D, pv, nspin, spgsym);
47+
ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR");
48+
return mRs_a2D;
49+
}
50+
51+
template<typename Tdata, typename Tmatrix>
52+
auto RI_2D_Comm::split_m2D_ktoR_gamma(const UnitCell& ucell,
53+
const std::vector<const Tmatrix*>& mks_2D,
54+
const Parallel_2D& pv,
55+
const int nspin)
56+
-> std::vector<std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>>
57+
{
58+
ModuleBase::TITLE("RI_2D_Comm","split_m2D_ktoR_gamma");
59+
ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR_gamma");
60+
4461
const std::map<int,int> nspin_k = {{1,1}, {2,2}, {4,1}};
4562
const double SPIN_multiple = std::map<int, double>{ {1,0.5}, {2,1}, {4,1} }.at(nspin); // why?
63+
const TC cell = {0, 0, 0};
4664

4765
std::vector<std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>> mRs_a2D(nspin);
66+
67+
#ifdef _OPENMP
68+
// pre-init all outer maps mRs_a2D[is_b][iat] to avoid concurrent std::map rebalancing
69+
for (int is_b = 0; is_b < nspin; ++is_b)
70+
for (int iat0 = 0; iat0 < ucell.nat; ++iat0)
71+
mRs_a2D[is_b][iat0];
72+
73+
std::vector<omp_lock_t> locks(ucell.nat);
74+
for (auto& l : locks)
75+
omp_init_lock(&l);
76+
#endif
77+
4878
for (int is_k = 0; is_k < nspin_k.at(nspin); ++is_k)
49-
{
50-
const std::vector<int> ik_list = RI_2D_Comm::get_ik_list(kv, is_k);
51-
for(const TC &cell : RI_Util::get_Born_von_Karmen_cells(period))
52-
{
53-
RI::Tensor<Tdata> mR_2D;
54-
int ik_full = 0;
55-
for (const int ik : ik_list)
79+
{
80+
using Tdata_m = typename Tmatrix::value_type;
81+
RI::Tensor<Tdata_m> mk_2D
82+
= RI_Util::Vector_to_Tensor<Tdata_m>(*mks_2D[is_k], pv.get_col_size(), pv.get_row_size());
83+
const Tdata_m frac = RI::Global_Func::convert<Tdata_m>(SPIN_multiple);
84+
RI::Tensor<Tdata> mR_2D = RI::Global_Func::convert<Tdata>(mk_2D * frac);
85+
86+
#ifdef _OPENMP
87+
#pragma omp parallel for schedule(dynamic)
88+
#endif
89+
for (int iwt0_2D = 0; iwt0_2D != mR_2D.shape[0]; ++iwt0_2D)
90+
{
91+
const int iwt0 = ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)
92+
? pv.local2global_col(iwt0_2D)
93+
: pv.local2global_row(iwt0_2D);
94+
int iat0, iw0_b, is0_b;
95+
std::tie(iat0, iw0_b, is0_b) = RI_2D_Comm::get_iat_iw_is_block(ucell, iwt0);
96+
const int it0 = ucell.iat2it[iat0];
97+
for (int iwt1_2D = 0; iwt1_2D != mR_2D.shape[1]; ++iwt1_2D)
5698
{
57-
auto set_mR_2D = [&mR_2D](auto&& mk_frac) {
58-
if (mR_2D.empty()) {
59-
mR_2D = RI::Global_Func::convert<Tdata>(mk_frac);
60-
} else {
61-
mR_2D
62-
= mR_2D + RI::Global_Func::convert<Tdata>(mk_frac);
63-
}
64-
};
65-
using Tdata_m = typename Tmatrix::value_type;
66-
if (!spgsym)
99+
const int iwt1 = ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)
100+
? pv.local2global_row(iwt1_2D)
101+
: pv.local2global_col(iwt1_2D);
102+
int iat1, iw1_b, is1_b;
103+
std::tie(iat1, iw1_b, is1_b) = RI_2D_Comm::get_iat_iw_is_block(ucell, iwt1);
104+
const int it1 = ucell.iat2it[iat1];
105+
106+
const int is_b = RI_2D_Comm::get_is_block(is_k, is0_b, is1_b);
107+
#ifdef _OPENMP
108+
omp_set_lock(&locks[iat0]);
109+
#endif
110+
RI::Tensor<Tdata>& mR_a2D = mRs_a2D[is_b][iat0][{iat1, cell}];
111+
if (mR_a2D.empty())
67112
{
68-
RI::Tensor<Tdata_m> mk_2D = RI_Util::Vector_to_Tensor<Tdata_m>(*mks_2D[ik], pv.get_col_size(), pv.get_row_size());
69-
const Tdata_m frac = SPIN_multiple
70-
* RI::Global_Func::convert<Tdata_m>(std::exp(
71-
-ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * (kv.kvec_c[ik] * (RI_Util::array3_to_Vector3(cell) * ucell.latvec))));
72-
if (static_cast<int>(std::round(SPIN_multiple * kv.wk[ik] * kv.get_nkstot_full())) == 2)
73-
{ set_mR_2D(mk_2D * (frac * 0.5) + tensor_conj(mk_2D * (frac * 0.5))); }
74-
else { set_mR_2D(mk_2D * frac); }
113+
mR_a2D = RI::Tensor<Tdata>(
114+
{static_cast<size_t>(ucell.atoms[it0].nw),
115+
static_cast<size_t>(ucell.atoms[it1].nw)});
75116
}
117+
mR_a2D(iw0_b, iw1_b) = mR_2D(iwt0_2D, iwt1_2D);
118+
#ifdef _OPENMP
119+
omp_unset_lock(&locks[iat0]);
120+
#endif
121+
}
122+
}
123+
}
124+
125+
#ifdef _OPENMP
126+
for (auto& l : locks)
127+
omp_destroy_lock(&l);
128+
129+
// prune empty inner maps created by pre-init
130+
for (int is_b = 0; is_b < nspin; ++is_b)
131+
for (auto it = mRs_a2D[is_b].begin(); it != mRs_a2D[is_b].end();)
132+
{
133+
if (it->second.empty())
134+
it = mRs_a2D[is_b].erase(it);
76135
else
77-
{ // traverse kstar, ik means ik_ibz
78-
for (auto& isym_kvd : kv.kstars[ik % ik_list.size()])
136+
++it;
137+
}
138+
#endif
139+
140+
ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR_gamma");
141+
return mRs_a2D;
142+
}
143+
144+
template<typename Tdata, typename Tmatrix>
145+
auto RI_2D_Comm::split_m2D_ktoR_k(const UnitCell& ucell,
146+
const K_Vectors& kv,
147+
const std::vector<const Tmatrix*>& mks_2D,
148+
const Parallel_2D& pv,
149+
const int nspin,
150+
const bool spgsym)
151+
-> std::vector<std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>>
152+
{
153+
ModuleBase::TITLE("RI_2D_Comm","split_m2D_ktoR_k");
154+
ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR_k");
155+
156+
const TC period = RI_Util::get_Born_vonKarmen_period(kv);
157+
const std::map<int,int> nspin_k = {{1,1}, {2,2}, {4,1}};
158+
const double SPIN_multiple = std::map<int, double>{ {1,0.5}, {2,1}, {4,1} }.at(nspin); // why?
159+
160+
std::vector<std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>> mRs_a2D(nspin);
161+
#ifdef _OPENMP
162+
#pragma omp parallel
163+
#endif
164+
{
165+
std::vector<std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>> mRs_a2D_thread(nspin);
166+
for (int is_k = 0; is_k < nspin_k.at(nspin); ++is_k)
167+
{
168+
const std::vector<int> ik_list = RI_2D_Comm::get_ik_list(kv, is_k);
169+
const auto cells = RI_Util::get_Born_von_Karmen_cells(period);
170+
#pragma omp for schedule(dynamic)
171+
for (size_t icell = 0; icell < cells.size(); ++icell)
172+
{
173+
const TC& cell = cells[icell];
174+
RI::Tensor<Tdata> mR_2D;
175+
int ik_full = 0;
176+
for (const int ik : ik_list)
177+
{
178+
using Tdata_m = typename Tmatrix::value_type;
179+
auto set_mR_2D = [&mR_2D](RI::Tensor<Tdata_m>&& mk_frac)
180+
{
181+
if (mR_2D.empty())
182+
{ mR_2D = RI::Global_Func::convert<Tdata>(mk_frac); }
183+
else
184+
{ mR_2D = mR_2D + RI::Global_Func::convert<Tdata>(mk_frac); }
185+
};
186+
if (!spgsym)
79187
{
80-
RI::Tensor<Tdata_m> mk_2D = RI_Util::Vector_to_Tensor<Tdata_m>(*mks_2D[ik_full + is_k * kv.get_nkstot_full()], pv.get_col_size(), pv.get_row_size());
188+
RI::Tensor<Tdata_m> mk_2D = RI_Util::Vector_to_Tensor<Tdata_m>(*mks_2D[ik], pv.get_col_size(), pv.get_row_size());
81189
const Tdata_m frac = SPIN_multiple
82190
* RI::Global_Func::convert<Tdata_m>(std::exp(
83-
-ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * ((isym_kvd.second * ucell.G) * (RI_Util::array3_to_Vector3(cell) * ucell.latvec))));
84-
set_mR_2D(mk_2D * frac);
85-
++ik_full;
191+
-ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * (kv.kvec_c[ik] * (RI_Util::array3_to_Vector3(cell) * ucell.latvec))));
192+
if (static_cast<int>(std::round(SPIN_multiple * kv.wk[ik] * kv.get_nkstot_full())) == 2)
193+
{ set_mR_2D(mk_2D * (frac * 0.5) + tensor_conj(mk_2D * (frac * 0.5))); }
194+
else
195+
{ set_mR_2D(mk_2D * frac); }
86196
}
87-
}
88-
}
89-
for(int iwt0_2D=0; iwt0_2D!=mR_2D.shape[0]; ++iwt0_2D)
90-
{
91-
const int iwt0 =ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)
92-
? pv.local2global_col(iwt0_2D)
93-
: pv.local2global_row(iwt0_2D);
94-
int iat0, iw0_b, is0_b;
95-
std::tie(iat0,iw0_b,is0_b) = RI_2D_Comm::get_iat_iw_is_block(ucell,iwt0);
96-
const int it0 = ucell.iat2it[iat0];
97-
for(int iwt1_2D=0; iwt1_2D!=mR_2D.shape[1]; ++iwt1_2D)
98-
{
99-
const int iwt1 =ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)
100-
? pv.local2global_row(iwt1_2D)
101-
: pv.local2global_col(iwt1_2D);
102-
int iat1, iw1_b, is1_b;
103-
std::tie(iat1,iw1_b,is1_b) = RI_2D_Comm::get_iat_iw_is_block(ucell,iwt1);
104-
const int it1 = ucell.iat2it[iat1];
105-
106-
const int is_b = RI_2D_Comm::get_is_block(is_k, is0_b, is1_b);
107-
RI::Tensor<Tdata> &mR_a2D = mRs_a2D[is_b][iat0][{iat1,cell}];
108-
if (mR_a2D.empty()) {
109-
mR_a2D = RI::Tensor<Tdata>(
110-
{static_cast<size_t>(ucell.atoms[it0].nw),
111-
static_cast<size_t>(
112-
ucell.atoms[it1].nw)});
197+
else
198+
{ // traverse kstar, ik means ik_ibz
199+
for (auto& isym_kvd : kv.kstars[ik % ik_list.size()])
200+
{
201+
RI::Tensor<Tdata_m> mk_2D = RI_Util::Vector_to_Tensor<Tdata_m>(*mks_2D[ik_full + is_k * kv.get_nkstot_full()], pv.get_col_size(), pv.get_row_size());
202+
const Tdata_m frac = SPIN_multiple
203+
* RI::Global_Func::convert<Tdata_m>(std::exp(
204+
-ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * ((isym_kvd.second * ucell.G) * (RI_Util::array3_to_Vector3(cell) * ucell.latvec))));
205+
set_mR_2D(mk_2D * frac);
206+
++ik_full;
207+
}
208+
}
209+
} // end for ik
210+
for(int iwt0_2D=0; iwt0_2D!=mR_2D.shape[0]; ++iwt0_2D)
211+
{
212+
const int iwt0 =ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)
213+
? pv.local2global_col(iwt0_2D)
214+
: pv.local2global_row(iwt0_2D);
215+
int iat0, iw0_b, is0_b;
216+
std::tie(iat0,iw0_b,is0_b) = RI_2D_Comm::get_iat_iw_is_block(ucell,iwt0);
217+
const int it0 = ucell.iat2it[iat0];
218+
for(int iwt1_2D=0; iwt1_2D!=mR_2D.shape[1]; ++iwt1_2D)
219+
{
220+
const int iwt1 =ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)
221+
? pv.local2global_row(iwt1_2D)
222+
: pv.local2global_col(iwt1_2D);
223+
int iat1, iw1_b, is1_b;
224+
std::tie(iat1,iw1_b,is1_b) = RI_2D_Comm::get_iat_iw_is_block(ucell,iwt1);
225+
const int it1 = ucell.iat2it[iat1];
226+
227+
const int is_b = RI_2D_Comm::get_is_block(is_k, is0_b, is1_b);
228+
RI::Tensor<Tdata>& mR_a2D = mRs_a2D_thread[is_b][iat0][{iat1, cell}];
229+
if (mR_a2D.empty())
230+
{
231+
mR_a2D = RI::Tensor<Tdata>(
232+
{static_cast<size_t>(ucell.atoms[it0].nw),
233+
static_cast<size_t>(ucell.atoms[it1].nw)});
234+
}
235+
mR_a2D(iw0_b, iw1_b) = mR_2D(iwt0_2D, iwt1_2D);
236+
} // for iwt1_2D
237+
} // end for iwt0_2D
238+
} // end for icell
239+
} // end for is_k
240+
241+
#ifdef _OPENMP
242+
#pragma omp critical
243+
#endif
244+
{
245+
for(int is=0; is<nspin; ++is)
246+
for(auto &mRs_A : mRs_a2D_thread[is])
247+
for(auto &mRs_B : mRs_A.second)
248+
{
249+
assert(mRs_a2D[is][mRs_A.first][mRs_B.first].empty());
250+
mRs_a2D[is][mRs_A.first][mRs_B.first] = std::move(mRs_B.second);
113251
}
114-
mR_a2D(iw0_b,iw1_b) = mR_2D(iwt0_2D, iwt1_2D);
115-
}
116-
}
117252
}
118-
}
119-
ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR");
253+
} // end #pragma omp parallel
254+
ModuleBase::timer::tick("RI_2D_Comm", "split_m2D_ktoR_k");
120255
return mRs_a2D;
121256
}
122257

0 commit comments

Comments
 (0)