perf(neighbors): tile partial naive neighbor lists

[zubatyuk]

ALCHEMI Toolkit-Ops Pull Request

Description

Accelerate native, JAX, and Torch naive neighbor-list rebuilds that update a compact subset of atoms. The previous implementation used scalar traversal in framework partial paths even when a CUDA tile-cooperative topology-only specialization was available, making large-system partial rebuilds disproportionately expensive. Expose the native specialization consistently through JAX and Torch, preserve scalar behavior for unsupported outputs, and retain optional JAX-only Warp graph replay for callers that keep all graph inputs and outputs at stable addresses.

The new compact-row path maps each output row to its global source atom:

ishift, row = wp.tid()
atom_i = row
if PARTIAL:
    atom_i = target_indices[row]

For batched selective rebuilds, reset compact output rows using the target atom's owning system rather than the compact-row index:

row = wp.tid()
atom = target_indices[row]
if rebuild_flags[batch_idx[atom]]:
    num_neighbors[row] = 0

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

N/A

Changes Made

• Add CUDA tiled naive neighbor-matrix support for topology-only compact target_indices rows, including no-PBC, wrapped PBC, and prewrapped PBC paths.
• Expose the same topology-only partial tile path through JAX and Torch single-system APIs; expose explicit tile through both batched APIs while retaining scalar batched strategy="auto".
• Add dtype-specific single-system automatic dispatch thresholds: tile at 256 atoms for float64 and 1,024 atoms for float16/float32; retain scalar automatic dispatch for batched partial and batched prewrapped PBC workloads.
• Fix batched PBC strategy propagation, scalar partial-PBC wrapping parity, target-aware selective reset behavior, and tiled-kernel input ABI coverage.
• Add JAX opt-in graph_mode="warp" replay for tiled partial topology calls with persistent caller-owned buffers; reject dynamic COO output in this mode.
• Add regression coverage for scalar/tile count equality and order-independent (neighbor, periodic-shift) parity, graph configuration isolation, persistent-graph constraints, batched routing, and target-aware resets.

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

Tile execution is CUDA-only and intentionally limited to topology-only results. Requests for distances, vectors, pair functions, or selective rebuilds continue to use scalar kernels. Scalar and tiled neighbor rows may have different entry order; their active neighbor and periodic-shift multisets are equivalent. The benchmark numbers below are JAX measurements; this PR makes no Torch performance claim.

Controlled RTX 6000 Ada benchmarks used a 9,660-atom periodic system, six target rows, a 12 Å cutoff, and a capacity of 4,096 neighbors per row. Isolated synchronized neighbor-list calls measured:

Dtype	Scalar	Tile	Speedup
float32	1,266.2 µs	234.8 µs	5.39x
float64	4,887.7 µs	257.8 µs	18.96x

Every scalar/tile run returned 2,074 active edges with exact count equality and exact order-independent (neighbor, periodic-shift) multiset equality.

With stable pointers, optional Warp graph replay reduced isolated tiled partial-call latency from 228.4 to 125.2 µs for float32 (1.82x) and from 255.2 to 153.5 µs for float64 (1.66x). Fresh input addresses made graph replay slower, so graph replay remains opt-in and requires persistent buffers.

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.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds CUDA tile-cooperative support for partial (compact target_indices) naive neighbor-list rebuilds, exposing the native specialization consistently through JAX and Torch single-system and batched APIs. It also introduces prewrapped-PBC tiled variants for both batched paths, fixes batched-PBC strategy propagation, corrects scalar partial-PBC wrapping parity, and adds an opt-in JAX graph_mode="warp" replay path for topology-only tiled partial calls.

• Kernel changes: _make_tile_kernel gains a PARTIAL static flag; each tile block maps its row dimension to target_indices[row], computes bidirectional shifts (2 * num_shifts - 1) for the full-fill case, and writes neighbors into the compact output row rather than the global atom row. A new selective_zero_partial_num_neighbors kernel correctly resets compact rows by looking up the owning system via batch_idx[target_indices[row]].
• Dispatch layer: Single-system auto dispatch selects tile at \u2265 256 atoms (float64) or \u2265 1,024 atoms (float32/float16); batched partial auto dispatch remains scalar. JAX warp-graph capture is guarded to topology-only partial calls with stable caller-owned buffers and return_neighbor_list=False.
• Tests: New tests verify scalar/tile count and order-independent multiset equality, warp-graph configuration isolation and replay correctness, persistent-graph buffer constraints, batched routing, and target-aware selective resets.

Important Files Changed

Filename	Overview
nvalchemiops/neighbors/naive/kernels.py	Adds PARTIAL static flag to tile kernel factory; correctly maps compact row to global atom_i, extends shift dimension to 2*num_shifts-1 for full-fill partial, and handles self-exclusion and half_fill exclusion. Minor test-coverage gap for half_fill+PARTIAL in eager tile mode.
nvalchemiops/neighbors/naive/launchers.py	Auto-dispatch thresholds and strategy routing for partial tile are correctly separated from geometry/pair-fn outputs; strategy propagation to _launch_naive_neighbor_matrix_pbc is fixed; selective_zero_partial_num_neighbors is wired correctly for batched partial+rebuild_flags paths.
nvalchemiops/neighbors/neighbor_utils.py	Adds selective_zero_partial_num_neighbors kernel that correctly uses batch_idx[target_indices[row]] to identify the owning system for each compact row before zeroing.
nvalchemiops/jax/neighbors/naive.py	Adds partial tile warp-graph callable family, topology_only_partial detection, strict validation for warp-graph constraints (persistent buffers, no COO output), and correct use_tile dispatch in _naive_pair_outputs_forward. Three pre-existing P2 nits flagged in prior reviews remain.
nvalchemiops/jax/neighbors/batch_naive.py	Adds prewrapped-PBC tiled variants (f32/f64), partial flag threading through all tile callables, and topology_only_partial routing in _batch_naive_pair_outputs_forward including correct prewrapped scalar parity.
nvalchemiops/jax/neighbors/_dispatch.py	_is_jax_cpu_array now catches ConcretizationTypeError and TracerIntegerConversionError, correctly falling back to jax.default_backend() during JIT tracing.
nvalchemiops/torch/neighbors/naive.py	Exposes tile partial path through the Torch single-system API with auto-dispatch thresholds matching the native launcher.
nvalchemiops/torch/neighbors/batch_naive.py	Threads target_indices into the batched Torch API; batched partial auto dispatch correctly remains scalar while explicit tile is opt-in.
test/neighbors/bindings/jax/test_naive.py	Comprehensive coverage: scalar/tile multiset equality, warp-graph isolation and replay, persistent-buffer constraint guards, COO output rejection; half_fill=True not tested in eager non-graph partial tile path.
test/neighbors/bindings/torch/test_naive.py	Adds GPU-marked tests for tile vs scalar parity and sorted multiset comparison helper.
test/neighbors/test_naive_kernel_getters.py	Adds kernel-getter tests for the new partial tile specialization.
test/neighbors/test_batch_naive_kernels.py	New tests cover batched prewrapped-PBC tile, target-aware selective reset, and batched partial routing.

_{Reviews (2): Last reviewed commit: "fix(neighbors): clean up partial tile di..." | Re-trigger Greptile}