|
Quantum 3xplorer (Q3)
An AI-native quantum simulation engine
|
#include <serial.hpp>
Inheritance diagram for q3::SerialPlanewaveSubspaceManager:Public Member Functions | |
| SerialPlanewaveSubspaceManager (const execution::ExecutionPolicy &execution_policy) | |
| void | apply_preconditioner (int number_of_bands, int number_of_plane_waves, const Real *kinetic_energy_squared, const Real *eigenvalues, const std::vector< bool > &band_converged, Complex *residual_wavefunctions) override |
| void | apply_preconditioner_to_columns (int number_of_columns, int number_of_plane_waves, const Real *kinetic_energy_squared, const Real *eigenvalues, Complex *columns) const |
| int | build_correction_vectors (int basis_size, int number_of_plane_waves, int number_of_bands, int maximum_candidates, bool is_gamma_only, const Complex *eigenvector_matrix, int eigenvector_leading_dimension, const Real *eigenvalues, const std::vector< bool > &band_converged, const Real *kinetic_energy_squared, const Complex *hamiltonian_basis_vectors, const Complex *overlap_basis_vectors, Complex *correction_columns) override |
| int | collapse_subspace (int basis_size, int number_of_plane_waves, int number_of_bands, int maximum_basis_size, bool is_gamma_only, bool is_identity_overlap, const OverlapOperator *overlap_operator, const Complex *subspace_hamiltonian, const Complex *subspace_overlap, Complex *basis_vectors, Complex *hamiltonian_basis_vectors, Complex *overlap_basis_vectors, Complex *rotation_scratch) override |
| int | diagonalize_subspace (int basis_size, int maximum_basis_size, bool is_gamma_only, const Complex *subspace_hamiltonian, const Complex *subspace_overlap, Real *eigenvalues, Complex *eigenvector_matrix) override |
| void | extract_eigenvectors (int basis_size, int number_of_plane_waves, int number_of_bands, int kpoint_index_state, bool is_gamma_only, const Complex *eigenvector_matrix, const Complex *basis_vectors, const Real *eigenvalues, Complex *rotated_scratch, PlanewaveKohnShamQuantumSystem &quantum_system) override |
| void | multiply_basis_by_coefficients_across_bands (int number_of_plane_waves, int number_of_bands, int basis_size, const std::complex< double > *basis, const std::complex< double > *coefficients, std::complex< double > *output) const |
| int | orthogonalize_batch (int starting_vector_index, int number_of_candidates, int number_of_plane_waves, bool is_gamma_only, const OverlapOperator *overlap_operator, Complex *basis_vectors, Complex *overlap_basis_vectors) override |
| void | set_gamma_origin_is_local (bool origin_is_local) override |
| void | update_subspace_matrices (int start_column, int new_columns, int total_basis_size, int maximum_basis_size, int number_of_plane_waves, bool is_gamma_only, const Complex *basis_vectors, const Complex *hamiltonian_basis_vectors, const Complex *overlap_basis_vectors, Complex *subspace_hamiltonian, Complex *subspace_overlap) override |
Public Member Functions inherited from q3::PlanewaveSubspaceManager | |
| virtual | ~PlanewaveSubspaceManager ()=default |
Definition at line 31 of file serial.hpp.
|
inlineexplicit |
Definition at line 33 of file serial.hpp.
|
inlineoverridevirtual |
Apply Davidson preconditioner to residual wavefunctions.
Implements q3::PlanewaveSubspaceManager.
Definition at line 1158 of file serial.hpp.
References q3::math::tolerances::preconditioner_denominator_floor.
|
inline |
Teter-Payne-Allan preconditioning of a compact block of columns, applied in place. The legacy per-band entry point below delegates to this.
Definition at line 1138 of file serial.hpp.
References q3::math::tolerances::preconditioner_denominator_floor.
Referenced by build_correction_vectors().
|
inlineoverridevirtual |
Compute residual vectors: R = Hψ - εSψ. Build preconditioned correction vectors straight into the free columns of the basis, and return how many were written.
R_n = (HV)c_n - eps_n (SV)c_n, formed by two GEMMs whose destination is the basis itself. Nothing else is materialised: no Ritz vectors, no H|Ritz>, no separate residual array. That is the layout Quantum ESPRESSO's Davidson uses, and it is three arrays of number_of_bands x number_of_plane_waves that Q3 no longer has to hold.
Only unconverged bands contribute, and they are compacted into consecutive columns as they are written, so the caller can hand the block straight to orthogonalize_batch.
Gather the unconverged bands' eigenvector columns, compacted, so the two GEMMs below write consecutive columns and the caller can hand the block straight to orthogonalisation.
R = (SV)c scaled by -eps, then (HV)c accumulated on top. Under band distribution each rank owns a slice of the candidate columns and both GEMMs plus the scaling touch only that slice, so one allgather at the end republishes the block.
Implements q3::PlanewaveSubspaceManager.
Definition at line 995 of file serial.hpp.
References q3::execution::ExecutionPolicy::allgather_across_bands(), apply_preconditioner_to_columns(), q3::execution::ExecutionPolicy::band_distribution_rank_index(), q3::execution::ExecutionPolicy::band_distribution_size(), q3::execution::BandPartition::counts, q3::execution::BandPartition::displacements, q3::execution::BandPartition::local_first_band, q3::execution::BandPartition::local_number_of_bands, q3::execution::partition_bands(), and q3::blas::zgemm_().
|
inlineoverridevirtual |
Collapse the subspace back to number_of_bands vectors by rotating the basis with the eigenvector matrix.
The rotation needs somewhere to build its result before overwriting the array it is reading, and the caller supplies one buffer of at least number_of_bands * number_of_plane_waves. One is enough because each rotation reads a different array than it writes. These used to be three members sized by maximum_basis_size, which made the collapse as expensive in memory as the basis it collapses.
The rotation matrix is the eigenvector matrix of the current subspace, and the caller already holds that subspace: subspace_hamiltonian and subspace_overlap are the incrementally updated matrices with leading dimension maximum_basis_size, complete through basis_size. Diagonalising them costs basis_size cubed; rebuilding them from the vectors first – which is what this used to do – costs two products of basis_size squared by number_of_plane_waves on top, and was the largest single compute term of a restart. Pass nullptr for either to force that rebuild, which is also the fallback when the passed matrices cannot be diagonalised: the Davidson loop asks for a collapse precisely when its own diagonalisation of them has just failed.
Returns the new basis size (= number_of_bands).
The subspace the caller holds is complete through basis_size, so its eigenvectors are the rotation. Rebuilding it from the vectors is kept only as the fallback, for the recovery collapse that follows a failed diagonalisation of these very matrices.
This diagonalisation is local to the collapse, so it is not covered by the Davidson loop's eigenvector broadcast. Without republishing it, each rank would rotate its band range by a slightly different matrix and the allgathered basis would be a stitch of inconsistent rotations.
Only the first number_of_bands columns survive a collapse, so only that much is rotated and only that much is copied back. Columns past it are left as they were rather than zeroed: every reader downstream is bounded by basis_size, which the caller resets to number_of_bands, and the expansion writes each new column before anything reads it.
One scratch buffer does for all three rotations, because each one reads a different array than it writes: rotate into the scratch, copy back, move on. The overlap array is the exception and needs no rotation at all when it aliases the basis – it already holds the rotated vectors by the time the first copy-back lands.
Skip the overlap rotation only when the overlap array is the basis itself. An identity overlap is not enough: the CPU workspace aliases S|V> to |V> for it, so rotating the basis has already rotated S|V>, but the device workspace keeps S|V> in a buffer of its own. Keying the skip on the flag left that buffer unrotated after every collapse, so the next orthogonalisation and subspace update read stale S|V>, and a run that collapsed often converged to the wrong answer – 1.77 Ry low on the two-atom silicon deck at expansion factor 2, which collapses every iteration. Factor 4 hid it: small decks converge before the basis fills.
Implements q3::PlanewaveSubspaceManager.
Definition at line 825 of file serial.hpp.
References q3::execution::ExecutionPolicy::broadcast_across_bands(), diagonalize_subspace(), and multiply_basis_by_coefficients_across_bands().
|
inlineoverridevirtual |
Diagonalize the subspace by solving the generalized eigenvalue problem H_sub · c = ε · S_sub · c.
Returns the LAPACK info code (0 = success).
The solver overwrites both inputs, and the caller's matrices carry the next iteration's incremental update, so it works on copies. The Hamiltonian's copy is the eigenvector buffer itself – the solver returns the eigenvectors where the matrix was – and only the overlap needs a buffer of its own. Both are compacted from the caller's leading dimension to basis_size. That is two matrices of basis_size squared here rather than three, on every rank.
Which ranks run the eigensolver, in three cases:
The two conditions must agree. Testing group_rank_index() == 0 to solve but group_size() > 1 to broadcast leaves the non-root ranks of a replicated group with neither a result nor a broadcast, and they then diverge from the root at the next collective.
Implements q3::PlanewaveSubspaceManager.
Definition at line 680 of file serial.hpp.
References q3::execution::ExecutionPolicy::broadcast_from_group_root(), q3::execution::ExecutionPolicy::group_rank_index(), q3::execution::ExecutionPolicy::group_size(), and q3::log::warn().
Referenced by collapse_subspace().
|
inlineoverridevirtual |
Extract converged eigenvectors from the subspace representation and store them in the quantum system.
Implements q3::PlanewaveSubspaceManager.
Definition at line 947 of file serial.hpp.
References q3::PlanewaveKohnShamQuantumSystem::set_eigenvalue(), q3::PlanewaveKohnShamQuantumSystem::wavefunction_band_data(), and q3::blas::zgemm_().
|
inline |
output = basis * coefficients, divided by band when bands are distributed.
The product is (plane_waves x basis_size) times (basis_size x bands), so its columns are bands and each rank can form its own contiguous range against the matching column slice of the coefficient matrix. One allgather then completes the array on every rank. Because every block has exactly one source, the assembled result is identical on all ranks – the property the replicated basis depends on and redundant ZGEMMs cannot give. The inputs must already agree across ranks, which the Davidson loop guarantees by broadcasting the basis and the eigenvector matrix.
Definition at line 770 of file serial.hpp.
References q3::execution::ExecutionPolicy::allgather_across_bands(), q3::execution::ExecutionPolicy::band_distribution_rank_index(), q3::execution::ExecutionPolicy::band_distribution_size(), q3::execution::partition_bands(), and q3::blas::zgemm_().
Referenced by collapse_subspace().
|
inlineoverridevirtual |
S-orthogonalize and normalize a batch of candidate vectors against the existing basis using modified Gram-Schmidt with S inner product.
Returns the number of accepted (linearly independent) vectors.
When the overlap operator is the identity – every norm-conserving run – S|v> is |v>, and the caller hands us the same buffer twice rather than storing a second copy of the basis. Reads then resolve correctly on their own; what must not happen is applying an update to the same memory twice, so every write below is guarded.
Already subtracted from the basis; repeating it here would double it.
Already subtracted from the basis; repeating it here would double it.
The basis was already solved against this factor.
The basis was already solved against this factor.
The basis copy above already moved these coefficients.
Already subtracted from the basis vector.
Already subtracted from the basis vector.
The basis vector carries the scaling already.
Implements q3::PlanewaveSubspaceManager.
Definition at line 41 of file serial.hpp.
References q3::OverlapOperator::apply_batch(), q3::blas::dgemm_(), q3::blas::dgemv_(), q3::blas::dpotrf_(), q3::blas::dtrsm_(), q3::math::tolerances::iterative_solver_norm_threshold, q3::execution::ExecutionPolicy::reduce_within_group(), q3::blas::zdscal_(), q3::blas::zgemm_(), q3::blas::zgemv_(), q3::blas::zpotrf_(), and q3::blas::ztrsm_().
|
inlineoverridevirtual |
Whether this rank's plane-wave list begins with G = 0.
Every gamma-only inner product below is 2 Re over the stored half sphere with the origin – held once, counted twice – removed, and then reduced across the ranks sharing the coefficients. Under a reciprocal-space distribution exactly one rank holds the origin, and only that rank may remove it. Replicated layouts always hold it, which is the default. Set once per k-point, before the solve.
Implements q3::PlanewaveSubspaceManager.
Definition at line 37 of file serial.hpp.
|
inlineoverridevirtual |
Incrementally update the subspace matrices H_sub and S_sub with new_columns vectors starting at start_column.
Implements q3::PlanewaveSubspaceManager.
Definition at line 555 of file serial.hpp.
References q3::blas::dgemm_(), q3::execution::ExecutionPolicy::reduce_within_group(), and q3::blas::zgemm_().