From e7884e43c5d85a70af99c4b64ed332eeadedea1c Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 15 Jul 2026 01:20:18 +0800 Subject: [PATCH 1/4] fix(tf): avoid empty AMix neighbor mapping pointers Pass a neighbor-index mapping only when periodic coordinate expansion actually populated it. Non-PBC, external, and tensor-provided neighbor-list modes now pass nullptr instead of forming an invalid pointer from an empty vector. Add direct CPU coverage for the standard external-pointer AMix path and the NVNMD no-PBC path, including neighbor indices, types, masks, valid entries, and padding. Fixes #5653. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/tf/prod_env_mat_multi_device.cc | 7 +- .../op/tf/prod_env_mat_multi_device_nvnmd.cc | 7 +- source/tests/tf/test_prod_env_mat.py | 122 ++++++++++++++++++ 3 files changed, 132 insertions(+), 4 deletions(-) diff --git a/source/op/tf/prod_env_mat_multi_device.cc b/source/op/tf/prod_env_mat_multi_device.cc index 66a0a45c88..8133954db8 100644 --- a/source/op/tf/prod_env_mat_multi_device.cc +++ b/source/op/tf/prod_env_mat_multi_device.cc @@ -1964,8 +1964,11 @@ class ProdEnvMatAMixOp : public OpKernel { deepmd::prod_env_mat_a_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, + // The mapping is populated only when PBC coordinate copies are made. + // Other neighbor-list modes must not form a pointer into the empty + // vector. + _map_nei_info_cpu(nlist, ntype, nmask, type, + b_nlist_map ? idx_mapping.data() : nullptr, nloc, nnei, ntypes, b_nlist_map); } } 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..797a75e715 100644 --- a/source/op/tf/prod_env_mat_multi_device_nvnmd.cc +++ b/source/op/tf/prod_env_mat_multi_device_nvnmd.cc @@ -798,8 +798,11 @@ class ProdEnvMatAMixNvnmdQuantizeOp : public OpKernel { 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, + // The mapping is populated only when PBC coordinate copies are made. + // Other neighbor-list modes must not form a pointer into the empty + // vector. + _map_nei_info_cpu(nlist, ntype, nmask, type, + b_nlist_map ? idx_mapping.data() : nullptr, nloc, nnei, ntypes, b_nlist_map); } } diff --git a/source/tests/tf/test_prod_env_mat.py b/source/tests/tf/test_prod_env_mat.py index a1ea2d7ce7..1f4b969eff 100644 --- a/source/tests/tf/test_prod_env_mat.py +++ b/source/tests/tf/test_prod_env_mat.py @@ -1110,6 +1110,128 @@ def test_nopbc_self_built_nlist(self) -> None: for ff in range(self.nframes): np.testing.assert_almost_equal(dem[ff], self.nopbc_expected_output, 5) + def test_nopbc_nvnmd_amix(self) -> None: + """Exercise NVNMD AMix without the PBC-only coordinate mapping.""" + outputs = op_module.prod_env_mat_a_mix_nvnmd_quantize( + self.tcoord, + self.ttype, + self.tnatoms, + self.tbox, + tf.constant(np.zeros(0, dtype=np.int32)), + self.t_avg, + self.t_std, + rcut_a=-1, + rcut_r=self.rcut, + rcut_r_smth=self.rcut_smth, + sel_a=[self.nnei], + sel_r=[0], + ) + self.sess.run(tf.global_variables_initializer()) + result = self.sess.run( + outputs, + feed_dict={ + self.tcoord: self.dcoord, + self.ttype: self.dtype, + self.tbox: self.dbox, + self.tnatoms: self.dnatoms, + }, + ) + expected_shapes = ( + (self.nframes, self.nloc * self.ndescrpt), + (self.nframes, self.nloc * self.ndescrpt * 3), + (self.nframes, self.nloc * self.nnei * 3), + (self.nframes, self.nloc * self.nnei), + (self.nframes, self.nloc * self.nnei), + (self.nframes, self.nloc * self.nnei), + ) + for value, expected_shape in zip(result, expected_shapes, strict=True): + self.assertEqual(value.shape, expected_shape) + + nlist, neighbor_types, neighbor_mask = result[3:] + valid_neighbors = nlist >= 0 + self.assertTrue(np.any(valid_neighbors)) + self.assertTrue(np.any(~valid_neighbors)) + source_types = np.take_along_axis(self.dtype, np.maximum(nlist, 0), axis=1) + expected_types = np.where(valid_neighbors, source_types, self.ntypes) + np.testing.assert_array_equal(neighbor_mask, valid_neighbors) + np.testing.assert_array_equal(neighbor_types, expected_types) + + def test_external_nlist_amix(self) -> None: + """Exercise the AMix CPU fallback with a pointer-based neighbor list.""" + ilist = np.arange(self.nloc, dtype=np.int32) + coordinates = self.dcoord[0].reshape(self.nall, 3) + neighbor_rows = [ + np.array( + [ + jj + for jj in range(self.nall) + if jj != ii + and np.linalg.norm(coordinates[ii] - coordinates[jj]) < self.rcut + ], + dtype=np.int32, + ) + for ii in range(self.nloc) + ] + numneigh = np.array([row.size for row in neighbor_rows], dtype=np.int32) + firstneigh = np.array( + [row.ctypes.data for row in neighbor_rows], dtype=np.uintp + ) + + # The legacy external-list ABI stores native pointers at int offsets + # 4, 8, and 12. Keep every backing array alive through sess.run below. + mesh = np.zeros(16, dtype=np.int32) + mesh[1] = self.nloc + mesh_bytes = mesh.view(np.uint8) + for offset, address in ( + (4, ilist.ctypes.data), + (8, numneigh.ctypes.data), + (12, firstneigh.ctypes.data), + ): + address_bytes = np.array([address], dtype=np.uintp).view(np.uint8) + byte_offset = offset * mesh.itemsize + mesh_bytes[byte_offset : byte_offset + address_bytes.size] = address_bytes + + with tf.device("/CPU:0"): + outputs = op_module.prod_env_mat_a_mix( + self.tcoord, + self.ttype, + self.tnatoms, + self.tbox, + tf.constant(mesh), + self.t_avg, + self.t_std, + rcut_a=-1, + rcut_r=self.rcut, + rcut_r_smth=self.rcut_smth, + sel_a=[self.nnei], + sel_r=[0], + ) + self.sess.run(tf.global_variables_initializer()) + nlist, neighbor_types, neighbor_mask = self.sess.run( + outputs[3:], + feed_dict={ + self.tcoord: self.dcoord[:1], + self.ttype: self.dtype[:1], + self.tbox: self.dbox[:1], + self.tnatoms: self.dnatoms, + }, + ) + + nlist = nlist.reshape(self.nloc, self.nnei) + neighbor_types = neighbor_types.reshape(self.nloc, self.nnei) + neighbor_mask = neighbor_mask.reshape(self.nloc, self.nnei) + valid_neighbors = nlist >= 0 + np.testing.assert_array_equal(valid_neighbors.sum(axis=1), numneigh) + for ii, row in enumerate(neighbor_rows): + np.testing.assert_array_equal(np.sort(nlist[ii][valid_neighbors[ii]]), row) + expected_types = np.where( + valid_neighbors, + self.dtype[0][np.maximum(nlist, 0)], + self.ntypes, + ) + np.testing.assert_array_equal(neighbor_mask, valid_neighbors) + np.testing.assert_array_equal(neighbor_types, expected_types) + def test_nopbc_self_built_nlist_deriv(self) -> None: hh = 1e-4 tem, tem_deriv, trij, tnlist = op_module.prod_env_mat_a( From 54abd29328a3870e346cdb247ae120227bb0c561 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 22:57:28 +0800 Subject: [PATCH 2/4] test(tf): cover hardened NVNMD AMix mapping Exercise the periodic NVNMD AMix remapping branch directly and run both affected AMix paths in the C++ sanitizer matrix with libstdc++ bounds assertions enabled. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- .github/workflows/test_cc.yml | 16 ++++++++++- source/tests/tf/test_prod_env_mat.py | 40 ++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml index 153eb09586..d0d03fb008 100644 --- a/.github/workflows/test_cc.yml +++ b/.github/workflows/test_cc.yml @@ -63,7 +63,7 @@ jobs: TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 CMAKE_GENERATOR: Ninja - CXXFLAGS: ${{ matrix.check_memleak && '-fsanitize=leak' || '' }} + CXXFLAGS: ${{ matrix.check_memleak && '-fsanitize=leak -D_GLIBCXX_ASSERTIONS' || '' }} LSAN_OPTIONS: suppressions=${{ github.workspace }}/.github/workflows/suppr.txt ENABLE_TENSORFLOW: ${{ matrix.enable_tensorflow && 'TRUE' || 'FALSE' }} ENABLE_PYTORCH: ${{ matrix.enable_pytorch && 'TRUE' || 'FALSE' }} @@ -87,6 +87,20 @@ jobs: ENABLE_JAX: ${{ matrix.enable_tensorflow && '1' || '0' }} ENABLE_PADDLE: ${{ matrix.enable_paddle && '1' || '0' }} if: ${{ !matrix.check_memleak }} + - name: Test AMix neighbor mapping with hardened libstdc++ + run: | + export LD_LIBRARY_PATH=${{ github.workspace }}/dp_test/lib:$LD_LIBRARY_PATH + pytest -q \ + source/tests/tf/test_prod_env_mat.py::TestProdEnvMat::test_nopbc_nvnmd_amix \ + source/tests/tf/test_prod_env_mat.py::TestProdEnvMat::test_pbc_nvnmd_amix_maps_copied_neighbors + env: + OMP_NUM_THREADS: 1 + TF_INTRA_OP_PARALLELISM_THREADS: 1 + TF_INTER_OP_PARALLELISM_THREADS: 1 + # This step targets libstdc++ bounds assertions. The main C++ test + # step above remains responsible for leak detection. + LSAN_OPTIONS: detect_leaks=0 + if: ${{ matrix.check_memleak && matrix.enable_tensorflow }} # test ipi - run: | export PATH=${{ github.workspace }}/dp_test/bin:$PATH diff --git a/source/tests/tf/test_prod_env_mat.py b/source/tests/tf/test_prod_env_mat.py index 1f4b969eff..2c61248e1c 100644 --- a/source/tests/tf/test_prod_env_mat.py +++ b/source/tests/tf/test_prod_env_mat.py @@ -1156,6 +1156,46 @@ def test_nopbc_nvnmd_amix(self) -> None: np.testing.assert_array_equal(neighbor_mask, valid_neighbors) np.testing.assert_array_equal(neighbor_types, expected_types) + def test_pbc_nvnmd_amix_maps_copied_neighbors(self) -> None: + """Exercise the NVNMD AMix branch that remaps periodic atom copies.""" + outputs = op_module.prod_env_mat_a_mix_nvnmd_quantize( + self.tcoord, + self.ttype, + self.tnatoms, + self.tbox, + tf.constant(np.zeros(6, dtype=np.int32)), + self.t_avg, + self.t_std, + rcut_a=-1, + rcut_r=self.rcut, + rcut_r_smth=self.rcut_smth, + sel_a=[self.nnei], + sel_r=[0], + ) + self.sess.run(tf.global_variables_initializer()) + nlist, neighbor_types, neighbor_mask = self.sess.run( + outputs[3:], + feed_dict={ + self.tcoord: self.dcoord, + self.ttype: self.dtype, + self.tbox: self.dbox, + self.tnatoms: self.dnatoms, + }, + ) + + nlist = nlist.reshape(self.nframes, self.nloc, self.nnei) + neighbor_types = neighbor_types.reshape(self.nframes, self.nloc, self.nnei) + neighbor_mask = neighbor_mask.reshape(self.nframes, self.nloc, self.nnei) + valid_neighbors = nlist >= 0 + self.assertTrue(np.any(valid_neighbors)) + self.assertTrue(np.all(nlist[valid_neighbors] < self.nall)) + source_types = np.take_along_axis( + self.dtype[:, None, :], np.maximum(nlist, 0), axis=2 + ) + expected_types = np.where(valid_neighbors, source_types, self.ntypes) + np.testing.assert_array_equal(neighbor_mask, valid_neighbors) + np.testing.assert_array_equal(neighbor_types, expected_types) + def test_external_nlist_amix(self) -> None: """Exercise the AMix CPU fallback with a pointer-based neighbor list.""" ilist = np.arange(self.nloc, dtype=np.int32) From 64260c562041f93446ad4dc1e36c76703bdc6d2c Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 23:31:09 +0800 Subject: [PATCH 3/4] ci(tf): harden standard AMix regression Run the standard external-neighbor-list AMix test alongside both NVNMD mapping cases in the libstdc++ assertions job. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- .github/workflows/test_cc.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml index d0d03fb008..c7cfce2b1d 100644 --- a/.github/workflows/test_cc.yml +++ b/.github/workflows/test_cc.yml @@ -92,7 +92,8 @@ jobs: export LD_LIBRARY_PATH=${{ github.workspace }}/dp_test/lib:$LD_LIBRARY_PATH pytest -q \ source/tests/tf/test_prod_env_mat.py::TestProdEnvMat::test_nopbc_nvnmd_amix \ - source/tests/tf/test_prod_env_mat.py::TestProdEnvMat::test_pbc_nvnmd_amix_maps_copied_neighbors + source/tests/tf/test_prod_env_mat.py::TestProdEnvMat::test_pbc_nvnmd_amix_maps_copied_neighbors \ + source/tests/tf/test_prod_env_mat.py::TestProdEnvMat::test_external_nlist_amix env: OMP_NUM_THREADS: 1 TF_INTRA_OP_PARALLELISM_THREADS: 1 From cb556ea75483f3757bba7c77c5dce90a9ad0be81 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 18 Jul 2026 14:25:52 +0800 Subject: [PATCH 4/4] ci(tf): drop AMix workflow hardening Keep the AMix pointer fix scoped to implementation and regression tests. The existing leak-check matrix should not be repurposed as a global libstdc++ assertions job. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- .github/workflows/test_cc.yml | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml index c7cfce2b1d..153eb09586 100644 --- a/.github/workflows/test_cc.yml +++ b/.github/workflows/test_cc.yml @@ -63,7 +63,7 @@ jobs: TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 CMAKE_GENERATOR: Ninja - CXXFLAGS: ${{ matrix.check_memleak && '-fsanitize=leak -D_GLIBCXX_ASSERTIONS' || '' }} + CXXFLAGS: ${{ matrix.check_memleak && '-fsanitize=leak' || '' }} LSAN_OPTIONS: suppressions=${{ github.workspace }}/.github/workflows/suppr.txt ENABLE_TENSORFLOW: ${{ matrix.enable_tensorflow && 'TRUE' || 'FALSE' }} ENABLE_PYTORCH: ${{ matrix.enable_pytorch && 'TRUE' || 'FALSE' }} @@ -87,21 +87,6 @@ jobs: ENABLE_JAX: ${{ matrix.enable_tensorflow && '1' || '0' }} ENABLE_PADDLE: ${{ matrix.enable_paddle && '1' || '0' }} if: ${{ !matrix.check_memleak }} - - name: Test AMix neighbor mapping with hardened libstdc++ - run: | - export LD_LIBRARY_PATH=${{ github.workspace }}/dp_test/lib:$LD_LIBRARY_PATH - pytest -q \ - source/tests/tf/test_prod_env_mat.py::TestProdEnvMat::test_nopbc_nvnmd_amix \ - source/tests/tf/test_prod_env_mat.py::TestProdEnvMat::test_pbc_nvnmd_amix_maps_copied_neighbors \ - source/tests/tf/test_prod_env_mat.py::TestProdEnvMat::test_external_nlist_amix - env: - OMP_NUM_THREADS: 1 - TF_INTRA_OP_PARALLELISM_THREADS: 1 - TF_INTER_OP_PARALLELISM_THREADS: 1 - # This step targets libstdc++ bounds assertions. The main C++ test - # step above remains responsible for leak detection. - LSAN_OPTIONS: detect_leaks=0 - if: ${{ matrix.check_memleak && matrix.enable_tensorflow }} # test ipi - run: | export PATH=${{ github.workspace }}/dp_test/bin:$PATH