Perf: reduce memory allocation and copy in Diago_DavSubspace::diag_zhegvx

Signed-off-by：Tianxiang Wang<[email protected]>, Contributed under MetaX Integrated Circuits (Shanghai) Co., Ltd.

Author: wangtianxiang

source/source_hsolver/diago_dav_subspace.cpp

@@ -76,6 +76,7 @@ Diago_DavSubspace<T, Device>::Diago_DavSubspace(const std::vector<Real>& precond
     {
         resmem_real_op()(this->d_precondition, nbasis_in);
         // syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->d_precondition, this->precondition.data(), nbasis_in);
+        base_device::memory::resize_memory_op<T, Device>()(this->d_scc, this->nbase_x * this->nbase_x);
         resmem_real_op()(this->d_eigenvalue, this->nbase_x);
     }
 #endif
@@ -95,6 +96,7 @@ Diago_DavSubspace<T, Device>::~Diago_DavSubspace()
     if (this->device == base_device::GpuDevice)
     {
         delmem_real_op()(this->d_precondition);
+        delmem_complex_op()(this->d_scc);
         delmem_real_op()(this->d_eigenvalue);
     }
 #endif
@@ -546,34 +548,9 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
 #if defined(__CUDA) || defined(__ROCM)
         if (this->diag_comm.rank == 0)
         {
-            Real* eigenvalue_gpu = nullptr;
-            resmem_real_op()(eigenvalue_gpu, this->nbase_x);
-
-            syncmem_var_h2d_op()(eigenvalue_gpu, (*eigenvalue_iter).data(), this->nbase_x);
-
-            T* hcc_gpu = nullptr;
-            T* scc_gpu = nullptr;
-            T* vcc_gpu = nullptr;
-            base_device::memory::resize_memory_op<T, Device>()(hcc_gpu, nbase * nbase);
-            base_device::memory::resize_memory_op<T, Device>()(scc_gpu, nbase * nbase);
-            base_device::memory::resize_memory_op<T, Device>()(vcc_gpu, nbase * nbase);
-            for(int i=0;i<nbase;i++)
-            {
-                base_device::memory::synchronize_memory_op<T, Device, Device>()(hcc_gpu + i * nbase, hcc + i * nbase_x, nbase);
-                base_device::memory::synchronize_memory_op<T, Device, Device>()(scc_gpu + i * nbase, scc + i * nbase_x, nbase);
-            }
-            dngvd_op<T, Device>()(this->ctx, nbase, nbase, hcc_gpu, scc_gpu, eigenvalue_gpu, vcc_gpu);
-            for(int i=0;i<nbase;i++)
-            {
-                base_device::memory::synchronize_memory_op<T, Device, Device>()(vcc + i * nbase_x, vcc_gpu + i * nbase, nbase);
-            }
-            delmem_complex_op()(hcc_gpu);
-            delmem_complex_op()(scc_gpu);
-            delmem_complex_op()(vcc_gpu);
-
-            syncmem_var_d2h_op()((*eigenvalue_iter).data(), eigenvalue_gpu, this->nbase_x);
-
-            delmem_real_op()(eigenvalue_gpu);
+            base_device::memory::synchronize_memory_op<T, Device, Device>()(this->d_scc, scc, nbase * this->nbase_x);
+            dngvd_op<T, Device>()(this->ctx, nbase, this->nbase_x, this->hcc, this->d_scc, this->d_eigenvalue, this->vcc);
+            syncmem_var_d2h_op()((*eigenvalue_iter).data(), this->d_eigenvalue, this->nbase_x);
         }
 #endif
     }

source/source_hsolver/diago_dav_subspace.h

@@ -94,6 +94,7 @@ class Diago_DavSubspace
     /// Eigenvectors on the reduced basis
     T* vcc = nullptr;
 
+    T* d_scc = nullptr;
     Real* d_eigenvalue = nullptr;
 
     /// device type of psi

source/source_hsolver/kernels/cuda/dngvd_op.cu

@@ -216,7 +216,7 @@ struct dngvd_op<T, base_device::DEVICE_GPU>
                     Real* W,    // eigenvalue
                     T* V)
     {
-        assert(nstart == ldh);
+        // assert(nstart == ldh);
         // A to V
         cudaErrcheck(cudaMemcpy(V, A, sizeof(T) * ldh * nstart, cudaMemcpyDeviceToDevice));
         xhegvd_wrapper(CUBLAS_FILL_MODE_UPPER, nstart, V, ldh,

source/source_hsolver/kernels/rocm/dngvd_op.hip.cu

@@ -8,7 +8,7 @@ namespace hsolver {
 // NOTE: mimicked from ../cuda/dngvd_op.cu for three dngvd_op
 
 static hipsolverHandle_t hipsolver_H = nullptr;
-// Test on DCU platform. When nstart is greater than 234, code on DCU performs better. 
+// Test on DCU platform. When nstart is greater than 234, code on DCU performs better.
 const int N_DCU = 234;
 
 void createGpuSolverHandle() {
@@ -97,7 +97,7 @@ void dngvd_op<double, base_device::DEVICE_GPU>::operator()(const base_device::DE
         hipErrcheck(hipMemcpy(_eigenvalue, eigenvalue.data(), sizeof(double) * eigenvalue.size(), hipMemcpyHostToDevice));
     }
 
-    
+
 }
 #endif // __LCAO
 
@@ -112,7 +112,7 @@ void dngvd_op<std::complex<float>, base_device::DEVICE_GPU>::operator()(const ba
 {
     // copied from ../cuda/dngvd_op.cu, "dngvd_op"
     assert(nstart == ldh);
-    
+
     if (nstart > N_DCU){
         hipErrcheck(hipMemcpy(_vcc, _hcc, sizeof(std::complex<float>) * ldh * nstart, hipMemcpyDeviceToDevice));
         // now vcc contains hcc
@@ -170,7 +170,7 @@ void dngvd_op<std::complex<float>, base_device::DEVICE_GPU>::operator()(const ba
         hipErrcheck(hipMemcpy(_eigenvalue, eigenvalue.data(), sizeof(float) * eigenvalue.size(), hipMemcpyHostToDevice));
     }
 
-    
+
 }
 
 template <>
@@ -184,7 +184,7 @@ void dngvd_op<std::complex<double>, base_device::DEVICE_GPU>::operator()(const b
                                                                         )
 {
     // copied from ../cuda/dngvd_op.cu, "dngvd_op"
-    assert(nstart == ldh);
+    // assert(nstart == ldh);
 
     // save a copy of scc in case the diagonalization fails
     if (nstart > N_DCU){
@@ -253,7 +253,7 @@ void dngvd_op<std::complex<double>, base_device::DEVICE_GPU>::operator()(const b
 
 
 
-    
+
 }
 
 #ifdef __LCAO