Neighbor List Benchmarks

Performance benchmarks for neighbor list algorithms in ALCHEMI Toolkit-Ops. Results show scaling behaviour across system sizes for multiple cutoff radii and algorithms.

Warning

These results are intended to be indicative only: your actual performance may vary depending on the atomic system topology, software and hardware configuration and we encourage users to benchmark on their own systems of interest.

How to Read These Charts

Time Scaling : Mean execution time (µs/atom) vs. system size. Lower is better. The naive strategies enumerate all pairs and have \(O(N^2)\) work. Cell-list strategies are expected to approach \(O(N)\) at fixed density and cutoff; finite-size timings also include launch, build/query, and output costs.

Throughput : Atoms processed per second (plotted as \(10^6\) atoms/s). Higher is better. This metric helps compare efficiency across different system sizes.

Memory : Peak memory reported by the Torch CUDA allocator vs. system size. Units switch between MB and GB automatically on the y-axis. JAX memory is not measured by this suite.

Batch Scaling : Compares the concrete batch APIs. A batch size of one resolves to the corresponding single-system API and is therefore kept out of these curves. The smaller fixed systems contain 250 atoms for CsCl and 256 atoms for NH₃, so their first plotted batch points contain 500 and 512 total atoms, respectively. The system-size panels include the 128-atom cases.

Performance Results

Cutoffs 6 Å 15 Å 25 Å

Torch

CsCl

System Size Scaling

Mean execution time vs. system size for naive and cell list algorithms.

Throughput (\(10^6\) atoms/s) vs. system size.

Peak GPU memory vs. system size.

Constant Workload

Execution time near the configured total-atom target, varying batch size.

Throughput near the configured total-atom target.

Peak GPU memory near the configured total-atom target.

Batch Scaling

Execution time vs. batch size (fixed atoms per system).

Throughput vs. batch size.

Peak GPU memory vs. batch size.

NH₃

System Size Scaling

Mean execution time vs. system size for naive and cell list algorithms (NH₃).

Throughput (\(10^6\) atoms/s) vs. system size (NH₃).

Peak GPU memory vs. system size (NH₃).

Constant Workload

Execution time near the configured total-atom target (NH₃).

Throughput near the configured total-atom target (NH₃).

Peak GPU memory near the configured total-atom target (NH₃).

Batch Scaling

Execution time vs. batch size (NH₃).

Throughput vs. batch size (NH₃).

Peak GPU memory vs. batch size (NH₃).

JAX

The default JAX reportable pass covers naive_scalar, naive_tile, cell_list_atom_centric, and cluster_tile (plus their batched forms). Pair-centric cell-list timing remains available as an explicit coverage-only request. The cluster-tile path requires CUDA, float32 positions, and fully periodic cells; unsupported inputs are retained as explicit policy rows.

CsCl

System Size Scaling

Mean execution time vs. system size (JAX).

Throughput (\(10^6\) atoms/s) vs. system size (JAX).

Constant Workload

Execution time near the configured total-atom target (JAX).

Throughput near the configured total-atom target (JAX).

Batch Scaling

Execution time vs. batch size (JAX).

Throughput vs. batch size (JAX).

Note

JAX memory plots are omitted. The suite does not measure JAX memory: XLA’s allocator pool and buffer reuse make per-call allocation deltas unreliable. The Torch panels report the Torch allocator only; they are not a proxy for JAX memory because the framework wrappers, buffer lifetimes, and allocators differ.

NH₃

System Size Scaling

Mean execution time vs. system size (JAX, NH₃).

Throughput (\(10^6\) atoms/s) vs. system size (JAX, NH₃).

Constant Workload

Execution time near the configured total-atom target (JAX, NH₃).

Throughput near the configured total-atom target (JAX, NH₃).

Batch Scaling

Execution time vs. batch size (JAX, NH₃).

Throughput vs. batch size (JAX, NH₃).

Note

JAX memory plots are omitted. The suite does not measure JAX memory: XLA’s allocator pool and buffer reuse make per-call allocation deltas unreliable. The Torch panels report the Torch allocator only; they are not a proxy for JAX memory because the framework wrappers, buffer lifetimes, and allocators differ.

Backend Comparison

These overlays use only matched, comparison-eligible Torch/JAX rows. The cutoff switches above control these panels as well as the individual backend panels. JAX allocator memory is not shown.

Note

In the two Constant Workload comparison tabs, selecting only 25 Å leaves the chart empty. Those JAX measurements had to wait after every timed call to fit their output storage, so they are not timed the same way as Torch and are deliberately excluded from the overlay. The individual Torch and JAX tabs still show the 25 Å measurements.

CsCl

System Size Scaling

Torch vs. JAX execution time comparison.

Torch vs. JAX throughput comparison.

Constant Workload

Torch vs. JAX execution time at constant workload.

Torch vs. JAX throughput at constant workload.

Batch Scaling

Torch vs. JAX execution time vs. total atoms.

Torch vs. JAX throughput vs. total atoms.

NH₃

System Size Scaling

Torch vs. JAX execution time comparison (NH₃).

Torch vs. JAX throughput comparison (NH₃).

Constant Workload

Torch vs. JAX execution time at constant workload (NH₃).

Torch vs. JAX throughput at constant workload (NH₃).

Batch Scaling

Torch vs. JAX execution time vs. total atoms (NH₃).

Torch vs. JAX throughput vs. total atoms (NH₃).

Comparison contract

Torch/JAX comparison eligibility is limited to naive_scalar and batch_naive_scalar. Both backends must succeed at the same x-value, and JAX serial-fallback rows are excluded.

The other strategies remain in the individual backend panels. They are not overlaid because the current public APIs use different execution contracts: JAX naive-tile runs through an eager public call, the cell-list wrappers can choose different cell grids or query paths, pair-centric has backend-specific launch guards, and cluster-tile has different traceability between single and batched APIs.

JAX normally dispatches timed calls back-to-back and blocks on the last result. If output storage cannot fit that queue, it blocks after each call and records timing_method=jax_wall_block_each; no atom-count or hardware-specific case limit is applied.

This is a framework-facing steady-state comparison, not a raw-kernel microbenchmark. Each CSV row records its caller-, API-, or JIT-managed allocation_boundary and its timing_scope so differing execution contracts remain explicit.

Method Variants

The NL suite treats each user-facing neighbor-list strategy as a separate method. The naive family contains naive_scalar and naive_tile; the cell-list family contains cell_list_atom_centric and cell_list_pair_centric. cluster_tile is the CUDA/float32, fully periodic path that Morton-sorts atoms into 32-atom clusters and evaluates candidate cluster pairs cooperatively. Batched inputs use the matching batch API where needed, but the CSV method column keeps the concrete strategy visible for plotting and comparison.

The shipped YAML config enables all five strategies. Torch runs all five by default. JAX omits pair-centric from its default expansion; an explicit pair-centric selection runs as coverage-only. JAX includes cluster-tile by default when its CUDA, float32, and fully periodic requirements are met. The bundled H100 snapshot includes current JAX cluster-tile coverage at the 6, 15, and 25 Angstrom cutoffs.

Benchmark Configuration

Parameter	Value
Cutoffs	6.0, 15.0, 25.0 Å
Configured Methods	naive_scalar, naive_tile, cell_list_atom_centric, cell_list_pair_centric, cluster_tile
System Type	CsCl (programmatic), NH₃ (PDB)
Warmup Iterations	3
Timing Iterations	10
Dtype	float32

Running Your Own Benchmarks

Run from the repository root:

RESULT_DIR="$BENCHMARK_SCRATCH/results/manual-nl-run"
python -m benchmarks.neighborlist.benchmark_neighborlist \
    --config benchmarks/neighborlist/benchmark_config.yaml \
    --output-dir "$RESULT_DIR"

For the JAX backend:

RESULT_DIR="$BENCHMARK_SCRATCH/results/manual-nl-run"
python -m benchmarks.neighborlist.benchmark_neighborlist \
    --config benchmarks/neighborlist/benchmark_config.yaml \
    --backend jax \
    --output-dir "$RESULT_DIR"

For direct Warp API timing:

RESULT_DIR="$BENCHMARK_SCRATCH/results/manual-nl-run"
python -m benchmarks.neighborlist.benchmark_neighborlist \
    --config benchmarks/neighborlist/benchmark_config.yaml \
    --backend warp \
    --method cell_list_atom_centric \
    --output-dir "$RESULT_DIR"

Use --method / --methods to restrict the benchmark to particular APIs, for example --method naive_tile cell_list_pair_centric, and --dry-run to inspect the expanded (system, mode, method, cutoff) plan before allocating GPU memory.