fix(neighbors): harden selective cluster-tile state - #141
Conversation
Dispatch the existing wrapped PBC kernels after position wrapping for normal and selective rebuilds. Add parity coverage against single-cutoff results for both dtypes and wrap modes. Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
Validate fixed-capacity COO state before launches, preserve JAX empty selective state, and support fullgraph single-system COO reuse. Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile SummaryThis PR hardens selective cluster-tile state handling across PyTorch, JAX, and Warp.
Important Files Changed
Reviews (2): Last reviewed commit: "update changelog." | Re-trigger Greptile |
Prevent malformed segmented COO metadata from writing past physical output capacity. Document the compiled cluster-tile requirement for eagerly validated, fully periodic PBC. Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
laserkelvin
left a comment
laserkelvin
left a comment
There was a problem hiding this comment.
Several correctness and state-lifecycle issues remain; details are inline. F1 already had an existing inline thread, so I did not duplicate it.
|
|
||
| if ( | ||
| rebuild_flags is not None | ||
| and format == "coo" |
There was a problem hiding this comment.
issue: This prior-state guard only covers COO. A matrix call with any false flag can allocate fresh outputs and tile state, then skip that system, so there is no previous topology to preserve. I assume False preserves the prior result for every supported format; please require the matrix triples (including dual-cutoff outputs) and tile state before launch when any flag is false.
| "neighbor_list_shifts": neighbor_list_shifts, | ||
| "pair_offsets": pair_offsets, | ||
| "pair_counts": pair_counts, | ||
| "rebuild_flags": rebuild_flags, |
There was a problem hiding this comment.
issue: Valid nonselective segmented COO reaches this helper with rebuild_flags=None; adding it to tensors makes the loop below dereference None.device and raise AttributeError. I assume pair_offsets/pair_counts may select segmented mode without selective rebuilding. Please make rebuild_flags optional here and validate it only when present.
| internally without replacing those state buffers. | ||
|
|
||
| Use two phases for PyTorch selective COO. First make an eager bootstrap call with | ||
| every flag true; it may allocate omitted topology and tile buffers and returns the |
There was a problem hiding this comment.
issue: This bootstrap path is not currently implementable: cluster_tile_neighbor_list requires every topology and tile buffer whenever rebuild_flags is set, even when all flags are true. I assume the first all-true call should allocate and return reusable state as documented. Please implement that path or replace this paragraph with complete explicit allocation/bootstrap steps.
| return ( | ||
| if ( | ||
| selective | ||
| and not torch.compiler.is_compiling() |
There was a problem hiding this comment.
suggestion: Excluding compiled calls from this fast return means a false fullgraph step still performs wrapping, Morton coding, argsort, gathers, and the query launch. I assume selective reuse should make an unchanged single system a no-op. Please gate preprocessing/query with a compiled conditional or flag-aware operation, and add a diagnostic or benchmark for the false path.
| or neighbor_list_shifts is None | ||
| ): | ||
| raise ValueError("selective COO requires fixed topology buffers") | ||
| max_pairs_from_buffers = _validate_segmented_coo_state( |
There was a problem hiding this comment.
suggestion: This eager validation is repeated inside query_cluster_tile_coo, so a high-level call copies the same offsets and counts to CPU twice and serializes GPU work twice. I assume direct low-level callers still need validation, while the convenience path should pay for it once. Please split tile/pair validation or dispatch through an already-validated internal query path.
| per_atom_cell_offsets, | ||
| launch_dims=(total_atoms,), | ||
| ) | ||
| if rebuild_flags is not None: |
There was a problem hiding this comment.
suggestion: This new wrapped-PBC selective branch has no matching regression test: the PBC tests omit rebuild_flags, while the selective tests omit PBC. I assume False preserves both cutoffs and shifts, and True matches a fresh nonselective rebuild. Please add both flag cases against those references.
| Explicit cluster-tile methods append their documented tile-state suffix | ||
| when ``return_state=True``. | ||
|
|
||
| Single-system selective COO calls to explicit ``method="cluster_tile"`` |
There was a problem hiding this comment.
suggestion: This documents only the single-system selective exception. The batched route returns four topology tensors plus six state tensors, and the JAX matrix docstrings also omit their selective state suffixes. I assume the primary API docs should make tuple arity and ordering explicit; please add those missing return forms here and in the JAX Returns sections.
2 participants
ALCHEMI Toolkit-Ops Pull Request
Description
Harden selective cluster-tile neighbor-list state handling across PyTorch and JAX, including fixed-capacity COO validation, empty-system semantics, and compiled single-system reuse.
Type of Change
Related Issues
None.
Changes Made
Testing
make pytest)make lint)Checklist
Additional Notes
None.
Tip
This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.