diff --git a/source/lib/src/gpu/tabulate.cu b/source/lib/src/gpu/tabulate.cu index a3d9cc0e3a..d978b5643c 100644 --- a/source/lib/src/gpu/tabulate.cu +++ b/source/lib/src/gpu/tabulate.cu @@ -398,8 +398,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( const int thread_idx = threadIdx.x; // KTILE * WARP_SIZE, usually 128 here~ int warp_idx = GpuShuffleSync(0xffffffff, threadIdx.x / WARP_SIZE, 0); int lane_idx = threadIdx.x % WARP_SIZE; - int breakpoint = nnei - 1; - bool unloop = false; + __shared__ int breakpoint; FPTYPE* iteratorA = (FPTYPE*)&_data[0]; // dy for (int ii = 0; ii < MTILE; ii++) { for (int jj = thread_idx; jj < last_layer_size; jj += blockDim.x) { @@ -408,71 +407,93 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( } } __syncthreads(); - FPTYPE ago = GpuShuffleSync(0xffffffff, em_x[block_idx * nnei + nnei - 1], 0); - for (int ii = warp_idx; ii < nnei; ii += KTILE) { - FPTYPE xx = em_x[block_idx * nnei + ii]; - if (ago == xx && em[block_idx * nnei * 4 + ii * 4 + 1] == 0. && - em[block_idx * nnei * 4 + ii * 4 + 2] == 0. && - em[block_idx * nnei * 4 + ii * 4 + 3] == 0. && is_sorted) { - unloop = true; - breakpoint = ii; + + // Sorted padding must be folded at the first sentinel for the whole atom, + // exactly as in the sequential CPU implementation. A warp-local search can + // select several later sentinels because neighbor indices are striped over + // KTILE warps, producing duplicate tail contributions and overlapping + // dy_dtwo writes. + if (thread_idx == 0) { + breakpoint = nnei; // nnei means that no padding sentinel was found. + if (is_sorted) { + const FPTYPE ago = em_x[block_idx * nnei + nnei - 1]; + for (int ii = 0; ii < nnei; ++ii) { + if (ago == em_x[block_idx * nnei + ii] && + em[block_idx * nnei * 4 + ii * 4 + 1] == 0. && + em[block_idx * nnei * 4 + ii * 4 + 2] == 0. && + em[block_idx * nnei * 4 + ii * 4 + 3] == 0.) { + breakpoint = ii; + break; + } + } } + } + __syncthreads(); - int table_idx = 0; - FPTYPE reg_em[MTILE] = {em[block_idx * nnei * MTILE + ii * 4 + 0], - em[block_idx * nnei * MTILE + ii * 4 + 1], - em[block_idx * nnei * MTILE + ii * 4 + 2], - em[block_idx * nnei * MTILE + ii * 4 + 3]}; + // Keep the tile loop uniform across the block. GpuSyncThreads is a warp + // barrier on CUDA but a block barrier on ROCm, so warp-specific early exits + // would deadlock HIP when the shared breakpoint falls inside a tile. + for (int tile = 0; tile < nnei && tile <= breakpoint; tile += KTILE) { + const int ii = tile + warp_idx; + const bool active = ii < nnei && ii <= breakpoint; FPTYPE Csub = (FPTYPE)0.; FPTYPE sum[MTILE] = {(FPTYPE)0.}; - FPTYPE extrapolate_delta = (FPTYPE)0.; - locate_xx_se_a(xx, table_idx, lower, upper, max, stride0, stride1, - extrapolate_delta); + if (active) { + const int repeat_count = ii == breakpoint ? nnei - breakpoint : 1; + FPTYPE xx = em_x[block_idx * nnei + ii]; + int table_idx = 0; + FPTYPE reg_em[MTILE] = {em[block_idx * nnei * MTILE + ii * MTILE + 0], + em[block_idx * nnei * MTILE + ii * MTILE + 1], + em[block_idx * nnei * MTILE + ii * MTILE + 2], + em[block_idx * nnei * MTILE + ii * MTILE + 3]}; + FPTYPE extrapolate_delta = (FPTYPE)0.; + locate_xx_se_a(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); - FPTYPE res_grad = polynomial5_grad(var, xx); - FPTYPE res = polynomial5(var, xx) + res_grad * extrapolate_delta; - FPTYPE oldres = res; - FPTYPE t; - if (enable_se_atten) { - t = two_embed[block_idx * nnei * last_layer_size + - ii * last_layer_size + jj]; - res = res * t + res; - } + 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); + FPTYPE res_grad = polynomial5_grad(var, xx); + FPTYPE res = polynomial5(var, xx) + res_grad * extrapolate_delta; + FPTYPE oldres = res; + FPTYPE t; + if (enable_se_atten) { + t = two_embed[block_idx * nnei * last_layer_size + + ii * last_layer_size + jj]; + res = res * t + res; + } - for (int kk = 0; kk < MTILE; kk++) { - sum[kk] += - (nnei - breakpoint) * iteratorA[kk * last_layer_size + jj] * res; - } - res = reg_em[0] * iteratorA[0 * last_layer_size + jj]; - res += reg_em[1] * iteratorA[1 * last_layer_size + jj]; - res += reg_em[2] * iteratorA[2 * last_layer_size + jj]; - res += reg_em[3] * iteratorA[3 * last_layer_size + jj]; - Csub += (nnei - breakpoint) * res_grad * - (enable_se_atten ? res * t + res : res); - if (enable_se_atten) { - // from ii to ii + (nnei - breakpoint) - for (int ii2 = ii; ii2 < ii + nnei - breakpoint; ii2++) { - dy_dtwo[block_idx * nnei * last_layer_size + ii2 * last_layer_size + - jj] = oldres * res; + for (int kk = 0; kk < MTILE; kk++) { + sum[kk] += repeat_count * iteratorA[kk * last_layer_size + jj] * res; + } + res = reg_em[0] * iteratorA[0 * last_layer_size + jj]; + res += reg_em[1] * iteratorA[1 * last_layer_size + jj]; + res += reg_em[2] * iteratorA[2 * last_layer_size + jj]; + res += reg_em[3] * iteratorA[3 * last_layer_size + jj]; + Csub += + repeat_count * res_grad * (enable_se_atten ? res * t + res : res); + if (enable_se_atten) { + // A real neighbor owns one entry; the first sentinel owns the full + // padding tail. No other warp writes these dy_dtwo positions. + for (int ii2 = ii; ii2 < ii + repeat_count; ii2++) { + dy_dtwo[block_idx * nnei * last_layer_size + ii2 * last_layer_size + + jj] = oldres * res; + } } } } GpuSyncThreads(); - for (int kk = 0; kk < MTILE; kk++) { - warp_reduce(sum[kk]); - } - warp_reduce(Csub); - if (lane_idx == 0) { + if (active) { for (int kk = 0; kk < MTILE; kk++) { - dy_dem[block_idx * nnei * MTILE + ii * 4 + kk] = sum[kk]; + warp_reduce(sum[kk]); + } + warp_reduce(Csub); + if (lane_idx == 0) { + for (int kk = 0; kk < MTILE; kk++) { + dy_dem[block_idx * nnei * MTILE + ii * MTILE + kk] = sum[kk]; + } + dy_dem_x[block_idx * nnei + ii] = Csub; } - dy_dem_x[block_idx * nnei + ii] = Csub; - } - if (unloop) { - break; } } } diff --git a/source/lib/tests/test_tabulate_se_a.cc b/source/lib/tests/test_tabulate_se_a.cc index 66a77f41fd..2b5b8faa00 100644 --- a/source/lib/tests/test_tabulate_se_a.cc +++ b/source/lib/tests/test_tabulate_se_a.cc @@ -856,4 +856,110 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu) { deepmd::delete_device_memory(dy_dev); deepmd::delete_device_memory(two_embed_dev); } + +TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_sorted_padding) { + constexpr int test_nloc = 1; + constexpr int test_nnei = 12; + constexpr int padding_begin = 4; + std::vector test_em_x = {0.04, 0.07, 0.11, 0.14}; + test_em_x.resize(test_nnei, 0.19); + std::vector test_em(test_nnei * 4, 0.0); + for (int ii = 0; ii < padding_begin; ++ii) { + test_em[ii * 4 + 0] = test_em_x[ii]; + test_em[ii * 4 + 1] = 0.1 * (ii + 1); + test_em[ii * 4 + 2] = -0.05 * (ii + 1); + test_em[ii * 4 + 3] = 0.025 * (ii + 1); + } + for (int ii = padding_begin; ii < test_nnei; ++ii) { + test_em[ii * 4] = test_em_x[ii]; + } + std::vector test_dy(4 * last_layer_size); + for (int ii = 0; ii < test_dy.size(); ++ii) { + test_dy[ii] = 0.01 * (ii + 1); + } + std::vector test_two_embed(test_nnei * last_layer_size); + for (int ii = 0; ii < test_two_embed.size(); ++ii) { + test_two_embed[ii] = 0.001 * (ii + 1); + } + + auto compare_cpu_gpu = [&](const std::vector* two_embed_host) { + std::vector expected_dy_dem_x(test_nnei); + std::vector expected_dy_dem(test_nnei * 4); + std::vector expected_dy_dtwo(test_nnei * last_layer_size); + deepmd::tabulate_fusion_se_a_grad_cpu( + expected_dy_dem_x.data(), expected_dy_dem.data(), + two_embed_host == nullptr ? nullptr : expected_dy_dtwo.data(), + table.data(), info.data(), test_em_x.data(), test_em.data(), + two_embed_host == nullptr ? nullptr : two_embed_host->data(), + test_dy.data(), test_nloc, test_nnei, last_layer_size, true); + + std::vector actual_dy_dem_x(test_nnei); + std::vector actual_dy_dem(test_nnei * 4); + std::vector actual_dy_dtwo(test_nnei * last_layer_size); + double *dy_dem_x_dev = nullptr, *dy_dem_dev = nullptr, + *dy_dtwo_dev = nullptr, *table_dev = nullptr, *em_x_dev = nullptr, + *em_dev = nullptr, *two_embed_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); + deepmd::malloc_device_memory_sync(dy_dev, test_dy); + if (two_embed_host != nullptr) { + deepmd::malloc_device_memory_sync(dy_dtwo_dev, actual_dy_dtwo); + deepmd::malloc_device_memory_sync(two_embed_dev, *two_embed_host); + } + + deepmd::tabulate_fusion_se_a_grad_gpu( + dy_dem_x_dev, dy_dem_dev, dy_dtwo_dev, table_dev, info.data(), em_x_dev, + em_dev, two_embed_dev, dy_dev, test_nloc, test_nnei, last_layer_size, + true); + deepmd::memcpy_device_to_host(dy_dem_x_dev, actual_dy_dem_x); + deepmd::memcpy_device_to_host(dy_dem_dev, actual_dy_dem); + if (two_embed_host != nullptr) { + deepmd::memcpy_device_to_host(dy_dtwo_dev, actual_dy_dtwo); + } + + for (int ii = 0; ii < actual_dy_dem_x.size(); ++ii) { + EXPECT_NEAR(actual_dy_dem_x[ii], expected_dy_dem_x[ii], 1e-10); + } + for (int ii = 0; ii < actual_dy_dem.size(); ++ii) { + EXPECT_NEAR(actual_dy_dem[ii], expected_dy_dem[ii], 1e-10); + } + if (two_embed_host != nullptr) { + for (int ii = 0; ii < actual_dy_dtwo.size(); ++ii) { + EXPECT_NEAR(actual_dy_dtwo[ii], expected_dy_dtwo[ii], 1e-10); + } + // The CPU contract copies the first sentinel's two-embedding gradient + // across the complete sorted-padding tail. + for (int ii = padding_begin + 1; ii < test_nnei; ++ii) { + for (int jj = 0; jj < last_layer_size; ++jj) { + EXPECT_NEAR(actual_dy_dtwo[ii * last_layer_size + jj], + actual_dy_dtwo[padding_begin * last_layer_size + jj], + 1e-10); + } + } + } + // Later sentinels are padding, not independent neighbors owned by other + // warps, and therefore must retain zero descriptor gradients. + for (int ii = padding_begin + 1; ii < test_nnei; ++ii) { + EXPECT_DOUBLE_EQ(actual_dy_dem_x[ii], 0.0); + for (int jj = 0; jj < 4; ++jj) { + EXPECT_DOUBLE_EQ(actual_dy_dem[ii * 4 + jj], 0.0); + } + } + + deepmd::delete_device_memory(dy_dem_x_dev); + deepmd::delete_device_memory(dy_dem_dev); + deepmd::delete_device_memory(dy_dtwo_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(two_embed_dev); + deepmd::delete_device_memory(dy_dev); + }; + + compare_cpu_gpu(nullptr); + compare_cpu_gpu(&test_two_embed); +} #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM