#include "diago_david.h"

#include "source_base/timer.h" // ModuleBase::timer::tick

#include "source_base/tool_title.h" // ModuleBase::TITLE

#include "source_base/module_device/device.h"

#include "source_hsolver/kernels/hegvd_op.h"

#include "source_base/kernels/math_kernel_op.h"

using namespace hsolver;

template <typename T, typename Device>

DiagoDavid<T, Device>::DiagoDavid(const Real* precondition_in,

const int nband_in,

const int dim_in,

const int david_ndim_in,

const bool use_paw_in,

const diag_comm_info& diag_comm_in)

: nband(nband_in), dim(dim_in), nbase_x(david_ndim_in * nband_in), david_ndim(david_ndim_in), use_paw(use_paw_in), diag_comm(diag_comm_in)

{

this->device = base_device::get_device_type(this->ctx);

this->precondition = precondition_in;

this->one = &one_;

this->zero = &zero_;

this->neg_one = &neg_one_;

test_david = 2;

// 1: check which function is called and which step is executed

// 2: check the eigenvalues of the result of each iteration

// 3: check the eigenvalues and errors of the last result

// default: no check

// set auxiliary memory

// !!!

// allocate operaions MUST be paired with deallocate operations

// in the destructor

assert(this->david_ndim > 1);

assert(this->david_ndim * nband < dim * diag_comm.nproc);

// qianrui change it 2021-7-25.

// In strictly speaking, it shoule be PW_DIAG_NDIM*nband < npw sum of all pools. We roughly estimate it here.

// However, in most cases, total number of plane waves should be much larger than nband*PW_DIAG_NDIM

/// initialize variables

/// k_first = 0 means that nks is more like a dimension of "basis" to be contracted in "HPsi".In LR-TDDFT the formula writes :

/// $$\sum_{ jb\mathbf{k}'}A^I_{ia\mathbf{k}, jb\mathbf{k}' }X ^ I_{ jb\mathbf{k}'}$$

/// In the code :

/// - "H" means A

/// - "Psi" means X

/// - "band" means the superscript I : the number of excited states to be solved

/// - k : k-points, the same meaning as the ground state

/// - "basis" : number of occupied ks-orbitals(subscripts i,j) * number of unoccupied ks-orbitals(subscripts a,b), corresponding to "bands" of the ground state

// the lowest N eigenvalues

base_device::memory::resize_memory_op<Real, base_device::DEVICE_CPU>()(this->eigenvalue, nbase_x, "DAV::eig");

base_device::memory::set_memory_op<Real, base_device::DEVICE_CPU>()(this->eigenvalue, 0, nbase_x);

// basis(dim, nbase_x), leading dimension = dim

resmem_complex_op()(basis, nbase_x * dim, "DAV::basis");

setmem_complex_op()(basis, 0, nbase_x * dim);

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

// hpsi(nbase_x, dim); // the product of H and psi in the reduced basis set

resmem_complex_op()(this->hpsi, nbase_x * dim, "DAV::hpsi");

setmem_complex_op()(this->hpsi, 0, nbase_x * dim);

// spsi(nbase_x, dim); // the Product of S and psi in the reduced basis set

resmem_complex_op()(this->spsi, nbase_x * dim, "DAV::spsi");

setmem_complex_op()(this->spsi, 0, nbase_x * dim);

// hcc(nbase_x, nbase_x); // Hamiltonian on the reduced basis

resmem_complex_op()(this->hcc, nbase_x * nbase_x, "DAV::hcc");

setmem_complex_op()(this->hcc, 0, nbase_x * nbase_x);

// vcc(nbase_x, nbase_x); // Eigenvectors of hcc

resmem_complex_op()(this->vcc, nbase_x * nbase_x, "DAV::vcc");

setmem_complex_op()(this->vcc, 0, nbase_x * nbase_x);

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

// lagrange_matrix(nband, nband); // for orthogonalization

resmem_complex_op()(this->lagrange_matrix, nband * nband);

setmem_complex_op()(this->lagrange_matrix, 0, nband * nband);

#if defined(__CUDA) || defined(__ROCM)

// device precondition array

if (this->device == base_device::GpuDevice)

{

resmem_var_op()(this->d_precondition, dim);

syncmem_var_h2d_op()(this->d_precondition, this->precondition, dim);

}

#endif

}

template <typename T, typename Device>

DiagoDavid<T, Device>::~DiagoDavid()

{

delmem_complex_op()(this->basis);

delmem_complex_op()(this->hpsi);

delmem_complex_op()(this->spsi);

delmem_complex_op()(this->hcc);

delmem_complex_op()(this->vcc);

delmem_complex_op()(this->lagrange_matrix);

base_device::memory::delete_memory_op<Real, base_device::DEVICE_CPU>()(this->eigenvalue);

// If the device is a GPU device, free the d_precondition array.

#if defined(__CUDA) || defined(__ROCM)

if (this->device == base_device::GpuDevice)

{

delmem_var_op()(this->d_precondition);

}

#endif

}

template <typename T, typename Device>

int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,

const SPsiFunc& spsi_func,

const int dim,

const int nband,

const int ld_psi,

T *psi_in,

Real* eigenvalue_in,

const std::vector<double>& ethr_band,

const int david_maxiter)

{

if (test_david == 1)

{

ModuleBase::TITLE("DiagoDavid", "diag_once");

}

ModuleBase::timer::start("DiagoDavid", "diag_once");

// convflag[m] = true if the m th band is converged

std::vector<bool> convflag(nband, false);

// unconv[m] store the number of the m th unconverged band

std::vector<int> unconv(nband);

int nbase = 0; // the dimension of the reduced basis set

this->notconv = nband; // the number of unconverged eigenvalues

for (int m = 0; m < nband; m++) {

unconv[m] = m;

}

ModuleBase::timer::start("DiagoDavid", "first");

// orthogonalise the initial trial psi(0~nband-1)

// plan for SchmidtOrth

std::vector<int> pre_matrix_mm_m(nband, 0);

std::vector<int> pre_matrix_mv_m(nband, 1);

this->planSchmidtOrth(nband, pre_matrix_mm_m, pre_matrix_mv_m);

for (int m = 0; m < nband; m++)

{

{

// phm_in->sPsi(psi_in + m*ld_psi, &this->spsi[m * dim], dim, dim, 1);

spsi_func(psi_in + m*ld_psi,&this->spsi[m*dim],dim, 1);

}

}

// begin SchmidtOrth

for (int m = 0; m < nband; m++)

{

syncmem_complex_op()(basis + dim*m, psi_in + m*ld_psi, dim);

this->SchmidtOrth(dim,

nband,

m,

this->spsi,

&this->lagrange_matrix[m * nband],

pre_matrix_mm_m[m],

pre_matrix_mv_m[m]);

{

// phm_in->sPsi(basis + dim*m, &this->spsi[m * dim], dim, dim, 1);

spsi_func(basis + dim*m, &this->spsi[m * dim], dim, 1);

}

}

// end of SchmidtOrth and calculate H|psi>

// hpsi_info dav_hpsi_in(&basis, psi::Range(true, 0, 0, nband - 1), this->hpsi);

// phm_in->ops->hPsi(dav_hpsi_in);

// hpsi[:, 0:nband] = H basis[:, 0:nband]

// slice index in this piece of code is in C manner. i.e. 0:id stands for [0,id)

hpsi_func(basis, hpsi, dim, nband);

this->cal_elem(dim, nbase, nbase_x, this->notconv, this->hpsi, this->spsi, this->hcc);

this->diag_zhegvx(nbase, nband, this->hcc, nbase_x, this->eigenvalue, this->vcc);

for (int m = 0; m < nband; m++)

{

eigenvalue_in[m] = this->eigenvalue[m];

}

ModuleBase::timer::end("DiagoDavid", "first");

int dav_iter = 0;

do

{

dav_iter++;

this->cal_grad(hpsi_func,

spsi_func,

dim,

nbase,

nbase_x,

this->notconv,

this->hpsi,

this->spsi,

this->vcc,

unconv.data(),

this->eigenvalue);

this->cal_elem(dim, nbase, nbase_x, this->notconv, this->hpsi, this->spsi, this->hcc);

this->diag_zhegvx(nbase, nband, this->hcc, nbase_x, this->eigenvalue, this->vcc);

// check convergence and update eigenvalues

ModuleBase::timer::start("DiagoDavid", "check_update");

this->notconv = 0;

for (int m = 0; m < nband; m++)

{

convflag[m] = (std::abs(this->eigenvalue[m] - eigenvalue_in[m]) < ethr_band[m]);

if (!convflag[m])

{

unconv[this->notconv] = m;

this->notconv++;

}

eigenvalue_in[m] = this->eigenvalue[m];

}

ModuleBase::timer::end("DiagoDavid", "check_update");

if (!this->notconv || (nbase + this->notconv > nbase_x)

|| (dav_iter == david_maxiter))

{

ModuleBase::timer::start("DiagoDavid", "last");

// update eigenvectors of Hamiltonian

setmem_complex_op()(psi_in, 0, nband * ld_psi);

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

ModuleBase::gemm_op<T, Device>()('N',

'N',

dim, // m: row of A,C

nband, // n: col of B,C

nbase, // k: col of A, row of B

this->one,

basis, // A dim * nbase

dim,

this->vcc, // B nbase * nband

nbase_x,

this->zero,

psi_in, // C dim * nband

ld_psi);

if (!this->notconv || (dav_iter == david_maxiter))

{

// overall convergence or last iteration: exit the iteration

ModuleBase::timer::end("DiagoDavid", "last");

break;

}

else

{

// if the dimension of the reduced basis set is becoming too large,

// then replace the first N (=nband) basis vectors with the current

// estimate of the eigenvectors and set the basis dimension to N;

this->refresh(dim,

nband,

nbase,

nbase_x,

eigenvalue_in,

psi_in,

ld_psi,

this->hpsi,

this->spsi,

this->hcc,

this->vcc);

ModuleBase::timer::end("DiagoDavid", "last");

}

} // end of if

} while (true);

ModuleBase::timer::end("DiagoDavid", "diag_once");

return dav_iter;

}

template <typename T, typename Device>

void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,

const SPsiFunc& spsi_func,

const int& dim,

const int& nbase, // current dimension of the reduced basis

const int nbase_x, // maximum dimension of the reduced basis set

const int& notconv,

T* hpsi,

T* spsi,

const T* vcc,

const int* unconv,

const Real* eigenvalue)

{

if (test_david == 1) {

ModuleBase::TITLE("DiagoDavid", "cal_grad");

}

// if all eigenpairs have converged, return

if (notconv == 0) {

return;

}

ModuleBase::timer::start("DiagoDavid", "cal_grad");

// use template pointer for accelerate

// std::complex<double>* spsi;

// std::complex<double>* ppsi;

// expand the reduced basis set with the new basis vectors P|Real(psi)>...

// in which psi are the last eigenvectors

// we define |Real(psi)> as (H-ES)*|Psi>, E = <psi|H|psi>/<psi|S|psi>

// vc_ev_vector(notconv, nbase);

// eigenvectors of unconverged index extracted from vcc

T* vc_ev_vector = nullptr;

resmem_complex_op()(vc_ev_vector, notconv * nbase);

setmem_complex_op()(vc_ev_vector, 0, notconv * nbase);

//>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

// for (int m = 0; m < notconv; m++)

// {

// for (int i = 0; i < nbase; i++)

// {

// // vc_ev_vector(m, i) = vc(i, unconv[m]);

// vc_ev_vector[m * nbase + i] = vcc[i * nbase_x + unconv[m]];

// }

// }

//>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

// set vc_ev_vector = vcc (unconverged index, leading dimension = nbase_x)

// vc_ev_vector(m, i) = vc(i, unconv[m]);

// vc_ev_vector[m * nbase + i] = vcc[i * nbase_x + unconv[m]];

for (int m = 0; m < notconv; m++)

{

syncmem_complex_op()(vc_ev_vector + m * nbase,

vcc + unconv[m] * nbase_x,

nbase);

}

// basis[nbase] = hpsi * vc_ev_vector = hpsi*vcc

// basis' = vc_ev_vector' * hpsi'

// (dim, notconv) (dim, nbase) (nbase, notconv)

if (notconv == 1){

//Reuse gemv for vector case to avoid potential bug using gemm call with n=1

ModuleBase::gemv_op<T, Device>()('N',

dim, // m: row of A

nbase, // n: col of A

this->one, // alpha

hpsi, // A dim * nbase

dim, // LDA: if(N) max(1,m)

vc_ev_vector, // X nbase

1, // incx

this->zero, // beta

basis + dim * nbase, // Y dim

1 // incy

);

}else

{

ModuleBase::gemm_op<T, Device>()('N',

'N',

dim, // m: row of A,C

notconv, // n: col of B,C

nbase, // k: col of A, row of B

this->one, // alpha

hpsi, // A dim * nbase

dim, // LDA: if(N) max(1,m) if(T) max(1,k)

vc_ev_vector, // B nbase * notconv

nbase, // LDB: if(N) max(1,k) if(T) max(1,n)

this->zero, // belta

basis + dim * nbase, // C dim * notconv

dim // LDC: if(N) max(1, m)

);

}

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

// for (int m = 0; m < notconv; m++)

// {

// for (int i = 0; i < nbase; i++)

// {

// vc_ev_vector[m * nbase + i] *= -1 * this->eigenvalue[unconv[m]];

// }

// }

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

// e_temp_cpu = {-lambda}

// vc_ev_vector[nbase] = vc_ev_vector[nbase] * e_temp_cpu

// now vc_ev_vector[nbase] = - lambda * ev = -lambda * vcc

for (int m = 0; m < notconv; m++)

{

std::vector<Real> e_temp_cpu(nbase, (-1.0 * this->eigenvalue[unconv[m]]));

if (this->device == base_device::GpuDevice)

{

#if defined(__CUDA) || defined(__ROCM)

Real* e_temp_gpu = nullptr;

resmem_var_op()(e_temp_gpu, nbase);

syncmem_var_h2d_op()(e_temp_gpu, e_temp_cpu.data(), nbase);

ModuleBase::vector_mul_vector_op<T, Device>()(nbase,

vc_ev_vector + m * nbase,

vc_ev_vector + m * nbase,

e_temp_gpu);

delmem_var_op()(e_temp_gpu);

#endif

}

else

{

ModuleBase::vector_mul_vector_op<T, Device>()(nbase,

vc_ev_vector + m * nbase,

vc_ev_vector + m * nbase,

e_temp_cpu.data());

}

}

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

// basis[nbase] = basis[nbase] - spsi * vc_ev_vector

// = hpsi - spsi * lambda * vcc

// = (H - lambda * S) * psi * vcc

// = (H - lambda * S) * psi_new

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

if (notconv == 1){

//Use gemv for vector case to avoid potential bug using gemm call with n=1

ModuleBase::gemv_op<T, Device>()('N',

dim, // m: row of A

nbase, // n: col of A

this->one, // alpha

spsi, // A dim * nbase

dim, // LDA: if(N) max(1,m)

vc_ev_vector, // X nbase

1, // incx

this->one, // beta

basis + dim * nbase, // Y dim

1 //incy

);

} else

{

ModuleBase::gemm_op<T, Device>()('N',

'N',

dim, // m: row of A,C

notconv, // n: col of B,C

nbase, // k: col of A, row of B

this->one, // alpha

spsi, // A

dim, // LDA: if(N) max(1,m) if(T) max(1,k)

vc_ev_vector, // B

nbase, // LDB: if(N) max(1,k) if(T) max(1,n)

this->one, // beta

basis + dim * nbase, // C dim * notconv

dim // LDC: if(N) max(1, m)

);

}

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

// Preconditioning

// basis[nbase] = T * basis[nbase] = T * (H - lambda * S) * psi

// where T, the preconditioner, is an approximate inverse of H

// T is a diagonal stored in array `precondition`

// to do preconditioning, divide each column of basis by the corresponding element of precondition

for (int m = 0; m < notconv; m++)

{

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

if (this->device == base_device::GpuDevice)

{

#if defined(__CUDA) || defined(__ROCM)

ModuleBase::vector_div_vector_op<T, Device>()(dim,

basis + dim * (nbase + m),

basis + dim * (nbase + m),

this->d_precondition);

#endif

}

else

{

ModuleBase::vector_div_vector_op<T, Device>()(dim,

basis + dim * (nbase + m),

basis + dim * (nbase + m),

this->precondition);

}

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

// for (int ig = 0; ig < dim; ig++)

// {

// ppsi[ig] /= this->precondition[ig];

// }

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

}

// there is a nbase to nbase + notconv band orthogonalise

// plan for SchmidtOrth

T* lagrange = nullptr;

resmem_complex_op()(lagrange, notconv * (nbase + notconv));

setmem_complex_op()(lagrange, 0, notconv * (nbase + notconv));

std::vector<int> pre_matrix_mm_m(notconv, 0);

std::vector<int> pre_matrix_mv_m(notconv, 1);

this->planSchmidtOrth(notconv, pre_matrix_mm_m, pre_matrix_mv_m);

for (int m = 0; m < notconv; m++)

{

{

// phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1);

spsi_func(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1);

}

}

// first nbase bands psi* dot notconv bands spsi to prepare lagrange_matrix

// calculate the square matrix for future lagranges

if (notconv == 1){

//Use gemv for vector case to avoid potential bug using gemm call with n=1

ModuleBase::gemv_op<T, Device>()('C',

dim, // m: row of A

nbase, // n: col of A

this->one, // alpha

basis, // A dim * nbase

dim, // LDA: if(N) max(1,m)

&spsi[nbase * dim], // X dim

1, // incx

this->zero, // beta

lagrange, // Y nbase

1

);

} else

{

ModuleBase::gemm_op<T, Device>()('C',

'N',

nbase, // m: row of A,C

notconv, // n: col of B,C

dim, // k: col of A, row of B

this->one, // alpha

basis, // A

dim, // LDA: if(N) max(1,m) if(T) max(1,k)

&spsi[nbase * dim], // B

dim, // LDB: if(N) max(1,k) if(T) max(1,n)

this->zero, // belta

lagrange, // C

nbase + notconv // LDC: if(N) max(1, m)

);

}

for (int m = 0; m < notconv; m++)

{

this->SchmidtOrth(dim,

nbase + notconv,

nbase + m,

spsi,

&lagrange[m * (nbase + notconv)],

pre_matrix_mm_m[m],

pre_matrix_mv_m[m]);

{

// phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1);

spsi_func(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1);

}

}

// calculate H|psi> for not convergence bands

// hpsi_info dav_hpsi_in(&basis,

// psi::Range(true, 0, nbase, nbase + notconv - 1),

// &hpsi[nbase * dim]); // &hp(nbase, 0)

// phm_in->ops->hPsi(dav_hpsi_in);

// hpsi[:, nbase:nbase+notcnv] = H basis[:, nbase:nbase+notcnv]

hpsi_func(basis + nbase * dim, hpsi + nbase * dim, dim, notconv);

delmem_complex_op()(lagrange);

delmem_complex_op()(vc_ev_vector);

ModuleBase::timer::end("DiagoDavid", "cal_grad");

return;

}

template <typename T, typename Device>

void DiagoDavid<T, Device>::cal_elem(const int& dim,

int& nbase, // current dimension of the reduced basis

const int nbase_x, // maximum dimension of the reduced basis set

const int& notconv, // number of newly added basis vectors

const T* hpsi,

const T* spsi,

T* hcc)

{

if (test_david == 1) {

ModuleBase::TITLE("DiagoDavid", "cal_elem");

}

if (notconv == 0) {

return;

}

ModuleBase::timer::start("DiagoDavid", "cal_elem");

// hcc[nbase](notconv, nbase + notconv)= basis[nbase]' * hpsi

ModuleBase::gemm_op<T, Device>()('C',

'N',

notconv,

nbase + notconv,

dim,

this->one,

basis + dim * nbase, // basis(:,nbase:) dim * notconv

dim,

hpsi, // dim * (nbase + notconv)

dim,

this->zero,

hcc + nbase, // notconv * (nbase + notconv)

nbase_x);

#ifdef __MPI

if (diag_comm.nproc > 1)

{

ModuleBase::matrixTranspose_op<T, Device>()(nbase_x, nbase_x, hcc, hcc);

auto* swap = new T[notconv * nbase_x];

syncmem_complex_op()(swap, hcc + nbase * nbase_x, notconv * nbase_x);

if (std::is_same<T, double>::value)

{

Parallel_Reduce::reduce_pool(hcc + nbase * nbase_x, notconv * nbase_x);

}

else

{

if (base_device::get_current_precision(swap) == "single") {

MPI_Reduce(swap, hcc + nbase * nbase_x, notconv * nbase_x, MPI_COMPLEX, MPI_SUM, 0, diag_comm.comm);

}

else {

MPI_Reduce(swap, hcc + nbase * nbase_x, notconv * nbase_x, MPI_DOUBLE_COMPLEX, MPI_SUM, 0, diag_comm.comm);

}

}

delete[] swap;

// Parallel_Reduce::reduce_complex_double_pool( hcc + nbase * nbase_x, notconv * nbase_x );

ModuleBase::matrixTranspose_op<T, Device>()(nbase_x, nbase_x, hcc, hcc);

}

#endif

nbase += notconv;

ModuleBase::timer::end("DiagoDavid", "cal_elem");

return;

}

template <typename T, typename Device>

void DiagoDavid<T, Device>::diag_zhegvx(const int& nbase,

const int& nband,

const T* hcc,

const int& nbase_x,

Real* eigenvalue, // in CPU

T* vcc)

{

ModuleBase::timer::start("DiagoDavid", "diag_zhegvx");

if (diag_comm.rank == 0)

{

assert(nbase_x >= std::max(1, nbase));

if (this->device == base_device::GpuDevice)

{

#if defined(__CUDA) || defined(__ROCM)

Real* eigenvalue_gpu = nullptr;

resmem_var_op()(eigenvalue_gpu, nbase_x);

syncmem_var_h2d_op()(eigenvalue_gpu, this->eigenvalue, nbase_x);

// heevx_op<T, Device>()(this->ctx, nbase, nbase_x, hcc, nband, eigenvalue_gpu, vcc);

ct::kernels::lapack_heevx<T, ct_Device>()(nbase, nbase_x, hcc, nband, eigenvalue_gpu, vcc);

syncmem_var_d2h_op()(this->eigenvalue, eigenvalue_gpu, nbase_x);

delmem_var_op()(eigenvalue_gpu);

#endif

}

else

{

//heevx_op<T, Device>()(this->ctx, nbase, nbase_x, hcc, nband, this->eigenvalue, vcc);

ct::kernels::lapack_heevx<T, ct_Device>()(nbase, nbase_x, hcc, nband, this->eigenvalue, vcc);

}

}

#ifdef __MPI

if (diag_comm.nproc > 1)

{

// vcc: nbase * nband

for (int i = 0; i < nband; i++)

{

MPI_Bcast(&vcc[i * nbase_x], nbase, MPI_DOUBLE_COMPLEX, 0, diag_comm.comm);

}

MPI_Bcast(this->eigenvalue, nband, MPI_DOUBLE, 0, diag_comm.comm);

}

#endif

ModuleBase::timer::end("DiagoDavid", "diag_zhegvx");

return;

}

template <typename T, typename Device>

void DiagoDavid<T, Device>::refresh(const int& dim,

const int& nband,

int& nbase,

const int nbase_x, // maximum dimension of the reduced basis set

const Real* eigenvalue_in,

const T *psi_in,

const int ld_psi,

T* hpsi,

T* spsi,

T* hcc,

T* vcc)

{

if (test_david == 1) {

ModuleBase::TITLE("DiagoDavid", "refresh");

}

ModuleBase::timer::start("DiagoDavid", "refresh");

// update hp,sp

setmem_complex_op()(basis , 0, nbase_x * dim);

// basis(dim, nband) = hpsi(dim, nbase) * vcc(nbase, nband)

ModuleBase::gemm_op<T, Device>()('N',

'N',

dim, // m: row of A,C

nband, // n: col of B,C

nbase, // k: col of A, row of B

this->one,

hpsi, // A dim * nbase

dim,

vcc, // B nbase * nband

nbase_x,

zero,

basis, // C dim * nband

dim);

//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

// basis[nband] = spsi * vcc

ModuleBase::gemm_op<T, Device>()('N',

'N',

dim, // m: row of A,C

nband, // n: col of B,C

nbase, // k: col of A, row of B

this->one,

spsi, // A dim * nbase

dim,

vcc, // B nbase * nband

nbase_x,

this->zero,

basis + dim * nband, // C dim * nband

dim);

// hpsi = basis, spsi = basis[nband]

syncmem_complex_op()(hpsi, basis, dim * nband);

syncmem_complex_op()(spsi, basis + dim*nband, dim * nband);

/*for (int m = 0; m < nband; m++) {

// update basis

setmem_complex_op()(basis , 0, nbase_x * dim);

for (int m = 0; m < nband; m++)

{

syncmem_complex_op()(basis + dim*m,psi_in + m*ld_psi, dim);

/*for (int ig = 0; ig < npw; ig++)

}

// update the reduced Hamiltonian

// basis set size reset to nband

nbase = nband;

setmem_complex_op()(hcc, 0, nbase_x * nbase_x);

if (this->device == base_device::GpuDevice)

{

#if defined(__CUDA) || defined(__ROCM)

T* hcc_cpu = nullptr;

T* vcc_cpu = nullptr;

base_device::memory::resize_memory_op<T, base_device::DEVICE_CPU>()(hcc_cpu,

nbase_x * nbase_x,

"DAV::hcc");

base_device::memory::resize_memory_op<T, base_device::DEVICE_CPU>()(vcc_cpu,

nbase_x * nbase_x,

"DAV::vcc");

syncmem_d2h_op()(hcc_cpu, hcc, nbase_x * nbase_x);

syncmem_d2h_op()(vcc_cpu, vcc, nbase_x * nbase_x);

for (int i = 0; i < nbase; i++)

{

hcc_cpu[i * nbase_x + i] = eigenvalue_in[i];

vcc_cpu[i * nbase_x + i] = this->one[0];

}

syncmem_h2d_op()(hcc, hcc_cpu, nbase_x * nbase_x);

syncmem_h2d_op()(vcc, vcc_cpu, nbase_x * nbase_x);

base_device::memory::delete_memory_op<T, base_device::DEVICE_CPU>()(hcc_cpu);

base_device::memory::delete_memory_op<T, base_device::DEVICE_CPU>()(vcc_cpu);

#endif

}

else

{

for (int i = 0; i < nbase; i++)

{

hcc[i * nbase_x + i] = eigenvalue_in[i];

// sc(i, i) = this->one;

// vc(i, i) = this->one;

vcc[i * nbase_x + i] = this->one[0];

}

}

ModuleBase::timer::end("DiagoDavid", "refresh");

return;

}

template <typename T, typename Device>

void DiagoDavid<T, Device>::SchmidtOrth(const int& dim,

const int nband,

const int m,

const T* spsi,

T* lagrange_m,

const int mm_size,

const int mv_size)

{

// if(test_david == 1) ModuleBase::TITLE("DiagoDavid","SchmidtOrth");

ModuleBase::timer::start("DiagoDavid", "SchmidtOrth");

// orthogonalize starting eigenfunction to those already calculated

// psi_m orthogonalize to psi(0) ~ psi(m-1)

// Attention, the orthogonalize here read as

// psi(m) -> psi(m) - \sum_{i < m} \langle psi(i)|S|psi(m) \rangle psi(i)

// so the orthogonalize is performed about S.

// assert(basis.get_nbands() >= nband);

assert(m >= 0);

assert(m < nband);

// psi_m = basis[m]

T* psi_m = basis + dim*m;

// std::complex<double> *lagrange = new std::complex<double>[m + 1];

// ModuleBase::GlobalFunc::ZEROS(lagrange, m + 1);

// calculate the square matrix for future lagranges

if (mm_size != 0)

{

// lagrange_m[m - mv_size + 1 - mm_size]

// = basis[m - mv_size + 1 - mm_size]' * spsi[m]

ModuleBase::gemm_op<T, Device>()('C',

'N',

mm_size, // m: row of A,C

mm_size, // n: col of B,C

dim, // k: col of A, row of B

this->one, // alpha

basis + dim * (m - mv_size + 1 - mm_size), // A

dim, // LDA: if(N) max(1,m) if(T) max(1,k)

&spsi[m * dim], // B

dim, // LDB: if(N) max(1,k) if(T) max(1,n)

this->zero, // belta

&lagrange_m[m - mv_size + 1 - mm_size], // C

nband // LDC: if(N) max(1, m)

);

}

// calculate other lagranges for this band

// lagrange_m[m - mv_size + 1]

// = basis[m - mv_size + 1]' * spsi[m]

ModuleBase::gemv_op<T, Device>()('C',

dim,

mv_size,

this->one,

basis + dim * (m - mv_size + 1),

dim,

&spsi[m * dim],

1,

this->zero,

&lagrange_m[m - mv_size + 1],

1);

Parallel_Reduce::reduce_pool(lagrange_m, m + 1);

T var = *this->zero;

syncmem_d2h_op()(&var, lagrange_m + m, 1);

double psi_norm = get_real(var);

assert(psi_norm > 0.0);

// / psi_m = psi_m - \sum_{i < m} \langle psi(i)|S|psi(m) \rangle psi(i)

// psi_m = psi_m - basis * lagrange_m

ModuleBase::gemv_op<T, Device>()('N',

dim,

m,

this->neg_one,

basis,

dim,

lagrange_m,

1,

this->one,

psi_m,

1);

// psi_norm = psi_norm - lagrange_m · lagrange_m

psi_norm -= ModuleBase::dot_real_op<T, Device>()(m, lagrange_m, lagrange_m, false);

// for (int j = 0; j < m; j++)

// {

// const std::complex<double> alpha = std::complex<double>(-1, 0) * lagrange_m[j];

// zaxpy_(&npw, &alpha, &psi(j,0), &inc, psi_m, &inc);

// /*for (int ig = 0; ig < npw; ig++)

// {

// psi_m[ig] -= lagrange[j] * psi(j, ig);

// }*/

// psi_norm -= (conj(lagrange_m[j]) * lagrange_m[j]).real();

// }

assert(psi_norm > 0.0);

psi_norm = sqrt(psi_norm);

if (psi_norm < 1.0e-12)

{

std::cout << "DiagoDavid::SchmidtOrth:aborted for psi_norm <1.0e-12" << std::endl;

std::cout << "nband = " << nband << std::endl;

std::cout << "m = " << m << std::endl;

exit(0);

}

else

{

// psi_m = psi_m / psi_norm

ModuleBase::vector_mul_real_op<T, Device>()(dim, psi_m, psi_m, Real(1.0 / psi_norm));

// for (int i = 0; i < npw; i++)

// {

// psi_m[i] /= psi_norm;

// }

}

// delete[] lagrange;

ModuleBase::timer::end("DiagoDavid", "SchmidtOrth");

return;

}

template <typename T, typename Device>

void DiagoDavid<T, Device>::planSchmidtOrth(const int nband, std::vector<int>& pre_matrix_mm_m, std::vector<int>& pre_matrix_mv_m)

{

if (nband <= 0) {

return;

}

std::fill(pre_matrix_mm_m.begin(), pre_matrix_mm_m.end(), 0);

std::fill(pre_matrix_mv_m.begin(), pre_matrix_mv_m.end(), 0);

int last_matrix_size = nband;

int matrix_size = int(nband / 2);

int divide_times = 0;

std::vector<int> divide_points(nband);

int res_nband = nband - matrix_size;

while (matrix_size > 1)

{

int index = nband - matrix_size;

if (divide_times == 0)

{

divide_points[0] = index;

pre_matrix_mm_m[index] = matrix_size;

if (res_nband == matrix_size) {

pre_matrix_mv_m[index] = 1;

} else {

pre_matrix_mv_m[index] = 2;

}

divide_times = 1;

}

else

{

for (int i = divide_times - 1; i >= 0; i--)

{

divide_points[i * 2] = divide_points[i] - matrix_size;

divide_points[i * 2 + 1] = divide_points[i * 2] + last_matrix_size;

pre_matrix_mm_m[divide_points[i * 2]] = matrix_size;

pre_matrix_mm_m[divide_points[i * 2 + 1]] = matrix_size;

if (res_nband == matrix_size)

{

pre_matrix_mv_m[divide_points[i * 2]] = 1;

pre_matrix_mv_m[divide_points[i * 2 + 1]] = 1;

}

else

{

pre_matrix_mv_m[divide_points[i * 2]] = 2;

pre_matrix_mv_m[divide_points[i * 2 + 1]] = 2;

}