fix(tabulate): preserve SE-T polynomial cache

[njzjz-bot]

Root cause

The SE-T grad-grad kernel cached mark_table_idx across the inner neighbor loop, but declared the matching FPTYPE var[6] coefficients inside that loop. When consecutive inputs mapped to the same table interval, the kernel skipped load_polynomial_params because the index matched, then consumed a newly scoped uninitialized array.

Any correct result depended on accidental register or stack reuse.

Introduction and history

The bug was introduced by 36275c81a, PR #1274, Accelerate model compression, in 2021. That change added the table-index cache but did not extend the coefficient cache lifetime to match it.

Test gap

There was no dedicated SE-T GPU grad-grad case with adjacent em_x values in the same spline interval. Broader numerical tests could also pass accidentally when the compiler reused the same registers for the newly scoped array.

Fix

Move the coefficient array outside the inner neighbor loop so its lifetime matches mark_table_idx. A cache hit now reuses initialized coefficients by construction.

Validation

• CUDA 12.4 compilation of the modified kernel;
• CPU/GPU grad-grad comparison with three consecutive inputs in one table interval;
• test run on an NVIDIA GeForce RTX 5090 allocation;
• ruff format .;
• ruff check ..

Fixes #5894

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 when multiple inputs share the same table interval.
  • Prevented coefficient handling from being reset during repeated calculations.

• Tests

  • Added GPU coverage comparing gradient results against the CPU implementation for shared table intervals.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 971c8cde-0e8e-4f32-9f80-f26010506c82

📥 Commits

Reviewing files that changed from the base of the PR and between 1e42611 and a3e5259.

📒 Files selected for processing (2)

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

🚧 Files skipped from review as they are similar to previous changes (2)

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

---

📝 Walkthrough

Walkthrough

The CUDA SE-T grad-grad kernel preserves polynomial coefficients across inner iterations. A GPU regression test uses repeated table-interval inputs and compares GPU results with CPU reference results.

Changes

SE-T grad-grad GPU correctness

Layer / File(s)	Summary
Polynomial cache lifetime and regression validation source/lib/src/gpu/tabulate.cu, source/lib/tests/test_tabulate_se_t.cc	The polynomial coefficient array persists across inner iterations. A GPU test validates repeated table-interval grad-grad results against CPU output with 1e-10 tolerance and releases device allocations.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• deepmodeling/deepmd-kit#5601: Modifies SE tabulation gradient and grad-grad logic in the same CUDA source area.
• deepmodeling/deepmd-kit#5933: Modifies related SE-T GPU kernels and adds SE-T GPU regression tests.

Suggested reviewers: wanghan-iapcm

🚥 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 clearly summarizes the main change: preserving the SE-T polynomial cache lifetime in tabulation.
Linked Issues check	✅ Passed	The code fix and targeted GPU regression test satisfy the cache-lifetime and same-table-interval requirements in [#5894].
Out of Scope Changes check	✅ Passed	The changes are limited to the SE-T cache-lifetime fix and its focused GPU regression test.
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_t.cc (1)

5330-5332: 🎯 Functional Correctness | 🔵 Trivial

Assert the repeated-interval precondition.

This regression depends on 0.011, 0.012, and 0.019 resolving to the same SE_T table index, but the test only states that in a comment. Add a direct check, such as resolving each input’s table index with the fixture metadata and asserting they are equal, so future metadata/table generation changes do not make the cache-hit test pass without covering reuse.

[loweffort_and_high_reward]

🤖 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_t.cc` around lines 5330 - 5332, In the
regression test around test_em_x, explicitly resolve the table index for each
input using the fixture metadata and assert that all three indices are equal
before exercising the cache behavior. Keep the existing inputs and cache-hit
assertions unchanged, using the test’s established metadata/index-resolution
symbols.

🤖 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_t.cc`:
- Around line 5330-5332: In the regression test around test_em_x, explicitly
resolve the table index for each input using the fixture metadata and assert
that all three indices are equal before exercising the cache behavior. Keep the
existing inputs and cache-hit assertions unchanged, using the test’s established
metadata/index-resolution symbols.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 7c73218b-e488-4bbe-bc02-a62213c9a4de

📥 Commits

Reviewing files that changed from the base of the PR and between 4f827cc and 93fc6da.

📒 Files selected for processing (2)

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.16%. Comparing base (9b2582f) to head (a3e5259).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5931      +/-   ##
==========================================
- Coverage   79.41%   79.16%   -0.26%     
==========================================
  Files        1072     1072              
  Lines      124893   124893              
  Branches     4531     4536       +5     
==========================================
- Hits        99187    98874     -313     
- Misses      24085    24398     +313     
  Partials     1621     1621              

☔ 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]

Correct and minimal. Hoisting var[6] to the ii scope gives it the same lifetime as mark_table_idx, so a cache hit reuses coefficients that were definitely loaded. mark_table_idx resets to -1 each outer iteration and locate_xx_se_t never returns a negative index, so the first inner iteration always takes the load branch and var is initialised before any read.

I spent most of the review on whether the fix is complete, since a half-applied fix is the usual failure mode in this file. It holds up two independent ways. Sweeping every tabulate kernel in tabulate.cu, se_t_grad_grad was the only one where the coefficient array and the cached index had mismatched scopes; the other cached kernels already pair them, and the uncached ones reload every iteration. Separately, the history agrees: 36275c81a added the cache to four kernels and hoisted var in three of them, missing only this one. mark_table_idx exists nowhere else in the tree, and the CPU path has no such cache, so there is no sibling left to patch.

The history also explains why this survived so long, which I think is worth recording. 450455f4e removed the unloop early break, so the loop now runs every nnei_j entry including zero-padded neighbours, which map to identical indices back to back and therefore hit the cache reliably. And c41c36631 fixed a for (int jj = 0; ii < nnei_j; jj++) typo in this same kernel, meaning the inner loop condition was broken until then. The bug only became properly reachable fairly recently. There is precedent too: 1503fc3a1 fixed two instances of this same family, which showed up as NaN output.

On the test: it does exercise the cache-hit path, since all three em_x values resolve to one table index and the second and third iterations take the hit branch. I want to be explicit that it is a guard rather than a reproducer, and I am not asking for more. Reading an uninitialized loop-scoped array is undefined behaviour, and the compiler will normally reuse the same slot, so the unfixed build would very likely pass this assertion. You already say as much in the description and in #5894, which is the right way to present it. If this class of defect comes up again, running the kernel under compute-sanitizer --tool initcheck is the tool that would actually catch it.

Two small things, neither blocking. CodeRabbit's open point is worth taking: the shared-interval property that makes this test meaningful is asserted only in a comment, so an ASSERT_EQ on the computed table indices would stop a future change to the fixture's table/info from silently downgrading it to a plain CPU-vs-GPU consistency check. And #5933 appends its new TEST_F at the same place in this file, so whichever of the two lands second will need a rebase.

[njzjz]

Addressed the non-blocking cache-test note in 1e42611. The regression now derives the stride-0 interval from the fixture metadata, verifies every test input is inside that range, and asserts all three resolve to the same interval before comparing the CUDA result with the CPU reference.

Validation: git diff --check, ruff format ., and ruff check . passed. The CUDA C++ test binary is not available in this worktree, so the GPU test remains CI-validated.

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