From aa7fdf53a77cae24941656b7431f46ce54a803a1 Mon Sep 17 00:00:00 2001 From: q00596439 Date: Thu, 25 Mar 2021 20:42:18 +0800 Subject: [PATCH] master 03252 sponge ops zoo cu --- .../sponge/angle/angle_atom_energy_impl.cu | 2 +- .../sponge/angle/angle_force_impl.cu | 2 +- .../angle_force_with_atom_energy_impl.cu | 2 +- .../bond/bond_atom_energy_cuda_gpu_impl.cu | 2 +- .../sponge/bond/bond_force_cuda_gpu_impl.cu | 2 +- .../bond/bond_force_with_atom_energy_impl.cu | 4 +- .../bond/bond_force_with_atom_virial_impl.cu | 4 +- .../dihedral/dihedral_atom_energy_impl.cu | 2 +- .../sponge/dihedral/dihedral_force_impl.cu | 2 +- .../dihedral_force_with_atom_energy_impl.cu | 2 +- .../gpu/cuda_impl/sponge/lj/lj_force_impl.cu | 2 +- .../lj/lj_force_with_pme_direct_force_impl.cu | 2 +- .../nb14/dihedral_14_cf_atom_energy_impl.cu | 2 +- .../sponge/nb14/dihedral_14_cf_energy_impl.cu | 15 +- .../nb14/dihedral_14_cf_energy_impl.cuh | 4 +- .../nb14/dihedral_14_lj_atom_energy_impl.cu | 2 +- ...al_14_lj_cf_force_with_atom_energy_impl.cu | 30 +- ...l_14_lj_cf_force_with_atom_energy_impl.cuh | 9 +- .../sponge/nb14/dihedral_14_lj_energy_impl.cu | 19 +- .../nb14/dihedral_14_lj_energy_impl.cuh | 6 +- ...ihedral_14_lj_force_with_direct_cf_impl.cu | 2 +- .../nvtit/md_iteration_leap_frog_impl.cu | 3 - .../sponge/pme/pme_excluded_force_impl.cu | 2 +- .../sponge/pme/pme_reciprocal_force_impl.cu | 2 +- .../nb14/dihedral_14_cf_energy_kernel.h | 16 +- ...l_14_lj_cf_force_with_atom_energy_kernel.h | 18 +- .../nb14/dihedral_14_lj_energy_kernel.h | 16 +- .../gpu/sponge/pme/pme_energy_kernel.h | 3 +- .../sponge/pme/pme_excluded_force_kernel.h | 3 +- mindspore/ops/operations/sponge_ops.py | 435 +++++++++++++----- model_zoo/research/hpc/sponge/src/__init__.py | 0 model_zoo/research/hpc/sponge/src/nb14.py | 2 +- .../hpc/sponge/src/simulation_initial.py | 22 +- 33 files changed, 441 insertions(+), 198 deletions(-) create mode 100644 model_zoo/research/hpc/sponge/src/__init__.py diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_atom_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_atom_energy_impl.cu index 7d1299e28a3..0d9ea99fb7a 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_atom_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_atom_energy_impl.cu @@ -50,7 +50,7 @@ __global__ void AngleAtomEnergyKernel(int angle_numbers, const UNSIGNED_INT_VECT 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<<(atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); + Reset_List<<(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(angle_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_impl.cu index a69942e3e9f..3eb87afad37 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_impl.cu @@ -69,7 +69,7 @@ __global__ void AngleForceKernel(int angle_numbers, const UNSIGNED_INT_VECTOR *u 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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(angle_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_with_atom_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_with_atom_energy_impl.cu index 049d9da5752..1ce1cff5438 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_with_atom_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/angle/angle_force_with_atom_energy_impl.cu @@ -73,7 +73,7 @@ __global__ void AngleForceWithAtomEnergyKernel(int angle_numbers, const UNSIGNED 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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(angle_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_atom_energy_cuda_gpu_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_atom_energy_cuda_gpu_impl.cu index f5734dd39d1..c38acc3571a 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_atom_energy_cuda_gpu_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_atom_energy_cuda_gpu_impl.cu @@ -41,7 +41,7 @@ __global__ void BondAtomEnergyCudaKernel(const int bond_numbers, const UNSIGNED_ 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<<(atom_numbers) / 128), 128>>>(atom_numbers, atom_ene, 0.); + Reset_List<<(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, atom_ene, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(bond_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_cuda_gpu_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_cuda_gpu_impl.cu index 2d06a9dee8e..aa4564926d0 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_cuda_gpu_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_cuda_gpu_impl.cu @@ -45,7 +45,7 @@ __global__ void BondForceCudaKernel(int bond_numbers, const UNSIGNED_INT_VECTOR 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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(bond_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_energy_impl.cu index 5a411146ee0..2dd19acc73e 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_energy_impl.cu @@ -52,8 +52,8 @@ __global__ void BondForceWithAtomEnergyKernel(int bond_numbers, const UNSIGNED_I 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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); - Reset_List<<(atom_numbers) / 128), 128>>>(atom_numbers, atom_e, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, atom_e, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(bond_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_virial_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_virial_impl.cu index 23760957ba6..bf04be130f2 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_virial_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/bond/bond_force_with_atom_virial_impl.cu @@ -52,8 +52,8 @@ __global__ void BondForceWithAtomVirialKernel(int bond_numbers, const UNSIGNED_I 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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); - Reset_List<<(atom_numbers) / 128), 128>>>(atom_numbers, atom_v, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, atom_v, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(bond_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_atom_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_atom_energy_impl.cu index 310055e4205..b313b3b78a6 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_atom_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_atom_energy_impl.cu @@ -67,7 +67,7 @@ void DihedralAtomEnergy(int dihedral_numbers, int atom_numbers, const int *uint_ 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<<(atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); + Reset_List<<(atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(dihedral_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_impl.cu index b9fba0ea74d..40abe15adef 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_impl.cu @@ -103,7 +103,7 @@ void DihedralForce(int dihedral_numbers, int atom_numbers, const int *uint_crd_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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(dihedral_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_with_atom_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_with_atom_energy_impl.cu index ad2e1c49dc0..a82c080945c 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_with_atom_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/dihedral/dihedral_force_with_atom_energy_impl.cu @@ -107,7 +107,7 @@ void DihedralForceWithAtomEnergy(int dihedral_numbers, int atom_numbers, const i 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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(dihedral_numbers) / 128); UNSIGNED_INT_VECTOR *uint_crd = diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_impl.cu index bf526e9f7b7..1cb040aff69 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_impl.cu @@ -92,7 +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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); VECTOR *frc = reinterpret_cast(frc_f); VECTOR *scaler = const_cast(reinterpret_cast(scaler_f)); int max_neighbor_numbers = 800; diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_with_pme_direct_force_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_with_pme_direct_force_impl.cu index 4b00b7e5e7c..e1995d596e1 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_with_pme_direct_force_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/lj/lj_force_with_pme_direct_force_impl.cu @@ -106,7 +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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); VECTOR *frc = reinterpret_cast(frc_f); VECTOR *scaler = const_cast(reinterpret_cast(scaler_f)); int max_neighbor_numbers = 800; diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_atom_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_atom_energy_impl.cu index 44219d215e7..64a4fad8225 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_atom_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_atom_energy_impl.cu @@ -66,7 +66,7 @@ void Dihedral14CFAtomEnergy(const int dihedral_14_numbers, const int atom_number atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); VECTOR *boxlength = const_cast(reinterpret_cast(boxlength_f)); - Reset_List<<(3. * atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.); Dihedral14CFAtomEnergyKernel<<>>( dihedral_14_numbers, uint_crd_with_LJ, boxlength, a_14, b_14, cf_scale_factor, ene); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cu index 1bbde8bfa8f..9ed64ea7117 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cu @@ -52,21 +52,20 @@ __global__ void Dihedral14CFEnergyKernel(const int dihedral_14_numbers, const UI } void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, - const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, - const float *cf_scale_factor, float *ene, cudaStream_t stream) { + const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14, + const int *b_14, const float *cf_scale_factor, float *ene, cudaStream_t stream) { size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(atom_numbers) / 128); - UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = NULL; - Cuda_Malloc_Safely(reinterpret_cast(&uint_crd_with_LJ), sizeof(UINT_VECTOR_LJ_TYPE) * atom_numbers); - UNSIGNED_INT_VECTOR *uint_crd = const_cast(reinterpret_cast(uint_crd_f)); + UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = reinterpret_cast(uint_crd_with_LJ_f); + Copy_Crd_To_New_Crd_Start<<(atom_numbers) / 32), 32, 0, stream>>>( atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); VECTOR *boxlength = const_cast(reinterpret_cast(boxlength_f)); - Reset_List<<(3. * atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.); Dihedral14CFEnergyKernel<<>>( dihedral_14_numbers, uint_crd_with_LJ, boxlength, a_14, b_14, cf_scale_factor, ene); @@ -76,5 +75,5 @@ void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, c } void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, - const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, - const float *cf_scale_factor, float *ene, cudaStream_t stream); + const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14, + const int *b_14, const float *cf_scale_factor, float *ene, cudaStream_t stream); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cuh index 0e14bd7534d..117b070c284 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_cf_energy_impl.cuh @@ -20,6 +20,6 @@ #include "runtime/device/gpu/cuda_common.h" void Dihedral14CFEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, - const float *charge, const float *boxlength, const int *a_14, const int *b_14, - const float *cf_scale_factor, float *ene, cudaStream_t stream); + const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14, + const int *b_14, const float *cf_scale_factor, float *ene, cudaStream_t stream); #endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_NB14_DIHEDRAL_14_CF_ENERGY_IMPL_H diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_atom_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_atom_energy_impl.cu index ccd4ba05957..3beadecdf06 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_atom_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_atom_energy_impl.cu @@ -87,7 +87,7 @@ void Dihedral14LJAtomEnergy(const int dihedral_14_numbers, const int atom_number atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); VECTOR *boxlength = const_cast(reinterpret_cast(boxlength_f)); - Reset_List<<(3. * atom_numbers) / 128), 128>>>(atom_numbers, ene, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, ene, 0.); Dihedral14LJAtomEnergyKernel<<>>( dihedral_14_numbers, uint_crd_with_LJ, boxlength, a_14, b_14, lj_scale_factor, LJ_type_A, LJ_type_B, ene); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cu index dd757caab7f..730310c89c2 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cu @@ -105,22 +105,23 @@ __global__ void Dihedral14LJCFForceWithAtomEnergyKernel(const int dihedral_14_nu } void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, - const int *LJtype, const float *charge, const float *boxlength_f, - const int *a_14, const int *b_14, const float *lj_scale_factor, - const float *cf_scale_factor, const float *LJ_type_A, const float *LJ_type_B, - float *frc_f, float *atom_energy, cudaStream_t stream) { + const int *LJtype, const float *charge, float *uint_crd_with_LJ_f, + const float *boxlength_f, const int *a_14, const int *b_14, + const float *lj_scale_factor, const float *cf_scale_factor, + const float *LJ_type_A, const float *LJ_type_B, float *frc_f, float *atom_energy, + cudaStream_t stream) { size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(atom_numbers) / 128); - UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = NULL; - Cuda_Malloc_Safely(reinterpret_cast(&uint_crd_with_LJ), sizeof(UINT_VECTOR_LJ_TYPE) * atom_numbers); - UNSIGNED_INT_VECTOR *uint_crd = const_cast(reinterpret_cast(uint_crd_f)); + UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = reinterpret_cast(uint_crd_with_LJ_f); + Copy_Crd_To_New_Crd_Start<<(atom_numbers) / 32), 32, 0, stream>>>( atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); - Reset_List<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); - Reset_List<<(3. * atom_numbers) / 128), 128>>>(atom_numbers, atom_energy, 0.); + + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(atom_numbers, atom_energy, 0.); VECTOR *boxlength = const_cast(reinterpret_cast(boxlength_f)); VECTOR *frc = const_cast(reinterpret_cast(frc_f)); @@ -133,8 +134,9 @@ void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int return; } -void Dihedral14LJForceWithDirectCF(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, - const int *LJtype, const float *charge, const float *boxlength_f, const int *a_14, - const int *b_14, const float *lj_scale_factor, const float *cf_scale_factor, - const float *LJ_type_A, const float *LJ_type_B, float *frc, float *atom_energy, - cudaStream_t stream); +void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, + const int *LJtype, const float *charge, float *uint_crd_with_LJ_f, + const float *boxlength_f, const int *a_14, const int *b_14, + const float *lj_scale_factor, const float *cf_scale_factor, + const float *LJ_type_A, const float *LJ_type_B, float *frc_f, float *atom_energy, + cudaStream_t stream); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cuh index 4ea8262b3ca..c565ea23f8e 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_impl.cuh @@ -20,8 +20,9 @@ #include "runtime/device/gpu/cuda_common.h" void Dihedral14LJCFForceWithAtomEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, - const int *LJtype, const float *charge, const float *boxlength_f, - const int *a_14, const int *b_14, const float *lj_scale_factor, - const float *cf_scale_factor, const float *LJ_type_A, const float *LJ_type_B, - float *frc, float *atom_energy, cudaStream_t stream); + const int *LJtype, const float *charge, float *uint_crd_with_LJ_f, + const float *boxlength_f, const int *a_14, const int *b_14, + const float *lj_scale_factor, const float *cf_scale_factor, + const float *LJ_type_A, const float *LJ_type_B, float *frc_f, float *atom_energy, + cudaStream_t stream); #endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_NB14_DIHEDRAL_14_LJ_CF_FORCE_WITH_ATOM_ENERGY_IMPL_H diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cu index 3436c356134..54a24a565ac 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cu @@ -72,20 +72,19 @@ __global__ void Dihedral14LJEnergyKernel(const int dihedral_14_numbers, const UI } void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, - const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, - const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B, float *ene, - cudaStream_t stream) { + const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14, + const int *b_14, const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B, + float *ene, cudaStream_t stream) { size_t thread_per_block = 128; size_t block_per_grid = ceilf(static_cast(atom_numbers) / 128); - UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = NULL; - Cuda_Malloc_Safely(reinterpret_cast(&uint_crd_with_LJ), sizeof(UINT_VECTOR_LJ_TYPE) * atom_numbers); - UNSIGNED_INT_VECTOR *uint_crd = const_cast(reinterpret_cast(uint_crd_f)); + UINT_VECTOR_LJ_TYPE *uint_crd_with_LJ = reinterpret_cast(uint_crd_with_LJ_f); + Copy_Crd_To_New_Crd_Start<<(atom_numbers) / 32), 32, 0, stream>>>( atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); - Reset_List<<(3. * atom_numbers) / 128), 128>>>(dihedral_14_numbers, ene, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(dihedral_14_numbers, ene, 0.); VECTOR *boxlength = const_cast(reinterpret_cast(boxlength_f)); Dihedral14LJEnergyKernel<<>>( @@ -97,6 +96,6 @@ void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, c } void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, - const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, - const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B, float *ene, - cudaStream_t stream); + const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14, + const int *b_14, const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B, + float *ene, cudaStream_t stream); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cuh index 4a132438b98..fbf21d27e11 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_energy_impl.cuh @@ -20,8 +20,8 @@ #include "runtime/device/gpu/cuda_common.h" void Dihedral14LJEnergy(const int dihedral_14_numbers, const int atom_numbers, const int *uint_crd_f, const int *LJtype, - const float *charge, const float *boxlength_f, const int *a_14, const int *b_14, - const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B, float *ene, - cudaStream_t stream); + const float *charge, float *uint_crd_with_LJ_f, const float *boxlength_f, const int *a_14, + const int *b_14, const float *lj_scale_factor, const float *LJ_type_A, const float *LJ_type_B, + float *ene, cudaStream_t stream); #endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_NB14_DIHEDRAL_14_LJ_ENERGY_IMPL_H diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_force_with_direct_cf_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_force_with_direct_cf_impl.cu index 592d3a52063..88dc3656a71 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_force_with_direct_cf_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nb14/dihedral_14_lj_force_with_direct_cf_impl.cu @@ -108,7 +108,7 @@ void Dihedral14LJForceWithDirectCF(const int dihedral_14_numbers, const int atom atom_numbers, uint_crd, uint_crd_with_LJ, LJtype, charge); cudaStreamSynchronize(stream); VECTOR *boxlength = const_cast(reinterpret_cast(boxlength_f)); - Reset_List<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); VECTOR *frc = const_cast(reinterpret_cast(frc_f)); Dihedral14LJForceWithDirectCFKernel<<>>( diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_impl.cu index 47e3e454f26..2b710212fd4 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_impl.cu @@ -97,9 +97,6 @@ void MDIterationLeapFrog(const int float4_numbers, const int atom_numbers, const const float exp_gamma, const int is_max_velocity, const float max_velocity, const float *d_mass_inverse, const float *d_sqrt_mass, float *vel_f, float *crd_f, float *frc_f, float *acc_f, cudaStream_t stream) { - Reset_List<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, vel_f, 0.); - Reset_List<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, crd_f, 0.); - Reset_List<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); Reset_List<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, acc_f, 0.); VECTOR *frc = const_cast(reinterpret_cast(frc_f)); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_excluded_force_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_excluded_force_impl.cu index 7de68cad587..413d1f1998d 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_excluded_force_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_excluded_force_impl.cu @@ -86,7 +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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, frc_f, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, frc_f, 0.); UNSIGNED_INT_VECTOR *uint_crd = const_cast(reinterpret_cast(uint_crd_f)); VECTOR *frc = reinterpret_cast(frc_f); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu index 592d66d10f1..c7e089c2f0c 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_reciprocal_force_impl.cu @@ -75,7 +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<<(3. * atom_numbers) / 128), 128>>>(3 * atom_numbers, force, 0.); + Reset_List<<(3. * atom_numbers) / 128), 128, 0, stream>>>(3 * atom_numbers, force, 0.); UNSIGNED_INT_VECTOR *uint_crd = const_cast(reinterpret_cast(uint_crd_f)); UNSIGNED_INT_VECTOR *PME_uxyz = reinterpret_cast(pme_uxyz); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_cf_energy_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_cf_energy_kernel.h index 3e38cf13ab6..c30eff3d1e4 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_cf_energy_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_cf_energy_kernel.h @@ -64,7 +64,7 @@ class Dihedral14CFEnergyGpuKernel : public GpuKernel { const std::vector &GetOutputSizeList() const override { return output_size_list_; } const std::vector &GetWorkspaceSizeList() const override { return workspace_size_list_; } - bool Launch(const std::vector &inputs, const std::vector &, + bool Launch(const std::vector &inputs, const std::vector &workspace, const std::vector &outputs, void *stream_ptr) override { auto uint_crd_f = GetDeviceAddress(inputs, 0); auto LJtype = GetDeviceAddress(inputs, 1); @@ -74,9 +74,10 @@ class Dihedral14CFEnergyGpuKernel : public GpuKernel { auto b_14 = GetDeviceAddress(inputs, 5); auto cf_scale_factor = GetDeviceAddress(inputs, 6); auto ene = GetDeviceAddress(outputs, 0); + auto uint_crd_with_LJ = GetDeviceAddress(workspace, 0); - Dihedral14CFEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, boxlength_f, a_14, b_14, - cf_scale_factor, ene, reinterpret_cast(stream_ptr)); + Dihedral14CFEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, uint_crd_with_LJ, boxlength_f, + a_14, b_14, cf_scale_factor, ene, reinterpret_cast(stream_ptr)); return true; } @@ -90,7 +91,7 @@ class Dihedral14CFEnergyGpuKernel : public GpuKernel { input_size_list_.push_back(ele_a_14 * sizeof(T1)); input_size_list_.push_back(ele_b_14 * sizeof(T1)); input_size_list_.push_back(ele_cf_scale_factor * sizeof(T)); - + workspace_size_list_.push_back(atom_numbers * sizeof(UINT_VECTOR_LJ_TYPE)); output_size_list_.push_back(atom_numbers * sizeof(T)); } @@ -108,6 +109,13 @@ class Dihedral14CFEnergyGpuKernel : public GpuKernel { std::vector workspace_size_list_; int dihedral_14_numbers; int atom_numbers; + struct UINT_VECTOR_LJ_TYPE { + unsigned int uint_x; + unsigned int uint_y; + unsigned int uint_z; + int LJ_type; + float charge; + }; }; } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_kernel.h index cf3e3b83133..d3a5e1cd0c8 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_cf_force_with_atom_energy_kernel.h @@ -70,7 +70,7 @@ class Dihedral14LJCFForceWithAtomEnergyGpuKernel : public GpuKernel { const std::vector &GetOutputSizeList() const override { return output_size_list_; } const std::vector &GetWorkspaceSizeList() const override { return workspace_size_list_; } - bool Launch(const std::vector &inputs, const std::vector &, + bool Launch(const std::vector &inputs, const std::vector &workspace, const std::vector &outputs, void *stream_ptr) override { auto uint_crd_f = GetDeviceAddress(inputs, 0); auto LJtype = GetDeviceAddress(inputs, 1); @@ -85,9 +85,11 @@ class Dihedral14LJCFForceWithAtomEnergyGpuKernel : public GpuKernel { auto frc_f = GetDeviceAddress(outputs, 0); auto atom_energy = GetDeviceAddress(outputs, 1); - Dihedral14LJCFForceWithAtomEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, boxlength_f, a_14, - b_14, lj_scale_factor, cf_scale_factor, LJ_type_A, LJ_type_B, frc_f, atom_energy, - reinterpret_cast(stream_ptr)); + auto uint_crd_with_LJ = GetDeviceAddress(workspace, 0); + + Dihedral14LJCFForceWithAtomEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, uint_crd_with_LJ, + boxlength_f, a_14, b_14, lj_scale_factor, cf_scale_factor, LJ_type_A, LJ_type_B, + frc_f, atom_energy, reinterpret_cast(stream_ptr)); return true; } @@ -104,6 +106,7 @@ class Dihedral14LJCFForceWithAtomEnergyGpuKernel : public GpuKernel { input_size_list_.push_back(ele_cf_scale_factor * sizeof(T)); input_size_list_.push_back(ele_LJ_type_A * sizeof(T)); input_size_list_.push_back(ele_LJ_type_B * sizeof(T)); + workspace_size_list_.push_back(atom_numbers * sizeof(UINT_VECTOR_LJ_TYPE)); output_size_list_.push_back(3 * atom_numbers * sizeof(T)); output_size_list_.push_back(atom_numbers * sizeof(T)); @@ -126,6 +129,13 @@ class Dihedral14LJCFForceWithAtomEnergyGpuKernel : public GpuKernel { std::vector workspace_size_list_; int dihedral_14_numbers; int atom_numbers; + struct UINT_VECTOR_LJ_TYPE { + unsigned int uint_x; + unsigned int uint_y; + unsigned int uint_z; + int LJ_type; + float charge; + }; }; } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_energy_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_energy_kernel.h index 95f1ca5b98a..82f81355e5b 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_energy_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nb14/dihedral_14_lj_energy_kernel.h @@ -68,7 +68,7 @@ class Dihedral14LJEnergyGpuKernel : public GpuKernel { const std::vector &GetOutputSizeList() const override { return output_size_list_; } const std::vector &GetWorkspaceSizeList() const override { return workspace_size_list_; } - bool Launch(const std::vector &inputs, const std::vector &, + bool Launch(const std::vector &inputs, const std::vector &workspace, const std::vector &outputs, void *stream_ptr) override { auto uint_crd_f = GetDeviceAddress(inputs, 0); auto LJtype = GetDeviceAddress(inputs, 1); @@ -80,9 +80,11 @@ class Dihedral14LJEnergyGpuKernel : public GpuKernel { auto LJ_type_A = GetDeviceAddress(inputs, 7); auto LJ_type_B = GetDeviceAddress(inputs, 8); auto ene = GetDeviceAddress(outputs, 0); + auto uint_crd_with_LJ = GetDeviceAddress(workspace, 0); - Dihedral14LJEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, boxlength_f, a_14, b_14, - lj_scale_factor, LJ_type_A, LJ_type_B, ene, reinterpret_cast(stream_ptr)); + Dihedral14LJEnergy(dihedral_14_numbers, atom_numbers, uint_crd_f, LJtype, charge, uint_crd_with_LJ, boxlength_f, + a_14, b_14, lj_scale_factor, LJ_type_A, LJ_type_B, ene, + reinterpret_cast(stream_ptr)); return true; } @@ -98,6 +100,7 @@ class Dihedral14LJEnergyGpuKernel : public GpuKernel { input_size_list_.push_back(ele_lj_scale_factor * sizeof(T)); input_size_list_.push_back(ele_LJ_type_A * sizeof(T)); input_size_list_.push_back(ele_LJ_type_B * sizeof(T)); + workspace_size_list_.push_back(atom_numbers * sizeof(UINT_VECTOR_LJ_TYPE)); output_size_list_.push_back(atom_numbers * sizeof(T)); } @@ -118,6 +121,13 @@ class Dihedral14LJEnergyGpuKernel : public GpuKernel { std::vector workspace_size_list_; int dihedral_14_numbers; int atom_numbers; + struct UINT_VECTOR_LJ_TYPE { + unsigned int uint_x; + unsigned int uint_y; + unsigned int uint_z; + int LJ_type; + float charge; + }; }; } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.h index 37834b078ca..51b49c41e24 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_energy_kernel.h @@ -93,7 +93,7 @@ class PMEEnergyGpuKernel : public GpuKernel { input_size_list_.push_back(atom_numbers * sizeof(VECTOR)); input_size_list_.push_back(atom_numbers * sizeof(T1)); - input_size_list_.push_back(atom_numbers * sizeof(T1)); + input_size_list_.push_back(excluded_numbers * sizeof(T1)); input_size_list_.push_back(atom_numbers * sizeof(T1)); workspace_size_list_.push_back(atom_numbers * sizeof(UNSIGNED_INT_VECTOR)); @@ -118,6 +118,7 @@ class PMEEnergyGpuKernel : public GpuKernel { std::vector output_size_list_; std::vector workspace_size_list_; int atom_numbers; + int excluded_numbers = 2719; int max_nl_numbers = 800; int fftx; int ffty; diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_excluded_force_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_excluded_force_kernel.h index 4eca3b7ffd4..f8b4486210b 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_excluded_force_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/pme/pme_excluded_force_kernel.h @@ -65,7 +65,7 @@ class PMEExcludedForceGpuKernel : public GpuKernel { input_size_list_.push_back(atom_numbers * sizeof(VECTOR)); input_size_list_.push_back(atom_numbers * sizeof(T)); input_size_list_.push_back(atom_numbers * sizeof(T1)); - input_size_list_.push_back(atom_numbers * sizeof(T1)); + input_size_list_.push_back(excluded_numbers * sizeof(T1)); input_size_list_.push_back(atom_numbers * sizeof(T1)); output_size_list_.push_back(atom_numbers * 3 * sizeof(T)); @@ -77,6 +77,7 @@ class PMEExcludedForceGpuKernel : public GpuKernel { std::vector output_size_list_; std::vector workspace_size_list_; int atom_numbers; + int excluded_numbers = 2719; float beta; struct VECTOR { float x; diff --git a/mindspore/ops/operations/sponge_ops.py b/mindspore/ops/operations/sponge_ops.py index d32242dd5d1..9ce217ff4ef 100644 --- a/mindspore/ops/operations/sponge_ops.py +++ b/mindspore/ops/operations/sponge_ops.py @@ -29,6 +29,9 @@ class BondForce(PrimitiveWithInfer): .. math:: dr = (x_1-x_2, y_1-y_2, z_1-z_2) + + .. math:: + F = (F_x, F_y, F_z) = 2*k*(1 - r_0/|dr|)*dr Args: @@ -53,8 +56,8 @@ class BondForce(PrimitiveWithInfer): @prim_attr_register def __init__(self, bond_numbers, atom_numbers): - assert isinstance(bond_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -66,8 +69,8 @@ class BondForce(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -93,6 +96,9 @@ class BondEnergy(PrimitiveWithInfer): .. math:: dr = (x_1-x_2, y_1-y_2, z_1-z_2) + + .. math:: + E = k*(|dr| - r_0)^2 Args: @@ -117,8 +123,8 @@ class BondEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, bond_numbers, atom_numbers): - assert isinstance(bond_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -130,8 +136,8 @@ class BondEnergy(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -179,8 +185,8 @@ class BondAtomEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, bond_numbers, atom_numbers): - assert isinstance(bond_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -192,8 +198,8 @@ class BondAtomEnergy(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -240,8 +246,8 @@ class BondForceWithAtomEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, bond_numbers, atom_numbers): - assert isinstance(bond_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -253,8 +259,8 @@ class BondForceWithAtomEnergy(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -286,6 +292,9 @@ class BondForceWithAtomVirial(PrimitiveWithInfer): .. math:: dr = (x_1-x_2, y_1-y_2, z_1-z_2) + + .. math:: + virial = |dr|*(|dr| - r_0)*k Args: @@ -311,8 +320,8 @@ class BondForceWithAtomVirial(PrimitiveWithInfer): @prim_attr_register def __init__(self, bond_numbers, atom_numbers): - assert isinstance(bond_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('bond_numbers', bond_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.bond_numbers = bond_numbers self.atom_numbers = atom_numbers self.add_prim_attr('bond_numbers', self.bond_numbers) @@ -324,8 +333,8 @@ class BondForceWithAtomVirial(PrimitiveWithInfer): cls_name = self.name N = self.atom_numbers M = self.bond_numbers - validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f", cls_name) - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "uint_crd_f", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -355,25 +364,36 @@ class DihedralForce(PrimitiveWithInfer): .. math:: dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a) + .. math:: dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c) + .. math:: dr_{cd} = (x_d-x_c, y_d-y_c, z_d-z_c) - + .. math:: r1 = dr_{ab}*dr_{cb} + .. math:: r2 = dr_{cd}*dr_{cb} - + .. math:: phi = pi - sign(inner_product(r1*r2), dr_{cb}) * arccos(inner_product(r1, r2)/|r1|/|r2|) + .. math:: dEdphi = n*phi*(k*cos(phi_0)*sin(n*phi) - k*sin(phi_0)*cos(n*phi))/sin(phi) + .. math:: dphidr1 = r2/|r1|/|r2| + cos(phi)/|r1|^2*r1 + .. math:: dphidr2 = r1/|r1|/|r2| + cos(phi)/|r2|^2*r2 - + .. math:: dEdra = dEdphi * dr_{cb} * dphidr1 + .. math:: dEdrd = dEdphi * dphi_dr2 * dr_{cb} + .. math:: dEdrjpart = dEdphi * ((dr_{ab} * dphidr1) + (dr_{cd} * dphidr2)) - + .. math:: F_a = dEdri + .. math:: F_b = dEdrjpart - dEdri + .. math:: F_c = - dEdrl - dEdrjpart + .. math:: F_d = dEdrl Args: @@ -405,7 +425,7 @@ class DihedralForce(PrimitiveWithInfer): @prim_attr_register def __init__(self, dihedral_numbers): - assert isinstance(dihedral_numbers, int) + validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) self.dihedral_numbers = dihedral_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', 'gamc', 'gams', 'pn'], @@ -415,8 +435,8 @@ class DihedralForce(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape, ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): cls_name = self.name - M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + M = self.dihedral_numbers + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -485,7 +505,7 @@ class DihedralEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, dihedral_numbers): - assert isinstance(dihedral_numbers, int) + validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) self.dihedral_numbers = dihedral_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', 'gamc', 'gams', 'pn'], @@ -495,8 +515,8 @@ class DihedralEnergy(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape, ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): cls_name = self.name - M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + M = self.dihedral_numbers + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -563,7 +583,7 @@ class DihedralAtomEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, dihedral_numbers): - assert isinstance(dihedral_numbers, int) + validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) self.dihedral_numbers = dihedral_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', 'gamc', 'gams', 'pn'], @@ -574,8 +594,8 @@ class DihedralAtomEnergy(PrimitiveWithInfer): ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): cls_name = self.name N = uint_crd_f_shape[0] - M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + M = self.dihedral_numbers + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -641,7 +661,7 @@ class DihedralForceWithAtomEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, dihedral_numbers): - assert isinstance(dihedral_numbers, int) + validator.check_value_type('dihedral_numbers', dihedral_numbers, (int), self.name) self.dihedral_numbers = dihedral_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk', 'gamc', 'gams', 'pn'], @@ -652,8 +672,8 @@ class DihedralForceWithAtomEnergy(PrimitiveWithInfer): ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): cls_name = self.name N = uint_crd_f_shape[0] - M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + M = self.dihedral_numbers + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -690,14 +710,18 @@ class AngleForce(PrimitiveWithInfer): number of atoms is N. .. math:: - dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a) + .. math:: dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c) + .. math:: theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|) + .. math:: F_a = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{ab}|^2*dr_{ab} - 1/|dr_{ab}|/|dr_{cb}|*dr_{cb}] + .. math:: F_c = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{cb}|^2*dr_{cb} - - 1/|dr_{cb}|/|dr_{ab}|*dr_{ab}] + - 1/|dr_{cb}|/|dr_{ab}|*dr_{ab}] + .. math:: F_b = -F_a - F_c Args: @@ -722,7 +746,7 @@ class AngleForce(PrimitiveWithInfer): @prim_attr_register def __init__(self, angle_numbers): - assert isinstance(angle_numbers, int) + validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) self.angle_numbers = angle_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', 'angle_theta0'], @@ -732,8 +756,8 @@ class AngleForce(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape, angle_theta0_shape): cls_name = self.name - M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + M = self.angle_numbers + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -756,13 +780,16 @@ class AngleForce(PrimitiveWithInfer): class AngleEnergy(PrimitiveWithInfer): """ - Calculate the energy caused by 3-atoms angle term. + Calculate the energy caused by 3-atoms angle term. Assume the number of angles is M and the + number of atoms is N. .. math:: - dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a) + .. math:: dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c) + .. math:: theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|) + .. math:: E = k*(theta - theta_0)^2 Args: @@ -787,7 +814,7 @@ class AngleEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, angle_numbers): - assert isinstance(angle_numbers, int) + validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) self.angle_numbers = angle_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', 'angle_theta0'], @@ -797,8 +824,8 @@ class AngleEnergy(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape, angle_theta0_shape): cls_name = self.name - M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + M = self.angle_numbers + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -822,7 +849,8 @@ class AngleEnergy(PrimitiveWithInfer): class AngleAtomEnergy(PrimitiveWithInfer): """ Add the potential energy caused by angle terms to the total potential - energy of each atom. + energy of each atom. Assume the number of angles is M and the + number of atoms is N. The calculation formula is the same as operator AngleEnergy(). @@ -848,7 +876,7 @@ class AngleAtomEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, angle_numbers): - assert isinstance(angle_numbers, int) + validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) self.angle_numbers = angle_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', 'angle_theta0'], @@ -859,8 +887,8 @@ class AngleAtomEnergy(PrimitiveWithInfer): angle_theta0_shape): cls_name = self.name N = uint_crd_f_shape[0] - M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + M = self.angle_numbers + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -883,7 +911,8 @@ class AngleAtomEnergy(PrimitiveWithInfer): class AngleForceWithAtomEnergy(PrimitiveWithInfer): """ - Calculate angle force and potential energy together. + Calculate angle force and potential energy together. Assume the number of angles is M and the + number of atoms is N. The calculation formula is the same as operator AngleForce() and AngleEnergy(). @@ -910,7 +939,7 @@ class AngleForceWithAtomEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, angle_numbers): - assert isinstance(angle_numbers, int) + validator.check_value_type('angle_numbers', angle_numbers, (int), self.name) self.angle_numbers = angle_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k', 'angle_theta0'], @@ -921,8 +950,8 @@ class AngleForceWithAtomEnergy(PrimitiveWithInfer): angle_theta0_shape): cls_name = self.name N = uint_crd_f_shape[0] - M = atom_a_shape[0] - validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f", cls_name) + M = self.angle_numbers + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f", cls_name) validator.check_int(atom_a_shape[0], M, Rel.EQ, "atom_a", cls_name) validator.check_int(atom_b_shape[0], M, Rel.EQ, "atom_b", cls_name) @@ -954,8 +983,8 @@ class Dihedral14LJForce(PrimitiveWithInfer): for all kinds of atom pairs. .. math:: - dr = (x_a-x_b, y_a-y_b, z_a-z_b) + .. math:: F = k*(-12*A/|dr|^{14} + 6*B/|dr|^{8})*dr Args: @@ -972,7 +1001,9 @@ class Dihedral14LJForce(PrimitiveWithInfer): - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the Lennard-Jones part of force correction of each dihedral 1,4 term. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. + Q is the number of atom pair. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. + Q is the number of atom pair. Outputs: - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. @@ -983,8 +1014,8 @@ class Dihedral14LJForce(PrimitiveWithInfer): @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): - assert isinstance(nb14_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers self.init_prim_io_names( @@ -996,6 +1027,19 @@ class Dihedral14LJForce(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): + cls_name = self.name + N = self.atom_numbers + M = self.dihedral_14_numbers + Q = LJ_type_A_shape[0] + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) + validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) + validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name) + validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name) + validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name) return uint_crd_f_shape def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, @@ -1019,9 +1063,9 @@ class Dihedral14LJEnergy(PrimitiveWithInfer): Calculate the Lennard-Jones part of 1,4 dihedral energy correction for each necessary dihedral terms on the corresponding atoms. - .. math:, the : - + .. math:: dr = (x_a-x_b, y_a-y_b, z_a-z-b) + .. math:: E = k*(A/|dr|^{12} - B/|dr|^{6}) Args: @@ -1038,7 +1082,9 @@ class Dihedral14LJEnergy(PrimitiveWithInfer): - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the Lennard-Jones part of force correction of each dihedral 1,4 term. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. + Q is the number of atom pair. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. + Q is the number of atom pair. Outputs: - **ene** (Tensor, float32) - [M,], the Lennard-Jones potential @@ -1050,8 +1096,8 @@ class Dihedral14LJEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): - assert isinstance(nb14_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1064,6 +1110,19 @@ class Dihedral14LJEnergy(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): + cls_name = self.name + N = self.atom_numbers + M = self.dihedral_14_numbers + Q = LJ_type_A_shape[0] + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) + validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) + validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name) + validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name) + validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name) return [self.dihedral_14_numbers,] def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, @@ -1090,9 +1149,9 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer): Dihedral14LJForce(), and the Coulomb part is as follows: .. math:: - - dr = (x_a-x_b, y_a-y_b, z_a-z_b) - F = -k*q_a*q_b/|r|^3*dr + dr = (x_a-x_b, y_a-y_b, z_a-z_b) + .. math:: + F = -k*q_a*q_b/|r|^3*dr Args: nb14_numbers (int32): the number of necessary dihedral 1,4 terms M. @@ -1110,7 +1169,9 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer): - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the Coulomb part of force correction for each dihedral 1,4 terms. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. + Q is the number of atom pair. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. + Q is the number of atom pair. Outputs: - **frc_f** (Tensor, float) - [N, 3], the force felt by each atom. @@ -1121,8 +1182,8 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer): @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): - assert isinstance(nb14_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1136,6 +1197,20 @@ class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, lj_scale_factor_shape, cf_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): + cls_name = self.name + N = self.atom_numbers + M = self.dihedral_14_numbers + Q = LJ_type_A_shape[0] + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) + validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) + validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name) + validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name) + validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name) + validator.check_int(cf_scale_factor_shape[0], M, Rel.EQ, "cf_scale_factor", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name) return [self.atom_numbers, 3] def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, @@ -1180,7 +1255,9 @@ class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer): - **cf_scale_factor** (Tensor, float) - [M,], the scale factor for the Coulomb part of force correction for each dihedral 1,4 terms. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. + Q is the number of atom pair. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. + Q is the number of atom pair. Outputs: - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. @@ -1192,8 +1269,8 @@ class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): - assert isinstance(nb14_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1207,6 +1284,20 @@ class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, lj_scale_factor_shape, cf_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): + cls_name = self.name + N = self.atom_numbers + M = self.dihedral_14_numbers + Q = LJ_type_A_shape[0] + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) + validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) + validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name) + validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name) + validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name) + validator.check_int(cf_scale_factor_shape[0], M, Rel.EQ, "cf_scale_factor", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name) return uint_crd_f_shape, charge_shape def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, @@ -1246,7 +1337,9 @@ class Dihedral14LJAtomEnergy(PrimitiveWithInfer): - **lj_scale_factor** (Tensor, float32) - [M,], the scale factor for the Lennard-Jones part of force correction of each dihedral 1,4 term. - **LJ_type_A** (Tensor, float32) - [Q,], the A parameter in Lennard-Jones scheme of each atom pair type. + Q is the number of atom pair. - **LJ_type_B** (Tensor, float32) - [Q,], the B parameter in Lennard-Jones shceme of each atom pair type. + Q is the number of atom pair. Outputs: - **ene** (Tensor, float32) - [N,], the accumulated potential energy of each atom. @@ -1257,8 +1350,8 @@ class Dihedral14LJAtomEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): - assert isinstance(nb14_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1271,6 +1364,19 @@ class Dihedral14LJAtomEnergy(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, lj_scale_factor_shape, LJ_type_A_shape, LJ_type_B_shape): + cls_name = self.name + N = self.atom_numbers + M = self.dihedral_14_numbers + Q = LJ_type_A_shape[0] + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) + validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) + validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name) + validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name) + validator.check_int(lj_scale_factor_shape[0], M, Rel.EQ, "lj_scale_factor", cls_name) + validator.check_int(LJ_type_B_shape[0], Q, Rel.EQ, "LJ_type_B", cls_name) return LJtype_shape def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, @@ -1297,6 +1403,8 @@ class Dihedral14CFEnergy(PrimitiveWithInfer): .. math:: dr = (x_a-x_b, y_a-y_b, z_a-z_b) + + .. math:: E = k*q_a*q_b/|dr| Args: @@ -1322,8 +1430,8 @@ class Dihedral14CFEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): - assert isinstance(nb14_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1335,6 +1443,17 @@ class Dihedral14CFEnergy(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, cf_scale_factor_shape): + cls_name = self.name + N = self.atom_numbers + M = self.dihedral_14_numbers + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) + validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) + validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name) + validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name) + validator.check_int(cf_scale_factor_shape[0], M, Rel.EQ, "cf_scale_factor", cls_name) return [self.dihedral_14_numbers,] def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, @@ -1382,8 +1501,8 @@ class Dihedral14CFAtomEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, nb14_numbers, atom_numbers): - assert isinstance(nb14_numbers, int) - assert isinstance(atom_numbers, int) + validator.check_value_type('nb14_numbers', nb14_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) self.dihedral_14_numbers = nb14_numbers self.atom_numbers = atom_numbers @@ -1395,6 +1514,17 @@ class Dihedral14CFAtomEnergy(PrimitiveWithInfer): def infer_shape(self, uint_crd_f_shape, LJtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape, cf_scale_factor_shape): + cls_name = self.name + N = self.atom_numbers + M = self.dihedral_14_numbers + validator.check_int(uint_crd_f_shape[0], N, Rel.EQ, "uint_crd_f[0]", cls_name) + validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name) + validator.check_int(LJtype_shape[0], 3, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge_shape[0], M, Rel.EQ, "charge", cls_name) + validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name) + validator.check_int(a_14_shape[0], M, Rel.EQ, "a_14", cls_name) + validator.check_int(b_14_shape[0], M, Rel.EQ, "b_14", cls_name) + validator.check_int(cf_scale_factor_shape[0], M, Rel.EQ, "cf_scale_factor", cls_name) return LJtype_shape def infer_dtype(self, uint_crd_f_dtype, LJtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type, @@ -1450,13 +1580,13 @@ class MDIterationLeapFrog(PrimitiveWithInfer): @prim_attr_register def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): - assert isinstance(float4_numbers, int) - assert isinstance(atom_numbers, int) - assert isinstance(half_dt, float) - assert isinstance(dt, float) - assert isinstance(exp_gamma, float) - assert isinstance(is_max_velocity, int) - assert isinstance(max_velocity, float) + validator.check_value_type('float4_numbers', float4_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('half_dt', half_dt, (float), self.name) + validator.check_value_type('dt', dt, (float), self.name) + validator.check_value_type('exp_gamma', exp_gamma, (float), self.name) + validator.check_value_type('is_max_velocity', is_max_velocity, (int), self.name) + validator.check_value_type('max_velocity', max_velocity, (float), self.name) self.float4_numbers = float4_numbers self.atom_numbers = atom_numbers self.half_dt = half_dt @@ -1477,6 +1607,10 @@ class MDIterationLeapFrog(PrimitiveWithInfer): self.add_prim_attr('max_velocity', self.max_velocity) def infer_shape(self, mass_inverse_shape, sqrt_mass_shape): + cls_name = self.name + N = self.atom_numbers + validator.check_int(mass_inverse_shape[0], N, Rel.EQ, "mass_inverse", cls_name) + validator.check_int(sqrt_mass_shape[0], N, Rel.EQ, "sqrt_mass", cls_name) return [self.atom_numbers, 3], [self.atom_numbers, 3], [self.atom_numbers, 3], [self.atom_numbers, 3] def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype): @@ -1517,11 +1651,11 @@ class PMEReciprocalForce(PrimitiveWithInfer): @prim_attr_register def __init__(self, atom_numbers, beta, fftx, ffty, fftz): - assert isinstance(atom_numbers, int) - assert isinstance(beta, float) - assert isinstance(fftx, int) - assert isinstance(ffty, int) - assert isinstance(fftz, int) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('beta', beta, (float), self.name) + validator.check_value_type('fftx', fftx, (int), self.name) + validator.check_value_type('ffty', ffty, (int), self.name) + validator.check_value_type('fftz', fftz, (int), self.name) self.atom_numbers = atom_numbers self.beta = beta self.fftx = fftx @@ -1536,6 +1670,12 @@ class PMEReciprocalForce(PrimitiveWithInfer): self.add_prim_attr('fftz', self.fftz) def infer_shape(self, boxlength_shape, uint_crd_shape, charge_shape): + cls_name = self.name + N = self.atom_numbers + validator.check_int(uint_crd_shape[0], N, Rel.EQ, "uint_crd[0]", cls_name) + validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd[1]", cls_name) + validator.check_int(boxlength_shape[0], 3, Rel.EQ, "boxlength", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge", cls_name) return uint_crd_shape def infer_dtype(self, boxlength_type, uint_crd_type, charge_type): @@ -1564,7 +1704,7 @@ class PMEExcludedForce(PrimitiveWithInfer): - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index in excluded list for each atom. - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded - list of each atom. + list of each atom. E is the number of excluded atoms. - **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded in excluded list for each atom. @@ -1577,8 +1717,8 @@ class PMEExcludedForce(PrimitiveWithInfer): @prim_attr_register def __init__(self, atom_numbers, beta): - assert isinstance(atom_numbers, int) - assert isinstance(beta, float) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('beta', beta, (float), self.name) self.atom_numbers = atom_numbers self.beta = beta self.init_prim_io_names( @@ -1589,6 +1729,14 @@ class PMEExcludedForce(PrimitiveWithInfer): def infer_shape(self, uint_crd_shape, sacler_shape, charge_shape, excluded_list_start_shape, excluded_list_shape, excluded_atom_numbers_shape): + cls_name = self.name + N = self.atom_numbers + validator.check_int(uint_crd_shape[0], N, Rel.EQ, "uint_crd[0]", cls_name) + validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd[1]", cls_name) + validator.check_int(sacler_shape[0], 3, Rel.EQ, "sacler", cls_name) + validator.check_int(charge_shape[0], N, Rel.EQ, "charge", cls_name) + validator.check_int(excluded_list_start_shape[0], N, Rel.EQ, "excluded_list_start", cls_name) + validator.check_int(excluded_atom_numbers_shape[0], N, Rel.EQ, "excluded_atom_numbers", cls_name) return uint_crd_shape def infer_dtype(self, uint_crd_type, sacler_type, charge_type, excluded_list_start_type, excluded_list_type, @@ -1632,7 +1780,7 @@ class PMEEnergy(PrimitiveWithInfer): - **excluded_list_start** (Tensor, int32) - [N,], the start excluded index in excluded list for each atom. - **excluded_list** (Tensor, int32) - [E,], the contiguous join of excluded - list of each atom. + list of each atom. E is the number of excluded atoms. - **excluded_atom_numbers** (Tensor, int32) - [N,], the number of atom excluded in excluded list for each atom. @@ -1648,11 +1796,11 @@ class PMEEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, atom_numbers, beta, fftx, ffty, fftz): - assert isinstance(atom_numbers, int) - assert isinstance(beta, float) - assert isinstance(fftx, int) - assert isinstance(ffty, int) - assert isinstance(fftz, int) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('beta', beta, (float), self.name) + validator.check_value_type('fftx', fftx, (int), self.name) + validator.check_value_type('ffty', ffty, (int), self.name) + validator.check_value_type('fftz', fftz, (int), self.name) self.atom_numbers = atom_numbers self.beta = beta self.fftx = fftx @@ -1670,6 +1818,18 @@ class PMEEnergy(PrimitiveWithInfer): def infer_shape(self, box_length, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start, excluded_list, excluded_atom_numbers): + cls_name = self.name + N = self.atom_numbers + validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd[0]", cls_name) + validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd[1]", cls_name) + validator.check_int(box_length[0], 3, Rel.EQ, "box_length", cls_name) + validator.check_int(charge[0], N, Rel.EQ, "charge", cls_name) + validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers[0]", cls_name) + validator.check_int(nl_serial[0], N, Rel.LE, "nl_serial[0]", cls_name) + validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial[1]", cls_name) + validator.check_int(excluded_list_start[0], N, Rel.EQ, "excluded_list_start", cls_name) + validator.check_int(excluded_atom_numbers[0], N, Rel.EQ, "excluded_atom_numbers", cls_name) + validator.check_int(excluded_list[0], 0, Rel.GE, "excluded_list", cls_name) return (1,), (1,), (1,), (1,) def infer_dtype(self, box_length, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start, @@ -1701,6 +1861,8 @@ class LJEnergy(PrimitiveWithInfer): .. math:: dr = (x_a-x_b, y_a-y_b, z_a-z_b) + + .. math:: E = A/|dr|^{12} - B/|dr|^{6} Agrs: @@ -1716,7 +1878,9 @@ class LJEnergy(PrimitiveWithInfer): - **nl_numbers** - (Tensor, int32) - [N,], the each atom. - **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800. - **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair. + Q is the number of atom pair. - **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair. + Q is the number of atom pair. Outputs: - **d_LJ_energy_atom** (Tensor, float32) - [N,], the Lennard-Jones potential energy of each atom. @@ -1728,8 +1892,8 @@ class LJEnergy(PrimitiveWithInfer): @prim_attr_register def __init__(self, atom_numbers, cutoff_square): - assert isinstance(atom_numbers, int) - assert isinstance(cutoff_square, float) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('cutoff_square', cutoff_square, (float), self.name) self.atom_numbers = atom_numbers self.cutoff_square = cutoff_square self.init_prim_io_names( @@ -1739,6 +1903,19 @@ class LJEnergy(PrimitiveWithInfer): self.add_prim_attr('cutoff_square', self.cutoff_square) def infer_shape(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): + cls_name = self.name + N = self.atom_numbers + Q = d_LJ_A[0] + validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd[0]", cls_name) + validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd[1]", cls_name) + validator.check_int(LJtype[0], N, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge[0], 3, Rel.EQ, "charge", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name) + validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers", cls_name) + validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial[0]", cls_name) + validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial[1]", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name) + validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B[0]", cls_name) return charge def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): @@ -1761,6 +1938,9 @@ class LJForce(PrimitiveWithInfer): .. math:: dr = (x_a-x_b, y_a-y_b, z_a-z_b) + + .. math:: + F = (-12*A/|dr|^{14} + 6*B/|dr|^{8}) * dr Agrs: @@ -1776,7 +1956,9 @@ class LJForce(PrimitiveWithInfer): - **nl_numbers** - (Tensor, int32) - [N,], the each atom. - **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800. - **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair. + Q is the number of atom pair. - **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair. + Q is the number of atom pair. outputs: - **frc** (Tensor, float32) - [N, 3], the force felt by each atom. @@ -1787,8 +1969,8 @@ class LJForce(PrimitiveWithInfer): @prim_attr_register def __init__(self, atom_numbers, cutoff_square): - assert isinstance(atom_numbers, int) - assert isinstance(cutoff_square, float) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('cutoff_square', cutoff_square, (float), self.name) self.atom_numbers = atom_numbers self.cutoff_square = cutoff_square self.init_prim_io_names( @@ -1798,6 +1980,19 @@ class LJForce(PrimitiveWithInfer): self.add_prim_attr('cutoff_square', self.cutoff_square) def infer_shape(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): + cls_name = self.name + N = self.atom_numbers + Q = d_LJ_A[0] + validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd[0]", cls_name) + validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd[1]", cls_name) + validator.check_int(LJtype[0], N, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge[0], 3, Rel.EQ, "charge", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name) + validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers", cls_name) + validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial[0]", cls_name) + validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial[1]", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name) + validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B[0]", cls_name) return uint_crd def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): @@ -1833,7 +2028,9 @@ class LJForceWithPMEDirectForce(PrimitiveWithInfer): - **nl_numbers** - (Tensor, int32) - [N,], the each atom. - **nl_serial** - (Tensor, int32) - [N, 800], the neighbor list of each atom, the max number is 800. - **d_LJ_A** (Tensor, float32) - [Q,], the Lennard-Jones A coefficient of each kind of atom pair. + Q is the number of atom pair. - **d_LJ_B** (Tensor, float32) - [Q,], the Lennard-Jones B coefficient of each kind of atom pair. + Q is the number of atom pair. Outputs: - **frc** (Tensor, float32), [N, 3], the force felt by each atom. @@ -1844,9 +2041,9 @@ class LJForceWithPMEDirectForce(PrimitiveWithInfer): @prim_attr_register def __init__(self, atom_numbers, cutoff, pme_beta): - assert isinstance(atom_numbers, int) - assert isinstance(cutoff, float) - assert isinstance(pme_beta, float) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('cutoff', cutoff, (float), self.name) + validator.check_value_type('pme_beta', pme_beta, (float), self.name) self.atom_numbers = atom_numbers self.cutoff = cutoff self.pme_beta = pme_beta @@ -1858,6 +2055,19 @@ class LJForceWithPMEDirectForce(PrimitiveWithInfer): self.add_prim_attr('pme_beta', self.pme_beta) def infer_shape(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): + cls_name = self.name + N = self.atom_numbers + Q = d_LJ_A[0] + validator.check_int(uint_crd[0], N, Rel.EQ, "uint_crd[0]", cls_name) + validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd[1]", cls_name) + validator.check_int(LJtype[0], N, Rel.EQ, "LJtype", cls_name) + validator.check_int(charge[0], 3, Rel.EQ, "charge", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name) + validator.check_int(nl_numbers[0], N, Rel.EQ, "nl_numbers", cls_name) + validator.check_int(nl_serial[0], N, Rel.EQ, "nl_serial[0]", cls_name) + validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial[1]", cls_name) + validator.check_int(scaler[0], 3, Rel.EQ, "scaler", cls_name) + validator.check_int(d_LJ_B[0], Q, Rel.EQ, "d_LJ_B[0]", cls_name) return uint_crd def infer_dtype(self, uint_crd, LJtype, charge, scaler, nl_numbers, nl_serial, d_LJ_A, d_LJ_B): @@ -1882,6 +2092,8 @@ class GetCenterOfGeometry(PrimitiveWithInfer): @prim_attr_register def __init__(self, center_numbers, center_numbers_inverse): + validator.check_value_type('center_numbers', center_numbers, (int), self.name) + validator.check_value_type('center_numbers_inverse', center_numbers_inverse, (float), self.name) self.center_numbers = center_numbers self.center_numbers_inverse = center_numbers_inverse self.add_prim_attr('center_numbers', self.center_numbers) @@ -1894,8 +2106,8 @@ class GetCenterOfGeometry(PrimitiveWithInfer): cls_name = self.name N = self.center_numbers validator.check_int(center_atoms_shape[0], N, Rel.EQ, "center_atoms", cls_name) - validator.check_int(crd_f_shape[0], N, Rel.EQ, "crd_f", cls_name) - validator.check_int(crd_f_shape[1], 3, Rel.EQ, "crd_f", cls_name) + validator.check_int(crd_f_shape[0], N, Rel.EQ, "crd_f[0]", cls_name) + validator.check_int(crd_f_shape[1], 3, Rel.EQ, "crd_f[1]", cls_name) return [3,] def infer_dtype(self, center_atoms_dtype, crd_f_dtype): @@ -1990,6 +2202,7 @@ class NeighborListUpdate(PrimitiveWithInfer): between the unsigned int value and the real space coordinates. - **uint_crd** (Tensor, uint32) - [N, 3], the unsigned int coordinates value fo each atom. - **gpointer** (Tensor, int32) - [G, 125], the 125 nearest neighbor grids (including self) of each grid. + G is the number of nearest neighbor grids. - **nl_atom_numbers** (Tensor, int32) - [N,], the number of atoms in neighbor list of each atom. - **nl_atom_serial** (Tensor, int32) - [N, L], the indices of atoms in neighbor list of each atom. - **uint_dr_to_dr_cof** (Tensor, float32) - [3,], the scale factor between @@ -2182,13 +2395,13 @@ class MDIterationLeapFrogWithRF(PrimitiveWithInfer): @prim_attr_register def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity): - assert isinstance(float4_numbers, int) - assert isinstance(atom_numbers, int) - assert isinstance(half_dt, float) - assert isinstance(dt, float) - assert isinstance(exp_gamma, float) - assert isinstance(is_max_velocity, int) - assert isinstance(max_velocity, float) + validator.check_value_type('float4_numbers', float4_numbers, (int), self.name) + validator.check_value_type('atom_numbers', atom_numbers, (int), self.name) + validator.check_value_type('half_dt', half_dt, (float), self.name) + validator.check_value_type('dt', dt, (float), self.name) + validator.check_value_type('exp_gamma', exp_gamma, (float), self.name) + validator.check_value_type('is_max_velocity', is_max_velocity, (int), self.name) + validator.check_value_type('max_velocity', max_velocity, (float), self.name) self.float4_numbers = float4_numbers self.atom_numbers = atom_numbers self.half_dt = half_dt diff --git a/model_zoo/research/hpc/sponge/src/__init__.py b/model_zoo/research/hpc/sponge/src/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/model_zoo/research/hpc/sponge/src/nb14.py b/model_zoo/research/hpc/sponge/src/nb14.py index 01f8f9213ec..fce4edd10a6 100644 --- a/model_zoo/research/hpc/sponge/src/nb14.py +++ b/model_zoo/research/hpc/sponge/src/nb14.py @@ -82,7 +82,7 @@ class NON_BOND_14(nn.Cell): self.cf_scale_type = [0] * self.dihedral_type_numbers self.lj_scale_type = [0] * self.dihedral_type_numbers - + self.process1(context) self.h_atom_a = [0] * self.dihedral_numbers self.h_atom_b = [0] * self.dihedral_numbers self.h_lj_scale_factor = [0] * self.dihedral_numbers diff --git a/model_zoo/research/hpc/sponge/src/simulation_initial.py b/model_zoo/research/hpc/sponge/src/simulation_initial.py index 5eed6cea9b8..f8b3846f2b7 100644 --- a/model_zoo/research/hpc/sponge/src/simulation_initial.py +++ b/model_zoo/research/hpc/sponge/src/simulation_initial.py @@ -18,15 +18,15 @@ import numpy as np import mindspore.common.dtype as mstype from mindspore import Tensor, nn -from Langevin_Liujian_md import Langevin_Liujian -from angle import Angle -from bond import Bond -from dihedral import Dihedral -from lennard_jones import Lennard_Jones_Information -from md_information import md_information -from nb14 import NON_BOND_14 -from neighbor_list import nb_infomation -from particle_mesh_ewald import Particle_Mesh_Ewald +from .Langevin_Liujian_md import Langevin_Liujian +from .angle import Angle +from .bond import Bond +from .dihedral import Dihedral +from .lennard_jones import Lennard_Jones_Information +from .md_information import md_information +from .nb14 import NON_BOND_14 +from .neighbor_list import nb_infomation +from .particle_mesh_ewald import Particle_Mesh_Ewald class controller: @@ -102,7 +102,8 @@ class Simulation(nn.Cell): nb_info = self.nb_info pme_method = self.pme_method bond_frc, _ = self.bond.Bond_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof) - frc_t = bond_frc.asnumpy() + frc_t = 0 + frc_t += bond_frc.asnumpy() angle_frc, _ = self.angle.Angle_Force_With_Atom_Energy(md_info.uint_crd, md_info.uint_dr_to_dr_cof) frc_t += angle_frc.asnumpy() @@ -180,6 +181,7 @@ class Simulation(nn.Cell): if md_info.mode > 0 and int(control.Command_Set["thermostat"]) == 1: md_info.vel, md_info.crd, md_info.frc, md_info.acc = liujian_info.MD_Iteration_Leap_Frog( md_info.d_mass_inverse, md_info.vel, md_info.crd, md_info.frc) + self.Main_After_Iteration() def Main_After_Iteration(self): """main after iteration"""