Skip to content

Commit 592e7d1

Browse files
author
abacus_fixer
committed
Refactor DeltaSpin, step 4: encapsulate PW subspace cache in SubspaceCache
Move the three ad-hoc public SpinConstrain members sub_h_save / sub_s_save / becp_save (raw TK* pointers) plus the lambda_in_sub_ snapshot into a new spinconstrain::pw::SubspaceCache (deltaspin_pw_cache.h), owned by value as SpinConstrain::pw_cache_. The class encapsulates the CPU vs GPU allocation/free difference: - allocate_cpu()/release_cpu() use new[]/delete[] on the host; - allocate_gpu()/release_gpu() use base_device resize/delete_memory_op on DEVICE_GPU, guarded by #if __CUDA/__ROCM. It still exposes raw per-k pointers h_k()/s_k()/becp_k() because the hsolver subspace routines and GPU memcpy ops require raw pointers, so std::vector is not applicable for the device buffers. The buffer element type is fixed to std::complex<double> (the PW path is always complex; the TK=double stub never allocates the cache). This also fixes a latent bug: the old SpinConstrain destructor called delete[] on sub_h_save/sub_s_save/becp_save unconditionally. In GPU runs those pointers are device memory allocated with resize_memory_op<DEVICE_GPU>, so delete[] on them is undefined behavior (and on some setups an invalid free). The destructor is now trivial (= default); device buffers are correctly freed via release_gpu() -> delete_memory_op<DEVICE_GPU>() in update_psi_charge_pw_gpu(), and host buffers via release_cpu() in update_psi_charge_pw_cpu(). The singleton lives for the whole program so no leak is introduced. Call sites updated: cal_mw_from_lambda.cpp (CPU+GPU allocation points and per-k views) and deltaspin_pw_impl.cpp (asserts, per-k views, CPU/GPU release, lambda_in_sub_ access). Buffer layout and reuse semantics are unchanged. No INPUT parameter behavior changes; docs update not required. Verification: - cmake --build build -j 16 (CPU, __LCAO): success - cmake --build build_pw_gpu --target module_pwdft (CUDA, nvcc 12.9): success (covers the __CUDA allocation/free branch) - OMP_NUM_THREADS=1 ctest --test-dir build -R deltaspin: 5/5 passed - python3 tools/03_code_analysis/agent_governance_check.py --staged: no findings
1 parent c49c67b commit 592e7d1

4 files changed

Lines changed: 172 additions & 54 deletions

File tree

‎source/source_lcao/module_deltaspin/cal_mw_from_lambda.cpp‎

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -164,23 +164,21 @@ void spinconstrain::SpinConstrain<std::complex<double>>::cal_mw_from_lambda(
164164
becp_tmp.resize(size_becp * nk);
165165
std::vector<std::complex<double>> h_tmp(nbands * nbands), s_tmp(nbands * nbands);
166166
int initial_hs = 0;
167-
if(this->sub_h_save == nullptr)
167+
if(!this->pw_cache_.allocated())
168168
{
169169
// FIRST CALL: save subspace data for reuse across lambda steps
170170
initial_hs = 1;
171-
this->sub_h_save = new std::complex<double>[nbands * nbands * nk];
172-
this->sub_s_save = new std::complex<double>[nbands * nbands * nk];
173-
this->becp_save = new std::complex<double>[size_becp * nk];
174-
this->lambda_in_sub_ = this->state_.lambda_;
171+
this->pw_cache_.allocate_cpu(nbands, nk, size_becp);
172+
this->pw_cache_.lambda_in_sub() = this->state_.lambda_;
175173
}
176174
for (int ik = 0; ik < nk; ++ik)
177175
{
178176

179177
psi_t->fix_k(ik);
180178

181-
std::complex<double>* h_k = this->sub_h_save + ik * nbands * nbands;
182-
std::complex<double>* s_k = this->sub_s_save + ik * nbands * nbands;
183-
std::complex<double>* becp_k = this->becp_save + ik * size_becp;
179+
std::complex<double>* h_k = this->pw_cache_.h_k(ik, nbands);
180+
std::complex<double>* s_k = this->pw_cache_.s_k(ik, nbands);
181+
std::complex<double>* becp_k = this->pw_cache_.becp_k(ik, size_becp);
184182
if(initial_hs)
185183
{
186184
/// Compute H(k) and extract subspace matrices for this k-point
@@ -224,23 +222,21 @@ void spinconstrain::SpinConstrain<std::complex<double>>::cal_mw_from_lambda(
224222
base_device::memory::resize_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(h_tmp, nbands * nbands);
225223
base_device::memory::resize_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(s_tmp, nbands * nbands);
226224
int initial_hs = 0;
227-
if(this->sub_h_save == nullptr)
225+
if(!this->pw_cache_.allocated())
228226
{
229227
initial_hs = 1;
230-
base_device::memory::resize_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(this->sub_h_save, nbands * nbands * nk);
231-
base_device::memory::resize_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(this->sub_s_save, nbands * nbands * nk);
232-
base_device::memory::resize_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(this->becp_save, size_becp * nk);
233-
this->lambda_in_sub_ = this->state_.lambda_;
228+
this->pw_cache_.allocate_gpu(nbands, nk, size_becp);
229+
this->pw_cache_.lambda_in_sub() = this->state_.lambda_;
234230
}
235231
std::complex<double>* becp_pointer = nullptr;
236232
base_device::memory::resize_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(becp_pointer, size_becp);
237233
for (int ik = 0; ik < nk; ++ik)
238234
{
239235
psi_t->fix_k(ik);
240236

241-
std::complex<double>* h_k = this->sub_h_save + ik * nbands * nbands;
242-
std::complex<double>* s_k = this->sub_s_save + ik * nbands * nbands;
243-
std::complex<double>* becp_k = this->becp_save + ik * size_becp;
237+
std::complex<double>* h_k = this->pw_cache_.h_k(ik, nbands);
238+
std::complex<double>* s_k = this->pw_cache_.s_k(ik, nbands);
239+
std::complex<double>* becp_k = this->pw_cache_.becp_k(ik, size_becp);
244240
if(initial_hs)
245241
{
246242
hamilt_t->updateHk(ik);

‎source/source_lcao/module_deltaspin/deltaspin_pw_cache.h‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
/**
2+
* @file deltaspin_pw_cache.h
3+
* @brief PW-basis subspace data cache for DeltaSpin, decoupled from SpinConstrain.
4+
*
5+
* @par Purpose
6+
* In the PW basis, the subspace Hamiltonian H_sub = <psi|H|psi>, overlap S_sub
7+
* and becp coefficients are expensive to compute. They are cached on the first
8+
* cal_mw_from_lambda() call and reused across multiple lambda steps within the
9+
* same SCF iteration, then freed after the final subspace diagonalization in
10+
* update_psi_charge_pw_{cpu,gpu}().
11+
*
12+
* This class owns the three raw device/host pointers plus the lambda snapshot
13+
* taken when the cache was filled, replacing the ad-hoc new[]/delete[] that used
14+
* to live as public SpinConstrain members. It encapsulates the CPU vs GPU
15+
* allocation/free difference behind allocate()/release(), while still exposing
16+
* raw per-k pointers (h_k/s_k/becp_k) because the hsolver subspace routines and
17+
* GPU memcpy ops require raw pointers.
18+
*
19+
* @par Layout (same as before, unchanged)
20+
* - h(ik)[i * nbands + j]: H_sub for k-point ik
21+
* - s(ik): same layout for overlap S_sub
22+
* - becp(ik)[ib * nkb * npol + ip]: becp coefficients
23+
*/
24+
#ifndef DELTASPIN_PW_CACHE_H
25+
#define DELTASPIN_PW_CACHE_H
26+
27+
#include <complex>
28+
#include <vector>
29+
30+
#include "source_base/vector3.h"
31+
#include "source_base/module_device/memory_op.h"
32+
33+
namespace spinconstrain
34+
{
35+
namespace pw
36+
{
37+
38+
/**
39+
* @brief Owning cache of PW subspace H/S/becp data plus the lambda snapshot.
40+
*
41+
* The buffer element type is std::complex<double> because the PW DeltaSpin path
42+
* is always instantiated on complex wavefunctions; the legacy TK=double stub
43+
* never allocates it.
44+
*/
45+
class SubspaceCache
46+
{
47+
public:
48+
SubspaceCache() = default;
49+
50+
// Owns raw memory; non-copyable, non-movable to keep ownership unambiguous.
51+
SubspaceCache(const SubspaceCache&) = delete;
52+
SubspaceCache& operator=(const SubspaceCache&) = delete;
53+
54+
/// True when the subspace buffers are allocated.
55+
bool allocated() const { return sub_h_save_ != nullptr; }
56+
57+
/// Lambda values captured when the cache was filled.
58+
std::vector<ModuleBase::Vector3<double>>& lambda_in_sub() { return lambda_in_sub_; }
59+
const std::vector<ModuleBase::Vector3<double>>& lambda_in_sub() const { return lambda_in_sub_; }
60+
61+
/// Raw base pointers (needed by hsolver subspace ops and GPU memcpy).
62+
std::complex<double>* h() { return sub_h_save_; }
63+
std::complex<double>* s() { return sub_s_save_; }
64+
std::complex<double>* becp() { return becp_save_; }
65+
66+
/// Per-k-point views.
67+
std::complex<double>* h_k(int ik, int nbands) { return sub_h_save_ + ik * nbands * nbands; }
68+
std::complex<double>* s_k(int ik, int nbands) { return sub_s_save_ + ik * nbands * nbands; }
69+
std::complex<double>* becp_k(int ik, int size_becp) { return becp_save_ + ik * size_becp; }
70+
71+
/**
72+
* @brief Allocate the three buffers on the host (CPU path) with new[].
73+
* No-op if already allocated.
74+
*/
75+
void allocate_cpu(int nbands, int nk, int size_becp)
76+
{
77+
if (allocated())
78+
{
79+
return;
80+
}
81+
sub_h_save_ = new std::complex<double>[nbands * nbands * nk];
82+
sub_s_save_ = new std::complex<double>[nbands * nbands * nk];
83+
becp_save_ = new std::complex<double>[size_becp * nk];
84+
}
85+
86+
/**
87+
* @brief Release the host (CPU) buffers with delete[].
88+
*/
89+
void release_cpu()
90+
{
91+
delete[] sub_h_save_;
92+
delete[] sub_s_save_;
93+
delete[] becp_save_;
94+
sub_h_save_ = nullptr;
95+
sub_s_save_ = nullptr;
96+
becp_save_ = nullptr;
97+
}
98+
99+
#if ((defined __CUDA) || (defined __ROCM))
100+
/**
101+
* @brief Allocate the three buffers on the device (GPU path).
102+
* No-op if already allocated.
103+
*/
104+
void allocate_gpu(int nbands, int nk, int size_becp)
105+
{
106+
if (allocated())
107+
{
108+
return;
109+
}
110+
using mem = base_device::memory::resize_memory_op<std::complex<double>, base_device::DEVICE_GPU>;
111+
mem()(sub_h_save_, nbands * nbands * nk);
112+
mem()(sub_s_save_, nbands * nbands * nk);
113+
mem()(becp_save_, size_becp * nk);
114+
}
115+
116+
/**
117+
* @brief Release the device (GPU) buffers.
118+
*/
119+
void release_gpu()
120+
{
121+
using del = base_device::memory::delete_memory_op<std::complex<double>, base_device::DEVICE_GPU>;
122+
del()(sub_h_save_);
123+
del()(sub_s_save_);
124+
del()(becp_save_);
125+
sub_h_save_ = nullptr;
126+
sub_s_save_ = nullptr;
127+
becp_save_ = nullptr;
128+
}
129+
#endif // __CUDA || __ROCM
130+
131+
private:
132+
std::complex<double>* sub_h_save_ = nullptr; ///< Cached subspace Hamiltonian for all k-points
133+
std::complex<double>* sub_s_save_ = nullptr; ///< Cached subspace overlap matrix for all k-points
134+
std::complex<double>* becp_save_ = nullptr; ///< Cached becp coefficients for all k-points
135+
std::vector<ModuleBase::Vector3<double>> lambda_in_sub_; ///< Lambda when the cache was saved
136+
};
137+
138+
} // namespace pw
139+
} // namespace spinconstrain
140+
141+
#endif // DELTASPIN_PW_CACHE_H

‎source/source_lcao/module_deltaspin/spin_constrain.h‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
#include "source_estate/elecstate.h"
6161

6262
#include "deltaspin_state.h"
63+
#include "deltaspin_pw_cache.h"
6364

6465
#ifdef __LCAO
6566
#include "source_estate/module_dm/density_matrix.h" // mohan add 2025-11-02
@@ -456,15 +457,10 @@ class SpinConstrain
456457

457458
private:
458459
SpinConstrain(){};
459-
~SpinConstrain()
460-
{
461-
delete[] sub_h_save;
462-
delete[] sub_s_save;
463-
delete[] becp_save;
464-
sub_h_save = nullptr;
465-
sub_s_save = nullptr;
466-
becp_save = nullptr;
467-
};
460+
// Subspace buffers are owned by `pw_cache_` (RAII via release_cpu/gpu in the
461+
// PW update paths). The destructor is trivial; the singleton lives for the
462+
// whole program and the cache is released by the PW update functions.
463+
~SpinConstrain() = default;
468464
SpinConstrain& operator=(SpinConstrain const&) = delete; ///< Copy assignment deleted
469465
SpinConstrain& operator=(SpinConstrain &&) = delete; ///< Move assignment deleted
470466

@@ -509,10 +505,9 @@ class SpinConstrain
509505
* update_psi_charge_pw_cpu/gpu() after final subspace diagonalization.
510506
*/
511507
public:
512-
TK* sub_h_save = nullptr; ///< Cached subspace Hamiltonian for all k-points
513-
TK* sub_s_save = nullptr; ///< Cached subspace overlap matrix for all k-points
514-
TK* becp_save = nullptr; ///< Cached becp coefficients for all k-points
515-
std::vector<ModuleBase::Vector3<double>> lambda_in_sub_; ///< Lambda values when subspace was saved
508+
/// PW subspace data cache (H_sub/S_sub/becp + lambda snapshot). Owned object;
509+
/// buffers are device/host memory managed via allocate_cpu/gpu + release_cpu/gpu.
510+
pw::SubspaceCache pw_cache_;
516511

517512
private:
518513
/// RMS error of the most recent lambda optimization loop; -1.0 if no loop has run.

‎source/source_pw/module_pwdft/deltaspin_pw_impl.cpp‎

Lines changed: 11 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ void SpinConstrain<std::complex<double>>::calculate_delta_hcc(std::complex<doubl
125125
actual_delta.resize(nat);
126126
for (int iat = 0; iat < nat; iat++)
127127
{
128-
actual_delta[iat] = delta_lambda[iat] - this->lambda_in_sub_[iat];
128+
actual_delta[iat] = delta_lambda[iat] - this->pw_cache_.lambda_in_sub()[iat];
129129
}
130130
effective_lambda = actual_delta.data();
131131
}
@@ -315,9 +315,7 @@ void SpinConstrain<std::complex<double>>::update_psi_charge_pw_cpu(const ModuleB
315315
std::vector<std::complex<double>> h_tmp(nbands * nbands), s_tmp(nbands * nbands);
316316

317317
// CRITICAL: subspace data must have been saved by cal_mw_from_lambda()
318-
assert(this->sub_h_save != nullptr);
319-
assert(this->sub_s_save != nullptr);
320-
assert(this->becp_save != nullptr);
318+
assert(this->pw_cache_.allocated());
321319

322320
// Determine which lambda to use for H correction
323321
const ModuleBase::Vector3<double>* lambda_for_hcc = delta_lambda;
@@ -332,9 +330,9 @@ void SpinConstrain<std::complex<double>>::update_psi_charge_pw_cpu(const ModuleB
332330
// =============================================================
333331
for (int ik = 0; ik < nk; ++ik)
334332
{
335-
std::complex<double>* h_k = this->sub_h_save + ik * nbands * nbands;
336-
std::complex<double>* s_k = this->sub_s_save + ik * nbands * nbands;
337-
std::complex<double>* becp_k = this->becp_save + ik * size_becp;
333+
std::complex<double>* h_k = this->pw_cache_.h_k(ik, nbands);
334+
std::complex<double>* s_k = this->pw_cache_.s_k(ik, nbands);
335+
std::complex<double>* becp_k = this->pw_cache_.becp_k(ik, size_becp);
338336

339337
psi_t->fix_k(ik);
340338

@@ -354,12 +352,7 @@ void SpinConstrain<std::complex<double>>::update_psi_charge_pw_cpu(const ModuleB
354352
}
355353

356354
// Free saved subspace data (allocated in cal_mw_from_lambda)
357-
delete[] this->sub_h_save;
358-
delete[] this->sub_s_save;
359-
delete[] this->becp_save;
360-
this->sub_h_save = nullptr;
361-
this->sub_s_save = nullptr;
362-
this->becp_save = nullptr;
355+
this->pw_cache_.release_cpu();
363356

364357
// =============================================================
365358
// STAGE 2: Full-space update
@@ -435,9 +428,7 @@ void SpinConstrain<std::complex<double>>::update_psi_charge_pw_gpu(const ModuleB
435428
base_device::memory::resize_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(h_tmp, nbands * nbands);
436429
base_device::memory::resize_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(s_tmp, nbands * nbands);
437430

438-
assert(this->sub_h_save != nullptr);
439-
assert(this->sub_s_save != nullptr);
440-
assert(this->becp_save != nullptr);
431+
assert(this->pw_cache_.allocated());
441432

442433
const ModuleBase::Vector3<double>* lambda_for_hcc = delta_lambda;
443434
std::vector<ModuleBase::Vector3<double>> computed_delta;
@@ -449,9 +440,9 @@ void SpinConstrain<std::complex<double>>::update_psi_charge_pw_gpu(const ModuleB
449440
// STAGE 1: Subspace diagonalization for each k-point (GPU)
450441
for (int ik = 0; ik < nk; ++ik)
451442
{
452-
std::complex<double>* h_k = this->sub_h_save + ik * nbands * nbands;
453-
std::complex<double>* s_k = this->sub_s_save + ik * nbands * nbands;
454-
std::complex<double>* becp_k = this->becp_save + ik * size_becp;
443+
std::complex<double>* h_k = this->pw_cache_.h_k(ik, nbands);
444+
std::complex<double>* s_k = this->pw_cache_.s_k(ik, nbands);
445+
std::complex<double>* becp_k = this->pw_cache_.becp_k(ik, size_becp);
455446

456447
psi_t->fix_k(ik);
457448

@@ -471,12 +462,7 @@ void SpinConstrain<std::complex<double>>::update_psi_charge_pw_gpu(const ModuleB
471462
base_device::memory::delete_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(s_tmp);
472463

473464
// Free GPU memory for saved subspace data
474-
base_device::memory::delete_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(sub_h_save);
475-
base_device::memory::delete_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(sub_s_save);
476-
base_device::memory::delete_memory_op<std::complex<double>, base_device::DEVICE_GPU>()(becp_save);
477-
this->sub_h_save = nullptr;
478-
this->sub_s_save = nullptr;
479-
this->becp_save = nullptr;
465+
this->pw_cache_.release_gpu();
480466

481467
// STAGE 2: Full-space update (GPU)
482468
if (pw_solve)

0 commit comments

Comments
 (0)