fix(gpu): preserve exact-capacity neighbor rows - #5937
Conversation
Compare sorted neighbor keys with the explicit padding sentinel so a fully occupied row keeps its final valid entry. Cover every supported GPU sort capacity with a boundary regression test. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughGPU neighbor iteration and formatting now identify empty slots by the ChangesGPU neighbor formatting
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@source/lib/tests/test_fmt_nlist.cc`:
- Around line 447-454: Update the test setup around the neighbor initialization
loop and sec so the final neighbor uses type 1 and the single type bucket
boundary is placed at the last neighbor slot. Ensure the resulting data
exercises the final-type boundary in fill_nei_iter while preserving coverage of
both modified kernels.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 2d5ea603-370f-4e21-b34b-ebf674bff839
📒 Files selected for processing (2)
source/lib/src/gpu/prod_env_mat.cusource/lib/tests/test_fmt_nlist.cc
Place the second type bucket at the last occupied sort slot so the exact-capacity regression requires both neighbor iteration and output formatting to process the final key. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
CI failure is unrelated. |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The fix is right and the reasoning behind it is right. The old predicate inferred "is this slot padding?" from the final sorted key, which is only sound while the row has slack; comparing against the sentinel directly removes the inference rather than patching around it. I checked that the two predicates really are equivalent on partial rows, since that is what makes this safe to change: format_nlist_fill_a leaves a slot untouched on all three skip paths — idy >= nsize, type < 0, and the rr <= rcut guard on the write — so every non-neighbor slot genuinely holds the memset value, and being UINT64_MAX it always sorts last. Nothing the old comparison caught is lost.
The sentinel matches too: gpuMemset(key, 0xffffffff, ...) fills bytewise with the low byte, so each uint_64 slot becomes 0xFFFFFFFFFFFFFFFF, exactly static_cast<uint_64>(-1). And the boundary is genuinely reachable rather than theoretical — deepmd::env_mat_nbor_update rounds with <= comparisons (_max_nbor_size <= 256 gives 256, and so on), so a row holding exactly 256 neighbors lands in a 256-slot buffer with no padding at all. The assert in format_nbor_list_gpu lists the same five capacities the test enumerates, so the test covers the supported set completely. The fix is complete as well: the terminal-key idiom appears at exactly two sites repo-wide and both are changed, and since ROCm compiles this same file after the .hip.cu fold-away in #2844, there is no parallel path left behind.
On coverage: the existing formatter tests round a six-atom fixture up to a 1024-slot buffer, so their rows always carry slack and the final sorted key is always a real sentinel, which is why the faulty comparison looked correct for years. The short-selection test bounds the formatted output rather than the candidate row, so it retains padding too. The untested cell was a candidate row whose valid-neighbor count exactly equals a supported sort capacity, and the new test crosses it: with sec = {0, N-1, N} and type[N] = 1 the single type-1 neighbor sits at sort position N-1, the final slot of a full row, so pre-fix format_nlist_fill_b skips it and formatted[N-1] stays -1 against an expected N. Putting the type boundary there is a good touch — it forces fill_nei_iter to process the last valid key as well, not just the fill kernel.
One correction to the description. The terminal-key assumption is considerably older than stated: git log -S 'key_out[MAGIC_NUMBER - 1]' bottoms out at 3d5901961 ("Add GPU support for tensorflow operations", 2019-11-10), which already contains for (unsigned int kk = 0; key_out[kk] != key_out[MAGIC_NUMBER - 1]; kk++) in source/op/cuda/descrpt_se_a.cu. 3c9dfc49c renamed MAGIC_NUMBER to MAX_NBOR_SIZE, and 76a1a2a1d (#379) only relocated the already-faulty loop into the consolidated formatter. Relatedly, fill_nei_iter did not get the assumption from #845 — it was written that way in e1c64ec06 (#832) on the ROCm side, and #845 mirrored it into CUDA. So the pattern is a 2019 original that survived two moves, which if anything strengthens the case for replacing the inference outright rather than special-casing the full row.
Last thing, on validation rather than code. The body reports "RTX 5090: 12 formatter/codec tests passed" under a CUDA 12.4 build, but #5826 — same setup — states that kernel-launching tests could not run on that host because CUDA 12.4 provides no sm_120 image for the RTX 5090, and notes its own new test passed only because it rejects the row before any launch. This test does launch kernels, so I would not treat that line as evidence. It does not change my read of the change: test_cuda.yml runs on merge_group, so the queue will execute this test for real and a failure blocks the merge. Worth dropping the claim from the body rather than leaving something on the record that cannot be reproduced.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5937 +/- ##
==========================================
- Coverage 79.21% 78.96% -0.25%
==========================================
Files 1069 1069
Lines 124070 124070
Branches 4522 4522
==========================================
- Hits 98278 97969 -309
- Misses 24171 24482 +311
+ Partials 1621 1619 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
63f5767
3 participants
Summary
Closes #5936.
Root cause
format_nbor_list_gpuinitializes unused key slots to the all-bits-one sentinel before sorting. Bothfill_nei_iterandformat_nlist_fill_bidentified padding by comparing each key with the final sorted key. That works only when the row is partially filled and the final slot is guaranteed to contain padding.For an exact-capacity row, every slot contains a real neighbor. The final sorted key is therefore valid, but comparing it with itself classifies it as padding and silently leaves the corresponding output slot as
-1. The callers permit these exact capacities, so this is a reachable non-overflow boundary case.The fix detects the actual sentinel value instead of inferring it from the final slot.
Introducing change
Git history traces the faulty terminal-key assumption to commit 76a1a2a in #379, which introduced the consolidated CUDA environment-matrix formatter in February 2021. #845 later parallelized the formatting step and copied the same assumption into
fill_nei_iter, but the bug was already present before that PR.Why tests missed it
The existing GPU formatter tests use a small six-atom fixture, then round every row with at most 1024 candidates up to a 1024-slot sort buffer. Those rows always contain many unused slots, so the final sorted key is always a real padding sentinel and the faulty comparison appears correct.
The short-selection test limits the formatted output capacity, not the candidate sort row, so it also retains padding. The encoding/decoding test exercises only key serialization. No test previously constructed a candidate row whose valid-neighbor count exactly equaled a supported sort capacity.
Validation
ruff check .ruff format .(1713 files unchanged)Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
Bug Fixes
Tests
max_nbor_sizevalues from 256 to 4096, ensuring each slot is preserved correctly without unintended padding.