feat(dpmodel): add adaptive Array API neighbor search#5889
Conversation
Add a Cartesian cell-list path for large systems while retaining the dense builder for backend-specific small-system fast paths. Cover NumPy, PyTorch, JAX eager, and TensorFlow graph execution, and include reproducible rcut=6 CPU/GPU benchmark data and plots. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughDefaultNeighborList now adaptively selects dense or Cartesian cell-list neighbor construction using backend- and device-aware thresholds. The cell-list path preserves ordering, exclusions, padding, and output contracts, with deterministic common and cross-backend tests plus TensorFlow shape handling updates. ChangesAdaptive neighbor-list implementation
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant DefaultNeighborList
participant DispatchChecks
participant CellListBuilder
participant DenseBuilder
participant PairExclusion
DefaultNeighborList->>DispatchChecks: inspect nloc, backend, device, periodicity
DispatchChecks-->>DefaultNeighborList: choose cell-list or dense construction
DefaultNeighborList->>CellListBuilder: build fixed-width neighbors
CellListBuilder->>PairExclusion: apply pair exclusions
DefaultNeighborList->>DenseBuilder: construct fallback neighbors when unsupported
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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
🧹 Nitpick comments (3)
doc/development/default-neighbor-list-benchmark.md (1)
84-101: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd an example command for the plotting script.
This section documents how to reproduce the CSV data but not the two embedded plots. Since
plot_default_neighbor_list_benchmark.pyis linked just above (line 38), consider adding a short example invocation here for completeness.📝 Proposed addition
srun --gres=gpu:1 python \ source/tests/common/dpmodel/benchmark_default_neighbor_list.py \ --backend jax --device gpu --scenario nonperiodic \ --output result_jax_gpu_nonperiodic.csv
+Generate the plots from a directory of
result_*.csvfiles:
+
+bash +python doc/development/plot_default_neighbor_list_benchmark.py <results-dir> +</details> <details> <summary>🤖 Prompt for AI Agents</summary>Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.In
@doc/development/default-neighbor-list-benchmark.mdaround lines 84 - 101,
Add a short plotting example to the “Example commands” section, using the linked
plot_default_neighbor_list_benchmark.pyscript with a results-directory
argument for generating plots fromresult_*.csvfiles. Keep the existing CPU
and GPU benchmark commands unchanged.</details> <!-- cr-comment:v1:f7cd7a2855387dfded186913 --> </blockquote></details> <details> <summary>doc/development/plot_default_neighbor_list_benchmark.py (2)</summary><blockquote> `95-96`: _🎯 Functional Correctness_ | _🔵 Trivial_ | _⚡ Quick win_ **Legend built from a single axis in both `plot_runtime` and `plot_speedup`.** Both functions collect the shared figure legend's handles/labels from only the `("cpu", "nonperiodic")` axis; any series absent from that one panel (but present in another) silently disappears from the legend even though its line is drawn elsewhere. - `doc/development/plot_default_neighbor_list_benchmark.py#L95-L96`: build the legend from the union of handles/labels across all four axes (e.g. iterate `axes.values()`, dedupe by label) instead of `axes[("cpu", "nonperiodic")]` alone. - `doc/development/plot_default_neighbor_list_benchmark.py#L139-L140`: apply the same all-axes, deduplicated handle collection here. <details> <summary>♻️ Proposed shared helper</summary> ```diff +def collect_legend(axes: dict[tuple[str, str], plt.Axes]) -> tuple[list, list]: + seen: dict[str, Any] = {} + for axis in axes.values(): + handles, labels = axis.get_legend_handles_labels() + for handle, label in zip(handles, labels): + seen.setdefault(label, handle) + return list(seen.values()), list(seen.keys())Then in
plot_runtimeandplot_speedup:- handles, labels = axes[("cpu", "nonperiodic")].get_legend_handles_labels() + handles, labels = collect_legend(axes)🤖 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 `@doc/development/plot_default_neighbor_list_benchmark.py` around lines 95 - 96, Update plot_runtime at doc/development/plot_default_neighbor_list_benchmark.py#L95-L96 and plot_speedup at doc/development/plot_default_neighbor_list_benchmark.py#L139-L140 to build shared legends from handles and labels collected across all axes.values(), deduplicated by label, rather than only the ("cpu", "nonperiodic") axis. Reuse a shared helper if appropriate, ensuring every plotted series appears once in both legends.
14-19: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
matplotlibandpandasto the doc requirements, and require Matplotlib >=3.7 forloc="outside lower center".doc/requirements.txtdoesn’t currently declare either package, and that legend placement was added in Matplotlib 3.7.0.🤖 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 `@doc/development/plot_default_neighbor_list_benchmark.py` around lines 14 - 19, Add matplotlib and pandas to doc/requirements.txt, pinning or constraining matplotlib to version 3.7.0 or newer so the benchmark’s outside-lower-center legend placement is supported.
🤖 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 `@doc/development/plot_default_neighbor_list_benchmark.py`:
- Around line 95-96: Update the shared legend construction near axes[("cpu",
"nonperiodic")] so it collects and combines handles and labels from every
plotted panel in axes, rather than relying on a single panel. Preserve the
existing figure.legend placement and formatting while ensuring each series
present in any panel appears once in the shared legend.
---
Nitpick comments:
In `@doc/development/default-neighbor-list-benchmark.md`:
- Around line 84-101: Add a short plotting example to the “Example commands”
section, using the linked `plot_default_neighbor_list_benchmark.py` script with
a results-directory argument for generating plots from `result_*.csv` files.
Keep the existing CPU and GPU benchmark commands unchanged.
In `@doc/development/plot_default_neighbor_list_benchmark.py`:
- Around line 95-96: Update plot_runtime at
doc/development/plot_default_neighbor_list_benchmark.py#L95-L96 and plot_speedup
at doc/development/plot_default_neighbor_list_benchmark.py#L139-L140 to build
shared legends from handles and labels collected across all axes.values(),
deduplicated by label, rather than only the ("cpu", "nonperiodic") axis. Reuse a
shared helper if appropriate, ensuring every plotted series appears once in both
legends.
- Around line 14-19: Add matplotlib and pandas to doc/requirements.txt, pinning
or constraining matplotlib to version 3.7.0 or newer so the benchmark’s
outside-lower-center legend placement is supported.
🪄 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
Run ID: 8b567895-debc-4682-b23e-68e76aa8b9cd
⛔ Files ignored due to path filters (3)
doc/development/default_neighbor_list_benchmark_rcut6.csvis excluded by!**/*.csvdoc/images/neighbor_list_rcut6_runtime_loglog.pngis excluded by!**/*.pngdoc/images/neighbor_list_rcut6_speedup_loglog.pngis excluded by!**/*.png
📒 Files selected for processing (6)
deepmd/dpmodel/utils/default_neighbor_list.pydoc/development/default-neighbor-list-benchmark.mddoc/development/plot_default_neighbor_list_benchmark.pydoc/index.rstsource/tests/common/dpmodel/benchmark_default_neighbor_list.pysource/tests/common/dpmodel/test_default_neighbor_list.py
There was a problem hiding this comment.
Pull request overview
This PR updates DeePMD-kit’s dpmodel default neighbor-list construction to dispatch between the historical dense all-pairs method and a new Array-API-based Cartesian cell-list candidate search, aiming to preserve existing neighbor-list semantics while improving scalability on larger systems across backends.
Changes:
- Added an adaptive dispatcher in
DefaultNeighborListwith backend/device/periodicity-specific crossover thresholds. - Implemented a compact cell-list neighbor search using Array API ops (NumPy/Torch/JAX eager/TF graph/array-api-strict) while preserving cutoff/virtual-atom/exclusion/ordering behavior.
- Added comprehensive tests plus reproducible benchmark scripts, raw data, and documentation/plotting utilities for
rcut=6.
Reviewed changes
Copilot reviewed 7 out of 9 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
deepmd/dpmodel/utils/default_neighbor_list.py |
Implements adaptive dense vs cell-list dispatch and the Array-API cell-list builder. |
source/tests/common/dpmodel/test_default_neighbor_list.py |
Validates dense vs cell-list equivalence, ordering, exclusions, and backend-specific behaviors. |
source/tests/common/dpmodel/benchmark_default_neighbor_list.py |
Provides a reproducible benchmark harness for dense vs cell-list runtime comparisons. |
doc/development/default-neighbor-list-benchmark.md |
Documents methodology and summarizes measured crossovers and scaling observations. |
doc/development/default_neighbor_list_benchmark_rcut6.csv |
Adds the raw benchmark results used for plots and threshold selection rationale. |
doc/development/plot_default_neighbor_list_benchmark.py |
Generates runtime and speedup plots from benchmark CSVs. |
doc/index.rst |
Wires the new benchmark doc into the development documentation index. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5889 +/- ##
==========================================
- Coverage 78.87% 78.70% -0.18%
==========================================
Files 1054 1054
Lines 121770 122042 +272
Branches 4409 4409
==========================================
+ Hits 96046 96053 +7
- Misses 24158 24417 +259
- Partials 1566 1572 +6 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Normalize legacy TensorFlow dynamic dimensions in ndtensorflow indexing. Move all backend-specific neighbor-list coverage into the consistent test suite, with TF2 remaining opt-in, and address the benchmark plotting review feedback. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Pushed 13b05a3 to address the failed Python 3.10 shard and review feedback.
Validation: common NumPy tests 7 passed; default consistent selection 4 passed and 2 TF2 tests skipped; TF2-only selection 2 passed; combined regular selection 11 passed and 2 skipped; all hooks passed on the six changed files; both benchmark plots regenerated successfully. Coding agent: Codex |
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/tests/consistent/test_default_neighbor_list.py`:
- Line 125: Update the test containing jax.config.update to preserve the prior
jax_enable_x64 setting and restore it during test cleanup, ensuring the
process-wide JAX dtype mode is unchanged for subsequent tests.
🪄 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
Run ID: 0f471318-4419-4204-a94c-eb8639ea261c
📒 Files selected for processing (6)
deepmd/_vendors/ndtensorflow/_namespace.pydoc/development/default-neighbor-list-benchmark.mddoc/development/plot_default_neighbor_list_benchmark.pydoc/requirements.txtsource/tests/common/dpmodel/test_default_neighbor_list.pysource/tests/consistent/test_default_neighbor_list.py
🚧 Files skipped from review as they are similar to previous changes (2)
- doc/development/default-neighbor-list-benchmark.md
- doc/development/plot_default_neighbor_list_benchmark.py
Use one unittest import style and restore the process-wide JAX x64 setting after the neighbor-list consistency test. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Sort eager candidates in padded per-center rows to avoid three global candidate sorts and the final selection scatter. Keep compact sorting for traced graphs, JAX accelerators, and imbalanced candidate distributions. Remove benchmark documentation and generated data from the PR scope. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Use the compact selection path when a PyTorch build does not expose torch.compiler.is_compiling. This preserves compatibility with older or reduced variants without risking a data-dependent allocation while tracing. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Keep placeholder coordinates from expanding or overflowing the Cartesian cell grid, including all-virtual frames. Reuse the backend consistency mixin to cover CPU and available accelerator implementations. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Generate fractional coordinates in the requested dtype so float32 common and backend consistency cases exercise the intended path. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Use bounded int32 candidate indices on measured eager backends, avoid scan-based iotas when shapes are concrete, and compact unreachable periodic images only beyond backend-specific crossover points. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Pushed c642e8a with only benchmark-positive optimizations. Retained changes:
Discarded experiments:
Representative rcut=6, density=0.05, nsel=128 results:
For periodic PyTorch CUDA at N=262,144, a same-process A/B measured full-image sorting at 29.65 ms / 7.394 GiB versus compact boundary-shell sorting at 27.56 ms / 6.925 GiB. The conservative compaction threshold is therefore 262,144 atoms; it remains disabled below that point and for JAX/TF accelerators. On the 31.36 GiB RTX 5090, the previous maximum successes were about 983,040 nonperiodic and 851,968 periodic atoms. The updated implementation successfully built 1,179,648 nonperiodic and 1,048,576 periodic atoms. Validation:
Coding agent: Codex |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Nice work — the dense/cell-list equivalence holds up under scrutiny (checked cutoff boundary, stable tie-break, self-exclusion, virtual atoms, 27-cell completeness, frame separation, and the compaction pruning bound), and the _namespace.py as_list() change is a correct fix for the legacy Dimension(None) case, not a regression. No correctness issues found. A few non-blocking notes below — one test-coverage gap worth closing and two stale docstrings.
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
| offsets = xp.asarray(_NEIGHBOR_CELL_OFFSETS, dtype=xp.int64, device=device) | ||
| query_cell = center_cell[:, :, None, :] + offsets[None, None, :, :] |
| if array_api_compat.is_torch_array(coord): | ||
| import torch | ||
|
|
||
| # torch.export/compile must keep the compact path: converting the maximum | ||
| # candidate count to ``int`` would specialize an unbacked symbolic value. | ||
| # Keep unmeasured accelerator implementations on the compact path too. | ||
| device = array_api_compat.device(coord) | ||
| is_compiling = getattr(getattr(torch, "compiler", None), "is_compiling", None) | ||
| # Older or reduced PyTorch builds may not expose torch.compiler. Without | ||
| # a reliable tracing-state query, keep the dynamic compact path instead | ||
| # of risking a data-dependent Python allocation during compilation. | ||
| return ( | ||
| device.type in ("cpu", "cuda") | ||
| and is_compiling is not None | ||
| and not is_compiling() | ||
| ) |
| center_start = xp.cumulative_sum(count_per_center) - count_per_center | ||
| edge_iota = _index_iota(center) | ||
| rank = edge_iota - xp.take(center_start, center, axis=0) | ||
| slot = center * max_candidates + rank | ||
|
|
||
| neighbor_rows = xp_scatter_sum( | ||
| xp.zeros((padded_size,), dtype=xp.int64, device=device), | ||
| 0, | ||
| slot, | ||
| neighbor_ext + 1, | ||
| ) |
Summary
source/tests/consistent, running it on CPU and on an available GPU; TF2 remains opt-in through the existing environment variablesJAX neighbor-list construction remains eager outside
jit; traced full-builder calls retain the dense fallback. TensorFlowtf.functionretains compact global sorting because a data-dependent maximum row width cannot safely define a graph allocation.rcut=6 optimization impact
Cell-list time before and after the per-center selection optimization, on an AMD EPYC 7K62 CPU and NVIDIA RTX 5090:
JAX GPU keeps the compact path because extracting the padded row width onto the host caused a small regression near its crossover. TensorFlow graph mode also keeps the compact path.
Current GPU dense/cell validation
Representative measurements at the automatic dispatch region:
tf.functiontf.functionNumPy has no GPU backend. The following plots are linked from the completed benchmark run for review only; the generated PNG files are not part of the final PR branch or diff.
End-to-end comparison with vesin
These measurements start from the same local coordinates and box at
rcut=6, density 0.05, and one float32 frame. DeePMD runsDefaultNeighborList.buildwithnsel=128, including periodic ghost extension and fixed-width stable distance ordering. vesin v0.6.0 runsfull_list=True,sorted=False, andquantities="ij"; its variable-length unsorted edge output is lighter, so this comparison represents the benefit of a specialized native implementation rather than identical output work. Million-scale DeePMD CPU points use one measured build because each call takes tens to hundreds of seconds; GPU and smaller points use warmed medians.CPU: vesin versus NumPy DeePMD
DeePMD peak RSS was about 69.8/77.1 GiB at 2,097,152 and 140.7/153.9 GiB at 4,194,304 nonperiodic/periodic atoms. The 8,388,608-atom DeePMD runs were not attempted because linear memory extrapolation exceeds the 251 GiB host, while vesin completed both in about 5.1 seconds.
GPU: vesin CUDA versus PyTorch DeePMD
On the 31.36 GiB RTX 5090, DeePMD succeeded through 983,040 nonperiodic atoms at about 29.2 GiB peak allocated memory and through 851,968 periodic atoms at about 28.4 GiB; the next measured sizes OOMed. The gap to vesin narrows near this limit because vesin caps the CUDA cell grid at 8,192 cells, after which its runtime becomes close to quadratic:
At 4,194,304 atoms, vesin CPU (1.75/1.77 s) is already slightly faster than CUDA (1.95/2.06 s) because of that grid cap. Neighbor counts matched exactly at smaller sizes and differed by at most 63 cutoff-boundary pairs among roughly 190 million directed edges at 4,194,304 atoms. vesin converts coordinates to float64 internally, while DeePMD uses float32 here, so bit-identical cutoff membership is not expected.
Validation
srun --gres=gpu:1: 6 passed, with 6 shared CPU/GPU subtests passedsrun --gres=gpu:1: 4 passed, with 6 shared CPU/GPU subtests passedDP_TEST_TF2_ONLY=1): 6 passed, with 6 shared CPU/GPU subtests passedCoding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh