Skip to content

Commit 3251840

Browse files
cheerly-pkucheerly-pku
committed
Reuse PPCG block subspace buffers
1 parent a910418 commit 3251840

2 files changed

Lines changed: 56 additions & 58 deletions

File tree

‎source/source_hsolver/diago_ppcg.h‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,19 @@ class DiagoPPCG
162162
std::vector<T> k; // K matrix (projected H)
163163
std::vector<T> m; // M matrix (projected S)
164164
std::vector<Real> eval; // eigenvalues
165-
std::vector<Real> w_scale;
165+
std::vector<T> psi_l;
166+
std::vector<T> spsi_l;
167+
std::vector<T> hpsi_l;
168+
std::vector<T> w_l;
169+
std::vector<T> sw_l;
170+
std::vector<T> hw_l;
171+
std::vector<T> basis;
172+
std::vector<T> hbasis;
173+
std::vector<T> sbasis;
174+
std::vector<T> coeff_state;
175+
std::vector<T> psi_new;
176+
std::vector<T> spsi_new;
177+
std::vector<T> hpsi_new;
166178
};
167179

168180
void lock_epairs(const std::vector<T>& residual,
@@ -178,7 +190,7 @@ class DiagoPPCG
178190
void update_one_block(T* psi,
179191
const std::vector<int>& cols,
180192
int l,
181-
const SmallSubspace& subspace);
193+
SmallSubspace& subspace);
182194

183195
bool is_s_orthonormal(const T* psi, const T* spsi, int ncol) const;
184196

‎source/source_hsolver/ppcg/diago_ppcg_subspace.hpp‎

Lines changed: 42 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -51,16 +51,13 @@ void DiagoPPCG<T, Device>::build_small_subspace(
5151
subspace.k.resize(dim * dim);
5252
subspace.m.resize(dim * dim);
5353
subspace.eval.resize(dim);
54-
subspace.w_scale.assign(l, static_cast<Real>(1));
5554

56-
std::vector<T> psi_l, spsi_l, hpsi_l;
57-
std::vector<T> w_l, sw_l, hw_l;
58-
copy_cols(psi, cols, psi_l);
59-
copy_cols(spsi_.data(), cols, spsi_l);
60-
copy_cols(hpsi_.data(), cols, hpsi_l);
61-
copy_cols(w_.data(), cols, w_l);
62-
copy_cols(sw_.data(), cols, sw_l);
63-
copy_cols(hw_.data(), cols, hw_l);
55+
copy_cols(psi, cols, subspace.psi_l);
56+
copy_cols(spsi_.data(), cols, subspace.spsi_l);
57+
copy_cols(hpsi_.data(), cols, subspace.hpsi_l);
58+
copy_cols(w_.data(), cols, subspace.w_l);
59+
copy_cols(sw_.data(), cols, subspace.sw_l);
60+
copy_cols(hw_.data(), cols, subspace.hw_l);
6461

6562
// ---------------------------------------------------------------------------
6663
// Normalize w columns to unit S-norm for numerical stability.
@@ -70,12 +67,12 @@ void DiagoPPCG<T, Device>::build_small_subspace(
7067
// sygvd to produce garbage eigenvectors.
7168
//
7269
// Scaling to unit S-norm keeps M well-conditioned (diagonal ~1) without
73-
// changing the subspace. The Ritz values are identical and the Ritz
74-
// vector coefficients in update_one_block automatically compensate.
70+
// changing the subspace. The same scaled basis is reused in update_one_block.
7571
// ---------------------------------------------------------------------------
76-
auto scale_to_unit_snorm = [this](std::vector<T>& x, std::vector<T>& sx,
77-
std::vector<T>& hx, int lcols,
78-
std::vector<Real>& scale) {
72+
auto scale_to_unit_snorm = [this](std::vector<T>& x,
73+
std::vector<T>& sx,
74+
std::vector<T>& hx,
75+
int lcols) {
7976
std::vector<double> sn2_all(lcols, 0.0);
8077
#ifdef _OPENMP
8178
#pragma omp parallel for schedule(static) if (n_dim_ * lcols > 4096)
@@ -95,7 +92,6 @@ void DiagoPPCG<T, Device>::build_small_subspace(
9592
// column is a converged band whose contribution is harmless.
9693
if (sn > static_cast<Real>(1e-15)) {
9794
Real inv = static_cast<Real>(1) / sn;
98-
scale[j] = inv;
9995
#ifdef _OPENMP
10096
#pragma omp parallel for schedule(static) if (n_dim_ > 4096)
10197
#endif
@@ -107,7 +103,10 @@ void DiagoPPCG<T, Device>::build_small_subspace(
107103
}
108104
}
109105
};
110-
scale_to_unit_snorm(w_l, sw_l, hw_l, l, subspace.w_scale);
106+
scale_to_unit_snorm(subspace.w_l,
107+
subspace.sw_l,
108+
subspace.hw_l,
109+
l);
111110

112111
auto copy_block = [&](const std::vector<T>& src,
113112
const int col0,
@@ -137,18 +136,18 @@ void DiagoPPCG<T, Device>::build_small_subspace(
137136
}
138137
};
139138

140-
std::vector<T> basis(ld_psi_ * dim);
141-
std::vector<T> hbasis(ld_psi_ * dim);
142-
std::vector<T> sbasis(ld_psi_ * dim);
143-
copy_block(psi_l, 0, basis);
144-
copy_block(hpsi_l, 0, hbasis);
145-
copy_block(spsi_l, 0, sbasis);
146-
copy_block(w_l, l, basis);
147-
copy_block(hw_l, l, hbasis);
148-
copy_block(sw_l, l, sbasis);
139+
subspace.basis.resize(ld_psi_ * dim);
140+
subspace.hbasis.resize(ld_psi_ * dim);
141+
subspace.sbasis.resize(ld_psi_ * dim);
142+
copy_block(subspace.psi_l, 0, subspace.basis);
143+
copy_block(subspace.hpsi_l, 0, subspace.hbasis);
144+
copy_block(subspace.spsi_l, 0, subspace.sbasis);
145+
copy_block(subspace.w_l, l, subspace.basis);
146+
copy_block(subspace.hw_l, l, subspace.hbasis);
147+
copy_block(subspace.sw_l, l, subspace.sbasis);
149148

150-
gram(basis.data(), hbasis.data(), dim, dim, subspace.k, dim);
151-
gram(basis.data(), sbasis.data(), dim, dim, subspace.m, dim);
149+
gram(subspace.basis.data(), subspace.hbasis.data(), dim, dim, subspace.k, dim);
150+
gram(subspace.basis.data(), subspace.sbasis.data(), dim, dim, subspace.m, dim);
152151
hermitize(subspace.k);
153152
hermitize(subspace.m);
154153
}
@@ -208,35 +207,25 @@ void DiagoPPCG<T, Device>::update_one_block(
208207
T* psi,
209208
const std::vector<int>& cols,
210209
int l,
211-
const SmallSubspace& subspace)
210+
SmallSubspace& subspace)
212211
{
213212
const int dim = 2 * l;
214213
const T* eigvec = subspace.k.data();
215214

216-
std::vector<T> psi_l, spsi_l, hpsi_l;
217-
std::vector<T> w_l, sw_l, hw_l;
218-
copy_cols(psi, cols, psi_l);
219-
copy_cols(spsi_.data(), cols, spsi_l);
220-
copy_cols(hpsi_.data(), cols, hpsi_l);
221-
copy_cols(w_.data(), cols, w_l);
222-
copy_cols(sw_.data(), cols, sw_l);
223-
copy_cols(hw_.data(), cols, hw_l);
215+
subspace.psi_new.assign(ld_psi_ * l, T(0));
216+
subspace.spsi_new.assign(ld_psi_ * l, T(0));
217+
subspace.hpsi_new.assign(ld_psi_ * l, T(0));
224218

225-
std::vector<T> psi_new(ld_psi_ * l, T(0));
226-
std::vector<T> spsi_new(ld_psi_ * l, T(0));
227-
std::vector<T> hpsi_new(ld_psi_ * l, T(0));
228-
229-
std::vector<T> coeff_state(dim * l, T(0));
219+
subspace.coeff_state.resize(dim * l);
230220
#ifdef _OPENMP
231221
#pragma omp parallel for schedule(static) if (l * l > 4096)
232222
#endif
233223
for (int j = 0; j < l; ++j)
234224
{
235225
for (int i = 0; i < l; ++i)
236226
{
237-
coeff_state[i + j * dim] = eigvec[i + j * dim];
238-
const T cw = eigvec[(l + i) + j * dim] * subspace.w_scale[i];
239-
coeff_state[(l + i) + j * dim] = cw;
227+
subspace.coeff_state[i + j * dim] = eigvec[i + j * dim];
228+
subspace.coeff_state[(l + i) + j * dim] = eigvec[(l + i) + j * dim];
240229
}
241230
}
242231

@@ -280,20 +269,17 @@ void DiagoPPCG<T, Device>::update_one_block(
280269
ld_psi_);
281270
};
282271

283-
std::vector<T> psi_basis;
284-
std::vector<T> spsi_basis;
285-
std::vector<T> hpsi_basis;
286-
fill_basis(psi_l, w_l, psi_basis);
287-
fill_basis(spsi_l, sw_l, spsi_basis);
288-
fill_basis(hpsi_l, hw_l, hpsi_basis);
272+
fill_basis(subspace.psi_l, subspace.w_l, subspace.basis);
273+
fill_basis(subspace.spsi_l, subspace.sw_l, subspace.sbasis);
274+
fill_basis(subspace.hpsi_l, subspace.hw_l, subspace.hbasis);
289275

290-
combine(psi_basis, coeff_state, psi_new);
291-
combine(spsi_basis, coeff_state, spsi_new);
292-
combine(hpsi_basis, coeff_state, hpsi_new);
276+
combine(subspace.basis, subspace.coeff_state, subspace.psi_new);
277+
combine(subspace.sbasis, subspace.coeff_state, subspace.spsi_new);
278+
combine(subspace.hbasis, subspace.coeff_state, subspace.hpsi_new);
293279

294-
scatter_cols(psi, cols, psi_new);
295-
scatter_cols(spsi_.data(), cols, spsi_new);
296-
scatter_cols(hpsi_.data(), cols, hpsi_new);
280+
scatter_cols(psi, cols, subspace.psi_new);
281+
scatter_cols(spsi_.data(), cols, subspace.spsi_new);
282+
scatter_cols(hpsi_.data(), cols, subspace.hpsi_new);
297283
}
298284

299285
} // namespace hsolver

0 commit comments

Comments
 (0)