refactor(klist): extract MPI k-point pack/unpack into KListIO

Move the this-free flatten/scatter loops in K_Vectors::mpi_k into
KListIO::pack_kpts (rank-0 -> contiguous broadcast buffers) and
KListIO::unpack_kpts (broadcast buffers -> this pool's k-point slice).
mpi_k keeps the broadcast calls and the renew/distribute orchestration.

pack_kpts sizes the kvec_c_full copy from the vector length so the full
grid is handled independently of nkstot.

Verified: cmake --build build_max_para_test -j30 (exit 0, __MPI on) and
ctest -R "MODULE_CELL_(klist|reciprocal_grid|qlist)" -> 5/5 (covers the
para1/para4 mpi_k paths).

Author: abacus_fixer

source/source_cell/klist.cpp

@@ -727,20 +727,17 @@ void K_Vectors::mpi_k(std::ofstream& ofs_running)
     // collect and process in rank 0
     if (GlobalV::MY_RANK == 0)
     {
-        for (int ik = 0; ik < this->nkstot; ik++)
-        {
-            isk_aux[ik] = this->isk[ik];
-            wk_aux[ik] = this->wk[ik];
-            kvec_c_aux[3 * ik] = this->kvec_c[ik].x;
-            kvec_c_aux[3 * ik + 1] = this->kvec_c[ik].y;
-            kvec_c_aux[3 * ik + 2] = this->kvec_c[ik].z;
-            kvec_d_aux[3 * ik] = this->kvec_d[ik].x;
-            kvec_d_aux[3 * ik + 1] = this->kvec_d[ik].y;
-            kvec_d_aux[3 * ik + 2] = this->kvec_d[ik].z;
-            kvec_c_full_aux[3 * ik] = this->kvec_c_full[ik].x;
-            kvec_c_full_aux[3 * ik + 1] = this->kvec_c_full[ik].y;
-            kvec_c_full_aux[3 * ik + 2] = this->kvec_c_full[ik].z;
-        }
+        KListIO::pack_kpts(this->isk,
+                           this->wk,
+                           this->kvec_c,
+                           this->kvec_d,
+                           this->kvec_c_full,
+                           this->nkstot,
+                           isk_aux,
+                           wk_aux,
+                           kvec_c_aux,
+                           kvec_d_aux,
+                           kvec_c_full_aux);
     }
 
     // broadcast k point data to all processors
@@ -755,24 +752,18 @@ void K_Vectors::mpi_k(std::ofstream& ofs_running)
     this->renew(this->nks * this->spin_mult);
 
     // distribute
-    int k_index = 0;
-
-    for (int i = 0; i < this->nks; i++)
-    {
-        // 3 is because each k point has three value:kx, ky, kz
-        k_index = i + this->para_k.startk_pool[GlobalV::MY_POOL];
-        this->kvec_c[i].x = kvec_c_aux[k_index * 3];
-        this->kvec_c[i].y = kvec_c_aux[k_index * 3 + 1];
-        this->kvec_c[i].z = kvec_c_aux[k_index * 3 + 2];
-        this->kvec_d[i].x = kvec_d_aux[k_index * 3];
-        this->kvec_d[i].y = kvec_d_aux[k_index * 3 + 1];
-        this->kvec_d[i].z = kvec_d_aux[k_index * 3 + 2];
-        this->kvec_c_full[i].x = kvec_c_full_aux[k_index * 3];
-        this->kvec_c_full[i].y = kvec_c_full_aux[k_index * 3 + 1];
-        this->kvec_c_full[i].z = kvec_c_full_aux[k_index * 3 + 2];
-        this->wk[i] = wk_aux[k_index];
-        this->isk[i] = isk_aux[k_index];
-    }
+    KListIO::unpack_kpts(isk_aux,
+                         wk_aux,
+                         kvec_c_aux,
+                         kvec_d_aux,
+                         kvec_c_full_aux,
+                         this->nks,
+                         this->para_k.startk_pool[GlobalV::MY_POOL],
+                         this->isk,
+                         this->wk,
+                         this->kvec_c,
+                         this->kvec_d,
+                         this->kvec_c_full);
 
 #ifdef __EXX
     if (ModuleSymmetry::Symmetry::symm_flag == 1)

source/source_cell/klist_io.cpp

@@ -178,4 +178,67 @@ void build_kstars(const std::vector<ModuleBase::Vector3<double>>& kvec_d,
     }
 }
 
+void pack_kpts(const std::vector<int>& isk,
+               const std::vector<double>& wk,
+               const std::vector<ModuleBase::Vector3<double>>& kvec_c,
+               const std::vector<ModuleBase::Vector3<double>>& kvec_d,
+               const std::vector<ModuleBase::Vector3<double>>& kvec_c_full,
+               const int nkstot,
+               std::vector<int>& isk_aux,
+               std::vector<double>& wk_aux,
+               std::vector<double>& kvec_c_aux,
+               std::vector<double>& kvec_d_aux,
+               std::vector<double>& kvec_c_full_aux)
+{
+    for (int ik = 0; ik < nkstot; ik++)
+    {
+        isk_aux[ik] = isk[ik];
+        wk_aux[ik] = wk[ik];
+        kvec_c_aux[3 * ik] = kvec_c[ik].x;
+        kvec_c_aux[3 * ik + 1] = kvec_c[ik].y;
+        kvec_c_aux[3 * ik + 2] = kvec_c[ik].z;
+        kvec_d_aux[3 * ik] = kvec_d[ik].x;
+        kvec_d_aux[3 * ik + 1] = kvec_d[ik].y;
+        kvec_d_aux[3 * ik + 2] = kvec_d[ik].z;
+    }
+    const int nkstot_full = static_cast<int>(kvec_c_full.size());
+    for (int ik = 0; ik < nkstot_full; ik++)
+    {
+        kvec_c_full_aux[3 * ik] = kvec_c_full[ik].x;
+        kvec_c_full_aux[3 * ik + 1] = kvec_c_full[ik].y;
+        kvec_c_full_aux[3 * ik + 2] = kvec_c_full[ik].z;
+    }
+}
+
+void unpack_kpts(const std::vector<int>& isk_aux,
+                 const std::vector<double>& wk_aux,
+                 const std::vector<double>& kvec_c_aux,
+                 const std::vector<double>& kvec_d_aux,
+                 const std::vector<double>& kvec_c_full_aux,
+                 const int nks,
+                 const int startk,
+                 std::vector<int>& isk,
+                 std::vector<double>& wk,
+                 std::vector<ModuleBase::Vector3<double>>& kvec_c,
+                 std::vector<ModuleBase::Vector3<double>>& kvec_d,
+                 std::vector<ModuleBase::Vector3<double>>& kvec_c_full)
+{
+    for (int i = 0; i < nks; i++)
+    {
+        // 3 is because each k point has three value:kx, ky, kz
+        const int k_index = i + startk;
+        kvec_c[i].x = kvec_c_aux[k_index * 3];
+        kvec_c[i].y = kvec_c_aux[k_index * 3 + 1];
+        kvec_c[i].z = kvec_c_aux[k_index * 3 + 2];
+        kvec_d[i].x = kvec_d_aux[k_index * 3];
+        kvec_d[i].y = kvec_d_aux[k_index * 3 + 1];
+        kvec_d[i].z = kvec_d_aux[k_index * 3 + 2];
+        kvec_c_full[i].x = kvec_c_full_aux[k_index * 3];
+        kvec_c_full[i].y = kvec_c_full_aux[k_index * 3 + 1];
+        kvec_c_full[i].z = kvec_c_full_aux[k_index * 3 + 2];
+        wk[i] = wk_aux[k_index];
+        isk[i] = isk_aux[k_index];
+    }
+}
+
 } // namespace KListIO

source/source_cell/klist_io.h

@@ -51,6 +51,35 @@ void build_kstars(const std::vector<ModuleBase::Vector3<double>>& kvec_d,
                   double epsilon,
                   const std::function<bool(double, double)>& equal,
                   std::vector<std::map<int, ModuleBase::Vector3<double>>>& kstars);
+
+/// Flatten k-point arrays into contiguous MPI buffers (x,y,z interleaved).
+/// this-free; used on rank 0 before broadcasting in K_Vectors::mpi_k.
+void pack_kpts(const std::vector<int>& isk,
+               const std::vector<double>& wk,
+               const std::vector<ModuleBase::Vector3<double>>& kvec_c,
+               const std::vector<ModuleBase::Vector3<double>>& kvec_d,
+               const std::vector<ModuleBase::Vector3<double>>& kvec_c_full,
+               int nkstot,
+               std::vector<int>& isk_aux,
+               std::vector<double>& wk_aux,
+               std::vector<double>& kvec_c_aux,
+               std::vector<double>& kvec_d_aux,
+               std::vector<double>& kvec_c_full_aux);
+
+/// Scatter the broadcast buffers into this pool's k-point slice, starting at
+/// global index `startk`. this-free; mirrors pack_kpts after the broadcast.
+void unpack_kpts(const std::vector<int>& isk_aux,
+                 const std::vector<double>& wk_aux,
+                 const std::vector<double>& kvec_c_aux,
+                 const std::vector<double>& kvec_d_aux,
+                 const std::vector<double>& kvec_c_full_aux,
+                 int nks,
+                 int startk,
+                 std::vector<int>& isk,
+                 std::vector<double>& wk,
+                 std::vector<ModuleBase::Vector3<double>>& kvec_c,
+                 std::vector<ModuleBase::Vector3<double>>& kvec_d,
+                 std::vector<ModuleBase::Vector3<double>>& kvec_c_full);
 } // namespace KListIO
 
 #endif // KLIST_IO_H