Skip to content

Refactor: inject nlocal/nbands into the LCAO dense diagonalizers #7721

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Critsium-xy wants to merge 2 commits into
base: develop
Choose a base branch
from
Draft

Refactor: inject nlocal/nbands into the LCAO dense diagonalizers #7721

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0152800
Refactor: inject nlocal/nbands into the LCAO dense diagonalizers
Critsium-xy Jul 31, 2026
d87b4b3
Merge branch 'develop' into refactor/remove-param-hsolver-d
Critsium-xy Aug 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion source/source_esolver/esolver_ks_lcao.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -450,7 +450,9 @@ void ESolver_KS_LCAO<TK, TR>::hamilt2rho_single(UnitCell& ucell, int istep, int
PARAM.inp.ks_solver,
PARAM.globalv.kpar_lcao,
PARAM.globalv.nlocal,
PARAM.inp.nelec);
PARAM.inp.nbands,
PARAM.inp.nelec,
PARAM.inp.device == "gpu");
hsolver_lcao_obj.solve(static_cast<hamilt::Hamilt<TK>*>(this->p_hamilt), this->psi[0], this->pelec, *this->dmat.dm,
this->chr, PARAM.inp.nspin, skip_charge);
}
Expand Down
4 changes: 3 additions & 1 deletion source/source_esolver/esolver_ks_lcao_tddft.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,7 +351,9 @@ void ESolver_KS_LCAO_TDDFT<TR, Device>::hamilt2rho_single(UnitCell& ucell,
PARAM.inp.ks_solver,
PARAM.globalv.kpar_lcao,
PARAM.globalv.nlocal,
PARAM.inp.nelec);
PARAM.inp.nbands,
PARAM.inp.nelec,
PARAM.inp.device == "gpu");
hsolver_lcao_obj.solve(static_cast<hamilt::Hamilt<std::complex<double>>*>(this->p_hamilt),
this->psi[0],
this->pelec,
Expand Down
7 changes: 3 additions & 4 deletions source/source_hsolver/diago_cusolver.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
#include "source_base/module_external/scalapack_connector.h"
#include "source_base/tool_title.h"
#include "source_base/timer.h"
#include "source_io/module_parameter/parameter.h"

#include <memory>
#include <type_traits>
Expand All @@ -22,7 +21,7 @@ template <typename T>
int DiagoCusolver<T>::DecomposedState = 0;

template <typename T>
DiagoCusolver<T>::DiagoCusolver()
DiagoCusolver<T>::DiagoCusolver(const int nlocal_in, const int nbands_in) : nlocal(nlocal_in), nbands(nbands_in)
{
}

Expand All @@ -42,13 +41,13 @@ void DiagoCusolver<T>::diag(
ModuleBase::TITLE("DiagoCusolver", "diag");
ModuleBase::timer::start("DiagoCusolver", "cusolver");
// Allocate memory for eigenvalues
std::vector<double> eigen(PARAM.globalv.nlocal, 0.0);
std::vector<double> eigen(this->nlocal, 0.0);
std::vector<T> eigenvectors(h_mat.row * h_mat.col);
this->dc.Dngvd(h_mat.row, h_mat.col, h_mat.p, s_mat.p, eigen.data(), eigenvectors.data());
const int size = psi.get_nbands() * psi.get_nbasis();
BlasConnector::copy(size, eigenvectors.data(), 1, psi.get_pointer(), 1);
const int inc = 1;
BlasConnector::copy(PARAM.inp.nbands, eigen.data(), inc, eigenvalue_in, inc);
BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc);
ModuleBase::timer::end("DiagoCusolver", "cusolver");
}

Expand Down
9 changes: 7 additions & 2 deletions source/source_hsolver/diago_cusolver.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,11 @@ class DiagoCusolver

public:

DiagoCusolver();
/// @param nlocal_in global dimension of the NAO Hamiltonian
/// @param nbands_in number of lowest eigenpairs to compute
DiagoCusolver(const int nlocal_in, const int nbands_in);
~DiagoCusolver();

// Override the diag function for CUSOLVER diagonalization
void diag(
hamilt::MatrixBlock<T>& h_mat,
Expand All @@ -40,6 +42,9 @@ class DiagoCusolver
// Function to check if ELPA handle needs to be created or reused in MPI settings
bool ifElpaHandle(const bool& newIteration, const bool& ifNSCF) const;
#endif

const int nlocal;
const int nbands;
};

} // namespace hsolver
Expand Down
5 changes: 2 additions & 3 deletions source/source_hsolver/diago_cusolvermp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
#ifdef __CUSOLVERMP

#include "source_io/module_parameter/parameter.h"
#include "diago_cusolvermp.h"

#include "source_base/module_external/blas_connector.h"
Expand All @@ -18,7 +17,7 @@ void DiagoCusolverMP<T>::diag(hamilt::Hamilt<T>* phm_in, psi::Psi<T>& psi, Real*
hamilt::MatrixBlock<T> h_mat, s_mat;
phm_in->matrix(h_mat, s_mat);

std::vector<Real> eigen(PARAM.globalv.nlocal, 0.0);
std::vector<Real> eigen(this->nlocal, 0.0);
std::vector<T> eigenvectors(h_mat.row * h_mat.col);

MPI_Comm COMM_DIAG = MPI_COMM_WORLD; // use all processes
Expand All @@ -30,7 +29,7 @@ void DiagoCusolverMP<T>::diag(hamilt::Hamilt<T>* phm_in, psi::Psi<T>& psi, Real*
ModuleBase::timer::end("DiagoCusolverMP", "Diag_CusolverMP_gvd");
}
const int inc = 1;
BlasConnector::copy(PARAM.inp.nbands, eigen.data(), inc, eigenvalue_in, inc);
BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc);
const int size = psi.get_nbands() * psi.get_nbasis();
BlasConnector::copy(size, eigenvectors.data(), inc, psi.get_pointer(), inc);
}
Expand Down
8 changes: 7 additions & 1 deletion source/source_hsolver/diago_cusolvermp.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,17 @@ class DiagoCusolverMP
using Real = typename GetTypeReal<T>::type;

public:
DiagoCusolverMP()
/// @param nlocal_in global dimension of the NAO Hamiltonian
/// @param nbands_in number of lowest eigenpairs to compute
DiagoCusolverMP(const int nlocal_in, const int nbands_in) : nlocal(nlocal_in), nbands(nbands_in)
{
}
// the diag function for CUSOLVERMP diagonalization
void diag(hamilt::Hamilt<T>* phm_in, psi::Psi<T>& psi, Real* eigenvalue_in);

private:
const int nlocal;
const int nbands;
};
} // namespace hsolver
#endif // __CUSOLVERMP
Expand Down
29 changes: 12 additions & 17 deletions source/source_hsolver/diago_elpa.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
#include "source_base/global_function.h"
#include "source_base/module_external/blas_connector.h"

#include "source_io/module_parameter/parameter.h"
#include "module_genelpa/elpa_solver.h"
#include "source_base/module_external/blacs_connector.h"
#include "source_base/global_variable.h"
Expand Down Expand Up @@ -75,13 +74,13 @@ void DiagoElpa<std::complex<double>>::diag(
matcd h_mat, s_mat;
phm_in->matrix(h_mat, s_mat);

std::vector<double> eigen(PARAM.globalv.nlocal, 0.0);
std::vector<double> eigen(this->nlocal, 0.0);

bool isReal = false;
MPI_Comm COMM_DIAG = setmpicomm(); // set mpi_comm needed
ELPA_Solver es((const bool)isReal,
COMM_DIAG,
(const int)PARAM.inp.nbands,
(const int)this->nbands,
(const int)h_mat.row,
(const int)h_mat.col,
(const int*)h_mat.desc);
Expand All @@ -97,7 +96,7 @@ void DiagoElpa<std::complex<double>>::diag(
es.exit();

const int inc = 1;
BlasConnector::copy(PARAM.inp.nbands, eigen.data(), inc, eigenvalue_in, inc);
BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc);
#else
ModuleBase::WARNING_QUIT("DiagoElpa",
"DiagoElpa only can be used with macro __MPI");
Expand All @@ -113,15 +112,13 @@ void DiagoElpa<double>::diag(hamilt::Hamilt<double>* phm_in,
matd h_mat, s_mat;
phm_in->matrix(h_mat, s_mat);

std::vector<double> eigen(PARAM.globalv.nlocal, 0.0);
std::vector<double> eigen(this->nlocal, 0.0);

bool isReal = true;
MPI_Comm COMM_DIAG = setmpicomm(); // set mpi_comm needed
// ELPA_Solver es(isReal, COMM_DIAG, PARAM.inp.nbands, h_mat.row, h_mat.col,
// h_mat.desc);
ELPA_Solver es((const bool)isReal,
COMM_DIAG,
(const int)PARAM.inp.nbands,
(const int)this->nbands,
(const int)h_mat.row,
(const int)h_mat.col,
(const int*)h_mat.desc);
Expand All @@ -135,7 +132,7 @@ void DiagoElpa<double>::diag(hamilt::Hamilt<double>* phm_in,
es.exit();

const int inc = 1;
BlasConnector::copy(PARAM.inp.nbands, eigen.data(), inc, eigenvalue_in, inc);
BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc);
#else
ModuleBase::WARNING_QUIT("DiagoElpa",
"DiagoElpa only can be used with macro __MPI");
Expand All @@ -151,11 +148,11 @@ void DiagoElpa<std::complex<double>>::diag_pool(hamilt::MatrixBlock<std::complex
Real* eigenvalue_in,
MPI_Comm& comm)
{
std::vector<double> eigen(PARAM.globalv.nlocal, 0.0);
std::vector<double> eigen(this->nlocal, 0.0);
bool isReal = false;
ELPA_Solver es((const bool)isReal,
comm,
(const int)PARAM.inp.nbands,
(const int)this->nbands,
(const int)h_mat.row,
(const int)h_mat.col,
(const int*)h_mat.desc);
Expand All @@ -170,7 +167,7 @@ void DiagoElpa<std::complex<double>>::diag_pool(hamilt::MatrixBlock<std::complex
ModuleBase::timer::end("DiagoElpa", "elpa_solve");
es.exit();
const int inc = 1;
BlasConnector::copy(PARAM.inp.nbands, eigen.data(), inc, eigenvalue_in, inc);
BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc);
}

template <>
Expand All @@ -180,14 +177,12 @@ void DiagoElpa<double>::diag_pool(hamilt::MatrixBlock<double>& h_mat,
Real* eigenvalue_in,
MPI_Comm& comm)
{
std::vector<double> eigen(PARAM.globalv.nlocal, 0.0);
std::vector<double> eigen(this->nlocal, 0.0);

bool isReal = true;
// ELPA_Solver es(isReal, COMM_DIAG, PARAM.inp.nbands, h_mat.row, h_mat.col,
// h_mat.desc);
ELPA_Solver es((const bool)isReal,
comm,
(const int)PARAM.inp.nbands,
(const int)this->nbands,
(const int)h_mat.row,
(const int)h_mat.col,
(const int*)h_mat.desc);
Expand All @@ -203,7 +198,7 @@ void DiagoElpa<double>::diag_pool(hamilt::MatrixBlock<double>& h_mat,
const int inc = 1;
ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running,
"K-S equation was solved by genelpa2");
BlasConnector::copy(PARAM.inp.nbands, eigen.data(), inc, eigenvalue_in, inc);
BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc);
ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running,
"eigenvalues were copied to ekb");
}
Expand Down
7 changes: 7 additions & 0 deletions source/source_hsolver/diago_elpa.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,10 @@ class DiagoElpa
using Real = typename GetTypeReal<T>::type;

public:
/// @param nlocal_in global dimension of the NAO Hamiltonian
/// @param nbands_in number of lowest eigenpairs to compute
DiagoElpa(const int nlocal_in, const int nbands_in) : nlocal(nlocal_in), nbands(nbands_in) {};

void diag(hamilt::Hamilt<T>* phm_in, psi::Psi<T>& psi, Real* eigenvalue_in);
#ifdef __MPI
// diagnolization used in parallel-k case
Expand All @@ -30,6 +34,9 @@ class DiagoElpa
bool ifElpaHandle(const bool& newIteration, const bool& ifNSCF) const;
static int lastmpinum; // last using mpi;
#endif

const int nlocal;
const int nbands;
};

template <typename T>
Expand Down
9 changes: 4 additions & 5 deletions source/source_hsolver/diago_elpa_native.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
#include "source_base/module_external/blas_connector.h"
#include "source_base/module_external/blacs_connector.h"
#include "source_base/global_variable.h"
#include "source_io/module_parameter/parameter.h"
#include "source_base/timer.h"
#include "source_base/tool_quit.h"
#include "source_hsolver/module_genelpa/elpa_new.h"
Expand Down Expand Up @@ -59,7 +58,7 @@ void DiagoElpaNative<T>::diag_pool(hamilt::MatrixBlock<T>& h_mat,

ModuleBase::timer::start("DiagoElpaNative", "elpa_solve");

int nev = PARAM.inp.nbands;
int nev = this->nbands;
int narows = h_mat.row;
int nacols = h_mat.col;

Expand All @@ -70,7 +69,7 @@ void DiagoElpaNative<T>::diag_pool(hamilt::MatrixBlock<T>& h_mat,
int nprows, npcols, myprow, mypcol;

Cblacs_gridinfo(cblacs_ctxt, &nprows, &npcols, &myprow, &mypcol);
std::vector<Real> eigen(PARAM.globalv.nlocal, 0.0);
std::vector<Real> eigen(this->nlocal, 0.0);
std::vector<T> eigenvectors(narows * nacols);

if (elpa_init(20210430) != ELPA_OK)
Expand Down Expand Up @@ -107,7 +106,7 @@ void DiagoElpaNative<T>::diag_pool(hamilt::MatrixBlock<T>& h_mat,
#define ELPA_WITH_SYCL_GPU_VERSION 0
*/
#if ELPA_WITH_NVIDIA_GPU_VERSION
if (PARAM.inp.device == "gpu")
if (this->use_gpu)
{
elpa_set(handle, "nvidia-gpu", 1, &success);
elpa_set(handle, "real_kernel", ELPA_2STAGE_REAL_NVIDIA_GPU, &success);
Expand Down Expand Up @@ -138,7 +137,7 @@ void DiagoElpaNative<T>::diag_pool(hamilt::MatrixBlock<T>& h_mat,
}

const int inc = 1;
BlasConnector::copy(PARAM.inp.nbands, eigen.data(), inc, eigenvalue_in, inc);
BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc);
const int size = psi.get_nbands() * psi.get_nbasis();
BlasConnector::copy(size, eigenvectors.data(), inc, psi.get_pointer(), inc);
}
Expand Down
10 changes: 10 additions & 0 deletions source/source_hsolver/diago_elpa_native.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@ class DiagoElpaNative
using Real = typename GetTypeReal<T>::type;

public:
/// @param nlocal_in global dimension of the NAO Hamiltonian
/// @param nbands_in number of lowest eigenpairs to compute
/// @param use_gpu_in offload to the NVIDIA-GPU ELPA kernels when ELPA was built with GPU support
DiagoElpaNative(const int nlocal_in, const int nbands_in, const bool use_gpu_in)
: nlocal(nlocal_in), nbands(nbands_in), use_gpu(use_gpu_in) {};

void diag(hamilt::Hamilt<T>* phm_in, psi::Psi<T>& psi, Real* eigenvalue_in);
#ifdef __MPI
// diagnolization used in parallel-k case
Expand All @@ -27,6 +33,10 @@ class DiagoElpaNative

static int DecomposedState;

private:
const int nlocal;
const int nbands;
const bool use_gpu;
};

template <typename T>
Expand Down
Loading
Loading