diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common_sponge.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common_sponge.cuh index 57f2a4b006d..bcedd2183dd 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common_sponge.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/common_sponge.cuh @@ -359,4 +359,8 @@ __global__ static void Print(const size_t size, const int *input_x) { return; } +__device__ static VECTOR Make_Vector_Not_Exceed_Value(VECTOR vector, const float value) { + return fminf(1.0, value * rnorm3df(vector.x, vector.y, vector.z)) * vector; +} + #endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_COMMON_SPONGE_H_ diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_gradient_descent_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_gradient_descent_impl.cu new file mode 100644 index 00000000000..f3aad0e1ae9 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_gradient_descent_impl.cu @@ -0,0 +1,41 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_gradient_descent_impl.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/util.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/common_sponge.cuh" + +__global__ void MD_Iteration_Gradient_Descent(const int atom_numbers, VECTOR *crd, VECTOR *frc, + const float learning_rate) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < atom_numbers) { + crd[i].x = crd[i].x + learning_rate * frc[i].x; + crd[i].y = crd[i].y + learning_rate * frc[i].y; + crd[i].z = crd[i].z + learning_rate * frc[i].z; + + frc[i].x = 0.; + frc[i].y = 0.; + frc[i].z = 0.; + } +} + +void MDIterationGradientDescent(const int atom_numbers, float *crd, float *frc, const float learning_rate, + cudaStream_t stream) { + VECTOR *d_crd = reinterpret_cast(crd); + VECTOR *d_frc = reinterpret_cast(frc); + MD_Iteration_Gradient_Descent<<(atom_numbers) / 128), 128, 0, stream>>>( + atom_numbers, d_crd, d_frc, learning_rate); +} diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_gradient_descent_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_gradient_descent_impl.cuh new file mode 100644 index 00000000000..a2d9b7eef87 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_gradient_descent_impl.cuh @@ -0,0 +1,25 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_GRADIENT_DESCENT_IMPL_H_ +#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_GRADIENT_DESCENT_IMPL_H_ + +#include +#include "runtime/device/gpu/cuda_common.h" +void MDIterationGradientDescent(const int atom_numbers, float *crd, float *frc, const float learning_rate, + cudaStream_t stream); + +#endif diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cu index 4c9bda4b3f2..d0b4e6653e0 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_liujian_gpu_impl.cu @@ -54,7 +54,7 @@ void MD_Iteration_Leap_Frog_With_LiuJian(const int atom_numbers, const float hal curandStatePhilox4_32_10_t *rand_state, float *rand_frc, float *output, cudaStream_t stream) { Rand_Normal<<(float4_numbers) / 32.), 32, 0, stream>>>(float4_numbers, rand_state, - reinterpret_cast(rand_frc)); + reinterpret_cast(rand_frc)); VECTOR *d_vel = reinterpret_cast(vel); VECTOR *d_crd = reinterpret_cast(crd); VECTOR *d_frc = reinterpret_cast(frc); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_with_max_vel_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_with_max_vel_impl.cu new file mode 100644 index 00000000000..fabb19f7550 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_with_max_vel_impl.cu @@ -0,0 +1,44 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_with_max_vel_impl.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/util.cuh" +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/common_sponge.cuh" + +__global__ void MD_Iteration_Leap_Frog_With_Max_Velocity(const int atom_numbers, VECTOR *vel, VECTOR *crd, VECTOR *frc, + VECTOR *acc, const float *inverse_mass, const float dt, + const float max_velocity) { + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < atom_numbers) { + VECTOR acc_i = inverse_mass[i] * frc[i]; + VECTOR vel_i = vel[i] + dt * acc_i; + vel_i = Make_Vector_Not_Exceed_Value(vel_i, max_velocity); + vel[i] = vel_i; + crd[i] = crd[i] + dt * vel_i; + frc[i] = {0.0f, 0.0f, 0.0f}; + } +} + +void MDIterationLeapFrogWithMaxVelocity(const int atom_numbers, float *vel, float *crd, float *frc, float *acc, + const float *inverse_mass, const float dt, const float max_velocity, + cudaStream_t stream) { + VECTOR *d_vel = reinterpret_cast(vel); + VECTOR *d_crd = reinterpret_cast(crd); + VECTOR *d_frc = reinterpret_cast(frc); + VECTOR *d_acc = reinterpret_cast(acc); + MD_Iteration_Leap_Frog_With_Max_Velocity<<(atom_numbers) / 128), 128, 0, stream>>>( + atom_numbers, d_vel, d_crd, d_frc, d_acc, inverse_mass, dt, max_velocity); +} diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_with_max_vel_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_with_max_vel_impl.cuh new file mode 100644 index 00000000000..477647c01eb --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_with_max_vel_impl.cuh @@ -0,0 +1,26 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_LEAP_FROG_WITH_MAX_VEL_IMPL_H_ +#define MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_LEAP_FROG_WITH_MAX_VEL_IMPL_H_ + +#include +#include "runtime/device/gpu/cuda_common.h" +void MDIterationLeapFrogWithMaxVelocity(const int atom_numbers, float *vel, float *crd, float *frc, float *acc, + const float *inverse_mass, const float dt, const float max_velocity, + cudaStream_t stream); + +#endif // MINDSPORE_CCSRC_KERNEL_GPU_CUDA_IMPL_SPONGE_MD_ITERATION_LEAP_FROG_WITH_MAX_VEL_IMPL_H_ diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cu index 1dbd34e39b2..5a2d487ca50 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_setup_random_state_gpu_impl.cu @@ -21,7 +21,7 @@ void MD_Iteration_Setup_Random_State(int float4_numbers, curandStatePhilox4_32_10_t *rand_state, int seed, cudaStream_t stream) { Setup_Rand_Normal_Kernel<<(float4_numbers) / 32.), 32, 0, stream>>>(float4_numbers, - rand_state, seed); + rand_state, seed); } void MD_Iteration_Setup_Random_State(int float4_numbers, curandStatePhilox4_32_10_t *rand_state, int seed, diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_common.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_common.cuh index 427e63e73e0..c8ef5052a9f 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_common.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/sponge/pme/pme_common.cuh @@ -108,16 +108,16 @@ __global__ static void PME_Atom_Near(const UNSIGNED_INT_VECTOR *uint_crd, int *P UNSIGNED_INT_VECTOR *temp_uxyz = &PME_uxyz[atom]; int k, tempux, tempuy, tempuz; float tempf; - tempf = static_cast (uint_crd[atom].uint_x) * periodic_factor_inverse_x; - tempux = static_cast (tempf); + tempf = static_cast(uint_crd[atom].uint_x) * periodic_factor_inverse_x; + tempux = static_cast(tempf); PME_frxyz[atom].x = tempf - tempux; - tempf = static_cast (uint_crd[atom].uint_y) * periodic_factor_inverse_y; - tempuy = static_cast (tempf); + tempf = static_cast(uint_crd[atom].uint_y) * periodic_factor_inverse_y; + tempuy = static_cast(tempf); PME_frxyz[atom].y = tempf - tempuy; - tempf = static_cast (uint_crd[atom].uint_z) * periodic_factor_inverse_z; - tempuz = static_cast (tempf); + tempf = static_cast(uint_crd[atom].uint_z) * periodic_factor_inverse_z; + tempuz = static_cast(tempf); PME_frxyz[atom].z = tempf - tempuz; if (tempux != (*temp_uxyz).uint_x || tempuy != (*temp_uxyz).uint_y || tempuz != (*temp_uxyz).uint_z) { diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_gradient_descent_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_gradient_descent_kernel.cc new file mode 100644 index 00000000000..d858839c7a9 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_gradient_descent_kernel.cc @@ -0,0 +1,26 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_gradient_descent_kernel.h" + +namespace mindspore { +namespace kernel { +MS_REG_GPU_KERNEL_ONE( + MDIterationGradientDescent, + KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), + MDIterationGradientDescentGpuKernel, float) +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_gradient_descent_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_gradient_descent_kernel.h new file mode 100644 index 00000000000..4a5c77ce2bb --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_gradient_descent_kernel.h @@ -0,0 +1,77 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_MD_ITERATION_GRADIENT_DESCENT_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_MD_ITERATION_GRADIENT_DESCENT_KERNEL_H_ + +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_gradient_descent_impl.cuh" +#include +#include +#include +#include + +#include "backend/kernel_compiler/gpu/gpu_kernel.h" +#include "backend/kernel_compiler/gpu/gpu_kernel_factory.h" +#include "runtime/device/gpu/cuda_common.h" + +namespace mindspore { +namespace kernel { +template +class MDIterationGradientDescentGpuKernel : public GpuKernel { + public: + MDIterationGradientDescentGpuKernel() {} + ~MDIterationGradientDescentGpuKernel() override = default; + + bool Init(const CNodePtr &kernel_node) override { + // get bond_numbers + kernel_node_ = kernel_node; + atom_numbers = static_cast(GetAttr(kernel_node, "atom_numbers")); + learning_rate = static_cast(GetAttr(kernel_node, "learning_rate")); + InitSizeLists(); + return true; + } + + const std::vector &GetInputSizeList() const override { return input_size_list_; } + 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 &workspace, + const std::vector &outputs, void *stream_ptr) override { + auto crd = GetDeviceAddress(inputs, 0); + auto frc = GetDeviceAddress(inputs, 1); + + MDIterationGradientDescent(atom_numbers, crd, frc, learning_rate, reinterpret_cast(stream_ptr)); + return true; + } + + protected: + void InitSizeLists() override { + input_size_list_.push_back(atom_numbers * 3 * sizeof(T)); + input_size_list_.push_back(atom_numbers * 3 * sizeof(T)); + + output_size_list_.push_back(sizeof(T)); + } + + private: + std::vector input_size_list_; + std::vector output_size_list_; + std::vector workspace_size_list_; + int atom_numbers; + float learning_rate; +}; +} // namespace kernel +} // namespace mindspore +#endif diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_with_max_vel_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_with_max_vel_kernel.cc new file mode 100644 index 00000000000..d05702b8e46 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_with_max_vel_kernel.cc @@ -0,0 +1,31 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_with_max_vel_kernel.h" + +namespace mindspore { +namespace kernel { +MS_REG_GPU_KERNEL_ONE(MDIterationLeapFrogWithMaxVel, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32), + MDIterationLeapFrogWithMaxVelGpuKernel, float) +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_with_max_vel_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_with_max_vel_kernel.h new file mode 100644 index 00000000000..d7e1e1c5efa --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/sponge/nvtit/md_iteration_leap_frog_with_max_vel_kernel.h @@ -0,0 +1,86 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_MD_ITERATION_LEAP_FROG_WITH_MAX_VEL_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_MD_ITERATION_LEAP_FROG_WITH_MAX_VEL_KERNEL_H_ + +#include "backend/kernel_compiler/gpu/cuda_impl/sponge/nvtit/md_iteration_leap_frog_with_max_vel_impl.cuh" +#include +#include +#include +#include + +#include "backend/kernel_compiler/gpu/gpu_kernel.h" +#include "backend/kernel_compiler/gpu/gpu_kernel_factory.h" +#include "runtime/device/gpu/cuda_common.h" + +namespace mindspore { +namespace kernel { +template +class MDIterationLeapFrogWithMaxVelGpuKernel : public GpuKernel { + public: + MDIterationLeapFrogWithMaxVelGpuKernel() {} + ~MDIterationLeapFrogWithMaxVelGpuKernel() override = default; + + bool Init(const CNodePtr &kernel_node) override { + // get bond_numbers + kernel_node_ = kernel_node; + atom_numbers = static_cast(GetAttr(kernel_node, "atom_numbers")); + dt = static_cast(GetAttr(kernel_node, "dt")); + max_velocity = static_cast(GetAttr(kernel_node, "max_velocity")); + InitSizeLists(); + return true; + } + + const std::vector &GetInputSizeList() const override { return input_size_list_; } + 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 &workspace, + const std::vector &outputs, void *stream_ptr) override { + auto vel = GetDeviceAddress(inputs, 0); + auto crd = GetDeviceAddress(inputs, 1); + auto frc = GetDeviceAddress(inputs, 2); + auto acc = GetDeviceAddress(inputs, 3); + auto inverse_mass = GetDeviceAddress(inputs, 4); + + MDIterationLeapFrogWithMaxVelocity(atom_numbers, vel, crd, frc, acc, inverse_mass, dt, max_velocity, + reinterpret_cast(stream_ptr)); + return true; + } + + protected: + void InitSizeLists() override { + input_size_list_.push_back(atom_numbers * 3 * sizeof(T)); + input_size_list_.push_back(atom_numbers * 3 * sizeof(T)); + input_size_list_.push_back(atom_numbers * 3 * sizeof(T)); + input_size_list_.push_back(atom_numbers * 3 * sizeof(T)); + input_size_list_.push_back(atom_numbers * sizeof(T)); + + output_size_list_.push_back(sizeof(T)); + } + + private: + std::vector input_size_list_; + std::vector output_size_list_; + std::vector workspace_size_list_; + int atom_numbers; + float dt; + float max_velocity; +}; +} // namespace kernel +} // namespace mindspore +#endif