high performance optimization for force calculation

Author: kunkunblueberry

source/module_hamilt_pw/hamilt_pwdft/forces.cpp

@@ -16,7 +16,7 @@
 #include "module_hamilt_general/module_surchem/surchem.h"
 #include "module_hamilt_general/module_vdw/vdw.h"
 #include "kernels/force_op.h"
-
+#include <type_traits>
 #ifdef _OPENMP
 #include <omp.h>
 #endif
@@ -579,7 +579,7 @@ void Forces<FPTYPE, Device>::cal_force_loc(const UnitCell& ucell,
         syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, forcelc_d, forcelc.c, this->nat * 3);
         syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, vloc_d, vloc.c, vloc.nr * vloc.nc);
 
-        hamilt::cal_force_loc_op<FPTYPE, Device>()(
+       hamilt::cal_force_loc_op<FPTYPE, Device>()(
             this->nat,
             rho_basis->npw,
             ucell.tpiba * ucell.omega,
@@ -591,6 +591,8 @@ void Forces<FPTYPE, Device>::cal_force_loc(const UnitCell& ucell,
             vloc_d,
             vloc.nc,
             forcelc_d);
+
+
         syncmem_var_d2h_op()(this->cpu_ctx, this->ctx, forcelc.c, forcelc_d, this->nat * 3);
 
         delmem_int_op()(this->ctx,iat2it_d);
@@ -799,6 +801,7 @@ void Forces<FPTYPE, Device>::cal_force_ew(const UnitCell& ucell,
             aux_d,
             forceion_d);
 
+
         syncmem_var_d2h_op()(this->cpu_ctx, this->ctx, forceion.c, forceion_d, this->nat * 3);
         delmem_int_op()(this->ctx,iat2it_d);
         delmem_var_op()(this->ctx,gcar_d);
@@ -917,8 +920,25 @@ void Forces<FPTYPE, Device>::cal_force_ew(const UnitCell& ucell,
     return;
 }
 
+namespace hamilt {
+
+#if defined(__ROCM) || defined(__HIP_PLATFORM_AMD__)
+    template struct cal_force_ew_sincos_op<double, base_device::DEVICE_GPU>;
+    template struct cal_force_ew_sincos_op<float, base_device::DEVICE_GPU>;
 
+    template struct cal_force_loc_sincos_op<double, base_device::DEVICE_GPU>;
+    template struct cal_force_loc_sincos_op<float, base_device::DEVICE_GPU>;
+#endif
+
+#if defined(__CUDA) || defined(__NVCC__)
+    template struct cal_force_ew_op<double, base_device::DEVICE_GPU>;
+    template struct cal_force_ew_op<float, base_device::DEVICE_GPU>;
+
+    template struct cal_force_loc_op<double, base_device::DEVICE_GPU>;
+    template struct cal_force_loc_op<float, base_device::DEVICE_GPU>;
+#endif
 
+} // namespace hamilt
 template class Forces<double, base_device::DEVICE_CPU>;
 #if ((defined __CUDA) || (defined __ROCM))
 template class Forces<double, base_device::DEVICE_GPU>;

source/module_hamilt_pw/hamilt_pwdft/kernels/force_op.h

@@ -179,6 +179,35 @@ struct cal_force_ew_op{
         FPTYPE* forceion
     ) {};
 };
+
+template <typename FPTYPE, typename Device>
+struct cal_force_loc_sincos_op{
+    void operator()(
+            const Device* ctx,
+            const int nat,
+            const int npw,
+            const int ntype,
+            const FPTYPE* gcar,
+            const FPTYPE* tau,
+            const FPTYPE* vloc_per_type,
+            const std::complex<FPTYPE>* aux,
+            const FPTYPE& scale_factor,
+            FPTYPE* force) {};
+};
+
+template <typename FPTYPE, typename Device>
+struct cal_force_ew_sincos_op{
+    void operator()(
+            const Device* ctx,
+            const int nat,
+            const int npw,
+            const int ig_gge0,
+            const FPTYPE* gcar,
+            const FPTYPE* tau,
+            const FPTYPE* it_facts,
+            const std::complex<FPTYPE>* aux,
+            FPTYPE* force) {};
+};
 #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
 template <typename FPTYPE>
 struct cal_vkb1_nl_op<FPTYPE, base_device::DEVICE_GPU>
@@ -335,6 +364,32 @@ struct cal_force_ew_op<FPTYPE, base_device::DEVICE_GPU>{
         FPTYPE* forceion
     );
 };
+template <typename FPTYPE>
+struct cal_force_loc_sincos_op<FPTYPE, base_device::DEVICE_GPU> {
+    void operator()(const base_device::DEVICE_GPU* ctx,
+            const int& nat,
+            const int& npw,
+            const int& ntype,
+            const FPTYPE* gcar,
+            const FPTYPE* tau,
+            const FPTYPE* vloc_per_type,
+            const std::complex<FPTYPE>* aux,
+            const FPTYPE& scale_factor,
+            FPTYPE* force);
+};
+
+template <typename FPTYPE>
+struct cal_force_ew_sincos_op<FPTYPE, base_device::DEVICE_GPU> {
+    void operator()(const base_device::DEVICE_GPU* ctx,
+            const int& nat,
+            const int& npw,
+            const int& ig_gge0,
+            const FPTYPE* gcar,
+            const FPTYPE* tau,
+            const FPTYPE* it_facts,
+            const std::complex<FPTYPE>* aux,
+            FPTYPE* force);
+};
 #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
 } // namespace hamilt
-#endif // W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_source_pw_HAMILT_PWDFT_KERNELS_FORCE_OP_H
+#endif // W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_source_pw_HAMILT_PWDFT_KERNELS_FORCE_OP_H