make bpcg support bndpar > 1

Author: Qianruipku

source/module_base/para_gemm.cpp

@@ -36,7 +36,7 @@ class PGemmCN
      * @param LDB leading dimension of B in each proc
      * @param nrow number of rows of A or B
      * @param LDC leading dimension of C. C can be C_local or C_global
-     * @param gatherC whether gather C_local to C_global
+     * @param mode 1: gather C_local to C_global, 2:C_local(nrow * ncol_loc), 3:C_global(nrow_loc * ncol)
      */
     void set_dimension(
 #ifdef __MPI
@@ -49,7 +49,7 @@ class PGemmCN
         const int LDB,
         const int nrow,
         const int LDC,
-        const bool gatherC = true);
+        const int mode = 1);
 
     /**
      * @brief calculate C = alpha * A^H * B + beta * C
@@ -67,14 +67,16 @@ class PGemmCN
 
     std::vector<int> colA_loc; ///< [col_nproc] number of columns of A matrix in each proc
     int max_colA = 0;          ///< maximum number of columns of A matrix in all procs
-    std::vector<int> colB_loc; ///<[col_nproc] number of columns of B matrix in each proc
+    std::vector<int> colB_loc; ///< [col_nproc] number of columns of B matrix in each proc
+    int max_colB = 0;          ///< maximum number of columns of B matrix in all procs
 
     std::vector<MPI_Request> requests; ///< MPI request
     std::vector<int> recv_counts;      ///< receive counts for gathering C_local to C_global
     std::vector<int> displs;           ///< displacements for gathering C_local to C_global
     int size_C_local = 0;              ///< size of C_local, which is a local matrix in each proc
     int size_C_global = 0;             ///< size of C_global, which is the global C matrix gathered from all procs
     bool gatherC = true;               ///< whether gather C_local to C_global
+    bool divideCrow = false;           ///< whether divide C_global to C_local
 #endif
     int ncolA = 0; ///< number of columns of A, which is a local matrix in each proc
     int ncolB = 0; ///< number of columns of B, which is a local matrix in each proc
@@ -83,6 +85,14 @@ class PGemmCN
     int LDB = 0;   ///< leading dimension of B in each proc
     int LDC = 0;   ///< leading dimension of C, which can be C_local or C_global
   private:
+    /// @brief for col_nproc == 1
+    void multiply_single(const T alpha, const T* A, const T* B, const T beta, T* C);
+#ifdef __MPI
+    /// @brief for mode = 1 or 2
+    void multiply_col(const T alpha, const T* A, const T* B, const T beta, T* C);
+    /// @brief for mode = 3
+    void multiply_row(const T alpha, const T* A, const T* B, const T beta, T* C);
+#endif
     using resmem_dev_op = base_device::memory::resize_memory_op<T, Device>;
     using delmem_dev_op = base_device::memory::delete_memory_op<T, Device>;
     using syncmem_dev_op = base_device::memory::synchronize_memory_op<T, Device, Device>;

source/module_base/para_gemm.h

@@ -3,10 +3,10 @@
 #include "mpi.h"
 #include "parallel_global.h"
 
-MPI_Comm POOL_WORLD;
-MPI_Comm INTER_POOL; // communicator among different pools
-MPI_Comm STO_WORLD;
-MPI_Comm PARAPW_WORLD;
+MPI_Comm POOL_WORLD; //groups for different plane waves. In this group, only plane waves are different. K-points and bands are the same.
+MPI_Comm KP_WORLD;   // groups for differnt k. In this group, only k-points are different. Bands and plane waves are the same.
+MPI_Comm BP_WORLD;   // groups for differnt bands. In this group, only bands are different. K-points and plane waves are the same.
+MPI_Comm INT_BGROUP; // internal comm groups for same bands. In this group, only bands are the same. K-points and plane waves are different.
 MPI_Comm GRID_WORLD; // mohan add 2012-01-13
 MPI_Comm DIAG_WORLD; // mohan add 2012-01-13
 

source/module_base/parallel_comm.cpp

@@ -4,9 +4,9 @@
 #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 KP_WORLD; // communicator among different pools
+extern MPI_Comm INT_BGROUP;
+extern MPI_Comm BP_WORLD;
 extern MPI_Comm GRID_WORLD; // mohan add 2012-01-13
 extern MPI_Comm DIAG_WORLD; // mohan add 2012-01-13
 

source/module_base/parallel_comm.h

@@ -237,12 +237,12 @@ void Parallel_Global::read_pal_param(int argc,
 void Parallel_Global::finalize_mpi()
 {
     MPI_Comm_free(&POOL_WORLD);
-    if (INTER_POOL != MPI_COMM_NULL)
+    if (KP_WORLD != MPI_COMM_NULL)
     {
-        MPI_Comm_free(&INTER_POOL);
+        MPI_Comm_free(&KP_WORLD);
     }
-    MPI_Comm_free(&STO_WORLD);
-    MPI_Comm_free(&PARAPW_WORLD);
+    MPI_Comm_free(&INT_BGROUP);
+    MPI_Comm_free(&BP_WORLD);
     MPI_Comm_free(&GRID_WORLD);
     MPI_Comm_free(&DIAG_WORLD);
     MPI_Finalize();
@@ -324,7 +324,7 @@ void Parallel_Global::divide_pools(const int& NPROC,
                                    int& MY_POOL)
 {
     // note: the order of k-point parallelization and band parallelization is important
-    //       The order will not change the behavior of INTER_POOL or PARAPW_WORLD, and MY_POOL
+    //       The order will not change the behavior of KP_WORLD or BP_WORLD, and MY_POOL
     //       and MY_BNDGROUP will be the same as well.
     if(BNDPAR > 1 && NPROC %(BNDPAR * KPAR) != 0)
     {
@@ -349,28 +349,28 @@ void Parallel_Global::divide_pools(const int& NPROC,
     MPI_Comm_dup(bndpar_group.group_comm, &POOL_WORLD);
     if(kpar_group.inter_comm != MPI_COMM_NULL)
     {
-        MPI_Comm_dup(kpar_group.inter_comm, &INTER_POOL);
+        MPI_Comm_dup(kpar_group.inter_comm, &KP_WORLD);
     }
     else
     {
-        INTER_POOL = MPI_COMM_NULL;
+        KP_WORLD = MPI_COMM_NULL;
     }
 
     if(BNDPAR > 1)
     {
         NPROC_IN_BNDGROUP = kpar_group.ngroups * bndpar_group.nprocs_in_group;
         RANK_IN_BPGROUP = kpar_group.my_group * bndpar_group.nprocs_in_group + bndpar_group.rank_in_group;
         MY_BNDGROUP = bndpar_group.my_group;
-        MPI_Comm_split(MPI_COMM_WORLD, MY_BNDGROUP, RANK_IN_BPGROUP, &STO_WORLD);
-        MPI_Comm_dup(bndpar_group.inter_comm, &PARAPW_WORLD);
+        MPI_Comm_split(MPI_COMM_WORLD, MY_BNDGROUP, RANK_IN_BPGROUP, &INT_BGROUP);
+        MPI_Comm_dup(bndpar_group.inter_comm, &BP_WORLD);
     }
     else
     {
         NPROC_IN_BNDGROUP = NPROC;
         RANK_IN_BPGROUP = MY_RANK;
         MY_BNDGROUP = 0;
-        MPI_Comm_dup(MPI_COMM_WORLD, &STO_WORLD);
-        MPI_Comm_split(MPI_COMM_WORLD, MY_RANK, 0, &PARAPW_WORLD);
+        MPI_Comm_dup(MPI_COMM_WORLD, &INT_BGROUP);
+        MPI_Comm_split(MPI_COMM_WORLD, MY_RANK, 0, &BP_WORLD);
     }
     return;
 }

source/module_base/parallel_global.cpp

@@ -110,9 +110,9 @@ void Parallel_Reduce::reduce_pool<double>(double* object, const int n)
 
 // (1) the value is same in each pool.
 // (2) we need to reduce the value from different pool.
-void Parallel_Reduce::reduce_double_allpool(const int& kpar, const int& nproc_in_pool, double& object)
+void Parallel_Reduce::reduce_double_allpool(const int& npool, const int& nproc_in_pool, double& object)
 {
-    if (kpar == 1) 
+    if (npool == 1) 
     {
         return;
     }
@@ -124,9 +124,9 @@ void Parallel_Reduce::reduce_double_allpool(const int& kpar, const int& nproc_in
 
 // (1) the value is same in each pool.
 // (2) we need to reduce the value from different pool.
-void Parallel_Reduce::reduce_double_allpool(const int& kpar, const int& nproc_in_pool, double* object, const int n)
+void Parallel_Reduce::reduce_double_allpool(const int& npool, const int& nproc_in_pool, double* object, const int n)
 {
-    if (kpar == 1) 
+    if (npool == 1) 
     {
         return;
     }

source/module_base/parallel_reduce.cpp

@@ -31,8 +31,8 @@ void reduce_int_grid(int* object, const int n); // mohan add 2012-01-12
 void reduce_double_grid(double* object, const int n);
 void reduce_double_diag(double* object, const int n);
 
-void reduce_double_allpool(const int& kpar, const int& nproc_in_pool, double& object);
-void reduce_double_allpool(const int& kpar, const int& nproc_in_pool, double* object, const int n);
+void reduce_double_allpool(const int& npool, const int& nproc_in_pool, double& object);
+void reduce_double_allpool(const int& npool, const int& nproc_in_pool, double* object, const int n);
 
 void gather_min_int_all(const int& nproc, int& v);
 void gather_max_double_all(const int& nproc, double& v);

source/module_base/parallel_reduce.h

@@ -367,7 +367,49 @@ TYPED_TEST(PgemmTest, odd_case)
     this->compare_result(ncolA_global, ncolB_global, LDC_global);
 }
 
-TYPED_TEST(PgemmTest, odd_case_not_gather)
+TYPED_TEST(PgemmTest, row_parallel)
+{
+    const int ncolA_global = 17, ncolB_global = 7, nrow_global = 13;
+    const int LDA_global = 17, LDB_global = 18, LDC_global = 19;
+
+    this->decide_ngroup(1, 4);
+    this->prepare(ncolA_global, ncolB_global, nrow_global, LDA_global, LDB_global, LDC_global);
+
+    this->pgemm.set_dimension(this->col_world,
+                              this->row_world,
+                              this->ncolA,
+                              this->LDA,
+                              this->ncolB,
+                              this->LDB,
+                              this->nrow,
+                              LDC_global);
+    this->pgemm.multiply(this->alpha, this->A_local.data(), this->B_local.data(), this->beta, this->C_global.data());
+
+    this->compare_result(ncolA_global, ncolB_global, LDC_global);
+}
+
+TYPED_TEST(PgemmTest, col_parallel)
+{
+    const int ncolA_global = 17, ncolB_global = 7, nrow_global = 13;
+    const int LDA_global = 17, LDB_global = 18, LDC_global = 19;
+
+    this->decide_ngroup(4, 1);
+    this->prepare(ncolA_global, ncolB_global, nrow_global, LDA_global, LDB_global, LDC_global);
+
+    this->pgemm.set_dimension(this->col_world,
+                              this->row_world,
+                              this->ncolA,
+                              this->LDA,
+                              this->ncolB,
+                              this->LDB,
+                              this->nrow,
+                              LDC_global);
+    this->pgemm.multiply(this->alpha, this->A_local.data(), this->B_local.data(), this->beta, this->C_global.data());
+
+    this->compare_result(ncolA_global, ncolB_global, LDC_global);
+}
+
+TYPED_TEST(PgemmTest, divide_col)
 {
     const int ncolA_global = 17, ncolB_global = 7, nrow_global = 13;
     const int LDA_global = 17, LDB_global = 18, LDC_global = 19;
@@ -392,7 +434,7 @@ TYPED_TEST(PgemmTest, odd_case_not_gather)
                               this->LDB,
                               this->nrow,
                               LDC_global,
-                              false);
+                              2);
     this->pgemm.multiply(this->alpha, this->A_local.data(), this->B_local.data(), this->beta, this->C_global.data()+ start);
 
 
@@ -408,34 +450,32 @@ TYPED_TEST(PgemmTest, odd_case_not_gather)
     }
 }
 
-TYPED_TEST(PgemmTest, row_parallel)
+TYPED_TEST(PgemmTest, divide_row)
 {
     const int ncolA_global = 17, ncolB_global = 7, nrow_global = 13;
     const int LDA_global = 17, LDB_global = 18, LDC_global = 19;
 
-    this->decide_ngroup(1, 4);
+    this->decide_ngroup(2, 2);
     this->prepare(ncolA_global, ncolB_global, nrow_global, LDA_global, LDB_global, LDC_global);
+    std::vector<int> colA_loc(this->nproc_col);
+    MPI_Allgather(&this->ncolA, 1, MPI_INT, colA_loc.data(), 1, MPI_INT, this->col_world);
+    std::vector<int> displs(this->nproc_col);
+    displs[0] = 0;
+    for (int i = 1; i < this->nproc_col; i++)
+    {
+        displs[i] = (displs[i - 1] + colA_loc[i - 1]);
+    }
+    int start = displs[this->rank_col];
 
-    this->pgemm.set_dimension(this->col_world,
-                              this->row_world,
-                              this->ncolA,
-                              this->LDA,
-                              this->ncolB,
-                              this->LDB,
-                              this->nrow,
-                              LDC_global);
-    this->pgemm.multiply(this->alpha, this->A_local.data(), this->B_local.data(), this->beta, this->C_global.data());
-
-    this->compare_result(ncolA_global, ncolB_global, LDC_global);
-}
-
-TYPED_TEST(PgemmTest, col_parallel)
-{
-    const int ncolA_global = 17, ncolB_global = 7, nrow_global = 13;
-    const int LDA_global = 17, LDB_global = 18, LDC_global = 19;
-
-    this->decide_ngroup(4, 1);
-    this->prepare(ncolA_global, ncolB_global, nrow_global, LDA_global, LDB_global, LDC_global);
+    int LDC_local = this->ncolA + 2;
+    std::vector<TypeParam> C_loc(LDC_local * ncolB_global, 0.0);
+    for(int i = 0; i < ncolB_global; i++)
+    {
+        for(int j = 0; j < this->ncolA; j++)
+        {
+            C_loc[i * LDC_local + j] = this->C_global[i * LDC_global + start + j];
+        }
+    }
 
     this->pgemm.set_dimension(this->col_world,
                               this->row_world,
@@ -444,10 +484,21 @@ TYPED_TEST(PgemmTest, col_parallel)
                               this->ncolB,
                               this->LDB,
                               this->nrow,
-                              LDC_global);
-    this->pgemm.multiply(this->alpha, this->A_local.data(), this->B_local.data(), this->beta, this->C_global.data());
+                              LDC_local,
+                              3);
+    this->pgemm.multiply(this->alpha, this->A_local.data(), this->B_local.data(), this->beta, C_loc.data());
 
-    this->compare_result(ncolA_global, ncolB_global, LDC_global);
+    
+
+    for (int i = 0; i < ncolB_global; i++)
+    {
+        for (int j = 0; j < this->ncolA; j++)
+        {
+            EXPECT_NEAR(get_double(this->Cref_global[i * LDC_global + start + j]),
+                        get_double(C_loc[i * LDC_local + j]),
+                        1e-10);
+        }
+    }
 }
 
 int main(int argc, char** argv)

source/module_base/test_parallel/test_para_gemm.cpp

@@ -163,8 +163,8 @@ void ElecState::calculate_weights()
                              this->klist->isk);
         }
 #ifdef __MPI
-        // qianrui fix a bug on 2021-7-21
-        Parallel_Reduce::reduce_double_allpool(GlobalV::KPAR, GlobalV::NPROC_IN_POOL, this->f_en.demet);
+        const int npool = GlobalV::KPAR * PARAM.inp.bndpar;
+        Parallel_Reduce::reduce_double_allpool(npool, GlobalV::NPROC_IN_POOL, this->f_en.demet);
 #endif
     }
     else if (Occupy::fixed_occupations)
@@ -192,16 +192,11 @@ void ElecState::calEBand()
         }
     }
     this->f_en.eband = eband;
-    if (GlobalV::KPAR != 1 && PARAM.inp.esolver_type != "sdft")
-    {
-        //==================================
-        // Reduce all the Energy in each cpu
-        //==================================
-        this->f_en.eband /= GlobalV::NPROC_IN_POOL;
+
 #ifdef __MPI
-        Parallel_Reduce::reduce_all(this->f_en.eband);
+    const int npool = GlobalV::KPAR * PARAM.inp.bndpar;
+    Parallel_Reduce::reduce_double_allpool(npool, GlobalV::NPROC_IN_POOL, this->f_en.eband);
 #endif
-    }
     return;
 }
 
@@ -253,8 +248,8 @@ void ElecState::init_ks(Charge* chg_in, // pointer for class Charge
     // init nelec_spin with nelec and nupdown
     this->init_nelec_spin();
     // initialize ekb and wg
-    this->ekb.create(nk_in, PARAM.inp.nbands);
-    this->wg.create(nk_in, PARAM.inp.nbands);
+    this->ekb.create(nk_in, PARAM.globalv.nbands_l);
+    this->wg.create(nk_in, PARAM.globalv.nbands_l);
 }
 
 } // namespace elecstate

source/module_elecstate/elecstate.cpp

@@ -19,7 +19,7 @@ void ElecStatePW_SDFT<T, Device>::psiToRho(const psi::Psi<T, Device>& psi)
         setmem_var_op()(this->rho[is], 0, this->charge->nrxx);
     }
 
-    if (GlobalV::MY_BNDGROUP == 0 || PARAM.inp.ks_solver == "bpcg")
+    if (PARAM.globalv.ks_run)
     {
         for (int ik = 0; ik < psi.get_nk(); ++ik)
         {