fix(rocm): keep tabulation barriers uniform

[njzjz-bot]

Root cause

The SE-T and SE-R gradient kernels distributed neighbors by wavefront with loop bounds derived from warp_idx, while GpuSyncThreads() remained inside those loops. On ROCm that helper is a block-wide barrier, so neighbor counts below four or not divisible by four let only a subset of wavefronts reach the final barrier. Participating wavefronts then wait forever.

CUDA does not deadlock in the same control flow because the helper maps to a warp-scoped synchronization there.

Introduction and history

• The divergent SE-T loop came from 1cf3e50b1, PR Model compression for se-t descriptor #1225, in 2021.
• The equivalent SE-R loop came from 05c785b18, PR Finish model compression for se_r descriptor! #1361, in 2022.
• The latent divergence became a ROCm block-barrier deadlock hazard in 0f07afa60, PR merge cuda and rocm files #2844, in 2023, when the CUDA and ROCm implementations were merged and the ROCm helper used __syncthreads().

Thus the loop structure was old, while the CUDA/ROCm synchronization unification made the failure mode observable on AMD GPUs.

Test gap

The existing GPU tests used nnei == 4, exactly the kernel tile count. All four wavefronts therefore entered the same number of iterations, hiding both the short final tile and the partial-wavefront cases.

Fix

Restructure both kernels around uniform tile loops. Every wavefront now executes every tile-level barrier, while an active predicate limits reads, reductions, and writes to valid neighbors.

Validation

• CUDA 12.4 compilation of the modified GPU kernels;
• CPU/GPU numerical comparisons for neighbor counts 1, 3, and 5 in both SE-T and SE-R;
• tests run on an NVIDIA GeForce RTX 5090 allocation;
• static control-flow verification that every ROCm wavefront reaches every block barrier;
• ruff format .;
• ruff check ..

No AMD GPU was available for a runtime ROCm test.

Fixes #5893

Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh

Summary by CodeRabbit

• Bug Fixes

  • Improved GPU gradient calculations for workloads with partial neighbor tiles.
  • Ensured consistent synchronization across supported GPU platforms.

• Tests

  • Added coverage for GPU gradient calculations across multiple partial neighbor configurations.
  • Verified GPU results against CPU references with high-precision comparisons.

[coderabbitai]

📝 Walkthrough

Walkthrough

SE-T and SE-R GPU gradient kernels now process neighbors through uniform KTILE iterations with active predicates around partial tiles and reductions. New GPU tests validate counts of 1, 3, and 5 against CPU references.

Changes

GPU gradient synchronization

Layer / File(s)	Summary
Uniform gradient tile iteration source/lib/src/gpu/tabulate.cu	SE-T and SE-R gradient kernels use uniform tile loops, synchronize every tile, and guard partial-tile accumulation and writes with active-lane predicates.
Partial-tile gradient tests source/lib/tests/test_tabulate_se_r.cc, source/lib/tests/test_tabulate_se_t.cc	GPU tests compare SE-T and SE-R results with CPU references for neighbor counts 1, 3, and 5, including device allocation and cleanup.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• deepmodeling/deepmd-kit#5601: Modifies the same SE gradient kernels in tabulate.cu.

Suggested labels: bug, Core, ROCM

Suggested reviewers: njzjz

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title is concise and accurately describes the main ROCm barrier-synchronization fix.
Linked Issues check	✅ Passed	The kernels now use uniform tile loops with barriers reached by all threads, and regression tests cover neighbor counts 1, 3, and 5.
Out of Scope Changes check	✅ Passed	The PR stays focused on the ROCm deadlock fix and matching regression tests; no unrelated changes are indicated.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

source/lib/tests/test_tabulate_se_r.cc (1)

658-695: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Both new regression tests use the same narrow neighbor-count set. They cover only partial first tiles with a single block, leaving the exact-tile and multi-tile paths (and per-block indexing) unexercised.

• source/lib/tests/test_tabulate_se_r.cc#L658-L695: extend neighbor_counts with 4 and 8, and use test_nloc > 1.
• source/lib/tests/test_tabulate_se_t.cc#L5327-L5374: extend neighbor_counts with 4 and 8, and add a case with test_nnei_i > 1.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@source/lib/tests/test_tabulate_se_r.cc` around lines 658 - 695, Broaden both
regression tests to cover exact-tile, multi-tile, and multi-block indexing: in
source/lib/tests/test_tabulate_se_r.cc lines 658-695, extend neighbor_counts
with 4 and 8 and set test_nloc greater than 1; in
source/lib/tests/test_tabulate_se_t.cc lines 5327-5374, extend neighbor_counts
with 4 and 8 and add a case where test_nnei_i is greater than 1. Update the
existing test setup and expected-value calculations consistently in each test.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@source/lib/tests/test_tabulate_se_r.cc`:
- Around line 658-695: Broaden both regression tests to cover exact-tile,
multi-tile, and multi-block indexing: in source/lib/tests/test_tabulate_se_r.cc
lines 658-695, extend neighbor_counts with 4 and 8 and set test_nloc greater
than 1; in source/lib/tests/test_tabulate_se_t.cc lines 5327-5374, extend
neighbor_counts with 4 and 8 and add a case where test_nnei_i is greater than 1.
Update the existing test setup and expected-value calculations consistently in
each test.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 46660d7c-4cdf-4bac-ba2f-4b144b7a601a

📥 Commits

Reviewing files that changed from the base of the PR and between 4f827cc and 70cec62.

📒 Files selected for processing (3)

• source/lib/src/gpu/tabulate.cu
• source/lib/tests/test_tabulate_se_r.cc
• source/lib/tests/test_tabulate_se_t.cc

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.96%. Comparing base (4f827cc) to head (70cec62).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5933      +/-   ##
==========================================
- Coverage   79.21%   78.96%   -0.25%     
==========================================
  Files        1069     1069              
  Lines      124070   124070              
  Branches     4522     4527       +5     
==========================================
- Hits        98278    97969     -309     
- Misses      24171    24482     +311     
+ Partials     1621     1619       -2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[wanghan-iapcm]

The kernel restructure is correct and, for this file, complete. I checked the parts that could have broken: warp_idx comes from GpuShuffleSync(0xffffffff, thread_idx / WARP_SIZE, 0), so it is warp-uniform and therefore active is uniform across a warp — warp_reduce's full-mask shuffle still gets full-warp participation, which is the thing a predicated reduction usually gets wrong. sum/Csub are declared and zeroed before the guard, the block dim is KK * WARP_SIZE so tile += KTILE covers every neighbor exactly once with no duplication, and the shared iteratorA is written once before the loop behind a real __syncthreads() and only read inside, so moving the in-loop barrier does not affect it. GpuSyncThreads() has exactly three call sites in the file — se_a grad (already converted in #5844), se_t grad and se_r grad — and after this PR all three use the uniform tile form, so no instance of the pattern is left behind.

On the test gap, which is the one thing I want on the record. The existing GPU tests used nnei == 4, exactly KTILE, so all four wavefronts ran identical trip counts and the divergence could not appear; the untested cell is ROCm build x nnei % KTILE != 0, and it needs both coordinates. The new grad_gpu_partial_neighbor_tiles tests cross the neighbor-count half but not the platform half: GpuSyncThreads is __syncwarp() under GOOGLE_CUDA, so the pre-fix kernels were already numerically correct for nnei = 1/3/5 and these tests pass identically with and without the kernel change; and the ROCm job in build_cc.yml installs rocm-dev hipcub-dev, builds, and only checks that the test files exist — it never runs runUnitTests_lib. So nothing in CI goes red if this fix is reverted. That is a genuine gap rather than a blocker: it is inherent to the repo having no AMD GPU runner, and #5844 shipped its SE-A counterpart under the same constraint. The tests are still worth having as CPU/GPU parity coverage. I would only ask that the "Test gap" section not read as though they close it — they narrow it on the neighbor-count axis and leave the platform axis open.

Two corrections to the history section, neither affecting the code. For SE-T the hazard is older than stated: git show 1cf3e50b1 -- source/lib/src/rocm/tabulate.hip.cu (#1225, 2021) added the ROCm SE-T grad kernel with the divergent for (jj = warp_idx; ...) loop and __syncthreads() inside it together, so the block-scoped barrier was there from the start; the 2023 CUDA/ROCm unification did not introduce it. The SE-R half of the claim does hold, and its real history is more useful than the PR says — at 0f07afa60^:source/lib/src/rocm/tabulate.hip.cu the ROCm SE-R grad kernel used a uniform for (ii = 0; ii < nnei; ii += KTILE) loop indexing ii + warp_idx with no bounds check, i.e. this PR's exact structure minus the guard, which is an out-of-bounds read and write whenever nnei % 4 != 0. #2844 traded that for the divergent form. Worth stating, since it is precisely why the old uniform loop cannot simply be restored and why the active predicate is load-bearing rather than cosmetic.

Minor, for whoever merges: this and #5931 both append a new TEST_F at the same anchor in source/lib/tests/test_tabulate_se_t.cc, so the second one in needs a rebase. The tabulate.cu hunks are far enough apart to merge cleanly.