Skip to content

Commit 1ac9fdc

Browse files
kunkunblueberrykunkunblueberry
committed
Fix: high performance optimization for force calculation
1 parent bc17385 commit 1ac9fdc

9 files changed

Lines changed: 164 additions & 54 deletions

File tree

‎source/module_base/module_device/device.cpp‎

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212

1313
#if defined(__CUDA)
1414
#include <cuda_runtime.h>
15-
#include <cuda.h>
1615
#endif
1716

1817
#if defined(__ROCM)
@@ -300,7 +299,6 @@ void print_device_info<base_device::DEVICE_GPU>(
300299
sprintf(msg, " CUDA Capability Major/Minor version number: %d.%d\n",
301300
deviceProp.major, deviceProp.minor);
302301
ofs_device << msg << std::endl;
303-
#if defined(CUDA_VERSION) && CUDA_VERSION < 13000
304302
sprintf(msg,
305303
" GPU Max Clock rate: %.0f MHz (%0.2f "
306304
"GHz)\n",
@@ -314,7 +312,6 @@ void print_device_info<base_device::DEVICE_GPU>(
314312
sprintf(msg, " Memory Bus Width: %d-bit\n",
315313
deviceProp.memoryBusWidth);
316314
ofs_device << msg << std::endl;
317-
#endif
318315
sprintf(msg,
319316
" Maximum Texture Dimension Size (x,y,z) 1D=(%d), 2D=(%d, "
320317
"%d), 3D=(%d, %d, %d)\n",
@@ -369,7 +366,6 @@ void print_device_info<base_device::DEVICE_GPU>(
369366
sprintf(msg, " Texture alignment: %zu bytes\n",
370367
deviceProp.textureAlignment);
371368
ofs_device << msg << std::endl;
372-
#if defined(CUDA_VERSION) && CUDA_VERSION < 13000
373369
sprintf(msg,
374370
" Concurrent copy and kernel execution: %s with %d copy "
375371
"engine(s)\n",
@@ -379,7 +375,6 @@ void print_device_info<base_device::DEVICE_GPU>(
379375
sprintf(msg, " Run time limit on kernels: %s\n",
380376
deviceProp.kernelExecTimeoutEnabled ? "Yes" : "No");
381377
ofs_device << msg << std::endl;
382-
#endif
383378
sprintf(msg, " Integrated GPU sharing Host Memory: %s\n",
384379
deviceProp.integrated ? "Yes" : "No");
385380
ofs_device << msg << std::endl;
@@ -404,15 +399,13 @@ void print_device_info<base_device::DEVICE_GPU>(
404399
sprintf(msg, " Supports Cooperative Kernel Launch: %s\n",
405400
deviceProp.cooperativeLaunch ? "Yes" : "No");
406401
ofs_device << msg << std::endl;
402+
sprintf(msg, " Supports MultiDevice Co-op Kernel Launch: %s\n",
403+
deviceProp.cooperativeMultiDeviceLaunch ? "Yes" : "No");
404+
ofs_device << msg << std::endl;
407405
sprintf(msg,
408406
" Device PCI Domain ID / Bus ID / location ID: %d / %d / %d\n",
409407
deviceProp.pciDomainID, deviceProp.pciBusID, deviceProp.pciDeviceID);
410408
ofs_device << msg << std::endl;
411-
#if defined(CUDA_VERSION) && CUDA_VERSION < 13000
412-
sprintf(msg, " Supports MultiDevice Co-op Kernel Launch: %s\n",
413-
deviceProp.cooperativeMultiDeviceLaunch ? "Yes" : "No");
414-
ofs_device << msg << std::endl;
415-
416409
const char *sComputeMode[] = {
417410
"Default (multiple host threads can use ::cudaSetDevice() with device "
418411
"simultaneously)",
@@ -428,7 +421,7 @@ void print_device_info<base_device::DEVICE_GPU>(
428421
ofs_device << msg << std::endl;
429422
ofs_device << " " << sComputeMode[deviceProp.computeMode] << std::endl
430423
<< std::endl;
431-
#endif
424+
432425
// If there are 2 or more GPUs, query to determine whether RDMA is supported
433426
if (deviceCount >= 2) {
434427
cudaDeviceProp prop[64];
@@ -718,4 +711,4 @@ void record_device_memory<base_device::DEVICE_GPU>(
718711
#endif
719712

720713
} // end of namespace information
721-
} // end of namespace base_device
714+
} // end of namespace base_device

‎source/module_hamilt_lcao/module_deepks/LCAO_deepks_torch.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ void LCAO_Deepks::cal_gevdm(const int nat, std::vector<torch::Tensor>& gevdm)
115115
// repeat each block for nm times in an additional dimension
116116
torch::Tensor tmp_x = this->pdm[inl].reshape({nm, nm}).unsqueeze(0).repeat({nm, 1, 1});
117117
// torch::Tensor tmp_y = std::get<0>(torch::symeig(tmp_x, true));
118-
torch::Tensor tmp_y = std::get<0>(torch::linalg_eigh(tmp_x, "U"));
118+
torch::Tensor tmp_y = std::get<0>(torch::linalg::eigh(tmp_x, "U"));
119119
torch::Tensor tmp_yshell = torch::eye(nm, torch::TensorOptions().dtype(torch::kFloat64));
120120
std::vector<torch::Tensor> tmp_rpt; // repeated-pdm-tensor (x)
121121
std::vector<torch::Tensor> tmp_rdt; // repeated-d-tensor (y)

‎source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ void LCAO_Deepks::cal_descriptor(const int nat)
6767
std::tuple<torch::Tensor, torch::Tensor> d_v(this->d_tensor[inl], vd);
6868
// d_v = torch::symeig(pdm[inl], /*eigenvalues=*/true,
6969
// /*upper=*/true);
70-
d_v = torch::linalg_eigh(pdm[inl], /*uplo*/ "U");
70+
d_v = torch::linalg::eigh(pdm[inl], /*uplo*/ "U");
7171
d_tensor[inl] = std::get<0>(d_v);
7272
}
7373
ModuleBase::timer::tick("LCAO_Deepks", "cal_descriptor");

‎source/module_hamilt_pw/hamilt_pwdft/forces.cpp‎

Lines changed: 74 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
#include "module_hamilt_general/module_surchem/surchem.h"
1717
#include "module_hamilt_general/module_vdw/vdw.h"
1818
#include "kernels/force_op.h"
19-
19+
#include <type_traits>
2020
#ifdef _OPENMP
2121
#include <omp.h>
2222
#endif
@@ -579,7 +579,7 @@ void Forces<FPTYPE, Device>::cal_force_loc(const UnitCell& ucell,
579579
syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, forcelc_d, forcelc.c, this->nat * 3);
580580
syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, vloc_d, vloc.c, vloc.nr * vloc.nc);
581581

582-
hamilt::cal_force_loc_op<FPTYPE, Device>()(
582+
/* hamilt::cal_force_loc_op<FPTYPE, Device>()(
583583
this->nat,
584584
rho_basis->npw,
585585
ucell.tpiba * ucell.omega,
@@ -590,7 +590,35 @@ void Forces<FPTYPE, Device>::cal_force_loc(const UnitCell& ucell,
590590
aux_d,
591591
vloc_d,
592592
vloc.nc,
593-
forcelc_d);
593+
forcelc_d);*/
594+
595+
#if defined(__ROCM) || defined(__UT_USE_ROCM)
596+
hamilt::cal_force_loc_sincos_op<FPTYPE, Device>()(
597+
this->ctx,
598+
this->nat,
599+
rho_basis->npw,
600+
ucell.ntype,
601+
gcar_d,
602+
tau_d,
603+
vloc_d,
604+
aux_d,
605+
static_cast<FPTYPE>(ucell.tpiba * ucell.omega),
606+
forcelc_d);
607+
#else
608+
hamilt::cal_force_loc_op<FPTYPE, Device>()(
609+
this->nat,
610+
rho_basis->npw,
611+
ucell.tpiba * ucell.omega,
612+
iat2it_d,
613+
ig2gg_d,
614+
gcar_d,
615+
tau_d,
616+
aux_d,
617+
vloc_d,
618+
vloc.nc,
619+
forcelc_d);
620+
#endif
621+
594622
syncmem_var_d2h_op()(this->cpu_ctx, this->ctx, forcelc.c, forcelc_d, this->nat * 3);
595623

596624
delmem_int_op()(this->ctx,iat2it_d);
@@ -788,7 +816,7 @@ void Forces<FPTYPE, Device>::cal_force_ew(const UnitCell& ucell,
788816
syncmem_complex_h2d_op()(this->ctx, this->cpu_ctx, aux_d, aux.data(), rho_basis->npw);
789817
syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, forceion_d, forceion.c, this->nat * 3);
790818

791-
hamilt::cal_force_ew_op<FPTYPE, Device>()(
819+
/* hamilt::cal_force_ew_op<FPTYPE, Device>()(
792820
this->nat,
793821
rho_basis->npw,
794822
rho_basis->ig_gge0,
@@ -798,7 +826,31 @@ void Forces<FPTYPE, Device>::cal_force_ew(const UnitCell& ucell,
798826
it_fact_d,
799827
aux_d,
800828
forceion_d);
801-
829+
*/
830+
831+
#if defined(__ROCM) || defined(__UT_USE_ROCM)
832+
hamilt::cal_force_ew_sincos_op<FPTYPE, Device>()(
833+
this->ctx,
834+
this->nat,
835+
rho_basis->npw,
836+
rho_basis->ig_gge0,
837+
gcar_d,
838+
tau_d,
839+
it_fact_d,
840+
aux_d,
841+
forceion_d);
842+
#else
843+
hamilt::cal_force_ew_op<FPTYPE, Device>()(
844+
this->nat,
845+
rho_basis->npw,
846+
rho_basis->ig_gge0,
847+
iat2it_d,
848+
gcar_d,
849+
tau_d,
850+
it_fact_d,
851+
aux_d,
852+
forceion_d);
853+
#endif
802854
syncmem_var_d2h_op()(this->cpu_ctx, this->ctx, forceion.c, forceion_d, this->nat * 3);
803855
delmem_int_op()(this->ctx,iat2it_d);
804856
delmem_var_op()(this->ctx,gcar_d);
@@ -917,8 +969,25 @@ void Forces<FPTYPE, Device>::cal_force_ew(const UnitCell& ucell,
917969
return;
918970
}
919971

972+
namespace hamilt {
973+
974+
#if defined(__ROCM) || defined(__HIP_PLATFORM_AMD__)
975+
template struct cal_force_ew_sincos_op<double, base_device::DEVICE_GPU>;
976+
template struct cal_force_ew_sincos_op<float, base_device::DEVICE_GPU>;
977+
978+
template struct cal_force_loc_sincos_op<double, base_device::DEVICE_GPU>;
979+
template struct cal_force_loc_sincos_op<float, base_device::DEVICE_GPU>;
980+
#endif
981+
982+
#if defined(__CUDA) || defined(__NVCC__)
983+
template struct cal_force_ew_op<double, base_device::DEVICE_GPU>;
984+
template struct cal_force_ew_op<float, base_device::DEVICE_GPU>;
920985

986+
template struct cal_force_loc_op<double, base_device::DEVICE_GPU>;
987+
template struct cal_force_loc_op<float, base_device::DEVICE_GPU>;
988+
#endif
921989

990+
} // namespace hamilt
922991
template class Forces<double, base_device::DEVICE_CPU>;
923992
#if ((defined __CUDA) || (defined __ROCM))
924993
template class Forces<double, base_device::DEVICE_GPU>;

‎source/module_hamilt_pw/hamilt_pwdft/global.h‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
#include "module_hamilt_general/module_xc/xc_functional.h"
1616
#ifdef __CUDA
1717
#include "cublas_v2.h"
18-
#include <cuda.h> // for CUDA_VERSION
1918
#include "cufft.h"
2019

2120
static const char* _cublasGetErrorString(cublasStatus_t error)
@@ -66,27 +65,22 @@ static const char* _cufftGetErrorString(cufftResult_t error)
6665
return "CUFFT_INVALID_SIZE";
6766
case CUFFT_UNALIGNED_DATA:
6867
return "CUFFT_UNALIGNED_DATA";
68+
case CUFFT_INCOMPLETE_PARAMETER_LIST:
69+
return "CUFFT_INCOMPLETE_PARAMETER_LIST";
6970
case CUFFT_INVALID_DEVICE:
7071
return "CUFFT_INVALID_DEVICE";
72+
case CUFFT_PARSE_ERROR:
73+
return "CUFFT_PARSE_ERROR";
7174
case CUFFT_NO_WORKSPACE:
7275
return "CUFFT_NO_WORKSPACE";
7376
case CUFFT_NOT_IMPLEMENTED:
7477
return "CUFFT_NOT_IMPLEMENTED";
75-
case CUFFT_NOT_SUPPORTED:
76-
return "CUFFT_NOT_SUPPORTED";
77-
78-
#if defined(CUDA_VERSION) && CUDA_VERSION < 13000
79-
case CUFFT_INCOMPLETE_PARAMETER_LIST:
80-
return "CUFFT_INCOMPLETE_PARAMETER_LIST";
81-
case CUFFT_PARSE_ERROR:
82-
return "CUFFT_PARSE_ERROR";
8378
case CUFFT_LICENSE_ERROR:
8479
return "CUFFT_LICENSE_ERROR";
85-
#endif
86-
87-
default:
88-
return "<unknown>";
80+
case CUFFT_NOT_SUPPORTED:
81+
return "CUFFT_NOT_SUPPORTED";
8982
}
83+
return "<unknown>";
9084
}
9185

9286
#define CHECK_CUDA(func) \

‎source/module_hamilt_pw/hamilt_pwdft/kernels/force_op.h‎

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,35 @@ struct cal_force_ew_op{
179179
FPTYPE* forceion
180180
) {};
181181
};
182+
183+
template <typename FPTYPE, typename Device>
184+
struct cal_force_loc_sincos_op{
185+
void operator()(
186+
const Device* ctx,
187+
const int nat,
188+
const int npw,
189+
const int ntype,
190+
const FPTYPE* gcar,
191+
const FPTYPE* tau,
192+
const FPTYPE* vloc_per_type,
193+
const std::complex<FPTYPE>* aux,
194+
const FPTYPE& scale_factor,
195+
FPTYPE* force) {};
196+
};
197+
198+
template <typename FPTYPE, typename Device>
199+
struct cal_force_ew_sincos_op{
200+
void operator()(
201+
const Device* ctx,
202+
const int nat,
203+
const int npw,
204+
const int ig_gge0,
205+
const FPTYPE* gcar,
206+
const FPTYPE* tau,
207+
const FPTYPE* it_facts,
208+
const std::complex<FPTYPE>* aux,
209+
FPTYPE* force) {};
210+
};
182211
#if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
183212
template <typename FPTYPE>
184213
struct cal_vkb1_nl_op<FPTYPE, base_device::DEVICE_GPU>
@@ -335,6 +364,32 @@ struct cal_force_ew_op<FPTYPE, base_device::DEVICE_GPU>{
335364
FPTYPE* forceion
336365
);
337366
};
367+
template <typename FPTYPE>
368+
struct cal_force_loc_sincos_op<FPTYPE, base_device::DEVICE_GPU> {
369+
void operator()(const base_device::DEVICE_GPU* ctx,
370+
const int& nat,
371+
const int& npw,
372+
const int& ntype,
373+
const FPTYPE* gcar,
374+
const FPTYPE* tau,
375+
const FPTYPE* vloc_per_type,
376+
const std::complex<FPTYPE>* aux,
377+
const FPTYPE& scale_factor,
378+
FPTYPE* force);
379+
};
380+
381+
template <typename FPTYPE>
382+
struct cal_force_ew_sincos_op<FPTYPE, base_device::DEVICE_GPU> {
383+
void operator()(const base_device::DEVICE_GPU* ctx,
384+
const int& nat,
385+
const int& npw,
386+
const int& ig_gge0,
387+
const FPTYPE* gcar,
388+
const FPTYPE* tau,
389+
const FPTYPE* it_facts,
390+
const std::complex<FPTYPE>* aux,
391+
FPTYPE* force);
392+
};
338393
#endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
339394
} // namespace hamilt
340-
#endif // W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_source_pw_HAMILT_PWDFT_KERNELS_FORCE_OP_H
395+
#endif // W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_source_pw_HAMILT_PWDFT_KERNELS_FORCE_OP_H

‎source/module_hamilt_pw/hamilt_pwdft/kernels/rocm/force_op.hip.cu‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010

1111
namespace hamilt {
1212

13+
__device__ __forceinline__ void sincos_(float x, float* s, float* c) { sincosf(x, s, c); }
14+
__device__ __forceinline__ void sincos_(double x, double* s, double* c) { sincos(x, s, c); }
1315
template <typename FPTYPE>
1416
__global__ void cal_vkb1_nl(
1517
const int npwx,
@@ -658,7 +660,7 @@ __global__ void cal_force_loc_sincos_kernel(
658660

659661
// Use HIP intrinsic for sincos
660662
FPTYPE sinp, cosp;
661-
sincos(phase, &sinp, &cosp);
663+
sincos_(phase, &sinp, &cosp);
662664

663665
// Calculate force factor
664666
const FPTYPE vloc_factor = vloc_per_type[iat * npw + ig];
@@ -718,7 +720,7 @@ __global__ void cal_force_ew_sincos_kernel(
718720

719721
// Use HIP intrinsic for sincos
720722
FPTYPE sinp, cosp;
721-
sincos(phase, &sinp, &cosp);
723+
sincos_(phase, &sinp, &cosp);
722724

723725
// Calculate Ewald sum contribution (fixed sign error)
724726
const FPTYPE factor = it_fact * (-cosp * aux[ig].imag() + sinp * aux[ig].real());

0 commit comments

Comments
 (0)