fix(rocm): keep tabulation barriers uniform

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;
+      }
     }
   }
 }

source/lib/tests/test_tabulate_se_r.cc

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 #include <gtest/gtest.h>
 
+#include <array>
 #include <iostream>
 #include <vector>
 
@@ -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<int, 3> neighbor_counts = {1, 3, 5};
+
+  for (const int test_nnei : neighbor_counts) {
+    std::vector<double> test_em(test_nnei);
+    for (int ii = 0; ii < test_nnei; ++ii) {
+      test_em[ii] = 0.01 * (ii + 1);
+    }
+    std::vector<double> test_dy(test_nnei * last_layer_size, 1.0);
+    std::vector<double> expected_dy_dem(test_nnei);
+    deepmd::tabulate_fusion_se_r_grad_cpu<double>(
+        expected_dy_dem.data(), table.data(), info.data(), test_em.data(),
+        test_dy.data(), test_nloc, test_nnei, last_layer_size);
+
+    std::vector<double> 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<double>(
+        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

source/lib/tests/test_tabulate_se_t.cc

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 #include <gtest/gtest.h>
 
+#include <array>
 #include <iostream>
 #include <vector>
 
@@ -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<int, 3> neighbor_counts = {1, 3, 5};
+  std::vector<double> test_dy(last_layer_size, 1.0);
+
+  for (const int test_nnei_j : neighbor_counts) {
+    std::vector<double> 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<double> expected_dy_dem_x(test_nnei_j);
+    std::vector<double> expected_dy_dem(test_nnei_j);
+    deepmd::tabulate_fusion_se_t_grad_cpu<double>(
+        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<double> actual_dy_dem_x(test_nnei_j);
+    std::vector<double> 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<double>(
+        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