diff --git a/source/lib/src/gpu/tabulate.cu b/source/lib/src/gpu/tabulate.cu index 6bab737739..0ab18aca3b 100644 --- a/source/lib/src/gpu/tabulate.cu +++ b/source/lib/src/gpu/tabulate.cu @@ -730,13 +730,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 aa1f2ee2f4..5b12fcbdcd 100644 --- a/source/lib/tests/test_tabulate_se_t.cc +++ b/source/lib/tests/test_tabulate_se_t.cc @@ -5324,6 +5324,58 @@ TEST_F(TestTabulateSeT, tabulate_fusion_se_a_grad_gpu) { } } +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}; + 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); + 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); +} + TEST_F(TestTabulateSeT, grad_gpu_partial_neighbor_tiles) { constexpr int test_nloc = 1; constexpr int test_nnei_i = 1;