update sponge

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_atom_energy_impl.cu

@@ -47,9 +47,10 @@ __global__ void AngleAtomEnergyKernel(int angle_numbers, const UNSIGNED_INT_VECT
   }
 }
 
-void AngleAtomEnergy(int angle_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                     const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0, float *ene,
-                     cudaStream_t stream) {
+void AngleAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                     const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k,
+                     const float *angle_theta0, float *ene, cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(angle_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -60,6 +61,6 @@ void AngleAtomEnergy(int angle_numbers, const int *uint_crd_f, const float *scal
                                                                          atom_b, atom_c, angle_k, angle_theta0, ene);
   return;
 }
-void AngleAtomEnergy(int angle_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                     const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0, float *ene,
-                     cudaStream_t stream);
+void AngleAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                     const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k,
+                     const float *angle_theta0, float *ene, cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_atom_energy_impl.cuh

@@ -20,7 +20,7 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void AngleAtomEnergy(int angle_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                     const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0, float *ene,
-                     cudaStream_t stream);
+void AngleAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                     const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k,
+                     const float *angle_theta0, float *ene, cudaStream_t stream);
 #endif

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_impl.cu

@@ -66,8 +66,10 @@ __global__ void AngleForceKernel(int angle_numbers, const UNSIGNED_INT_VECTOR *u
   }
 }
 
-void AngleForce(int angle_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, const int *atom_b,
-                const int *atom_c, const float *angle_k, const float *angle_theta0, float *frc_f, cudaStream_t stream) {
+void AngleForce(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+                const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0, float *frc_f,
+                cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(angle_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -79,5 +81,6 @@ void AngleForce(int angle_numbers, const int *uint_crd_f, const float *scaler_f,
                                                                     atom_c, angle_k, angle_theta0, frc);
   return;
 }
-void AngleForce(int angle_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, const int *atom_b,
-                const int *atom_c, const float *angle_k, const float *angle_theta0, float *frc_f, cudaStream_t stream);
+void AngleForce(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+                const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0, float *frc_f,
+                cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_impl.cuh

@@ -20,6 +20,7 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void AngleForce(int angle_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, const int *atom_b,
-                const int *atom_c, const float *angle_k, const float *angle_theta0, float *frc_f, cudaStream_t stream);
+void AngleForce(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+                const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0, float *frc_f,
+                cudaStream_t stream);
 #endif

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_with_atom_energy_impl.cu

@@ -70,9 +70,10 @@ __global__ void AngleForceWithAtomEnergyKernel(int angle_numbers, const UNSIGNED
   }
 }
 
-void AngleForceWithAtomEnergy(int angle_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                              const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0,
-                              float *frc_f, float *ene, cudaStream_t stream) {
+void AngleForceWithAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                              const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k,
+                              const float *angle_theta0, float *frc_f, float *ene, cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(angle_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -84,6 +85,6 @@ void AngleForceWithAtomEnergy(int angle_numbers, const int *uint_crd_f, const fl
     angle_numbers, uint_crd, scaler, atom_a, atom_b, atom_c, angle_k, angle_theta0, frc, ene);
   return;
 }
-void AngleForceWithAtomEnergy(int angle_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                              const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0,
-                              float *frc_f, float *ene, cudaStream_t stream);
+void AngleForceWithAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                              const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k,
+                              const float *angle_theta0, float *frc_f, float *ene, cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_with_atom_energy_impl.cuh

@@ -20,7 +20,7 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void AngleForceWithAtomEnergy(int angle_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                              const int *atom_b, const int *atom_c, const float *angle_k, const float *angle_theta0,
-                              float *frc_f, float *ene, cudaStream_t stream);
+void AngleForceWithAtomEnergy(int angle_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                              const int *atom_a, const int *atom_b, const int *atom_c, const float *angle_k,
+                              const float *angle_theta0, float *frc_f, float *ene, cudaStream_t stream);
 #endif

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_atom_energy_cuda_gpu_impl.cu

@@ -38,9 +38,10 @@ __global__ void BondAtomEnergyCudaKernel(const int bond_numbers, const UNSIGNED_
   }
 }
 
-void BondAtomEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+void BondAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
                     const int *atom_b, const float *bond_k, const float *bond_r0, float *atom_ene,
                     cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, atom_ene, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -52,5 +53,5 @@ void BondAtomEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler
   return;
 }
 
-void BondAtomEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+void BondAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
                     const int *atom_b, const float *bond_k, const float *bond_r0, float *atom_ene, cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_atom_energy_cuda_gpu_impl.cuh

@@ -20,7 +20,7 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void BondAtomEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+void BondAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
                     const int *atom_b, const float *bond_k, const float *bond_r0, float *atom_ene, cudaStream_t stream);
 
 #endif  // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_BOND_ATOM_ENERGY_GPU_IMPL_H_

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_energy_cuda_gpu_impl.cu

@@ -38,8 +38,8 @@ __global__ void BondEnergyCudaKernel(const int bond_numbers, const UNSIGNED_INT_
   }
 }
 
-void BondEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, const int *atom_b,
-                const float *bond_k, const float *bond_r0, float *bond_ene, cudaStream_t stream) {
+void BondEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+                const int *atom_b, const float *bond_k, const float *bond_r0, float *bond_ene, cudaStream_t stream) {
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -51,5 +51,5 @@ void BondEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f,
   return;
 }
 
-void BondEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, const int *atom_b,
-                const float *bond_k, const float *bond_r0, float *bond_ene, cudaStream_t stream);
+void BondEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+                const int *atom_b, const float *bond_k, const float *bond_r0, float *bond_ene, cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_energy_cuda_gpu_impl.cuh

@@ -20,7 +20,7 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void BondEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, const int *atom_b,
-                const float *bond_k, const float *bond_r0, float *bond_ene, cudaStream_t stream);
+void BondEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+                const int *atom_b, const float *bond_k, const float *bond_r0, float *bond_ene, cudaStream_t stream);
 
 #endif  // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_BOND_ENERGY_CUDA_GPU_IMPL_H_

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_cuda_gpu_impl.cu

@@ -43,8 +43,9 @@ __global__ void BondForceCudaKernel(int bond_numbers, const UNSIGNED_INT_VECTOR
   }
 }
 
-void BondForce(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, const int *atom_b,
-               const float *bond_k, const float *bond_r0, float *frc_f, cudaStream_t stream) {
+void BondForce(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+               const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -56,5 +57,5 @@ void BondForce(int bond_numbers, const int *uint_crd_f, const float *scaler_f, c
   return;
 }
 
-void BondForce(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, const int *atom_b,
-               const float *bond_k, const float *bond_r0, float *frc_f, cudaStream_t stream);
+void BondForce(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+               const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_cuda_gpu_impl.cuh

@@ -20,6 +20,6 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void BondForce(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a, const int *atom_b,
-               const float *bond_k, const float *bond_r0, float *frc_f, cudaStream_t stream);
+void BondForce(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
+               const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, cudaStream_t stream);
 #endif  // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_BOND_FORCE_CUDA_GPU_IMPL_H_

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_energy_impl.cu

@@ -49,9 +49,11 @@ __global__ void BondForceWithAtomEnergyKernel(int bond_numbers, const UNSIGNED_I
   }
 }
 
-void BondForceWithAtomEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                             const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, float *atom_e,
-                             cudaStream_t stream) {
+void BondForceWithAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                             const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0,
+                             float *frc_f, float *atom_e, cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
+  Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, atom_e, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -62,6 +64,6 @@ void BondForceWithAtomEnergy(int bond_numbers, const int *uint_crd_f, const floa
                                                                                  atom_b, bond_k, bond_r0, frc, atom_e);
   return;
 }
-void BondForceWithAtomEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                             const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, float *atom_e,
-                             cudaStream_t stream);
+void BondForceWithAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                             const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0,
+                             float *frc_f, float *atom_e, cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_energy_impl.cuh

@@ -20,7 +20,7 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void BondForceWithAtomEnergy(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                             const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, float *atom_e,
-                             cudaStream_t stream);
+void BondForceWithAtomEnergy(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                             const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0,
+                             float *frc_f, float *atom_e, cudaStream_t stream);
 #endif

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_virial_impl.cu

@@ -49,9 +49,11 @@ __global__ void BondForceWithAtomVirialKernel(int bond_numbers, const UNSIGNED_I
   }
 }
 
-void BondForceWithAtomVirial(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                             const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, float *atom_v,
-                             cudaStream_t stream) {
+void BondForceWithAtomVirial(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                             const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0,
+                             float *frc_f, float *atom_v, cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
+  Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, atom_v, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(bond_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -62,6 +64,6 @@ void BondForceWithAtomVirial(int bond_numbers, const int *uint_crd_f, const floa
                                                                                  atom_b, bond_k, bond_r0, frc, atom_v);
   return;
 }
-void BondForceWithAtomVirial(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                             const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, float *atom_v,
-                             cudaStream_t stream);
+void BondForceWithAtomVirial(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                             const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0,
+                             float *frc_f, float *atom_v, cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_virial_impl.cuh

@@ -20,7 +20,7 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void BondForceWithAtomVirial(int bond_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                             const int *atom_b, const float *bond_k, const float *bond_r0, float *frc_f, float *atom_v,
-                             cudaStream_t stream);
+void BondForceWithAtomVirial(int bond_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                             const int *atom_a, const int *atom_b, const float *bond_k, const float *bond_r0,
+                             float *frc_f, float *atom_v, cudaStream_t stream);
 #endif

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_atom_energy_impl.cu

@@ -63,9 +63,11 @@ __global__ void DihedralAtomEnergyKernel(int dihedral_numbers, const UNSIGNED_IN
   }
 }
 
-void DihedralAtomEnergy(int dihedral_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                        const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn, const float *pk,
-                        const float *gamc, const float *gams, const float *pn, float *ene, cudaStream_t stream) {
+void DihedralAtomEnergy(int dihedral_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                        const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
+                        const float *pk, const float *gamc, const float *gams, const float *pn, float *ene,
+                        cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(dihedral_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -76,6 +78,7 @@ void DihedralAtomEnergy(int dihedral_numbers, const int *uint_crd_f, const float
     dihedral_numbers, uint_crd, scaler, atom_a, atom_b, atom_c, atom_d, ipn, pk, gamc, gams, pn, ene);
   return;
 }
-void DihedralAtomEnergy(int dihedral_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                        const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn, const float *pk,
-                        const float *gamc, const float *gams, const float *pn, float *ene, cudaStream_t stream);
+void DihedralAtomEnergy(int dihedral_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                        const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
+                        const float *pk, const float *gamc, const float *gams, const float *pn, float *ene,
+                        cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_atom_energy_impl.cuh

@@ -20,7 +20,8 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void DihedralAtomEnergy(int dihedral_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                        const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn, const float *pk,
-                        const float *gamc, const float *gams, const float *pn, float *ene, cudaStream_t stream);
+void DihedralAtomEnergy(int dihedral_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                        const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
+                        const float *pk, const float *gamc, const float *gams, const float *pn, float *ene,
+                        cudaStream_t stream);
 #endif

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_impl.cu

@@ -99,9 +99,11 @@ __global__ void DihedralForceKernel(int dihedral_numbers, const UNSIGNED_INT_VEC
   }
 }
 
-void DihedralForce(int dihedral_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                   const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn, const float *pk,
-                   const float *gamc, const float *gams, const float *pn, float *frc_f, cudaStream_t stream) {
+void DihedralForce(int dihedral_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                   const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
+                   const float *pk, const float *gamc, const float *gams, const float *pn, float *frc_f,
+                   cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(dihedral_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -113,6 +115,7 @@ void DihedralForce(int dihedral_numbers, const int *uint_crd_f, const float *sca
     dihedral_numbers, uint_crd, scaler, atom_a, atom_b, atom_c, atom_d, ipn, pk, gamc, gams, pn, frc);
   return;
 }
-void DihedralForce(int dihedral_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                   const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn, const float *pk,
-                   const float *gamc, const float *gams, const float *pn, float *frc_f, cudaStream_t stream);
+void DihedralForce(int dihedral_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                   const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
+                   const float *pk, const float *gamc, const float *gams, const float *pn, float *frc_f,
+                   cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_impl.cuh

@@ -20,7 +20,8 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void DihedralForce(int dihedral_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                   const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn, const float *pk,
-                   const float *gamc, const float *gams, const float *pn, float *frc_f, cudaStream_t stream);
+void DihedralForce(int dihedral_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                   const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
+                   const float *pk, const float *gamc, const float *gams, const float *pn, float *frc_f,
+                   cudaStream_t stream);
 #endif

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_with_atom_energy_impl.cu

@@ -103,10 +103,11 @@ __global__ void DihedralForceWithAtomEnergyKernel(int dihedral_numbers, const UN
   }
 }
 
-void DihedralForceWithAtomEnergy(int dihedral_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                                 const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
-                                 const float *pk, const float *gamc, const float *gams, const float *pn, float *frc_f,
-                                 float *ene, cudaStream_t stream) {
+void DihedralForceWithAtomEnergy(int dihedral_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                                 const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d,
+                                 const int *ipn, const float *pk, const float *gamc, const float *gams, const float *pn,
+                                 float *frc_f, float *ene, cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
   size_t thread_per_block = 128;
   size_t block_per_grid = ceilf(static_cast<float>(dihedral_numbers) / 128);
   UNSIGNED_INT_VECTOR *uint_crd =
@@ -118,7 +119,7 @@ void DihedralForceWithAtomEnergy(int dihedral_numbers, const int *uint_crd_f, co
     dihedral_numbers, uint_crd, scaler, atom_a, atom_b, atom_c, atom_d, ipn, pk, gamc, gams, pn, frc, ene);
   return;
 }
-void DihedralForceWithAtomEnergy(int dihedral_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                                 const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
-                                 const float *pk, const float *gamc, const float *gams, const float *pn, float *frc_f,
-                                 float *ene, cudaStream_t stream);
+void DihedralForceWithAtomEnergy(int dihedral_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                                 const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d,
+                                 const int *ipn, const float *pk, const float *gamc, const float *gams, const float *pn,
+                                 float *frc_f, float *ene, cudaStream_t stream);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_with_atom_energy_impl.cuh

@@ -20,8 +20,8 @@
 #include <curand_kernel.h>
 #include "runtime/device/gpu/cuda_common.h"
 
-void DihedralForceWithAtomEnergy(int dihedral_numbers, const int *uint_crd_f, const float *scaler_f, const int *atom_a,
-                                 const int *atom_b, const int *atom_c, const int *atom_d, const int *ipn,
-                                 const float *pk, const float *gamc, const float *gams, const float *pn, float *frc_f,
-                                 float *ene, cudaStream_t stream);
+void DihedralForceWithAtomEnergy(int dihedral_numbers, int atom_numbers, const int *uint_crd_f, const float *scaler_f,
+                                 const int *atom_a, const int *atom_b, const int *atom_c, const int *atom_d,
+                                 const int *ipn, const float *pk, const float *gamc, const float *gams, const float *pn,
+                                 float *frc_f, float *ene, cudaStream_t stream);
 #endif

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_impl.cu

@@ -92,6 +92,7 @@ void LJForce(const int atom_numbers, const float cutoff_square, const int *uint_
              const float *charge, const float *scaler_f, float *uint_crd_with_LJ, int *nl_atom_numbers,
              int *nl_atom_serial, int *nl, const float *d_LJ_A, const float *d_LJ_B, float *frc_f,
              cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
   VECTOR *frc = reinterpret_cast<VECTOR *>(frc_f);
   VECTOR *scaler = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(scaler_f));
   int max_neighbor_numbers = 800;

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_with_pme_direct_force_impl.cu

@@ -106,6 +106,7 @@ void LJForceWithPMEDirectForce(const int atom_numbers, const float cutoff, const
                                const int *LJtype, const float *charge, const float *scaler_f, float *uint_crd_with_LJ,
                                int *nl_atom_numbers, int *nl_atom_serial, int *nl, const float *d_LJ_A,
                                const float *d_LJ_B, float *frc_f, cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
   VECTOR *frc = reinterpret_cast<VECTOR *>(frc_f);
   VECTOR *scaler = const_cast<VECTOR *>(reinterpret_cast<const VECTOR *>(scaler_f));
   int max_neighbor_numbers = 800;

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/neighbor_list/neighbor_list_impl.cu

@@ -330,6 +330,17 @@ void Construct_Neighbor_List(int atom_numbers, int max_neighbor_numbers, int *nl
     atom_numbers, max_neighbor_numbers, nl_atom_numbers, nl_atom_serial, nl);
 }
 
+__global__ void copy_neighbor_list_atom_number(int atom_numbers, NEIGHBOR_LIST *nl, int *nl_atom_numbers) {
+  for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < atom_numbers; i += gridDim.x * blockDim.x) {
+    nl_atom_numbers[i] = nl[i].atom_numbers;
+  }
+}
+
+void CopyNeighborListAtomNumber(int atom_numbers, NEIGHBOR_LIST *nl, int *nl_atom_numbers, cudaStream_t stream) {
+  copy_neighbor_list_atom_number<<<ceilf(static_cast<float>(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, nl,
+                                                                                                    nl_atom_numbers);
+}
+
 void Refresh_Neighbor_List_No_Check(int grid_numbers, int atom_numbers, float skin, int Nxy, float cutoff_skin_square,
                                     int *grid_N, float *box_length, int *atom_numbers_in_grid_bucket,
                                     float *grid_length_inverse, int *atom_in_grid_serial, GRID_BUCKET *bucket,