Fix: bug when using scan with kpar > 1

source/module_base/parallel_global.cpp

@@ -19,6 +19,7 @@
 
 #if defined __MPI
 MPI_Comm POOL_WORLD;
+MPI_Comm INTER_POOL;
 MPI_Comm STO_WORLD;
 MPI_Comm PARAPW_WORLD; // qianrui add it for sto-dft 2021-4-14
 MPI_Comm GRID_WORLD; // mohan add 2012-01-13z
@@ -261,6 +262,7 @@ void Parallel_Global::read_mpi_parameters(int argc,char **argv)
 void Parallel_Global::finalize_mpi()
 {
 	MPI_Comm_free(&POOL_WORLD);
+    MPI_Comm_free(&INTER_POOL);
 	MPI_Comm_free(&STO_WORLD);
 	MPI_Comm_free(&PARAPW_WORLD);
 	MPI_Comm_free(&GRID_WORLD);
@@ -367,7 +369,13 @@ void Parallel_Global::divide_pools(void)
     // Note: The color must be non-negative or MPI_UNDEFINED.
     //========================================================
     MPI_Comm_split(STO_WORLD,GlobalV::MY_POOL,key,&POOL_WORLD);
-	int color = GlobalV::MY_RANK % GlobalV::NPROC_IN_STOGROUP;
+
+    if (GlobalV::NPROC_IN_STOGROUP % GlobalV::KPAR == 0)
+    {
+        MPI_Comm_split(STO_WORLD, GlobalV::RANK_IN_POOL, key, &INTER_POOL);
+    }
+
+    int color = GlobalV::MY_RANK % GlobalV::NPROC_IN_STOGROUP;
 	MPI_Comm_split(MPI_COMM_WORLD, color, key, &PARAPW_WORLD);
 
     return;

source/module_base/parallel_global.h

@@ -11,6 +11,7 @@
 #ifdef __MPI
 #include "mpi.h"
 extern MPI_Comm POOL_WORLD;
+extern MPI_Comm INTER_POOL; //communicator among different pools
 extern MPI_Comm STO_WORLD;
 extern MPI_Comm PARAPW_WORLD;
 extern MPI_Comm GRID_WORLD; //mohan add 2012-01-13

source/module_basis/module_pw/test_serial/CMakeLists.txt

@@ -5,45 +5,30 @@ remove_definitions(-D__UT_USE_CUDA)
 remove_definitions(-D__ROCM)
 remove_definitions(-D__UT_USE_ROCM)
 remove_definitions(-D__DEEPKS)
+
+add_library(
+    planewave_serial
+    OBJECT
+    ../fft.cpp
+    ../pw_basis.cpp
+    ../pw_basis_k.cpp
+    ../pw_distributeg.cpp
+    ../pw_distributeg_method1.cpp
+    ../pw_distributeg_method2.cpp
+    ../pw_distributer.cpp
+    ../pw_init.cpp
+    ../pw_transform.cpp
+    ../pw_transform_k.cpp
+)
+
 AddTest(
   TARGET basis_pw_serial
-  LIBS ${math_libs}
-  SOURCES pw_basis_test.cpp ../pw_basis.cpp ../pw_init.cpp
-  		../pw_distributeg.cpp
-  		../pw_distributer.cpp
-		../pw_distributeg_method1.cpp
-		../pw_distributeg_method2.cpp
-  		../fft.cpp
-		../../../module_base/blas_connector.h
-		../../../module_base/matrix.cpp
-		../../../module_base/matrix3.cpp
-		../../../module_base/timer.cpp
-		../../../module_base/mymath.cpp
-		../../../module_base/tool_quit.cpp
-		../../../module_base/memory.cpp
-		../../../module_base/global_variable.cpp
-		../../../module_base/global_function.cpp
-		../../../module_base/global_file.cpp
+  LIBS ${math_libs} planewave_serial device base
+  SOURCES pw_basis_test.cpp
 )
 
 AddTest(
   TARGET basis_pw_k_serial
-  LIBS ${math_libs}
-  SOURCES pw_basis_k_test.cpp ../pw_basis_k.cpp ../pw_basis.cpp ../pw_init.cpp
-  		../pw_distributeg.cpp
-  		../pw_distributer.cpp
-		../pw_distributeg_method1.cpp
-		../pw_distributeg_method2.cpp
-  		../fft.cpp
-		../../../module_psi/kernels/memory_op.cpp
-		../../../module_base/blas_connector.h
-		../../../module_base/matrix.cpp
-		../../../module_base/matrix3.cpp
-		../../../module_base/timer.cpp
-		../../../module_base/mymath.cpp
-		../../../module_base/tool_quit.cpp
-		../../../module_base/memory.cpp
-		../../../module_base/global_variable.cpp
-		../../../module_base/global_function.cpp
-		../../../module_base/global_file.cpp
+  LIBS ${math_libs} planewave_serial device base
+  SOURCES pw_basis_k_test.cpp
 )

source/module_elecstate/CMakeLists.txt

@@ -48,4 +48,7 @@ endif()
 
 if(BUILD_TESTING)
   add_subdirectory(test)
+  if(ENABLE_MPI)
+    add_subdirectory(test_mpi)
+  endif()
 endif()

source/module_elecstate/elecstate.cpp

@@ -210,7 +210,7 @@ void ElecState::init_scf(const int istep, const ModuleBase::ComplexMatrix& struc
     //--------------------------------------------------------------------
     if (istep == 0)
     {
-        this->charge->init_rho(this->eferm, strucfac);
+        this->charge->init_rho(this->eferm, strucfac, this->bigpw->nbz, this->bigpw->bz);
     }
 
     // renormalize the charge density

source/module_elecstate/elecstate_pw.cpp

@@ -112,7 +112,7 @@ template<typename FPTYPE, typename Device>
 void ElecStatePW<FPTYPE, Device>::parallelK()
 {
 #ifdef __MPI
-    this->charge->rho_mpi(this->bigpw->nbz, this->bigpw->bz);
+    this->charge->rho_mpi();
     if(GlobalV::ESOLVER_TYPE == "sdft") //qinarui add it 2021-7-21
 	{
         this->f_en.eband /= GlobalV::NPROC_IN_POOL;

source/module_elecstate/module_charge/charge.cpp

@@ -41,6 +41,10 @@ Charge::Charge()
 Charge::~Charge()
 {
     this->destroy();
+#ifdef __MPI
+    delete[] rec;
+    delete[] dis;
+#endif
 }
 
 void Charge::set_rhopw(ModulePW::PW_Basis* rhopw_in)

source/module_elecstate/module_charge/charge.h

@@ -50,8 +50,16 @@ class Charge
 
     void set_rhopw(ModulePW::PW_Basis* rhopw_in);
 
-    void init_rho(elecstate::efermi& eferm_iout, const ModuleBase::ComplexMatrix& strucFac);
-    // mohan update 2021-02-20
+    /**
+     * @brief Init charge density from file or atomic pseudo-wave-functions
+     * 
+     * @param eferm_iout fermi energy to be initialized
+     * @param strucFac [in] structure factor 
+     * @param nbz [in] number of big grids in z direction
+     * @param bz [in] number of small grids in big grids for z dirction
+     */
+    void init_rho(elecstate::efermi& eferm_iout, const ModuleBase::ComplexMatrix& strucFac, const int& nbz, const int& bz);
+
     void allocate(const int &nspin_in);
 
     void atomic_rho(const int spin_number_need,
@@ -82,8 +90,28 @@ class Charge
     void init_final_scf(); //LiuXh add 20180619
 
 	public:
-
-    void rho_mpi(const int& nbz, const int& bz);
+    /**
+     * @brief init some arrays for mpi_inter_pools, rho_mpi
+     * 
+     * @param nbz number of bigz in big grids
+     * @param bz  number of z for each bigz
+     */
+    void init_chgmpi(const int& nbz, const int& bz);
+
+    /**
+     * @brief Sum rho at different pools (k-point parallelism).
+     *        Only used when GlobalV::KPAR > 1
+     */
+    void rho_mpi();
+
+	  /**
+	   * @brief 	Reduce among different pools 
+     *          If NPROC_IN_POOLs are all the same, use GlobalV::INTER_POOL
+     *          else, gather rho in a POOL, and then reduce among different POOLs
+	   * 
+	   * @param array_rho f(rho): an array [nrxx]
+	   */
+	  void reduce_diff_pools(double* array_rho) const;
 
     // mohan add 2021-02-20
     int nrxx; // number of r vectors in this processor
@@ -99,6 +127,13 @@ class Charge
     bool allocate_rho;
 
     bool allocate_rho_final_scf; // LiuXh add 20180606
+#ifdef __MPI
+  private:
+    bool use_intel_pool = false; //use INTER_POOL when NPROC_IN_POOLs are all the same
+    int *rec = nullptr; //The number of elements each process should receive into the receive buffer.
+    int *dis = nullptr; //The displacement (relative to recvbuf) for each process in the receive buffer.
+#endif
+
 };
 
 #endif // charge

source/module_elecstate/module_charge/charge_init.cpp

@@ -15,7 +15,7 @@
 #include "module_hamilt_pw/hamilt_pwdft/parallel_grid.h"
 #include "module_io/rho_io.h"
 
-void Charge::init_rho(elecstate::efermi& eferm_iout, const ModuleBase::ComplexMatrix& strucFac)
+void Charge::init_rho(elecstate::efermi& eferm_iout, const ModuleBase::ComplexMatrix& strucFac, const int& nbz, const int& bz)
 {
     ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running, "init_chg", GlobalV::init_chg);
 
@@ -152,6 +152,9 @@ void Charge::init_rho(elecstate::efermi& eferm_iout, const ModuleBase::ComplexMa
         }
         GlobalC::restart.info_load.load_charge_finish = true;
     }
+#ifdef __MPI
+    this->init_chgmpi(nbz, bz);
+#endif
 }
 
 //==========================================================