Skip to content

refactor(rt): make block-cyclic layout access explicit via para_2d() #8254

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

Open
kirk0830 wants to merge 3 commits into
base: develop
Choose a base branch
from
+296 −280
Open

refactor(rt): make block-cyclic layout access explicit via para_2d() #8254

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
77163dd
C1: add para_2d() accessor and get_blacs_ctxt() to Parallel_Orbitals
kirk0830 Oct 27, 2026
c1dc33f
D2(module_rt): drop Parallel_2D inheritance in Parallel_Orbitals, mig…
kirk0830 Oct 28, 2026
f748d15
Merge branch 'develop' into pr/01-rt-explicit-para2d
Critsium-xy Nov 2, 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
16 changes: 16 additions & 0 deletions source/source_basis/module_ao/parallel_orbitals.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,22 @@ class Parallel_Orbitals : public Parallel_2D
Parallel_Orbitals();
~Parallel_Orbitals();

/// Access to the block-cyclic distribution information. The non-const
/// overload exists only for the two cases a const getter cannot serve:
/// non-const mutation by setup_2d_division, and test-side descriptor writes.
// DEPRECATED-ALIAS: para_2d() -> get_block_cyclic_distribution()
// DELETION-COMMIT: PR#TBD/3 commits after merge; trigger: git grep -w para_2d == 0
Parallel_2D& para_2d() { return *this; }

const Parallel_2D& para_2d() const { return *this; }

#ifdef __MPI
/// BLACS context of the 2D process grid, completing the getter family
/// (get_row_size / get_col_size / get_local_size / get_block_size /
/// get_dim0 / get_dim1 / get_desc) inherited from Parallel_2D.
int get_blacs_ctxt() const { return blacs_ctxt; }
#endif

/// local size of bands, used for 2d wavefunction
/// must divided on dim1 because of elpa interface
int ncol_bands;
Expand Down
26 changes: 13 additions & 13 deletions source/source_lcao/module_rt/band_energy.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ void gather_eij(const Parallel_Orbitals& pv,
const EijOutputInfo& output_info,
std::ofstream& ofs)
{
const bool is_root = pv.blacs_in_this_processor(0, 0);
const bool is_root = pv.para_2d().blacs_in_this_processor(0, 0);
const size_t buffer_size = is_root ? static_cast<size_t>(nband) * nband : 1;
std::vector<std::complex<double>> gathered(buffer_size);
int source_desc[9];
Expand Down Expand Up @@ -57,16 +57,16 @@ void compute_ekb(const Parallel_Orbitals* pv,
std::ofstream& ofs_running,
const EijOutputInfo& output_info)
{
assert(pv->nloc_wfc > 0 && pv->nloc > 0);
assert(pv->nloc_wfc > 0 && pv->para_2d().get_local_size() > 0);

std::complex<double>* tmp1 = new std::complex<double>[pv->nloc_wfc];
ModuleBase::GlobalFunc::ZEROS(tmp1, pv->nloc_wfc);

std::complex<double>* eij = new std::complex<double>[pv->nloc];
ModuleBase::GlobalFunc::ZEROS(eij, pv->nloc);
std::complex<double>* eij = new std::complex<double>[pv->para_2d().get_local_size()];
ModuleBase::GlobalFunc::ZEROS(eij, pv->para_2d().get_local_size());

ScalapackConnector::
gemm('N', 'N', nlocal, nband, nlocal, 1.0, Htmp, 1, 1, pv->desc, psi_k, 1, 1, pv->desc_wfc, 0.0, tmp1, 1, 1, pv->desc_wfc);
gemm('N', 'N', nlocal, nband, nlocal, 1.0, Htmp, 1, 1, pv->para_2d().get_desc(), psi_k, 1, 1, pv->desc_wfc, 0.0, tmp1, 1, 1, pv->desc_wfc);

ScalapackConnector::
gemm('C', 'N', nband, nband, nlocal, 1.0, psi_k, 1, 1, pv->desc_wfc, tmp1, 1, 1, pv->desc_wfc, 0.0, eij, 1, 1, pv->desc_Eij);
Expand All @@ -83,24 +83,24 @@ void compute_ekb(const Parallel_Orbitals* pv,
double* eii = new double[nband];
ModuleBase::GlobalFunc::ZEROS(eii, nband);

for (int iprow = 0; iprow < pv->dim0; ++iprow)
for (int iprow = 0; iprow < pv->para_2d().get_dim0(); ++iprow)
{
for (int ipcol = 0; ipcol < pv->dim1; ++ipcol)
for (int ipcol = 0; ipcol < pv->para_2d().get_dim1(); ++ipcol)
{
if (pv->blacs_in_this_processor(iprow, ipcol))
if (pv->para_2d().blacs_in_this_processor(iprow, ipcol))
{
naroc[0] = pv->nrow;
naroc[1] = pv->ncol;
naroc[0] = pv->para_2d().get_row_size();
naroc[1] = pv->para_2d().get_col_size();
for (int j = 0; j < naroc[1]; ++j)
{
int igcol = globalIndex(j, pv->nb, pv->dim1, ipcol);
int igcol = globalIndex(j, pv->para_2d().get_block_size(), pv->para_2d().get_dim1(), ipcol);
if (igcol >= nband)
{
continue;
}
for (int i = 0; i < naroc[0]; ++i)
{
int igrow = globalIndex(i, pv->nb, pv->dim0, iprow);
int igrow = globalIndex(i, pv->para_2d().get_block_size(), pv->para_2d().get_dim0(), iprow);
if (igrow >= nband)
{
continue;
Expand All @@ -114,7 +114,7 @@ void compute_ekb(const Parallel_Orbitals* pv,
}
} // loop ipcol
} // loop iprow
info = MPI_Allreduce(eii, ekb, nband, MPI_DOUBLE, MPI_SUM, pv->comm());
info = MPI_Allreduce(eii, ekb, nband, MPI_DOUBLE, MPI_SUM, pv->para_2d().comm());

delete[] tmp1;
delete[] eij;
Expand Down
12 changes: 6 additions & 6 deletions source/source_lcao/module_rt/boundary_fix.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,23 +86,23 @@ void boundary_shift_mat(const std::complex<double>& phase,
const std::complex<double> phase_conj = std::conj(phase);
size_t row0 = pv->atom_begin_row[iat];
size_t col0 = pv->atom_begin_col[iat];
std::complex<double>* p_matkc = matk + col0 * pv->get_row_size();
std::complex<double>* p_matkc = matk + col0 * pv->para_2d().get_row_size();
for(size_t nu = 0; nu < pv->get_ncol_atom(iat); ++nu)
{

BlasConnector::scal(pv->get_row_size(),
BlasConnector::scal(pv->para_2d().get_row_size(),
phase,
p_matkc,
1);
p_matkc += pv->get_row_size();
p_matkc += pv->para_2d().get_row_size();
}
std::complex<double>* p_matkr = matk + row0;
for(size_t mu = 0; mu < pv->get_nrow_atom(iat); ++mu)
{
BlasConnector::scal(pv->get_col_size(),
BlasConnector::scal(pv->para_2d().get_col_size(),
phase_conj,
p_matkr,
pv->get_row_size());
pv->para_2d().get_row_size());
p_matkr += 1;
}
return;
Expand All @@ -121,7 +121,7 @@ void boundary_shift_c(const std::complex<double>& phase,
BlasConnector::scal(pv->ncol_bands,
phase_conj,
p_ck,
pv->get_row_size());
pv->para_2d().get_row_size());
p_ck+=1;
}
return;
Expand Down
8 changes: 4 additions & 4 deletions source/source_lcao/module_rt/evolve_elec.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,13 +59,13 @@ void Evolve_elec<Device>::solve_psi(const std::vector<EijOutputInfo>& eij_output
if (!use_tensor)
{
// Construct the local P_k matrix for moving spatial gauge, CPU only for now
std::vector<std::complex<double>> P_k_local(para_orb.nloc, {0.0, 0.0});
std::vector<std::complex<double>> P_k_local(para_orb.para_2d().get_local_size(), {0.0, 0.0});
if (use_td_moving_gauge && td_mg != nullptr)
{
td_mg->get_P_k(ucell, kvec_d[ik], P_k_local.data(), para_orb.nloc, para_orb.ncol);
td_mg->get_P_k(ucell, kvec_d[ik], P_k_local.data(), para_orb.para_2d().get_local_size(), para_orb.para_2d().get_col_size());
}

const int len_HS_laststep = use_lapack ? nlocal * nlocal : para_orb.nloc;
const int len_HS_laststep = use_lapack ? nlocal * nlocal : para_orb.para_2d().get_local_size();
evolve_psi(nband,
nlocal,
&(para_orb),
Expand All @@ -90,7 +90,7 @@ void Evolve_elec<Device>::solve_psi(const std::vector<EijOutputInfo>& eij_output

const int len_psi_k_1 = use_lapack ? nband : psi->get_nbands();
const int len_psi_k_2 = use_lapack ? nlocal : psi->get_nbasis();
const int len_HS_laststep = use_lapack ? nlocal * nlocal : para_orb.nloc;
const int len_HS_laststep = use_lapack ? nlocal * nlocal : para_orb.para_2d().get_local_size();

// Create Tensor for psi_k, psi_k_laststep, H_laststep, S_laststep, ekb
ct::Tensor psi_k_tensor(ct::DataType::DT_COMPLEX_DOUBLE, ct_device_type, ct::TensorShape({len_psi_k_1, len_psi_k_2}));
Expand Down
24 changes: 12 additions & 12 deletions source/source_lcao/module_rt/evolve_psi.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,20 +44,20 @@ void evolve_psi(const int nband,
hamilt::MatrixBlock<std::complex<double>> s_mat;
p_hamilt->matrix(h_mat, s_mat);

std::complex<double>* Stmp = new std::complex<double>[pv->nloc];
ModuleBase::GlobalFunc::ZEROS(Stmp, pv->nloc);
BlasConnector::copy(pv->nloc, s_mat.p, 1, Stmp, 1);
std::complex<double>* Stmp = new std::complex<double>[pv->para_2d().get_local_size()];
ModuleBase::GlobalFunc::ZEROS(Stmp, pv->para_2d().get_local_size());
BlasConnector::copy(pv->para_2d().get_local_size(), s_mat.p, 1, Stmp, 1);

std::complex<double>* Htmp = new std::complex<double>[pv->nloc];
ModuleBase::GlobalFunc::ZEROS(Htmp, pv->nloc);
BlasConnector::copy(pv->nloc, h_mat.p, 1, Htmp, 1);
std::complex<double>* Htmp = new std::complex<double>[pv->para_2d().get_local_size()];
ModuleBase::GlobalFunc::ZEROS(Htmp, pv->para_2d().get_local_size());
BlasConnector::copy(pv->para_2d().get_local_size(), h_mat.p, 1, Htmp, 1);

std::complex<double>* Hold = new std::complex<double>[pv->nloc];
ModuleBase::GlobalFunc::ZEROS(Hold, pv->nloc);
BlasConnector::copy(pv->nloc, h_mat.p, 1, Hold, 1);
std::complex<double>* Hold = new std::complex<double>[pv->para_2d().get_local_size()];
ModuleBase::GlobalFunc::ZEROS(Hold, pv->para_2d().get_local_size());
BlasConnector::copy(pv->para_2d().get_local_size(), h_mat.p, 1, Hold, 1);

std::complex<double>* U_operator = new std::complex<double>[pv->nloc];
ModuleBase::GlobalFunc::ZEROS(U_operator, pv->nloc);
std::complex<double>* U_operator = new std::complex<double>[pv->para_2d().get_local_size()];
ModuleBase::GlobalFunc::ZEROS(U_operator, pv->para_2d().get_local_size());

// (1)->>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

Expand Down Expand Up @@ -204,7 +204,7 @@ void evolve_psi_tensor(const int nband,
s_src = s_mat.p;
}

const int len_HS = use_lapack ? nlocal * nlocal : pv->nloc;
const int len_HS = use_lapack ? nlocal * nlocal : pv->para_2d().get_local_size();

ct::Tensor Stmp(ct::DataType::DT_COMPLEX_DOUBLE, ct_device_type, ct::TensorShape({len_HS}));

Expand Down
40 changes: 20 additions & 20 deletions source/source_lcao/module_rt/force_rt_overlap.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,8 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
// get dS/dR_{x,y,z}
std::vector<hamilt::HContainer<double>*> dsxr = td_info->get_grad_overlap();
// allocate matrix
const long nloc = pv.nloc;
const int nrow = pv.nrow;
const long nloc = pv.para_2d().get_local_size();
const int nrow = pv.para_2d().get_row_size();
std::vector<std::complex<double>> Htmp(nloc);
std::vector<std::complex<double>> Sinv(nloc);
std::vector<std::complex<double>> dsxk(nloc);
Expand Down Expand Up @@ -63,7 +63,7 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
const std::complex<double> mone_complex = {-1.0, 0.0};
const std::complex<double> zero_complex = {0.0, 0.0};

ScalapackConnector::getrf(nlocal, nlocal, Sinv.data(), one_int, one_int, pv.desc, ipiv.data(), &info);
ScalapackConnector::getrf(nlocal, nlocal, Sinv.data(), one_int, one_int, pv.para_2d().get_desc(), ipiv.data(), &info);

int lwork = -1;
int liwork = -1;
Expand All @@ -78,7 +78,7 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
Sinv.data(),
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
ipiv.data(),
work.data(),
&lwork,
Expand All @@ -95,7 +95,7 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
Sinv.data(),
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
ipiv.data(),
work.data(),
&lwork,
Expand All @@ -116,16 +116,16 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
tmp_dmk,
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
Htmp.data(),
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
zero_complex,
tmp1.data(),
one_int,
one_int,
pv.desc);
pv.para_2d().get_desc());

ScalapackConnector::gemm(N_char,
N_char,
Expand All @@ -136,16 +136,16 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
tmp1.data(),
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
Sinv.data(),
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
zero_complex,
tmp2.data(),
one_int,
one_int,
pv.desc);
pv.para_2d().get_desc());
for (int dir = 0; dir < 3; dir++)
{
ModuleBase::GlobalFunc::ZEROS(dsxk.data(), nloc);
Expand All @@ -171,16 +171,16 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
tmp2.data(),
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
dsxk.data(),
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
one_complex,
tmp_out[dir].data(),
one_int,
one_int,
pv.desc);
pv.para_2d().get_desc());

ScalapackConnector::gemm(T_char,
N_char,
Expand All @@ -191,16 +191,16 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
tmp_dmk,
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
pdsxk.data(),
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
one_complex,
tmp3.data(),
one_int,
one_int,
pv.desc);
pv.para_2d().get_desc());

ScalapackConnector::geadd(N_char,
nlocal,
Expand All @@ -209,12 +209,12 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
tmp3.data(),
one_int,
one_int,
pv.desc,
pv.para_2d().get_desc(),
one_complex,
tmp_out[dir].data(),
one_int,
one_int,
pv.desc);
pv.para_2d().get_desc());
}
}
// std::string filename = "process_debug_" + std::to_string(GlobalV::MY_RANK) + ".txt";
Expand Down Expand Up @@ -245,7 +245,7 @@ void cal_foverlap_rt(ModuleBase::matrix& foverlap,
double* force_tmp1 = &foverlap(iat, 0);
int row0 = pv.atom_begin_row[iat];
int col0 = pv.atom_begin_col[iat];
const int row_size = pv.get_row_size();
const int row_size = pv.para_2d().get_row_size();
std::vector<std::complex<double>*> p_diag = {tmp_out[0].data(), tmp_out[1].data(), tmp_out[2].data()};
for (int mu = 0; mu < pv.get_nrow_atom(iat); ++mu)
{
Expand Down
22 changes: 11 additions & 11 deletions source/source_lcao/module_rt/middle_hamilt.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,10 @@ void half_Hmatrix(const Parallel_Orbitals* pv,
ofs_running << std::setprecision(10);
ofs_running << std::endl;
ofs_running << " H(t+dt) :" << std::endl;
for (int i = 0; i < pv->nrow; i++)
for (int i = 0; i < pv->para_2d().get_row_size(); i++)
{
const int in = i * pv->ncol;
for (int j = 0; j < pv->ncol; j++)
const int in = i * pv->para_2d().get_col_size();
for (int j = 0; j < pv->para_2d().get_col_size(); j++)
{
ofs_running << Htmp[in + j].real() << "+" << Htmp[in + j].imag() << "i ";
}
Expand All @@ -41,10 +41,10 @@ void half_Hmatrix(const Parallel_Orbitals* pv,
ofs_running << std::endl;
ofs_running << std::endl;
ofs_running << " H(t):" << std::endl;
for (int i = 0; i < pv->nrow; i++)
for (int i = 0; i < pv->para_2d().get_row_size(); i++)
{
const int in = i * pv->ncol;
for (int j = 0; j < pv->ncol; j++)
const int in = i * pv->para_2d().get_col_size();
for (int j = 0; j < pv->para_2d().get_col_size(); j++)
{
ofs_running << H_laststep[in + j].real() << "+" << H_laststep[in + j].imag() << "i ";
}
Expand All @@ -55,17 +55,17 @@ void half_Hmatrix(const Parallel_Orbitals* pv,

std::complex<double> alpha = {0.5, 0.0};
std::complex<double> beta = {0.5, 0.0};
ScalapackConnector::geadd('N', nlocal, nlocal, alpha, H_laststep, 1, 1, pv->desc, beta, Htmp, 1, 1, pv->desc);
ScalapackConnector::geadd('N', nlocal, nlocal, alpha, S_laststep, 1, 1, pv->desc, beta, Stmp, 1, 1, pv->desc);
ScalapackConnector::geadd('N', nlocal, nlocal, alpha, H_laststep, 1, 1, pv->para_2d().get_desc(), beta, Htmp, 1, 1, pv->para_2d().get_desc());
ScalapackConnector::geadd('N', nlocal, nlocal, alpha, S_laststep, 1, 1, pv->para_2d().get_desc(), beta, Stmp, 1, 1, pv->para_2d().get_desc());

if (print_matrix)
{
ofs_running << std::endl;
ofs_running << " H (t+dt/2) :" << std::endl;
for (int i = 0; i < pv->nrow; i++)
for (int i = 0; i < pv->para_2d().get_row_size(); i++)
{
const int in = i * pv->ncol;
for (int j = 0; j < pv->ncol; j++)
const int in = i * pv->para_2d().get_col_size();
for (int j = 0; j < pv->para_2d().get_col_size(); j++)
{
ofs_running << Htmp[in + j].real() << "+" << Htmp[in + j].imag() << "i ";
}
Expand Down
Loading
Loading