!13097 update sponge module and fix some bugs

From: @zhangxinfeng3
Reviewed-by: @wang_zi_dong,@ljl0711
Signed-off-by: @ljl0711

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,

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

@@ -46,6 +46,8 @@ struct GRID_POINTER {
 void Construct_Neighbor_List(int grid_numbers, int max_neighbor_numbers, int *nl_atom_numbers, int *nl_atom_serial,
                              NEIGHBOR_LIST *nl, cudaStream_t stream);
 
+void CopyNeighborListAtomNumber(int atom_numbers, NEIGHBOR_LIST *nl, int *nl_atom_numbers, cudaStream_t stream);
+
 void Neighbor_List_Update(int grid_numbers, int atom_numbers, int refresh_count, int refresh_interval,
                           int not_first_time, float skin, int Nxy, float cutoff_square, float cutoff_with_skin_square,
                           int *grid_N, float *box_length, int *atom_numbers_in_grid_bucket, float *grid_length_inverse,

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_excluded_force_impl.cu

@@ -86,6 +86,7 @@ __global__ void PME_Excluded_Force_Correction(const int atom_numbers, const UNSI
 void PMEExcludedForce(const int atom_numbers, const float pme_beta, const int *uint_crd_f, const float *sacler_f,
                       const float *charge, const int *excluded_list_start, const int *excluded_list,
                       const int *excluded_atom_numbers, float *frc_f, cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.);
   UNSIGNED_INT_VECTOR *uint_crd =
     const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f));
   VECTOR *frc = reinterpret_cast<VECTOR *>(frc_f);

mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu

@@ -75,6 +75,7 @@ void PMEReciprocalForce(int fftx, int ffty, int fftz, int atom_numbers, float be
                         float *pme_frxyz, float *PME_Q, float *pme_fq, int *PME_atom_near, int *pme_kxyz,
                         const float *box_length_f, const int *uint_crd_f, const float *charge, float *force,
                         cudaStream_t stream) {
+  Reset_List<<<ceilf(static_cast<float>(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, force, 0.);
   UNSIGNED_INT_VECTOR *uint_crd =
     const_cast<UNSIGNED_INT_VECTOR *>(reinterpret_cast<const UNSIGNED_INT_VECTOR *>(uint_crd_f));
   UNSIGNED_INT_VECTOR *PME_uxyz = reinterpret_cast<UNSIGNED_INT_VECTOR *>(pme_uxyz);

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/angle/angle_atom_energy_kernel.h

@@ -69,8 +69,8 @@ class AngleAtomEnergyGpuKernel : public GpuKernel {
     auto angle_theta0 = GetDeviceAddress<T>(inputs, 6);
 
     auto ene = GetDeviceAddress<T>(outputs, 0);
-    AngleAtomEnergy(angle_numbers, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, angle_k, angle_theta0, ene,
-                    reinterpret_cast<cudaStream_t>(stream_ptr));
+    AngleAtomEnergy(angle_numbers, ele_uint_crd, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, angle_k, angle_theta0,
+                    ene, reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
 

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/angle/angle_force_kernel.h

@@ -69,7 +69,7 @@ class AngleForceGpuKernel : public GpuKernel {
     auto angle_theta0 = GetDeviceAddress<T>(inputs, 6);
 
     auto frc_f = GetDeviceAddress<T>(outputs, 0);
-    AngleForce(angle_numbers, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, angle_k, angle_theta0, frc_f,
+    AngleForce(angle_numbers, ele_uint_crd, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, angle_k, angle_theta0, frc_f,
                reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/angle/angle_force_with_atom_energy_kernel.h

@@ -70,8 +70,8 @@ class AngleForceWithAtomEnergyGpuKernel : public GpuKernel {
 
     auto frc_f = GetDeviceAddress<T>(outputs, 0);
     auto ene = GetDeviceAddress<T>(outputs, 1);
-    AngleForceWithAtomEnergy(angle_numbers, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, angle_k, angle_theta0, frc_f,
-                             ene, reinterpret_cast<cudaStream_t>(stream_ptr));
+    AngleForceWithAtomEnergy(angle_numbers, ele_uint_crd, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, angle_k,
+                             angle_theta0, frc_f, ene, reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
 

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/bond/bond_atom_energy_cuda_gpu_kernel.h

@@ -39,7 +39,7 @@ class BondAtomEnergyCudaGpuKernel : public GpuKernel {
   bool Init(const CNodePtr &kernel_node) override {
     kernel_node_ = kernel_node;
     bond_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "bond_numbers"));
-
+    atom_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "atom_numbers"));
     auto shape_uint_crd = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
     auto shape_scaler = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
     auto shape_atom_a = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 2);
@@ -73,7 +73,7 @@ class BondAtomEnergyCudaGpuKernel : public GpuKernel {
 
     auto atom_ene = GetDeviceAddress<T>(outputs, 0);
 
-    BondAtomEnergy(bond_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, atom_ene,
+    BondAtomEnergy(bond_numbers, atom_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, atom_ene,
                    reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
@@ -87,7 +87,7 @@ class BondAtomEnergyCudaGpuKernel : public GpuKernel {
     input_size_list_.push_back(ele_bond_k * sizeof(T));
     input_size_list_.push_back(ele_bond_r0 * sizeof(T));
 
-    output_size_list_.push_back(bond_numbers * sizeof(T));
+    output_size_list_.push_back(atom_numbers * sizeof(T));
   }
 
  private:
@@ -102,6 +102,7 @@ class BondAtomEnergyCudaGpuKernel : public GpuKernel {
   std::vector<size_t> output_size_list_;
   std::vector<size_t> workspace_size_list_;
   int bond_numbers;
+  int atom_numbers;
 };
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/bond/bond_energy_cuda_gpu_kernel.h

@@ -40,7 +40,7 @@ class BondEnergyCudaGpuKernel : public GpuKernel {
     // get bond_numbers
     kernel_node_ = kernel_node;
     bond_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "bond_numbers"));
-
+    atom_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "atom_numbers"));
     auto shape_uint_crd = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
     auto shape_scaler = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
     auto shape_atom_a = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 2);
@@ -74,7 +74,7 @@ class BondEnergyCudaGpuKernel : public GpuKernel {
 
     auto bond_ene = GetDeviceAddress<T>(outputs, 0);
 
-    BondEnergy(bond_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, bond_ene,
+    BondEnergy(bond_numbers, atom_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, bond_ene,
                reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
@@ -103,6 +103,7 @@ class BondEnergyCudaGpuKernel : public GpuKernel {
   std::vector<size_t> output_size_list_;
   std::vector<size_t> workspace_size_list_;
   int bond_numbers;
+  int atom_numbers;
 };
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/bond/bond_force_cuda_gpu_kernel.h

@@ -40,7 +40,7 @@ class BondForceCudaGpuKernel : public GpuKernel {
     // get bond_numbers
     kernel_node_ = kernel_node;
     bond_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "bond_numbers"));
-
+    atom_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "atom_numbers"));
     auto shape_uint_crd = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
     auto shape_scaler = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
     auto shape_atom_a = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 2);
@@ -74,7 +74,7 @@ class BondForceCudaGpuKernel : public GpuKernel {
 
     auto frc_f = GetDeviceAddress<T>(outputs, 0);
 
-    BondForce(bond_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, frc_f,
+    BondForce(bond_numbers, atom_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, frc_f,
               reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
@@ -88,7 +88,7 @@ class BondForceCudaGpuKernel : public GpuKernel {
     input_size_list_.push_back(ele_bond_k * sizeof(T));
     input_size_list_.push_back(ele_bond_r0 * sizeof(T));
 
-    output_size_list_.push_back(bond_numbers * 3 * sizeof(T));
+    output_size_list_.push_back(atom_numbers * 3 * sizeof(T));
   }
 
  private:
@@ -103,6 +103,7 @@ class BondForceCudaGpuKernel : public GpuKernel {
   std::vector<size_t> output_size_list_;
   std::vector<size_t> workspace_size_list_;
   int bond_numbers;
+  int atom_numbers;
 };
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/bond/bond_force_with_atom_energy_kernel.h

@@ -40,6 +40,7 @@ class BondForceWithAtomEnergyGpuKernel : public GpuKernel {
     // get bond_numbers
     kernel_node_ = kernel_node;
     bond_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "bond_numbers"));
+    atom_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "atom_numbers"));
     auto shape_uint_crd = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
     auto shape_scaler = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
     auto shape_atom_a = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 2);
@@ -72,8 +73,8 @@ class BondForceWithAtomEnergyGpuKernel : public GpuKernel {
 
     auto frc_f = GetDeviceAddress<T>(outputs, 0);
     auto atom_e = GetDeviceAddress<T>(outputs, 1);
-    BondForceWithAtomEnergy(bond_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, frc_f, atom_e,
-                            reinterpret_cast<cudaStream_t>(stream_ptr));
+    BondForceWithAtomEnergy(bond_numbers, atom_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, frc_f,
+                            atom_e, reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
 
@@ -86,8 +87,8 @@ class BondForceWithAtomEnergyGpuKernel : public GpuKernel {
     input_size_list_.push_back(ele_bond_k * sizeof(T));
     input_size_list_.push_back(ele_bond_r0 * sizeof(T));
 
-    output_size_list_.push_back(bond_numbers * 3 * sizeof(T));
-    output_size_list_.push_back(bond_numbers * sizeof(T));
+    output_size_list_.push_back(atom_numbers * 3 * sizeof(T));
+    output_size_list_.push_back(atom_numbers * sizeof(T));
   }
 
  private:
@@ -102,6 +103,7 @@ class BondForceWithAtomEnergyGpuKernel : public GpuKernel {
   std::vector<size_t> output_size_list_;
   std::vector<size_t> workspace_size_list_;
   int bond_numbers;
+  int atom_numbers;
 };
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/bond/bond_force_with_atom_virial_kernel.h

@@ -40,6 +40,7 @@ class BondForceWithAtomVirialGpuKernel : public GpuKernel {
     // get bond_numbers
     kernel_node_ = kernel_node;
     bond_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "bond_numbers"));
+    atom_numbers = static_cast<int>(GetAttr<int64_t>(kernel_node, "atom_numbers"));
     auto shape_uint_crd = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
     auto shape_scaler = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
     auto shape_atom_a = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 2);
@@ -72,8 +73,8 @@ class BondForceWithAtomVirialGpuKernel : public GpuKernel {
 
     auto frc_f = GetDeviceAddress<T>(outputs, 0);
     auto atom_v = GetDeviceAddress<T>(outputs, 1);
-    BondForceWithAtomVirial(bond_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, frc_f, atom_v,
-                            reinterpret_cast<cudaStream_t>(stream_ptr));
+    BondForceWithAtomVirial(bond_numbers, atom_numbers, uint_crd_f, scaler_f, atom_a, atom_b, bond_k, bond_r0, frc_f,
+                            atom_v, reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
 
@@ -86,8 +87,8 @@ class BondForceWithAtomVirialGpuKernel : public GpuKernel {
     input_size_list_.push_back(ele_bond_k * sizeof(T));
     input_size_list_.push_back(ele_bond_r0 * sizeof(T));
 
-    output_size_list_.push_back(bond_numbers * 3 * sizeof(T));
-    output_size_list_.push_back(bond_numbers * sizeof(T));
+    output_size_list_.push_back(atom_numbers * 3 * sizeof(T));
+    output_size_list_.push_back(atom_numbers * sizeof(T));
   }
 
  private:
@@ -102,6 +103,7 @@ class BondForceWithAtomVirialGpuKernel : public GpuKernel {
   std::vector<size_t> output_size_list_;
   std::vector<size_t> workspace_size_list_;
   int bond_numbers;
+  int atom_numbers;
 };
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/dihedral/dihedral_atom_energy_kernel.h

@@ -81,8 +81,8 @@ class DihedralAtomEnergyGpuKernel : public GpuKernel {
     auto pn = GetDeviceAddress<T>(inputs, 10);
 
     auto ene = GetDeviceAddress<T>(outputs, 0);
-    DihedralAtomEnergy(dihedral_numbers, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, atom_d, ipn, pk, gamc, gams, pn,
-                       ene, reinterpret_cast<cudaStream_t>(stream_ptr));
+    DihedralAtomEnergy(dihedral_numbers, ele_uint_crd, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, atom_d, ipn, pk,
+                       gamc, gams, pn, ene, reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
 

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/dihedral/dihedral_force_kernel.h

@@ -81,8 +81,8 @@ class DihedralForceGpuKernel : public GpuKernel {
     auto pn = GetDeviceAddress<T>(inputs, 10);
 
     auto frc_f = GetDeviceAddress<T>(outputs, 0);
-    DihedralForce(dihedral_numbers, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, atom_d, ipn, pk, gamc, gams, pn,
-                  frc_f, reinterpret_cast<cudaStream_t>(stream_ptr));
+    DihedralForce(dihedral_numbers, ele_uint_crd, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, atom_d, ipn, pk, gamc,
+                  gams, pn, frc_f, reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
 

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/dihedral/dihedral_force_with_atom_energy_kernel.h

@@ -82,8 +82,8 @@ class DihedralForceWithAtomEnergyGpuKernel : public GpuKernel {
 
     auto frc_f = GetDeviceAddress<T>(outputs, 0);
     auto ene = GetDeviceAddress<T>(outputs, 1);
-    DihedralForceWithAtomEnergy(dihedral_numbers, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, atom_d, ipn, pk, gamc,
-                                gams, pn, frc_f, ene, reinterpret_cast<cudaStream_t>(stream_ptr));
+    DihedralForceWithAtomEnergy(dihedral_numbers, ele_uint_crd, uint_crd_f, scaler_f, atom_a, atom_b, atom_c, atom_d,
+                                ipn, pk, gamc, gams, pn, frc_f, ene, reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
 

mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/neighbor_list/neighbor_list_update_kernel.h

@@ -104,7 +104,7 @@ class NeighborListUpdateGpuKernel : public GpuKernel {
                          half_crd_to_uint_crd_cof, uint_crd, uint_dr_to_dr_cof, d_gpointer, nl, excluded_list_start,
                          excluded_list, excluded_numbers, half_skin_square, need_refresh_flag,
                          reinterpret_cast<cudaStream_t>(stream_ptr));
-
+    CopyNeighborListAtomNumber(atom_numbers, nl, nl_atom_numbers, reinterpret_cast<cudaStream_t>(stream_ptr));
     return true;
   }
 

model_zoo/research/hpc/sponge/README.md

@@ -0,0 +1 @@
+# Contents