Refactor: math_kernel_op (deepmodeling#6020)

* make KG support GPU

* fix bug
add tests for KG-GPU

* add tests

* modify reference Onsager.txt

* Refactor: math_kernel_op
split op files
re-order ops

* fix compile

* fix compile

* fix bug

* update results

Author: Qianruipku

python/pyabacus/CONTRIBUTING.md

@@ -190,8 +190,10 @@ list(APPEND _diago
     ${HSOLVER_PATH}/diag_const_nums.cpp
     ${HSOLVER_PATH}/diago_iter_assist.cpp
     ${HSOLVER_PATH}/kernels/dngvd_op.cpp
+    ${HSOLVER_PATH}/kernels/bpcg_kernel_op.cpp
     ${BASE_PATH}/kernels/math_kernel_op.cpp
-    ${BASE_PATH}/kernels/math_op.cpp
+    ${BASE_PATH}/kernels/math_kernel_op_vec.cpp
+    ${BASE_PATH}/kernels/math_ylm_op.cpp
     ${BASE_PATH}/module_device/device.cpp
     ${BASE_PATH}/module_device/memory_op.cpp
     ${PSI_PATH}/psi.cpp

python/pyabacus/src/ModuleBase/CMakeLists.txt

@@ -1,7 +1,8 @@
 list(APPEND pymodule_base
     ${PROJECT_SOURCE_DIR}/src/ModuleBase/py_base_math.cpp
-    ${BASE_PATH}/kernels/math_op.cpp
+    ${BASE_PATH}/kernels/math_ylm_op.cpp
     ${BASE_PATH}/kernels/math_kernel_op.cpp
+    ${BASE_PATH}/kernels/math_kernel_op_vec.cpp
     ${BASE_PATH}/module_device/memory_op.cpp
     ${BASE_PATH}/module_device/device.cpp
     )

python/pyabacus/src/ModuleNAO/CMakeLists.txt

@@ -13,8 +13,9 @@ list(APPEND _naos
     ${NAO_PATH}/two_center_integrator.cpp
     ${NAO_PATH}/two_center_table.cpp
     # dependency
-    ${ABACUS_SOURCE_DIR}/module_base/kernels/math_op.cpp
+    ${ABACUS_SOURCE_DIR}/module_base/kernels/math_ylm_op.cpp
     ${ABACUS_SOURCE_DIR}/module_base/kernels/math_kernel_op.cpp
+    ${ABACUS_SOURCE_DIR}/module_base/kernels/math_kernel_op_vec.cpp
     # ${ABACUS_SOURCE_DIR}/module_psi/kernels/psi_memory_op.cpp
     ${ABACUS_SOURCE_DIR}/module_base/module_device/memory_op.cpp
     ${ABACUS_SOURCE_DIR}/module_base/module_device/device.cpp

python/pyabacus/src/hsolver/CMakeLists.txt

@@ -10,9 +10,11 @@ list(APPEND _diago
 
 
     ${HSOLVER_PATH}/kernels/dngvd_op.cpp
+    ${HSOLVER_PATH}/kernels/bpcg_kernel_op.cpp
     # dependency
     ${BASE_PATH}/kernels/math_kernel_op.cpp
-    ${BASE_PATH}/kernels/math_op.cpp
+    ${BASE_PATH}/kernels/math_kernel_op_vec.cpp
+    ${BASE_PATH}/kernels/math_ylm_op.cpp
     ${BASE_PATH}/module_device/device.cpp
     ${BASE_PATH}/module_device/memory_op.cpp
 

source/CMakeLists.txt

@@ -36,6 +36,7 @@ list(APPEND device_srcs
   module_hamilt_pw/hamilt_stodft/kernels/hpsi_norm_op.cpp
   module_basis/module_pw/kernels/pw_op.cpp
   module_hsolver/kernels/dngvd_op.cpp
+  module_hsolver/kernels/bpcg_kernel_op.cpp
   module_elecstate/kernels/elecstate_op.cpp
 
   # module_psi/kernels/psi_memory_op.cpp
@@ -44,13 +45,14 @@ list(APPEND device_srcs
   module_base/module_device/device.cpp
   module_base/module_device/memory_op.cpp
   module_base/kernels/math_kernel_op.cpp
+  module_base/kernels/math_kernel_op_vec.cpp
 
   module_hamilt_pw/hamilt_pwdft/kernels/force_op.cpp
   module_hamilt_pw/hamilt_pwdft/kernels/stress_op.cpp
   module_hamilt_pw/hamilt_pwdft/kernels/onsite_op.cpp
   module_hamilt_pw/hamilt_pwdft/kernels/wf_op.cpp
   module_hamilt_pw/hamilt_pwdft/kernels/vnl_op.cpp
-  module_base/kernels/math_op.cpp
+  module_base/kernels/math_ylm_op.cpp
   module_hamilt_general/module_xc/kernels/xc_functional_op.cpp
 )
 
@@ -64,6 +66,7 @@ if(USE_CUDA)
     module_hamilt_pw/hamilt_pwdft/kernels/cuda/onsite_op.cu
     module_basis/module_pw/kernels/cuda/pw_op.cu
     module_hsolver/kernels/cuda/dngvd_op.cu
+    module_hsolver/kernels/cuda/bpcg_kernel_op.cu
     module_elecstate/kernels/cuda/elecstate_op.cu
 
     # module_psi/kernels/cuda/memory_op.cu
@@ -73,8 +76,9 @@ if(USE_CUDA)
     module_hamilt_pw/hamilt_pwdft/kernels/cuda/stress_op.cu
     module_hamilt_pw/hamilt_pwdft/kernels/cuda/wf_op.cu
     module_hamilt_pw/hamilt_pwdft/kernels/cuda/vnl_op.cu
-    module_base/kernels/cuda/math_op.cu
+    module_base/kernels/cuda/math_ylm_op.cu
     module_base/kernels/cuda/math_kernel_op.cu
+    module_base/kernels/cuda/math_kernel_op_vec.cu
     module_hamilt_general/module_xc/kernels/cuda/xc_functional_op.cu
   )
 endif()
@@ -89,6 +93,7 @@ if(USE_ROCM)
     module_hamilt_pw/hamilt_stodft/kernels/rocm/hpsi_norm_op.hip.cu
     module_basis/module_pw/kernels/rocm/pw_op.hip.cu
     module_hsolver/kernels/rocm/dngvd_op.hip.cu
+    module_hsolver/kernels/rocm/bpcg_kernel_op.hip.cu
     module_elecstate/kernels/rocm/elecstate_op.hip.cu
 
     # module_psi/kernels/rocm/memory_op.hip.cu
@@ -99,7 +104,8 @@ 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_op.hip.cu
+    module_base/kernels/rocm/math_kernel_op.hip_vec.cu
+    module_base/kernels/rocm/math_ylm_op.hip.cu
     module_hamilt_general/module_xc/kernels/rocm/xc_functional_op.hip.cu
   )
 endif()

source/Makefile.Objects

@@ -146,8 +146,9 @@ OBJS_BASE=abfs-vector3_order.o\
     math_ylmreal.o\
     math_bspline.o\
     math_chebyshev.o\
-    math_op.o\
+    math_ylm_op.o\
     math_kernel_op.o\
+    math_kernel_op_vec.o\
     mathzone_add1.o\
     matrix.o\
     matrix3.o\
@@ -349,6 +350,7 @@ OBJS_HSOLVER=diago_cg.o\
     hsolver_pw_sdft.o\
     diago_iter_assist.o\
     dngvd_op.o\
+    bpcg_kernel_op.o\
     diag_const_nums.o\
     diag_hs_para.o\
     diago_pxxxgvx.o\