From 4cf48fdba62bcc18c69997c3ecb32441cfc23ca7 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 12:34:38 +0800 Subject: [PATCH 1/3] fix(tf): validate flattened multi-device op widths Reject partial per-atom rows and descriptor/neighbor stride mismatches before TensorFlow force and virial kernels allocate outputs or access raw pointers. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/tf/custom_op.h | 43 ++++++ source/op/tf/prod_force_grad_multi_device.cc | 22 ++- source/op/tf/prod_force_multi_device.cc | 24 +++- source/op/tf/prod_virial_grad_multi_device.cc | 22 ++- .../tf/test_multi_device_shape_validation.py | 134 ++++++++++++++++++ 5 files changed, 233 insertions(+), 12 deletions(-) create mode 100644 source/tests/tf/test_multi_device_shape_validation.py diff --git a/source/op/tf/custom_op.h b/source/op/tf/custom_op.h index 538bfeb2cc..a10b4f4d1b 100644 --- a/source/op/tf/custom_op.h +++ b/source/op/tf/custom_op.h @@ -1,6 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #pragma once #include +#include #include #include #include @@ -51,6 +52,48 @@ inline Status InvalidArgument(Args&&... args) { return tensorflow::errors::InvalidArgument(std::forward(args)...); #endif } + +/** + * @brief Derive a dense tensor's per-atom width without truncating division. + * + * Several low-level TensorFlow ops flatten atom and feature dimensions into a + * single axis. Validate the flattened width before dividing by `nloc`; raw + * CPU/GPU kernels cannot safely consume a leftover partial atom row. + * + * @param per_atom_width Receives the validated feature width for one atom. + * @param shape Rank-two tensor shape whose second dimension is flattened. + * @param nloc Number of local atoms encoded in the flattened dimension. + * @param tensor_name Human-readable input name used in validation errors. + * @return An OK status, or InvalidArgument when the width is incompatible. + */ +inline Status GetPerAtomWidth(int* per_atom_width, + const TensorShape& shape, + const int nloc, + const char* tensor_name) { + const int64_t flattened_width = shape.dim_size(1); + if (nloc < 0) { + return InvalidArgument("number of local atoms should be non-negative"); + } + if (nloc == 0) { + if (flattened_width != 0) { + return InvalidArgument(tensor_name, + " width should be zero when nloc is zero"); + } + *per_atom_width = 0; + return Status(); + } + if (flattened_width % nloc != 0) { + return InvalidArgument(tensor_name, " width ", flattened_width, + " should be divisible by nloc ", nloc); + } + const int64_t width = flattened_width / nloc; + if (width > std::numeric_limits::max()) { + return InvalidArgument(tensor_name, + " width per atom exceeds the supported int range"); + } + *per_atom_width = static_cast(width); + return Status(); +} } // namespace tf_compat } // namespace deepmd diff --git a/source/op/tf/prod_force_grad_multi_device.cc b/source/op/tf/prod_force_grad_multi_device.cc index 6678c1b628..39a26cde68 100644 --- a/source/op/tf/prod_force_grad_multi_device.cc +++ b/source/op/tf/prod_force_grad_multi_device.cc @@ -73,8 +73,12 @@ class ProdForceSeAGradOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = natoms(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_shape, nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_shape, nloc, "nlist")); // check the sizes OP_REQUIRES( @@ -97,6 +101,9 @@ class ProdForceSeAGradOp : public OpKernel { OP_REQUIRES( context, (nnei == n_a_sel + n_r_sel), deepmd::tf_compat::InvalidArgument("number of neighbors should match")); + OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should be four times neighbor width")); // Create an output tensor TensorShape grad_net_shape; @@ -194,8 +201,12 @@ class ProdForceSeRGradOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = natoms(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_shape, nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_shape, nloc, "nlist")); // check the sizes OP_REQUIRES( @@ -215,6 +226,9 @@ class ProdForceSeRGradOp : public OpKernel { (int_64(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), deepmd::tf_compat::InvalidArgument( "number of descriptors should match")); + OP_REQUIRES(context, (nnei == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should equal neighbor width")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_force_multi_device.cc b/source/op/tf/prod_force_multi_device.cc index ae66f13503..3dcce46821 100644 --- a/source/op/tf/prod_force_multi_device.cc +++ b/source/op/tf/prod_force_multi_device.cc @@ -92,8 +92,13 @@ class ProdForceSeAOp : public OpKernel { int nloc = natoms[0]; int nall = natoms[1]; int nframes = net_deriv_tensor.shape().dim_size(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, + deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_tensor.shape(), nloc, "nlist")); // check the sizes OP_REQUIRES( context, (nframes == in_deriv_tensor.shape().dim_size(0)), @@ -106,6 +111,9 @@ class ProdForceSeAOp : public OpKernel { (int_64(nloc) * ndescrpt * 3 == in_deriv_tensor.shape().dim_size(1)), deepmd::tf_compat::InvalidArgument( "number of descriptors should match")); + OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should be four times neighbor width")); // Create an output tensor TensorShape force_shape; force_shape.AddDim(nframes); @@ -199,8 +207,13 @@ class ProdForceSeROp : public OpKernel { int nloc = natoms[0]; int nall = natoms[1]; int nframes = net_deriv_tensor.shape().dim_size(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, + deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_tensor.shape(), nloc, "nlist")); // check the sizes OP_REQUIRES( context, (nframes == in_deriv_tensor.shape().dim_size(0)), @@ -213,6 +226,9 @@ class ProdForceSeROp : public OpKernel { in_deriv_tensor.shape().dim_size(1)), deepmd::tf_compat::InvalidArgument( "number of descriptors should match")); + OP_REQUIRES(context, (nnei == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should equal neighbor width")); // Create an output tensor TensorShape force_shape; force_shape.AddDim(nframes); diff --git a/source/op/tf/prod_virial_grad_multi_device.cc b/source/op/tf/prod_virial_grad_multi_device.cc index fb011cabbd..18c00cc0de 100644 --- a/source/op/tf/prod_virial_grad_multi_device.cc +++ b/source/op/tf/prod_virial_grad_multi_device.cc @@ -79,8 +79,12 @@ class ProdVirialSeAGradOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = natoms(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_shape, nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_shape, nloc, "nlist")); // check the sizes OP_REQUIRES( @@ -109,6 +113,9 @@ class ProdVirialSeAGradOp : public OpKernel { OP_REQUIRES( context, (nnei == n_a_sel + n_r_sel), deepmd::tf_compat::InvalidArgument("number of neighbors should match")); + OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should be four times neighbor width")); // Create an output tensor TensorShape grad_net_shape; @@ -223,8 +230,12 @@ class ProdVirialSeRGradOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = natoms(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_shape, nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_shape, nloc, "nlist")); // check the sizes OP_REQUIRES( @@ -250,6 +261,9 @@ class ProdVirialSeRGradOp : public OpKernel { OP_REQUIRES( context, (int_64(nloc) * nnei * 3 == rij_shape.dim_size(1)), deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); + OP_REQUIRES(context, (nnei == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should equal neighbor width")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/tests/tf/test_multi_device_shape_validation.py b/source/tests/tf/test_multi_device_shape_validation.py new file mode 100644 index 0000000000..56840f8fc3 --- /dev/null +++ b/source/tests/tf/test_multi_device_shape_validation.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for flattened TensorFlow custom-op input dimensions.""" + +from deepmd.tf.env import ( + GLOBAL_TF_FLOAT_PRECISION, + op_grads_module, + op_module, + tf, +) + + +class TestMultiDeviceShapeValidation(tf.test.TestCase): + """Ensure malformed flattened widths fail before native kernel dispatch.""" + + def setUp(self) -> None: + self.sess = self.cached_session().__enter__() + self.nloc = 2 + self.nnei = 1 + self.ndescrpt = 4 + self.natoms = tf.constant([self.nloc, self.nloc, 1], dtype=tf.int32) + + def _floats(self, width: int): + """Create one frame of flattened floating-point custom-op input.""" + return tf.zeros([1, width], dtype=GLOBAL_TF_FLOAT_PRECISION) + + def _nlist(self, width: int): + """Create one frame of flattened neighbor indices.""" + return tf.zeros([1, width], dtype=tf.int32) + + def test_prod_force_rejects_partial_net_deriv_atom(self) -> None: + # The old integer division truncated 9 / 2 to four descriptors and + # allowed the extra value to survive until raw pointer dispatch. + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"net deriv width 9 should be divisible by nloc 2", + ): + self.sess.run( + op_module.prod_force_se_a( + self._floats(self.nloc * self.ndescrpt + 1), + self._floats(self.nloc * self.ndescrpt * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_force_rejects_in_deriv_width_mismatch(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, r"number of descriptors should match" + ): + self.sess.run( + op_module.prod_force_se_a( + self._floats(self.nloc * self.ndescrpt), + self._floats(self.nloc * self.ndescrpt * 3 - 1), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_force_r_rejects_descriptor_stride_mismatch(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"descriptor width should equal neighbor width", + ): + self.sess.run( + op_module.prod_force_se_r( + self._floats(self.nloc * (self.nnei + 1)), + self._floats(self.nloc * (self.nnei + 1) * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + ) + ) + + def test_prod_force_grad_rejects_partial_nlist_atom(self) -> None: + # Fixed-width placeholders in the original tests rejected this feed + # before the custom op ran, leaving its release-build checks untested. + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"nlist width 3 should be divisible by nloc 2", + ): + self.sess.run( + op_grads_module.prod_force_se_a_grad( + self._floats(self.nloc * 3), + self._floats(self.nloc * self.ndescrpt), + self._floats(self.nloc * self.ndescrpt * 3), + self._nlist(self.nloc * self.nnei + 1), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_virial_grad_rejects_descriptor_stride_mismatch(self) -> None: + mismatched_ndescrpt = self.ndescrpt * 2 + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"descriptor width should be four times neighbor width", + ): + self.sess.run( + op_grads_module.prod_virial_se_a_grad( + self._floats(9), + self._floats(self.nloc * mismatched_ndescrpt), + self._floats(self.nloc * mismatched_ndescrpt * 3), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_virial_grad_rejects_rij_width_mismatch(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, r"dim of rij should be nnei \* 3" + ): + self.sess.run( + op_grads_module.prod_virial_se_a_grad( + self._floats(9), + self._floats(self.nloc * self.ndescrpt), + self._floats(self.nloc * self.ndescrpt * 3), + self._floats(self.nloc * self.nnei * 3 - 1), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + +if __name__ == "__main__": + tf.test.main() From e736a90b9064223261bc1cce014a8c568375e608 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 23 Jul 2026 20:31:29 +0800 Subject: [PATCH 2/3] test(tf): cover multi-device width guards Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- .../tf/test_multi_device_shape_validation.py | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/source/tests/tf/test_multi_device_shape_validation.py b/source/tests/tf/test_multi_device_shape_validation.py index 56840f8fc3..5a7f0bf2be 100644 --- a/source/tests/tf/test_multi_device_shape_validation.py +++ b/source/tests/tf/test_multi_device_shape_validation.py @@ -27,6 +27,58 @@ def _nlist(self, width: int): """Create one frame of flattened neighbor indices.""" return tf.zeros([1, width], dtype=tf.int32) + def test_negative_nloc_is_rejected(self) -> None: + natoms = tf.constant([-1, 0, 0], dtype=tf.int32) + + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"number of local atoms should be non-negative", + ): + self.sess.run( + op_module.prod_force_se_a( + self._floats(0), + self._floats(0), + self._nlist(0), + natoms, + n_a_sel=0, + n_r_sel=0, + ) + ) + + def test_zero_nloc_rejects_nonempty_flattened_width(self) -> None: + natoms = tf.constant([0, 0, 0], dtype=tf.int32) + + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"net deriv width should be zero when nloc is zero", + ): + self.sess.run( + op_module.prod_force_se_a( + self._floats(1), + self._floats(0), + self._nlist(0), + natoms, + n_a_sel=0, + n_r_sel=0, + ) + ) + + def test_zero_nloc_accepts_empty_flattened_widths(self) -> None: + natoms = tf.constant([0, 0, 0], dtype=tf.int32) + + result = self.sess.run( + op_module.prod_force_se_a( + self._floats(0), + self._floats(0), + self._nlist(0), + natoms, + n_a_sel=0, + n_r_sel=0, + ) + ) + + self.assertEqual(result.shape, (1, 0)) + def test_prod_force_rejects_partial_net_deriv_atom(self) -> None: # The old integer division truncated 9 / 2 to four descriptors and # allowed the extra value to survive until raw pointer dispatch. @@ -93,6 +145,21 @@ def test_prod_force_grad_rejects_partial_nlist_atom(self) -> None: ) ) + def test_prod_force_r_grad_rejects_partial_nlist_atom(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"nlist width 3 should be divisible by nloc 2", + ): + self.sess.run( + op_grads_module.prod_force_se_r_grad( + self._floats(self.nloc * 3), + self._floats(self.nloc * self.nnei), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei + 1), + self.natoms, + ) + ) + def test_prod_virial_grad_rejects_descriptor_stride_mismatch(self) -> None: mismatched_ndescrpt = self.ndescrpt * 2 with self.assertRaisesRegex( @@ -129,6 +196,22 @@ def test_prod_virial_grad_rejects_rij_width_mismatch(self) -> None: ) ) + def test_prod_virial_r_grad_rejects_partial_net_deriv_atom(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"net deriv width 3 should be divisible by nloc 2", + ): + self.sess.run( + op_grads_module.prod_virial_se_r_grad( + self._floats(9), + self._floats(self.nloc * self.nnei + 1), + self._floats(self.nloc * self.nnei * 3), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + ) + ) + if __name__ == "__main__": tf.test.main() From 1777fb211b8cd2eb50e3be4f9d35087187f11572 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 11 Aug 2026 03:46:29 +0800 Subject: [PATCH 3/3] fix(tf): harden multi-device zero-work paths Validate forward force and virial atom/descriptor dimensions, and make GPU force and virial kernels return safely after zeroing outputs when no work is present. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/src/gpu/prod_force.cu | 18 +- source/lib/src/gpu/prod_force_grad.cu | 16 +- source/lib/src/gpu/prod_virial.cu | 14 +- source/lib/src/gpu/prod_virial_grad.cu | 14 +- source/op/tf/prod_force_multi_device.cc | 6 + source/op/tf/prod_virial_multi_device.cc | 30 ++- .../tf/test_multi_device_shape_validation.py | 174 ++++++++++++++++++ 7 files changed, 258 insertions(+), 14 deletions(-) diff --git a/source/lib/src/gpu/prod_force.cu b/source/lib/src/gpu/prod_force.cu index 7b1359b3b0..993b6650cf 100644 --- a/source/lib/src/gpu/prod_force.cu +++ b/source/lib/src/gpu/prod_force.cu @@ -113,7 +113,14 @@ void prod_force_a_gpu(FPTYPE* force, DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); const int ndescrpt = nnei * 4; - DPErrcheck(gpuMemset(force, 0, sizeof(FPTYPE) * nframes * nall * 3)); + const size_t output_size = + sizeof(FPTYPE) * static_cast(nframes) * nall * 3; + if (output_size > 0) { + DPErrcheck(gpuMemset(force, 0, output_size)); + } + if (nframes == 0 || nloc == 0 || nnei == 0) { + return; + } force_deriv_wrt_center_atom<<>>( force, net_deriv, in_deriv, ndescrpt, nloc, nall); @@ -142,7 +149,14 @@ void prod_force_r_gpu(FPTYPE* force, DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); const int ndescrpt = nnei * 1; - DPErrcheck(gpuMemset(force, 0, sizeof(FPTYPE) * nframes * nall * 3)); + const size_t output_size = + sizeof(FPTYPE) * static_cast(nframes) * nall * 3; + if (output_size > 0) { + DPErrcheck(gpuMemset(force, 0, output_size)); + } + if (nframes == 0 || nloc == 0 || nnei == 0) { + return; + } force_deriv_wrt_center_atom<<>>( force, net_deriv, in_deriv, ndescrpt, nloc, nall); diff --git a/source/lib/src/gpu/prod_force_grad.cu b/source/lib/src/gpu/prod_force_grad.cu index c784d6ba65..bcda4d2b7b 100644 --- a/source/lib/src/gpu/prod_force_grad.cu +++ b/source/lib/src/gpu/prod_force_grad.cu @@ -91,8 +91,12 @@ void prod_force_grad_a_gpu(FPTYPE* grad_net, DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); const int ndescrpt = nnei * 4; - DPErrcheck( - gpuMemset(grad_net, 0, sizeof(FPTYPE) * nframes * nloc * ndescrpt)); + const size_t output_size = + sizeof(FPTYPE) * static_cast(nframes) * nloc * ndescrpt; + if (output_size == 0) { + return; + } + DPErrcheck(gpuMemset(grad_net, 0, output_size)); const int nblock = (ndescrpt + TPB - 1) / TPB; dim3 block_grid(nframes * nloc, nblock); dim3 thread_grid(TPB, 1); @@ -122,8 +126,12 @@ void prod_force_grad_r_gpu(FPTYPE* grad_net, DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); const int ndescrpt = nnei * 1; - DPErrcheck( - gpuMemset(grad_net, 0, sizeof(FPTYPE) * nframes * nloc * ndescrpt)); + const size_t output_size = + sizeof(FPTYPE) * static_cast(nframes) * nloc * ndescrpt; + if (output_size == 0) { + return; + } + DPErrcheck(gpuMemset(grad_net, 0, output_size)); const int nblock = (ndescrpt + TPB - 1) / TPB; dim3 block_grid(nframes * nloc, nblock); dim3 thread_grid(TPB, 1); diff --git a/source/lib/src/gpu/prod_virial.cu b/source/lib/src/gpu/prod_virial.cu index ab9c5326e3..55b9353922 100644 --- a/source/lib/src/gpu/prod_virial.cu +++ b/source/lib/src/gpu/prod_virial.cu @@ -116,7 +116,12 @@ void prod_virial_a_gpu(FPTYPE* virial, DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); DPErrcheck(gpuMemset(virial, 0, sizeof(FPTYPE) * 9)); - DPErrcheck(gpuMemset(atom_virial, 0, sizeof(FPTYPE) * 9 * nall)); + if (nall > 0) { + DPErrcheck(gpuMemset(atom_virial, 0, sizeof(FPTYPE) * 9 * nall)); + } + if (nloc == 0 || nnei == 0) { + return; + } const int LEN = 16; int nblock = (nnei + LEN - 1) / LEN; @@ -146,7 +151,12 @@ void prod_virial_r_gpu(FPTYPE* virial, DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); DPErrcheck(gpuMemset(virial, 0, sizeof(FPTYPE) * 9)); - DPErrcheck(gpuMemset(atom_virial, 0, sizeof(FPTYPE) * 9 * nall)); + if (nall > 0) { + DPErrcheck(gpuMemset(atom_virial, 0, sizeof(FPTYPE) * 9 * nall)); + } + if (nloc == 0 || nnei == 0) { + return; + } const int LEN = 16; int nblock = (nnei + LEN - 1) / LEN; diff --git a/source/lib/src/gpu/prod_virial_grad.cu b/source/lib/src/gpu/prod_virial_grad.cu index dac5b20ba8..65471f0402 100644 --- a/source/lib/src/gpu/prod_virial_grad.cu +++ b/source/lib/src/gpu/prod_virial_grad.cu @@ -95,7 +95,12 @@ void prod_virial_grad_a_gpu(FPTYPE* grad_net, DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); const int ndescrpt = nnei * 4; - DPErrcheck(gpuMemset(grad_net, 0, sizeof(FPTYPE) * nloc * ndescrpt)); + const size_t output_size = + sizeof(FPTYPE) * static_cast(nloc) * ndescrpt; + if (output_size == 0) { + return; + } + DPErrcheck(gpuMemset(grad_net, 0, output_size)); const int LEN = 128; const int nblock = (nloc + LEN - 1) / LEN; dim3 block_grid(nblock, nnei); @@ -117,7 +122,12 @@ void prod_virial_grad_r_gpu(FPTYPE* grad_net, DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); const int ndescrpt = nnei; - DPErrcheck(gpuMemset(grad_net, 0, sizeof(FPTYPE) * nloc * ndescrpt)); + const size_t output_size = + sizeof(FPTYPE) * static_cast(nloc) * ndescrpt; + if (output_size == 0) { + return; + } + DPErrcheck(gpuMemset(grad_net, 0, output_size)); const int LEN = 128; const int nblock = (nloc + LEN - 1) / LEN; dim3 block_grid(nblock, nnei); diff --git a/source/op/tf/prod_force_multi_device.cc b/source/op/tf/prod_force_multi_device.cc index 3dcce46821..b98c6d3593 100644 --- a/source/op/tf/prod_force_multi_device.cc +++ b/source/op/tf/prod_force_multi_device.cc @@ -94,6 +94,9 @@ class ProdForceSeAOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int ndescrpt; int nnei; + OP_REQUIRES(context, (nall >= nloc), + deepmd::tf_compat::InvalidArgument( + "number of all atoms should be at least nloc")); OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv")); @@ -209,6 +212,9 @@ class ProdForceSeROp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int ndescrpt; int nnei; + OP_REQUIRES(context, (nall >= nloc), + deepmd::tf_compat::InvalidArgument( + "number of all atoms should be at least nloc")); OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv")); diff --git a/source/op/tf/prod_virial_multi_device.cc b/source/op/tf/prod_virial_multi_device.cc index 5fff5cac8e..6f1c32628f 100644 --- a/source/op/tf/prod_virial_multi_device.cc +++ b/source/op/tf/prod_virial_multi_device.cc @@ -73,9 +73,17 @@ class ProdVirialSeAOp : public OpKernel { const int* natoms = natoms_tensor.flat().data(); int nloc = natoms[0]; int nall = natoms[1]; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; int nframes = net_deriv_tensor.shape().dim_size(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES(context, (nall >= nloc), + deepmd::tf_compat::InvalidArgument( + "number of all atoms should be at least nloc")); + OP_REQUIRES_OK(context, + deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_tensor.shape(), nloc, "nlist")); // check the sizes OP_REQUIRES( context, (nframes == in_deriv_tensor.shape().dim_size(0)), @@ -94,6 +102,9 @@ class ProdVirialSeAOp : public OpKernel { OP_REQUIRES( context, (int_64(nloc) * nnei * 3 == rij_tensor.shape().dim_size(1)), deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); + OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should be four times neighbor width")); // Create an output tensor TensorShape virial_shape; virial_shape.AddDim(nframes); @@ -179,9 +190,17 @@ class ProdVirialSeROp : public OpKernel { const int* natoms = natoms_tensor.flat().data(); int nloc = natoms[0]; int nall = natoms[1]; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; int nframes = net_deriv_tensor.shape().dim_size(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES(context, (nall >= nloc), + deepmd::tf_compat::InvalidArgument( + "number of all atoms should be at least nloc")); + OP_REQUIRES_OK(context, + deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_tensor.shape(), nloc, "nlist")); // check the sizes OP_REQUIRES( context, (nframes == in_deriv_tensor.shape().dim_size(0)), @@ -200,6 +219,9 @@ class ProdVirialSeROp : public OpKernel { OP_REQUIRES( context, (int_64(nloc) * nnei * 3 == rij_tensor.shape().dim_size(1)), deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); + OP_REQUIRES(context, (nnei == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should equal neighbor width")); // Create an output tensor TensorShape virial_shape; virial_shape.AddDim(nframes); diff --git a/source/tests/tf/test_multi_device_shape_validation.py b/source/tests/tf/test_multi_device_shape_validation.py index 5a7f0bf2be..8cf8fa801a 100644 --- a/source/tests/tf/test_multi_device_shape_validation.py +++ b/source/tests/tf/test_multi_device_shape_validation.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Regression tests for flattened TensorFlow custom-op input dimensions.""" +import unittest + from deepmd.tf.env import ( GLOBAL_TF_FLOAT_PRECISION, op_grads_module, @@ -126,6 +128,97 @@ def test_prod_force_r_rejects_descriptor_stride_mismatch(self) -> None: ) ) + def test_forward_ops_reject_nall_smaller_than_nloc(self) -> None: + natoms = tf.constant([self.nloc, self.nloc - 1, 1], dtype=tf.int32) + operations = ( + op_module.prod_force_se_a( + self._floats(self.nloc * self.ndescrpt), + self._floats(self.nloc * self.ndescrpt * 3), + self._nlist(self.nloc * self.nnei), + natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ), + op_module.prod_force_se_r( + self._floats(self.nloc * self.nnei), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + natoms, + ), + op_module.prod_virial_se_a( + self._floats(self.nloc * self.ndescrpt), + self._floats(self.nloc * self.ndescrpt * 3), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ), + op_module.prod_virial_se_r( + self._floats(self.nloc * self.nnei), + self._floats(self.nloc * self.nnei * 3), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + natoms, + ), + ) + for operation in operations: + with ( + self.subTest(operation=operation), + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"number of all atoms should be at least nloc", + ), + ): + self.sess.run(operation) + + def test_prod_virial_rejects_partial_net_deriv_atom(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"net deriv width 9 should be divisible by nloc 2", + ): + self.sess.run( + op_module.prod_virial_se_a( + self._floats(self.nloc * self.ndescrpt + 1), + self._floats(self.nloc * self.ndescrpt * 3), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_virial_rejects_descriptor_stride_mismatch(self) -> None: + operations = ( + op_module.prod_virial_se_a( + self._floats(self.nloc * self.ndescrpt * 2), + self._floats(self.nloc * self.ndescrpt * 2 * 3), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ), + op_module.prod_virial_se_r( + self._floats(self.nloc * (self.nnei + 1)), + self._floats(self.nloc * (self.nnei + 1) * 3), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + ), + ) + messages = ( + r"descriptor width should be four times neighbor width", + r"descriptor width should equal neighbor width", + ) + for operation, message in zip(operations, messages, strict=True): + with ( + self.subTest(message=message), + self.assertRaisesRegex(tf.errors.InvalidArgumentError, message), + ): + self.sess.run(operation) + def test_prod_force_grad_rejects_partial_nlist_atom(self) -> None: # Fixed-width placeholders in the original tests rejected this feed # before the custom op ran, leaving its release-build checks untested. @@ -212,6 +305,87 @@ def test_prod_virial_r_grad_rejects_partial_net_deriv_atom(self) -> None: ) ) + @unittest.skipUnless(tf.test.is_gpu_available(), "GPU is required") + def test_gpu_zero_neighbor_work_returns_zero_outputs(self) -> None: + natoms = tf.constant([1, 1, 1], dtype=tf.int32) + with tf.device("/GPU:0"): + empty_floats = self._floats(0) + empty_nlist = self._nlist(0) + force_grad = self._floats(3) + virial_grad = self._floats(9) + force_a = op_module.prod_force_se_a( + empty_floats, + empty_floats, + empty_nlist, + natoms, + n_a_sel=0, + n_r_sel=0, + ) + force_r = op_module.prod_force_se_r( + empty_floats, empty_floats, empty_nlist, natoms + ) + force_grad_a = op_grads_module.prod_force_se_a_grad( + force_grad, + empty_floats, + empty_floats, + empty_nlist, + natoms, + n_a_sel=0, + n_r_sel=0, + ) + force_grad_r = op_grads_module.prod_force_se_r_grad( + force_grad, empty_floats, empty_floats, empty_nlist, natoms + ) + virial_a = op_module.prod_virial_se_a( + empty_floats, + empty_floats, + empty_floats, + empty_nlist, + natoms, + n_a_sel=0, + n_r_sel=0, + ) + virial_r = op_module.prod_virial_se_r( + empty_floats, + empty_floats, + empty_floats, + empty_nlist, + natoms, + ) + virial_grad_a = op_grads_module.prod_virial_se_a_grad( + virial_grad, + empty_floats, + empty_floats, + empty_floats, + empty_nlist, + natoms, + n_a_sel=0, + n_r_sel=0, + ) + virial_grad_r = op_grads_module.prod_virial_se_r_grad( + virial_grad, + empty_floats, + empty_floats, + empty_floats, + empty_nlist, + natoms, + ) + + tensors = ( + force_a, + force_r, + force_grad_a, + force_grad_r, + *virial_a, + *virial_r, + virial_grad_a, + virial_grad_r, + ) + for tensor in tensors: + self.assertIn("GPU:0", tensor.device) + result = self.sess.run(tensor) + self.assertAllEqual(result, result * 0) + if __name__ == "__main__": tf.test.main()