Skip to content

Feature: Output real space wavefunction and partial charge density when device=gpu #6391

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

Merged
mohanchen merged 2 commits into from
Jul 18, 2025
Merged

Feature: Output real space wavefunction and partial charge density when device=gpu #6391

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
2fb9105
Fix GPU output of out_pchg and out_wfc_norm, out_wfc_re_im
AsTonyshment Jul 17, 2025
d2b9869
GPU integrate test is functional again
AsTonyshment Jul 17, 2025
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
98 changes: 56 additions & 42 deletions source/source_esolver/esolver_ks_pw.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,27 +11,27 @@
#include "source_estate/module_charge/symmetry_rho.h"
#include "source_hamilt/module_ewald/H_Ewald_pw.h"
#include "source_hamilt/module_vdw/vdw.h"
#include "source_lcao/module_deltaspin/spin_constrain.h"
#include "source_lcao/module_dftu/dftu.h"
#include "source_pw/module_pwdft/elecond.h"
#include "source_pw/module_pwdft/forces.h"
#include "source_pw/module_pwdft/hamilt_pw.h"
#include "source_pw/module_pwdft/onsite_projector.h"
#include "source_pw/module_pwdft/stress_pw.h"
#include "source_hsolver/diago_iter_assist.h"
#include "source_hsolver/hsolver_pw.h"
#include "source_hsolver/kernels/dngvd_op.h"
#include "source_io/berryphase.h"
#include "source_io/cal_ldos.h"
#include "source_io/get_pchg_pw.h"
#include "source_io/get_wf_pw.h"
#include "source_io/module_parameter/parameter.h"
#include "source_io/numerical_basis.h"
#include "source_io/numerical_descriptor.h"
#include "source_io/to_wannier90_pw.h"
#include "source_io/winput.h"
#include "source_io/write_dos_pw.h"
#include "source_io/write_wfc_pw.h"
#include "source_io/module_parameter/parameter.h"
#include "source_lcao/module_deltaspin/spin_constrain.h"
#include "source_lcao/module_dftu/dftu.h"
#include "source_pw/module_pwdft/elecond.h"
#include "source_pw/module_pwdft/forces.h"
#include "source_pw/module_pwdft/hamilt_pw.h"
#include "source_pw/module_pwdft/onsite_projector.h"
#include "source_pw/module_pwdft/stress_pw.h"

#include <iostream>

Expand Down Expand Up @@ -717,16 +717,25 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
//------------------------------------------------------------------
// 4) calculate band-decomposed (partial) charge density in pw basis
//------------------------------------------------------------------
const std::vector<int> out_pchg = PARAM.inp.out_pchg;
if (out_pchg.size() > 0)
if (PARAM.inp.out_pchg.size() > 0)
{
ModuleIO::get_pchg_pw(out_pchg,
if (this->__kspw_psi != nullptr && PARAM.inp.precision == "single")
{
delete reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(this->__kspw_psi);
}

// Refresh __kspw_psi
this->__kspw_psi = PARAM.inp.precision == "single"
? new psi::Psi<std::complex<double>, Device>(this->kspw_psi[0])
: reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(this->kspw_psi);

ModuleIO::get_pchg_pw(PARAM.inp.out_pchg,
this->kspw_psi->get_nbands(),
PARAM.inp.nspin,
this->pw_rhod->nxyz,
this->chr.ngmc,
&ucell,
this->psi,
this->__kspw_psi,
this->pw_rhod,
this->pw_wfc,
this->ctx,
Expand Down Expand Up @@ -943,20 +952,25 @@ void ESolver_KS_PW<T, Device>::after_all_runners(UnitCell& ucell)
//----------------------------------------------------------
//! 5) Print out electronic wave functions in real space
//----------------------------------------------------------
const std::vector<int> out_wfc_norm = PARAM.inp.out_wfc_norm;
const std::vector<int> out_wfc_re_im = PARAM.inp.out_wfc_re_im;
if (out_wfc_norm.size() > 0 || out_wfc_re_im.size() > 0)
if (PARAM.inp.out_wfc_norm.size() > 0 || PARAM.inp.out_wfc_re_im.size() > 0)
{
ModuleIO::get_wf_pw(out_wfc_norm,
out_wfc_re_im,
if (this->__kspw_psi != nullptr && PARAM.inp.precision == "single")
{
delete reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(this->__kspw_psi);
}

// Refresh __kspw_psi
this->__kspw_psi = PARAM.inp.precision == "single"
? new psi::Psi<std::complex<double>, Device>(this->kspw_psi[0])
: reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(this->kspw_psi);

ModuleIO::get_wf_pw(PARAM.inp.out_wfc_norm,
PARAM.inp.out_wfc_re_im,
this->kspw_psi->get_nbands(),
PARAM.inp.nspin,
this->pw_rhod->nx,
this->pw_rhod->ny,
this->pw_rhod->nz,
this->pw_rhod->nxyz,
&ucell,
this->psi,
this->__kspw_psi,
this->pw_wfc,
this->ctx,
this->Pgrid,
Expand Down Expand Up @@ -991,29 +1005,29 @@ void ESolver_KS_PW<T, Device>::after_all_runners(UnitCell& ucell)

ModuleIO::Write_MLKEDF_Descriptors write_mlkedf_desc;
write_mlkedf_desc.cal_tool->set_para(this->chr.nrxx,
PARAM.inp.nelec,
PARAM.inp.of_tf_weight,
PARAM.inp.of_vw_weight,
PARAM.inp.of_ml_chi_p,
PARAM.inp.of_ml_chi_q,
PARAM.inp.of_ml_chi_xi,
PARAM.inp.of_ml_chi_pnl,
PARAM.inp.of_ml_chi_qnl,
PARAM.inp.of_ml_nkernel,
PARAM.inp.of_ml_kernel,
PARAM.inp.of_ml_kernel_scaling,
PARAM.inp.of_ml_yukawa_alpha,
PARAM.inp.of_ml_kernel_file,
ucell.omega,
this->pw_rho);
PARAM.inp.nelec,
PARAM.inp.of_tf_weight,
PARAM.inp.of_vw_weight,
PARAM.inp.of_ml_chi_p,
PARAM.inp.of_ml_chi_q,
PARAM.inp.of_ml_chi_xi,
PARAM.inp.of_ml_chi_pnl,
PARAM.inp.of_ml_chi_qnl,
PARAM.inp.of_ml_nkernel,
PARAM.inp.of_ml_kernel,
PARAM.inp.of_ml_kernel_scaling,
PARAM.inp.of_ml_yukawa_alpha,
PARAM.inp.of_ml_kernel_file,
ucell.omega,
this->pw_rho);

write_mlkedf_desc.generateTrainData_KS(PARAM.globalv.global_mlkedf_descriptor_dir,
this->kspw_psi,
this->pelec,
this->pw_wfc,
this->pw_rho,
ucell,
this->pelec->pot->get_effective_v(0));
this->kspw_psi,
this->pelec,
this->pw_wfc,
this->pw_rho,
ucell,
this->pelec->pot->get_effective_v(0));
}
#endif
}
Expand Down
44 changes: 39 additions & 5 deletions source/source_io/get_pchg_pw.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ void get_pchg_pw(const std::vector<int>& out_pchg,
const int nxyz,
const int chr_ngmc,
UnitCell* ucell,
const psi::Psi<std::complex<double>>* psi,
const psi::Psi<std::complex<double>, Device>* kspw_psi,
const ModulePW::PW_Basis* pw_rhod,
const ModulePW::PW_Basis_K* pw_wfc,
const Device* ctx,
Expand Down Expand Up @@ -67,9 +67,17 @@ void get_pchg_pw(const std::vector<int>& out_pchg,
bands_picked[i] = static_cast<int>(out_pchg[i]);
}

// Allocate host memory
std::vector<std::complex<double>> wfcr(nxyz);
std::vector<std::vector<double>> rho_band(nspin, std::vector<double>(nxyz));

// Allocate device memory
std::complex<double>* wfcr_device = nullptr;
if (!std::is_same<Device, base_device::DEVICE_CPU>::value)
{
base_device::memory::resize_memory_op<std::complex<double>, Device>()(wfcr_device, nxyz);
}

for (int ib = 0; ib < nbands; ++ib)
{
// Skip the loop iteration if bands_picked[ib] is 0
Expand All @@ -91,8 +99,21 @@ void get_pchg_pw(const std::vector<int>& out_pchg,
const int spin_index = kv.isk[ik]; // spin index
const int k_number = ikstot % (nkstot / nspin) + 1; // k-point number, starting from 1

psi->fix_k(ik);
pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr.data(), ik);
kspw_psi->fix_k(ik);

// FFT on device and copy result back to host
if (std::is_same<Device, base_device::DEVICE_CPU>::value)
{
pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr.data(), ik);
}
else
{
pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr_device, ik);

base_device::memory::synchronize_memory_op<std::complex<double>,
base_device::DEVICE_CPU,
Device>()(wfcr.data(), wfcr_device, nxyz);
}

// To ensure the normalization of charge density in multi-k calculation (if if_separate_k is true)
double wg_sum_k = 0.0;
Expand Down Expand Up @@ -142,8 +163,21 @@ void get_pchg_pw(const std::vector<int>& out_pchg,
const int spin_index = kv.isk[ik]; // spin index
const int k_number = ikstot % (nkstot / nspin) + 1; // k-point number, starting from 1

psi->fix_k(ik);
pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr.data(), ik);
kspw_psi->fix_k(ik);

// FFT on device and copy result back to host
if (std::is_same<Device, base_device::DEVICE_CPU>::value)
{
pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr.data(), ik);
}
else
{
pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr_device, ik);

base_device::memory::synchronize_memory_op<std::complex<double>,
base_device::DEVICE_CPU,
Device>()(wfcr.data(), wfcr_device, nxyz);
}

double w1 = static_cast<double>(kv.wk[ik] / ucell->omega);

Expand Down
71 changes: 49 additions & 22 deletions source/source_io/get_wf_pw.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,33 +1,16 @@
#ifndef GET_WF_PW_H
#define GET_WF_PW_H

#include "cube_io.h"
#include "source_base/module_device/device.h"
#include "source_base/tool_quit.h"
#include "source_basis/module_pw/pw_basis.h"
#include "source_basis/module_pw/pw_basis_k.h"
#include "source_cell/unitcell.h"
#include "source_estate/elecstate.h"
#include "source_estate/module_charge/symmetry_rho.h"
#include "source_pw/module_pwdft/parallel_grid.h"
#include "source_psi/psi.h"

#include <string>
#include <vector>

namespace ModuleIO
{
template <typename Device>
void get_wf_pw(const std::vector<int>& out_wfc_norm,
const std::vector<int>& out_wfc_re_im,
const int nbands,
const int nspin,
const int nx,
const int ny,
const int nz,
const int nxyz,
UnitCell* ucell,
const psi::Psi<std::complex<double>>* psi,
const psi::Psi<std::complex<double>, Device>* kspw_psi,
const ModulePW::PW_Basis_K* pw_wfc,
const Device* ctx,
const Parallel_Grid& pgrid,
Expand Down Expand Up @@ -94,9 +77,17 @@ void get_wf_pw(const std::vector<int>& out_wfc_norm,
bands_picked_re_im[i] = static_cast<int>(out_wfc_re_im[i]);
}

// Allocate host memory
std::vector<std::complex<double>> wfcr_norm(nxyz);
std::vector<std::vector<double>> rho_band_norm(nspin, std::vector<double>(nxyz));

// Allocate device memory
std::complex<double>* wfcr_norm_device = nullptr;
if (!std::is_same<Device, base_device::DEVICE_CPU>::value)
{
base_device::memory::resize_memory_op<std::complex<double>, Device>()(wfcr_norm_device, nxyz);
}

for (int ib = 0; ib < nbands; ++ib)
{
// Skip the loop iteration if bands_picked[ib] is 0
Expand All @@ -115,8 +106,22 @@ void get_wf_pw(const std::vector<int>& out_wfc_norm,
const int spin_index = kv.isk[ik]; // spin index
const int k_number = ikstot % (nkstot / nspin) + 1; // k-point number, starting from 1

psi->fix_k(ik);
pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr_norm.data(), ik);
kspw_psi->fix_k(ik);

// FFT on device and copy result back to host
if (std::is_same<Device, base_device::DEVICE_CPU>::value)
{
pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr_norm.data(), ik);
}
else
{
pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr_norm_device, ik);

base_device::memory::synchronize_memory_op<std::complex<double>, base_device::DEVICE_CPU, Device>()(
wfcr_norm.data(),
wfcr_norm_device,
nxyz);
}

// To ensure the normalization of charge density in multi-k calculation
double wg_sum_k = 0.0;
Expand Down Expand Up @@ -159,10 +164,18 @@ void get_wf_pw(const std::vector<int>& out_wfc_norm,
}
}

// Allocate host memory
std::vector<std::complex<double>> wfc_re_im(nxyz);
std::vector<std::vector<double>> rho_band_re(nspin, std::vector<double>(nxyz));
std::vector<std::vector<double>> rho_band_im(nspin, std::vector<double>(nxyz));

// Allocate device memory
std::complex<double>* wfc_re_im_device = nullptr;
if (!std::is_same<Device, base_device::DEVICE_CPU>::value)
{
base_device::memory::resize_memory_op<std::complex<double>, Device>()(wfc_re_im_device, nxyz);
}

for (int ib = 0; ib < nbands; ++ib)
{
// Skip the loop iteration if bands_picked[ib] is 0
Expand All @@ -182,8 +195,22 @@ void get_wf_pw(const std::vector<int>& out_wfc_norm,
const int spin_index = kv.isk[ik]; // spin index
const int k_number = ikstot % (nkstot / nspin) + 1; // k-point number, starting from 1

psi->fix_k(ik);
pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfc_re_im.data(), ik);
kspw_psi->fix_k(ik);

// FFT on device and copy result back to host
if (std::is_same<Device, base_device::DEVICE_CPU>::value)
{
pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfc_re_im.data(), ik);
}
else
{
pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfc_re_im_device, ik);

base_device::memory::synchronize_memory_op<std::complex<double>, base_device::DEVICE_CPU, Device>()(
wfc_re_im.data(),
wfc_re_im_device,
nxyz);
}

// To ensure the normalization of charge density in multi-k calculation
double wg_sum_k = 0.0;
Expand Down
22 changes: 22 additions & 0 deletions tests/11_PW_GPU/006_PW_get_wf_GPU/INPUT
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
INPUT_PARAMETERS
#Parameters (1.General)
suffix autotest
calculation scf

nbands 6
symmetry 0
pseudo_dir ../../PP_ORB

device gpu

#Parameters (2.Iteration)
ecutwfc 100
scf_thr 1e-9
scf_nmax 100

#Parameters (3.Basis)
basis_type pw

out_wfc_norm 1*1

pw_seed 1
0 tests/11_PW_GPU/BUG_PW_outwfcr_GPU/KPT → tests/11_PW_GPU/006_PW_get_wf_GPU/KPT
View file
Open in desktop
File renamed without changes.
1 change: 1 addition & 0 deletions tests/11_PW_GPU/006_PW_get_wf_GPU/README
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
GPU version / test the output of out_wfc_norm
0 tests/11_PW_GPU/BUG_PW_outwfcr_GPU/STRU → tests/11_PW_GPU/006_PW_get_wf_GPU/STRU
View file
Open in desktop
File renamed without changes.
11 changes: 11 additions & 0 deletions tests/11_PW_GPU/006_PW_get_wf_GPU/result.ref
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
etotref -211.878925376034
etotperatomref -105.9394626880
wfi1s1k1.cube 21.94721795
wfi1s1k2.cube 19.65847203
wfi1s1k3.cube 19.65846259
wfi1s1k4.cube 19.16322464
wfi1s1k5.cube 19.65846987
wfi1s1k6.cube 19.17198153
wfi1s1k7.cube 18.92365819
wfi1s1k8.cube 19.6584683
totaltimeref 3.26
Loading
Loading