Skip to content

Commit 7e2c8b5

Browse files
cheerly-pkucheerly-pku
committed
Parallelize local PPCG vector operations
1 parent 377fe68 commit 7e2c8b5

4 files changed

Lines changed: 102 additions & 24 deletions

File tree

‎source/source_hsolver/ppcg/diago_ppcg_cg.hpp‎

Lines changed: 54 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@ void DiagoPPCG<T, Device>::calc_gradient(
1717
std::vector<T>& grad) const
1818
{
1919
grad.assign(ld_psi_ * n_band_, T(0));
20+
#ifdef _OPENMP
21+
#pragma omp parallel for schedule(static) if (n_dim_ * n_band_ > 4096)
22+
#endif
2023
for (int j = 0; j < n_band_; ++j)
2124
{
2225
const Real ej = eigenvalue[j];
@@ -88,6 +91,9 @@ void DiagoPPCG<T, Device>::update_polak_ribiere(
8891
Real beta_num_zr = 0;
8992
Real beta_num_zo = 0;
9093

94+
#ifdef _OPENMP
95+
#pragma omp parallel for reduction(+ : beta_num_zr, beta_num_zo) schedule(static) if (n_dim_ > 4096)
96+
#endif
9197
for (int ig = 0; ig < n_dim_; ++ig)
9298
{
9399
// z_new = -P^{-1} * grad
@@ -113,6 +119,9 @@ void DiagoPPCG<T, Device>::update_polak_ribiere(
113119
}
114120

115121
// d_new = z_new + beta * d_old
122+
#ifdef _OPENMP
123+
#pragma omp parallel for schedule(static) if (n_dim_ > 4096)
124+
#endif
116125
for (int ig = 0; ig < n_dim_; ++ig)
117126
pj[ig] = zn[ig] + beta * pj[ig];
118127

@@ -230,9 +239,12 @@ void DiagoPPCG<T, Device>::line_minimize(
230239
alpha = alpha_linear;
231240
}
232241

242+
const T step = T(alpha) * phase;
243+
#ifdef _OPENMP
244+
#pragma omp parallel for schedule(static) if (n_dim_ > 4096)
245+
#endif
233246
for (int ig = 0; ig < n_dim_; ++ig)
234247
{
235-
const T step = T(alpha) * phase;
236248
pj[ig] += step * pp[ig];
237249
hj[ig] += step * hpp[ig];
238250
sj[ig] += step * spp[ig];
@@ -269,31 +281,52 @@ void DiagoPPCG<T, Device>::orth_cholesky(
269281
HermitianLapack<T>::potrf(ncol, gram_s.data());
270282
HermitianLapack<T>::trtri(ncol, gram_s.data());
271283

284+
const T one = T(1);
285+
const T zero = T(0);
272286
std::vector<T> tmp(ld_psi_ * ncol, T(0));
273-
for (int j = 0; j < ncol; ++j)
274-
for (int i = 0; i < ncol; ++i) {
275-
const T uinv = gram_s[i + j * ncol];
276-
for (int ig = 0; ig < n_dim_; ++ig)
277-
tmp[idx(ig, j, ld_psi_)] += psi[idx(ig, i, ld_psi_)] * uinv;
278-
}
287+
ModuleBase::gemm_op<T, Device>()('N',
288+
'N',
289+
n_dim_,
290+
ncol,
291+
ncol,
292+
&one,
293+
psi,
294+
ld_psi_,
295+
gram_s.data(),
296+
ncol,
297+
&zero,
298+
tmp.data(),
299+
ld_psi_);
279300
std::copy(tmp.begin(), tmp.end(), psi);
280301

281-
set_zero(tmp);
282-
for (int j = 0; j < ncol; ++j)
283-
for (int i = 0; i < ncol; ++i) {
284-
const T uinv = gram_s[i + j * ncol];
285-
for (int ig = 0; ig < n_dim_; ++ig)
286-
tmp[idx(ig, j, ld_psi_)] += hpsi[idx(ig, i, ld_psi_)] * uinv;
287-
}
302+
ModuleBase::gemm_op<T, Device>()('N',
303+
'N',
304+
n_dim_,
305+
ncol,
306+
ncol,
307+
&one,
308+
hpsi,
309+
ld_psi_,
310+
gram_s.data(),
311+
ncol,
312+
&zero,
313+
tmp.data(),
314+
ld_psi_);
288315
std::copy(tmp.begin(), tmp.end(), hpsi);
289316

290-
set_zero(tmp);
291-
for (int j = 0; j < ncol; ++j)
292-
for (int i = 0; i < ncol; ++i) {
293-
const T uinv = gram_s[i + j * ncol];
294-
for (int ig = 0; ig < n_dim_; ++ig)
295-
tmp[idx(ig, j, ld_psi_)] += spsi[idx(ig, i, ld_psi_)] * uinv;
296-
}
317+
ModuleBase::gemm_op<T, Device>()('N',
318+
'N',
319+
n_dim_,
320+
ncol,
321+
ncol,
322+
&one,
323+
spsi,
324+
ld_psi_,
325+
gram_s.data(),
326+
ncol,
327+
&zero,
328+
tmp.data(),
329+
ld_psi_);
297330
std::copy(tmp.begin(), tmp.end(), spsi);
298331

299332
cholesky_ok = is_s_orthonormal(psi, spsi, ncol);

‎source/source_hsolver/ppcg/diago_ppcg_ops.hpp‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,9 @@ void DiagoPPCG<T, Device>::apply_s(const SPsiFunc& spsi_func,
7373
if (spsi_func)
7474
spsi_func(psi_in, spsi_out, ld_psi_, ncol);
7575
else
76+
#ifdef _OPENMP
77+
#pragma omp parallel for schedule(static) if (ld_psi_ * ncol > 4096)
78+
#endif
7679
for (int j = 0; j < ncol; ++j)
7780
std::copy(psi_in + j * ld_psi_, psi_in + (j + 1) * ld_psi_,
7881
spsi_out + j * ld_psi_);
@@ -144,7 +147,11 @@ void DiagoPPCG<T, Device>::copy_cols(const T* src,
144147
std::vector<T>& dst) const
145148
{
146149
dst.assign(ld_psi_ * cols.size(), T(0));
147-
for (int j = 0; j < static_cast<int>(cols.size()); ++j)
150+
const int ncols = static_cast<int>(cols.size());
151+
#ifdef _OPENMP
152+
#pragma omp parallel for schedule(static) if (ld_psi_ * ncols > 4096)
153+
#endif
154+
for (int j = 0; j < ncols; ++j)
148155
{
149156
const int c = cols[j];
150157
std::copy(src + c * ld_psi_, src + c * ld_psi_ + ld_psi_,
@@ -161,7 +168,11 @@ void DiagoPPCG<T, Device>::scatter_cols(
161168
const std::vector<int>& cols,
162169
const std::vector<T>& src) const
163170
{
164-
for (int j = 0; j < static_cast<int>(cols.size()); ++j)
171+
const int ncols = static_cast<int>(cols.size());
172+
#ifdef _OPENMP
173+
#pragma omp parallel for schedule(static) if (ld_psi_ * ncols > 4096)
174+
#endif
175+
for (int j = 0; j < ncols; ++j)
165176
{
166177
const int c = cols[j];
167178
std::copy(src.begin() + j * ld_psi_,
@@ -214,10 +225,17 @@ void DiagoPPCG<T, Device>::divide_by_preconditioner(
214225
const Real* prec,
215226
std::vector<T>& x) const
216227
{
217-
for (const int c : active_cols)
228+
const int ncols = static_cast<int>(active_cols.size());
229+
#ifdef _OPENMP
230+
#pragma omp parallel for schedule(static) if (n_dim_ * ncols > 4096)
231+
#endif
232+
for (int j = 0; j < ncols; ++j)
233+
{
234+
const int c = active_cols[j];
218235
for (int ig = 0; ig < n_dim_; ++ig)
219236
x[idx(ig, c, ld_psi_)] /=
220237
std::max(prec[ig], static_cast<Real>(1.0e-12));
238+
}
221239
}
222240

223241
} // namespace hsolver

‎source/source_hsolver/ppcg/diago_ppcg_orth.hpp‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@ void DiagoPPCG<T, Device>::right_solve_upper(
88
const std::vector<T>& r, int n, std::vector<T>& x) const
99
{
1010
std::vector<T> b = x;
11+
#ifdef _OPENMP
12+
#pragma omp parallel for schedule(static) if (n_dim_ * n > 4096)
13+
#endif
1114
for (int row = 0; row < n_dim_; ++row)
1215
{
1316
for (int j = 0; j < n; ++j)
@@ -59,6 +62,9 @@ void DiagoPPCG<T, Device>::s_gram_schmidt(
5962
{
6063
T coeff = complex_dot(psi + k * ld_psi_,
6164
spsi + j * ld_psi_);
65+
#ifdef _OPENMP
66+
#pragma omp parallel for schedule(static) if (n_dim_ > 4096)
67+
#endif
6268
for (int ig = 0; ig < n_dim_; ++ig)
6369
{
6470
psi [idx(ig, j, ld_psi_)] -= coeff * psi [idx(ig, k, ld_psi_)];
@@ -72,6 +78,9 @@ void DiagoPPCG<T, Device>::s_gram_schmidt(
7278
gamma_dot(psi + j * ld_psi_, spsi + j * ld_psi_),
7379
static_cast<Real>(1e-30)));
7480
Real inv_nrm = static_cast<Real>(1) / nrm;
81+
#ifdef _OPENMP
82+
#pragma omp parallel for schedule(static) if (n_dim_ > 4096)
83+
#endif
7584
for (int ig = 0; ig < n_dim_; ++ig)
7685
{
7786
psi [idx(ig, j, ld_psi_)] *= inv_nrm;
@@ -223,6 +232,9 @@ void DiagoPPCG<T, Device>::rayleigh_ritz(
223232

224233
// Compute residual: w_i = H|psi_i> - eps_i * S|psi_i>
225234
set_zero(w_);
235+
#ifdef _OPENMP
236+
#pragma omp parallel for collapse(2) schedule(static) if (n_dim_ * n_band_ > 4096)
237+
#endif
226238
for (int j = 0; j < n_band_; ++j)
227239
for (int ig = 0; ig < n_dim_; ++ig)
228240
w_[idx(ig, j, ld_psi_)] = hpsi_[idx(ig, j, ld_psi_)]

‎source/source_hsolver/ppcg/diago_ppcg_subspace.hpp‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@ void DiagoPPCG<T, Device>::lock_epairs(
1717
for (int j = 0; j < n_band_; ++j)
1818
{
1919
Real nrm2 = 0;
20+
#ifdef _OPENMP
21+
#pragma omp parallel for reduction(+ : nrm2) schedule(static) if (n_dim_ > 4096)
22+
#endif
2023
for (int ig = 0; ig < n_dim_; ++ig)
2124
nrm2 += static_cast<Real>(std::norm(residual[idx(ig, j, ld_psi_)]));
2225
reduce_pool_if_mpi_ready(nrm2);
@@ -78,6 +81,9 @@ void DiagoPPCG<T, Device>::build_small_subspace(
7881
std::vector<Real>& scale) {
7982
for (int j = 0; j < lcols; ++j) {
8083
Real sn2 = 0;
84+
#ifdef _OPENMP
85+
#pragma omp parallel for reduction(+ : sn2) schedule(static) if (n_dim_ > 4096)
86+
#endif
8187
for (int ig = 0; ig < n_dim_; ++ig)
8288
sn2 += std::real(std::conj(x[idx(ig, j, ld_psi_)])
8389
* sx[idx(ig, j, ld_psi_)]);
@@ -88,6 +94,9 @@ void DiagoPPCG<T, Device>::build_small_subspace(
8894
if (sn > static_cast<Real>(1e-15)) {
8995
Real inv = static_cast<Real>(1) / sn;
9096
scale[j] = inv;
97+
#ifdef _OPENMP
98+
#pragma omp parallel for schedule(static) if (n_dim_ > 4096)
99+
#endif
91100
for (int ig = 0; ig < n_dim_; ++ig) {
92101
x[ idx(ig, j, ld_psi_)] *= inv;
93102
sx[idx(ig, j, ld_psi_)] *= inv;
@@ -217,6 +226,9 @@ void DiagoPPCG<T, Device>::update_one_block(
217226

218227
std::vector<T> coeff_state(dim * l, T(0));
219228
std::vector<T> coeff_dir(dim * l, T(0));
229+
#ifdef _OPENMP
230+
#pragma omp parallel for schedule(static) if (l * l > 4096)
231+
#endif
220232
for (int j = 0; j < l; ++j)
221233
{
222234
for (int i = 0; i < l; ++i)
@@ -240,6 +252,9 @@ void DiagoPPCG<T, Device>::update_one_block(
240252
std::vector<T>& basis)
241253
{
242254
basis.assign(ld_psi_ * dim, T(0));
255+
#ifdef _OPENMP
256+
#pragma omp parallel for schedule(static) if (ld_psi_ * l > 4096)
257+
#endif
243258
for (int j = 0; j < l; ++j)
244259
{
245260
std::copy(a.begin() + j * ld_psi_,

0 commit comments

Comments
 (0)