diff --git a/source/module_esolver/esolver_gets.cpp b/source/module_esolver/esolver_gets.cpp index 842b0340e1a..bbffde9d165 100644 --- a/source/module_esolver/esolver_gets.cpp +++ b/source/module_esolver/esolver_gets.cpp @@ -72,11 +72,8 @@ void ESolver_GetS::before_all_runners(UnitCell& ucell, const Input_para& inp) two_center_bundle_, orb_); - // 4) initialize the density matrix - // DensityMatrix is allocated here, DMK is also initialized here - // DMR is not initialized here, it will be constructed in each before_scf - dynamic_cast>*>(this->pelec) - ->init_DM(&this->kv, &(this->pv), inp.nspin); + // get_S only builds and writes overlap matrices. Allocating DMK here is + // unused and can dominate memory for large LCAO/SOC systems. ModuleBase::timer::tick("ESolver_GetS", "before_all_runners"); } diff --git a/source/module_io/write_HS_R.cpp b/source/module_io/write_HS_R.cpp index 0ed7cfa1490..82d0abec031 100644 --- a/source/module_io/write_HS_R.cpp +++ b/source/module_io/write_HS_R.cpp @@ -1,5 +1,6 @@ #include "write_HS_R.h" +#include "module_base/parallel_reduce.h" #include "module_parameter/parameter.h" #include "module_base/timer.h" #include "module_hamilt_lcao/hamilt_lcaodft/LCAO_HS_arrays.hpp" @@ -8,6 +9,248 @@ #include "module_hamilt_lcao/hamilt_lcaodft/spar_st.h" #include "write_HS_sparse.h" +#ifdef __MPI +#include +#endif + +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +template +struct GetSEntry +{ + size_t row; + size_t col; + T value; +}; + +void write_gets_data(std::ofstream& ofs, const double& data) +{ + ofs << " " << std::fixed << std::scientific << std::setprecision(8) << data; +} + +void write_gets_data(std::ofstream& ofs, const std::complex& data) +{ + ofs << " (" << std::fixed << std::scientific << std::setprecision(8) << data.real() << "," + << std::fixed << std::scientific << std::setprecision(8) << data.imag() << ")"; +} + +template +void save_gets_sparse_fast( + const std::map, std::map>>& smat, + const std::set>& all_R_coor, + const double& sparse_thr, + const std::string& filename, + const std::string& label, + const int& istep) +{ + const int nlocal = PARAM.globalv.nlocal; + const int total_R_num = all_R_coor.size(); + + std::vector nonzero_num(total_R_num, 0); + int count = 0; + for (const auto& R_coor : all_R_coor) + { + auto iter = smat.find(R_coor); + if (iter != smat.end()) + { + for (const auto& row_loop : iter->second) + { + for (const auto& value : row_loop.second) + { + if (std::abs(value.second) > sparse_thr) + { + ++nonzero_num[count]; + } + } + } + } + ++count; + } +#ifdef __MPI + Parallel_Reduce::reduce_all(nonzero_num.data(), total_R_num); +#endif + + int output_R_number = 0; + for (const int nnz : nonzero_num) + { + if (nnz != 0) + { + ++output_R_number; + } + } + + int myrank = 0; + int nprocs = 1; +#ifdef __MPI + MPI_Comm_rank(MPI_COMM_WORLD, &myrank); + MPI_Comm_size(MPI_COMM_WORLD, &nprocs); +#endif + + std::ofstream ofs; + if (myrank == 0) + { + ofs.open(filename.c_str()); + ofs << "STEP: " << std::max(istep, 0) << std::endl; + ofs << "Matrix Dimension of " + label + "(R): " << nlocal << std::endl; + ofs << "Matrix number of " + label + "(R): " << output_R_number << std::endl; + } + + count = 0; + for (const auto& R_coor : all_R_coor) + { + if (nonzero_num[count] == 0) + { + ++count; + continue; + } + + std::ofstream ofs_indices; + std::ifstream ifs_indices; + std::stringstream indices_filename; + indices_filename << PARAM.globalv.global_out_dir << std::to_string(myrank) << "temp_gets_sparse_indices.dat"; + + std::vector indptr; + if (myrank == 0) + { + ofs << R_coor.x << " " << R_coor.y << " " << R_coor.z << " " + << nonzero_num[count] << std::endl; + ofs_indices.open(indices_filename.str().c_str()); + indptr.reserve(nlocal + 1); + indptr.push_back(0); + } + + const int row_block_size = 4096; + for (int row_begin = 0; row_begin < nlocal; row_begin += row_block_size) + { + const int row_end = std::min(row_begin + row_block_size, nlocal); + + std::vector> local_entries; + auto iter = smat.find(R_coor); + if (iter != smat.end()) + { + auto row_iter = iter->second.lower_bound(static_cast(row_begin)); + while (row_iter != iter->second.end() && row_iter->first < static_cast(row_end)) + { + for (const auto& value : row_iter->second) + { + if (std::abs(value.second) > sparse_thr) + { + local_entries.push_back({row_iter->first, value.first, value.second}); + } + } + ++row_iter; + } + } + + std::vector> entries; +#ifdef __MPI + const int local_bytes = static_cast(local_entries.size() * sizeof(GetSEntry)); + std::vector recv_counts(nprocs, 0); + MPI_Gather(&local_bytes, 1, MPI_INT, recv_counts.data(), 1, MPI_INT, 0, MPI_COMM_WORLD); + + std::vector displs; + int total_bytes = 0; + if (myrank == 0) + { + displs.resize(nprocs, 0); + for (int ip = 1; ip < nprocs; ++ip) + { + displs[ip] = displs[ip - 1] + recv_counts[ip - 1]; + } + total_bytes = std::accumulate(recv_counts.begin(), recv_counts.end(), 0); + entries.resize(total_bytes / sizeof(GetSEntry)); + } + + MPI_Gatherv(local_entries.empty() ? nullptr : local_entries.data(), + local_bytes, + MPI_BYTE, + entries.empty() ? nullptr : entries.data(), + recv_counts.empty() ? nullptr : recv_counts.data(), + displs.empty() ? nullptr : displs.data(), + MPI_BYTE, + 0, + MPI_COMM_WORLD); +#else + entries.swap(local_entries); +#endif + + if (myrank == 0) + { + std::sort(entries.begin(), entries.end(), [](const auto& lhs, const auto& rhs) { + return lhs.row == rhs.row ? lhs.col < rhs.col : lhs.row < rhs.row; + }); + + std::vector> merged_entries; + merged_entries.reserve(entries.size()); + for (size_t i = 0; i < entries.size();) + { + GetSEntry merged = entries[i]; + ++i; + while (i < entries.size() && entries[i].row == merged.row && entries[i].col == merged.col) + { + merged.value += entries[i].value; + ++i; + } + if (std::abs(merged.value) > sparse_thr) + { + merged_entries.push_back(merged); + } + } + + size_t entry_index = 0; + for (int row = row_begin; row < row_end; ++row) + { + int row_nnz = 0; + while (entry_index < merged_entries.size() + && merged_entries[entry_index].row == static_cast(row)) + { + write_gets_data(ofs, merged_entries[entry_index].value); + ofs_indices << " " << merged_entries[entry_index].col; + ++row_nnz; + ++entry_index; + } + indptr.push_back(indptr.back() + row_nnz); + } + } + } + + if (myrank == 0) + { + ofs << std::endl; + + ofs_indices << std::endl; + ofs_indices.close(); + ifs_indices.open(indices_filename.str().c_str()); + ofs << ifs_indices.rdbuf(); + ifs_indices.close(); + + for (const auto& pointer : indptr) + { + ofs << " " << pointer; + } + ofs << std::endl; + + std::remove(indices_filename.str().c_str()); + } + + ++count; + } + + if (myrank == 0) + { + ofs.close(); + } +} +} // namespace + // if 'binary=true', output binary file. // The 'sparse_thr' is the accuracy of the sparse matrix. // If the absolute value of the matrix element is less than or equal to the @@ -178,25 +421,49 @@ void ModuleIO::output_SR(Parallel_Orbitals& pv, if (PARAM.inp.nspin == 4) { - ModuleIO::save_sparse(HS_Arrays.SR_soc_sparse, - HS_Arrays.all_R_coor, - sparse_thr, - binary, - SR_filename, - pv, - "S", - istep); + if (PARAM.inp.calculation == "get_S" && !binary) + { + save_gets_sparse_fast(HS_Arrays.SR_soc_sparse, + HS_Arrays.all_R_coor, + sparse_thr, + SR_filename, + "S", + istep); + } + else + { + ModuleIO::save_sparse(HS_Arrays.SR_soc_sparse, + HS_Arrays.all_R_coor, + sparse_thr, + binary, + SR_filename, + pv, + "S", + istep); + } } else { - ModuleIO::save_sparse(HS_Arrays.SR_sparse, - HS_Arrays.all_R_coor, - sparse_thr, - binary, - SR_filename, - pv, - "S", - istep); + if (PARAM.inp.calculation == "get_S" && !binary) + { + save_gets_sparse_fast(HS_Arrays.SR_sparse, + HS_Arrays.all_R_coor, + sparse_thr, + SR_filename, + "S", + istep); + } + else + { + ModuleIO::save_sparse(HS_Arrays.SR_sparse, + HS_Arrays.all_R_coor, + sparse_thr, + binary, + SR_filename, + pv, + "S", + istep); + } } sparse_format::destroy_HS_R_sparse(HS_Arrays); diff --git a/source/module_io/write_HS_sparse.cpp b/source/module_io/write_HS_sparse.cpp index 026d0ab8363..c108d618ccc 100644 --- a/source/module_io/write_HS_sparse.cpp +++ b/source/module_io/write_HS_sparse.cpp @@ -765,7 +765,13 @@ void ModuleIO::save_sparse( } } - output_single_R(ofs, smat.at(R_coor), sparse_thr, binary, pv, reduce); + // nonzero_num is reduced across all MPI ranks, but a given rank may + // not own local sparse data for this R block. It still has to enter + // output_single_R so the row-wise reductions remain matched. + const std::map> empty_smat; + auto iter = smat.find(R_coor); + const auto& local_smat = (iter == smat.end()) ? empty_smat : iter->second; + output_single_R(ofs, local_smat, sparse_thr, binary, pv, reduce); ++count; } if (!reduce || GlobalV::DRANK == 0) { @@ -797,4 +803,4 @@ template void ModuleIO::save_sparse>( const Parallel_Orbitals&, const std::string&, const int&, - const bool&); \ No newline at end of file + const bool&);