From 667181a8e3ab89cc2b071f2ae34c96b77fd2c39b Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 12:30:14 +0800 Subject: [PATCH 1/2] fix(gpu): preserve exact-capacity neighbor rows Compare sorted neighbor keys with the explicit padding sentinel so a fully occupied row keeps its final valid entry. Cover every supported GPU sort capacity with a boundary regression test. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/src/gpu/prod_env_mat.cu | 10 ++++- source/lib/tests/test_fmt_nlist.cc | 62 ++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/source/lib/src/gpu/prod_env_mat.cu b/source/lib/src/gpu/prod_env_mat.cu index 6e8fc1df85..66ee7d15b8 100644 --- a/source/lib/src/gpu/prod_env_mat.cu +++ b/source/lib/src/gpu/prod_env_mat.cu @@ -110,6 +110,12 @@ __device__ inline void decoding_nbor_info(int& type, index = key & 0xFFFFFF; } +__device__ inline bool is_padding_nbor_info(const uint_64 key) { + // Empty sort slots are initialized bytewise to 0xff. Compare with that + // sentinel directly because a full row has a valid key in its final slot. + return key == static_cast(-1); +} + __global__ void get_i_idx(int* i_idx, const int nloc, const int nframes, @@ -175,7 +181,7 @@ __global__ void fill_nei_iter(int* nei_iter_dev, const FPTYPE* key_out = key + nloc * max_nbor_size + row * max_nbor_size; int nei_type_cur = -1, nbor_idx_cur = 0; int nei_type_pre = -1, nbor_idx_pre = 0; - if (col < max_nbor_size && key_out[col] != key_out[max_nbor_size - 1]) { + if (col < max_nbor_size && !is_padding_nbor_info(key_out[col])) { if (col >= 1) { decoding_nbor_info(nei_type_pre, nbor_idx_pre, key_out[col - 1]); } @@ -201,7 +207,7 @@ __global__ void format_nlist_fill_b(int* nlist, FPTYPE* key_out = key + nloc * max_nbor_size + row * max_nbor_size; int* row_nlist = nlist + row * nlist_size; if (col < max_nbor_size) { - if (key_out[col] != key_out[max_nbor_size - 1]) { + if (!is_padding_nbor_info(key_out[col])) { int nei_type = 0, nbor_idx = 0; decoding_nbor_info(nei_type, nbor_idx, key_out[col]); int out_indx = col - nei_iter[nei_type] + sec[nei_type]; diff --git a/source/lib/tests/test_fmt_nlist.cc b/source/lib/tests/test_fmt_nlist.cc index 511a2be949..d2dfc8c365 100644 --- a/source/lib/tests/test_fmt_nlist.cc +++ b/source/lib/tests/test_fmt_nlist.cc @@ -1,6 +1,8 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include +#include + #include "fmt_nlist.h" #include "neighbor_list.h" @@ -434,6 +436,66 @@ TEST_F(TestFormatNlistShortSel, gpu) { } } +TEST(FormatNlistGpu, preserves_exact_capacity_rows) { + // Exercise every supported radix-sort size without leaving a padding slot. + // These exact boundaries are reachable after the caller rounds row capacity. + for (const int max_nbor_size : {256, 512, 1024, 2048, 4096}) { + SCOPED_TRACE("max_nbor_size=" + std::to_string(max_nbor_size)); + const int nloc = 1; + const int nall = max_nbor_size + 1; + const float rcut = 2.0f; + const std::vector sec = {0, max_nbor_size}; + std::vector coord(static_cast(nall) * 3, 0.0); + std::vector type(nall, 0); + std::vector neighbors(max_nbor_size); + for (int ii = 0; ii < max_nbor_size; ++ii) { + neighbors[ii] = ii + 1; + coord[static_cast(ii + 1) * 3] = 1.0; + } + + std::vector ilist = {0}; + std::vector numneigh = {max_nbor_size}; + std::vector firstneigh = {neighbors.data()}; + deepmd::InputNlist in_nlist(nloc, ilist.data(), numneigh.data(), + firstneigh.data()), + gpu_inlist; + std::vector formatted(max_nbor_size, -1); + + double* coord_dev = NULL; + int *type_dev = NULL, *nlist_dev = NULL, *array_int_dev = NULL, + *memory_dev = NULL; + uint_64* array_longlong_dev = NULL; + deepmd::malloc_device_memory_sync(coord_dev, coord); + deepmd::malloc_device_memory_sync(type_dev, type); + deepmd::malloc_device_memory_sync(nlist_dev, formatted); + deepmd::malloc_device_memory(array_int_dev, + sec.size() + nloc * sec.size() + nloc); + deepmd::malloc_device_memory(array_longlong_dev, + static_cast(nloc) * max_nbor_size * 2); + deepmd::malloc_device_memory(memory_dev, + static_cast(nloc) * max_nbor_size); + deepmd::convert_nlist_gpu_device(gpu_inlist, in_nlist, memory_dev, + max_nbor_size); + + format_nbor_list_gpu(nlist_dev, coord_dev, type_dev, gpu_inlist, + array_int_dev, array_longlong_dev, max_nbor_size, nloc, + nall, 1, rcut, sec); + deepmd::memcpy_device_to_host(nlist_dev, formatted); + + deepmd::delete_device_memory(nlist_dev); + deepmd::delete_device_memory(coord_dev); + deepmd::delete_device_memory(type_dev); + deepmd::delete_device_memory(array_int_dev); + deepmd::delete_device_memory(array_longlong_dev); + deepmd::delete_device_memory(memory_dev); + deepmd::free_nlist_gpu_device(gpu_inlist); + + for (int ii = 0; ii < max_nbor_size; ++ii) { + EXPECT_EQ(formatted[ii], ii + 1) << "neighbor slot " << ii; + } + } +} + TEST_F(TestEncodingDecodingNborInfo, valid_nbor_info_gpu) { int *valid_type_dev = NULL, *valid_index_dev = NULL, *out_type_dev = NULL, *out_index_dev = NULL; From 3533e6db4d56ef1ffac6a6bf6d14583ccc63a6ce Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 12:42:49 +0800 Subject: [PATCH 2/2] test(gpu): cover final neighbor type boundary Place the second type bucket at the last occupied sort slot so the exact-capacity regression requires both neighbor iteration and output formatting to process the final key. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/tests/test_fmt_nlist.cc | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/source/lib/tests/test_fmt_nlist.cc b/source/lib/tests/test_fmt_nlist.cc index d2dfc8c365..89f671ef5e 100644 --- a/source/lib/tests/test_fmt_nlist.cc +++ b/source/lib/tests/test_fmt_nlist.cc @@ -444,7 +444,7 @@ TEST(FormatNlistGpu, preserves_exact_capacity_rows) { const int nloc = 1; const int nall = max_nbor_size + 1; const float rcut = 2.0f; - const std::vector sec = {0, max_nbor_size}; + const std::vector sec = {0, max_nbor_size - 1, max_nbor_size}; std::vector coord(static_cast(nall) * 3, 0.0); std::vector type(nall, 0); std::vector neighbors(max_nbor_size); @@ -452,6 +452,9 @@ TEST(FormatNlistGpu, preserves_exact_capacity_rows) { neighbors[ii] = ii + 1; coord[static_cast(ii + 1) * 3] = 1.0; } + // Put the second type boundary in the final occupied slot so the test also + // requires fill_nei_iter to process the last valid key. + type[max_nbor_size] = 1; std::vector ilist = {0}; std::vector numneigh = {max_nbor_size};