From 21ef7734127899b3c4d59e781fcdc70ca9dcf679 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 10:40:43 +0800 Subject: [PATCH] fix(tabulate): handle empty SE-A neighbors Treat a zero neighbor capacity as an empty reduction on CPU and GPU, and cover forward and derivative entry points. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/src/gpu/tabulate.cu | 17 +++++++- source/lib/src/tabulate.cc | 17 ++++++++ source/lib/tests/test_tabulate_se_a.cc | 60 ++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 1 deletion(-) diff --git a/source/lib/src/gpu/tabulate.cu b/source/lib/src/gpu/tabulate.cu index d978b5643c..d3bb134b72 100644 --- a/source/lib/src/gpu/tabulate.cu +++ b/source/lib/src/gpu/tabulate.cu @@ -1054,6 +1054,13 @@ void tabulate_fusion_se_a_gpu(FPTYPE* out, if (nloc <= 0) { return; } + if (nnei <= 0) { + // The descriptor does not carry the empty neighbor dimension, so its + // mathematically empty reduction must be materialized explicitly. + DPErrcheck(gpuMemset(out, 0, sizeof(FPTYPE) * nloc * MM * last_layer_size)); + DPErrcheck(gpuDeviceSynchronize()); + return; + } DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); tabulate_fusion_se_a_fifth_order_polynomial @@ -1085,7 +1092,7 @@ void tabulate_fusion_se_a_grad_gpu(FPTYPE* dy_dem_x, const int nnei, const int last_layer_size, const bool is_sorted) { - if (nloc <= 0) { + if (nloc <= 0 || nnei <= 0) { return; } DPErrcheck(gpuGetLastError()); @@ -1119,6 +1126,14 @@ void tabulate_fusion_se_a_grad_grad_gpu(FPTYPE* dz_dy, if (nloc <= 0) { return; } + if (nnei <= 0) { + // Unlike the neighbor-shaped inputs, dz_dy remains non-empty and must be + // initialized to the zero second derivative of an empty reduction. + DPErrcheck( + gpuMemset(dz_dy, 0, sizeof(FPTYPE) * nloc * MM * last_layer_size)); + DPErrcheck(gpuDeviceSynchronize()); + return; + } DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); DPErrcheck(gpuMemset(dz_dy, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size)); diff --git a/source/lib/src/tabulate.cc b/source/lib/src/tabulate.cc index 19fcae13ba..bbfe058143 100644 --- a/source/lib/src/tabulate.cc +++ b/source/lib/src/tabulate.cc @@ -165,6 +165,12 @@ void deepmd::tabulate_fusion_se_a_cpu(FPTYPE* out, const bool is_sorted) { bool enable_se_atten = two_embed != nullptr; memset(out, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size); + // An empty neighbor axis is a valid empty reduction. Return after + // initializing the non-empty descriptor output instead of inspecting the + // nonexistent last neighbor below. + if (nnei <= 0) { + return; + } const FPTYPE lower = table_info[0]; const FPTYPE upper = table_info[1]; const FPTYPE _max = table_info[2]; @@ -246,6 +252,12 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x, const int nnei, const int last_layer_size, const bool is_sorted) { + // Every gradient output has a zero-sized neighbor axis in this case. Avoid + // both zero-length memory operations on potentially null tensor pointers and + // the last-neighbor lookup in the atom loop. + if (nnei <= 0) { + return; + } bool enable_se_atten = two_embed != nullptr; memset(dy_dem_x, 0, sizeof(FPTYPE) * nloc * nnei); memset(dy_dem, 0, sizeof(FPTYPE) * nloc * nnei * 4); @@ -352,6 +364,11 @@ void deepmd::tabulate_fusion_se_a_grad_grad_cpu(FPTYPE* dz_dy, const bool is_sorted) { bool enable_se_atten = two_embed != nullptr; memset(dz_dy, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size); + // The second-order output retains the descriptor shape, so initialize the + // empty reduction to zero before returning. + if (nnei <= 0) { + return; + } const FPTYPE lower = table_info[0]; const FPTYPE upper = table_info[1]; const FPTYPE _max = table_info[2]; diff --git a/source/lib/tests/test_tabulate_se_a.cc b/source/lib/tests/test_tabulate_se_a.cc index 966909a60e..d7b7d30ffc 100644 --- a/source/lib/tests/test_tabulate_se_a.cc +++ b/source/lib/tests/test_tabulate_se_a.cc @@ -1,6 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include +#include #include #include @@ -756,6 +757,32 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_cpu) { } } +TEST_F(TestTabulateSeA, empty_neighbors_cpu) { + constexpr int empty_nnei = 0; + std::vector descriptor(nloc * 4 * last_layer_size, 1.0); + + deepmd::tabulate_fusion_se_a_cpu( + descriptor.data(), table.data(), info.data(), nullptr, nullptr, nullptr, + nloc, empty_nnei, last_layer_size); + for (const double value : descriptor) { + EXPECT_DOUBLE_EQ(value, 0.0); + } + + // First-order outputs all carry the empty neighbor dimension, so null data + // pointers are valid and must not be dereferenced. + deepmd::tabulate_fusion_se_a_grad_cpu( + nullptr, nullptr, nullptr, table.data(), info.data(), nullptr, nullptr, + nullptr, nullptr, nloc, empty_nnei, last_layer_size); + + std::fill(descriptor.begin(), descriptor.end(), 1.0); + deepmd::tabulate_fusion_se_a_grad_grad_cpu( + descriptor.data(), table.data(), info.data(), nullptr, nullptr, nullptr, + nullptr, nullptr, nullptr, nloc, empty_nnei, last_layer_size); + for (const double value : descriptor) { + EXPECT_DOUBLE_EQ(value, 0.0); + } +} + #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM TEST_F(TestTabulateSeA, tabulate_fusion_se_a_gpu) { std::vector xyz_scatter(nloc * nnei * last_layer_size, 0.0); @@ -857,6 +884,39 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu) { deepmd::delete_device_memory(two_embed_dev); } +TEST_F(TestTabulateSeA, empty_neighbors_gpu) { + constexpr int empty_nnei = 0; + std::vector descriptor(nloc * 4 * last_layer_size, 1.0); + double* descriptor_dev = nullptr; + deepmd::malloc_device_memory_sync(descriptor_dev, descriptor); + + deepmd::tabulate_fusion_se_a_gpu(descriptor_dev, nullptr, info.data(), + nullptr, nullptr, nullptr, nloc, + empty_nnei, last_layer_size); + deepmd::memcpy_device_to_host(descriptor_dev, descriptor); + for (const double value : descriptor) { + EXPECT_DOUBLE_EQ(value, 0.0); + } + + // First-order outputs are all empty, and the entry point must return before + // launching a kernel or touching their null pointers. + deepmd::tabulate_fusion_se_a_grad_gpu( + nullptr, nullptr, nullptr, nullptr, info.data(), nullptr, nullptr, + nullptr, nullptr, nloc, empty_nnei, last_layer_size); + + std::fill(descriptor.begin(), descriptor.end(), 1.0); + deepmd::memcpy_host_to_device(descriptor_dev, descriptor); + deepmd::tabulate_fusion_se_a_grad_grad_gpu( + descriptor_dev, nullptr, info.data(), nullptr, nullptr, nullptr, nullptr, + nullptr, nullptr, nloc, empty_nnei, last_layer_size); + deepmd::memcpy_device_to_host(descriptor_dev, descriptor); + for (const double value : descriptor) { + EXPECT_DOUBLE_EQ(value, 0.0); + } + + deepmd::delete_device_memory(descriptor_dev); +} + TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_padding_modes) { constexpr int test_nloc = 1; constexpr int test_nnei = 12;