From d965ba1338f1d00c2019daaa53445609d3ad78c0 Mon Sep 17 00:00:00 2001 From: Qianruipku Date: Fri, 21 Mar 2025 16:48:08 +0800 Subject: [PATCH 1/8] Feature: Make sKG support GPU --- source/module_base/blas_connector.cpp | 8 +- .../kernels/cuda/math_kernel_op_vec.cu | 10 +- source/module_base/kernels/math_kernel_op.h | 4 +- .../kernels/math_kernel_op_vec.cpp | 8 +- .../kernels/rocm/math_kernel_op_vec.hip.cu | 10 +- .../kernels/test/math_kernel_test.cpp | 28 +- source/module_base/math_chebyshev.cpp | 3 + source/module_base/para_gemm.cpp | 5 - source/module_base/parallel_device.h | 29 + source/module_basis/module_pw/pw_basis.cpp | 2 + source/module_basis/module_pw/pw_basis.h | 7 +- source/module_basis/module_pw/pw_basis_k.cpp | 37 +- source/module_basis/module_pw/pw_basis_k.h | 1 - .../potentials/potential_new.cpp | 82 +- .../potentials/potential_new.h | 1 + source/module_esolver/esolver_ks.cpp | 7 + source/module_esolver/esolver_sdft_pw.cpp | 42 +- .../hamilt_pwdft/VNL_in_pw.cpp | 201 +--- .../module_hamilt_pw/hamilt_pwdft/VNL_in_pw.h | 3 +- .../hamilt_pwdft/fs_nonlocal_tools.cpp | 2 +- .../hamilt_pwdft/operator_pw/op_exx_pw.h | 2 +- .../hamilt_pwdft/structure_factor.cpp | 26 +- .../hamilt_stodft/hamilt_sdft_pw.cpp | 9 +- .../hamilt_stodft/hamilt_sdft_pw.h | 19 +- .../hamilt_stodft/kernels/hpsi_norm_op.cpp | 2 +- .../hamilt_stodft/sto_dos.cpp | 55 +- .../module_hamilt_pw/hamilt_stodft/sto_dos.h | 3 +- .../hamilt_stodft/sto_elecond.cpp | 875 ++++++++++-------- .../hamilt_stodft/sto_elecond.h | 82 +- .../hamilt_stodft/sto_tool.cpp | 93 +- .../module_hamilt_pw/hamilt_stodft/sto_tool.h | 61 +- source/module_hsolver/diago_cg.cpp | 60 +- .../kernels/test/perf_math_kernel.cpp | 24 +- .../module_hsolver/para_linear_transform.cpp | 5 - source/module_io/read_set_globalv.cpp | 10 +- source/module_parameter/system_parameter.h | 10 +- 36 files changed, 924 insertions(+), 902 deletions(-) diff --git a/source/module_base/blas_connector.cpp b/source/module_base/blas_connector.cpp index 38448a18225..0cc4dc736eb 100644 --- a/source/module_base/blas_connector.cpp +++ b/source/module_base/blas_connector.cpp @@ -820,7 +820,7 @@ void vector_add_vector(const int& dim, float *result, const float *vector1, cons } else if (device_type == base_device::GpuDevice){ #ifdef __CUDA - ModuleBase::constantvector_addORsub_constantVector_op()(dim, result, vector1, constant1, vector2, constant2); + ModuleBase::vector_add_vector_op()(dim, result, vector1, constant1, vector2, constant2); #endif } } @@ -838,7 +838,7 @@ void vector_add_vector(const int& dim, double *result, const double *vector1, co } else if (device_type == base_device::GpuDevice){ #ifdef __CUDA - ModuleBase::constantvector_addORsub_constantVector_op()(dim, result, vector1, constant1, vector2, constant2); + ModuleBase::vector_add_vector_op()(dim, result, vector1, constant1, vector2, constant2); #endif } } @@ -856,7 +856,7 @@ void vector_add_vector(const int& dim, std::complex *result, const std::c } else if (device_type == base_device::GpuDevice){ #ifdef __CUDA - ModuleBase::constantvector_addORsub_constantVector_op, base_device::DEVICE_GPU>()(dim, result, vector1, constant1, vector2, constant2); + ModuleBase::vector_add_vector_op, base_device::DEVICE_GPU>()(dim, result, vector1, constant1, vector2, constant2); #endif } } @@ -874,7 +874,7 @@ void vector_add_vector(const int& dim, std::complex *result, const std:: } else if (device_type == base_device::GpuDevice){ #ifdef __CUDA - ModuleBase::constantvector_addORsub_constantVector_op, base_device::DEVICE_GPU>()(dim, result, vector1, constant1, vector2, constant2); + ModuleBase::vector_add_vector_op, base_device::DEVICE_GPU>()(dim, result, vector1, constant1, vector2, constant2); #endif } } \ No newline at end of file diff --git a/source/module_base/kernels/cuda/math_kernel_op_vec.cu b/source/module_base/kernels/cuda/math_kernel_op_vec.cu index ecc5e8b898e..c0a3b102d7f 100644 --- a/source/module_base/kernels/cuda/math_kernel_op_vec.cu +++ b/source/module_base/kernels/cuda/math_kernel_op_vec.cu @@ -225,7 +225,7 @@ void vector_div_vector_op, base_device::DEVICE_GPU>::operat // vector operator: result[i] = vector1[i] * constant1 + vector2[i] * constant2 template -void constantvector_addORsub_constantVector_op::operator()(const int& dim, +void vector_add_vector_op::operator()(const int& dim, T* result, const T* vector1, const Real constant1, @@ -314,10 +314,10 @@ template struct vector_div_vector_op, base_device::DEVICE_GP template struct vector_div_vector_op; template struct vector_div_vector_op, base_device::DEVICE_GPU>; -template struct constantvector_addORsub_constantVector_op; -template struct constantvector_addORsub_constantVector_op, base_device::DEVICE_GPU>; -template struct constantvector_addORsub_constantVector_op; -template struct constantvector_addORsub_constantVector_op, base_device::DEVICE_GPU>; +template struct vector_add_vector_op; +template struct vector_add_vector_op, base_device::DEVICE_GPU>; +template struct vector_add_vector_op; +template struct vector_add_vector_op, base_device::DEVICE_GPU>; template struct dot_real_op, base_device::DEVICE_GPU>; template struct dot_real_op; diff --git a/source/module_base/kernels/math_kernel_op.h b/source/module_base/kernels/math_kernel_op.h index fdb6d244572..bae77b23d67 100644 --- a/source/module_base/kernels/math_kernel_op.h +++ b/source/module_base/kernels/math_kernel_op.h @@ -134,7 +134,7 @@ template struct axpy_op { // vector operator: result[i] = vector1[i] * constant1 + vector2[i] * constant2 template -struct constantvector_addORsub_constantVector_op { +struct vector_add_vector_op { using Real = typename GetTypeReal::type; /// @brief result[i] = vector1[i] * constant1 + vector2[i] * constant2 /// @@ -315,7 +315,7 @@ template struct vector_div_vector_op { // vector operator: result[i] = vector1[i] * constant1 + vector2[i] * constant2 template -struct constantvector_addORsub_constantVector_op { +struct vector_add_vector_op { using Real = typename GetTypeReal::type; void operator()(const int &dim, T *result, const T *vector1, const Real constant1, const T *vector2, diff --git a/source/module_base/kernels/math_kernel_op_vec.cpp b/source/module_base/kernels/math_kernel_op_vec.cpp index 747857c5807..4030edbb481 100644 --- a/source/module_base/kernels/math_kernel_op_vec.cpp +++ b/source/module_base/kernels/math_kernel_op_vec.cpp @@ -92,7 +92,7 @@ struct axpy_op template -struct constantvector_addORsub_constantVector_op +struct vector_add_vector_op { using Real = typename GetTypeReal::type; void operator()(const int& dim, @@ -167,9 +167,9 @@ template struct axpy_op, base_device::DEVICE_CPU>; template struct axpy_op, base_device::DEVICE_CPU>; template struct axpy_op; -template struct constantvector_addORsub_constantVector_op, base_device::DEVICE_CPU>; -template struct constantvector_addORsub_constantVector_op; -template struct constantvector_addORsub_constantVector_op, base_device::DEVICE_CPU>; +template struct vector_add_vector_op, base_device::DEVICE_CPU>; +template struct vector_add_vector_op; +template struct vector_add_vector_op, base_device::DEVICE_CPU>; template struct dot_real_op, base_device::DEVICE_CPU>; template struct dot_real_op, base_device::DEVICE_CPU>; diff --git a/source/module_base/kernels/rocm/math_kernel_op_vec.hip.cu b/source/module_base/kernels/rocm/math_kernel_op_vec.hip.cu index 4ba8ae8925c..a88be379988 100644 --- a/source/module_base/kernels/rocm/math_kernel_op_vec.hip.cu +++ b/source/module_base/kernels/rocm/math_kernel_op_vec.hip.cu @@ -275,7 +275,7 @@ void vector_div_vector_op, base_device::DEVICE_GPU>::operat // vector operator: result[i] = vector1[i] * constant1 + vector2[i] * constant2 template -void constantvector_addORsub_constantVector_op::operator()(const int& dim, +void vector_add_vector_op::operator()(const int& dim, T* result, const T* vector1, const Real constant1, @@ -365,10 +365,10 @@ template struct vector_div_vector_op, base_device::DEVICE_GP template struct vector_div_vector_op; template struct vector_div_vector_op, base_device::DEVICE_GPU>; -template struct constantvector_addORsub_constantVector_op; -template struct constantvector_addORsub_constantVector_op, base_device::DEVICE_GPU>; -template struct constantvector_addORsub_constantVector_op; -template struct constantvector_addORsub_constantVector_op, base_device::DEVICE_GPU>; +template struct vector_add_vector_op; +template struct vector_add_vector_op, base_device::DEVICE_GPU>; +template struct vector_add_vector_op; +template struct vector_add_vector_op, base_device::DEVICE_GPU>; template struct dot_real_op, base_device::DEVICE_GPU>; template struct dot_real_op; diff --git a/source/module_base/kernels/test/math_kernel_test.cpp b/source/module_base/kernels/test/math_kernel_test.cpp index b8cfb68387e..38a7b6f9d97 100644 --- a/source/module_base/kernels/test/math_kernel_test.cpp +++ b/source/module_base/kernels/test/math_kernel_test.cpp @@ -75,8 +75,8 @@ class TestModuleHsolverMathKernel : public ::testing::Test using vector_mul_real_op_cpu = ModuleBase::vector_mul_real_op, base_device::DEVICE_CPU>; using vector_mul_vector_op_cpu = ModuleBase::vector_mul_vector_op, base_device::DEVICE_CPU>; using vector_div_vector_op_cpu = ModuleBase::vector_div_vector_op, base_device::DEVICE_CPU>; - using constantvector_addORsub_constantVector_op_cpu - = ModuleBase::constantvector_addORsub_constantVector_op, base_device::DEVICE_CPU>; + using vector_add_vector_op_cpu + = ModuleBase::vector_add_vector_op, base_device::DEVICE_CPU>; using axpy_op_cpu = ModuleBase::axpy_op, base_device::DEVICE_CPU>; using scal_op_cpu = ModuleBase::scal_op; using gemv_op_cpu = ModuleBase::gemv_op, base_device::DEVICE_CPU>; @@ -84,8 +84,8 @@ class TestModuleHsolverMathKernel : public ::testing::Test using vector_mul_real_op_gpu = ModuleBase::vector_mul_real_op, base_device::DEVICE_GPU>; using vector_mul_vector_op_gpu = ModuleBase::vector_mul_vector_op, base_device::DEVICE_GPU>; using vector_div_vector_op_gpu = ModuleBase::vector_div_vector_op, base_device::DEVICE_GPU>; - using constantvector_addORsub_constantVector_op_gpu - = ModuleBase::constantvector_addORsub_constantVector_op, base_device::DEVICE_GPU>; + using vector_add_vector_op_gpu + = ModuleBase::vector_add_vector_op, base_device::DEVICE_GPU>; using axpy_op_gpu = ModuleBase::axpy_op, base_device::DEVICE_GPU>; using scal_op_gpu = ModuleBase::scal_op; using gemv_op_gpu = ModuleBase::gemv_op, base_device::DEVICE_GPU>; @@ -174,12 +174,12 @@ class TestModuleHsolverMathKernel : public ::testing::Test {2.05256102, -1.39373474}, {-0.10166335, -0.49934031}}; - // (3) for test constantvector_addORsub_constantVector_op + // (3) for test vector_add_vector_op const double constant1 = 6.6; const double constant2 = 4.4; const std::vector> input1 = L; const std::vector> input2 = R; - const std::vector> output_constantvector_addORsub_constantVector_op + const std::vector> output_vector_add_vector_op = {{-5.05571797, -5.64586374}, {-14.76279273, 4.05181248}, {21.81709620, -17.11884992}, @@ -294,10 +294,10 @@ TEST_F(TestModuleHsolverMathKernel, vector_div_vector_op_cpu) } } -TEST_F(TestModuleHsolverMathKernel, constantvector_addORsub_constantVector_op_cpu) +TEST_F(TestModuleHsolverMathKernel, vector_add_vector_op_cpu) { std::vector> output(input.size()); - constantvector_addORsub_constantVector_op_cpu()(dim, + vector_add_vector_op_cpu()(dim, output.data(), input1.data(), constant1, @@ -305,8 +305,8 @@ TEST_F(TestModuleHsolverMathKernel, constantvector_addORsub_constantVector_op_cp constant2); for (int i = 0; i < input.size(); i++) { - EXPECT_LT(fabs(output[i].imag() - output_constantvector_addORsub_constantVector_op[i].imag()), 1e-8); - EXPECT_LT(fabs(output[i].real() - output_constantvector_addORsub_constantVector_op[i].real()), 1e-8); + EXPECT_LT(fabs(output[i].imag() - output_vector_add_vector_op[i].imag()), 1e-8); + EXPECT_LT(fabs(output[i].real() - output_vector_add_vector_op[i].real()), 1e-8); } } @@ -478,7 +478,7 @@ TEST_F(TestModuleHsolverMathKernel, vector_div_vector_op_gpu) delete_memory_op()(output_dev); } -TEST_F(TestModuleHsolverMathKernel, constantvector_addORsub_constantVector_op_gpu) +TEST_F(TestModuleHsolverMathKernel, vector_add_vector_op_gpu) { // in CPU std::vector> output(input.size()); @@ -498,7 +498,7 @@ TEST_F(TestModuleHsolverMathKernel, constantvector_addORsub_constantVector_op_gp synchronize_memory_op()(input2_dev, input2.data(), input.size()); // run - constantvector_addORsub_constantVector_op_gpu()(dim, + vector_add_vector_op_gpu()(dim, output_dev, input1_dev, constant1, @@ -510,8 +510,8 @@ TEST_F(TestModuleHsolverMathKernel, constantvector_addORsub_constantVector_op_gp for (int i = 0; i < input.size(); i++) { - EXPECT_LT(fabs(output[i].imag() - output_constantvector_addORsub_constantVector_op[i].imag()), 1e-8); - EXPECT_LT(fabs(output[i].real() - output_constantvector_addORsub_constantVector_op[i].real()), 1e-8); + EXPECT_LT(fabs(output[i].imag() - output_vector_add_vector_op[i].imag()), 1e-8); + EXPECT_LT(fabs(output[i].real() - output_vector_add_vector_op[i].real()), 1e-8); } delete_memory_op()(input1_dev); diff --git a/source/module_base/math_chebyshev.cpp b/source/module_base/math_chebyshev.cpp index b2cc6aadea5..60bb560f05b 100644 --- a/source/module_base/math_chebyshev.cpp +++ b/source/module_base/math_chebyshev.cpp @@ -767,6 +767,9 @@ template class Chebyshev; #endif #if ((defined __CUDA) || (defined __ROCM)) template class Chebyshev; +#ifdef __ENABLE_FLOAT_FFTW +template class Chebyshev; +#endif #endif } // namespace ModuleBase diff --git a/source/module_base/para_gemm.cpp b/source/module_base/para_gemm.cpp index 79fdbf0de7e..abfacf5154a 100644 --- a/source/module_base/para_gemm.cpp +++ b/source/module_base/para_gemm.cpp @@ -256,11 +256,6 @@ void PGemmCN::multiply_col(const T alpha, const T* A, const T* B, con int m = colA_loc[ip]; int size = m * LDA; MPI_Status status; -#ifdef __CUDA_MPI - // If the memory is not set to zero, it may cause the result to be wrong when using CUDA Aware MPI - // I am not sure if it is due to CUDA Aware MPI or not - base_device::memory::set_memory_op()(Atmp_device, 0, size); -#endif Parallel_Common::recv_dev(Atmp_device, size, ip, 0, col_world, &status, A_tmp_.data()); MPI_Wait(&requests[ip], &status); ModuleBase::gemm_op()('C', diff --git a/source/module_base/parallel_device.h b/source/module_base/parallel_device.h index fbf127d07d5..3c7094f6059 100644 --- a/source/module_base/parallel_device.h +++ b/source/module_base/parallel_device.h @@ -144,6 +144,35 @@ void reduce_dev(T* object, const int& n, const MPI_Comm& comm, T* tmp_space = nu #endif return; } + +template +void gatherv_dev(const T* sendbuf, + int sendcount, + T* recvbuf, + const int* recvcounts, + const int* displs, + MPI_Comm& comm, + T* tmp_sspace = nullptr, + T* tmp_rspace = nullptr) +{ +#ifdef __CUDA_MPI + gatherv_data(sendbuf, sendcount, recvbuf, recvcounts, displs, comm); +#else + object_cpu_point o1, o2; + int size = 0; + MPI_Comm_size(comm, &size); + int gather_space = displs[size - 1] + recvcounts[size - 1]; + T* sendbuf_cpu = o1.get(sendbuf, sendcount, tmp_sspace); + T* recvbuf_cpu = o2.get(recvbuf, gather_space, tmp_rspace); + o1.sync_d2h(sendbuf_cpu, sendbuf, sendcount); + gatherv_data(sendbuf_cpu, sendcount, recvbuf, recvcounts, displs, comm); + o2.sync_h2d(recvbuf, recvbuf_cpu, gather_space); + o1.del(sendbuf_cpu); + o2.del(recvbuf_cpu); +#endif + return; +} + } diff --git a/source/module_basis/module_pw/pw_basis.cpp b/source/module_basis/module_pw/pw_basis.cpp index d5183db90c8..c7cd1cc1f04 100644 --- a/source/module_basis/module_pw/pw_basis.cpp +++ b/source/module_basis/module_pw/pw_basis.cpp @@ -16,6 +16,8 @@ PW_Basis::PW_Basis() PW_Basis::PW_Basis(std::string device_, std::string precision_) : device(std::move(device_)), precision(std::move(precision_)) { classname="PW_Basis"; this->fft_bundle.setfft("cpu",this->precision); + this->double_data_ = (this->precision == "double") || (this->precision == "mixing"); + this->float_data_ = (this->precision == "single") || (this->precision == "mixing"); } PW_Basis:: ~PW_Basis() diff --git a/source/module_basis/module_pw/pw_basis.h b/source/module_basis/module_pw/pw_basis.h index 47827000eb7..d988305d31f 100644 --- a/source/module_basis/module_pw/pw_basis.h +++ b/source/module_basis/module_pw/pw_basis.h @@ -289,10 +289,11 @@ class PW_Basis void set_precision(std::string precision_); protected: - std::string device = "cpu"; - std::string precision = "double"; + std::string device = "cpu"; ///< cpu or gpu + std::string precision = "double"; ///< single, double, mixing + bool double_data_ = true; ///< if has double data + bool float_data_ = false; ///< if has float data }; - } #endif // PWBASIS_H diff --git a/source/module_basis/module_pw/pw_basis_k.cpp b/source/module_basis/module_pw/pw_basis_k.cpp index 08391242eae..91343d61a49 100644 --- a/source/module_basis/module_pw/pw_basis_k.cpp +++ b/source/module_basis/module_pw/pw_basis_k.cpp @@ -25,17 +25,11 @@ PW_Basis_K::~PW_Basis_K() #if defined(__CUDA) || defined(__ROCM) if (this->device == "gpu") { - if (this->precision == "single") - { - delmem_sd_op()(this->s_kvec_c); - delmem_sd_op()(this->s_gcar); - delmem_sd_op()(this->s_gk2); - } - else - { - delmem_dd_op()(this->d_gcar); - delmem_dd_op()(this->d_gk2); - } + delmem_sd_op()(this->s_kvec_c); + delmem_sd_op()(this->s_gcar); + delmem_sd_op()(this->s_gk2); + delmem_dd_op()(this->d_gcar); + delmem_dd_op()(this->d_gk2); delmem_dd_op()(this->d_kvec_c); delmem_int_op()(this->ig2ixyz_k); delmem_int_op()(this->d_igl2isz_k); @@ -43,12 +37,9 @@ PW_Basis_K::~PW_Basis_K() else { #endif - if (this->precision == "single") - { - delmem_sh_op()(this->s_kvec_c); - delmem_sh_op()(this->s_gcar); - delmem_sh_op()(this->s_gk2); - } + delmem_sh_op()(this->s_kvec_c); + delmem_sh_op()(this->s_gcar); + delmem_sh_op()(this->s_gk2); // There's no need to delete double pointers while in a CPU environment. #if defined(__CUDA) || defined(__ROCM) } @@ -113,7 +104,7 @@ void PW_Basis_K::initparameters(const bool gamma_only_in, #if defined(__CUDA) || defined(__ROCM) if (this->device == "gpu") { - if (this->precision == "single") + if (this->float_data_) { resmem_sd_op()(this->s_kvec_c, this->nks * 3); castmem_d2s_h2d_op()(this->s_kvec_c, reinterpret_cast(&this->kvec_c[0][0]), this->nks * 3); @@ -124,7 +115,7 @@ void PW_Basis_K::initparameters(const bool gamma_only_in, else { #endif - if (this->precision == "single") + if (this->float_data_) { resmem_sh_op()(this->s_kvec_c, this->nks * 3); castmem_d2s_h2h_op()(this->s_kvec_c, reinterpret_cast(&this->kvec_c[0][0]), this->nks * 3); @@ -307,7 +298,7 @@ void PW_Basis_K::collect_local_pw(const double& erf_ecut_in, const double& erf_h #if defined(__CUDA) || defined(__ROCM) if (this->device == "gpu") { - if (this->precision == "single") + if (this->float_data_) { resmem_sd_op()(this->s_gk2, this->npwk_max * this->nks); resmem_sd_op()(this->s_gcar, this->npwk_max * this->nks * 3); @@ -316,7 +307,7 @@ void PW_Basis_K::collect_local_pw(const double& erf_ecut_in, const double& erf_h reinterpret_cast(&this->gcar[0][0]), this->npwk_max * this->nks * 3); } - else + if (this->double_data_) { resmem_dd_op()(this->d_gk2, this->npwk_max * this->nks); resmem_dd_op()(this->d_gcar, this->npwk_max * this->nks * 3); @@ -329,7 +320,7 @@ void PW_Basis_K::collect_local_pw(const double& erf_ecut_in, const double& erf_h else { #endif - if (this->precision == "single") + if (this->float_data_) { resmem_sh_op()(this->s_gk2, this->npwk_max * this->nks, "PW_B_K::s_gk2"); resmem_sh_op()(this->s_gcar, this->npwk_max * this->nks * 3, "PW_B_K::s_gcar"); @@ -338,7 +329,7 @@ void PW_Basis_K::collect_local_pw(const double& erf_ecut_in, const double& erf_h reinterpret_cast(&this->gcar[0][0]), this->npwk_max * this->nks * 3); } - else + if (this->double_data_) { this->d_gcar = reinterpret_cast(&this->gcar[0][0]); this->d_gk2 = this->gk2; diff --git a/source/module_basis/module_pw/pw_basis_k.h b/source/module_basis/module_pw/pw_basis_k.h index ae5076bba9e..defa0cb5cc5 100644 --- a/source/module_basis/module_pw/pw_basis_k.h +++ b/source/module_basis/module_pw/pw_basis_k.h @@ -205,7 +205,6 @@ class PW_Basis_K : public PW_Basis private: float * s_gcar = nullptr, * s_kvec_c = nullptr; double * d_gcar = nullptr, * d_kvec_c = nullptr; - }; } diff --git a/source/module_elecstate/potentials/potential_new.cpp b/source/module_elecstate/potentials/potential_new.cpp index 0f6b7fd9a15..617e18a1788 100644 --- a/source/module_elecstate/potentials/potential_new.cpp +++ b/source/module_elecstate/potentials/potential_new.cpp @@ -34,6 +34,7 @@ Potential::Potential(const ModulePW::PW_Basis* rho_basis_in, this->rho_basis_smooth_ = rho_basis_smooth_in; this->fixed_mode = true; this->dynamic_mode = true; + this->use_gpu_ = (PARAM.inp.basis_type == "pw" && PARAM.inp.device == "gpu"); // allocate memory for Potential. this->allocate(); @@ -49,21 +50,17 @@ Potential::~Potential() } this->components.clear(); } - if (PARAM.inp.basis_type == "pw" && PARAM.inp.device == "gpu") { - if (PARAM.inp.precision == "single") { - delmem_sd_op()(s_veff_smooth); - delmem_sd_op()(s_vofk_smooth); - } - else { - delmem_dd_op()(d_veff_smooth); - delmem_dd_op()(d_vofk_smooth); - } + if (use_gpu_) + { + delmem_sd_op()(s_veff_smooth); + delmem_sd_op()(s_vofk_smooth); + delmem_dd_op()(d_veff_smooth); + delmem_dd_op()(d_vofk_smooth); } - else { - if (PARAM.inp.precision == "single") { - delmem_sh_op()(s_veff_smooth); - delmem_sh_op()(s_vofk_smooth); - } + else + { + delmem_sh_op()(s_veff_smooth); + delmem_sh_op()(s_vofk_smooth); } } @@ -134,22 +131,28 @@ void Potential::allocate() this->vofk_smooth.create(PARAM.inp.nspin, nrxx_smooth); ModuleBase::Memory::record("Pot::vofk_smooth", sizeof(double) * PARAM.inp.nspin * nrxx_smooth); } - if (PARAM.inp.basis_type == "pw" && PARAM.inp.device == "gpu") { - if (PARAM.inp.precision == "single") { + if (use_gpu_) + { + if (PARAM.globalv.has_float_data) + { resmem_sd_op()(s_veff_smooth, PARAM.inp.nspin * nrxx_smooth); resmem_sd_op()(s_vofk_smooth, PARAM.inp.nspin * nrxx_smooth); } - else { + if (PARAM.globalv.has_double_data) + { resmem_dd_op()(d_veff_smooth, PARAM.inp.nspin * nrxx_smooth); resmem_dd_op()(d_vofk_smooth, PARAM.inp.nspin * nrxx_smooth); } } - else { - if (PARAM.inp.precision == "single") { + else + { + if (PARAM.globalv.has_float_data) + { resmem_sh_op()(s_veff_smooth, PARAM.inp.nspin * nrxx_smooth, "POT::sveff_smooth"); resmem_sh_op()(s_vofk_smooth, PARAM.inp.nspin * nrxx_smooth, "POT::svofk_smooth"); } - else { + if (PARAM.globalv.has_double_data) + { this->d_veff_smooth = this->veff_smooth.c; this->d_vofk_smooth = this->vofk_smooth.c; } @@ -182,32 +185,25 @@ void Potential::update_from_charge(const Charge*const chg, const UnitCell*const } #endif - if (PARAM.inp.basis_type == "pw" && PARAM.inp.device == "gpu") { - if (PARAM.inp.precision == "single") { - castmem_d2s_h2d_op()(s_veff_smooth, - this->veff_smooth.c, - this->veff_smooth.nr * this->veff_smooth.nc); - castmem_d2s_h2d_op()(s_vofk_smooth, - this->vofk_smooth.c, - this->vofk_smooth.nr * this->vofk_smooth.nc); + if (this->use_gpu_) + { + if (PARAM.globalv.has_float_data) + { + castmem_d2s_h2d_op()(s_veff_smooth, this->veff_smooth.c, this->veff_smooth.nr * this->veff_smooth.nc); + castmem_d2s_h2d_op()(s_vofk_smooth, this->vofk_smooth.c, this->vofk_smooth.nr * this->vofk_smooth.nc); } - else { - syncmem_d2d_h2d_op()(d_veff_smooth, - this->veff_smooth.c, - this->veff_smooth.nr * this->veff_smooth.nc); - syncmem_d2d_h2d_op()(d_vofk_smooth, - this->vofk_smooth.c, - this->vofk_smooth.nr * this->vofk_smooth.nc); + if (PARAM.globalv.has_double_data) + { + syncmem_d2d_h2d_op()(d_veff_smooth, this->veff_smooth.c, this->veff_smooth.nr * this->veff_smooth.nc); + syncmem_d2d_h2d_op()(d_vofk_smooth, this->vofk_smooth.c, this->vofk_smooth.nr * this->vofk_smooth.nc); } } - else { - if (PARAM.inp.precision == "single") { - castmem_d2s_h2h_op()(s_veff_smooth, - this->veff_smooth.c, - this->veff_smooth.nr * this->veff_smooth.nc); - castmem_d2s_h2h_op()(s_vofk_smooth, - this->vofk_smooth.c, - this->vofk_smooth.nr * this->vofk_smooth.nc); + else + { + if (PARAM.globalv.has_float_data) + { + castmem_d2s_h2h_op()(s_veff_smooth, this->veff_smooth.c, this->veff_smooth.nr * this->veff_smooth.nc); + castmem_d2s_h2h_op()(s_vofk_smooth, this->vofk_smooth.c, this->vofk_smooth.nr * this->vofk_smooth.nc); } // There's no need to synchronize memory for double precision pointers while in a CPU environment } diff --git a/source/module_elecstate/potentials/potential_new.h b/source/module_elecstate/potentials/potential_new.h index 5d421548b32..103a46e0647 100644 --- a/source/module_elecstate/potentials/potential_new.h +++ b/source/module_elecstate/potentials/potential_new.h @@ -220,6 +220,7 @@ class Potential : public PotBase const ModuleBase::matrix* vloc_ = nullptr; Structure_Factor* structure_factors_ = nullptr; surchem* solvent_ = nullptr; + bool use_gpu_ = false; }; } // namespace elecstate diff --git a/source/module_esolver/esolver_ks.cpp b/source/module_esolver/esolver_ks.cpp index 9d485f43063..cc9ee5fe5a2 100644 --- a/source/module_esolver/esolver_ks.cpp +++ b/source/module_esolver/esolver_ks.cpp @@ -59,6 +59,13 @@ ESolver_KS::ESolver_KS() { fft_device = "cpu"; } + std::string fft_precision = PARAM.inp.precision; +#ifdef __ENABLE_FLOAT_FFTW + if (PARAM.inp.cal_cond && PARAM.inp.esolver_type == "sdft") + { + fft_precision = "mixing"; + } +#endif pw_wfc = new ModulePW::PW_Basis_K_Big(fft_device, PARAM.inp.precision); ModulePW::PW_Basis_K_Big* tmp = static_cast(pw_wfc); diff --git a/source/module_esolver/esolver_sdft_pw.cpp b/source/module_esolver/esolver_sdft_pw.cpp index e013372a119..d1c9da420ed 100644 --- a/source/module_esolver/esolver_sdft_pw.cpp +++ b/source/module_esolver/esolver_sdft_pw.cpp @@ -262,17 +262,6 @@ void ESolver_SDFT_PW::after_all_runners(UnitCell& ucell) GlobalV::ofs_running << " !FINAL_ETOT_IS " << this->pelec->f_en.etot * ModuleBase::Ry_to_eV << " eV" << std::endl; GlobalV::ofs_running << " --------------------------------------------\n\n" << std::endl; ModuleIO::write_istate_info(this->pelec->ekb, this->pelec->wg, this->kv); -} - -template <> -void ESolver_SDFT_PW, base_device::DEVICE_CPU>::after_all_runners(UnitCell& ucell) -{ - - GlobalV::ofs_running << "\n\n --------------------------------------------" << std::endl; - GlobalV::ofs_running << std::setprecision(16); - GlobalV::ofs_running << " !FINAL_ETOT_IS " << this->pelec->f_en.etot * ModuleBase::Ry_to_eV << " eV" << std::endl; - GlobalV::ofs_running << " --------------------------------------------\n\n" << std::endl; - ModuleIO::write_istate_info(this->pelec->ekb, this->pelec->wg, this->kv); if (this->method_sto == 2) { @@ -280,7 +269,18 @@ void ESolver_SDFT_PW, base_device::DEVICE_CPU>::after_all_r } if (PARAM.inp.out_dos) { - Sto_DOS sto_dos(this->pw_wfc, &this->kv, this->pelec, this->psi, this->p_hamilt, this->stoche, &stowf); + if(!std::is_same>::value || !std::is_same::value) + { + ModuleBase::WARNING_QUIT("ESolver_SDFT_PW", "DOS does not support complex float or GPU yet."); + } + Sto_DOS sto_dos( + this->pw_wfc, + &this->kv, + this->pelec, + reinterpret_cast>*>(this->psi), + reinterpret_cast>*>(this->p_hamilt), + this->stoche, + reinterpret_cast, base_device::DEVICE_CPU>*>(&stowf)); sto_dos.decide_param(PARAM.inp.dos_nche, PARAM.inp.emin_sto, PARAM.inp.emax_sto, @@ -295,15 +295,15 @@ void ESolver_SDFT_PW, base_device::DEVICE_CPU>::after_all_r // sKG cost memory, and it should be placed at the end of the program if (PARAM.inp.cal_cond) { - Sto_EleCond sto_elecond(&ucell, - &this->kv, - this->pelec, - this->pw_wfc, - this->psi, - &this->ppcell, - this->p_hamilt, - this->stoche, - &stowf); + Sto_EleCond sto_elecond(&ucell, + &this->kv, + this->pelec, + this->pw_wfc, + this->kspw_psi, + &this->ppcell, + this->p_hamilt, + this->stoche, + &stowf); sto_elecond.decide_nche(PARAM.inp.cond_dt, 1e-8, this->nche_sto, PARAM.inp.emin_sto, PARAM.inp.emax_sto); sto_elecond.sKG(PARAM.inp.cond_smear, PARAM.inp.cond_fwhm, diff --git a/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp b/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp index bcd0cba74da..83225843778 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp @@ -17,6 +17,7 @@ pseudopot_cell_vnl::pseudopot_cell_vnl() { + this->use_gpu_ = (this->use_gpu_); } pseudopot_cell_vnl::~pseudopot_cell_vnl() @@ -32,25 +33,19 @@ void pseudopot_cell_vnl::release_memory() if (this->nhm <= 0 || memory_released) { return; } - if (PARAM.inp.device == "gpu") + if (this->use_gpu_) { - if (PARAM.inp.precision == "single") - { - delmem_sd_op()(this->s_deeq); - delmem_sd_op()(this->s_nhtol); - delmem_sd_op()(this->s_nhtolm); - delmem_sd_op()(this->s_indv); - delmem_sd_op()(this->s_tab); - delmem_sd_op()(this->s_qq_nt); - delmem_cd_op()(this->c_deeq_nc); - delmem_cd_op()(this->c_vkb); - delmem_cd_op()(this->c_qq_so); - } - else - { - delmem_zd_op()(this->z_deeq_nc); - delmem_zd_op()(this->z_qq_so); - } + delmem_sd_op()(this->s_deeq); + delmem_sd_op()(this->s_nhtol); + delmem_sd_op()(this->s_nhtolm); + delmem_sd_op()(this->s_indv); + delmem_sd_op()(this->s_tab); + delmem_sd_op()(this->s_qq_nt); + delmem_cd_op()(this->c_deeq_nc); + delmem_cd_op()(this->c_vkb); + delmem_cd_op()(this->c_qq_so); + delmem_zd_op()(this->z_deeq_nc); + delmem_zd_op()(this->z_qq_so); delmem_dd_op()(this->d_deeq); delmem_zd_op()(this->z_vkb); delmem_dd_op()(this->d_tab); @@ -61,18 +56,15 @@ void pseudopot_cell_vnl::release_memory() } else { - if (PARAM.inp.precision == "single") - { - delmem_sh_op()(this->s_deeq); - delmem_sh_op()(this->s_nhtol); - delmem_sh_op()(this->s_nhtolm); - delmem_sh_op()(this->s_indv); - delmem_sh_op()(this->s_tab); - delmem_sh_op()(this->s_qq_nt); - delmem_ch_op()(this->c_deeq_nc); - delmem_ch_op()(this->c_vkb); - delmem_ch_op()(this->c_qq_so); - } + delmem_sh_op()(this->s_deeq); + delmem_sh_op()(this->s_nhtol); + delmem_sh_op()(this->s_nhtolm); + delmem_sh_op()(this->s_indv); + delmem_sh_op()(this->s_tab); + delmem_sh_op()(this->s_qq_nt); + delmem_ch_op()(this->c_deeq_nc); + delmem_ch_op()(this->c_vkb); + delmem_ch_op()(this->c_qq_so); // There's no need to delete double precision pointers while in a CPU environment. } memory_released = true; @@ -154,9 +146,9 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, this->deeq_nc.create(PARAM.inp.nspin, ucell.nat, this->nhm, this->nhm); this->qq_nt.create(ntype, this->nhm, this->nhm); this->qq_so.create(ntype, 4, this->nhm, this->nhm); - if (PARAM.inp.device == "gpu") + if (this->use_gpu_) { - if (PARAM.inp.precision == "single") + if (PARAM.globalv.has_float_data) { resmem_sd_op()(s_deeq, PARAM.inp.nspin * ucell.nat * this->nhm * this->nhm); resmem_sd_op()(s_nhtol, ntype * this->nhm); @@ -166,7 +158,7 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, resmem_cd_op()(c_deeq_nc, PARAM.inp.nspin * ucell.nat * this->nhm * this->nhm); resmem_cd_op()(c_qq_so, ntype * 4 * this->nhm * this->nhm); } - else + if (PARAM.globalv.has_double_data) { resmem_zd_op()(z_deeq_nc, PARAM.inp.nspin * ucell.nat * this->nhm * this->nhm); resmem_zd_op()(z_qq_so, ntype * 4 * this->nhm * this->nhm); @@ -179,7 +171,7 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, } else { - if (PARAM.inp.precision == "single") + if (PARAM.globalv.has_float_data) { resmem_sh_op()(s_deeq, PARAM.inp.nspin * ucell.nat * this->nhm * this->nhm, @@ -193,7 +185,7 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, "VNL::c_deeq_nc"); resmem_ch_op()(c_qq_so, ntype * 4 * this->nhm * this->nhm, "VNL::c_qq_so"); } - else + if (PARAM.globalv.has_double_data) { this->z_deeq_nc = this->deeq_nc.ptr; this->z_qq_so = this->qq_so.ptr; @@ -269,9 +261,9 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, ModuleBase::Memory::record("VNL::tab_at", ntype * nchix_nc * PARAM.globalv.nqx * sizeof(double)); } } - if (PARAM.inp.device == "gpu") + if (this->use_gpu_) { - if (PARAM.inp.precision == "single") + if (PARAM.globalv.has_float_data) { resmem_sd_op()(s_tab, this->tab.getSize()); resmem_cd_op()(c_vkb, nkb * npwx); @@ -281,7 +273,7 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, } else { - if (PARAM.inp.precision == "single") + if (PARAM.globalv.has_float_data) { resmem_sh_op()(s_tab, this->tab.getSize()); resmem_ch_op()(c_vkb, nkb * npwx); @@ -299,115 +291,6 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, // Calculates beta functions (Kleinman-Bylander projectors), // with structure factor, for all atoms, in reciprocal space //---------------------------------------------------------- -void pseudopot_cell_vnl::getvnl(const int& ik, const UnitCell& ucell, ModuleBase::ComplexMatrix& vkb_in) const -{ - if (PARAM.inp.use_paw) { - return; -} - if (PARAM.inp.test_pp) { - ModuleBase::TITLE("pseudopot_cell_vnl", "getvnl"); -} - ModuleBase::timer::tick("pp_cell_vnl", "getvnl"); - - if (lmaxkb < 0) - { - return; - } - - const int npw = this->wfcpw->npwk[ik]; - - // When the internal memory is large enough, it is better to make vkb1 be the number of pseudopot_cell_vnl. - // We only need to initialize it once as long as the cell is unchanged. - ModuleBase::matrix vkb1(nhm, npw); - double* vq = new double[npw]; - const int x1 = (lmaxkb + 1) * (lmaxkb + 1); - - ModuleBase::matrix ylm(x1, npw); - ModuleBase::Memory::record("VNL::ylm", x1 * npw * sizeof(double)); - ModuleBase::Vector3* gk = new ModuleBase::Vector3[npw]; - for (int ig = 0; ig < npw; ig++) - { - gk[ig] = this->wfcpw->getgpluskcar(ik, ig); - } - - ModuleBase::YlmReal::Ylm_Real(cpu_ctx, x1, npw, reinterpret_cast(gk), ylm.c); - - using Device = base_device::DEVICE_CPU; - Device* ctx = {}; - using resmem_complex_op = base_device::memory::resize_memory_op, Device>; - using delmem_complex_op = base_device::memory::delete_memory_op, Device>; - std::complex* sk = nullptr; - resmem_complex_op()(sk, ucell.nat * npw, "VNL::sk"); - this->psf->get_sk(ctx, ik, this->wfcpw, sk); - - int jkb = 0, iat = 0; - for (int it = 0; it < ucell.ntype; it++) - { - // calculate beta in G-space using an interpolation table - const int nbeta = ucell.atoms[it].ncpp.nbeta; - const int nh = ucell.atoms[it].ncpp.nh; - - if (PARAM.inp.test_pp > 1) { - ModuleBase::GlobalFunc::OUT("nbeta", nbeta); -} - - for (int nb = 0; nb < nbeta; nb++) - { - if (PARAM.inp.test_pp > 1) { - ModuleBase::GlobalFunc::OUT("ib", nb); -} - for (int ig = 0; ig < npw; ig++) - { - const double gnorm = gk[ig].norm() * ucell.tpiba; - - vq[ig] = ModuleBase::PolyInt::Polynomial_Interpolation(this->tab, - it, - nb, - PARAM.globalv.nqx, - PARAM.globalv.dq, - gnorm); - } - - // add spherical harmonic part - for (int ih = 0; ih < nh; ih++) - { - if (nb == this->indv(it, ih)) - { - const int lm = static_cast(nhtolm(it, ih)); - for (int ig = 0; ig < npw; ig++) - { - vkb1(ih, ig) = ylm(lm, ig) * vq[ig]; - } - } - } // end ih - } // end nbeta - - // vkb1 contains all betas including angular part for type nt - // now add the structure factor and factor (-i)^l - for (int ia = 0; ia < ucell.atoms[it].na; ia++) - { - for (int ih = 0; ih < nh; ih++) - { - std::complex pref = pow(ModuleBase::NEG_IMAG_UNIT, nhtol(it, ih)); //? - std::complex* pvkb = &vkb_in(jkb, 0); - for (int ig = 0; ig < npw; ig++) - { - pvkb[ig] = vkb1(ih, ig) * sk[iat * npw + ig] * pref; - } - ++jkb; - } // end ih - iat++; - } // end ia - } // enddo - - delete[] gk; - delete[] vq; - delmem_complex_op()(sk); - ModuleBase::timer::tick("pp_cell_vnl", "getvnl"); - - return; -} // end subroutine getvnl - template void pseudopot_cell_vnl::getvnl(Device* ctx, const UnitCell& ucell, @@ -466,7 +349,7 @@ void pseudopot_cell_vnl::getvnl(Device* ctx, { _gk[ig] = this->wfcpw->getgpluskcar(ik, ig); } - if (PARAM.inp.device == "gpu") + if (this->use_gpu_) { resmem_int_op()(atom_nh, ucell.ntype); resmem_int_op()(atom_nb, ucell.ntype); @@ -483,7 +366,7 @@ void pseudopot_cell_vnl::getvnl(Device* ctx, atom_nh = h_atom_nh; atom_nb = h_atom_nb; atom_na = h_atom_na; - if (PARAM.inp.precision == "single") + if (std::is_same::value) { resmem_var_op()(gk, npw * 3); castmem_var_h2h_op()(gk, reinterpret_cast(_gk), npw * 3); @@ -530,11 +413,11 @@ void pseudopot_cell_vnl::getvnl(Device* ctx, delmem_var_op()(ylm); delmem_var_op()(vkb1); delmem_complex_op()(sk); - if (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single") + if (this->use_gpu_ || std::is_same::value) { delmem_var_op()(gk); } - if (PARAM.inp.device == "gpu") + if (this->use_gpu_) { delmem_int_op()(atom_nh); delmem_int_op()(atom_nb); @@ -868,9 +751,9 @@ void pseudopot_cell_vnl::init_vnl(UnitCell& cell, const ModulePW::PW_Basis* rho_ delete[] aux; delete[] jl; } - if (PARAM.inp.device == "gpu") + if (this->use_gpu_) { - if (PARAM.inp.precision == "single") + if (PARAM.globalv.has_float_data) { castmem_d2s_h2d_op()(this->s_indv, this->indv.c, this->indv.nr * this->indv.nc); castmem_d2s_h2d_op()(this->s_nhtol, this->nhtol.c, this->nhtol.nr * this->nhtol.nc); @@ -879,7 +762,7 @@ void pseudopot_cell_vnl::init_vnl(UnitCell& cell, const ModulePW::PW_Basis* rho_ castmem_d2s_h2d_op()(this->s_qq_nt, this->qq_nt.ptr, this->qq_nt.getSize()); castmem_z2c_h2d_op()(this->c_qq_so, this->qq_so.ptr, this->qq_so.getSize()); } - else + if (PARAM.globalv.has_double_data) { syncmem_z2z_h2d_op()(this->z_qq_so, this->qq_so.ptr, this->qq_so.getSize()); } @@ -894,7 +777,7 @@ void pseudopot_cell_vnl::init_vnl(UnitCell& cell, const ModulePW::PW_Basis* rho_ } else { - if (PARAM.inp.precision == "single") + if (PARAM.globalv.has_float_data) { castmem_d2s_h2h_op()(this->s_indv, this->indv.c, this->indv.nr * this->indv.nc); castmem_d2s_h2h_op()(this->s_nhtol, this->nhtol.c, this->nhtol.nr * this->nhtol.nc); @@ -1486,9 +1369,9 @@ void pseudopot_cell_vnl::cal_effective_D(const ModuleBase::matrix& veff, } } } - if (PARAM.inp.device == "gpu") + if (this->use_gpu_) { - if (PARAM.inp.precision == "single") + if (PARAM.globalv.has_float_data) { castmem_d2s_h2d_op()(this->s_deeq, this->deeq.ptr, @@ -1497,7 +1380,7 @@ void pseudopot_cell_vnl::cal_effective_D(const ModuleBase::matrix& veff, this->deeq_nc.ptr, PARAM.inp.nspin * cell.nat * this->nhm * this->nhm); } - else + if (PARAM.globalv.has_double_data) { syncmem_z2z_h2d_op()(this->z_deeq_nc, this->deeq_nc.ptr, @@ -1509,7 +1392,7 @@ void pseudopot_cell_vnl::cal_effective_D(const ModuleBase::matrix& veff, } else { - if (PARAM.inp.precision == "single") + if (PARAM.globalv.has_float_data) { castmem_d2s_h2h_op()(this->s_deeq, this->deeq.ptr, diff --git a/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.h b/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.h index e5b80ed792e..1e9fdaa7622 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.h +++ b/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.h @@ -41,8 +41,6 @@ class pseudopot_cell_vnl template void getvnl(Device* ctx, const UnitCell& ucell, const int& ik, std::complex* vkb_in) const; - void getvnl(const int& ik, const UnitCell& ucell, ModuleBase::ComplexMatrix& vkb_in) const; - // void getvnl_alpha(const int &ik); void init_vnl_alpha(const UnitCell& cell); @@ -203,6 +201,7 @@ class pseudopot_cell_vnl Soc soc; double omega_old = 0; + bool use_gpu_ = false; /** * @brief Compute interpolation table qrad diff --git a/source/module_hamilt_pw/hamilt_pwdft/fs_nonlocal_tools.cpp b/source/module_hamilt_pw/hamilt_pwdft/fs_nonlocal_tools.cpp index e59f7d4dab6..6f70a9c39cb 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/fs_nonlocal_tools.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/fs_nonlocal_tools.cpp @@ -293,7 +293,7 @@ void FS_Nonlocal_tools::reduce_pool_becp(const int& npm) #ifdef __MPI if (GlobalV::NPROC_IN_POOL > 1) { - Parallel_Common::reduce_dev, Device>(this->becp, size_becp_act, POOL_WORLD); + Parallel_Common::reduce_data(this->becp, size_becp_act, POOL_WORLD); } #endif } diff --git a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.h b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.h index af2d65d4a8b..7114f43daa4 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.h +++ b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.h @@ -141,7 +141,7 @@ class OperatorEXXPW : public OperatorPW using resmem_real_op = base_device::memory::resize_memory_op; using delmem_real_op = base_device::memory::delete_memory_op; using gemm_complex_op = ModuleBase::gemm_op; - using vec_add_vec_complex_op = ModuleBase::constantvector_addORsub_constantVector_op; + using vec_add_vec_complex_op = ModuleBase::vector_add_vector_op; using dot_op = ModuleBase::dot_real_op; }; diff --git a/source/module_hamilt_pw/hamilt_pwdft/structure_factor.cpp b/source/module_hamilt_pw/hamilt_pwdft/structure_factor.cpp index 4e328c1fdaf..09d0b56a05f 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/structure_factor.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/structure_factor.cpp @@ -25,22 +25,20 @@ Structure_Factor::Structure_Factor() Structure_Factor::~Structure_Factor() { - if (device == "gpu") { - if (PARAM.inp.precision == "single") { - delmem_cd_op()(this->c_eigts1); - delmem_cd_op()(this->c_eigts2); - delmem_cd_op()(this->c_eigts3); - } + if (device == "gpu") + { + delmem_cd_op()(this->c_eigts1); + delmem_cd_op()(this->c_eigts2); + delmem_cd_op()(this->c_eigts3); delmem_zd_op()(this->z_eigts1); delmem_zd_op()(this->z_eigts2); delmem_zd_op()(this->z_eigts3); } - else { - if (PARAM.inp.precision == "single") { - delmem_ch_op()(this->c_eigts1); - delmem_ch_op()(this->c_eigts2); - delmem_ch_op()(this->c_eigts3); - } + else + { + delmem_ch_op()(this->c_eigts1); + delmem_ch_op()(this->c_eigts2); + delmem_ch_op()(this->c_eigts3); // There's no need to delete double precision pointers while in a CPU environment. } } @@ -150,7 +148,7 @@ void Structure_Factor::setup_structure_factor(const UnitCell* Ucell, const Paral } } if (device == "gpu") { - if (PARAM.inp.precision == "single") { + if (PARAM.globalv.has_float_data) { resmem_cd_op()(this->c_eigts1, Ucell->nat * (2 * rho_basis->nx + 1)); resmem_cd_op()(this->c_eigts2, Ucell->nat * (2 * rho_basis->ny + 1)); resmem_cd_op()(this->c_eigts3, Ucell->nat * (2 * rho_basis->nz + 1)); @@ -166,7 +164,7 @@ void Structure_Factor::setup_structure_factor(const UnitCell* Ucell, const Paral syncmem_z2z_h2d_op()(this->z_eigts3, this->eigts3.c, Ucell->nat * (2 * rho_basis->nz + 1)); } else { - if (PARAM.inp.precision == "single") { + if (PARAM.globalv.has_float_data) { resmem_ch_op()(this->c_eigts1, Ucell->nat * (2 * rho_basis->nx + 1)); resmem_ch_op()(this->c_eigts2, Ucell->nat * (2 * rho_basis->ny + 1)); resmem_ch_op()(this->c_eigts3, Ucell->nat * (2 * rho_basis->nz + 1)); diff --git a/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.cpp b/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.cpp index 0fad7014bde..4bf6033477b 100644 --- a/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.cpp @@ -12,8 +12,8 @@ HamiltSdftPW::HamiltSdftPW(elecstate::Potential* pot_in, pseudopot_cell_vnl* nlpp, const UnitCell* ucell, const int& npol, - double* emin_in, - double* emax_in) + Real* emin_in, + Real* emax_in) : HamiltPW(pot_in, wfc_basis, p_kv, nlpp,ucell), ngk(p_kv->ngk) { this->classname = "HamiltSdftPW"; @@ -53,7 +53,6 @@ void HamiltSdftPW::hPsi_norm(const T* psi_in, T* hpsi_norm, const int const int ik = this->ops->get_ik(); const int npwk_max = this->npwk_max; const int npwk = this->ngk[ik]; - using Real = typename GetTypeReal::type; const Real emin = *this->emin; const Real emax = *this->emax; const Real Ebar = (emin + emax) / 2; @@ -63,10 +62,10 @@ void HamiltSdftPW::hPsi_norm(const T* psi_in, T* hpsi_norm, const int ModuleBase::timer::tick("HamiltSdftPW", "hPsi_norm"); } -// template class HamiltSdftPW, base_device::DEVICE_CPU>; +template class HamiltSdftPW, base_device::DEVICE_CPU>; template class HamiltSdftPW, base_device::DEVICE_CPU>; #if ((defined __CUDA) || (defined __ROCM)) -// template class HamiltSdftPW, base_device::DEVICE_GPU>; +template class HamiltSdftPW, base_device::DEVICE_GPU>; template class HamiltSdftPW, base_device::DEVICE_GPU>; #endif diff --git a/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.h b/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.h index 709c87cfa48..288d5241e45 100644 --- a/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.h +++ b/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.h @@ -10,8 +10,9 @@ template class HamiltSdftPW : public HamiltPW { public: + using Real = typename GetTypeReal::type; /** - * @brief Construct a new Hamilt Sdft P W object + * @brief Construct a new HamiltSdftPW object * * @param pot_in potential * @param wfc_basis pw basis for wave functions @@ -26,20 +27,14 @@ class HamiltSdftPW : public HamiltPW pseudopot_cell_vnl* nlpp, const UnitCell* ucell, const int& npol, - double* emin_in, - double* emax_in); + Real* emin_in, + Real* emax_in); /** - * @brief Destroy the Hamilt Sdft P W object + * @brief Destroy the HamiltSdftPW object * */ ~HamiltSdftPW(){}; - // void update_emin_emax(const double& emin, const double& emax) - // { - // this->emin = &emin; - // this->emax = &emax; - // } - /** * @brief Calculate \hat{H}|psi> * @@ -58,8 +53,8 @@ class HamiltSdftPW : public HamiltPW */ void hPsi_norm(const T* psi_in, T* hpsi, const int& nbands = 1); - double* emin = nullptr; ///< Emin of the Hamiltonian - double* emax = nullptr; ///< Emax of the Hamiltonian + Real* emin = nullptr; ///< Emin of the Hamiltonian + Real* emax = nullptr; ///< Emax of the Hamiltonian private: int npwk_max = 0; ///< maximum number of plane waves diff --git a/source/module_hamilt_pw/hamilt_stodft/kernels/hpsi_norm_op.cpp b/source/module_hamilt_pw/hamilt_stodft/kernels/hpsi_norm_op.cpp index 0457c0e98e2..b75da7900c8 100644 --- a/source/module_hamilt_pw/hamilt_stodft/kernels/hpsi_norm_op.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/kernels/hpsi_norm_op.cpp @@ -29,7 +29,7 @@ struct hpsi_norm_op } }; -// template struct hpsi_norm_op; +template struct hpsi_norm_op; template struct hpsi_norm_op; } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_dos.cpp b/source/module_hamilt_pw/hamilt_stodft/sto_dos.cpp index f647bb20105..54b614764ba 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_dos.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/sto_dos.cpp @@ -4,17 +4,19 @@ #include "module_base/tool_title.h" #include "module_parameter/parameter.h" #include "sto_tool.h" -Sto_DOS::~Sto_DOS() +template +Sto_DOS::~Sto_DOS() { } -Sto_DOS::Sto_DOS(ModulePW::PW_Basis_K* p_wfcpw_in, - K_Vectors* p_kv_in, - elecstate::ElecState* p_elec_in, - psi::Psi>* p_psi_in, - hamilt::Hamilt>* p_hamilt_in, - StoChe& stoche, - Stochastic_WF, base_device::DEVICE_CPU>* p_stowf_in) +template +Sto_DOS::Sto_DOS(ModulePW::PW_Basis_K* p_wfcpw_in, + K_Vectors* p_kv_in, + elecstate::ElecState* p_elec_in, + psi::Psi>* p_psi_in, + hamilt::Hamilt>* p_hamilt_in, + StoChe& stoche, + Stochastic_WF, base_device::DEVICE_CPU>* p_stowf_in) { this->p_wfcpw = p_wfcpw_in; this->p_kv = p_kv_in; @@ -26,19 +28,28 @@ Sto_DOS::Sto_DOS(ModulePW::PW_Basis_K* p_wfcpw_in, this->nbands_ks = p_psi_in->get_nbands(); this->nbands_sto = p_stowf_in->nchi; this->method_sto = stoche.method_sto; - this->stofunc.set_E_range(&stoche.emin_sto, &stoche.emax_sto); + this->stofunc.set_E_range((double*)&stoche.emin_sto, (double*)&stoche.emax_sto); } -void Sto_DOS::decide_param(const int& dos_nche, - const double& emin_sto, - const double& emax_sto, - const bool& dos_setemin, - const bool& dos_setemax, - const double& dos_emin_ev, - const double& dos_emax_ev, - const double& dos_scale) + +template +void Sto_DOS::decide_param(const int& dos_nche, + const double& emin_sto, + const double& emax_sto, + const bool& dos_setemin, + const bool& dos_setemax, + const double& dos_emin_ev, + const double& dos_emax_ev, + const double& dos_scale) { this->dos_nche = dos_nche; - check_che(this->dos_nche, emin_sto, emax_sto, this->nbands_sto, this->p_kv, this->p_stowf, this->p_hamilt_sto); + check_che_op()( + this->dos_nche, + emin_sto, + emax_sto, + this->nbands_sto, + this->p_kv, + reinterpret_cast, Device>*>(this->p_stowf), + reinterpret_cast, Device>*>(this->p_hamilt_sto)); if (dos_setemax) { this->emax = dos_emax_ev; @@ -64,7 +75,8 @@ void Sto_DOS::decide_param(const int& dos_nche, } } -void Sto_DOS::caldos(const double sigmain, const double de, const int npart) +template +void Sto_DOS::caldos(const double sigmain, const double de, const int npart) { ModuleBase::TITLE("Sto_DOS", "caldos"); ModuleBase::timer::tick("Sto_DOS", "caldos"); @@ -255,3 +267,8 @@ void Sto_DOS::caldos(const double sigmain, const double de, const int npart) ModuleBase::timer::tick("Sto_DOS", "caldos"); return; } + +template class Sto_DOS; +#if ((defined __CUDA) || (defined __ROCM)) +template class Sto_DOS; +#endif \ No newline at end of file diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_dos.h b/source/module_hamilt_pw/hamilt_stodft/sto_dos.h index 7b717df07c9..e0e145862d0 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_dos.h +++ b/source/module_hamilt_pw/hamilt_stodft/sto_dos.h @@ -6,6 +6,7 @@ #include "module_hamilt_pw/hamilt_stodft/sto_func.h" #include "module_hamilt_pw/hamilt_stodft/sto_wf.h" +template class Sto_DOS { public: @@ -14,7 +15,7 @@ class Sto_DOS elecstate::ElecState* p_elec_in, psi::Psi>* p_psi_in, hamilt::Hamilt>* p_hamilt_in, - StoChe& stoche, + StoChe& stoche, Stochastic_WF, base_device::DEVICE_CPU>* p_stowf_in); ~Sto_DOS(); diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp index 597c84041b7..483e7a14ce1 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp @@ -3,6 +3,8 @@ #include "module_base/complexmatrix.h" #include "module_base/constants.h" #include "module_base/memory.h" +#include "module_base/module_container/ATen/tensor.h" +#include "module_base/parallel_device.h" #include "module_base/timer.h" #include "module_base/vector3.h" #include "module_parameter/parameter.h" @@ -13,38 +15,45 @@ #define TWOSQRT2LN2 2.354820045030949 // FWHM = 2sqrt(2ln2) * \sigma #define FACTOR 1.839939223835727e7 -Sto_EleCond::Sto_EleCond(UnitCell* p_ucell_in, - K_Vectors* p_kv_in, - elecstate::ElecState* p_elec_in, - ModulePW::PW_Basis_K* p_wfcpw_in, - psi::Psi>* p_psi_in, - pseudopot_cell_vnl* p_ppcell_in, - hamilt::Hamilt>* p_hamilt_in, - StoChe& stoche, - Stochastic_WF, base_device::DEVICE_CPU>* p_stowf_in) - : EleCond(p_ucell_in, p_kv_in, p_elec_in, p_wfcpw_in, p_psi_in, p_ppcell_in) +template +Sto_EleCond::Sto_EleCond(UnitCell* p_ucell_in, + K_Vectors* p_kv_in, + elecstate::ElecState* p_elec_in, + ModulePW::PW_Basis_K* p_wfcpw_in, + psi::Psi, Device>* p_psi_in, + pseudopot_cell_vnl* p_ppcell_in, + hamilt::Hamilt, Device>* p_hamilt_in, + StoChe& stoche, + Stochastic_WF, Device>* p_stowf_in) + : EleCond(p_ucell_in, p_kv_in, p_elec_in, p_wfcpw_in, p_psi_in, p_ppcell_in) { this->p_hamilt = p_hamilt_in; - this->p_hamilt_sto = static_cast>*>(p_hamilt_in); + this->p_hamilt_sto = static_cast, Device>*>(p_hamilt_in); this->p_stowf = p_stowf_in; this->nbands_ks = p_psi_in->get_nbands(); this->nbands_sto = p_stowf_in->nchi; this->stofunc.set_E_range(&stoche.emin_sto, &stoche.emax_sto); +#ifdef __ENABLE_FLOAT_FFTW + if(!std::is_same::value) + { + this->hamilt_sto_ = new hamilt::HamiltSdftPW, Device>(p_elec_in->pot, p_wfcpw_in, p_kv_in, p_ppcell_in, p_ucell_in, 1, &this->emin_sto_, &this->emax_sto_); + } +#endif } -void Sto_EleCond::decide_nche(const double dt, - const double cond_thr, - const int& fd_nche, - double try_emin, - double try_emax) +template +void Sto_EleCond::decide_nche(const FPTYPE dt, + const FPTYPE cond_thr, + const int& fd_nche, + FPTYPE try_emin, + FPTYPE try_emax) { int nche_guess = 1000; - ModuleBase::Chebyshev chet(nche_guess); - const double mu = this->p_elec->eferm.ef; - this->stofunc.mu = mu; + ModuleBase::Chebyshev chet(nche_guess); + this->stofunc.mu = static_cast(this->p_elec->eferm.ef); int& nbatch = this->cond_dtbatch; - auto ncos = std::bind(&Sto_Func::ncos, &this->stofunc, std::placeholders::_1); - auto n_sin = std::bind(&Sto_Func::n_sin, &this->stofunc, std::placeholders::_1); + auto ncos = std::bind(&Sto_Func::ncos, &this->stofunc, std::placeholders::_1); + auto n_sin = std::bind(&Sto_Func::n_sin, &this->stofunc, std::placeholders::_1); // try to find nbatch if (nbatch == 0) { @@ -53,16 +62,16 @@ void Sto_EleCond::decide_nche(const double dt, nbatch = test_nbatch; this->stofunc.t = 0.5 * dt * nbatch; chet.calcoef_pair(ncos, n_sin); - double minerror = std::abs(chet.coef_complex[nche_guess - 1] / chet.coef_complex[0]); + FPTYPE minerror = std::abs(chet.coef_complex[nche_guess - 1] / chet.coef_complex[0]); if (minerror < cond_thr) { for (int i = 1; i < nche_guess; ++i) { - double error = std::abs(chet.coef_complex[i] / chet.coef_complex[0]); + FPTYPE error = std::abs(chet.coef_complex[i] / chet.coef_complex[0]); if (error < cond_thr) { // To make nche to around 100 - nbatch = ceil(double(test_nbatch) / i * 100.0); + nbatch = ceil(float(test_nbatch) / i * 100.0); std::cout << "set cond_dtbatch to " << nbatch << std::endl; break; } @@ -78,7 +87,7 @@ void Sto_EleCond::decide_nche(const double dt, chet.calcoef_pair(ncos, n_sin); for (int i = 1; i < nche_guess; ++i) { - double error = std::abs(chet.coef_complex[i] / chet.coef_complex[0]); + FPTYPE error = std::abs(chet.coef_complex[i] / chet.coef_complex[0]); if (error < cond_thr) { nche = i + 1; @@ -92,13 +101,13 @@ void Sto_EleCond::decide_nche(const double dt, int nche_new = 0; loop: // re-set Emin & Emax both in p_hamilt_sto & stofunc - check_che(std::max(nche_old * 2, fd_nche), - try_emin, - try_emax, - this->nbands_sto, - this->p_kv, - this->p_stowf, - this->p_hamilt_sto); + check_che_op()(std::max(nche_old * 2, fd_nche), + try_emin, + try_emax, + this->nbands_sto, + this->p_kv, + this->p_stowf, + this->p_hamilt_sto); // second try to find nche with new Emin & Emax getnche(nche_new); @@ -129,38 +138,38 @@ void Sto_EleCond::decide_nche(const double dt, this->fd_nche = fd_nche; } -void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, - const psi::Psi>& vkspsi, - const double* en, - const double* en_all, - std::complex* leftfact, - std::complex* rightfact, - const psi::Psi>& leftchi, - psi::Psi>& rightchi, - psi::Psi>& left_hchi, - psi::Psi>& batch_vchi, - psi::Psi>& batch_vhchi, +template +void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPW, Device>* hamilt, + const psi::Psi, Device>& kspsi_all, + const psi::Psi, Device>& vkspsi, + const double* en, + const double* en_all, + std::complex* leftfact, + std::complex* rightfact, + psi::Psi, Device>& leftchi, + psi::Psi, Device>& rightchi, + psi::Psi, Device>& left_hchi, + psi::Psi, Device>& batch_vchi, + psi::Psi, Device>& batch_vhchi, #ifdef __MPI - psi::Psi>& chi_all, - psi::Psi>& hchi_all, - void* gatherinfo_ks, - void* gatherinfo_sto, + psi::Psi, Device>& chi_all, + psi::Psi, Device>& hchi_all, + void* gatherinfo_ks, + void* gatherinfo_sto, #endif - const int& bsize_psi, - std::vector>& j1, - std::vector>& j2, - hamilt::Velocity& velop, - const int& ik, - const std::complex& factor, - const int bandinfo[6]) + const int& bsize_psi, + std::complex* j1, + std::complex* j2, + hamilt::Velocity& velop, + const int& ik, + const std::complex& factor, + const int bandinfo[6]) { ModuleBase::timer::tick("Sto_EleCond", "cal_jmatrix"); - const char transa = 'C'; - const char transb = 'N'; - const std::complex float_factor = static_cast>(factor); - const std::complex conjfactor = std::conj(float_factor); - const float mu = static_cast(this->p_elec->eferm.ef); - const std::complex zero = 0.0; + const std::complex float_factor = factor; + const std::complex conjfactor = std::conj(float_factor); + const lowTYPE mu = static_cast(this->p_elec->eferm.ef); + const std::complex zero = 0.0; const int npwx = this->p_wfcpw->npwk_max; const int npw = this->p_wfcpw->npwk[ik]; const int ndim = 3; @@ -172,44 +181,33 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, const int allbands = bandinfo[5]; const int dim_jmatrix = perbands_ks * allbands_sto + perbands_sto * allbands; - psi::Psi> right_hchi(1, perbands_sto, npwx, npw, true); - psi::Psi> f_rightchi(1, perbands_sto, npwx, npw, true); - psi::Psi> f_right_hchi(1, perbands_sto, npwx, npw, true); + psi::Psi, Device> right_hchi(1, perbands_sto, npwx, npw, true); - this->p_hamilt_sto->hPsi(leftchi.get_pointer(), left_hchi.get_pointer(), perbands_sto); - this->p_hamilt_sto->hPsi(rightchi.get_pointer(), right_hchi.get_pointer(), perbands_sto); - convert_psi(rightchi, f_rightchi); - convert_psi(right_hchi, f_right_hchi); - right_hchi.resize(1, 1, 1); + hamilt->hPsi(leftchi.get_pointer(), left_hchi.get_pointer(), perbands_sto); + hamilt->hPsi(rightchi.get_pointer(), right_hchi.get_pointer(), perbands_sto); - psi::Psi>* rightchi_all = &f_rightchi; - psi::Psi>* righthchi_all = &f_right_hchi; - std::vector> vec_rightf_all; - std::complex* rightf_all = rightfact; + psi::Psi, Device>* rightchi_all = &rightchi; + psi::Psi, Device>* righthchi_all = &right_hchi; + std::vector> vec_rightf_all; + std::complex* rightf_all = rightfact; #ifdef __MPI info_gatherv* ks_fact = static_cast(gatherinfo_ks); info_gatherv* sto_npwx = static_cast(gatherinfo_sto); - rightchi_all = gatherchi(f_rightchi, chi_all, npwx, sto_npwx->nrecv, sto_npwx->displs, perbands_sto); - righthchi_all = gatherchi(f_right_hchi, hchi_all, npwx, sto_npwx->nrecv, sto_npwx->displs, perbands_sto); + rightchi_all = gatherchi_op()(rightchi, chi_all, npwx, sto_npwx->nrecv, sto_npwx->displs, perbands_sto); + righthchi_all = gatherchi_op()(right_hchi, hchi_all, npwx, sto_npwx->nrecv, sto_npwx->displs, perbands_sto); if (PARAM.inp.bndpar > 1 && rightfact != nullptr) { vec_rightf_all.resize(allbands_ks); rightf_all = vec_rightf_all.data(); - MPI_Allgatherv(rightfact, - perbands_ks, - MPI_DOUBLE_COMPLEX, - rightf_all, - ks_fact->nrecv, - ks_fact->displs, - MPI_DOUBLE_COMPLEX, - BP_WORLD); + Parallel_Common::gatherv_dev, Device>(rightfact, + perbands_ks, + rightf_all, + ks_fact->nrecv, + ks_fact->displs, + BP_WORLD); } #endif - psi::Psi> f_batch_vchi(1, bsize_psi * ndim, npwx, npw, true); - psi::Psi> f_batch_vhchi(1, bsize_psi * ndim, npwx, npw, true); - std::vector> tmpj(ndim * allbands_sto * perbands_sto); - // 1. (<\psi|J|\chi>)^T // (allbands_sto, perbands_ks) if (perbands_ks > 0) @@ -218,23 +216,22 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, { const int idnb = id * perbands_ks; const int jbais = 0; - std::complex* j1mat = &j1[id * dim_jmatrix]; - std::complex* j2mat = &j2[id * dim_jmatrix]; + std::complex* j1mat = &j1[id * dim_jmatrix]; + std::complex* j2mat = &j2[id * dim_jmatrix]; //(<\psi|v|\chi>)^T - cgemm_(&transa, - &transb, - &allbands_sto, - &perbands_ks, - &npw, - &conjfactor, - rightchi_all->get_pointer(), - &npwx, - &vkspsi(idnb, 0), - &npwx, - &zero, - j1mat, - &allbands_sto); - + ModuleBase::gemm_op, Device>()('C', + 'N', + allbands_sto, + perbands_ks, + npw, + &conjfactor, + rightchi_all->get_pointer(), + npwx, + &vkspsi(idnb, 0), + npwx, + &zero, + j1mat, + allbands_sto); //(<\psi|Hv|\chi>)^T // for(int i = 0 ; i < perbands_ks ; ++i) // { @@ -248,22 +245,24 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, // } //(<\psi|vH|\chi>)^T - cgemm_(&transa, - &transb, - &allbands_sto, - &perbands_ks, - &npw, - &conjfactor, - righthchi_all->get_pointer(), - &npwx, - &vkspsi(idnb, 0), - &npwx, - &zero, - j2mat, - &allbands_sto); + ModuleBase::gemm_op, Device>()('C', + 'N', + allbands_sto, + perbands_ks, + npw, + &conjfactor, + righthchi_all->get_pointer(), + npwx, + &vkspsi(idnb, 0), + npwx, + &zero, + j2mat, + allbands_sto); } } + std::complex* tmpj = nullptr; + resmem_lcomplex_op()(tmpj, allbands_sto * perbands_sto); int remain = perbands_sto; int startnb = 0; while (remain > 0) @@ -271,10 +270,8 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, int tmpnb = std::min(remain, bsize_psi); // v|\chi> velop.act(&leftchi, tmpnb, &leftchi(0, startnb, 0), batch_vchi.get_pointer()); - convert_psi(batch_vchi, f_batch_vchi); // v|H\chi> velop.act(&leftchi, tmpnb, &left_hchi(0, startnb, 0), batch_vhchi.get_pointer()); - convert_psi(batch_vhchi, f_batch_vhchi); // 2. <\chi|J|\psi> // (perbands_sto, allbands_ks) if (allbands_ks > 0) @@ -283,22 +280,22 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, { const int idnb = id * tmpnb; const int jbais = perbands_ks * allbands_sto + startnb; - std::complex* j1mat = &j1[id * dim_jmatrix + jbais]; - std::complex* j2mat = &j2[id * dim_jmatrix + jbais]; + std::complex* j1mat = &j1[id * dim_jmatrix + jbais]; + std::complex* j2mat = &j2[id * dim_jmatrix + jbais]; //<\chi|v|\psi> - cgemm_(&transa, - &transb, - &tmpnb, - &allbands_ks, - &npw, - &float_factor, - &f_batch_vchi(idnb, 0), - &npwx, - kspsi_all.get_pointer(), - &npwx, - &zero, - j1mat, - &perbands_sto); + ModuleBase::gemm_op, Device>()('C', + 'N', + tmpnb, + allbands_ks, + npw, + &float_factor, + batch_vchi.get_pointer(), + npwx, + kspsi_all.get_pointer(), + npwx, + &zero, + j1mat, + perbands_sto); //<\chi|vH|\psi> = \epsilon * <\chi|v|\psi> // for(int i = 0 ; i < allbands_ks ; ++i) @@ -313,19 +310,19 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, // } //<\chi|Hv|\psi> - cgemm_(&transa, - &transb, - &tmpnb, - &allbands_ks, - &npw, - &float_factor, - &f_batch_vhchi(idnb, 0), - &npwx, - kspsi_all.get_pointer(), - &npwx, - &zero, - j2mat, - &perbands_sto); + ModuleBase::gemm_op, Device>()('C', + 'N', + tmpnb, + allbands_ks, + npw, + &float_factor, + batch_vhchi.get_pointer(), + npwx, + kspsi_all.get_pointer(), + npwx, + &zero, + j2mat, + perbands_sto); } } @@ -335,53 +332,53 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, { const int idnb = id * tmpnb; const int jbais = perbands_ks * allbands_sto + perbands_sto * allbands_ks + startnb; - std::complex* j1mat = &j1[id * dim_jmatrix + jbais]; - std::complex* j2mat = &j2[id * dim_jmatrix + jbais]; - std::complex* tmpjmat = &tmpj[id * allbands_sto * perbands_sto + startnb]; + std::complex* j1mat = &j1[id * dim_jmatrix + jbais]; + std::complex* j2mat = &j2[id * dim_jmatrix + jbais]; + std::complex* tmpjmat = &tmpj[id * allbands_sto * perbands_sto + startnb]; //<\chi|v|\chi> - cgemm_(&transa, - &transb, - &tmpnb, - &allbands_sto, - &npw, - &float_factor, - &f_batch_vchi(idnb, 0), - &npwx, - rightchi_all->get_pointer(), - &npwx, - &zero, - j1mat, - &perbands_sto); + ModuleBase::gemm_op()('C', + 'N', + tmpnb, + allbands_sto, + npw, + &float_factor, + batch_vchi.get_pointer(), + npwx, + rightchi_all->get_pointer(), + npwx, + &zero, + j1mat, + perbands_sto); //<\chi|Hv|\chi> - cgemm_(&transa, - &transb, - &tmpnb, - &allbands_sto, - &npw, - &float_factor, - &f_batch_vhchi(idnb, 0), - &npwx, - rightchi_all->get_pointer(), - &npwx, - &zero, - j2mat, - &perbands_sto); + ModuleBase::gemm_op()('C', + 'N', + tmpnb, + allbands_sto, + npw, + &float_factor, + batch_vhchi.get_pointer(), + npwx, + righthchi_all->get_pointer(), + npwx, + &zero, + j2mat, + perbands_sto); //<\chi|vH|\chi> - cgemm_(&transa, - &transb, - &tmpnb, - &allbands_sto, - &npw, - &float_factor, - &f_batch_vchi(idnb, 0), - &npwx, - righthchi_all->get_pointer(), - &npwx, - &zero, - tmpjmat, - &perbands_sto); + ModuleBase::gemm_op()('C', + 'N', + tmpnb, + allbands_sto, + npw, + &float_factor, + batch_vchi.get_pointer(), + npwx, + rightchi_all->get_pointer(), + npwx, + &zero, + tmpjmat, + perbands_sto); } remain -= tmpnb; @@ -392,71 +389,86 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, } } + const lowTYPE half = static_cast(0.5); + const lowTYPE one = static_cast(1.0); for (int id = 0; id < ndim; ++id) { for (int i = 0; i < perbands_ks; ++i) { - const float ei = static_cast(en[i]); + const lowTYPE ei = static_cast(en[i]); const int jst = i * allbands_sto; - std::complex* j2mat = j2.data() + id * dim_jmatrix + jst; - std::complex* j1mat = j1.data() + id * dim_jmatrix + jst; + lcomplex* j2mat = j2 + id * dim_jmatrix + jst; + lcomplex* j1mat = j1 + id * dim_jmatrix + jst; if (leftfact == nullptr) { - for (int j = 0; j < allbands_sto; ++j) - { - j2mat[j] = 0.5f * j2mat[j] + (0.5f * ei - mu) * j1mat[j]; - } + // for (int j = 0; j < allbands_sto; ++j) + // { + // j2mat[j] = 0.5f * j2mat[j] + (0.5f * ei - mu) * j1mat[j]; + // } + ModuleBase::vector_add_vector_op()(allbands_sto, j2mat, j2mat, half, j1mat, half * ei - mu); } else { - const std::complex jfac = static_cast>(leftfact[i]); - for (int j = 0; j < allbands_sto; ++j) - { - j2mat[j] = jfac * (0.5f * j2mat[j] + (0.5f * ei - mu) * j1mat[j]); - j1mat[j] *= jfac; - } + const lcomplex jfac = static_cast(leftfact[i]); + // for (int j = 0; j < allbands_sto; ++j) + // { + // j2mat[j] = jfac * (0.5f * j2mat[j] + (0.5f * ei - mu) * j1mat[j]); + // j1mat[j] *= jfac; + // } + ModuleBase::vector_add_vector_op()(allbands_sto, j2mat, j2mat, half, j1mat, half * ei - mu); + ModuleBase::scal_op()(allbands_sto, &jfac, j2mat, 1); + ModuleBase::scal_op()(allbands_sto, &jfac, j1mat, 1); } } for (int i = 0; i < allbands_ks; ++i) { - const float ei = static_cast(en_all[i]); + const lowTYPE ei = static_cast(en_all[i]); const int jst = perbands_ks * allbands_sto + i * perbands_sto; - std::complex* j2mat = j2.data() + id * dim_jmatrix + jst; - std::complex* j1mat = j1.data() + id * dim_jmatrix + jst; + lcomplex* j2mat = j2 + id * dim_jmatrix + jst; + lcomplex* j1mat = j1 + id * dim_jmatrix + jst; if (rightfact == nullptr) { - for (int j = 0; j < perbands_sto; ++j) - { - j2mat[j] = 0.5f * j2mat[j] + (0.5f * ei - mu) * j1mat[j]; - } + // for (int j = 0; j < perbands_sto; ++j) + // { + // j2mat[j] = 0.5f * j2mat[j] + (0.5f * ei - mu) * j1mat[j]; + // } + ModuleBase::vector_add_vector_op()(perbands_sto, j2mat, j2mat, half, j1mat, half * ei - mu); } else { - const std::complex jfac = static_cast>(rightf_all[i]); - for (int j = 0; j < perbands_sto; ++j) - { - j2mat[j] = jfac * (0.5f * j2mat[j] + (0.5f * ei - mu) * j1mat[j]); - j1mat[j] *= jfac; - } + const lcomplex jfac = static_cast(rightf_all[i]); + // for (int j = 0; j < perbands_sto; ++j) + // { + // j2mat[j] = jfac * (0.5f * j2mat[j] + (0.5f * ei - mu) * j1mat[j]); + // j1mat[j] *= jfac; + // } + ModuleBase::vector_add_vector_op()(perbands_sto, j2mat, j2mat, half, j1mat, half * ei - mu); + ModuleBase::scal_op()(perbands_sto, &jfac, j2mat, 1); + ModuleBase::scal_op()(perbands_sto, &jfac, j1mat, 1); } } const int jst = perbands_ks * allbands_sto + perbands_sto * allbands_ks; const int ed = dim_jmatrix - jst; - std::complex* j2mat = j2.data() + id * dim_jmatrix + jst; - std::complex* j1mat = j1.data() + id * dim_jmatrix + jst; - std::complex* tmpjmat = tmpj.data() + id * allbands_sto * perbands_sto; - - for (int j = 0; j < ed; ++j) - { - j2mat[j] = 0.5f * (j2mat[j] + tmpjmat[j]) - mu * j1mat[j]; - } + lcomplex* j2mat = j2 + id * dim_jmatrix + jst; + lcomplex* j1mat = j1 + id * dim_jmatrix + jst; + lcomplex* tmpjmat = tmpj + id * allbands_sto * perbands_sto; + + // for (int j = 0; j < ed; ++j) + // { + // j2mat[j] = 0.5f * (j2mat[j] + tmpjmat[j]) - mu * j1mat[j]; + // } + ModuleBase::vector_add_vector_op()(ed, j2mat, j2mat, one, tmpjmat, one); + ModuleBase::vector_add_vector_op()(ed, j2mat, j2mat, half, j1mat, -mu); } #ifdef __MPI - MPI_Allreduce(MPI_IN_PLACE, j1.data(), ndim * dim_jmatrix, MPI_COMPLEX, MPI_SUM, POOL_WORLD); - MPI_Allreduce(MPI_IN_PLACE, j2.data(), ndim * dim_jmatrix, MPI_COMPLEX, MPI_SUM, POOL_WORLD); + if (GlobalV::NPROC_IN_POOL > 1) + { + Parallel_Common::reduce_data(j1, ndim * dim_jmatrix, POOL_WORLD); + Parallel_Common::reduce_data(j2, ndim * dim_jmatrix, POOL_WORLD); + } #endif ModuleBase::timer::tick("Sto_EleCond", "cal_jmatrix"); @@ -464,13 +476,14 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi>& kspsi_all, return; } -void Sto_EleCond::sKG(const int& smear_type, - const double& fwhmin, - const double& wcut, - const double& dw_in, - const double& dt_in, - const bool& nonlocal, - const int& npart_sto) +template +void Sto_EleCond::sKG(const int& smear_type, + const double& fwhmin, + const double& wcut, + const double& dw_in, + const double& dt_in, + const bool& nonlocal, + const int& npart_sto) { ModuleBase::TITLE("Sto_EleCond", "sKG"); ModuleBase::timer::tick("Sto_EleCond", "sKG"); @@ -509,21 +522,21 @@ void Sto_EleCond::sKG(const int& smear_type, assert(nw >= 1); assert(nt >= 1); const int ndim = 3; - const int nk = p_kv->get_nks(); - const int npwx = p_wfcpw->npwk_max; - const double tpiba = p_wfcpw->tpiba; - psi::Psi>* stopsi; + const int nk = this->p_kv->get_nks(); + const int npwx = this->p_wfcpw->npwk_max; + const double tpiba = this->p_wfcpw->tpiba; + psi::Psi, Device>* stopsi; if (this->nbands_ks > 0) { - stopsi = p_stowf->chiortho; + stopsi = this->p_stowf->chiortho; // clean memories //Note shchi is different from \sqrt(fH_here)|chi>, since veffs are different - p_stowf->shchi->resize(1, 1, 1); - p_stowf->chi0->resize(1, 1, 1); // clean memories + this->p_stowf->shchi->resize(1, 1, 1); + this->p_stowf->chi0->resize(1, 1, 1); // clean memories } else { - stopsi = p_stowf->chi0; - p_stowf->shchi->resize(1, 1, 1); // clean memories + stopsi = this->p_stowf->chi0; + this->p_stowf->shchi->resize(1, 1, 1); // clean memories } const double dEcut = (wcut + fwhmin) / ModuleBase::Ry_to_eV; @@ -532,71 +545,101 @@ void Sto_EleCond::sKG(const int& smear_type, std::vector ct12(nt, 0); std::vector ct22(nt, 0); + // Convert to lowTYPE + std::complex one = static_cast>(1.0); + std::complex zero = static_cast>(0.0); + std::complex imag_one = static_cast>(ModuleBase::IMAG_UNIT); + Sto_Func lowfunc; + lowTYPE low_emin = static_cast(*this->stofunc.Emin); + lowTYPE low_emax = static_cast(*this->stofunc.Emax); + lowfunc.set_E_range(&low_emin, &low_emax); + hamilt::HamiltSdftPW* p_low_hamilt = nullptr; + if(hamilt_sto_ != nullptr) + { + p_low_hamilt = hamilt_sto_; + } + else + { + p_low_hamilt = reinterpret_cast, Device>*>(this->p_hamilt_sto); + } + // Init Chebyshev - ModuleBase::Chebyshev che(fd_nche); - ModuleBase::Chebyshev chet(cond_nche); - ModuleBase::Chebyshev chemt(cond_nche); + ModuleBase::Chebyshev che(fd_nche); + ModuleBase::Chebyshev chet(cond_nche); + ModuleBase::Chebyshev chemt(cond_nche); //------------------------------------------------------------------ // Calculate //------------------------------------------------------------------ // Prepare Chebyshev coefficients for exp(-i H/\hbar t) - const double mu = this->p_elec->eferm.ef; - this->stofunc.mu = mu; - this->stofunc.t = 0.5 * dt * nbatch; - auto ncos = std::bind(&Sto_Func::ncos, &this->stofunc, std::placeholders::_1); - auto nsin = std::bind(&Sto_Func::nsin, &this->stofunc, std::placeholders::_1); - auto n_sin = std::bind(&Sto_Func::n_sin, &this->stofunc, std::placeholders::_1); + lowfunc.mu = static_cast(this->p_elec->eferm.ef); + lowfunc.t = static_cast(0.5 * dt * nbatch); + auto ncos = std::bind(&Sto_Func::ncos, &lowfunc, std::placeholders::_1); + auto nsin = std::bind(&Sto_Func::nsin, &lowfunc, std::placeholders::_1); + auto n_sin = std::bind(&Sto_Func::n_sin, &lowfunc, std::placeholders::_1); chet.calcoef_pair(ncos, nsin); chemt.calcoef_pair(ncos, n_sin); - std::vector> batchcoef, batchmcoef; + lcomplex* batchcoef_ = nullptr; + lcomplex* batchmcoef_ = nullptr; + ct::DeviceType device_type = ct::DeviceTypeToEnum::value; + ct::DataType t_type = ct::DataTypeToEnum::value; + ct::Tensor batchcoef(t_type, device_type, {1}); + ct::Tensor batchmcoef(t_type, device_type, {1}); if (nbatch > 1) { - batchcoef.resize(cond_nche * nbatch); - std::complex* tmpcoef = batchcoef.data() + (nbatch - 1) * cond_nche; - batchmcoef.resize(cond_nche * nbatch); - std::complex* tmpmcoef = batchmcoef.data() + (nbatch - 1) * cond_nche; - for (int i = 0; i < cond_nche; ++i) - { - tmpcoef[i] = chet.coef_complex[i]; - tmpmcoef[i] = chemt.coef_complex[i]; - } + + // resmem_lcomplex_op()(batchcoef_, cond_nche * nbatch); + // std::complex* tmpcoef = batchcoef_ + (nbatch - 1) * cond_nche; + // resmem_lcomplex_op()(batchmcoef_, cond_nche * nbatch); + // std::complex* tmpmcoef = batchmcoef_ + (nbatch - 1) * cond_nche; + batchcoef.reshape({nbatch, cond_nche}); + lcomplex* tmpcoef = batchcoef[nbatch-1].data(); + batchmcoef.reshape({nbatch, cond_nche}); + lcomplex* tmpmcoef = batchmcoef[nbatch-1].data(); + + cpymem_lcomplex_op()(tmpcoef, chet.coef_complex, cond_nche); + cpymem_lcomplex_op()(tmpmcoef, chemt.coef_complex, cond_nche); for (int ib = 0; ib < nbatch - 1; ++ib) { - tmpcoef = batchcoef.data() + ib * cond_nche; - tmpmcoef = batchmcoef.data() + ib * cond_nche; - this->stofunc.t = 0.5 * dt * (ib + 1); + // tmpcoef = batchcoef.data() + ib * cond_nche; + // tmpmcoef = batchmcoef.data() + ib * cond_nche; + tmpcoef = batchcoef[ib].data(); + tmpmcoef = batchmcoef[ib].data(); + lowfunc.t = 0.5 * dt * (ib + 1); chet.calcoef_pair(ncos, nsin); chemt.calcoef_pair(ncos, n_sin); - for (int i = 0; i < cond_nche; ++i) - { - tmpcoef[i] = chet.coef_complex[i]; - tmpmcoef[i] = chemt.coef_complex[i]; - } + cpymem_lcomplex_op()(tmpcoef, chet.coef_complex, cond_nche); + cpymem_lcomplex_op()(tmpmcoef, chemt.coef_complex, cond_nche); } - this->stofunc.t = 0.5 * dt * nbatch; + lowfunc.t = 0.5 * dt * nbatch; } // ik loop ModuleBase::timer::tick("Sto_EleCond", "kloop"); - hamilt::Velocity velop(p_wfcpw, p_kv->isk.data(), p_ppcell, p_ucell, nonlocal); + hamilt::Velocity velop(this->p_wfcpw, this->p_kv->isk.data(), this->p_ppcell, this->p_ucell, nonlocal); + hamilt::Velocity low_velop(this->p_wfcpw, this->p_kv->isk.data(), this->p_ppcell, this->p_ucell, nonlocal); for (int ik = 0; ik < nk; ++ik) { velop.init(ik); + low_velop.init(ik); stopsi->fix_k(ik); - p_psi->fix_k(ik); + this->p_psi->fix_k(ik); if (nk > 1) { this->p_hamilt->updateHk(ik); } - const int npw = p_kv->ngk[ik]; + p_low_hamilt->updateHk(ik); + + const int npw = this->p_kv->ngk[ik]; // get allbands_ks int cutib0 = 0; + double emin = static_cast(*this->stofunc.Emin); + double emax = static_cast(*this->stofunc.Emax); if (this->nbands_ks > 0) { - double Emax_KS = std::max(*this->stofunc.Emin, this->p_elec->ekb(ik, this->nbands_ks - 1)); + double Emax_KS = std::max(emin, this->p_elec->ekb(ik, this->nbands_ks - 1)); for (cutib0 = this->nbands_ks - 1; cutib0 >= 0; --cutib0) { if (Emax_KS - this->p_elec->ekb(ik, cutib0) > dEcut) @@ -605,17 +648,16 @@ void Sto_EleCond::sKG(const int& smear_type, } } ++cutib0; - double Emin_KS = (cutib0 < this->nbands_ks) ? this->p_elec->ekb(ik, cutib0) : *this->stofunc.Emin; - double dE = *this->stofunc.Emax - Emin_KS + wcut / ModuleBase::Ry_to_eV; + double Emin_KS = (cutib0 < this->nbands_ks) ? this->p_elec->ekb(ik, cutib0) : emin; + double dE = emax - Emin_KS + wcut / ModuleBase::Ry_to_eV; std::cout << "Emin_KS(" << cutib0 + 1 << "): " << Emin_KS * ModuleBase::Ry_to_eV - << " eV; Emax: " << *this->stofunc.Emax * ModuleBase::Ry_to_eV - << " eV; Recommended max dt: " << 2 * M_PI / dE << " a.u." << std::endl; + << " eV; Emax: " << emax * ModuleBase::Ry_to_eV << " eV; Recommended max dt: " << 2 * M_PI / dE + << " a.u." << std::endl; } else { - double dE = *this->stofunc.Emax - *this->stofunc.Emin + wcut / ModuleBase::Ry_to_eV; - std::cout << "Emin: " << *this->stofunc.Emin * ModuleBase::Ry_to_eV - << " eV; Emax: " << *this->stofunc.Emax * ModuleBase::Ry_to_eV + double dE = emax - emin + wcut / ModuleBase::Ry_to_eV; + std::cout << "Emin: " << emin * ModuleBase::Ry_to_eV << " eV; Emax: " << emax * ModuleBase::Ry_to_eV << " eV; Recommended max dt: " << 2 * M_PI / dE << " a.u." << std::endl; } // Parallel for bands @@ -662,38 +704,34 @@ void Sto_EleCond::sKG(const int& smear_type, // sto conductivity //----------------------------------------------------------- //------------------- allocate ------------------------- - size_t ks_memory_cost = perbands_ks * npwx * sizeof(std::complex); - psi::Psi> kspsi(1, perbands_ks, npwx, npw, true); - psi::Psi> vkspsi(1, perbands_ks * ndim, npwx, npw, true); - std::vector> expmtmf_fact(perbands_ks), expmtf_fact(perbands_ks); - psi::Psi> f_kspsi(1, perbands_ks, npwx, npw, true); + size_t ks_memory_cost = perbands_ks * npwx * sizeof(lcomplex); + psi::Psi, Device> kspsi(1, perbands_ks, npwx, npw, true); + psi::Psi, Device> vkspsi(1, perbands_ks * ndim, npwx, npw, true); + std::vector> expmtmf_fact(perbands_ks), expmtf_fact(perbands_ks); + psi::Psi f_kspsi(1, perbands_ks, npwx, npw, true); ModuleBase::Memory::record("SDFT::kspsi", ks_memory_cost); - psi::Psi> f_vkspsi(1, perbands_ks * ndim, npwx, npw, true); + psi::Psi f_vkspsi(1, perbands_ks * ndim, npwx, npw, true); ModuleBase::Memory::record("SDFT::vkspsi", ks_memory_cost); - psi::Psi>* kspsi_all = &f_kspsi; + psi::Psi* kspsi_all = &f_kspsi; - size_t sto_memory_cost = perbands_sto * npwx * sizeof(std::complex); - psi::Psi> sfchi(1, perbands_sto, npwx, npw, true); + size_t sto_memory_cost = perbands_sto * npwx * sizeof(std::complex); + psi::Psi, Device> sfchi(1, perbands_sto, npwx, npw, true); ModuleBase::Memory::record("SDFT::sfchi", sto_memory_cost); - psi::Psi> smfchi(1, perbands_sto, npwx, npw, true); + psi::Psi, Device> smfchi(1, perbands_sto, npwx, npw, true); ModuleBase::Memory::record("SDFT::smfchi", sto_memory_cost); #ifdef __MPI - psi::Psi> chi_all, hchi_all, psi_all; + psi::Psi chi_all, hchi_all, psi_all; if (PARAM.inp.bndpar > 1) { chi_all.resize(1, allbands_sto, npwx); hchi_all.resize(1, allbands_sto, npwx); - ModuleBase::Memory::record("SDFT::chi_all", allbands_sto * npwx * sizeof(std::complex)); - ModuleBase::Memory::record("SDFT::hchi_all", allbands_sto * npwx * sizeof(std::complex)); + ModuleBase::Memory::record("SDFT::chi_all", allbands_sto * npwx * sizeof(lcomplex)); + ModuleBase::Memory::record("SDFT::hchi_all", allbands_sto * npwx * sizeof(lcomplex)); psi_all.resize(1, allbands_ks, npwx); - ModuleBase::Memory::record("SDFT::kspsi_all", allbands_ks * npwx * sizeof(std::complex)); + ModuleBase::Memory::record("SDFT::kspsi_all", allbands_ks * npwx * sizeof(lcomplex)); for (int ib = 0; ib < allbands_ks; ++ib) { - for (int ig = 0; ig < npw; ++ig) - { - psi_all(0, ib, ig) - = static_cast>(p_psi[0](this->nbands_ks - allbands_ks + ib, ig)); - } + castmem_lcomplex_op()(&psi_all(0, ib, 0), &this->p_psi[0](this->nbands_ks - allbands_ks + ib, 0), npw); } kspsi_all = &psi_all; f_kspsi.resize(1, 1, 1); @@ -702,9 +740,9 @@ void Sto_EleCond::sKG(const int& smear_type, const int nbatch_psi = npart_sto; const int bsize_psi = ceil(double(perbands_sto) / nbatch_psi); - psi::Psi> batch_vchi(1, bsize_psi * ndim, npwx, npw, true); - psi::Psi> batch_vhchi(1, bsize_psi * ndim, npwx, npw, true); - ModuleBase::Memory::record("SDFT::batchjpsi", 3 * bsize_psi * ndim * npwx * sizeof(std::complex)); + psi::Psi, Device> batch_vchi(1, bsize_psi * ndim, npwx, npw, true); + psi::Psi, Device> batch_vhchi(1, bsize_psi * ndim, npwx, npw, true); + ModuleBase::Memory::record("SDFT::batchjpsi", 3 * bsize_psi * ndim * npwx * sizeof(std::complex)); //------------------- sqrt(f)|psi> sqrt(1-f)|psi> --------------- if (perbands_ks > 0) @@ -713,9 +751,9 @@ void Sto_EleCond::sKG(const int& smear_type, { for (int ig = 0; ig < npw; ++ig) { - kspsi(0, ib, ig) = p_psi[0](ib0_ks + ib, ig); + kspsi(0, ib, ig) = this->p_psi[0](ib0_ks + ib, ig); } - double fi = this->stofunc.fd(en[ib]); + FPTYPE fi = this->stofunc.fd(FPTYPE(en[ib])); expmtmf_fact[ib] = 1 - fi; expmtf_fact[ib] = fi; } @@ -724,64 +762,72 @@ void Sto_EleCond::sKG(const int& smear_type, // convert to complex if (PARAM.inp.bndpar == 1) { - convert_psi(kspsi, f_kspsi); + convert_psi_op()(kspsi, f_kspsi); } - convert_psi(vkspsi, f_vkspsi); + convert_psi_op()(vkspsi, f_vkspsi); kspsi.resize(1, 1, 1); vkspsi.resize(1, 1, 1); } - auto nroot_fd = std::bind(&Sto_Func::nroot_fd, &this->stofunc, std::placeholders::_1); + auto nroot_fd = std::bind(&Sto_Func::nroot_fd, &this->stofunc, std::placeholders::_1); che.calcoef_real(nroot_fd); - auto hchi_norm = std::bind(&hamilt::HamiltSdftPW>::hPsi_norm, + auto hchi_norm = std::bind(&hamilt::HamiltSdftPW, Device>::hPsi_norm, p_hamilt_sto, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3); + auto hchi_norm_low = std::bind(&hamilt::HamiltSdftPW::hPsi_norm, + p_low_hamilt, + std::placeholders::_1, + std::placeholders::_2, + std::placeholders::_3); che.calfinalvec_real(hchi_norm, stopsi->get_pointer(), sfchi.get_pointer(), npw, npwx, perbands_sto); - auto nroot_mfd = std::bind(&Sto_Func::nroot_mfd, &this->stofunc, std::placeholders::_1); + auto nroot_mfd = std::bind(&Sto_Func::nroot_mfd, &this->stofunc, std::placeholders::_1); che.calcoef_real(nroot_mfd); che.calfinalvec_real(hchi_norm, stopsi->get_pointer(), smfchi.get_pointer(), npw, npwx, perbands_sto); //------------------------ allocate ------------------------ - psi::Psi>& expmtsfchi = sfchi; - psi::Psi>& expmtsmfchi = smfchi; - psi::Psi> exptsfchi = expmtsfchi; + psi::Psi expmtsfchi(1, perbands_sto, npwx, npw, true); + convert_psi_op()(sfchi, expmtsfchi); + sfchi.resize(1, 1, 1); + psi::Psi expmtsmfchi(1, perbands_sto, npwx, npw, true); + convert_psi_op()(smfchi, expmtsmfchi); + smfchi.resize(1, 1, 1); + psi::Psi exptsfchi = expmtsfchi; ModuleBase::Memory::record("SDFT::exptsfchi", sto_memory_cost); - psi::Psi> exptsmfchi = expmtsmfchi; + psi::Psi exptsmfchi = expmtsmfchi; ModuleBase::Memory::record("SDFT::exptsmfchi", sto_memory_cost); - psi::Psi> poly_expmtsfchi, poly_expmtsmfchi; - psi::Psi> poly_exptsfchi, poly_exptsmfchi; + psi::Psi poly_expmtsfchi, poly_expmtsmfchi; + psi::Psi poly_exptsfchi, poly_exptsmfchi; if (nbatch > 1) { poly_exptsfchi.resize(cond_nche, perbands_sto, npwx); - ModuleBase::Memory::record("SDFT::poly_exptsfchi", - sizeof(std::complex) * cond_nche * perbands_sto * npwx); + ModuleBase::Memory::record("SDFT::poly_exptsfchi", sizeof(lcomplex) * cond_nche * perbands_sto * npwx); poly_exptsmfchi.resize(cond_nche, perbands_sto, npwx); - ModuleBase::Memory::record("SDFT::poly_exptsmfchi", - sizeof(std::complex) * cond_nche * perbands_sto * npwx); + ModuleBase::Memory::record("SDFT::poly_exptsmfchi", sizeof(lcomplex) * cond_nche * perbands_sto * npwx); poly_expmtsfchi.resize(cond_nche, perbands_sto, npwx); - ModuleBase::Memory::record("SDFT::poly_expmtsfchi", - sizeof(std::complex) * cond_nche * perbands_sto * npwx); + ModuleBase::Memory::record("SDFT::poly_expmtsfchi", sizeof(lcomplex) * cond_nche * perbands_sto * npwx); poly_expmtsmfchi.resize(cond_nche, perbands_sto, npwx); - ModuleBase::Memory::record("SDFT::poly_expmtsmfchi", - sizeof(std::complex) * cond_nche * perbands_sto * npwx); + ModuleBase::Memory::record("SDFT::poly_expmtsmfchi", sizeof(lcomplex) * cond_nche * perbands_sto * npwx); } const int dim_jmatrix = perbands_ks * allbands_sto + perbands_sto * allbands; parallel_distribution parajmat(ndim * dim_jmatrix, GlobalV::NPROC_IN_POOL, GlobalV::RANK_IN_POOL); - std::vector> j1l(ndim * dim_jmatrix), j2l(ndim * dim_jmatrix); - ModuleBase::Memory::record("SDFT::j1l", sizeof(std::complex) * ndim * dim_jmatrix); - ModuleBase::Memory::record("SDFT::j2l", sizeof(std::complex) * ndim * dim_jmatrix); - std::vector> j1r(ndim * dim_jmatrix), j2r(ndim * dim_jmatrix); - ModuleBase::Memory::record("SDFT::j1r", sizeof(std::complex) * ndim * dim_jmatrix); - ModuleBase::Memory::record("SDFT::j2r", sizeof(std::complex) * ndim * dim_jmatrix); - psi::Psi> tmphchil(1, perbands_sto, npwx, npw, true); + + ct::Tensor j1l(t_type, device_type, {ndim, dim_jmatrix}); + ct::Tensor j2l(t_type, device_type, {ndim, dim_jmatrix}); + ct::Tensor j1r(t_type, device_type, {ndim, dim_jmatrix}); + ct::Tensor j2r(t_type, device_type, {ndim, dim_jmatrix}); + ModuleBase::Memory::record("SDFT::j1l", sizeof(lcomplex) * ndim * dim_jmatrix); + ModuleBase::Memory::record("SDFT::j2l", sizeof(lcomplex) * ndim * dim_jmatrix); + ModuleBase::Memory::record("SDFT::j1r", sizeof(lcomplex) * ndim * dim_jmatrix); + ModuleBase::Memory::record("SDFT::j2r", sizeof(lcomplex) * ndim * dim_jmatrix); + psi::Psi tmphchil(1, perbands_sto, npwx, npw, true); ModuleBase::Memory::record("SDFT::tmphchil/r", sto_memory_cost * 2); //------------------------ t loop -------------------------- @@ -815,32 +861,32 @@ void Sto_EleCond::sKG(const int& smear_type, for (int ib = 0; ib < perbands_ks; ++ib) { double eigen = en[ib]; - const std::complex expmfactor = exp(ModuleBase::NEG_IMAG_UNIT * eigen * dt); + const std::complex expmfactor = static_cast>(exp(ModuleBase::NEG_IMAG_UNIT * eigen * dt)); expmtf_fact[ib] *= expmfactor; expmtmf_fact[ib] *= expmfactor; } // Sto if (nbatch == 1) { - chemt.calfinalvec_complex(hchi_norm, + chemt.calfinalvec_complex(hchi_norm_low, expmtsfchi.get_pointer(), expmtsfchi.get_pointer(), npw, npwx, perbands_sto); - chemt.calfinalvec_complex(hchi_norm, + chemt.calfinalvec_complex(hchi_norm_low, expmtsmfchi.get_pointer(), expmtsmfchi.get_pointer(), npw, npwx, perbands_sto); - chet.calfinalvec_complex(hchi_norm, + chet.calfinalvec_complex(hchi_norm_low, exptsfchi.get_pointer(), exptsfchi.get_pointer(), npw, npwx, perbands_sto); - chet.calfinalvec_complex(hchi_norm, + chet.calfinalvec_complex(hchi_norm_low, exptsmfchi.get_pointer(), exptsmfchi.get_pointer(), npw, @@ -849,79 +895,81 @@ void Sto_EleCond::sKG(const int& smear_type, } else { - std::complex* tmppolyexpmtsfchi = poly_expmtsfchi.get_pointer(); - std::complex* tmppolyexpmtsmfchi = poly_expmtsmfchi.get_pointer(); - std::complex* tmppolyexptsfchi = poly_exptsfchi.get_pointer(); - std::complex* tmppolyexptsmfchi = poly_exptsmfchi.get_pointer(); - std::complex* stoexpmtsfchi = expmtsfchi.get_pointer(); - std::complex* stoexpmtsmfchi = expmtsmfchi.get_pointer(); - std::complex* stoexptsfchi = exptsfchi.get_pointer(); - std::complex* stoexptsmfchi = exptsmfchi.get_pointer(); + lcomplex* tmppolyexpmtsfchi = poly_expmtsfchi.get_pointer(); + lcomplex* tmppolyexpmtsmfchi = poly_expmtsmfchi.get_pointer(); + lcomplex* tmppolyexptsfchi = poly_exptsfchi.get_pointer(); + lcomplex* tmppolyexptsmfchi = poly_exptsmfchi.get_pointer(); + lcomplex* stoexpmtsfchi = expmtsfchi.get_pointer(); + lcomplex* stoexpmtsmfchi = expmtsmfchi.get_pointer(); + lcomplex* stoexptsfchi = exptsfchi.get_pointer(); + lcomplex* stoexptsmfchi = exptsmfchi.get_pointer(); if ((it - 1) % nbatch == 0) { - chet.calpolyvec_complex(hchi_norm, stoexptsfchi, tmppolyexptsfchi, npw, npwx, perbands_sto); - chet.calpolyvec_complex(hchi_norm, stoexptsmfchi, tmppolyexptsmfchi, npw, npwx, perbands_sto); - chemt.calpolyvec_complex(hchi_norm, stoexpmtsfchi, tmppolyexpmtsfchi, npw, npwx, perbands_sto); - chemt.calpolyvec_complex(hchi_norm, stoexpmtsmfchi, tmppolyexpmtsmfchi, npw, npwx, perbands_sto); + chet.calpolyvec_complex(hchi_norm_low, stoexptsfchi, tmppolyexptsfchi, npw, npwx, perbands_sto); + chet.calpolyvec_complex(hchi_norm_low, stoexptsmfchi, tmppolyexptsmfchi, npw, npwx, perbands_sto); + chemt.calpolyvec_complex(hchi_norm_low, stoexpmtsfchi, tmppolyexpmtsfchi, npw, npwx, perbands_sto); + chemt.calpolyvec_complex(hchi_norm_low, stoexpmtsmfchi, tmppolyexpmtsmfchi, npw, npwx, perbands_sto); } - std::complex* tmpcoef = batchcoef.data() + (it - 1) % nbatch * cond_nche; - std::complex* tmpmcoef = batchmcoef.data() + (it - 1) % nbatch * cond_nche; - const char transa = 'N'; + // std::complex* tmpcoef = batchcoef.data() + (it - 1) % nbatch * cond_nche; + // std::complex* tmpmcoef = batchmcoef.data() + (it - 1) % nbatch * cond_nche; + lcomplex* tmpcoef = batchcoef[(it - 1) % nbatch].data(); + lcomplex* tmpmcoef = batchmcoef[(it - 1) % nbatch].data(); const int LDA = perbands_sto * npwx; const int M = perbands_sto * npwx; const int N = cond_nche; const int inc = 1; - zgemv_(&transa, - &M, - &N, - &ModuleBase::ONE, - tmppolyexptsfchi, - &LDA, - tmpcoef, - &inc, - &ModuleBase::ZERO, - stoexptsfchi, - &inc); - zgemv_(&transa, - &M, - &N, - &ModuleBase::ONE, - tmppolyexptsmfchi, - &LDA, - tmpcoef, - &inc, - &ModuleBase::ZERO, - stoexptsmfchi, - &inc); - zgemv_(&transa, - &M, - &N, - &ModuleBase::ONE, - tmppolyexpmtsfchi, - &LDA, - tmpmcoef, - &inc, - &ModuleBase::ZERO, - stoexpmtsfchi, - &inc); - zgemv_(&transa, - &M, - &N, - &ModuleBase::ONE, - tmppolyexpmtsmfchi, - &LDA, - tmpmcoef, - &inc, - &ModuleBase::ZERO, - stoexpmtsmfchi, - &inc); + ModuleBase::gemv_op()('N', + M, + N, + &one, + tmppolyexptsfchi, + LDA, + tmpcoef, + inc, + &zero, + stoexptsfchi, + inc); + ModuleBase::gemv_op()('N', + M, + N, + &one, + tmppolyexptsmfchi, + LDA, + tmpcoef, + inc, + &zero, + stoexptsmfchi, + inc); + ModuleBase::gemv_op()('N', + M, + N, + &one, + tmppolyexpmtsfchi, + LDA, + tmpmcoef, + inc, + &zero, + stoexpmtsfchi, + inc); + ModuleBase::gemv_op()('N', + M, + N, + &one, + tmppolyexpmtsmfchi, + LDA, + tmpmcoef, + inc, + &zero, + stoexpmtsmfchi, + inc); } ModuleBase::timer::tick("Sto_EleCond", "evolution"); // calculate i<\psi|sqrt(f) exp(-iHt/2)*J*exp(iHt/2) sqrt(1-f)|\psi>^+ // = i<\psi|sqrt(1-f) exp(-iHt/2)*J*exp(iHt/2) sqrt(f)|\psi> - cal_jmatrix(*kspsi_all, + cal_jmatrix(p_low_hamilt, + *kspsi_all, f_vkspsi, en.data(), en_all, @@ -939,15 +987,16 @@ void Sto_EleCond::sKG(const int& smear_type, (void*)&sto_npwx, #endif bsize_psi, - j1l, - j2l, - velop, + j1l.data(), + j2l.data(), + low_velop, ik, - ModuleBase::IMAG_UNIT, + imag_one, bandsinfo); // calculate <\psi|sqrt(1-f) exp(iHt/2)*J*exp(-iHt/2) sqrt(f)|\psi> - cal_jmatrix(*kspsi_all, + cal_jmatrix(p_low_hamilt, + *kspsi_all, f_vkspsi, en.data(), en_all, @@ -965,11 +1014,11 @@ void Sto_EleCond::sKG(const int& smear_type, (void*)&sto_npwx, #endif bsize_psi, - j1r, - j2r, - velop, + j1r.data(), + j2r.data(), + low_velop, ik, - ModuleBase::ONE, + one, bandsinfo); // prepare for parallel @@ -980,19 +1029,19 @@ void Sto_EleCond::sKG(const int& smear_type, // ddot_real = real(A_i^* * B_i) ModuleBase::timer::tick("Sto_EleCond", "ddot_real"); ct11[it] += static_cast( - ModuleBase::GlobalFunc::ddot_real(num_per, j1l.data() + st_per, j1r.data() + st_per, false) - * p_kv->wk[ik] / 2.0); + ModuleBase::GlobalFunc::ddot_real(num_per, j1l.data() + st_per, j1r.data() + st_per, false) + * this->p_kv->wk[ik] / 2.0); double tmp12 = static_cast( - ModuleBase::GlobalFunc::ddot_real(num_per, j1l.data() + st_per, j2r.data() + st_per, false)); + ModuleBase::GlobalFunc::ddot_real(num_per, j1l.data() + st_per, j2r.data() + st_per, false)); double tmp21 = static_cast( - ModuleBase::GlobalFunc::ddot_real(num_per, j2l.data() + st_per, j1r.data() + st_per, false)); + ModuleBase::GlobalFunc::ddot_real(num_per, j2l.data() + st_per, j1r.data() + st_per, false)); - ct12[it] -= 0.5 * (tmp12 + tmp21) * p_kv->wk[ik] / 2.0; + ct12[it] -= 0.5 * (tmp12 + tmp21) * this->p_kv->wk[ik] / 2.0; ct22[it] += static_cast( - ModuleBase::GlobalFunc::ddot_real(num_per, j2l.data() + st_per, j2r.data() + st_per, false) - * p_kv->wk[ik] / 2.0); + ModuleBase::GlobalFunc::ddot_real(num_per, j2l.data() + st_per, j2r.data() + st_per, false) + * this->p_kv->wk[ik] / 2.0); ModuleBase::timer::tick("Sto_EleCond", "ddot_real"); } @@ -1010,7 +1059,13 @@ void Sto_EleCond::sKG(const int& smear_type, //------------------------------------------------------------------ if (GlobalV::MY_RANK == 0) { - calcondw(nt, dt, smear_type, fwhmin, wcut, dw_in, ct11.data(), ct12.data(), ct22.data()); + this->calcondw(nt, dt, smear_type, fwhmin, wcut, dw_in, ct11.data(), ct12.data(), ct22.data()); } ModuleBase::timer::tick("Sto_EleCond", "sKG"); } + +template class Sto_EleCond; +#if ((defined __CUDA) || (defined __ROCM)) +template class Sto_EleCond; +#endif + diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h index 4cb55ae3280..1080c86059b 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h +++ b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h @@ -6,19 +6,33 @@ #include "module_hamilt_pw/hamilt_stodft/sto_wf.h" #include "module_hsolver/hsolver_pw_sdft.h" -class Sto_EleCond : protected EleCond +template +class Sto_EleCond : protected EleCond { public: +#ifdef __ENABLE_FLOAT_FFTW + using lowTYPE = float; // Here we use float to accelerate the calculation, which is enough for the accuracy +#else + using lowTYPE = double; +#endif + using lcomplex = std::complex; + using resmem_lcomplex_op = base_device::memory::resize_memory_op, Device>; + using delmem_lcomplex_op = base_device::memory::delete_memory_op, Device>; + using cpymem_lcomplex_op = base_device::memory::synchronize_memory_op, Device, Device>; + using castmem_lcomplex_op = base_device::memory::cast_memory_op, std::complex, Device, Device>; + public: Sto_EleCond(UnitCell* p_ucell_in, K_Vectors* p_kv_in, elecstate::ElecState* p_elec_in, ModulePW::PW_Basis_K* p_wfcpw_in, - psi::Psi>* p_psi_in, + psi::Psi, Device>* p_psi_in, pseudopot_cell_vnl* p_ppcell_in, - hamilt::Hamilt>* p_hamilt_in, - StoChe& stoche, - Stochastic_WF, base_device::DEVICE_CPU>* p_stowf_in); - ~Sto_EleCond(){}; + hamilt::Hamilt, Device>* p_hamilt_in, + StoChe& stoche, + Stochastic_WF, Device>* p_stowf_in); + ~Sto_EleCond(){ + delete hamilt_sto_; + }; /** * @brief Set the N order of Chebyshev expansion for conductivities * It will change class member : fd_nche, cond_nche @@ -31,7 +45,7 @@ class Sto_EleCond : protected EleCond * @param try_emax trial Emax * */ - void decide_nche(const double dt, const double cond_thr, const int& fd_nche, double try_emin, double try_emax); + void decide_nche(const FPTYPE dt, const FPTYPE cond_thr, const int& fd_nche, FPTYPE try_emin, FPTYPE try_emax); /** * @brief calculate Stochastic Kubo-Greenwood * @@ -53,46 +67,48 @@ class Sto_EleCond : protected EleCond const int& npart_sto); protected: - int nbands_ks = 0; ///< number of KS bands - int nbands_sto = 0; ///< number of stochastic bands - int cond_nche = 0; ///< number of Chebyshev orders for conductivities - int fd_nche = 0; ///< number of Chebyshev orders for Fermi-Dirac function - int cond_dtbatch = 0; ///< number of time steps in a batch - hamilt::Hamilt>* p_hamilt; ///< pointer to the Hamiltonian - Stochastic_WF, base_device::DEVICE_CPU>* p_stowf - = nullptr; ///< pointer to the stochastic wavefunctions - Sto_Func stofunc; ///< functions - - hamilt::HamiltSdftPW>* p_hamilt_sto = nullptr; ///< pointer to the Hamiltonian for sDFT + int nbands_ks = 0; ///< number of KS bands + int nbands_sto = 0; ///< number of stochastic bands + int cond_nche = 0; ///< number of Chebyshev orders for conductivities + int fd_nche = 0; ///< number of Chebyshev orders for Fermi-Dirac function + int cond_dtbatch = 0; ///< number of time steps in a batch + hamilt::Hamilt, Device>* p_hamilt = nullptr; ///< pointer to the Hamiltonian + Stochastic_WF, Device>* p_stowf = nullptr; ///< pointer to the stochastic wavefunctions + Sto_Func stofunc; ///< functions + hamilt::HamiltSdftPW, Device>* p_hamilt_sto = nullptr; ///< pointer to the Hamiltonian for sDFT + hamilt::HamiltSdftPW, Device>* hamilt_sto_ = nullptr; ///< pointer to the Hamiltonian for sDFT + lowTYPE emin_sto_ = 0; ///< Emin of the Hamiltonian for sDFT + lowTYPE emax_sto_ = 0; ///< Emax of the Hamiltonian for sDFT protected: /** * @brief calculate Jmatrix * */ - void cal_jmatrix(const psi::Psi>& kspsi_all, - const psi::Psi>& vkspsi, + void cal_jmatrix(hamilt::HamiltSdftPW, Device>* hamilt, + const psi::Psi, Device>& kspsi_all, + const psi::Psi, Device>& vkspsi, const double* en, const double* en_all, - std::complex* leftfact, - std::complex* rightfact, - const psi::Psi>& leftchi, - psi::Psi>& rightchi, - psi::Psi>& left_hchi, - psi::Psi>& batch_vchi, - psi::Psi>& batch_vhchi, + std::complex* leftfact, + std::complex* rightfact, + psi::Psi, Device>& leftchi, + psi::Psi, Device>& rightchi, + psi::Psi, Device>& left_hchi, + psi::Psi, Device>& batch_vchi, + psi::Psi, Device>& batch_vhchi, #ifdef __MPI - psi::Psi>& chi_all, - psi::Psi>& hchi_all, + psi::Psi, Device>& chi_all, + psi::Psi, Device>& hchi_all, void* gatherinfo_ks, void* gatherinfo_sto, #endif const int& bsize_psi, - std::vector>& j1, - std::vector>& j2, - hamilt::Velocity& velop, + std::complex* j1, + std::complex* j2, + hamilt::Velocity& velop, const int& ik, - const std::complex& factor, + const std::complex& factor, const int bandinfo[6]); }; #endif // ELECOND_H \ No newline at end of file diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_tool.cpp b/source/module_hamilt_pw/hamilt_stodft/sto_tool.cpp index 98f346f1887..aaf15d960d7 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_tool.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/sto_tool.cpp @@ -1,6 +1,7 @@ #include "sto_tool.h" #include "module_base/math_chebyshev.h" +#include "module_base/parallel_device.h" #include "module_base/timer.h" #include "module_parameter/parameter.h" #ifdef __MPI @@ -8,20 +9,21 @@ #endif #include -void check_che(const int& nche_in, - const double& try_emin, - const double& try_emax, - const int& nbands_sto, - K_Vectors* p_kv, - Stochastic_WF, base_device::DEVICE_CPU>* p_stowf, - hamilt::HamiltSdftPW>* p_hamilt_sto) +template +void check_che_op::operator()(const int& nche_in, + const double& try_emin, + const double& try_emax, + const int& nbands_sto, + K_Vectors* p_kv, + Stochastic_WF, Device>* p_stowf, + hamilt::HamiltSdftPW, Device>* p_hamilt_sto) { //------------------------------ // Convergence test //------------------------------ bool change = false; const int nk = p_kv->get_nks(); - ModuleBase::Chebyshev chetest(nche_in); + ModuleBase::Chebyshev chetest(nche_in); int ntest0 = 5; *p_hamilt_sto->emax = try_emax; *p_hamilt_sto->emin = try_emin; @@ -42,21 +44,26 @@ void check_che(const int& nche_in, { p_hamilt_sto->updateHk(ik); const int npw = p_kv->ngk[ik]; - std::complex* pchi = nullptr; - std::vector> randchi; + std::complex* pchi = nullptr; + psi::Psi, Device> randchi_d; + if (nbands_sto == 0) // For case: PARAM.inp.nbands_sto = "all" + { + randchi_d.resize(1, 1, npw); + } int ntest = std::min(ntest0, p_stowf->nchip[ik]); for (int i = 0; i < ntest; ++i) { if (nbands_sto == 0) { - randchi.resize(npw); - pchi = &randchi[0]; + std::vector> randchi(npw); for (int ig = 0; ig < npw; ++ig) { - double rr = std::rand() / double(RAND_MAX); - double arg = std::rand() / double(RAND_MAX); - pchi[ig] = std::complex(rr * cos(arg), rr * sin(arg)); + FPTYPE rr = std::rand() / FPTYPE(RAND_MAX); + FPTYPE arg = std::rand() / FPTYPE(RAND_MAX); + randchi[ig] = std::complex(rr * cos(arg), rr * sin(arg)); } + syncmem_complex_h2d_op()(randchi_d.get_pointer(), randchi.data(), npw); + pchi = randchi_d.get_pointer(); } else if (PARAM.inp.nbands > 0) { @@ -69,7 +76,7 @@ void check_che(const int& nche_in, while (true) { bool converge; - auto hchi_norm = std::bind(&hamilt::HamiltSdftPW>::hPsi_norm, + auto hchi_norm = std::bind(&hamilt::HamiltSdftPW, Device>::hPsi_norm, p_hamilt_sto, std::placeholders::_1, std::placeholders::_2, @@ -106,41 +113,41 @@ void check_che(const int& nche_in, } } -void convert_psi(const psi::Psi>& psi_in, psi::Psi>& psi_out) +template +psi::Psi, Device>* gatherchi_op::operator()( + psi::Psi, Device>& chi, + psi::Psi, Device>& chi_all, + const int& npwx, + int* nrecv_sto, + int* displs_sto, + const int perbands_sto) { - psi_in.fix_k(0); - psi_out.fix_k(0); - for (int i = 0; i < psi_in.size(); ++i) - { - psi_out.get_pointer()[i] = static_cast>(psi_in.get_pointer()[i]); - } - return; -} - -psi::Psi>* gatherchi(psi::Psi>& chi, - psi::Psi>& chi_all, - const int& npwx, - int* nrecv_sto, - int* displs_sto, - const int perbands_sto) -{ - psi::Psi>* p_chi; + psi::Psi, Device>* p_chi; p_chi = χ #ifdef __MPI if (PARAM.inp.bndpar > 1) { p_chi = &chi_all; ModuleBase::timer::tick("sKG", "bands_gather"); - MPI_Allgatherv(chi.get_pointer(), - perbands_sto * npwx, - MPI_COMPLEX, - chi_all.get_pointer(), - nrecv_sto, - displs_sto, - MPI_COMPLEX, - BP_WORLD); + Parallel_Common::gatherv_dev, Device>(chi.get_pointer(), + perbands_sto * npwx, + chi_all.get_pointer(), + nrecv_sto, + displs_sto, + BP_WORLD); ModuleBase::timer::tick("sKG", "bands_gather"); } #endif return p_chi; -} \ No newline at end of file +} + +template struct check_che_op; +template struct check_che_op; +template struct gatherchi_op; +template struct gatherchi_op; +#if ((defined __CUDA) || (defined __ROCM)) +template struct check_che_op; +template struct check_che_op; +template struct gatherchi_op; +template struct gatherchi_op; +#endif \ No newline at end of file diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_tool.h b/source/module_hamilt_pw/hamilt_stodft/sto_tool.h index 83b941141e3..0e9fe1d41de 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_tool.h +++ b/source/module_hamilt_pw/hamilt_stodft/sto_tool.h @@ -1,7 +1,11 @@ +#ifndef STO_TOOL_H +#define STO_TOOL_H #include "module_cell/klist.h" #include "module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.h" #include "module_hamilt_pw/hamilt_stodft/sto_wf.h" +#include "module_base/module_device/memory_op.h" #include "module_psi/psi.h" + /** * @brief Check if Emin and Emax are converged * @@ -10,16 +14,19 @@ * @param try_emax trial Emax * @param nbands_sto number of stochastic bands */ -void check_che(const int& nche_in, - const double& try_emin, - const double& try_emax, - const int& nbands_sto, - K_Vectors* p_kv, - Stochastic_WF, base_device::DEVICE_CPU>* p_stowf, - hamilt::HamiltSdftPW>* p_hamilt_sto); +template +struct check_che_op +{ + using syncmem_complex_h2d_op = base_device::memory::synchronize_memory_op, Device, base_device::DEVICE_CPU>; + void operator()(const int& nche_in, + const double& try_emin, + const double& try_emax, + const int& nbands_sto, + K_Vectors* p_kv, + Stochastic_WF, Device>* p_stowf, + hamilt::HamiltSdftPW, Device>* p_hamilt_sto); +}; -#ifndef PARALLEL_DISTRIBUTION -#define PARALLEL_DISTRIBUTION /** * @brief structure to distribute calculation among processors * @@ -49,11 +56,8 @@ struct parallel_distribution int start; int num_per; }; -#endif #ifdef __MPI -#ifndef INFO_GATHERV -#define INFO_GATHERV /** * @brief gather information from all processors * @@ -85,7 +89,6 @@ struct info_gatherv int* displs = nullptr; }; #endif -#endif /** * @brief convert psi from double to float @@ -93,7 +96,19 @@ struct info_gatherv * @param psi_in input psi of double * @param psi_out output psi of float */ -void convert_psi(const psi::Psi>& psi_in, psi::Psi>& psi_out); +template +struct convert_psi_op +{ + using castmem_complex_op + = base_device::memory::cast_memory_op, std::complex, Device, Device>; + void operator()(const psi::Psi, Device>& psi_in, + psi::Psi, Device>& psi_out) + { + psi_in.fix_k(0); + psi_out.fix_k(0); + castmem_complex_op()(psi_out.get_pointer(), psi_in.get_pointer(), psi_in.size()); + } +}; /** * @brief gather chi from all processors @@ -107,9 +122,15 @@ void convert_psi(const psi::Psi>& psi_in, psi::Psi> pointer to gathered stochastic wave function * */ -psi::Psi>* gatherchi(psi::Psi>& chi, - psi::Psi>& chi_all, - const int& npwx, - int* nrecv_sto, - int* displs_sto, - const int perbands_sto); +template +struct gatherchi_op +{ + psi::Psi, Device>* operator()(psi::Psi, Device>& chi, + psi::Psi, Device>& chi_all, + const int& npwx, + int* nrecv_sto, + int* displs_sto, + const int perbands_sto); +}; + +#endif \ No newline at end of file diff --git a/source/module_hsolver/diago_cg.cpp b/source/module_hsolver/diago_cg.cpp index 9eba7c2bda9..b0eaaad008a 100644 --- a/source/module_hsolver/diago_cg.cpp +++ b/source/module_hsolver/diago_cg.cpp @@ -248,12 +248,12 @@ void DiagoCG::calc_grad(const ct::Tensor& prec, // grad.data()[i] -= lambda * this->pphi[i]; // } // haozhihan replace this 2022-10-6 - ModuleBase::constantvector_addORsub_constantVector_op()(this->n_basis_, - grad.data(), - grad.data(), - 1.0, - pphi.data(), - (-lambda)); + ModuleBase::vector_add_vector_op()(this->n_basis_, + grad.data(), + grad.data(), + 1.0, + pphi.data(), + (-lambda)); } template @@ -367,12 +367,12 @@ void DiagoCG::calc_gamma_cg(const int& iter, // pcg[i] = gamma * pcg[i] + grad.data()[i]; // } // haozhihan replace this 2022-10-6 - ModuleBase::constantvector_addORsub_constantVector_op()(this->n_basis_, - cg.data(), - cg.data(), - gamma, - grad.data(), - 1.0); + ModuleBase::vector_add_vector_op()(this->n_basis_, + cg.data(), + cg.data(), + gamma, + grad.data(), + 1.0); const Real norma = gamma * cg_norm * sin(theta); T znorma = static_cast(norma * -1); @@ -435,12 +435,12 @@ bool DiagoCG::update_psi(const ct::Tensor& pphi, // } // haozhihan replace this 2022-10-6 - ModuleBase::constantvector_addORsub_constantVector_op()(this->n_basis_, - phi_m.data(), - phi_m.data(), - cost, - cg.data(), - sint_norm); + ModuleBase::vector_add_vector_op()(this->n_basis_, + phi_m.data(), + phi_m.data(), + cost, + cg.data(), + sint_norm); if (std::abs(eigen - e0) < ethreshold) { @@ -456,18 +456,18 @@ bool DiagoCG::update_psi(const ct::Tensor& pphi, // } // haozhihan replace this 2022-10-6 - ModuleBase::constantvector_addORsub_constantVector_op()(this->n_basis_, - sphi.data(), - sphi.data(), - cost, - scg.data(), - sint_norm); - ModuleBase::constantvector_addORsub_constantVector_op()(this->n_basis_, - hphi.data(), - hphi.data(), - cost, - pphi.data(), - sint_norm); + ModuleBase::vector_add_vector_op()(this->n_basis_, + sphi.data(), + sphi.data(), + cost, + scg.data(), + sint_norm); + ModuleBase::vector_add_vector_op()(this->n_basis_, + hphi.data(), + hphi.data(), + cost, + pphi.data(), + sint_norm); return false; } } diff --git a/source/module_hsolver/kernels/test/perf_math_kernel.cpp b/source/module_hsolver/kernels/test/perf_math_kernel.cpp index f9f0228c1fc..f88de522722 100644 --- a/source/module_hsolver/kernels/test/perf_math_kernel.cpp +++ b/source/module_hsolver/kernels/test/perf_math_kernel.cpp @@ -20,7 +20,7 @@ * - vector_mul_real_op_cpu * - vector_mul_vector_op_cpu * - vector_div_vector_op_cpu - * - constantvector_addORsub_constantVector_op_cpu + * - vector_add_vector_op_cpu * - axpy_cpu * - scal_cpu @@ -28,7 +28,7 @@ * - vector_mul_real_op_gpu * - vector_mul_vector_op_gpu * - vector_div_vector_op_gpu - * - constantvector_addORsub_constantVector_op_gpu + * - vector_add_vector_op_gpu * - axpy_gpu * - scal_gpu */ @@ -137,8 +137,8 @@ class PerfModuleHsolverMathKernel : public benchmark::Fixture { using vector_mul_real_op_cpu = ModuleBase::vector_mul_real_op, base_device::DEVICE_CPU>; using vector_mul_vector_op_cpu = ModuleBase::vector_mul_vector_op, base_device::DEVICE_CPU>; using vector_div_vector_op_cpu = ModuleBase::vector_div_vector_op, base_device::DEVICE_CPU>; - using constantvector_addORsub_constantVector_op_cpu - = ModuleBase::constantvector_addORsub_constantVector_op, base_device::DEVICE_CPU>; + using vector_add_vector_op_cpu + = ModuleBase::vector_add_vector_op, base_device::DEVICE_CPU>; using axpy_op_cpu = ModuleBase::axpy_op, base_device::DEVICE_CPU>; using scal_op_cpu = ModuleBase::scal_op; using gemv_op_cpu = ModuleBase::gemv_op, base_device::DEVICE_CPU>; @@ -151,8 +151,8 @@ class PerfModuleHsolverMathKernel : public benchmark::Fixture { using vector_mul_real_op_gpu = ModuleBase::vector_mul_real_op, base_device::DEVICE_GPU>; using vector_mul_vector_op_gpu = ModuleBase::vector_mul_vector_op, base_device::DEVICE_GPU>; using vector_div_vector_op_gpu = ModuleBase::vector_div_vector_op, base_device::DEVICE_GPU>; - using constantvector_addORsub_constantVector_op_gpu - = ModuleBase::constantvector_addORsub_constantVector_op, base_device::DEVICE_GPU>; + using vector_add_vector_op_gpu + = ModuleBase::vector_add_vector_op, base_device::DEVICE_GPU>; using axpy_op_gpu = ModuleBase::axpy_op, base_device::DEVICE_GPU>; using scal_op_gpu = ModuleBase::scal_op; @@ -185,9 +185,9 @@ BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_cpu)(ben } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_constantvector_addORsub_constantVector_op_cpu)(benchmark::State& state) { +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_cpu)(benchmark::State& state) { for (auto _ : state) { - constantvector_addORsub_constantVector_op_cpu()(dim_vector, result_zvector, test_zvector_a, dconstant_a ,test_zvector_b, dconstant_b); + vector_add_vector_op_cpu()(dim_vector, result_zvector, test_zvector_a, dconstant_a ,test_zvector_b, dconstant_b); } } @@ -208,7 +208,7 @@ BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_zdot_real_cpu_op)->RangeMul BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_constantvector_addORsub_constantVector_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_axpy_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_scal_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); @@ -254,9 +254,9 @@ BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_gpu)(ben } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_constantvector_addORsub_constantVector_op_gpu)(benchmark::State& state) { +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_gpu)(benchmark::State& state) { for (auto _ : state) { - constantvector_addORsub_constantVector_op_gpu()(dim_vector, result_zvector_gpu, test_zvector_a_gpu, dconstant_a ,test_zvector_b_gpu, dconstant_b); + vector_add_vector_op_gpu()(dim_vector, result_zvector_gpu, test_zvector_a_gpu, dconstant_a ,test_zvector_b_gpu, dconstant_b); } } @@ -279,7 +279,7 @@ BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_zdot_real_gpu_op)->RangeMul BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_constantvector_addORsub_constantVector_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_axpy_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_scal_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); diff --git a/source/module_hsolver/para_linear_transform.cpp b/source/module_hsolver/para_linear_transform.cpp index d6111203468..5a6c8def273 100644 --- a/source/module_hsolver/para_linear_transform.cpp +++ b/source/module_hsolver/para_linear_transform.cpp @@ -120,11 +120,6 @@ void PLinearTransform::act(const T alpha, const T* A, const T* U, con int size = LDA * ncolA_ip; MPI_Status status; -#ifdef __CUDA_MPI - // If the memory is not set to zero, it may cause the result to be wrong when using CUDA Aware MPI - // I am not sure if it is due to CUDA Aware MPI or not - base_device::memory::set_memory_op()(Atmp_device, 0, size); -#endif Parallel_Common::recv_dev(Atmp_device, size, ip, 0, col_world, &status, A_tmp_.data()); ModuleBase::gemm_op()('N', 'N', diff --git a/source/module_io/read_set_globalv.cpp b/source/module_io/read_set_globalv.cpp index 960d3ad48f4..ca840f65e5a 100755 --- a/source/module_io/read_set_globalv.cpp +++ b/source/module_io/read_set_globalv.cpp @@ -65,9 +65,17 @@ void ReadInput::set_globalv(const Input_para& inp, System_para& sys) { sys.all_ks_run = false; } + // set the has_double_data and has_float_data +#ifdef __ENABLE_FLOAT_FFTW + bool float_cond = inp.cal_cond && inp.esolver_type == "sdft"; +#else + bool float_cond = false; +#endif + sys.has_double_data = (inp.precision == "double") || (inp.precision == "mixing") || float_cond; + sys.has_float_data = (inp.precision == "float") || (inp.precision == "mixing") || float_cond; } -/// @note Here para.inp has been synchronized of all ranks. +/// @note Here para.inp has not been synchronized of all ranks. /// Only para.inp in rank 0 is right. /// So we need to broadcast the results to all ranks. void ReadInput::set_global_dir(const Input_para& inp, System_para& sys) diff --git a/source/module_parameter/system_parameter.h b/source/module_parameter/system_parameter.h index e41ed6b823a..47edb9e8e6e 100755 --- a/source/module_parameter/system_parameter.h +++ b/source/module_parameter/system_parameter.h @@ -47,15 +47,19 @@ struct System_para bool deepks_setorb = false; ///< true if "deepks" is set int npol = 1; ///< number of polarization - bool domag = false; /// 1 : calculate the magnetism with x, y, z component - bool domag_z = false; /// 1 : constrain the magnetism to z axis + bool domag = false; ///< 1 : calculate the magnetism with x, y, z component + bool domag_z = false; ///< 1 : constrain the magnetism to z axis bool double_grid = false; ///< true if "ndx,ndy,ndz" is larger than "nx,ny,nz" - double uramping = -10.0 / 13.6; /// U-Ramping method (Ry) + double uramping = -10.0 / 13.6; ///< U-Ramping method (Ry) std::vector hubbard_u = {}; ///< Hubbard Coulomb interaction parameter U (Ry) int kpar_lcao = 1; ///< global number of pools for LCAO diagonalization only int nbands_l = 0; ///< number of bands of each band parallel calculation, same to nbands when bndpar=1 bool ks_run = false; ///< true if current process runs KS calculation bool all_ks_run = true; ///< true if only all processes run KS calculation + + bool has_double_data = true; ///< hamiltonian has double data + bool has_float_data = false; ///< hamiltonian has float data + }; #endif \ No newline at end of file From d2d9970f697ed81d5abcd16db014f83aa88dd326 Mon Sep 17 00:00:00 2001 From: Qianruipku Date: Sat, 22 Mar 2025 15:02:41 +0800 Subject: [PATCH 2/8] fix bug --- .../module_pw/module_fft/fft_bundle.cpp | 15 +--- .../module_pw/module_fft/fft_bundle.h | 1 - source/module_esolver/esolver_ks.cpp | 2 +- .../hamilt_stodft/sto_elecond.cpp | 80 +++++++++++-------- .../hamilt_stodft/sto_elecond.h | 6 +- 5 files changed, 56 insertions(+), 48 deletions(-) diff --git a/source/module_basis/module_pw/module_fft/fft_bundle.cpp b/source/module_basis/module_pw/module_fft/fft_bundle.cpp index 7289e8ab029..98889c69db6 100644 --- a/source/module_basis/module_pw/module_fft/fft_bundle.cpp +++ b/source/module_basis/module_pw/module_fft/fft_bundle.cpp @@ -47,24 +47,17 @@ void FFT_Bundle::initfft(int nx_in, assert(this->device == "cpu" || this->device == "gpu" || this->device == "dsp"); assert(this->precision == "single" || this->precision == "double" || this->precision == "mixing"); - if (this->precision == "single") + if (this->precision == "single" || this->precision == "mixing") { + float_flag = true; #if not defined(__ENABLE_FLOAT_FFTW) if (this->device == "cpu") { - float_define = false; + ModuleBase::WARNING_QUIT("FFT_Bundle", "Please enable float fftw in the cmake to use float fft"); } #endif -#if defined(__CUDA) || defined(__ROCM) - if (this->device == "gpu") - { - float_flag = float_define; - } -#endif - float_flag = float_define; - double_flag = true; } - if (this->precision == "double") + if (this->precision == "double" || this->precision == "mixing") { double_flag = true; } diff --git a/source/module_basis/module_pw/module_fft/fft_bundle.h b/source/module_basis/module_pw/module_fft/fft_bundle.h index 1982a79a0c7..1cef050884e 100644 --- a/source/module_basis/module_pw/module_fft/fft_bundle.h +++ b/source/module_basis/module_pw/module_fft/fft_bundle.h @@ -196,7 +196,6 @@ class FFT_Bundle private: int fft_mode = 0; bool float_flag = false; - bool float_define = true; bool double_flag = false; std::shared_ptr> fft_float = nullptr; std::shared_ptr> fft_double = nullptr; diff --git a/source/module_esolver/esolver_ks.cpp b/source/module_esolver/esolver_ks.cpp index cc9ee5fe5a2..b4e3bcd349e 100644 --- a/source/module_esolver/esolver_ks.cpp +++ b/source/module_esolver/esolver_ks.cpp @@ -66,7 +66,7 @@ ESolver_KS::ESolver_KS() fft_precision = "mixing"; } #endif - pw_wfc = new ModulePW::PW_Basis_K_Big(fft_device, PARAM.inp.precision); + pw_wfc = new ModulePW::PW_Basis_K_Big(fft_device, fft_precision); ModulePW::PW_Basis_K_Big* tmp = static_cast(pw_wfc); // should not use INPUT here, mohan 2024-05-12 diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp index 483e7a14ce1..c28bf3e4ea5 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp @@ -33,10 +33,11 @@ Sto_EleCond::Sto_EleCond(UnitCell* p_ucell_in, this->nbands_ks = p_psi_in->get_nbands(); this->nbands_sto = p_stowf_in->nchi; this->stofunc.set_E_range(&stoche.emin_sto, &stoche.emax_sto); + this->cond_dtbatch = PARAM.inp.cond_dtbatch; #ifdef __ENABLE_FLOAT_FFTW if(!std::is_same::value) { - this->hamilt_sto_ = new hamilt::HamiltSdftPW, Device>(p_elec_in->pot, p_wfcpw_in, p_kv_in, p_ppcell_in, p_ucell_in, 1, &this->emin_sto_, &this->emax_sto_); + this->hamilt_sto_ = new hamilt::HamiltSdftPW, Device>(p_elec_in->pot, p_wfcpw_in, p_kv_in, p_ppcell_in, p_ucell_in, 1, &this->low_emin_, &this->low_emax_); } #endif } @@ -149,6 +150,7 @@ void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPW, Device>& leftchi, psi::Psi, Device>& rightchi, psi::Psi, Device>& left_hchi, + psi::Psi, Device>& right_hchi, psi::Psi, Device>& batch_vchi, psi::Psi, Device>& batch_vhchi, #ifdef __MPI @@ -160,6 +162,7 @@ void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPW* j1, std::complex* j2, + std::complex* tmpj, hamilt::Velocity& velop, const int& ik, const std::complex& factor, @@ -181,8 +184,6 @@ void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPW, Device> right_hchi(1, perbands_sto, npwx, npw, true); - hamilt->hPsi(leftchi.get_pointer(), left_hchi.get_pointer(), perbands_sto); hamilt->hPsi(rightchi.get_pointer(), right_hchi.get_pointer(), perbands_sto); @@ -261,8 +262,6 @@ void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPW* tmpj = nullptr; - resmem_lcomplex_op()(tmpj, allbands_sto * perbands_sto); int remain = perbands_sto; int startnb = 0; while (remain > 0) @@ -289,7 +288,7 @@ void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPW::cal_jmatrix(hamilt::HamiltSdftPW::cal_jmatrix(hamilt::HamiltSdftPWget_pointer(), npwx, @@ -357,9 +356,9 @@ void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPWget_pointer(), + rightchi_all->get_pointer(), npwx, &zero, j2mat, @@ -372,9 +371,9 @@ void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPWget_pointer(), + righthchi_all->get_pointer(), npwx, &zero, tmpjmat, @@ -470,7 +469,6 @@ void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPW::sKG(const int& smear_type, std::complex zero = static_cast>(0.0); std::complex imag_one = static_cast>(ModuleBase::IMAG_UNIT); Sto_Func lowfunc; - lowTYPE low_emin = static_cast(*this->stofunc.Emin); - lowTYPE low_emax = static_cast(*this->stofunc.Emax); - lowfunc.set_E_range(&low_emin, &low_emax); + this->low_emin_ = static_cast(*this->stofunc.Emin); + this->low_emax_ = static_cast(*this->stofunc.Emax); + lowfunc.set_E_range(&low_emin_, &low_emax_); hamilt::HamiltSdftPW* p_low_hamilt = nullptr; if(hamilt_sto_ != nullptr) { @@ -593,9 +591,9 @@ void Sto_EleCond::sKG(const int& smear_type, // std::complex* tmpcoef = batchcoef_ + (nbatch - 1) * cond_nche; // resmem_lcomplex_op()(batchmcoef_, cond_nche * nbatch); // std::complex* tmpmcoef = batchmcoef_ + (nbatch - 1) * cond_nche; - batchcoef.reshape({nbatch, cond_nche}); + batchcoef.resize({nbatch, cond_nche}); lcomplex* tmpcoef = batchcoef[nbatch-1].data(); - batchmcoef.reshape({nbatch, cond_nche}); + batchmcoef.resize({nbatch, cond_nche}); lcomplex* tmpmcoef = batchmcoef[nbatch-1].data(); cpymem_lcomplex_op()(tmpcoef, chet.coef_complex, cond_nche); @@ -635,8 +633,8 @@ void Sto_EleCond::sKG(const int& smear_type, // get allbands_ks int cutib0 = 0; - double emin = static_cast(*this->stofunc.Emin); - double emax = static_cast(*this->stofunc.Emax); + const double emin = static_cast(*this->stofunc.Emin); + const double emax = static_cast(*this->stofunc.Emax); if (this->nbands_ks > 0) { double Emax_KS = std::max(emin, this->p_elec->ekb(ik, this->nbands_ks - 1)); @@ -697,7 +695,7 @@ void Sto_EleCond::sKG(const int& smear_type, //----------------------------------------------------------- if (GlobalV::MY_BNDGROUP == 0 && allbands_ks > 0) { - jjresponse_ks(ik, nt, dt, dEcut, this->p_elec->wg, velop, ct11.data(), ct12.data(), ct22.data()); + this->jjresponse_ks(ik, nt, dt, dEcut, this->p_elec->wg, velop, ct11.data(), ct12.data(), ct22.data()); } //----------------------------------------------------------- @@ -823,11 +821,14 @@ void Sto_EleCond::sKG(const int& smear_type, ct::Tensor j2l(t_type, device_type, {ndim, dim_jmatrix}); ct::Tensor j1r(t_type, device_type, {ndim, dim_jmatrix}); ct::Tensor j2r(t_type, device_type, {ndim, dim_jmatrix}); + ct::Tensor tmpj(t_type, device_type, {ndim, allbands_sto * perbands_sto}); ModuleBase::Memory::record("SDFT::j1l", sizeof(lcomplex) * ndim * dim_jmatrix); ModuleBase::Memory::record("SDFT::j2l", sizeof(lcomplex) * ndim * dim_jmatrix); ModuleBase::Memory::record("SDFT::j1r", sizeof(lcomplex) * ndim * dim_jmatrix); ModuleBase::Memory::record("SDFT::j2r", sizeof(lcomplex) * ndim * dim_jmatrix); + ModuleBase::Memory::record("SDFT::tmpj", sizeof(lcomplex) * ndim * allbands_sto * perbands_sto); psi::Psi tmphchil(1, perbands_sto, npwx, npw, true); + psi::Psi tmphchir(1, perbands_sto, npwx, npw, true); ModuleBase::Memory::record("SDFT::tmphchil/r", sto_memory_cost * 2); //------------------------ t loop -------------------------- @@ -978,6 +979,7 @@ void Sto_EleCond::sKG(const int& smear_type, exptsmfchi, exptsfchi, tmphchil, + tmphchir, batch_vchi, batch_vhchi, #ifdef __MPI @@ -989,6 +991,7 @@ void Sto_EleCond::sKG(const int& smear_type, bsize_psi, j1l.data(), j2l.data(), + tmpj.data(), low_velop, ik, imag_one, @@ -1005,6 +1008,7 @@ void Sto_EleCond::sKG(const int& smear_type, expmtsmfchi, expmtsfchi, tmphchil, + tmphchir, batch_vchi, batch_vhchi, #ifdef __MPI @@ -1016,6 +1020,7 @@ void Sto_EleCond::sKG(const int& smear_type, bsize_psi, j1r.data(), j2r.data(), + tmpj.data(), low_velop, ik, one, @@ -1028,21 +1033,30 @@ void Sto_EleCond::sKG(const int& smear_type, // Im(l_ij*r_ji) = Re(-il_ij * r_ji) = Re( ((il)^+_ji)^* * r_ji)=Re(((il)^+_i)^* * r^+_i) // ddot_real = real(A_i^* * B_i) ModuleBase::timer::tick("Sto_EleCond", "ddot_real"); - ct11[it] += static_cast( - ModuleBase::GlobalFunc::ddot_real(num_per, j1l.data() + st_per, j1r.data() + st_per, false) - * this->p_kv->wk[ik] / 2.0); - double tmp12 = static_cast( - ModuleBase::GlobalFunc::ddot_real(num_per, j1l.data() + st_per, j2r.data() + st_per, false)); - - double tmp21 = static_cast( - ModuleBase::GlobalFunc::ddot_real(num_per, j2l.data() + st_per, j1r.data() + st_per, false)); + ct11[it] += static_cast(ModuleBase::dot_real_op()(num_per, + j1l.data() + st_per, + j1r.data() + st_per, + false) + * this->p_kv->wk[ik] / 2.0); + double tmp12 + = static_cast(ModuleBase::dot_real_op()(num_per, + j1l.data() + st_per, + j2r.data() + st_per, + false)); + + double tmp21 + = static_cast(ModuleBase::dot_real_op()(num_per, + j2l.data() + st_per, + j1r.data() + st_per, + false)); ct12[it] -= 0.5 * (tmp12 + tmp21) * this->p_kv->wk[ik] / 2.0; - ct22[it] += static_cast( - ModuleBase::GlobalFunc::ddot_real(num_per, j2l.data() + st_per, j2r.data() + st_per, false) - * this->p_kv->wk[ik] / 2.0); - + ct22[it] += static_cast(ModuleBase::dot_real_op()(num_per, + j2l.data() + st_per, + j2r.data() + st_per, + false) + * this->p_kv->wk[ik] / 2.0); ModuleBase::timer::tick("Sto_EleCond", "ddot_real"); } std::cout << std::endl; diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h index 1080c86059b..5894713c1dc 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h +++ b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h @@ -78,8 +78,8 @@ class Sto_EleCond : protected EleCond hamilt::HamiltSdftPW, Device>* p_hamilt_sto = nullptr; ///< pointer to the Hamiltonian for sDFT hamilt::HamiltSdftPW, Device>* hamilt_sto_ = nullptr; ///< pointer to the Hamiltonian for sDFT - lowTYPE emin_sto_ = 0; ///< Emin of the Hamiltonian for sDFT - lowTYPE emax_sto_ = 0; ///< Emax of the Hamiltonian for sDFT + lowTYPE low_emin_ = 0; ///< Emin of the Hamiltonian for sDFT + lowTYPE low_emax_ = 0; ///< Emax of the Hamiltonian for sDFT protected: /** * @brief calculate Jmatrix @@ -95,6 +95,7 @@ class Sto_EleCond : protected EleCond psi::Psi, Device>& leftchi, psi::Psi, Device>& rightchi, psi::Psi, Device>& left_hchi, + psi::Psi, Device>& right_hchi, psi::Psi, Device>& batch_vchi, psi::Psi, Device>& batch_vhchi, #ifdef __MPI @@ -106,6 +107,7 @@ class Sto_EleCond : protected EleCond const int& bsize_psi, std::complex* j1, std::complex* j2, + std::complex* tmpj, hamilt::Velocity& velop, const int& ik, const std::complex& factor, From 6ffb2a36c4aa4a7c12e73936359d3444ec9aa87e Mon Sep 17 00:00:00 2001 From: Qianruipku Date: Sun, 23 Mar 2025 13:29:08 +0800 Subject: [PATCH 3/8] fix GPU bug --- source/module_base/parallel_device.h | 2 +- .../hamilt_pwdft/VNL_in_pw.cpp | 2 +- .../module_hamilt_pw/hamilt_pwdft/elecond.cpp | 7 +-- .../hamilt_stodft/sto_elecond.cpp | 12 +----- .../hamilt_stodft/sto_elecond.h | 1 + tests/integrate/186_PW_SKG_MALL_GPU/INPUT | 43 +++++++++++++++++++ tests/integrate/186_PW_SKG_MALL_GPU/KPT | 4 ++ tests/integrate/186_PW_SKG_MALL_GPU/README | 8 ++++ tests/integrate/186_PW_SKG_MALL_GPU/STRU | 18 ++++++++ .../186_PW_SKG_MALL_GPU/refOnsager.txt | 11 +++++ .../integrate/186_PW_SKG_MALL_GPU/result.ref | 7 +++ tests/integrate/CASES_GPU.txt | 1 + 12 files changed, 101 insertions(+), 15 deletions(-) create mode 100644 tests/integrate/186_PW_SKG_MALL_GPU/INPUT create mode 100644 tests/integrate/186_PW_SKG_MALL_GPU/KPT create mode 100644 tests/integrate/186_PW_SKG_MALL_GPU/README create mode 100644 tests/integrate/186_PW_SKG_MALL_GPU/STRU create mode 100644 tests/integrate/186_PW_SKG_MALL_GPU/refOnsager.txt create mode 100644 tests/integrate/186_PW_SKG_MALL_GPU/result.ref diff --git a/source/module_base/parallel_device.h b/source/module_base/parallel_device.h index 3c7094f6059..3ecc2f6aae3 100644 --- a/source/module_base/parallel_device.h +++ b/source/module_base/parallel_device.h @@ -165,7 +165,7 @@ void gatherv_dev(const T* sendbuf, T* sendbuf_cpu = o1.get(sendbuf, sendcount, tmp_sspace); T* recvbuf_cpu = o2.get(recvbuf, gather_space, tmp_rspace); o1.sync_d2h(sendbuf_cpu, sendbuf, sendcount); - gatherv_data(sendbuf_cpu, sendcount, recvbuf, recvcounts, displs, comm); + gatherv_data(sendbuf_cpu, sendcount, recvbuf_cpu, recvcounts, displs, comm); o2.sync_h2d(recvbuf, recvbuf_cpu, gather_space); o1.del(sendbuf_cpu); o2.del(recvbuf_cpu); diff --git a/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp b/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp index 83225843778..7d0f3b313e5 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp @@ -17,7 +17,7 @@ pseudopot_cell_vnl::pseudopot_cell_vnl() { - this->use_gpu_ = (this->use_gpu_); + this->use_gpu_ = (PARAM.inp.device == "gpu"); } pseudopot_cell_vnl::~pseudopot_cell_vnl() diff --git a/source/module_hamilt_pw/hamilt_pwdft/elecond.cpp b/source/module_hamilt_pw/hamilt_pwdft/elecond.cpp index 5997a139cf8..6c8d163d4e4 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/elecond.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/elecond.cpp @@ -155,15 +155,16 @@ void EleCond::jjresponse_ks(const int ik, nbands); std::complex* pij_c = nullptr; + std::vector> pij_h_; if(std::is_same::value) { pij_c = pij_d; } else { - std::vector> pij_h(nbands * nbands); - syncmem_complex_d2h_op()(pij_h.data(), pij_d, nbands * nbands); - pij_c = pij_h.data(); + pij_h_.resize(nbands * nbands); + syncmem_complex_d2h_op()(pij_h_.data(), pij_d, nbands * nbands); + pij_c = pij_h_.data(); } #ifdef __MPI diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp index c28bf3e4ea5..6c6db3a03a1 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp @@ -200,12 +200,7 @@ void Sto_EleCond::cal_jmatrix(hamilt::HamiltSdftPW, Device>(rightfact, - perbands_ks, - rightf_all, - ks_fact->nrecv, - ks_fact->displs, - BP_WORLD); + Parallel_Common::gatherv_data(rightfact, perbands_ks, rightf_all, ks_fact->nrecv, ks_fact->displs, BP_WORLD); } #endif @@ -747,10 +742,7 @@ void Sto_EleCond::sKG(const int& smear_type, { for (int ib = 0; ib < perbands_ks; ++ib) { - for (int ig = 0; ig < npw; ++ig) - { - kspsi(0, ib, ig) = this->p_psi[0](ib0_ks + ib, ig); - } + cpymem_complex_op()(&kspsi(0, ib, 0), &this->p_psi[0](ib0_ks + ib, 0), npw); FPTYPE fi = this->stofunc.fd(FPTYPE(en[ib])); expmtmf_fact[ib] = 1 - fi; expmtf_fact[ib] = fi; diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h index 5894713c1dc..f61040cc8ef 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h +++ b/source/module_hamilt_pw/hamilt_stodft/sto_elecond.h @@ -20,6 +20,7 @@ class Sto_EleCond : protected EleCond using delmem_lcomplex_op = base_device::memory::delete_memory_op, Device>; using cpymem_lcomplex_op = base_device::memory::synchronize_memory_op, Device, Device>; using castmem_lcomplex_op = base_device::memory::cast_memory_op, std::complex, Device, Device>; + using cpymem_complex_op = base_device::memory::synchronize_memory_op, Device, Device>; public: Sto_EleCond(UnitCell* p_ucell_in, K_Vectors* p_kv_in, diff --git a/tests/integrate/186_PW_SKG_MALL_GPU/INPUT b/tests/integrate/186_PW_SKG_MALL_GPU/INPUT new file mode 100644 index 00000000000..6b9a901fd80 --- /dev/null +++ b/tests/integrate/186_PW_SKG_MALL_GPU/INPUT @@ -0,0 +1,43 @@ +INPUT_PARAMETERS +#Parameters (1.General) +suffix autotest +calculation scf +esolver_type sdft +method_sto 1 + +nbands 5 +nbands_sto all +nche_sto 120 +emax_sto 0 +emin_sto 0 +seed_sto 20000 +pseudo_dir ../../PP_ORB +symmetry 1 +kpar 1 +bndpar 2 +device gpu + +#Parameters (2.Iteration) +ecutwfc 20 +scf_thr 1e-6 +scf_nmax 20 + + +#Parameters (3.Basis) +basis_type pw + +#Parameters (4.Smearing) +smearing_method fd +smearing_sigma 0.6 + +#Parameters (5.Mixing) +mixing_type broyden +mixing_beta 0.4 +mixing_gg0 0.0 + +cal_cond 1 +cond_fwhm 16 +cond_wcut 10 +cond_dw 1 +cond_dt 0.237464 +cond_nonlocal 1 diff --git a/tests/integrate/186_PW_SKG_MALL_GPU/KPT b/tests/integrate/186_PW_SKG_MALL_GPU/KPT new file mode 100644 index 00000000000..c289c0158aa --- /dev/null +++ b/tests/integrate/186_PW_SKG_MALL_GPU/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +1 1 1 0 0 0 diff --git a/tests/integrate/186_PW_SKG_MALL_GPU/README b/tests/integrate/186_PW_SKG_MALL_GPU/README new file mode 100644 index 00000000000..a22ce0d6879 --- /dev/null +++ b/tests/integrate/186_PW_SKG_MALL_GPU/README @@ -0,0 +1,8 @@ +This test for: +*MDFT +*Al +*kpoints 2*1*1 +*complete basis +*kpar 1 +*bndpar 2 +*cal_cond 1 diff --git a/tests/integrate/186_PW_SKG_MALL_GPU/STRU b/tests/integrate/186_PW_SKG_MALL_GPU/STRU new file mode 100644 index 00000000000..0dbc74b3433 --- /dev/null +++ b/tests/integrate/186_PW_SKG_MALL_GPU/STRU @@ -0,0 +1,18 @@ +ATOMIC_SPECIES +Si 14 Si.pz-vbc.UPF + +LATTICE_CONSTANT +5 // add lattice constant + +LATTICE_VECTORS +1 0 0 +0 1 0 +0 0 1 + +ATOMIC_POSITIONS +Direct + +Si // Element type +0.0 // magnetism +1 +0.00 0.00 0.00 1 1 1 diff --git a/tests/integrate/186_PW_SKG_MALL_GPU/refOnsager.txt b/tests/integrate/186_PW_SKG_MALL_GPU/refOnsager.txt new file mode 100644 index 00000000000..d4dc47fbfea --- /dev/null +++ b/tests/integrate/186_PW_SKG_MALL_GPU/refOnsager.txt @@ -0,0 +1,11 @@ +## w(eV) sigma(Sm^-1) kappa(W(mK)^-1) L12/e(Am^-1) L22/e^2(Wm^-1) + 0.5 145953 177.614 -3.70776e+06 1.11017e+08 + 1.5 144420 175.681 -3.67128e+06 1.0997e+08 + 2.5 141416 171.896 -3.59972e+06 1.07914e+08 + 3.5 137057 166.415 -3.49581e+06 1.0493e+08 + 4.5 131513 159.459 -3.36351e+06 1.01129e+08 + 5.5 124997 151.302 -3.20782e+06 9.66558e+07 + 6.5 117754 142.252 -3.03448e+06 9.16732e+07 + 7.5 110048 132.637 -2.84972e+06 8.63592e+07 + 8.5 102146 122.779 -2.65987e+06 8.08944e+07 + 9.5 94302.5 112.982 -2.47103e+06 7.54522e+07 diff --git a/tests/integrate/186_PW_SKG_MALL_GPU/result.ref b/tests/integrate/186_PW_SKG_MALL_GPU/result.ref new file mode 100644 index 00000000000..c5fd5b62e19 --- /dev/null +++ b/tests/integrate/186_PW_SKG_MALL_GPU/result.ref @@ -0,0 +1,7 @@ +etotref -150.093373857575 +etotperatomref -150.0933738576 +CompareH_Failed 0 +pointgroupref O_h +spacegroupref O_h +nksibzref 1 +totaltimeref 26.09 diff --git a/tests/integrate/CASES_GPU.txt b/tests/integrate/CASES_GPU.txt index b034e955e25..3d2b1b09a1c 100644 --- a/tests/integrate/CASES_GPU.txt +++ b/tests/integrate/CASES_GPU.txt @@ -3,6 +3,7 @@ 102_PW_BPCG_GPU 107_PW_outWfcPw_GPU 186_PW_KG_100_GPU +186_PW_SKG_MALL_GPU 187_PW_SDFT_ALL_GPU 187_PW_SDFT_MALL_GPU 187_PW_SDFT_MALL_BPCG_GPU From 786398d7ed3b76fac717ca4c2975dd083109d2e9 Mon Sep 17 00:00:00 2001 From: Qianruipku Date: Sun, 23 Mar 2025 16:01:29 +0800 Subject: [PATCH 4/8] fix compile --- .../hamilt_stodft/sto_tool.cpp | 4 ++++ .../186_PW_SKG_MALL_GPU/refOnsager.txt | 20 +++++++++---------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_tool.cpp b/source/module_hamilt_pw/hamilt_stodft/sto_tool.cpp index aaf15d960d7..2f99ef0e819 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_tool.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/sto_tool.cpp @@ -142,12 +142,16 @@ psi::Psi, Device>* gatherchi_op::operator() } template struct check_che_op; +#ifdef __ENABLE_FLOAT_FFTW template struct check_che_op; +#endif template struct gatherchi_op; template struct gatherchi_op; #if ((defined __CUDA) || (defined __ROCM)) template struct check_che_op; +#ifdef __ENABLE_FLOAT_FFTW template struct check_che_op; +#endif template struct gatherchi_op; template struct gatherchi_op; #endif \ No newline at end of file diff --git a/tests/integrate/186_PW_SKG_MALL_GPU/refOnsager.txt b/tests/integrate/186_PW_SKG_MALL_GPU/refOnsager.txt index d4dc47fbfea..ed601a37fba 100644 --- a/tests/integrate/186_PW_SKG_MALL_GPU/refOnsager.txt +++ b/tests/integrate/186_PW_SKG_MALL_GPU/refOnsager.txt @@ -1,11 +1,11 @@ ## w(eV) sigma(Sm^-1) kappa(W(mK)^-1) L12/e(Am^-1) L22/e^2(Wm^-1) - 0.5 145953 177.614 -3.70776e+06 1.11017e+08 - 1.5 144420 175.681 -3.67128e+06 1.0997e+08 - 2.5 141416 171.896 -3.59972e+06 1.07914e+08 - 3.5 137057 166.415 -3.49581e+06 1.0493e+08 - 4.5 131513 159.459 -3.36351e+06 1.01129e+08 - 5.5 124997 151.302 -3.20782e+06 9.66558e+07 - 6.5 117754 142.252 -3.03448e+06 9.16732e+07 - 7.5 110048 132.637 -2.84972e+06 8.63592e+07 - 8.5 102146 122.779 -2.65987e+06 8.08944e+07 - 9.5 94302.5 112.982 -2.47103e+06 7.54522e+07 + 0.5 145953 177.613 -3.70776e+06 1.11017e+08 + 1.5 144420 175.68 -3.67128e+06 1.09969e+08 + 2.5 141416 171.896 -3.59971e+06 1.07914e+08 + 3.5 137057 166.414 -3.4958e+06 1.0493e+08 + 4.5 131513 159.458 -3.36351e+06 1.01129e+08 + 5.5 124997 151.301 -3.20781e+06 9.66556e+07 + 6.5 117754 142.252 -3.03447e+06 9.1673e+07 + 7.5 110048 132.636 -2.84971e+06 8.6359e+07 + 8.5 102146 122.778 -2.65987e+06 8.08943e+07 + 9.5 94302.4 112.981 -2.47103e+06 7.5452e+07 From 3ac2eea512b26f904d38013df9afe8ffcac973ab Mon Sep 17 00:00:00 2001 From: Qianruipku Date: Sun, 23 Mar 2025 17:18:55 +0800 Subject: [PATCH 5/8] fix compile --- .../hamilt_stodft/test/test_sto_tool.cpp | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/source/module_hamilt_pw/hamilt_stodft/test/test_sto_tool.cpp b/source/module_hamilt_pw/hamilt_stodft/test/test_sto_tool.cpp index f343c93c814..9d27626ae6a 100644 --- a/source/module_hamilt_pw/hamilt_stodft/test/test_sto_tool.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/test/test_sto_tool.cpp @@ -18,13 +18,16 @@ void hamilt::HamiltPW::sPsi(T const*, T*, const int, const int, const template hamilt::HamiltSdftPW::HamiltSdftPW(elecstate::Potential* pot_in, - ModulePW::PW_Basis_K* wfc_basis, - K_Vectors* p_kv, - pseudopot_cell_vnl* nlpp, - const UnitCell* ucell, - const int& npol, - double* emin_in, - double* emax_in): HamiltPW(pot_in, wfc_basis, p_kv, nlpp,ucell), ngk(p_kv->ngk){} + ModulePW::PW_Basis_K* wfc_basis, + K_Vectors* p_kv, + pseudopot_cell_vnl* nlpp, + const UnitCell* ucell, + const int& npol, + Real* emin_in, + Real* emax_in) + : HamiltPW(pot_in, wfc_basis, p_kv, nlpp, ucell), ngk(p_kv->ngk) +{ +} template void hamilt::HamiltSdftPW::hPsi_norm(const T* psi_in, T* hpsi, const int& nbands){} @@ -74,7 +77,7 @@ TEST_F(TestStoTool, convert_psi) { psi_in.get_pointer()[i] = std::complex(i, i); } - convert_psi(psi_in, psi_out); + convert_psi_op()(psi_in, psi_out); for (int i = 0; i < 10; ++i) { EXPECT_EQ(psi_out.get_pointer()[i], std::complex(i, i)); @@ -89,6 +92,7 @@ TEST_F(TestStoTool, gatherchi) int nrecv_sto[4] = {1, 2, 3, 4}; int displs_sto[4] = {0, 1, 3, 6}; int perbands_sto = 1; - psi::Psi>* p_chi = gatherchi(chi, chi_all, npwx, nrecv_sto, displs_sto, perbands_sto); + psi::Psi>* p_chi + = gatherchi_op()(chi, chi_all, npwx, nrecv_sto, displs_sto, perbands_sto); EXPECT_EQ(p_chi, &chi); } From da566e566c60ba5e66574737640dc8cedbd6eff2 Mon Sep 17 00:00:00 2001 From: Qianruipku Date: Sun, 23 Mar 2025 17:57:32 +0800 Subject: [PATCH 6/8] fix compile of DCU --- source/CMakeLists.txt | 2 +- source/module_base/kernels/rocm/math_kernel_op_vec.hip.cu | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index 43e8fcdc16d..18f0e226c83 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -104,7 +104,7 @@ if(USE_ROCM) module_hamilt_pw/hamilt_pwdft/kernels/rocm/wf_op.hip.cu module_hamilt_pw/hamilt_pwdft/kernels/rocm/vnl_op.hip.cu module_base/kernels/rocm/math_kernel_op.hip.cu - module_base/kernels/rocm/math_kernel_op.hip_vec.cu + module_base/kernels/rocm/math_kernel_op_vec.hip.cu module_base/kernels/rocm/math_ylm_op.hip.cu module_hamilt_general/module_xc/kernels/rocm/xc_functional_op.hip.cu ) diff --git a/source/module_base/kernels/rocm/math_kernel_op_vec.hip.cu b/source/module_base/kernels/rocm/math_kernel_op_vec.hip.cu index a88be379988..a9009fe5c6e 100644 --- a/source/module_base/kernels/rocm/math_kernel_op_vec.hip.cu +++ b/source/module_base/kernels/rocm/math_kernel_op_vec.hip.cu @@ -87,7 +87,7 @@ void vector_mul_real_op::operator()(const int d { int thread = 1024; int block = (dim + thread - 1) / thread; - hipLaunchKernelGGL(HIP_KERNEL_NAME(vector_div_constant_kernel), + hipLaunchKernelGGL(HIP_KERNEL_NAME(vector_mul_real_kernel), dim3(block), dim3(thread), 0, From 8c130a8cfb95ad586bc1d48dbc5a2695486b9d9b Mon Sep 17 00:00:00 2001 From: Qianruipku Date: Sun, 23 Mar 2025 20:00:56 +0800 Subject: [PATCH 7/8] fix test --- source/module_basis/module_pw/test_serial/pw_basis_k_test.cpp | 4 ++-- source/module_hamilt_pw/hamilt_stodft/test/test_sto_tool.cpp | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/source/module_basis/module_pw/test_serial/pw_basis_k_test.cpp b/source/module_basis/module_pw/test_serial/pw_basis_k_test.cpp index 75f352f4747..78c18fa0280 100644 --- a/source/module_basis/module_pw/test_serial/pw_basis_k_test.cpp +++ b/source/module_basis/module_pw/test_serial/pw_basis_k_test.cpp @@ -150,7 +150,7 @@ TEST_F(PWBasisKTEST, Initparameters) TEST_F(PWBasisKTEST, SetupTransform) { - ModulePW::PW_Basis_K basis_k(device_flag, precision_single); + ModulePW::PW_Basis_K basis_k(device_flag, precision_double); double lat0 = 1.8897261254578281; ModuleBase::Matrix3 latvec(10.0,0.0,0.0, 0.0,10.0,0.0, @@ -170,7 +170,7 @@ TEST_F(PWBasisKTEST, SetupTransform) TEST_F(PWBasisKTEST, CollectLocalPW) { - ModulePW::PW_Basis_K basis_k(device_flag, precision_single); + ModulePW::PW_Basis_K basis_k(device_flag, precision_double); double lat0 = 1.8897261254578281; ModuleBase::Matrix3 latvec(10.0,0.0,0.0, 0.0,10.0,0.0, diff --git a/source/module_hamilt_pw/hamilt_stodft/test/test_sto_tool.cpp b/source/module_hamilt_pw/hamilt_stodft/test/test_sto_tool.cpp index 9d27626ae6a..465fd9c27f0 100644 --- a/source/module_hamilt_pw/hamilt_stodft/test/test_sto_tool.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/test/test_sto_tool.cpp @@ -33,6 +33,10 @@ void hamilt::HamiltSdftPW::hPsi_norm(const T* psi_in, T* hpsi, const template class hamilt::HamiltPW, base_device::DEVICE_CPU>; template class hamilt::HamiltSdftPW, base_device::DEVICE_CPU>; +#if ((defined __CUDA) || (defined __ROCM)) +template class hamilt::HamiltPW, base_device::DEVICE_GPU>; +template class hamilt::HamiltSdftPW, base_device::DEVICE_GPU>; +#endif /** * - Tested Functions: From 254a8885236211fb054e6be136b8eb3ad932fee8 Mon Sep 17 00:00:00 2001 From: Qianruipku Date: Sun, 23 Mar 2025 20:51:58 +0800 Subject: [PATCH 8/8] update --- tests/integrate/186_PW_SKG_MALL_GPU/result.ref | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integrate/186_PW_SKG_MALL_GPU/result.ref b/tests/integrate/186_PW_SKG_MALL_GPU/result.ref index c5fd5b62e19..fe942e2d335 100644 --- a/tests/integrate/186_PW_SKG_MALL_GPU/result.ref +++ b/tests/integrate/186_PW_SKG_MALL_GPU/result.ref @@ -1,5 +1,5 @@ -etotref -150.093373857575 -etotperatomref -150.0933738576 +etotref -150.09337298 +etotperatomref -150.09337298 CompareH_Failed 0 pointgroupref O_h spacegroupref O_h