[HSolver] PPCG Solver (CPU/CUDA), Parallel Optimization, and Diagonalization Refactor

.gitignore

@@ -1,5 +1,6 @@
 /build*
 build_info.h
+CMakeFiles/
 bin
 obj
 *.o

source/source_base/kernels/math_kernel_op_vec.cpp

@@ -25,7 +25,7 @@ struct vector_mul_real_op<T, base_device::DEVICE_CPU>
     void operator()(const int dim, T* result, const T* vector, const Real constant)
     {
 #ifdef _OPENMP
-#pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static) if(dim > 256)
 #endif
         for (int i = 0; i < dim; i++)
         {
@@ -43,7 +43,7 @@ struct vector_mul_vector_op<T, base_device::DEVICE_CPU>
         if (add)
         {
 #ifdef _OPENMP
-#pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static) if(dim > 256)
 #endif
             for (int i = 0; i < dim; i++)
             {
@@ -53,7 +53,7 @@ struct vector_mul_vector_op<T, base_device::DEVICE_CPU>
         else
         {
 #ifdef _OPENMP
-#pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static) if(dim > 256)
 #endif
             for (int i = 0; i < dim; i++)
             {
@@ -70,7 +70,7 @@ struct vector_div_constant_op<T, base_device::DEVICE_CPU>
     void operator()(const int& dim, T* result, const T* vector, const Real constant)
     {
 #ifdef _OPENMP
-#pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static) if(dim > 256)
 #endif
         for (int i = 0; i < dim; i++)
         {
@@ -86,7 +86,7 @@ struct vector_div_vector_op<T, base_device::DEVICE_CPU>
     void operator()(const int& dim, T* result, const T* vector1, const Real* vector2)
     {
 #ifdef _OPENMP
-#pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static) if(dim > 256)
 #endif
         for (int i = 0; i < dim; i++)
         {
@@ -122,7 +122,7 @@ struct vector_add_vector_op<T, base_device::DEVICE_CPU>
                     const Real constant2)
     {
 #ifdef _OPENMP
-#pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static) if(dim > 256)
 #endif
         for (int i = 0; i < dim; i++)
         {

source/source_base/module_container/ATen/kernels/cuda/memory.cu

@@ -94,7 +94,7 @@ struct synchronize_memory<T, DEVICE_GPU, DEVICE_GPU> {
         const T *arr_in,
         const size_t& size)
     {
-        CHECK_CUDA(cudaMemcpy(arr_out, arr_in, sizeof(T) * size, cudaMemcpyHostToDevice));
+        CHECK_CUDA(cudaMemcpy(arr_out, arr_in, sizeof(T) * size, cudaMemcpyDeviceToDevice));
     }
 };
 

source/source_base/test/CMakeLists.txt

@@ -1,5 +1,5 @@
 remove_definitions(-D__MPI)
-install(DIRECTORY data DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/data DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 AddTest(
   TARGET MODULE_BASE_blas_connector
   LIBS parameter  ${math_libs} base device

source/source_base/test_parallel/CMakeLists.txt

@@ -16,9 +16,9 @@ AddTest(
   SOURCES parallel_reduce_test.cpp ../global_variable.cpp ../parallel_global.cpp ../parallel_comm.cpp ../parallel_common.cpp ../parallel_reduce.cpp ../tool_quit.cpp ../global_file.cpp ../global_function.cpp ../memory_recorder.cpp ../timer.cpp
 )
 
-install(FILES parallel_common_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES parallel_global_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES parallel_reduce_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/parallel_common_test.sh ${CMAKE_CURRENT_BINARY_DIR}/parallel_common_test.sh COPYONLY)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/parallel_global_test.sh ${CMAKE_CURRENT_BINARY_DIR}/parallel_global_test.sh COPYONLY)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/parallel_reduce_test.sh ${CMAKE_CURRENT_BINARY_DIR}/parallel_reduce_test.sh COPYONLY)
 
 find_program(BASH bash)
 add_test(NAME MODULE_BASE_parallel_common_test
@@ -57,7 +57,7 @@ AddTest(
     LIBS parameter ${math_libs}
 )
 
-install(FILES parallel_2d_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/parallel_2d_test.sh ${CMAKE_CURRENT_BINARY_DIR}/parallel_2d_test.sh COPYONLY)
 find_program(BASH bash)
 add_test(NAME MODULE_BASE_parallel_2d_test_para
       COMMAND ${BASH} parallel_2d_test.sh
@@ -77,7 +77,7 @@ add_test(NAME MODULE_BASE_parallel_2d_test_para
     LIBS parameter MPI::MPI_CXX ${BLACS_LIB}
     SOURCES blacs_connector_test.cpp
   )
-  install(FILES blacs_connector_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+  configure_file(${CMAKE_CURRENT_SOURCE_DIR}/blacs_connector_test.sh ${CMAKE_CURRENT_BINARY_DIR}/blacs_connector_test.sh COPYONLY)
   add_test(NAME MODULE_BASE_blacs_connector_test
         COMMAND ${BASH} blacs_connector_test.sh
         WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}

source/source_basis/module_ao/test/CMakeLists.txt

@@ -42,7 +42,7 @@ list(APPEND depend_files
   )
 
 install(DIRECTORY GaAs DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/../../../tests)
-install(DIRECTORY GaAs DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/GaAs DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 
 
 AddTest(
@@ -83,13 +83,13 @@ AddTest(
     LIBS parameter ${math_libs} device base
 )
 
-install(FILES parallel_orbitals_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/parallel_orbitals_test.sh ${CMAKE_CURRENT_BINARY_DIR}/parallel_orbitals_test.sh COPYONLY)
 find_program(BASH bash)
 add_test(NAME MODULE_AO_parallel_orbitals_test_para
       COMMAND ${BASH} parallel_orbitals_test.sh
       WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
 )
 
-install(DIRECTORY lcao_H2O DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/lcao_H2O DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 install(DIRECTORY lcao_H2O DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/../../../tests)
 

source/source_cell/test/CMakeLists.txt

@@ -4,14 +4,14 @@ remove_definitions(-D__ROCM)
 remove_definitions(-D__EXX)
 
 find_program(BASH bash)
-install(DIRECTORY support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES bcast_atom_pseudo_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES bcast_atom_spec_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES parallel_kpoints_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES klist_test_para.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES unitcell_test_parallel.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES bcast_read_sep_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES bcast_sep_cell_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/bcast_atom_pseudo_test.sh ${CMAKE_CURRENT_BINARY_DIR}/bcast_atom_pseudo_test.sh COPYONLY)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/bcast_atom_spec_test.sh ${CMAKE_CURRENT_BINARY_DIR}/bcast_atom_spec_test.sh COPYONLY)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/parallel_kpoints_test.sh ${CMAKE_CURRENT_BINARY_DIR}/parallel_kpoints_test.sh COPYONLY)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/klist_test_para.sh ${CMAKE_CURRENT_BINARY_DIR}/klist_test_para.sh COPYONLY)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/unitcell_test_parallel.sh ${CMAKE_CURRENT_BINARY_DIR}/unitcell_test_parallel.sh COPYONLY)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/bcast_read_sep_test.sh ${CMAKE_CURRENT_BINARY_DIR}/bcast_read_sep_test.sh COPYONLY)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/bcast_sep_cell_test.sh ${CMAKE_CURRENT_BINARY_DIR}/bcast_sep_cell_test.sh COPYONLY)
 
 list(APPEND cell_simple_srcs
     ../unitcell.cpp

source/source_cell/test_pw/CMakeLists.txt

@@ -4,8 +4,8 @@ remove_definitions(-D__ROCM)
 remove_definitions(-D__EXX)
 remove_definitions(-D__LCAO)
 
-install(DIRECTORY support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES unitcell_test_pw_para.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/unitcell_test_pw_para.sh ${CMAKE_CURRENT_BINARY_DIR}/unitcell_test_pw_para.sh COPYONLY)
 
 AddTest(
   TARGET MODULE_CELL_unitcell_test_pw

source/source_esolver/test/CMakeLists.txt

@@ -1,7 +1,7 @@
 remove_definitions(-D__MPI)
 remove_definitions(-D__LCAO)
 
-install(DIRECTORY support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 
 AddTest(
   TARGET MODULE_ESOLVER_esolver_dp_test

source/source_estate/elecstate_print.cpp

@@ -58,6 +58,7 @@ void print_scf_iterinfo(const std::string& ks_solver,
            {"scalapack_gvx", "GV"},
            {"cusolver", "CU"},
            {"bpcg", "BP"},
+           {"ppcg", "PP"},
            {"pexsi", "PE"},
            {"cusolvermp", "CM"},
            {"sdft", "CT"}}; // CT = Chebyshev Trace, for pure SDFT (nbands=0) where no H diagonalization is performed

source/source_estate/module_dm/test/CMakeLists.txt

@@ -2,7 +2,7 @@ remove_definitions(-D__MLALGO)
 remove_definitions(-D__CUDA)
 remove_definitions(-D__ROCM)
 
-install(DIRECTORY support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 
 if(TARGET MODULE_ESTATE_dm_io_test_serial)
   remove_definitions(-D__MPI)

source/source_estate/test/CMakeLists.txt

@@ -9,7 +9,7 @@ remove_definitions(-D_OPENMP)
 
 if (ENABLE_MPI)
 
-install(DIRECTORY support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 
 AddTest(
   TARGET MODULE_ESTATE_Elecstate_Op_UTs

source/source_hamilt/module_surchem/test/CMakeLists.txt

@@ -1,6 +1,6 @@
 remove_definitions(-D__LCAO )
 remove_definitions(-DUSE_LIBXC)
-install(DIRECTORY support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/support DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 
 list(APPEND depend_files 
 

source/source_hamilt/module_vdw/test/CMakeLists.txt

@@ -2,8 +2,8 @@ remove_definitions(-D__MLALGO)
 remove_definitions(-D__CUDA)
 remove_definitions(-D__ROCM)
 
-install(FILES c6.txt DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
-install(FILES r0.txt DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/c6.txt ${CMAKE_CURRENT_BINARY_DIR}/c6.txt COPYONLY)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/r0.txt ${CMAKE_CURRENT_BINARY_DIR}/r0.txt COPYONLY)
 
 AddTest(
   TARGET MODULE_HAMILT_vdwTest

source/source_hsolver/CMakeLists.txt

@@ -4,6 +4,7 @@ list(APPEND objects
     diago_david.cpp
     diago_dav_subspace.cpp
     diago_bpcg.cpp
+    diago_ppcg.cpp
     para_linear_transform.cpp
     hsolver_pw.cpp
     hsolver_lcaopw.cpp

source/source_hsolver/diago_bpcg.cpp

@@ -1,16 +1,20 @@
 #include "source_hsolver/diago_bpcg.h"
+#include "source_hsolver/module_diag/diag_orthogonalizer.h"
 
 #include "diago_iter_assist.h"
 #include "source_base/global_function.h"
 #include "source_base/kernels/math_kernel_op.h"
 #include "source_base/parallel_comm.h" // different MPI worlds
 #include "source_hsolver/kernels/bpcg_kernel_op.h"
+#include "source_hsolver/module_diag/diago_trace.h"
 #include "para_linear_transform.h"
 
 #include <ATen/kernels/blas.h>
 #include <ATen/kernels/lapack.h>
 #include <ATen/ops/einsum_op.h>
+#include <algorithm>
 #include <limits>
+#include <vector>
 
 namespace hsolver {
 
@@ -80,6 +84,9 @@ bool DiagoBPCG<T, Device>::test_error(const ct::Tensor& err_in, const std::vecto
         _err_st = tmp_cpu.data();
         syncmem_var_d2h_op()(_err_st, err_in.data<Real>(), this->n_band_l);
     }
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static) reduction(||:not_conv) if(this->n_band_l > 64)
+#endif
     for (int ii = 0; ii < this->n_band_l; ii++) {
         if (_err_st[ii] > ethr_band[ii]) {
             not_conv = true;
@@ -117,20 +124,14 @@ void DiagoBPCG<T, Device>::orth_cholesky(
 		ct::Tensor& hpsi_out,
 		ct::Tensor& hsub_out)
 {
-    // gemm: hsub_out(n_band x n_band) = psi_out^T(n_band x n_basis) * psi_out(n_basis x n_band)
-    this->pmmcn.multiply(1.0, psi_out.data<T>(), psi_out.data<T>(), 0.0, hsub_out.data<T>());
-
-    // set hsub matrix to lower format;
-    ct::kernels::set_matrix<T, ct_Device>()(
-        'L', hsub_out.data<T>(), this->n_band);
-
-    ct::kernels::lapack_potrf<T, ct_Device>()(
-        'U', this->n_band, hsub_out.data<T>(), this->n_band);
-    ct::kernels::lapack_trtri<T, ct_Device>()(
-        'U', 'N', this->n_band, hsub_out.data<T>(), this->n_band);
-
-    this->rotate_wf(hsub_out, psi_out, workspace_in);
-    this->rotate_wf(hsub_out, hpsi_out, workspace_in);
+    DiagOrthogonalizer<T, Device>(this->n_dim, this->n_basis)
+        .cholesky_orth_parallel(workspace_in.data<T>(),
+                                psi_out.data<T>(),
+                                hpsi_out.data<T>(),
+                                hsub_out.data<T>(),
+                                this->n_band,
+                                this->pmmcn,
+                                this->plintrans);
 }
 
 template<typename T, typename Device>
@@ -167,13 +168,12 @@ void DiagoBPCG<T, Device>::orth_projection(
         ct::Tensor& hsub_in,
         ct::Tensor& grad_out)
 {
-    // gemm: hsub_in(n_band x n_band) = psi_in^T(n_band x n_basis) * grad_out(n_basis x n_band)
-    this->pmmcn.multiply(1.0, psi_in.data<T>(), grad_out.data<T>(), 0.0, hsub_in.data<T>());
-
-    // grad_out(n_basis x n_band) = 1.0 * grad_out(n_basis x n_band) - psi_in(n_basis x n_band) * hsub_in(n_band x
-    // n_band)
-    this->plintrans.act(-1.0, psi_in.data<T>(), hsub_in.data<T>(), 1.0, grad_out.data<T>());
-    return;
+    DiagOrthogonalizer<T, Device>(this->n_dim, this->n_basis)
+        .project_out_parallel(psi_in.data<T>(),
+                              grad_out.data<T>(),
+                              hsub_in.data<T>(),
+                              this->pmmcn,
+                              this->plintrans);
 }
 
 template<typename T, typename Device>
@@ -265,6 +265,7 @@ void DiagoBPCG<T, Device>::diag(const HPsiFunc& hpsi_func,
                                 const std::vector<double>& ethr_band)
 {
     const int current_scf_iter = hsolver::DiagoIterAssist<T, Device>::SCF_ITER;
+    DiagoTrace trace("BPCG");
     // Get the pointer of the input psi
     this->psi = std::move(ct::TensorMap(psi_in /*psi_in.get_pointer()*/, t_type, device_type, {this->n_band_l, this->n_basis}));
 
@@ -294,6 +295,40 @@ void DiagoBPCG<T, Device>::diag(const HPsiFunc& hpsi_func,
         this->calc_grad_with_block(this->prec, this->err_st, this->beta,
                                  this->psi, this->hpsi, this->grad, this->grad_old);
 
+        if (trace.enabled())
+        {
+            std::vector<Real> err_host(this->n_band_l);
+            const Real* err_ptr = this->err_st.template data<Real>();
+            if (this->err_st.device_type() == ct::DeviceType::GpuDevice)
+            {
+                syncmem_var_d2h_op()(err_host.data(), this->err_st.template data<Real>(), this->n_band_l);
+                err_ptr = err_host.data();
+            }
+
+            Real max_residual = Real(0);
+            Real avg_residual = Real(0);
+            int n_converged = 0;
+            for (int ib = 0; ib < this->n_band_l; ++ib)
+            {
+                max_residual = std::max(max_residual, err_ptr[ib]);
+                avg_residual += err_ptr[ib];
+                if (err_ptr[ib] <= ethr_band[ib])
+                {
+                    ++n_converged;
+                }
+            }
+            if (this->n_band_l > 0)
+            {
+                avg_residual /= this->n_band_l;
+            }
+            trace.record_iteration(ntry,
+                                   this->n_band_l,
+                                   max_residual,
+                                   avg_residual,
+                                   n_converged,
+                                   Real(-1));
+        }
+
         // Orthogonalize column vectors g_i in matrix grad to column vectors p_j in matrix psi
         // for all 'j less or equal to i'.
         // Note: hsub and work are only used to store intermediate variables of gemm operator.