Benchmarking#

The benchmark/ directory ships a curated set of representative material models together with a sweep harness and a vectorisation-efficiency fitter. The tooling answers the question that comes up whenever a NEML2 user sizes a production run:

How does compute time scale with batch size for a given model on a given device?

The same tooling is the substrate for tracking performance changes across NEML2 releases; runtime-touching PRs are expected to keep the headline numbers from regressing.

Scenarios#

The committed scenarios cover the breadth of the solid-mechanics module, from trivial elasticity to fully-coupled crystal plasticity:

Code name	Description
`elasticity`	Linear isotropic elasticity without model composition
`radret`	$J_2$ viscoplasticity with radial-return integration
`isoharden`	Viscoplasticity with isotropic hardening
`chaboche2`	Chaboche viscoplasticity, 2 back stresses
`chaboche4`	Chaboche viscoplasticity, 4 back stresses
`chaboche6`	Chaboche viscoplasticity, 6 back stresses
`gtntheig`	GTN visco-poroplasticity with thermal eigenstrain
`scpcoup`	Fully-coupled implicit single-crystal plasticity
`scpcoupmult`	Coupled SCP with multiplicative $F = F^e F^p$ decomposition
`scpdecoup`	Decoupled implicit single-crystal plasticity
`scpdecoupexp`	Decoupled SCP with explicit orientation update
`tcpsingle`	Taylor polycrystal with a single orientation
`tcprandom`	Taylor polycrystal with pseudo-random FCC orientations
`mxpc`	Mixed-control polycrystal: `nbatch` scales the sub-batch (grain) axis with a shared global mixed-control state, solved via `SchurComplement`

Each scenario lives under benchmark/<name>/model.i as an HIT input file with two HIT-substituted parameters: ${nbatch} (the dynamic batch size) and ${device} (cpu or cuda). All scenarios except tcprandom advance 100 prescribed time steps; tcprandom advances 500 to simulate severe plastic deformation from the standard rolling experiment.

Vectorisation-efficiency model#

For a fixed device and model, compute time as a function of batch size $N$ is fitted to ANL-24/43 eq 7.1:

\[ t(N;\, t_0,\, N^*) = t_0 \left( \frac{\langle N - N^* \rangle}{N^*} + 1 \right) \]

where $\langle x \rangle = \max(x, 0)$ is the Macaulay bracket, $t_0$ is the per-call cost when the device’s vectorisation bandwidth is not yet saturated (a constant time across all small batches), and $N^*$ is the optimal batch size at which the device starts trading constant time for linear-in-$N$ time.

Two limits of the model are intuitive:

$N \ll N^*$: $t \to t_0$. Vectorisation lanes are unsaturated; per-call cost is dominated by launch overhead, fixed setup, and whatever per-step constant work the model carries. Increasing batch costs nothing.
$N \gg N^*$: $t \to t_0 \cdot N / N^*$. The device is saturated; per-call cost scales linearly with the batch.

The fitted $(t_0,\, N^*)$ pair characterises the model+device pair in one of two useful ways:

On a fixed device, $t_0$ sets the floor on per-call cost — every measurement at $N \le N^*$ sits there, so a smaller $t_0$ means a cheaper baseline.
For a fixed model, $N^*$ is the largest batch that still pays only $t_0$, so a larger $N^*$ means more useful work per call before the device falls out of the free region.

The peak throughput of a model+device pair is

\[ \mathrm{throughput} = \frac{N^*}{t_0} \]

(material updates per ms), since by construction $t(N^*) = t_0$. This is the right single-number summary when comparing two runtimes — pick each runtime’s own optimal batch, and the user can run the model at that batch to get the headline rate. Comparing at a fixed batch can flatter or penalise either side depending on whether that batch happens to land in one runtime’s free region.

Running the sweep#

The sweep tool is invoked through python -m benchmark.sweep:

# Init a fresh result folder (just records sweep config + env snapshot)
python -m benchmark.sweep init --device cuda

# Run every scenario into that folder (or one folder at a time)
python -m benchmark.sweep run --output-dir benchmark/results/<TS>_cuda_float64_aoti

# Rebuild summary.csv from per-scenario CSVs
python -m benchmark.sweep summarize --output-dir <DIR>

# Fit eq 7.1 per scenario into fitting.csv
python -m benchmark.sweep fit --output-dir <DIR>

# End-to-end shortcut (init + run every discovered scenario + summarize)
python -m benchmark.sweep all --device cuda

Per-scenario stop conditions trip the inner adaptive-batch loop whichever comes first: (a) log-log slope $\ge 0.85$ for two consecutive points (the model has entered its linear regime so more data points add little information), (b) measured median call time exceeds 30 s (absolute wall-time saturation), © CUDA OOM / host MemoryError, or (d) batch exceeds the --max-batch cap (default 65536).

Each batch is compiled separately, with the scenario’s example batch shape resolved at the measured batch size. Per-batch compile mirrors what a production user gets from neml2-compile invoked at the batch they actually intend to run; reusing one compile across batches would be cheaper but can mis-report scaling behaviour at batches far from the example.

Output#

Each run populates four file kinds under <output-dir>:

File	Contents
`metadata.json`	Host + GPU snapshot, sweep config, per-scenario batch ranges, last-run wall time
`<scenario>.csv`	One row per measured batch: median / mean / std / min / max / p10 / p90 / RSS / peak CUDA memory
`<scenario>.runs.csv`	Raw per-iteration timings (warmup discarded, all measurement runs kept)
`summary.csv`	Aggregate view across every scenario, regenerated by `summarize`

fit adds a fifth file:

File	Contents
`fitting.csv`	Per-scenario $(t_0, N^*)$, log-space RMS residual, and number of fit points

The CUDA compile pipeline saturates available CPU cores while running; co-locating a CPU sweep on the same box double-occupies the cores and taints both. Run one device at a time on a quiet machine. The harness itself does not detect contention.