From 29cc7b5292967ff68ee816cc47641bb5b21b99ea Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 05:03:40 +0800 Subject: [PATCH] fix(tf): honor NVNMD tensor neighbor lists Parse and validate mode-4 mesh tensors instead of rebuilding their neighbor lists in the NVNMD environment-matrix operators. Reject truncated lists, invalid local indices, and neighbor indices outside the current frame before downstream coordinate or type access. Cover both operator variants, malformed tensor input, and positive out-of-range neighbor indices. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- .../op/tf/prod_env_mat_multi_device_nvnmd.cc | 203 ++++++++++++------ source/tests/tf/test_nvnmd_op.py | 68 ++++++ 2 files changed, 211 insertions(+), 60 deletions(-) diff --git a/source/op/tf/prod_env_mat_multi_device_nvnmd.cc b/source/op/tf/prod_env_mat_multi_device_nvnmd.cc index 95bcb71b3c..0540aaa05e 100644 --- a/source/op/tf/prod_env_mat_multi_device_nvnmd.cc +++ b/source/op/tf/prod_env_mat_multi_device_nvnmd.cc @@ -16,6 +16,8 @@ date: 2021-12-6 */ +#include + #include "coord.h" #include "custom_op.h" #include "errors.h" @@ -112,28 +114,29 @@ static void _map_nei_info_cpu(int* nlist, const bool& b_nlist_map); template -static void _prepare_coord_nlist_cpu(OpKernelContext* context, - FPTYPE const** coord, - std::vector& coord_cpy, - int const** type, - std::vector& type_cpy, - std::vector& idx_mapping, - deepmd::InputNlist& inlist, - std::vector& ilist, - std::vector& numneigh, - std::vector& firstneigh, - std::vector>& jlist, - int& new_nall, - int& mem_cpy, - int& mem_nnei, - int& max_nbor_size, - const FPTYPE* box, - const int* mesh_tensor_data, - const int& nloc, - const int& nei_mode, - const float& rcut_r, - const int& max_cpy_trial, - const int& max_nnei_trial); +static tensorflow::Status _prepare_coord_nlist_cpu( + FPTYPE const** coord, + std::vector& coord_cpy, + int const** type, + std::vector& type_cpy, + std::vector& idx_mapping, + deepmd::InputNlist& inlist, + std::vector& ilist, + std::vector& numneigh, + std::vector& firstneigh, + std::vector>& jlist, + int& new_nall, + int& mem_cpy, + int& mem_nnei, + int& max_nbor_size, + const FPTYPE* box, + const int* mesh_tensor_data, + const int_64 mesh_tensor_size, + const int& nloc, + const int& nei_mode, + const float& rcut_r, + const int& max_cpy_trial, + const int& max_nnei_trial); // instance of function @@ -228,39 +231,104 @@ static void _map_nei_info_cpu(int* nlist, ntypes, b_nlist_map); } +static tensorflow::Status _prepare_mesh_nlist_cpu( + deepmd::InputNlist& inlist, + std::vector& ilist, + std::vector& numneigh, + std::vector& firstneigh, + std::vector>& jlist, + int& max_nbor_size, + const int* mesh_tensor_data, + const int_64 mesh_tensor_size, + const int nloc, + const int nall) { + // Tensor-stored lists use a 16-int header followed by ilist, numneigh, + // and the concatenated neighbor rows. Validate every offset before copying + // because this input can be constructed directly by callers. + const int_64 header_size = 16 + static_cast(2) * nloc; + if (mesh_tensor_size < header_size) { + return errors::InvalidArgument("invalid mesh tensor"); + } + + const int* ilist_in = mesh_tensor_data + 16; + const int* numneigh_in = mesh_tensor_data + 16 + nloc; + const int* neighbors_in = mesh_tensor_data + header_size; + std::vector neighbor_offset(nloc + 1, 0); + int supplied_max_nbor_size = 0; + for (int ii = 0; ii < nloc; ++ii) { + const int_64 neighbor_count = numneigh_in[ii]; + if (neighbor_count < 0 || + neighbor_offset[ii] > mesh_tensor_size - header_size - neighbor_count) { + return errors::InvalidArgument("invalid mesh tensor"); + } + neighbor_offset[ii + 1] = neighbor_offset[ii] + neighbor_count; + supplied_max_nbor_size = std::max(supplied_max_nbor_size, numneigh_in[ii]); + } + + std::vector seen(static_cast(nloc), 0); + for (int ii = 0; ii < nloc; ++ii) { + const int i_idx = ilist_in[ii]; + if (i_idx < 0 || i_idx >= nloc || seen[static_cast(i_idx)]) { + return errors::InvalidArgument("invalid mesh tensor"); + } + seen[static_cast(i_idx)] = 1; + for (int_64 jj = neighbor_offset[ii]; jj < neighbor_offset[ii + 1]; ++jj) { + // Downstream environment-matrix kernels index both coordinates and + // atom types with this value, so accept only live frame atoms. + if (neighbors_in[jj] < 0 || neighbors_in[jj] >= nall) { + return errors::InvalidArgument("invalid mesh tensor"); + } + } + } + + for (int ii = 0; ii < nloc; ++ii) { + ilist[ii] = ilist_in[ii]; + numneigh[ii] = numneigh_in[ii]; + jlist[ii].assign(neighbors_in + neighbor_offset[ii], + neighbors_in + neighbor_offset[ii + 1]); + firstneigh[ii] = jlist[ii].data(); + } + inlist = deepmd::InputNlist(nloc, ilist.data(), numneigh.data(), + firstneigh.data()); + max_nbor_size = std::max(max_nbor_size, supplied_max_nbor_size); + return tensorflow::Status(); +} + template -static void _prepare_coord_nlist_cpu(OpKernelContext* context, - FPTYPE const** coord, - std::vector& coord_cpy, - int const** type, - std::vector& type_cpy, - std::vector& idx_mapping, - deepmd::InputNlist& inlist, - std::vector& ilist, - std::vector& numneigh, - std::vector& firstneigh, - std::vector>& jlist, - int& new_nall, - int& mem_cpy, - int& mem_nnei, - int& max_nbor_size, - const FPTYPE* box, - const int* mesh_tensor_data, - const int& nloc, - const int& nei_mode, - const float& rcut_r, - const int& max_cpy_trial, - const int& max_nnei_trial) { +static tensorflow::Status _prepare_coord_nlist_cpu( + FPTYPE const** coord, + std::vector& coord_cpy, + int const** type, + std::vector& type_cpy, + std::vector& idx_mapping, + deepmd::InputNlist& inlist, + std::vector& ilist, + std::vector& numneigh, + std::vector& firstneigh, + std::vector>& jlist, + int& new_nall, + int& mem_cpy, + int& mem_nnei, + int& max_nbor_size, + const FPTYPE* box, + const int* mesh_tensor_data, + const int_64 mesh_tensor_size, + const int& nloc, + const int& nei_mode, + const float& rcut_r, + const int& max_cpy_trial, + const int& max_nnei_trial) { inlist.inum = nloc; - if (nei_mode != 3) { + if (nei_mode != 3 && nei_mode != 4) { // build nlist by myself // normalize and copy coord if (nei_mode == 1) { int copy_ok = _norm_copy_coord_cpu(coord_cpy, type_cpy, idx_mapping, new_nall, mem_cpy, *coord, box, *type, nloc, max_cpy_trial, rcut_r); - OP_REQUIRES(context, copy_ok, - errors::Aborted("cannot allocate mem for copied coords")); + if (!copy_ok) { + return errors::Aborted("cannot allocate mem for copied coords"); + } *coord = &coord_cpy[0]; *type = &type_cpy[0]; } @@ -268,11 +336,19 @@ static void _prepare_coord_nlist_cpu(OpKernelContext* context, int build_ok = _build_nlist_cpu(ilist, numneigh, firstneigh, jlist, max_nbor_size, mem_nnei, *coord, nloc, new_nall, max_nnei_trial, rcut_r); - OP_REQUIRES(context, build_ok, - errors::Aborted("cannot allocate mem for nlist")); + if (!build_ok) { + return errors::Aborted("cannot allocate mem for nlist"); + } inlist.ilist = &ilist[0]; inlist.numneigh = &numneigh[0]; inlist.firstneigh = &firstneigh[0]; + } else if (nei_mode == 4) { + tensorflow::Status status = _prepare_mesh_nlist_cpu( + inlist, ilist, numneigh, firstneigh, jlist, max_nbor_size, + mesh_tensor_data, mesh_tensor_size, nloc, new_nall); + if (!status.ok()) { + return status; + } } else { // copy pointers to nlist data memcpy(&inlist.ilist, 4 + mesh_tensor_data, sizeof(int*)); @@ -280,6 +356,7 @@ static void _prepare_coord_nlist_cpu(OpKernelContext* context, memcpy(&inlist.firstneigh, 12 + mesh_tensor_data, sizeof(int**)); max_nbor_size = max_numneigh(inlist); } + return tensorflow::Status(); } /* @@ -506,11 +583,14 @@ class ProdEnvMatANvnmdQuantizeOp : public OpKernel { std::vector type_cpy; int frame_nall = nall; // prepare coord and nlist - _prepare_coord_nlist_cpu( - context, &coord, coord_cpy, &type, type_cpy, idx_mapping, inlist, - ilist, numneigh, firstneigh, jlist, frame_nall, mem_cpy, mem_nnei, - max_nbor_size, box, mesh_tensor.flat().data(), nloc, nei_mode, - rcut_r, max_cpy_trial, max_nnei_trial); + OP_REQUIRES_OK( + context, + _prepare_coord_nlist_cpu( + &coord, coord_cpy, &type, type_cpy, idx_mapping, inlist, ilist, + numneigh, firstneigh, jlist, frame_nall, mem_cpy, mem_nnei, + max_nbor_size, box, mesh_tensor.flat().data(), + mesh_tensor.NumElements(), nloc, nei_mode, rcut_r, + max_cpy_trial, max_nnei_trial)); // launch the cpu compute function deepmd::prod_env_mat_a_nvnmd_quantize_cpu( em, em_deriv, rij, nlist, coord, type, inlist, max_nbor_size, avg, @@ -789,17 +869,20 @@ class ProdEnvMatAMixNvnmdQuantizeOp : public OpKernel { std::vector type_cpy; int frame_nall = nall; // prepare coord and nlist - _prepare_coord_nlist_cpu( - context, &coord, coord_cpy, &f_type, type_cpy, idx_mapping, inlist, - ilist, numneigh, firstneigh, jlist, frame_nall, mem_cpy, mem_nnei, - max_nbor_size, box, mesh_tensor.flat().data(), nloc, nei_mode, - rcut_r, max_cpy_trial, max_nnei_trial); + OP_REQUIRES_OK( + context, + _prepare_coord_nlist_cpu( + &coord, coord_cpy, &f_type, type_cpy, idx_mapping, inlist, + ilist, numneigh, firstneigh, jlist, frame_nall, mem_cpy, + mem_nnei, max_nbor_size, box, mesh_tensor.flat().data(), + mesh_tensor.NumElements(), nloc, nei_mode, rcut_r, + max_cpy_trial, max_nnei_trial)); // launch the cpu compute function deepmd::prod_env_mat_a_nvnmd_quantize_cpu( em, em_deriv, rij, nlist, coord, type, inlist, max_nbor_size, avg, std, nloc, frame_nall, rcut_r, rcut_r_smth, sec_a, f_type); // do nlist mapping if coords were copied - _map_nei_info_cpu(nlist, ntype, nmask, type, &idx_mapping[0], nloc, + _map_nei_info_cpu(nlist, ntype, nmask, type, idx_mapping.data(), nloc, nnei, ntypes, b_nlist_map); } } diff --git a/source/tests/tf/test_nvnmd_op.py b/source/tests/tf/test_nvnmd_op.py index acf345b61d..9cafd8594e 100644 --- a/source/tests/tf/test_nvnmd_op.py +++ b/source/tests/tf/test_nvnmd_op.py @@ -421,5 +421,73 @@ def test_op(self) -> None: tf.reset_default_graph() +class TestOpProdEnvMatNvnmdTensorNlist(tf.test.TestCase): + """Verify NVNMD env-mat ops honor neighbor lists stored in mesh tensors.""" + + def setUp(self) -> None: + self.coord = np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0]]) + self.atype = np.array([[0, 0, 0]], dtype=np.int32) + self.natoms = np.array([3, 3, 3], dtype=np.int32) + self.box = np.array([[10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0]]) + self.sel = [2] + self.ndescrpt = 4 * sum(self.sel) + self.avg = np.zeros((1, self.ndescrpt)) + self.std = np.ones((1, self.ndescrpt)) + + # All atoms are mutually within rcut, but this caller-provided list + # deliberately retains only one cyclic neighbor per local atom. + header = np.zeros(16, dtype=np.int32) + ilist = np.array([0, 1, 2], dtype=np.int32) + numneigh = np.ones(3, dtype=np.int32) + jlist = np.array([2, 0, 1], dtype=np.int32) + self.mesh = np.concatenate((header, ilist, numneigh, jlist)) + self.expected_nlist = np.array([[2, -1, 0, -1, 1, -1]], dtype=np.int32) + + def _run_op(self, mixed_types: bool, mesh: np.ndarray | None = None) -> np.ndarray: + op = ( + op_module.prod_env_mat_a_mix_nvnmd_quantize + if mixed_types + else op_module.prod_env_mat_a_nvnmd_quantize + ) + outputs = op( + tf.constant(self.coord, dtype=tf.float64), + tf.constant(self.atype), + tf.constant(self.natoms), + tf.constant(self.box, dtype=tf.float64), + tf.constant(self.mesh if mesh is None else mesh), + tf.constant(self.avg, dtype=tf.float64), + tf.constant(self.std, dtype=tf.float64), + rcut_a=-1.0, + rcut_r=3.0, + rcut_r_smth=0.5, + sel_a=self.sel, + sel_r=[0], + ) + with self.cached_session() as sess: + return sess.run(outputs[3]) + + def test_standard_op_uses_tensor_neighbor_list(self) -> None: + np.testing.assert_array_equal(self._run_op(False), self.expected_nlist) + + def test_mixed_type_op_uses_tensor_neighbor_list(self) -> None: + np.testing.assert_array_equal(self._run_op(True), self.expected_nlist) + + def test_rejects_truncated_tensor_neighbor_list(self) -> None: + """Reject mode-4 tensors before reading an incomplete list header.""" + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "invalid mesh tensor" + ): + self._run_op(False, self.mesh[:17]) + + def test_rejects_out_of_range_tensor_neighbor_index(self) -> None: + """Reject neighbors that would index beyond this frame's atoms.""" + invalid_mesh = self.mesh.copy() + invalid_mesh[-1] = self.natoms[1] + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "invalid mesh tensor" + ): + self._run_op(False, invalid_mesh) + + if __name__ == "__main__": unittest.main()