# Benchmarks

Performance benchmarks for ALCHEMI Toolkit-Ops GPU kernels, covering both
the **Torch** and **JAX** backends. These pages present scaling behaviour
across system sizes, constant-workload sweeps, and batch-scaling modes for
each module.

```{toctree}
:maxdepth: 1

neighborlist
electrostatics
dftd3
dynamics
```

## Methodology

### Timing

| Backend | Mechanism | Details |
|---------|-----------|---------|
| **Torch** | CUDA Events | `start_event.record()` / `end_event.record()` around the kernel call, then `torch.cuda.synchronize()`. No host-side sync inside the timed loop. |
| **JAX** | Wall-clock + `block_until_ready` | Warmups block to exclude tracing/compilation. The standard timed loop dispatches all calls and blocks on the final result; allocation-heavy rows may retry with a block after each call. Rows record the exact `timing_method`. |

The shipped benchmark YAMLs use the publication timing protocol by default:
three warmup iterations and ten timed iterations on the full benchmark grid.
Both values are configurable via `parameters.warmup_runs` and
`parameters.timing_runs`, or via the shared `--warmup-runs` / `--timing-runs`
CLI flags. The reported value is the mean steady-state time per call after
warmup/compile/load work is excluded.

### Scaling modes

Each benchmark page presents three scaling modes:

System Size Scaling
: Fix batch size = 1, sweep the number of atoms per system.

Constant Workload
: Target a total atom count while sweeping batch size (fewer atoms per system,
  more systems). Discrete system sizes and integer batch counts mean the actual
  total can be below the target, especially for CsCl supercells. Plots use the
  actual `total_atoms` recorded in each CSV row.

Batch Scaling
: Fix atoms per system, sweep batch size. Reveals how well the kernel
  amortises fixed overhead across systems.

### Systems

| System | Source | Description |
|--------|--------|-------------|
| **CsCl** | Programmatic supercell builder in `benchmarks.suite_systems` | CsCl/B2 primitive-cubic supercells (2 atoms/unit cell, cubic 4.119 Å). |
| **NH₃** | PDB files generated by `benchmarks/nh3/generate_pbc_pdbs.sh` (requires [packmol](https://m3g.github.io/packmol/)) | Periodic ammonia boxes at powers-of-two atom counts. |

### Hardware

| GPU | VRAM | Architecture |
|-----|------|--------------|
| NVIDIA H100 80GB HBM3 | 80 GB HBM3 | Hopper |

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

```{note}
**Precision**: NL and DFT-D3 use ``float32``; electrostatics
(PME / Ewald) use ``float64``; charges are always ``float64``.

**Missing plotted points** are rows where the benchmark wrote `success=False`
with a concrete `error_type` such as `OutOfMemoryError`, `SkippedByPolicy`,
`SkippedAfterOOM`, or `UnsupportedConfiguration`. The latter records a public
API capability or launch-safety guard; it does not substitute another method.
Hardware-dependent OOM limits are not hardcoded into the suite; failed and
skipped rows remain in the CSVs, are not drawn as data, and break plotted lines
at their coordinates.

Successful JAX rows timed with per-call blocking are shown as separate line
segments from batch-dispatch timings; the CSV ``timing_method`` column records
the boundary used for every point.
```

## Running the full suite

The three modules share a single entry point,
``benchmarks.benchmark_suite``, which loads each per-module
``benchmark_config.yaml`` and dispatches in-process. Use a fresh scratch
directory for each reportable run, then promote a complete, reviewed CSV set
into the documentation only after validation.

::::{dropdown} Commands and controls

```bash
RESULT_DIR="$BENCHMARK_SCRATCH/results/reportable-run"

# All three modules, Torch backend (default)
python -m benchmarks.benchmark_suite --benchmark all \
    --run-dir "$RESULT_DIR"

# All three modules, JAX backend
python -m benchmarks.benchmark_suite --benchmark all --backend jax \
    --run-dir "$RESULT_DIR"

# Or a subset:
python -m benchmarks.benchmark_suite --benchmark nl d3 --backend torch
```

Common controls:

- `--method` / `--methods` selects named APIs.
- `--dry-run`, `--list`, and `--count` inspect the expanded plan without GPU
  allocation.
- `--plot-only RESULTS_DIR` rebuilds figures from an existing complete CSV set.
- `--backend warp` is available for neighbor lists; D3 and electrostatics use
  Torch or JAX.
- Electrostatics component profiling is optional; reportable runs time the full
  energy, force, and charge-gradient workload once per iteration.
::::

Before publishing plots, the suite checks case coverage and confirms that the
run, source, hardware, inputs, and runtime settings belong together. Backend
reruns replace only that backend's rows inside the same validated run rather
than mixing unrelated CSV snapshots.

## Published snapshot

The bundled H100 snapshot uses the full YAML grids with three warmups and ten
timed iterations. Failed cases remain explicit CSV rows; plots leave those
coordinates empty and do not connect lines across them. Collection may be
sharded across compatible GPUs; CSV timings always
represent steady-state calls, not queue, compilation, or environment setup.

See [Benchmark Results](benchmark_results/README.md) for the current run ID,
row and failure counts, CSV schema, scratch-only cluster setup, sharding rules,
and the detailed reproducibility runbook. Module-specific commands remain on
the neighbor-list, DFT-D3, and electrostatics pages.