From 70cec62fc0896d33b554a5e9a7a27856ccefa658 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 10:40:43 +0800 Subject: [PATCH] fix(rocm): keep tabulation barriers uniform Make every wavefront traverse the same SE-T and SE-R gradient tiles so ROCm block barriers cannot diverge. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/src/gpu/tabulate.cu | 87 +++++++++++++++----------- source/lib/tests/test_tabulate_se_r.cc | 40 ++++++++++++ source/lib/tests/test_tabulate_se_t.cc | 50 +++++++++++++++ 3 files changed, 140 insertions(+), 37 deletions(-) diff --git a/source/lib/src/gpu/tabulate.cu b/source/lib/src/gpu/tabulate.cu index d978b5643c..aa4203f0d9 100644 --- a/source/lib/src/gpu/tabulate.cu +++ b/source/lib/src/gpu/tabulate.cu @@ -670,30 +670,38 @@ __global__ void tabulate_fusion_se_t_grad_fifth_order_polynomial( __syncthreads(); for (int ii = 0; ii < nnei_i; ii++) { - for (int jj = warp_idx; jj < nnei_j; jj += KTILE) { - FPTYPE xx = em_x[block_idx * nnei_i * nnei_j + ii * nnei_j + jj]; - FPTYPE tmp = xx; - int table_idx = 0; - FPTYPE extrapolate_delta = (FPTYPE)0.; - locate_xx_se_t(xx, table_idx, lower, upper, -max, max, stride0, stride1, - extrapolate_delta); + // GpuSyncThreads is block-wide on ROCm, so every wavefront must execute + // the same number of tile iterations even when the last tile is partial. + for (int tile = 0; tile < nnei_j; tile += KTILE) { + const int jj = tile + warp_idx; + const bool active = jj < nnei_j; FPTYPE sum = (FPTYPE)0.; FPTYPE Csub = (FPTYPE)0.; - for (int kk = lane_idx; kk < last_layer_size; kk += WARP_SIZE) { - FPTYPE var[6]; - load_polynomial_params(var, table, table_idx, kk, last_layer_size); - FPTYPE res_grad = polynomial5_grad(var, xx); - FPTYPE res = polynomial5(var, xx) + res_grad * extrapolate_delta; - - sum += iteratorA[kk] * res; - Csub += iteratorA[kk] * tmp * res_grad; + if (active) { + FPTYPE xx = em_x[block_idx * nnei_i * nnei_j + ii * nnei_j + jj]; + FPTYPE tmp = xx; + int table_idx = 0; + FPTYPE extrapolate_delta = (FPTYPE)0.; + locate_xx_se_t(xx, table_idx, lower, upper, -max, max, stride0, stride1, + extrapolate_delta); + for (int kk = lane_idx; kk < last_layer_size; kk += WARP_SIZE) { + FPTYPE var[6]; + load_polynomial_params(var, table, table_idx, kk, last_layer_size); + FPTYPE res_grad = polynomial5_grad(var, xx); + FPTYPE res = polynomial5(var, xx) + res_grad * extrapolate_delta; + + sum += iteratorA[kk] * res; + Csub += iteratorA[kk] * tmp * res_grad; + } } GpuSyncThreads(); - warp_reduce(sum); - warp_reduce(Csub); - if (lane_idx == 0) { - dy_dem[block_idx * nnei_i * nnei_j + ii * nnei_j + jj] = sum; - dy_dem_x[block_idx * nnei_i * nnei_j + ii * nnei_j + jj] = Csub; + if (active) { + warp_reduce(sum); + warp_reduce(Csub); + if (lane_idx == 0) { + dy_dem[block_idx * nnei_i * nnei_j + ii * nnei_j + jj] = sum; + dy_dem_x[block_idx * nnei_i * nnei_j + ii * nnei_j + jj] = Csub; + } } } } @@ -975,27 +983,32 @@ __global__ void tabulate_fusion_se_r_grad_fifth_order_polynomial( int warp_idx = GpuShuffleSync(0xffffffff, thread_idx / WARP_SIZE, 0); int lane_idx = thread_idx % WARP_SIZE; __syncthreads(); - for (int ii = warp_idx; ii < nnei; ii += KTILE) { - FPTYPE xx = em[block_idx * nnei + ii]; - - int table_idx = 0; + // Keep all wavefronts on uniform control flow around the ROCm block barrier. + for (int tile = 0; tile < nnei; tile += KTILE) { + const int ii = tile + warp_idx; + const bool active = ii < nnei; FPTYPE Csub = (FPTYPE)0.; - FPTYPE extrapolate_delta = (FPTYPE)0.; - locate_xx_se_r(xx, table_idx, lower, upper, max, stride0, stride1, - extrapolate_delta); + if (active) { + FPTYPE xx = em[block_idx * nnei + ii]; + int table_idx = 0; + FPTYPE extrapolate_delta = (FPTYPE)0.; + locate_xx_se_r(xx, table_idx, lower, upper, max, stride0, stride1, + extrapolate_delta); - FPTYPE var[6]; - for (int jj = lane_idx; jj < last_layer_size; jj += WARP_SIZE) { - load_polynomial_params(var, table, table_idx, jj, last_layer_size); - Csub += - polynomial5_grad(var, xx) * - dy[block_idx * nnei * last_layer_size + ii * last_layer_size + jj]; + FPTYPE var[6]; + for (int jj = lane_idx; jj < last_layer_size; jj += WARP_SIZE) { + load_polynomial_params(var, table, table_idx, jj, last_layer_size); + Csub += + polynomial5_grad(var, xx) * + dy[block_idx * nnei * last_layer_size + ii * last_layer_size + jj]; + } } GpuSyncThreads(); - - warp_reduce(Csub); - if (lane_idx == 0) { - dy_dem[block_idx * nnei + ii] = Csub; + if (active) { + warp_reduce(Csub); + if (lane_idx == 0) { + dy_dem[block_idx * nnei + ii] = Csub; + } } } } diff --git a/source/lib/tests/test_tabulate_se_r.cc b/source/lib/tests/test_tabulate_se_r.cc index 8ac7e13c96..712efc1cbe 100644 --- a/source/lib/tests/test_tabulate_se_r.cc +++ b/source/lib/tests/test_tabulate_se_r.cc @@ -1,6 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include +#include #include #include @@ -653,4 +654,43 @@ TEST_F(TestTabulateSeR, tabulate_fusion_se_r_grad_gpu) { EXPECT_LT(fabs(dy_dem[jj] - expected_dy_dem[jj]), 1e-5); } } + +TEST_F(TestTabulateSeR, grad_gpu_partial_neighbor_tiles) { + constexpr int test_nloc = 1; + const std::array neighbor_counts = {1, 3, 5}; + + for (const int test_nnei : neighbor_counts) { + std::vector test_em(test_nnei); + for (int ii = 0; ii < test_nnei; ++ii) { + test_em[ii] = 0.01 * (ii + 1); + } + std::vector test_dy(test_nnei * last_layer_size, 1.0); + std::vector expected_dy_dem(test_nnei); + deepmd::tabulate_fusion_se_r_grad_cpu( + expected_dy_dem.data(), table.data(), info.data(), test_em.data(), + test_dy.data(), test_nloc, test_nnei, last_layer_size); + + std::vector actual_dy_dem(test_nnei); + double *dy_dem_dev = nullptr, *table_dev = nullptr, *em_dev = nullptr, + *dy_dev = nullptr; + deepmd::malloc_device_memory_sync(dy_dem_dev, actual_dy_dem); + deepmd::malloc_device_memory_sync(table_dev, table); + deepmd::malloc_device_memory_sync(em_dev, test_em); + deepmd::malloc_device_memory_sync(dy_dev, test_dy); + + deepmd::tabulate_fusion_se_r_grad_gpu( + dy_dem_dev, table_dev, info.data(), em_dev, dy_dev, test_nloc, + test_nnei, last_layer_size); + deepmd::memcpy_device_to_host(dy_dem_dev, actual_dy_dem); + + for (int ii = 0; ii < test_nnei; ++ii) { + EXPECT_NEAR(actual_dy_dem[ii], expected_dy_dem[ii], 1e-10); + } + + deepmd::delete_device_memory(dy_dem_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(dy_dev); + } +} #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/source/lib/tests/test_tabulate_se_t.cc b/source/lib/tests/test_tabulate_se_t.cc index be82a07ba6..aa1f2ee2f4 100644 --- a/source/lib/tests/test_tabulate_se_t.cc +++ b/source/lib/tests/test_tabulate_se_t.cc @@ -1,6 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include +#include #include #include @@ -5322,4 +5323,53 @@ 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_gpu_partial_neighbor_tiles) { + constexpr int test_nloc = 1; + constexpr int test_nnei_i = 1; + const std::array neighbor_counts = {1, 3, 5}; + std::vector test_dy(last_layer_size, 1.0); + + for (const int test_nnei_j : neighbor_counts) { + std::vector test_em_x(test_nnei_j); + for (int ii = 0; ii < test_nnei_j; ++ii) { + test_em_x[ii] = 0.01 * (ii + 1); + } + std::vector expected_dy_dem_x(test_nnei_j); + std::vector expected_dy_dem(test_nnei_j); + deepmd::tabulate_fusion_se_t_grad_cpu( + expected_dy_dem_x.data(), expected_dy_dem.data(), table.data(), + info.data(), test_em_x.data(), test_em_x.data(), test_dy.data(), + test_nloc, test_nnei_i, test_nnei_j, last_layer_size); + + std::vector actual_dy_dem_x(test_nnei_j); + std::vector actual_dy_dem(test_nnei_j); + double *dy_dem_x_dev = nullptr, *dy_dem_dev = nullptr, *table_dev = nullptr, + *em_x_dev = nullptr, *em_dev = nullptr, *dy_dev = nullptr; + deepmd::malloc_device_memory_sync(dy_dem_x_dev, actual_dy_dem_x); + deepmd::malloc_device_memory_sync(dy_dem_dev, actual_dy_dem); + 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(dy_dev, test_dy); + + deepmd::tabulate_fusion_se_t_grad_gpu( + dy_dem_x_dev, dy_dem_dev, table_dev, info.data(), em_x_dev, em_dev, + dy_dev, test_nloc, test_nnei_i, test_nnei_j, last_layer_size); + deepmd::memcpy_device_to_host(dy_dem_x_dev, actual_dy_dem_x); + deepmd::memcpy_device_to_host(dy_dem_dev, actual_dy_dem); + + for (int ii = 0; ii < test_nnei_j; ++ii) { + EXPECT_NEAR(actual_dy_dem_x[ii], expected_dy_dem_x[ii], 1e-10); + EXPECT_NEAR(actual_dy_dem[ii], expected_dy_dem[ii], 1e-10); + } + + deepmd::delete_device_memory(dy_dem_x_dev); + deepmd::delete_device_memory(dy_dem_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(dy_dev); + } +} #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM