Merge branch 'develop' of https://github.com/mohanchen/abacus-mc into develop

Author: mohanchen

source/source_base/module_device/device.h

@@ -65,8 +65,8 @@ std::string get_device_flag(const std::string& device,
 /**
  * @brief Get the rank of current node
  *        Note that GPU can only be binded with CPU in the same node
- * 
- * @return int 
+ *
+ * @return int
  */
 int get_node_rank();
 int get_node_rank_with_mpi_shared(const MPI_Comm mpi_comm = MPI_COMM_WORLD);
@@ -91,6 +91,14 @@ void record_device_memory(const Device* dev, std::ofstream& ofs_device, std::str
     return;
 }
 
+#if defined(__CUDA) || defined(__ROCM)
+template <>
+void print_device_info<base_device::DEVICE_GPU>(const base_device::DEVICE_GPU *ctx, std::ofstream &ofs_device);
+
+template <>
+void record_device_memory<base_device::DEVICE_GPU>(const base_device::DEVICE_GPU* dev, std::ofstream& ofs_device, std::string str, size_t size);
+#endif
+
 } // end of namespace information
 } // end of namespace base_device
 

source/source_base/module_device/output_device.cpp

@@ -190,7 +190,7 @@ void print_device_info<base_device::DEVICE_GPU>(
     ofs_device << "Detected " << deviceCount << " CUDA Capable device(s)\n";
   }
   int dev = 0, driverVersion = 0, runtimeVersion = 0;
-  cudaErrcheck(cudaSetDevice(dev));
+  cudaErrcheck(cudaGetDevice(&dev));
   cudaDeviceProp deviceProp;
   cudaErrcheck(cudaGetDeviceProperties(&deviceProp, dev));
   ofs_device << "\nDevice " << dev << ":\t " << deviceProp.name << std::endl;
@@ -429,7 +429,7 @@ void print_device_info<base_device::DEVICE_GPU>(
     ofs_device << "Detected " << deviceCount << " CUDA Capable device(s)\n";
   }
   int dev = 0, driverVersion = 0, runtimeVersion = 0;
-  hipErrcheck(hipSetDevice(dev));
+  hipErrcheck(hipGetDevice(&dev));
   hipDeviceProp_t deviceProp;
   hipErrcheck(hipGetDeviceProperties(&deviceProp, dev));
   ofs_device << "\nDevice " << dev << ":\t " << deviceProp.name << std::endl;

source/source_base/timer.cpp

@@ -15,7 +15,11 @@
 #include "source_base/formatter.h"
 
 #if defined(__CUDA) && defined(__USE_NVTX)
-#include <nvToolsExt.h>
+#if CUDA_VERSION < 12090
+#include "nvToolsExt.h"
+#else
+#include "nvtx3/nvToolsExt.h"
+#endif
 #include "source_io/module_parameter/parameter.h"
 #endif
 

source/source_hsolver/test/CMakeLists.txt

@@ -153,12 +153,19 @@ install(FILES diago_pexsi_parallel_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DI
 install(FILES parallel_k2d_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 
 
-
-AddTest(
-  TARGET MODULE_HSOLVER_diago_hs_parallel
-  LIBS parameter  ${math_libs} ELPA::ELPA base device MPI::MPI_CXX genelpa psi
-  SOURCES test_diago_hs_para.cpp ../diag_hs_para.cpp ../diago_pxxxgvx.cpp ../diago_elpa.cpp ../diago_scalapack.cpp 
-)
+if (USE_ELPA)
+  AddTest(
+    TARGET MODULE_HSOLVER_diago_hs_parallel
+    LIBS parameter  ${math_libs} ELPA::ELPA base device MPI::MPI_CXX genelpa psi
+    SOURCES test_diago_hs_para.cpp ../diag_hs_para.cpp ../diago_pxxxgvx.cpp ../diago_elpa.cpp ../diago_scalapack.cpp 
+  )
+else()
+  AddTest(
+      TARGET MODULE_HSOLVER_diago_hs_parallel
+      LIBS parameter  ${math_libs} base device MPI::MPI_CXX psi
+      SOURCES test_diago_hs_para.cpp ../diag_hs_para.cpp ../diago_pxxxgvx.cpp ../diago_scalapack.cpp 
+    )
+endif()
 
 AddTest(
   TARGET MODULE_HSOLVER_linear_trans

source/source_hsolver/test/test_diago_hs_para.cpp

@@ -160,7 +160,9 @@ void test_performance(int lda, int nb, int nbands, MPI_Comm comm,int case_numb,
     MPI_Comm_size(comm, &nproc);
 
     std::vector<T> h_mat, s_mat, wfc, h_psi, s_psi;
+#ifdef __ELPA
     std::vector<typename GetTypeReal<T>::type> ekb_elpa(lda);
+#endif
     std::vector<typename GetTypeReal<T>::type> ekb_scalap(lda);
     std::vector<typename GetTypeReal<T>::type> ekb_lapack(lda);
 
@@ -176,32 +178,36 @@ void test_performance(int lda, int nb, int nbands, MPI_Comm comm,int case_numb,
     }
 
     // store all the times in a vector
+#ifdef __ELPA
     std::vector<double> time_elpa(case_numb, 0);
+#endif
     std::vector<double> time_scalap(case_numb, 0);
     std::vector<double> time_lapack(case_numb, 0);
 
     if (my_rank == 0) { std::cout << "Random matrix ";
 }
-    for (int randomi = 0; randomi < case_numb; ++randomi) 
+    for (int randomi = 0; randomi < case_numb; ++randomi)
     {
-        
+
         if (my_rank == 0) {
             std::cout << randomi << " ";
             generate_random_hs(lda, randomi, h_mat, s_mat);
         }
-
+        auto start = std::chrono::high_resolution_clock::now();
+        auto end = std::chrono::high_resolution_clock::now();
+#ifdef __ELPA
         // ELPA
         MPI_Barrier(comm);
-        auto start = std::chrono::high_resolution_clock::now();
+        start = std::chrono::high_resolution_clock::now();
         for (int j=0;j<loop_numb;j++)
         {
             hsolver::diago_hs_para<T>(h_mat.data(), s_mat.data(), lda, nbands,ekb_elpa.data(), wfc.data(), comm, 1, nb);
             MPI_Barrier(comm);
         }
         MPI_Barrier(comm);
-        auto end = std::chrono::high_resolution_clock::now();
+        end = std::chrono::high_resolution_clock::now();
         time_elpa[randomi] = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
-
+#endif
 
         // scalapack
         start = std::chrono::high_resolution_clock::now();
@@ -215,8 +221,8 @@ void test_performance(int lda, int nb, int nbands, MPI_Comm comm,int case_numb,
         time_scalap[randomi] = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
 
         //LApack
-        if (my_rank == 0) 
-        { 
+        if (my_rank == 0)
+        {
             std::vector<T> h_tmp, s_tmp;
             start = std::chrono::high_resolution_clock::now();
             base_device::DEVICE_CPU* ctx = {};
@@ -239,26 +245,34 @@ void test_performance(int lda, int nb, int nbands, MPI_Comm comm,int case_numb,
 
             //COMPARE EKB
             for (int i = 0; i < nbands; ++i) {
-                typename GetTypeReal<T>::type diff_elpa_lapack = std::abs(ekb_elpa[i] - ekb_lapack[i]);
                 typename GetTypeReal<T>::type diff_scalap_lapack = std::abs(ekb_scalap[i] - ekb_lapack[i]);
+#ifdef __ELPA
+                typename GetTypeReal<T>::type diff_elpa_lapack = std::abs(ekb_elpa[i] - ekb_lapack[i]);
                 if (diff_elpa_lapack > 1e-6 || diff_scalap_lapack > 1e-6)
+#else
+                if (diff_scalap_lapack > 1e-6)
+#endif
                 {
+#ifdef __ELPA
                     std::cout << "eigenvalue " << i << " by ELPA: " << ekb_elpa[i] << std::endl;
+#endif
                     std::cout << "eigenvalue " << i << " by Scalapack: " << ekb_scalap[i] << std::endl;
                     std::cout << "eigenvalue " << i << " by Lapack: " << ekb_lapack[i] << std::endl;
                 }
             }
         }
-        MPI_Barrier(comm);   
+        MPI_Barrier(comm);
 
     }
 
     if (my_rank == 0)
     {
+#ifdef __ELPA
         std::cout << "\nELPA Time     : ";
         for (int i=0; i < case_numb;i++)
         {std::cout << time_elpa[i] << " ";}
         std::cout << std::endl;
+#endif
 
         std::cout << "scalapack Time: ";
         for (int i=0; i < case_numb;i++)
@@ -271,21 +285,29 @@ void test_performance(int lda, int nb, int nbands, MPI_Comm comm,int case_numb,
         std::cout << std::endl;
 
         // print out the average time and speedup
+#ifdef __ELPA
         double avg_time_elpa = 0;
+#endif
         double avg_time_scalap = 0;
         double avg_time_lapack = 0;
         for (int i=0; i < case_numb;i++)
         {
+#ifdef __ELPA
             avg_time_elpa += time_elpa[i];
+#endif
             avg_time_scalap += time_scalap[i];
             avg_time_lapack += time_lapack[i];
         }
 
+#ifdef __ELPA
         avg_time_elpa /= case_numb;
+#endif
         avg_time_scalap /= case_numb;
         avg_time_lapack /= case_numb;
         std::cout << "Average Lapack Time   : " << avg_time_lapack << " ms" << std::endl;
+#ifdef __ELPA
         std::cout << "Average ELPA Time     : " << avg_time_elpa << " ms, Speedup: " << avg_time_lapack / avg_time_elpa << std::endl;
+#endif
         std::cout << "Average Scalapack Time: " << avg_time_scalap << " ms, Speedup: " << avg_time_lapack / avg_time_scalap << std::endl;
     }
 }

source/source_lcao/module_deepks/LCAO_deepks_interface.cpp

@@ -618,6 +618,24 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
 // 7. atom.npy, box.npy, overlap.npy
 //================================================================================
 
+    if(  ((PARAM.inp.deepks_out_labels == 2) && is_after_scf )
+      || ( PARAM.inp.deepks_out_freq_elec ) )// need overlap when deepks_out_freq_elec
+    {
+        if (PARAM.inp.deepks_v_delta > 0)
+        {
+            // prepare for overlap.npy, very much like h_tot except for p_ham->getSk()
+            std::vector<TH> s_tot(nks);
+            DeePKS_domain::get_h_tot<TK, TH, TR>(*ParaV, p_ham, s_tot, nlocal, nks, 'S');
+            const std::string file_stot = get_filename("overlap", PARAM.inp.deepks_out_labels, iter);
+            LCAO_deepks_io::save_npy_h<TK, TH>(s_tot,
+                                            file_stot,
+                                            nlocal,
+                                            nks,
+                                            rank,
+                                            1.0); // don't need unit_scale for overlap
+        }
+    }
+
     if ( is_after_scf ) // don't need to output in multiple electronic steps
     {
         if (PARAM.inp.deepks_out_labels == 2)
@@ -632,20 +650,6 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
             DeePKS_domain::prepare_box(ucell, box_out);
             const std::string file_box = PARAM.globalv.global_out_dir + "deepks_box.npy";
             LCAO_deepks_io::save_tensor2npy<double>(file_box, box_out, rank);
-
-            if (PARAM.inp.deepks_v_delta > 0)
-            {
-                // prepare for overlap.npy, very much like h_tot except for p_ham->getSk()
-                std::vector<TH> s_tot(nks);
-                DeePKS_domain::get_h_tot<TK, TH, TR>(*ParaV, p_ham, s_tot, nlocal, nks, 'S');
-                const std::string file_stot = PARAM.globalv.global_out_dir + "deepks_overlap.npy";
-                LCAO_deepks_io::save_npy_h<TK, TH>(s_tot,
-                                                file_stot,
-                                                nlocal,
-                                                nks,
-                                                rank,
-                                                1.0); // don't need unit_scale for overlap
-            }
         }
 
 //================================================================================

source/source_lcao/module_gint/gint.h

@@ -265,7 +265,7 @@ class Gint {
     std::vector<hamilt::HContainer<double>*> DMRGint; 
 
     //! tmp tools used in transfer_DM2DtoGrid 
-    hamilt::HContainer<double>* DMRGint_full = nullptr;
+    hamilt::HContainer<double>* dm2d_tmp = nullptr;
 
     std::vector<hamilt::HContainer<double>> pvdpRx_reduced;
     std::vector<hamilt::HContainer<double>> pvdpRy_reduced;

source/source_lcao/module_gint/gint_old.cpp

@@ -33,7 +33,7 @@ Gint::~Gint() {
         delete this->hRGint_tmp[is];
     }
 #ifdef __MPI
-    delete this->DMRGint_full;
+    delete this->dm2d_tmp;
 #endif
 }
 
@@ -171,10 +171,9 @@ void Gint::initialize_pvpR(const UnitCell& ucell_in, const Grid_Driver* gd, cons
             this->hRGint_tmp[is] = new hamilt::HContainer<double>(ucell_in.nat);
         }
 #ifdef __MPI
-        if (this->DMRGint_full != nullptr) {
-            delete this->DMRGint_full;
+        if (this->dm2d_tmp != nullptr) {
+            delete this->dm2d_tmp;
         }
-        this->DMRGint_full = new hamilt::HContainer<double>(ucell_in.nat);
 #endif
     }
 
@@ -210,12 +209,6 @@ void Gint::initialize_pvpR(const UnitCell& ucell_in, const Grid_Driver* gd, cons
         ModuleBase::Memory::record("Gint::DMRGint",
                                        this->DMRGint[0]->get_memory_size()
                                            * this->DMRGint.size()*nspin);
-#ifdef __MPI
-        this->DMRGint_full->insert_ijrs(this->gridt->get_ijr_info(), ucell_in, npol);
-        this->DMRGint_full->allocate(nullptr, true);
-        ModuleBase::Memory::record("Gint::DMRGint_full",
-                                   this->DMRGint_full->get_memory_size());
-#endif
     }
 }
 
@@ -231,9 +224,7 @@ void Gint::reset_DMRGint(const int& nspin)
         {
             for (auto& d : this->DMRGint) { d->allocate(nullptr, false); }
 #ifdef __MPI
-            delete this->DMRGint_full;
-            this->DMRGint_full = new hamilt::HContainer<double>(*this->hRGint);
-            this->DMRGint_full->allocate(nullptr, false);
+            delete this->dm2d_tmp;
 #endif
         }
     }
@@ -262,37 +253,46 @@ void Gint::transfer_DM2DtoGrid(std::vector<hamilt::HContainer<double>*> DM2D) {
     } else // NSPIN=4 case
     {
 #ifdef __MPI
-        hamilt::transferParallels2Serials(*DM2D[0], this->DMRGint_full);
-#else
-        this->DMRGint_full = DM2D[0];
-#endif
-        std::vector<double*> tmp_pointer(4, nullptr);
-        for (int iap = 0; iap < this->DMRGint_full->size_atom_pairs(); ++iap) {
-            auto& ap = this->DMRGint_full->get_atom_pair(iap);
-            int iat1 = ap.get_atom_i();
-            int iat2 = ap.get_atom_j();
-            for (int ir = 0; ir < ap.get_R_size(); ++ir) {
-                const ModuleBase::Vector3<int> r_index = ap.get_R_index(ir);
-                for (int is = 0; is < 4; is++) {
-                    tmp_pointer[is] = this->DMRGint[is]
-                                          ->find_matrix(iat1, iat2, r_index)
-                                          ->get_pointer();
-                }
-                double* data_full = ap.get_pointer(ir);
-                for (int irow = 0; irow < ap.get_row_size(); irow += 2) {
-                    for (int icol = 0; icol < ap.get_col_size(); icol += 2) {
-                        *(tmp_pointer[0])++ = data_full[icol];
-                        *(tmp_pointer[1])++ = data_full[icol + 1];
-                    }
-                    data_full += ap.get_col_size();
-                    for (int icol = 0; icol < ap.get_col_size(); icol += 2) {
-                        *(tmp_pointer[2])++ = data_full[icol];
-                        *(tmp_pointer[3])++ = data_full[icol + 1];
+        // is=0:↑↑, 1:↑↓, 2:↓↑, 3:↓↓
+        const int row_set[4] = {0, 0, 1, 1};
+        const int col_set[4] = {0, 1, 0, 1};
+        int mg = DM2D[0]->get_paraV()->get_global_row_size()/2;
+        int ng = DM2D[0]->get_paraV()->get_global_col_size()/2;
+        int nb = DM2D[0]->get_paraV()->get_block_size()/2;
+        int blacs_ctxt = DM2D[0]->get_paraV()->blacs_ctxt;
+        std::vector<int> iat2iwt(ucell->nat);
+        for (int iat = 0; iat < ucell->nat; iat++) {
+            iat2iwt[iat] = ucell->get_iat2iwt()[iat]/2;
+        }
+        Parallel_Orbitals *pv = new Parallel_Orbitals();
+        pv->set(mg, ng, nb, blacs_ctxt);
+        pv->set_atomic_trace(iat2iwt.data(), ucell->nat, mg);
+        auto ijr_info = DM2D[0]->get_ijr_info();
+        this-> dm2d_tmp = new hamilt::HContainer<double>(pv, nullptr, &ijr_info);
+        ModuleBase::Memory::record("Gint::dm2d_tmp", this->dm2d_tmp->get_memory_size());
+        for (int is = 0; is < 4; is++){
+            for (int iap = 0; iap < DM2D[0]->size_atom_pairs(); ++iap) {
+                auto& ap = DM2D[0]->get_atom_pair(iap);
+                int iat1 = ap.get_atom_i();
+                int iat2 = ap.get_atom_j();
+                for (int ir = 0; ir < ap.get_R_size(); ++ir) {
+                    const ModuleBase::Vector3<int> r_index = ap.get_R_index(ir);
+                    double* matrix_out = this -> dm2d_tmp -> find_matrix(iat1, iat2, r_index)->get_pointer();
+                    double* matrix_in = ap.get_pointer(ir);
+                    for (int irow = 0; irow < ap.get_row_size()/2; irow ++) {
+                        for (int icol = 0; icol < ap.get_col_size()/2; icol++){
+                            int index_i = irow* ap.get_col_size()/2 + icol;
+                            int index_j = (irow*2+row_set[is]) * ap.get_col_size() + icol*2+col_set[is];
+                            matrix_out[index_i] = matrix_in[index_j];
+                        }
                     }
-                    data_full += ap.get_col_size();
                 }
             }
+            hamilt::transferParallels2Serials( *(this->dm2d_tmp), this->DMRGint[is]);
         }
+#else
+        //this->DMRGint_full = DM2D[0];
+#endif
     }
     ModuleBase::timer::tick("Gint", "transfer_DMR");
 }