Skip to content

Refactor: Preconditioner function library #5482

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

Closed
wants to merge 6 commits into from
Closed

Refactor: Preconditioner function library #5482

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a8bb37f
extract precondition function
maki49 Nov 14, 2024
4f4f356
apply to dav
maki49 Nov 14, 2024
17cac27
fix UTs
maki49 Nov 14, 2024
0687643
separate sorce and destination pointer
maki49 Nov 15, 2024
3ac0de2
fix sdft test compile bug on icpx
maki49 Nov 15, 2024
a0571fb
fix error passing host pointer to gpu kernel
maki49 Nov 15, 2024
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 python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,8 +130,10 @@ class PyDiagoDavSubspace
std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out);
};

hsolver::PreOP<std::complex<double>, base_device::DEVICE_CPU, hsolver::fvec::DivTransMinusEigKernel<std::complex<double>, base_device::DEVICE_CPU>>
Copy link
Collaborator

@mohanchen mohanchen Nov 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The name "DivTransMinusEigKernel" is a disaster.

Copy link
Collaborator Author

@maki49 maki49 Nov 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry but do you have any better idea?

Copy link
Collaborator

@kirk0830 kirk0830 Nov 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The function should be named with following considerations: precisely describe what it does. In present case, you mix the mathematical operation with some rare numerical trick description like "kernel". So there will be quite few people can know what this function does from its name. You should either name the function purely with names mathematical operations, or use human natural language describe the action that what quantity is calculated, additionally from what quantity.

Copy link
Collaborator

@kirk0830 kirk0830 Nov 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot see any significance of the necessarity of instantiating the hsolver::PreOP. The name of PreOP is also confusing quite much. From the rest of your PR, I can understand it is about preconditioning, but unfortunately you did not name them with much care and enough considerations.
You should spend at least 60% of your time on coding to carefully consider the name of variable, function, class, namespace, and the design of interface.

Copy link
Collaborator

@kirk0830 kirk0830 Nov 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

presently the lines 133, 134 are quite confusing and it is quite difficult for others to know what the operation is. It wastes other's time. You should always consider a clearer way to organize your code like:

// instantiate hsolver::PreOP
hsolver::PreOP<std::complex<double>, 
               base_device::DEVICE_CPU, 
               hsolver::fvec::DivTransMinusEigKernel<
                   std::complex<double>, 
                   base_device::DEVICE_CPU>
               > pre_op(precond_vec, 
                        hsolver::fvec::div_trans_prevec_minus_eigen<std::complex<double>>, 
                        hsolver::fval::qe_pw<double>);

however, do you think naming the instance of PreOP as pre_op is a good idea?

pre_op(precond_vec, hsolver::fvec::div_trans_prevec_minus_eigen<std::complex<double>>, hsolver::fval::qe_pw<double>);
obj = std::make_unique<hsolver::Diago_DavSubspace<std::complex<double>, base_device::DEVICE_CPU>>(
Copy link
Collaborator

@kirk0830 kirk0830 Nov 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it is the first time that I find there is an instance even named obj. Informativeless.

precond_vec,
pre_op.get(),
nband,
nbasis,
dav_ndim,
Expand Down
3 changes: 2 additions & 1 deletion python/pyabacus/src/hsolver/py_diago_david.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,8 +137,9 @@ class PyDiagoDavid
syncmem_op()(this->ctx, this->ctx, spsi_out, psi_in, static_cast<size_t>(nbands * nrow));
};

hsolver::PreOP<std::complex<double>> pre_op(precond_vec);
Copy link
Collaborator

@mohanchen mohanchen Nov 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

其它地方命名都是下划线命名法，PreOP这个名字起得能不能有点默契！

Copy link
Collaborator Author

@maki49 maki49 Nov 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://github.com/MCresearch/abacus-user-guide/blob/master/develop-C++.md
根据这个文档，变量名和命名空间用下划线命名法，类型名用的是每个单词首字母大写的命名法。
image

obj = std::make_unique<hsolver::DiagoDavid<std::complex<double>, base_device::DEVICE_CPU>>(
precond_vec.data(),
pre_op.get(),
nband,
nbasis,
dav_ndim,
Expand Down
51 changes: 4 additions & 47 deletions source/module_hsolver/diago_dav_subspace.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,16 +13,16 @@
using namespace hsolver;

template <typename T, typename Device>
Diago_DavSubspace<T, Device>::Diago_DavSubspace(const std::vector<Real>& precondition_in,
Diago_DavSubspace<T, Device>::Diago_DavSubspace(PreFunc&& precondition_in,
Copy link
Collaborator

@kirk0830 kirk0830 Nov 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could you explain the necessarity of using rvalue instead of lvalue?

const int& nband_in,
const int& nbasis_in,
const int& david_ndim_in,
const double& diag_thr_in,
const int& diag_nmax_in,
const bool& need_subspace_in,
const diag_comm_info& diag_comm_in)
: precondition(precondition_in), n_band(nband_in), dim(nbasis_in), nbase_x(nband_in * david_ndim_in),
diag_thr(diag_thr_in), iter_nmax(diag_nmax_in), is_subspace(need_subspace_in), diag_comm(diag_comm_in)
: precondition(precondition_in), n_band(nband_in), dim(nbasis_in), nbase_x(nband_in* david_ndim_in),
diag_thr(diag_thr_in), iter_nmax(diag_nmax_in), is_subspace(need_subspace_in), diag_comm(diag_comm_in)
{
this->device = base_device::get_device_type<Device>(this->ctx);

Expand Down Expand Up @@ -55,14 +55,6 @@ Diago_DavSubspace<T, Device>::Diago_DavSubspace(const std::vector<Real>& precond
resmem_complex_op()(this->ctx, this->vcc, this->nbase_x * this->nbase_x, "DAV::vcc");
setmem_complex_op()(this->ctx, this->vcc, 0, this->nbase_x * this->nbase_x);
//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

#if defined(__CUDA) || defined(__ROCM)
if (this->device == base_device::GpuDevice)
{
resmem_real_op()(this->ctx, this->d_precondition, nbasis_in);
// syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->d_precondition, this->precondition.data(), nbasis_in);
}
#endif
}

template <typename T, typename Device>
Expand All @@ -74,13 +66,6 @@ Diago_DavSubspace<T, Device>::~Diago_DavSubspace()
delmem_complex_op()(this->ctx, this->hcc);
delmem_complex_op()(this->ctx, this->scc);
delmem_complex_op()(this->ctx, this->vcc);

#if defined(__CUDA) || defined(__ROCM)
if (this->device == base_device::GpuDevice)
{
delmem_real_op()(this->ctx, this->d_precondition);
}
#endif
}

template <typename T, typename Device>
Expand Down Expand Up @@ -334,35 +319,7 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
this->dim);

// "precondition!!!"
std::vector<Real> pre(this->dim, 0.0);
for (int m = 0; m < notconv; m++)
{
for (size_t i = 0; i < this->dim; i++)
{
// pre[i] = std::abs(this->precondition[i] - (*eigenvalue_iter)[m]);
double x = std::abs(this->precondition[i] - (*eigenvalue_iter)[m]);
pre[i] = 0.5 * (1.0 + x + sqrt(1 + (x - 1.0) * (x - 1.0)));
}
#if defined(__CUDA) || defined(__ROCM)
if (this->device == base_device::GpuDevice)
{
syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->d_precondition, pre.data(), this->dim);
vector_div_vector_op<T, Device>()(this->ctx,
this->dim,
psi_iter + (nbase + m) * this->dim,
psi_iter + (nbase + m) * this->dim,
this->d_precondition);
}
else
#endif
{
vector_div_vector_op<T, Device>()(this->ctx,
this->dim,
psi_iter + (nbase + m) * this->dim,
psi_iter + (nbase + m) * this->dim,
pre.data());
}
}
this->precondition(psi_iter + nbase * this->dim, eigenvalue_iter->data(), this->dim, notconv);

// "normalize!!!" in order to improve numerical stability of subspace diagonalization
std::vector<Real> psi_norm(notconv, 0.0);
Expand Down
12 changes: 7 additions & 5 deletions source/module_hsolver/diago_dav_subspace.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

#include <vector>
#include <functional>
#include "module_hsolver/precondition_funcs.h"

namespace hsolver
{
Expand All @@ -23,8 +24,9 @@ class Diago_DavSubspace
// otherwise return the real type of T(complex<float>, complex<double>)
using Real = typename GetTypeReal<T>::type;

public:
Diago_DavSubspace(const std::vector<Real>& precondition_in,
using PreFunc = fvec::DivTransMinusEig<T>;
public:
Diago_DavSubspace(PreFunc&& precondition_in, /// pass in a function, lambda or PreOP object
const int& nband_in,
const int& nbasis_in,
const int& david_ndim_in,
Expand Down Expand Up @@ -67,9 +69,9 @@ class Diago_DavSubspace
/// the maximum dimension of the reduced basis set
const int nbase_x = 0;

/// precondition for diag
const std::vector<Real>& precondition;
Real* d_precondition = nullptr;
/// The precondition operation, can be a function, lambda or PreOP object
const PreFunc precondition;
// note that lambdas can only passed by value

/// record for how many bands not have convergence eigenvalues
int notconv = 0;
Expand Down
54 changes: 6 additions & 48 deletions source/module_hsolver/diago_david.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,16 +32,16 @@ using namespace hsolver;
* @note Auxiliary memory is allocated in the constructor and deallocated in the destructor.
*/
template <typename T, typename Device>
DiagoDavid<T, Device>::DiagoDavid(const Real* precondition_in,
DiagoDavid<T, Device>::DiagoDavid(PreFunc&& 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)
: 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),
precondition(precondition_in)
{
this->device = base_device::get_device_type<Device>(this->ctx);
this->precondition = precondition_in;

this->one = &one_;
this->zero = &zero_;
Expand Down Expand Up @@ -110,15 +110,6 @@ DiagoDavid<T, Device>::DiagoDavid(const Real* precondition_in,
// lagrange_matrix(nband, nband); // for orthogonalization
resmem_complex_op()(this->ctx, this->lagrange_matrix, nband * nband);
setmem_complex_op()(this->ctx, this->lagrange_matrix, 0, nband * nband);

#if defined(__CUDA) || defined(__ROCM)
// device precondition array
if (this->device == base_device::GpuDevice)
{
resmem_var_op()(this->ctx, this->d_precondition, dim);
syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->d_precondition, this->precondition, dim);
}
#endif
}

/**
Expand All @@ -139,13 +130,6 @@ DiagoDavid<T, Device>::~DiagoDavid()
delmem_complex_op()(this->ctx, this->vcc);
delmem_complex_op()(this->ctx, this->lagrange_matrix);
base_device::memory::delete_memory_op<Real, base_device::DEVICE_CPU>()(this->cpu_ctx, 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->ctx, this->d_precondition);
}
#endif
}

template <typename T, typename Device>
Expand Down Expand Up @@ -499,40 +483,14 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
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)
vector_div_vector_op<T, Device>()(this->ctx,
dim,
basis + dim*(nbase + m),
basis + dim*(nbase + m),
this->d_precondition);
#endif
}
else
{
vector_div_vector_op<T, Device>()(this->ctx,
dim,
basis + dim*(nbase + m),
basis + dim*(nbase + m),
this->precondition);
}
//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
// for (int ig = 0; ig < dim; ig++)
// {
// ppsi[ig] /= this->precondition[ig];
// }
//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
}
//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
this->precondition(basis + dim * nbase, dim, notconv);
//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

// there is a nbase to nbase + notconv band orthogonalise
// plan for SchmidtOrth
Expand Down
12 changes: 6 additions & 6 deletions source/module_hsolver/diago_david.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
#include "module_base/module_device/memory_op.h"// base_device::memory

#include "module_hsolver/diag_comm_info.h"

#include "module_hsolver/precondition_funcs.h"
#include <vector>
#include <functional>

Expand All @@ -21,10 +21,11 @@ class DiagoDavid
// return T if T is real type(float, double),
// otherwise return the real type of T(complex<float>, complex<double>)
using Real = typename GetTypeReal<T>::type;

public:

DiagoDavid(const Real* precondition_in,
using PreFunc = fvec::Div<T>;
Copy link
Collaborator

@Cstandardlib Cstandardlib Nov 15, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to give PrecondFunc the interface like HPsi or SPsi?

// in precondition lib
using Div = std::function<void(T*, const size_t&, const size_t&)>;

Maybe we can give it the form

using HPsiFunc = std::function<void(T*, T*, const int, const int)>;
using SPsiFunc = std::function<void(T*, T*, const int, const int)>;
using PrecndFunc = std::function<void(T*, T*, const int, const int)>;
// maybe it's better to change int to size_t, I'll do it on current HPsiFunc and SPsiFunc

and pass the preconditioners to this universal function wrapper.

Sometimes the source and destination of preconditioning may be different. For example from a temporary block residual r to a block of total search space X[:, nband:nband+notconv]. The results will be directly stored to a new array in such case.

Copy link
Collaborator Author

@maki49 maki49 Nov 15, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, see the latest change.

And seems on HPsiFunc and SPsiFunc can "const" also be added like this:

using HPsiFunc = std::function<void(T*const, const T*const, const int, const int)>;

public:

DiagoDavid(PreFunc&& precondition_in,
const int nband_in,
const int dim_in,
const int david_ndim_in,
Expand Down Expand Up @@ -102,8 +103,7 @@ class DiagoDavid
int notconv = 0;

/// precondition for diag, diagonal approximation of matrix A(i.e. Hamilt)
const Real* precondition = nullptr;
Real* d_precondition = nullptr;
const PreFunc precondition;

/// eigenvalue results
Real* eigenvalue = nullptr;
Expand Down
8 changes: 6 additions & 2 deletions source/module_hsolver/hsolver_pw.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,7 +503,9 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
};
bool scf = this->calculation_type == "nscf" ? false : true;

Diago_DavSubspace<T, Device> dav_subspace(pre_condition,
// const auto pre_op = make_pre_op(pre_condition, fvec::div_trans_prevec_minus_eigen<T, Device>, fval::qe_pw<Real>);
const PreOP<T, Device, fvec::DivTransMinusEigKernel<T, Device>> pre_op(pre_condition, fvec::div_trans_prevec_minus_eigen<T, Device>, fval::qe_pw<Real>);
Diago_DavSubspace<T, Device> dav_subspace(pre_op.get(),
psi.get_nbands(),
psi.get_k_first() ? psi.get_current_nbas()
: psi.get_nk() * psi.get_nbasis(),
Expand Down Expand Up @@ -572,7 +574,9 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
ModuleBase::timer::tick("David", "spsi_func");
};

DiagoDavid<T, Device> david(pre_condition.data(),
// const auto pre_op = make_pre_op(pre_condition, fvec::div_prevec<T, Device>);
const PreOP<T, Device> pre_op(pre_condition);
DiagoDavid<T, Device> david(pre_op.get(),
nband,
dim,
PARAM.inp.pw_diag_ndim,
Comment on lines +579 to 582
Copy link
Collaborator

@kirk0830 kirk0830 Nov 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

besides the PreOP, now I realize the DiagoDavid class has the similar problem. It is confusing to instantiate the class DiagoDavid then name as david. A factory function will be better and more consistent, but will not be work of this PR:

std::unique_ptr<Diago> dav = make_diago("dav", /* other args */); // or use object directly without pointer

Expand Down
Loading
Loading