diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md index a03b3c4e51..f0ccd7e208 100644 --- a/docs/advanced/input_files/input-main.md +++ b/docs/advanced/input_files/input-main.md @@ -671,13 +671,22 @@ ### kpar - **Type**: Integer -- **Description**: Divide all processors into kpar groups, and k points will be distributed among each group. The value taken should be less than or equal to the number of k points as well as the number of MPI processes. +- **Description**: Controls k-point parallelism. The value must be positive and should not exceed either the number of k-points or the number of MPI processes. + - For PW calculations, divide all MPI processes into persistent k-point pools. Each pool stores and processes a subset of the k-points. + - For LCAO calculations with lapack, genelpa, elpa, or scalapack_gvx, divide the diagonalization work into temporary k-point pools. After diagonalization, the eigenvalues and distributed wavefunctions are restored for all k-points before occupations, density matrices, and output are evaluated. + - Multi-process LCAO cusolver uses its own active-GPU distribution and does not use this value to define its k-point layout. Other LCAO eigensolvers do not use the temporary k-point-pool implementation. - **Default**: 1 ### bndpar - **Type**: Integer -- **Description**: Divide all processors into bndpar groups for SDFT or the BPCG solver. bndpar must be positive, no greater than the number of MPI processes, and kpar * bndpar must divide the number of MPI processes exactly. +- **Availability**: *([`basis_type`](#basis_type)==pw and [`esolver_type`](#esolver_type)==sdft) or ([`basis_type`](#basis_type)==pw and [`esolver_type`](#esolver_type)==ksdft and [`ks_solver`](#ks_solver)==bpcg)* +- **Description**: Controls band-group parallelism for PW SDFT and PW KSDFT calculations using the BPCG eigensolver. + - Within each k-point pool, divide the MPI processes into bndpar band groups. Each group contains NPROC / (kpar * bndpar) processes when bndpar is greater than 1. + - With BPCG, distribute contiguous ranges of global Kohn-Sham bands among the band groups. nbands does not need to be divisible by bndpar, but bndpar cannot exceed a positive nbands. Groups with lower indices receive one additional band when necessary. + - In SDFT, distribute stochastic orbitals among the band groups. When the deterministic Kohn-Sham eigensolver is not BPCG, band group 0 calculates the deterministic orbitals and broadcasts them to the other groups. + - bndpar must be positive and no greater than the number of MPI processes. When bndpar is greater than 1, kpar * bndpar must divide the number of MPI processes exactly. + > Note: For PW calculations on GPU, if the input kpar * bndpar differs from the number of MPI processes, ABACUS automatically sets the effective kpar to NPROC / bndpar. - **Default**: 1 ### latname @@ -4492,9 +4501,9 @@ ### ocp_set - **Type**: String -- **Description**: Fixed occupation weights used when ocp is true. Values are assigned band by band for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x. +- **Description**: Fixed occupation weights used when ocp is true. Values are assigned in band order for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x. - Example: 1 10*1 0 1 expands to 13 values, with the 12th value equal to 0 and all other values equal to 1. - - After expansion, the array length must equal nks * nbands. + - After expansion, provide one block of nbands values for each k-point. If nspin is 2, provide all k-point blocks for spin up followed by all k-point blocks for spin down; otherwise, provide one block per k-point. - The sum of all weights must equal nelec; otherwise the calculation terminates with an error. - **Default**: None diff --git a/docs/parameters.yaml b/docs/parameters.yaml index 78b5997b84..80cc83b3b5 100644 --- a/docs/parameters.yaml +++ b/docs/parameters.yaml @@ -100,7 +100,10 @@ parameters: category: System variables type: Integer description: | - Divide all processors into kpar groups, and k points will be distributed among each group. The value taken should be less than or equal to the number of k points as well as the number of MPI processes. + Controls k-point parallelism. The value must be positive and should not exceed either the number of k-points or the number of MPI processes. + * For PW calculations, divide all MPI processes into persistent k-point pools. Each pool stores and processes a subset of the k-points. + * For LCAO calculations with lapack, genelpa, elpa, or scalapack_gvx, divide the diagonalization work into temporary k-point pools. After diagonalization, the eigenvalues and distributed wavefunctions are restored for all k-points before occupations, density matrices, and output are evaluated. + * Multi-process LCAO cusolver uses its own active-GPU distribution and does not use this value to define its k-point layout. Other LCAO eigensolvers do not use the temporary k-point-pool implementation. default_value: "1" unit: "" availability: "" @@ -108,10 +111,15 @@ parameters: category: System variables type: Integer description: | - Divide all processors into bndpar groups for SDFT or the BPCG solver. bndpar must be positive, no greater than the number of MPI processes, and kpar * bndpar must divide the number of MPI processes exactly. + Controls band-group parallelism for PW SDFT and PW KSDFT calculations using the BPCG eigensolver. + * Within each k-point pool, divide the MPI processes into bndpar band groups. Each group contains NPROC / (kpar * bndpar) processes when bndpar is greater than 1. + * With BPCG, distribute contiguous ranges of global Kohn-Sham bands among the band groups. nbands does not need to be divisible by bndpar, but bndpar cannot exceed a positive nbands. Groups with lower indices receive one additional band when necessary. + * In SDFT, distribute stochastic orbitals among the band groups. When the deterministic Kohn-Sham eigensolver is not BPCG, band group 0 calculates the deterministic orbitals and broadcasts them to the other groups. + * bndpar must be positive and no greater than the number of MPI processes. When bndpar is greater than 1, kpar * bndpar must divide the number of MPI processes exactly. + [NOTE] For PW calculations on GPU, if the input kpar * bndpar differs from the number of MPI processes, ABACUS automatically sets the effective kpar to NPROC / bndpar. default_value: "1" unit: "" - availability: "" + availability: (basis_type==pw and esolver_type==sdft) or (basis_type==pw and esolver_type==ksdft and ks_solver==bpcg) - name: latname category: System variables type: String @@ -2744,9 +2752,9 @@ parameters: category: "RT-TDDFT: Real-Time Time-Dependent Density Functional Theory" type: String description: | - Fixed occupation weights used when ocp is true. Values are assigned band by band for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x. + Fixed occupation weights used when ocp is true. Values are assigned in band order for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x. * Example: 1 10*1 0 1 expands to 13 values, with the 12th value equal to 0 and all other values equal to 1. - * After expansion, the array length must equal nks * nbands. + * After expansion, provide one block of nbands values for each k-point. If nspin is 2, provide all k-point blocks for spin up followed by all k-point blocks for spin down; otherwise, provide one block per k-point. * The sum of all weights must equal nelec; otherwise the calculation terminates with an error. default_value: None unit: "" diff --git a/source/Makefile.Objects b/source/Makefile.Objects index e2e6c10491..860977acfc 100644 --- a/source/Makefile.Objects +++ b/source/Makefile.Objects @@ -552,6 +552,7 @@ OBJS_IO=module_parameter/input_conv.o\ cal_test.o\ module_dos/write_dos_pw.o\ nscf_fermi_surf.o\ + band_parallel_output.o\ write_bands.o\ module_dos/cal_dos.o\ module_dos/cal_pdos_gamma.o\ diff --git a/source/source_base/parallel_grid.cpp b/source/source_base/parallel_grid.cpp index fb43a2cb70..a2c0d2c237 100644 --- a/source/source_base/parallel_grid.cpp +++ b/source/source_base/parallel_grid.cpp @@ -331,60 +331,81 @@ void Parallel_Grid::zpiece_to_stogroup(double* zpiece, const int& iz, double* rh // ofs_running << "\n iz = " << iz << " Done."; return; } + +// Taoni modified on 2026-08-21, fixed BPCG out_chg MPI_ERR_RANK void Parallel_Grid::reduce(double* rhotot, const double* const rhoin, const bool reduce_all_pool) const { - // ModuleBase::TITLE("Parallel_Grid","reduce"); - - // if not the first pool, wait here until processpr 0 - // send the Barrier command. + // POOL_WORLD communicators are disjoint, so inactive pools may return + // without skipping a collective required by an active pool. if (!reduce_all_pool && GlobalV::MY_POOL != 0) { return; } - double* zpiece = new double[this->ncxy]; + assert(rhoin != nullptr); + assert(this->nrxx == this->ncxy * this->nczp); - for (int iz = 0; iz < this->ncz; iz++) - { - const int znow = iz - this->startz[GlobalV::MY_POOL][GlobalV::RANK_IN_POOL]; - const int proc = this->whichpro[GlobalV::MY_POOL][iz]; - const int proc_loc = this->whichpro_loc[GlobalV::MY_POOL][iz]; // Obtain the local processor index in the pool - ModuleBase::GlobalFunc::ZEROS(zpiece, this->ncxy); - int tag = iz; - MPI_Status ierror; - - // Local processor 0 collects data from all other processors in the pool - // proc = proc_loc if GlobalV::MY_POOL == 0 - if (proc_loc == GlobalV::RANK_IN_POOL) - { - for (int ir = 0; ir < ncxy; ir++) - { - zpiece[ir] = rhoin[ir * this->nczp + znow]; - } - // Send data to the root of the pool - if (GlobalV::RANK_IN_POOL != 0) - { - MPI_Send(zpiece, ncxy, MPI_DOUBLE, 0, tag, POOL_WORLD); - } - } + int pool_size = 0; + MPI_Comm_size(POOL_WORLD, &pool_size); + + // Derive slab ownership from the communicator used below. Precomputed + // source ranks may belong to a larger BPCG layout and are not necessarily + // valid ranks in the current POOL_WORLD. + std::vector local_z_counts(pool_size); + MPI_Allgather(&this->nczp, 1, MPI_INT, local_z_counts.data(), 1, MPI_INT, POOL_WORLD); - // The root of the pool receives data from other processors - if (GlobalV::RANK_IN_POOL == 0 && proc_loc != GlobalV::RANK_IN_POOL) + // MPI_Gatherv concatenates complete [xy][local_z] buffers by rank. + // These displacements therefore address rank blocks, not global z coordinates. + std::vector receive_counts(pool_size); + std::vector displacements(pool_size, 0); + int total_z = 0; + for (int rank = 0; rank < pool_size; ++rank) + { + receive_counts[rank] = local_z_counts[rank] * this->ncxy; + if (rank > 0) { - MPI_Recv(zpiece, ncxy, MPI_DOUBLE, proc_loc, tag, POOL_WORLD, &ierror); + displacements[rank] = displacements[rank - 1] + receive_counts[rank - 1]; } + total_z += local_z_counts[rank]; + } + assert(total_z == this->ncz); - if (GlobalV::RANK_IN_POOL == 0) + // Only the pool root needs storage for the gathered global grid. + // MPI ignores the receive buffer on all non-root ranks. + std::vector gathered_data; + if (GlobalV::RANK_IN_POOL == 0) + { + assert(rhotot != nullptr); + gathered_data.resize(this->ncxyz); + } + MPI_Gatherv(rhoin, + this->nrxx, + MPI_DOUBLE, + gathered_data.data(), + receive_counts.data(), + displacements.data(), + MPI_DOUBLE, + 0, + POOL_WORLD); + + if (GlobalV::RANK_IN_POOL == 0) + { + // Rank blocks cannot be copied directly to rhotot: each block stores + // [xy][local_z], whereas Cube output expects [xy][global_z]. + // The slab decomposition is contiguous and ordered by POOL_WORLD rank. + int global_z_start = 0; + for (int rank = 0; rank < pool_size; ++rank) { + const int local_z = local_z_counts[rank]; for (int ixy = 0; ixy < this->ncxy; ++ixy) { - rhotot[ixy * ncz + iz] = zpiece[ixy]; + for (int iz = 0; iz < local_z; ++iz) + { + rhotot[ixy * this->ncz + global_z_start + iz] = gathered_data[displacements[rank] + ixy * local_z + iz]; + } } + global_z_start += local_z; } } - - delete[] zpiece; - - return; } #endif diff --git a/source/source_cell/cal_atoms_info.h b/source/source_cell/cal_atoms_info.h index 950aa0d19c..506d164149 100644 --- a/source/source_cell/cal_atoms_info.h +++ b/source/source_cell/cal_atoms_info.h @@ -143,6 +143,11 @@ class CalAtomsInfo result.nbands = nbands; unitcell::cal_nbands(static_cast(result.nelec), result.nlocal, nelec_spin, result.nbands, esolver_type, lspinorb, nspin, basis_type, smearing_method); + // Taoni add check for bndpar > nbands on 2026-08-21 + if (ks_solver == "bpcg" && result.nbands > 0 && bndpar > result.nbands) + { + ModuleBase::WARNING_QUIT("CalAtomsInfo", "bndpar cannot exceed nbands when ks_solver is bpcg"); + } // calculate the number of nbands_local result.nbands_l = result.nbands; diff --git a/source/source_cell/parallel_kpoints.cpp b/source/source_cell/parallel_kpoints.cpp index 8288cc01ba..ebe868ef50 100644 --- a/source/source_cell/parallel_kpoints.cpp +++ b/source/source_cell/parallel_kpoints.cpp @@ -113,16 +113,20 @@ void Parallel_Kpoints::set_startpro_pool() return; } - // gather kpoints from all processor pools, only need to be called by the first processor of each pool. void Parallel_Kpoints::gatherkvec(const std::vector>& vec_local, std::vector>& vec_global) const { + int world_rank = 0; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + // Taoni fix bndpar on 2026-08-21 + // rank_in_pool is repeated in every band group under BPCG. The first + // world rank of the k-point pool instead identifies one unique contributor. + const bool is_pool_root = (world_rank == this->startpro_pool[this->my_pool]); vec_global.resize(this->nkstot_np, ModuleBase::Vector3(0.0, 0.0, 0.0)); for (int i = 0; i < this->nks_np; ++i) { - - if (this->rank_in_pool == 0) + if (is_pool_root) { vec_global[i + startk_pool[this->my_pool]] = vec_local[i]; } @@ -154,7 +158,6 @@ void Parallel_Kpoints::pool_collection(double& value, const double* wk, const in { MPI_Status ierror; MPI_Recv(&value, 1, MPI_DOUBLE, this->startpro_pool[pool], ik, MPI_COMM_WORLD, &ierror); - } } else @@ -215,7 +218,7 @@ void Parallel_Kpoints::pool_collection_aux(T* value, const V& w, const int& dim, T* p = &w.ptr[begin]; // temprary restrict kpar=1 for NSPIN=2 case for generating_orbitals int pool = 0; - if (this->nspin != 2) + if (this->nspin != 2) { pool = this->whichpool[ik]; } diff --git a/source/source_estate/elecstate_tools.cpp b/source/source_estate/elecstate_tools.cpp index 3076c6b437..87abd69fef 100644 --- a/source/source_estate/elecstate_tools.cpp +++ b/source/source_estate/elecstate_tools.cpp @@ -1,6 +1,20 @@ #include "elecstate_tools.h" #include "occupy.h" +#include "source_base/parallel_comm.h" #include "source_base/parallel_reduce.h" + +namespace +{ +int get_band_offset(const int local_nbands) +{ + int band_end = local_nbands; +#ifdef __MPI + MPI_Scan(&local_nbands, &band_end, 1, MPI_INT, MPI_SUM, BP_WORLD); +#endif + return band_end - local_nbands; +} +} // namespace + namespace elecstate { void calEBand(const ModuleBase::matrix& ekb,const ModuleBase::matrix& wg,fenergy& f_en) @@ -42,11 +56,14 @@ namespace elecstate const int nks = ekb.nr; if (!(Occupy::use_gaussian_broadening || Occupy::fixed_occupations)) { + // Taoni fix smearing_method=fixed for BPCG on 2026-08-21 + const int band_offset = get_band_offset(nbands); if (PARAM.globalv.two_fermi) { Occupy::iweights(nks, klist->wk, nbands, + band_offset, nelec_spin[0], ekb, eferm.ef_up, @@ -56,6 +73,7 @@ namespace elecstate Occupy::iweights(nks, klist->wk, nbands, + band_offset, nelec_spin[1], ekb, eferm.ef_dw, @@ -70,6 +88,7 @@ namespace elecstate Occupy::iweights(nks, klist->wk, nbands, + band_offset, PARAM.inp.nelec, ekb, eferm.ef, @@ -151,15 +170,15 @@ namespace elecstate const double ne_thr = 1.0e-5; - const int num = klist->get_nks() * nbands; - if (num != ocp_kb.size()) + const std::size_t expected_size = static_cast(klist->get_nkstot()) * static_cast(nbands); + if (expected_size != ocp_kb.size()) { ModuleBase::WARNING_QUIT("ElecState::fixed_weights", "size of occupation array is wrong , please check ocp_set"); } double num_elec = 0.0; - for (int i = 0; i < ocp_kb.size(); ++i) + for (std::size_t i = 0; i < ocp_kb.size(); ++i) { num_elec += ocp_kb[i]; } @@ -170,11 +189,28 @@ namespace elecstate "total number of occupations is wrong , please check ocp_set"); } + const int band_offset = get_band_offset(wg.nc); + if (klist->ik2iktot.size() < static_cast(wg.nr) + || band_offset < 0 + || band_offset + wg.nc > nbands) + { + ModuleBase::WARNING_QUIT("ElecState::fixed_weights", "invalid distributed occupation layout"); + } + for (int ik = 0; ik < wg.nr; ++ik) { + const int global_k = klist->ik2iktot[ik]; + if (global_k < 0 || global_k >= klist->get_nkstot()) + { + ModuleBase::WARNING_QUIT("ElecState::fixed_weights", "invalid global k-point index"); + } for (int ib = 0; ib < wg.nc; ++ib) { - wg(ik, ib) = ocp_kb[ik * wg.nc + ib]; + // Taoni fix ocp_set under kpar and bndpar on 2026-08-21 + const int global_band = band_offset + ib; + const std::size_t occupation_index + = static_cast(global_k) * static_cast(nbands) + static_cast(global_band); + wg(ik, ib) = ocp_kb[occupation_index]; } } skip_weights = true; diff --git a/source/source_estate/occupy.cpp b/source/source_estate/occupy.cpp index fa50d1520d..0a492735c6 100644 --- a/source/source_estate/occupy.cpp +++ b/source/source_estate/occupy.cpp @@ -120,7 +120,8 @@ void Occupy::decision(const std::string& name, const std::string& smearing_metho * * @param nks number of k points. * @param wk weight of each k point (consider symmetry). - * @param nbands number of bands. + * @param nbands number of locally owned bands. + * @param band_offset global index of the first locally owned band. * @param nelec number of electrons for this spin direction. * @param ekb the array save the band energy. * @param ef output: the highest occupied Kohn-Sham level. @@ -132,6 +133,7 @@ void Occupy::iweights( const int nks, const std::vector& wk, const int nbands, + const int band_offset, const double& nelec, const ModuleBase::matrix& ekb, double& ef, @@ -168,7 +170,8 @@ void Occupy::iweights( for (int ib = 0; ib < nbands; ++ib) { - if (ib < ib_min) + const int global_band = band_offset + ib; + if (global_band < ib_min) { wg(ik, ib) = wk[ik]; ef = std::max(ef, ekb(ik, ib)); diff --git a/source/source_estate/occupy.h b/source/source_estate/occupy.h index 24626e0d16..0a9f7e585c 100644 --- a/source/source_estate/occupy.h +++ b/source/source_estate/occupy.h @@ -38,6 +38,7 @@ class Occupy static void iweights(const int nks, const std::vector& wk, const int nband, + const int band_offset, const double& nelec, const ModuleBase::matrix& ekb, double& ef, diff --git a/source/source_estate/test/elecstate_base_test.cpp b/source/source_estate/test/elecstate_base_test.cpp index e917bfe749..3dba3092fe 100644 --- a/source/source_estate/test/elecstate_base_test.cpp +++ b/source/source_estate/test/elecstate_base_test.cpp @@ -249,6 +249,8 @@ TEST_F(ElecStateTest, FixedWeights) PARAM.input.nelec = 30; K_Vectors* klist = new K_Vectors; klist->set_nks(5); + klist->set_nkstot(5); + klist->ik2iktot = {0, 1, 2, 3, 4}; elecstate->klist = klist; elecstate->wg.create(klist->get_nks(), PARAM.input.nbands); std::vector ocp_kb; @@ -269,6 +271,8 @@ TEST_F(ElecStateDeathTest, FixedWeightsWarning1) PARAM.input.nelec = 30; K_Vectors* klist = new K_Vectors; klist->set_nks(5); + klist->set_nkstot(5); + klist->ik2iktot = {0, 1, 2, 3, 4}; elecstate->klist = klist; elecstate->wg.create(klist->get_nks(), PARAM.input.nbands); std::vector ocp_kb; @@ -291,6 +295,8 @@ TEST_F(ElecStateDeathTest, FixedWeightsWarning2) PARAM.input.nelec = 29; K_Vectors* klist = new K_Vectors; klist->set_nks(5); + klist->set_nkstot(5); + klist->ik2iktot = {0, 1, 2, 3, 4}; elecstate->klist = klist; elecstate->wg.create(klist->get_nks(), PARAM.input.nbands); std::vector ocp_kb; diff --git a/source/source_estate/test/elecstate_occupy_test.cpp b/source/source_estate/test/elecstate_occupy_test.cpp index f4c17ce24d..8f034ee93b 100644 --- a/source/source_estate/test/elecstate_occupy_test.cpp +++ b/source/source_estate/test/elecstate_occupy_test.cpp @@ -189,7 +189,7 @@ TEST_F(OccupyTest, IweightsNOSPIN) ModuleBase::matrix ekb(1, 1); std::vector isk(1); ekb(0, 0) = 0.1; - occupy.iweights(1, wk, 1, 2.0, ekb, ef, wg, 0, isk); + occupy.iweights(1, wk, 1, 0, 2.0, ekb, ef, wg, 0, isk); EXPECT_DOUBLE_EQ(wg(0, 0), 2.0); EXPECT_DOUBLE_EQ(ef, 0.1); } @@ -207,8 +207,8 @@ TEST_F(OccupyTest, IweightsSPIN) isk[1] = 1; ekb(0, 0) = 0.1; ekb(1, 0) = 0.2; - occupy.iweights(2, wk, 1, 1.0, ekb, ef_up, wg, 0, isk); - occupy.iweights(2, wk, 1, 1.0, ekb, ef_dw, wg, 1, isk); + occupy.iweights(2, wk, 1, 0, 1.0, ekb, ef_up, wg, 0, isk); + occupy.iweights(2, wk, 1, 0, 1.0, ekb, ef_dw, wg, 1, isk); EXPECT_DOUBLE_EQ(wg(0, 0), 1.0); EXPECT_DOUBLE_EQ(wg(1, 0), 1.0); EXPECT_DOUBLE_EQ(ef_up, 0.1); @@ -226,7 +226,7 @@ TEST_F(OccupyTest, IweightsWarning) ekb(0, 0) = 0.1; testing::internal::CaptureStdout(); - EXPECT_EXIT(occupy.iweights(1, wk, 1, 1.0, ekb, ef, wg, -1, isk);, ::testing::ExitedWithCode(1), ""); + EXPECT_EXIT(occupy.iweights(1, wk, 1, 0, 1.0, ekb, ef, wg, -1, isk);, ::testing::ExitedWithCode(1), ""); output = testing::internal::GetCapturedStdout(); EXPECT_THAT(output, testing::HasSubstr("It is not a semiconductor or insulator. Please do not set 'smearing_method=fixed', and try other options.")); } diff --git a/source/source_io/CMakeLists.txt b/source/source_io/CMakeLists.txt index f215d68ff6..18cf58759c 100644 --- a/source/source_io/CMakeLists.txt +++ b/source/source_io/CMakeLists.txt @@ -12,6 +12,7 @@ list(APPEND objects module_dos/cal_ldos.cpp module_ml/cal_mlkedf_desc.cpp module_dos/write_dos_pw.cpp + module_energy/band_parallel_output.cpp module_energy/write_bands.cpp module_energy/nscf_fermi_surf.cpp module_energy/write_eig_occ.cpp diff --git a/source/source_io/module_energy/band_parallel_output.cpp b/source/source_io/module_energy/band_parallel_output.cpp new file mode 100644 index 0000000000..2429c7a090 --- /dev/null +++ b/source/source_io/module_energy/band_parallel_output.cpp @@ -0,0 +1,79 @@ +#include "band_parallel_output.h" + +#include "source_base/global_function.h" +#include "source_base/parallel_comm.h" + +#include +#include +#include + +#ifdef __MPI +#include +#endif + +ModuleBase::matrix ModuleIO::gather_band_matrix(const ModuleBase::matrix& local_matrix, const int global_nbands) +{ + if (global_nbands < 0) + { + ModuleBase::WARNING_QUIT("ModuleIO::gather_band_matrix", "global band count cannot be negative"); + } +#ifndef __MPI + if (local_matrix.nc != global_nbands) + { + ModuleBase::WARNING_QUIT("ModuleIO::gather_band_matrix", "distributed bands require an MPI build"); + } + return local_matrix; +#else + int band_groups = 0; + MPI_Comm_size(BP_WORLD, &band_groups); + std::vector band_counts(band_groups); + std::vector row_counts(band_groups); + MPI_Allgather(&local_matrix.nc, 1, MPI_INT, band_counts.data(), 1, MPI_INT, BP_WORLD); + MPI_Allgather(&local_matrix.nr, 1, MPI_INT, row_counts.data(), 1, MPI_INT, BP_WORLD); + + // SDFT without BPCG replicates all deterministic bands across band groups, + // whereas BPCG stores complementary column ranges. All ranks inspect the + // same gathered shapes before choosing a path, preventing rank-divergent + // collectives for uneven band distributions. + const bool rows_match + = std::all_of(row_counts.begin(), row_counts.end(), [&local_matrix](const int rows) { return rows == local_matrix.nr; }); + if (!rows_match) + { + ModuleBase::WARNING_QUIT("ModuleIO::gather_band_matrix", "band groups have inconsistent row counts"); + } + + const bool bands_are_replicated + = std::all_of(band_counts.begin(), band_counts.end(), [global_nbands](const int bands) { return bands == global_nbands; }); + if (bands_are_replicated) + { + return local_matrix; + } + + const int gathered_nbands = std::accumulate(band_counts.begin(), band_counts.end(), 0); + if (gathered_nbands != global_nbands) + { + ModuleBase::WARNING_QUIT("ModuleIO::gather_band_matrix", "local band counts do not match global nbands"); + } + + std::vector band_offsets(band_groups, 0); + for (int group = 1; group < band_groups; ++group) + { + band_offsets[group] = band_offsets[group - 1] + band_counts[group - 1]; + } + + ModuleBase::matrix global_matrix(local_matrix.nr, global_nbands, false); + for (int ik = 0; ik < local_matrix.nr; ++ik) + { + const double* local_row = local_matrix.nc > 0 ? local_matrix.c + ik * local_matrix.nc : nullptr; + MPI_Allgatherv(local_row, + local_matrix.nc, + MPI_DOUBLE, + global_matrix.c + ik * global_nbands, + band_counts.data(), + band_offsets.data(), + MPI_DOUBLE, + BP_WORLD); + } + return global_matrix; +#endif +} diff --git a/source/source_io/module_energy/band_parallel_output.h b/source/source_io/module_energy/band_parallel_output.h new file mode 100644 index 0000000000..7443dcb410 --- /dev/null +++ b/source/source_io/module_energy/band_parallel_output.h @@ -0,0 +1,24 @@ +#ifndef BAND_PARALLEL_OUTPUT_H_ +#define BAND_PARALLEL_OUTPUT_H_ + +#include "source_base/matrix.h" + +namespace ModuleIO +{ + +/** + * @brief Obtain a complete band matrix from replicated or distributed input. + * + * Band groups may either hold identical complete column ranges, as in SDFT + * without BPCG, or complementary contiguous ranges, as in BPCG. An empty + * matrix is valid when the calculation has no deterministic bands. + * + * @param local_matrix Complete or locally owned band columns. + * @param global_nbands Expected global number of bands. + * @return Complete matrix on every rank in BP_WORLD. + */ +ModuleBase::matrix gather_band_matrix(const ModuleBase::matrix& local_matrix, const int global_nbands); + +} // namespace ModuleIO + +#endif // BAND_PARALLEL_OUTPUT_H_ diff --git a/source/source_io/module_energy/write_bands.cpp b/source/source_io/module_energy/write_bands.cpp index f7cdcee530..d63df3bb80 100644 --- a/source/source_io/module_energy/write_bands.cpp +++ b/source/source_io/module_energy/write_bands.cpp @@ -1,4 +1,5 @@ #include "write_bands.h" +#include "band_parallel_output.h" #include "source_base/global_function.h" #include "source_base/global_variable.h" #include "source_base/timer.h" @@ -18,6 +19,8 @@ void ModuleIO::write_bands(const Input_para& inp, // write band information to band.txt if (inp.out_band[0]) { + // Taoni fix out_band under bndpar on 2026-08-21 + const ModuleBase::matrix global_ekb = ModuleIO::gather_band_matrix(ekb, inp.nbands); const int nspin0 = (inp.nspin == 2) ? 2 : 1; for (int is = 0; is < nspin0; is++) { @@ -38,7 +41,7 @@ void ModuleIO::write_bands(const Input_para& inp, const double eshift = 0.0; nscf_bands(is, ss.str(), inp.nbands, eshift, inp.out_band[1], // precision - ekb, kv); + global_ekb, kv); } } } @@ -113,7 +116,7 @@ void ModuleIO::nscf_bands( const int ik_now = ik - kv.para_k.startk_pool[GlobalV::MY_POOL]; //! if present kpoint corresponds the spin of the present one assert( kv.isk[ik_now+is*nks_np] == is ); - if ( GlobalV::RANK_IN_POOL == 0) + if (GlobalV::RANK_IN_POOL == 0 && GlobalV::MY_BNDGROUP == 0) { std::ofstream ofs(eig_file.c_str(), std::ios::app); ofs << FmtCore::format("%4d", ik+1); diff --git a/source/source_io/module_energy/write_eig_occ.cpp b/source/source_io/module_energy/write_eig_occ.cpp index ff5d9ea4cd..ecd1568290 100644 --- a/source/source_io/module_energy/write_eig_occ.cpp +++ b/source/source_io/module_energy/write_eig_occ.cpp @@ -1,5 +1,6 @@ #include "write_eig_occ.h" +#include "band_parallel_output.h" #include "source_io/module_parameter/parameter.h" #include "source_base/global_function.h" #include "source_base/global_variable.h" @@ -17,6 +18,10 @@ void ModuleIO::write_eig_iter(const ModuleBase::matrix &ekb,const ModuleBase::ma GlobalV::ofs_running << "\n PRINT #EIGENVALUES# AND #OCCUPATIONS#" << std::endl; + // Taoni fix bndpar on 2026-08-21 + const ModuleBase::matrix global_ekb = ModuleIO::gather_band_matrix(ekb, PARAM.inp.nbands); + const ModuleBase::matrix global_wg = ModuleIO::gather_band_matrix(wg, PARAM.inp.nbands); + const int nspin = PARAM.inp.nspin; const int nks = kv.get_nks(); const int nkstot = kv.get_nkstot(); @@ -51,11 +56,11 @@ void ModuleIO::write_eig_iter(const ModuleBase::matrix &ekb,const ModuleBase::ma << std::setw(18) << "Eigenvalues(eV)" << std::setw(18) << "Occupations" << std::endl; - for (int ib = 0; ib < ekb.nc; ib++) + for (int ib = 0; ib < global_ekb.nc; ib++) { GlobalV::ofs_running << std::setw(8) << ib + 1 - << std::setw(18) << ekb(ik, ib) * ModuleBase::Ry_to_eV - << std::setw(18) << wg(ik, ib) << std::endl; + << std::setw(18) << global_ekb(ik, ib) * ModuleBase::Ry_to_eV + << std::setw(18) << global_wg(ik, ib) << std::endl; } GlobalV::ofs_running << std::endl; } @@ -118,7 +123,7 @@ void ModuleIO::write_eig_iter(const ModuleBase::matrix &ekb,const ModuleBase::ma // for the current spin channel [is] and pool [ip] // MPI_Send the size of matrix, ik2iktot, ekb and wg to RANK=0 const int send_nks_np = nks_np; - const int send_nbands = ekb.nc; + const int send_nbands = global_ekb.nc; const int is_offset = is * nks_np; int* send_ik2iktot = new int[send_nks_np]; for (int ik = 0; ik < send_nks_np; ++ik) @@ -128,10 +133,10 @@ void ModuleIO::write_eig_iter(const ModuleBase::matrix &ekb,const ModuleBase::ma MPI_Send(&send_nks_np, 1, MPI_INT, 0, 0, MPI_COMM_WORLD); MPI_Send(&send_nbands, 1, MPI_INT, 0, 1, MPI_COMM_WORLD); MPI_Send(send_ik2iktot, send_nks_np, MPI_INT, 0, 2, MPI_COMM_WORLD); - MPI_Send(ekb.c + is_offset * send_nbands, + MPI_Send(global_ekb.c + is_offset * send_nbands, send_nks_np * send_nbands, MPI_DOUBLE, 0, 3, MPI_COMM_WORLD); - MPI_Send(wg.c + is_offset * send_nbands, + MPI_Send(global_wg.c + is_offset * send_nbands, send_nks_np * send_nbands, MPI_DOUBLE, 0, 4, MPI_COMM_WORLD); @@ -173,6 +178,9 @@ void ModuleIO::write_eig_file(const ModuleBase::matrix &ekb, */ const int nspin = PARAM.inp.nspin; + // Taoni fix bndpar on 2026-08-21 + const ModuleBase::matrix global_ekb = ModuleIO::gather_band_matrix(ekb, PARAM.inp.nbands); + const ModuleBase::matrix global_wg = ModuleIO::gather_band_matrix(wg, PARAM.inp.nbands); const int nks = kv.get_nks(); const int nkstot = kv.get_nkstot(); @@ -244,7 +252,7 @@ void ModuleIO::write_eig_file(const ModuleBase::matrix &ekb, MPI_Barrier(MPI_COMM_WORLD); #endif - bool ip_flag = PARAM.inp.out_alllog || (GlobalV::RANK_IN_POOL == 0 && GlobalV::MY_BNDGROUP == 0); + const bool ip_flag = GlobalV::RANK_IN_POOL == 0 && GlobalV::MY_BNDGROUP == 0; if (GlobalV::MY_POOL == ip && ip_flag) { @@ -263,14 +271,14 @@ void ModuleIO::write_eig_file(const ModuleBase::matrix &ekb, ofs_eig << std::setprecision(16); ofs_eig << std::setiosflags(std::ios::showpoint); - for (int ib = 0; ib < ekb.nc; ib++) + for (int ib = 0; ib < global_ekb.nc; ib++) { - double occupation = wg(ik, ib); + double occupation = global_wg(ik, ib); if (std::abs(occupation) < 1.0e-15) { occupation = 0.0; } - ofs_eig << " " << ib + 1 << " " << ekb(ik, ib) * ModuleBase::Ry_to_eV + ofs_eig << " " << ib + 1 << " " << global_ekb(ik, ib) * ModuleBase::Ry_to_eV << " " << occupation << std::endl; } ofs_eig << std::endl; diff --git a/source/source_io/module_parameter/read_inp_sys.cpp b/source/source_io/module_parameter/read_inp_sys.cpp index 1d207658e4..3c9bff3935 100644 --- a/source/source_io/module_parameter/read_inp_sys.cpp +++ b/source/source_io/module_parameter/read_inp_sys.cpp @@ -300,9 +300,12 @@ void ReadInput::item_system() "will be distributed among"; item.category = "System variables"; item.type = "Integer"; - item.description = "Divide all processors into kpar groups, and k points will be distributed among each group. " - "The value taken should be less than or equal to the number of k points as well as the number of MPI processes."; + item.description = R"(Controls k-point parallelism. The value must be positive and should not exceed either the number of k-points or the number of MPI processes. +* For PW calculations, divide all MPI processes into persistent k-point pools. Each pool stores and processes a subset of the k-points. +* For LCAO calculations with lapack, genelpa, elpa, or scalapack_gvx, divide the diagonalization work into temporary k-point pools. After diagonalization, the eigenvalues and distributed wavefunctions are restored for all k-points before occupations, density matrices, and output are evaluated. +* Multi-process LCAO cusolver uses its own active-GPU distribution and does not use this value to define its k-point layout. Other LCAO eigensolvers do not use the temporary k-point-pool implementation.)"; item.default_value = "1"; + item.unit = ""; read_sync_int(input.kpar); item.reset_value = [](const Input_Item& item, Parameter& para) { #ifdef __LCAO @@ -326,25 +329,24 @@ void ReadInput::item_system() } #endif }; - item.check_value = [](const Input_Item& item, const Parameter& para) { - if (para.input.basis_type == "lcao" && para.input.kpar > 1) - { - ModuleBase::WARNING("ReadInput", "kpar > 1 has not been supported for lcao calculation."); - } - }; this->add_item(item); add_int_bcast(sys.kpar_lcao); } { Input_Item item("bndpar"); - item.annotation = "devide all processors into bndpar groups and bands " - "will be distributed among each group"; + item.annotation = "divide each k-point pool into band-parallel groups"; item.category = "System variables"; item.type = "Integer"; - item.description = "Divide all processors into bndpar groups for SDFT or the BPCG solver. bndpar must be " - "positive, no greater than the number of MPI processes, and kpar * bndpar must divide " - "the number of MPI processes exactly."; + item.description = R"(Controls band-group parallelism for PW SDFT and PW KSDFT calculations using the BPCG eigensolver. +* Within each k-point pool, divide the MPI processes into bndpar band groups. Each group contains NPROC / (kpar * bndpar) processes when bndpar is greater than 1. +* With BPCG, distribute contiguous ranges of global Kohn-Sham bands among the band groups. nbands does not need to be divisible by bndpar, but bndpar cannot exceed a positive nbands. Groups with lower indices receive one additional band when necessary. +* In SDFT, distribute stochastic orbitals among the band groups. When the deterministic Kohn-Sham eigensolver is not BPCG, band group 0 calculates the deterministic orbitals and broadcasts them to the other groups. +* bndpar must be positive and no greater than the number of MPI processes. When bndpar is greater than 1, kpar * bndpar must divide the number of MPI processes exactly. +[NOTE] For PW calculations on GPU, if the input kpar * bndpar differs from the number of MPI processes, ABACUS automatically sets the effective kpar to NPROC / bndpar.)"; item.default_value = "1"; + item.unit = ""; + item.set_availability("(basis_type==pw and esolver_type==sdft) or " + "(basis_type==pw and esolver_type==ksdft and ks_solver==bpcg)"); read_sync_int(input.bndpar); item.check_value = [](const Input_Item& item, const Parameter& para) { if (para.input.bndpar <= 0) diff --git a/source/source_io/module_parameter/read_inp_tddft.cpp b/source/source_io/module_parameter/read_inp_tddft.cpp index 695d9d36ff..9f98bb4364 100644 --- a/source/source_io/module_parameter/read_inp_tddft.cpp +++ b/source/source_io/module_parameter/read_inp_tddft.cpp @@ -737,9 +737,9 @@ In the velocity and hybrid gauges, ABACUS obtains the vector potential actually item.annotation = "set occupation"; item.category = "RT-TDDFT: Real-Time Time-Dependent Density Functional Theory"; item.type = "String"; - item.description = R"(Fixed occupation weights used when ocp is true. Values are assigned band by band for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x. + item.description = R"(Fixed occupation weights used when ocp is true. Values are assigned in band order for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x. * Example: 1 10*1 0 1 expands to 13 values, with the 12th value equal to 0 and all other values equal to 1. -* After expansion, the array length must equal nks * nbands. +* After expansion, provide one block of nbands values for each k-point. If nspin is 2, provide all k-point blocks for spin up followed by all k-point blocks for spin down; otherwise, provide one block per k-point. * The sum of all weights must equal nelec; otherwise the calculation terminates with an error.)"; item.default_value = "None"; item.unit = ""; diff --git a/source/source_io/test/CMakeLists.txt b/source/source_io/test/CMakeLists.txt index c0bd352830..7ab92b6ae4 100644 --- a/source/source_io/test/CMakeLists.txt +++ b/source/source_io/test/CMakeLists.txt @@ -63,7 +63,7 @@ add_test( AddTest( TARGET MODULE_IO_write_eig_occ_test LIBS parameter base device symmetry - SOURCES write_eig_occ_test.cpp ../module_energy/write_eig_occ.cpp ../../source_cell/parallel_kpoints.cpp ../../source_cell/klist.cpp ../../source_cell/k_vector_utils.cpp + SOURCES write_eig_occ_test.cpp ../module_energy/band_parallel_output.cpp ../module_energy/write_eig_occ.cpp ../../source_cell/parallel_kpoints.cpp ../../source_cell/klist.cpp ../../source_cell/k_vector_utils.cpp ../../source_cell/cif_io.cpp ) diff --git a/source/source_io/test_serial/CMakeLists.txt b/source/source_io/test_serial/CMakeLists.txt index 09e509dff2..db79fbd41a 100644 --- a/source/source_io/test_serial/CMakeLists.txt +++ b/source/source_io/test_serial/CMakeLists.txt @@ -56,7 +56,7 @@ AddTest( AddTest( TARGET MODULE_IO_write_bands LIBS parameter base device - SOURCES write_bands_test.cpp ../module_energy/write_bands.cpp + SOURCES write_bands_test.cpp ../module_energy/band_parallel_output.cpp ../module_energy/write_bands.cpp ) AddTest(