From 93fc6dacbb2e35c88fa9723a1c454ce882f508e0 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 10:40:43 +0800 Subject: [PATCH 1/2] fix(tabulate): preserve SE-T polynomial cache Keep cached polynomial coefficients alive for the same lifetime as their cached table index in SE-T grad-grad. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/src/gpu/tabulate.cu | 5 ++- source/lib/tests/test_tabulate_se_t.cc | 43 ++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/source/lib/src/gpu/tabulate.cu b/source/lib/src/gpu/tabulate.cu index d978b5643c..d106cbeeab 100644 --- a/source/lib/src/gpu/tabulate.cu +++ b/source/lib/src/gpu/tabulate.cu @@ -721,13 +721,16 @@ __global__ void tabulate_fusion_se_t_grad_grad_fifth_order_polynomial( FPTYPE sum = (FPTYPE)0.; for (int ii = 0; ii < nnei_i; ii++) { int mark_table_idx = -1; + // The cached table index and coefficients must have the same lifetime. + // Keeping var outside the neighbor loop makes a cache hit reuse initialized + // coefficients instead of a newly scoped, uninitialized local array. + FPTYPE var[6]; for (int jj = 0; jj < nnei_j; jj++) { FPTYPE xx = em_x[block_idx * nnei_i * nnei_j + ii * nnei_j + jj]; FPTYPE tmp = xx; FPTYPE dz_xx = dz_dy_dem_x[block_idx * nnei_i * nnei_j + ii * nnei_j + jj]; FPTYPE dz_em = dz_dy_dem[block_idx * nnei_i * nnei_j + ii * nnei_j + jj]; - FPTYPE var[6]; int table_idx = 0; FPTYPE extrapolate_delta = (FPTYPE)0.; diff --git a/source/lib/tests/test_tabulate_se_t.cc b/source/lib/tests/test_tabulate_se_t.cc index be82a07ba6..7b15fb8fa6 100644 --- a/source/lib/tests/test_tabulate_se_t.cc +++ b/source/lib/tests/test_tabulate_se_t.cc @@ -5322,4 +5322,47 @@ TEST_F(TestTabulateSeT, tabulate_fusion_se_a_grad_gpu) { EXPECT_LT(fabs(dy_dem[jj] - expected_dy_dem[jj]), 1e-5); } } + +TEST_F(TestTabulateSeT, grad_grad_gpu_reuses_same_table_interval) { + constexpr int test_nloc = 1; + constexpr int test_nnei_i = 1; + constexpr int test_nnei_j = 3; + // All inputs lie in the same stride-0 interval. The GPU cache therefore + // must retain the coefficients loaded for the first neighbor. + std::vector test_em_x = {0.011, 0.012, 0.019}; + std::vector test_dz_dy_dem_x = {0.5, -0.25, 0.75}; + std::vector test_dz_dy_dem = {-0.2, 0.4, 0.1}; + std::vector expected(last_layer_size); + deepmd::tabulate_fusion_se_t_grad_grad_cpu( + expected.data(), table.data(), info.data(), test_em_x.data(), + test_em_x.data(), test_dz_dy_dem_x.data(), test_dz_dy_dem.data(), + test_nloc, test_nnei_i, test_nnei_j, last_layer_size); + + std::vector actual(last_layer_size); + double *actual_dev = nullptr, *table_dev = nullptr, *em_x_dev = nullptr, + *em_dev = nullptr, *dz_dy_dem_x_dev = nullptr, + *dz_dy_dem_dev = nullptr; + deepmd::malloc_device_memory_sync(actual_dev, actual); + deepmd::malloc_device_memory_sync(table_dev, table); + deepmd::malloc_device_memory_sync(em_x_dev, test_em_x); + deepmd::malloc_device_memory_sync(em_dev, test_em_x); + deepmd::malloc_device_memory_sync(dz_dy_dem_x_dev, test_dz_dy_dem_x); + deepmd::malloc_device_memory_sync(dz_dy_dem_dev, test_dz_dy_dem); + + deepmd::tabulate_fusion_se_t_grad_grad_gpu( + actual_dev, table_dev, info.data(), em_x_dev, em_dev, dz_dy_dem_x_dev, + dz_dy_dem_dev, test_nloc, test_nnei_i, test_nnei_j, last_layer_size); + deepmd::memcpy_device_to_host(actual_dev, actual); + + for (int ii = 0; ii < last_layer_size; ++ii) { + EXPECT_NEAR(actual[ii], expected[ii], 1e-10); + } + + deepmd::delete_device_memory(actual_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(dz_dy_dem_x_dev); + deepmd::delete_device_memory(dz_dy_dem_dev); +} #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM From 1e4261149a32b7398d706f0b50ce2c7ee7024ae1 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 31 Jul 2026 13:15:30 +0800 Subject: [PATCH 2/2] test(tabulate): assert SE-T cache interval Pin the repeated-interval precondition in the CUDA regression test so fixture metadata changes cannot silently stop exercising coefficient reuse. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/tests/test_tabulate_se_t.cc | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/source/lib/tests/test_tabulate_se_t.cc b/source/lib/tests/test_tabulate_se_t.cc index 7b15fb8fa6..a60126ee44 100644 --- a/source/lib/tests/test_tabulate_se_t.cc +++ b/source/lib/tests/test_tabulate_se_t.cc @@ -5330,6 +5330,15 @@ TEST_F(TestTabulateSeT, grad_grad_gpu_reuses_same_table_interval) { // All inputs lie in the same stride-0 interval. The GPU cache therefore // must retain the coefficients loaded for the first neighbor. std::vector test_em_x = {0.011, 0.012, 0.019}; + const auto stride0_interval = [this](const double value) { + EXPECT_GE(value, info[0]); + EXPECT_LT(value, info[1]); + return static_cast((value - info[0]) / info[3]); + }; + const int expected_interval = stride0_interval(test_em_x.front()); + for (const double value : test_em_x) { + EXPECT_EQ(stride0_interval(value), expected_interval); + } std::vector test_dz_dy_dem_x = {0.5, -0.25, 0.75}; std::vector test_dz_dy_dem = {-0.2, 0.4, 0.1}; std::vector expected(last_layer_size);