Fix and optimize get_S sparse output

Author: sgyang345

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<elecstate::ElecStateLCAO<std::complex<double>>*>(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");
 }

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 <mpi.h>
+#endif
+
+#include <algorithm>
+#include <cstdio>
+#include <fstream>
+#include <iomanip>
+#include <numeric>
+#include <sstream>
+#include <vector>
+
+namespace
+{
+template <typename T>
+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<double>& data)
+{
+    ofs << " (" << std::fixed << std::scientific << std::setprecision(8) << data.real() << ","
+        << std::fixed << std::scientific << std::setprecision(8) << data.imag() << ")";
+}
+
+template <typename Tdata>
+void save_gets_sparse_fast(
+    const std::map<Abfs::Vector3_Order<int>, std::map<size_t, std::map<size_t, Tdata>>>& smat,
+    const std::set<Abfs::Vector3_Order<int>>& 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<int> 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<int> 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<GetSEntry<Tdata>> local_entries;
+            auto iter = smat.find(R_coor);
+            if (iter != smat.end())
+            {
+                auto row_iter = iter->second.lower_bound(static_cast<size_t>(row_begin));
+                while (row_iter != iter->second.end() && row_iter->first < static_cast<size_t>(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<GetSEntry<Tdata>> entries;
+#ifdef __MPI
+            const int local_bytes = static_cast<int>(local_entries.size() * sizeof(GetSEntry<Tdata>));
+            std::vector<int> recv_counts(nprocs, 0);
+            MPI_Gather(&local_bytes, 1, MPI_INT, recv_counts.data(), 1, MPI_INT, 0, MPI_COMM_WORLD);
+
+            std::vector<int> 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<Tdata>));
+            }
+
+            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<GetSEntry<Tdata>> merged_entries;
+                merged_entries.reserve(entries.size());
+                for (size_t i = 0; i < entries.size();)
+                {
+                    GetSEntry<Tdata> 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<size_t>(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);

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<size_t, std::map<size_t, Tdata>> 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<std::complex<double>>(
     const Parallel_Orbitals&,
     const std::string&,
     const int&,
-    const bool&);
+    const bool&);