fix ROCm tabulate_fusion_se_a_grad_fifth_order_polynomial (#2357)

Follow #830 to fix #2303. Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
deepmodeling · Mar 1, 2023 · 719005c · 719005c
1 parent a16ac9f
commit 719005c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/source/lib/src/rocm/tabulate.hip.cu b/source/lib/src/rocm/tabulate.hip.cu
@@ -162,7 +162,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial(
   }
   __syncthreads();
   FPTYPE ago = __shfl(em_x[block_idx * nnei + nnei - 1], 0);
-  for (int ii = 0; ii < nnei; ii += KTILE) {
+  for (int ii = 0; ii < nnei - warp_idx; ii += KTILE) {
     FPTYPE xx = em_x[block_idx * nnei + ii + warp_idx];
     if (ago == xx) {
       unloop = true;