From 1db22740a5b47ead3d4c6b2825e4a1af9fcfac13 Mon Sep 17 00:00:00 2001 From: James Kermode Date: Tue, 9 Jun 2026 23:15:20 +0100 Subject: [PATCH] perf: buffer per-atom output in the C writer (byte-identical) Profiling the write path: our C writer already beats ASE's built-in extxyz writer by ~2.5x and extxyz-ng by ~1.5x, but writing a 200k-atom frame took ~266 ms vs ~33-50 ms to read it. The cost is the per-atom output loop, which did one fprintf per value (~15 per atom incl. separators/newline for a 7-column frame) -- each paying format-string parsing and a flockfile/funlockfile pair. A micro-benchmark showed the dominant cost is the %.8f float formatting itself, not I/O: buffering alone is ~1.14x while a custom integer formatter would be ~2x but cannot be byte-identical to printf's round-half-to-even (mismatches ~1 in 5M) without a fiddly correctly-rounded formatter. So this lands the safe, byte- identical win only: build each line in a growable memory buffer with snprintf and fwrite it in blocks, flushing at line boundaries. Same output bytes, far fewer locked stdio calls -- ~1.26x on a 200k write (266 -> 212 ms). Output is byte-identical (verified old-vs-new across signs, -0.0, rounding edges, ints/bools/strings, multi-column, and a custom format_dict); a golden + round-trip test locks it. The format_dict override (#22) path is unchanged. New benchmarks/bench_write.py reproduces the comparison vs ASE and extxyz-ng. Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 10 ++++ benchmarks/bench_write.py | 112 +++++++++++++++++++++++++++++++++++++ libextxyz/extxyz.c | 68 +++++++++++++++------- tests/test_write_buffer.py | 65 +++++++++++++++++++++ 4 files changed, 234 insertions(+), 21 deletions(-) create mode 100644 benchmarks/bench_write.py create mode 100644 tests/test_write_buffer.py diff --git a/README.md b/README.md index 7ac7128..471d3f8 100644 --- a/README.md +++ b/README.md @@ -110,6 +110,16 @@ python benchmarks/bench_read.py --max-atoms 200000 --repeats 3 python benchmarks/plot_bench.py ``` +### Writing + +Writing is bounded by formatting the per-atom floats (`"%16.8f"`), not I/O. The C +writer builds each line in a memory buffer and `fwrite`s it in blocks rather than +one `fprintf` per value (same output bytes, fewer locked stdio calls). It writes a +200k-atom frame in ~210 ms — about **2.5× faster than ASE's built-in `extxyz` +writer** (and faster than `extxyz-ng`); the pure-Python (`np.savetxt`) writer +matches ASE. `benchmarks/bench_write.py` reproduces the comparison (and times +`extxyz-ng` if `EXTXYZ_NG_PYTHON` points at a venv with it installed). + ## `libextxyz` C library and standalone executables The C parser, the standalone `libextxyz` shared library, and the C-only diff --git a/benchmarks/bench_write.py b/benchmarks/bench_write.py new file mode 100644 index 0000000..50af3b6 --- /dev/null +++ b/benchmarks/bench_write.py @@ -0,0 +1,112 @@ +"""Benchmark: extxyz writers vs ASE built-in and the Rust extxyz-ng. + +Mirrors ``bench_read.py``. Generates synthetic frames, then times writing each +with: + +* ``extxyz.write_dicts(use_cextxyz=True)`` — this repo's C writer +* ``extxyz.write_dicts(use_cextxyz=False)`` — the pure-Python (np.savetxt) writer +* ``ase.io.write(format='extxyz')`` — ASE's built-in writer (baseline) +* ``ase.io.write(format='cextxyz')`` — the ase-extxyz plugin +* ``extxyz-ng`` ``write_frame`` — the Rust port, if ``EXTXYZ_NG_PYTHON`` + points at a venv that has it (run in a subprocess, like ``test_bench_vs_rust``). + +Run:: + + python benchmarks/bench_write.py [--max-atoms 200000] [--repeats 5] + EXTXYZ_NG_PYTHON=/path/to/ng/bin/python python benchmarks/bench_write.py +""" +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +import tempfile +import textwrap +import time +from pathlib import Path + +import numpy as np +import ase.io + +import extxyz +import ase_extxyz.io # noqa: F401 (registers the 'cextxyz' format) + +from bench_read import make_xyz # reuse the fixture generator + + +def _best(fn, repeats): + best = float('inf') + for _ in range(repeats): + t0 = time.perf_counter() + fn() + best = min(best, time.perf_counter() - t0) + return best + + +def _time_ng(src_path, repeats): + """Time extxyz-ng's write_frame in its own interpreter (read a frame from + ``src_path``, then time writing it). Returns ms or None if unavailable.""" + ng = os.environ.get('EXTXYZ_NG_PYTHON') + if not ng: + return None + snippet = textwrap.dedent(f""" + import time, tempfile, os, extxyz + frame = extxyz.read_frame_from_file({str(src_path)!r}) + out = tempfile.mktemp(suffix='.xyz') + def run(): + with open(out, 'wb') as fh: + extxyz.write_frame(fh, frame) + run() + best = None + for _ in range({repeats}): + t0 = time.perf_counter(); run(); dt = time.perf_counter()-t0 + best = dt if best is None else min(best, dt) + os.path.exists(out) and os.remove(out) + print(best) + """) + env = {k: v for k, v in os.environ.items() if k != 'PYTHONPATH'} + out = subprocess.run([ng, '-c', snippet], env=env, + capture_output=True, text=True, timeout=300) + if out.returncode != 0: + return None + return float(out.stdout.strip().splitlines()[-1]) * 1e3 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('--max-atoms', type=int, default=200_000) + ap.add_argument('--repeats', type=int, default=5) + args = ap.parse_args() + + sizes = [n for n in (1000, 4000, 16000, 64000, args.max_atoms) if n <= args.max_atoms] + + hdr = (f'{"N atoms":>9} {"MB":>6} {"our C":>8} {"pyPython":>9} ' + f'{"ASE":>8} {"plugin":>8} {"ng":>8} {"vs ASE":>7} {"vs ng":>6}') + print(hdr); print('-' * len(hdr)) + + with tempfile.TemporaryDirectory() as tmp: + for n in sizes: + src = Path(tmp) / f'src_{n}.xyz' + mb = make_xyz(src, n) / 1e6 + frames = extxyz.read_dicts(str(src)) + frames = frames if isinstance(frames, list) else [frames] + atoms = ase.io.read(str(src), format='cextxyz', index=':') + out = Path(tmp) / 'out.xyz' + + t_c = _best(lambda: extxyz.write_dicts(str(out), frames, use_cextxyz=True), args.repeats) + t_py = _best(lambda: extxyz.write_dicts(str(out), frames, use_cextxyz=False), args.repeats) + t_ase = _best(lambda: ase.io.write(str(out), atoms, format='extxyz'), args.repeats) + t_plug = _best(lambda: ase.io.write(str(out), atoms, format='cextxyz'), args.repeats) + t_ng = _time_ng(src, args.repeats) + + ng_s = f'{t_ng*1e3:8.1f}' if t_ng else f'{"n/a":>8}' + vs_ng = f'{t_ng*1e3/(t_c*1e3):6.2f}x' if t_ng else f'{"-":>6}' + print(f'{n:>9} {mb:>6.2f} {t_c*1e3:>8.1f} {t_py*1e3:>9.1f} ' + f'{t_ase*1e3:>8.1f} {t_plug*1e3:>8.1f} {ng_s} ' + f'{t_ase/t_c:>6.2f}x {vs_ng}') + + +if __name__ == '__main__': + sys.path.insert(0, str(Path(__file__).resolve().parent)) + main() diff --git a/libextxyz/extxyz.c b/libextxyz/extxyz.c index eb08e65..a5007aa 100644 --- a/libextxyz/extxyz.c +++ b/libextxyz/extxyz.c @@ -1241,7 +1241,37 @@ int extxyz_write_ll_fmt(FILE *fp, int nat, DictEntry *info, DictEntry *arrays, free(quoted_properties_str); free(properties_str); - // write per-atom data + // write per-atom data. Build each line in a growable memory buffer with + // snprintf and fwrite it in blocks, instead of one (FILE-locked) fprintf per + // value — same output bytes, fewer locked stdio calls. Flush at line + // boundaries once the buffer passes WBUF_FLUSH so a line is never split. + size_t wbuf_cap = 1u << 16, wbuf_n = 0; + const size_t WBUF_FLUSH = 1u << 15; + char *wbuf = (char *) malloc(wbuf_cap); + if (! wbuf) { return 7; } + // append `fmt`-formatted `val`, growing the buffer (and re-formatting) only + // if it didn't fit — for a flushed buffer it almost always fits first time. + #define WB_FMT(fmt, val) do { \ + int _l = snprintf(wbuf + wbuf_n, wbuf_cap - wbuf_n, (fmt), (val)); \ + if (_l < 0) { free(wbuf); return 7; } \ + if ((size_t)_l >= wbuf_cap - wbuf_n) { \ + while (wbuf_n + (size_t)_l + 1 > wbuf_cap) wbuf_cap *= 2; \ + char *_nb = (char *) realloc(wbuf, wbuf_cap); \ + if (! _nb) { free(wbuf); return 7; } \ + wbuf = _nb; \ + snprintf(wbuf + wbuf_n, wbuf_cap - wbuf_n, (fmt), (val)); \ + } \ + wbuf_n += (size_t)_l; \ + } while (0) + #define WB_CH(c) do { \ + if (wbuf_n + 1 > wbuf_cap) { \ + wbuf_cap *= 2; \ + char *_nb = (char *) realloc(wbuf, wbuf_cap); \ + if (! _nb) { free(wbuf); return 7; } \ + wbuf = _nb; \ + } \ + wbuf[wbuf_n++] = (c); \ + } while (0) for (int i_at=0; i_at < nat; i_at++) { for (DictEntry *entry = arrays; entry; entry = entry->next) { @@ -1249,26 +1279,20 @@ int extxyz_write_ll_fmt(FILE *fp, int nat, DictEntry *info, DictEntry *arrays, switch(entry->data_t) { case data_i: for (int i_col=0; i_col < ncols; i_col++) { - fprintf(fp, FMT_I, ((int *)(entry->data))[i_at*ncols+i_col]); - if (i_col < ncols-1) { - fprintf(fp, " "); - } + WB_FMT(FMT_I, ((int *)(entry->data))[i_at*ncols+i_col]); + if (i_col < ncols-1) { WB_CH(' '); } } break; case data_f: for (int i_col=0; i_col < ncols; i_col++) { - fprintf(fp, FMT_F, ((double *)(entry->data))[i_at*ncols+i_col]); - if (i_col < ncols-1) { - fprintf(fp, " "); - } + WB_FMT(FMT_F, ((double *)(entry->data))[i_at*ncols+i_col]); + if (i_col < ncols-1) { WB_CH(' '); } } break; case data_b: for (int i_col=0; i_col < ncols; i_col++) { - fprintf(fp, FMT_B, ((int *)(entry->data))[i_at*ncols+i_col] ? "T" : "F"); - if (i_col < ncols-1) { - fprintf(fp, " "); - } + WB_FMT(FMT_B, ((int *)(entry->data))[i_at*ncols+i_col] ? "T" : "F"); + if (i_col < ncols-1) { WB_CH(' '); } } break; case data_s: @@ -1279,21 +1303,23 @@ int extxyz_write_ll_fmt(FILE *fp, int nat, DictEntry *info, DictEntry *arrays, const char *s = (entry->n_in_row < 0) ? (const char *)entry->data + (size_t)(i_at*ncols+i_col)*(-entry->n_in_row) : ((char **)(entry->data))[i_at*ncols+i_col]; - fprintf(fp, FMT_S, s); - if (i_col < ncols-1) { - fprintf(fp, " "); - } + WB_FMT(FMT_S, s); + if (i_col < ncols-1) { WB_CH(' '); } } break; default: + free(wbuf); return 6; } - if (entry->next) { - fprintf(fp, " "); - } + if (entry->next) { WB_CH(' '); WB_CH(' '); WB_CH(' '); } } - fprintf(fp, "\n"); + WB_CH('\n'); + if (wbuf_n >= WBUF_FLUSH) { fwrite(wbuf, 1, wbuf_n, fp); wbuf_n = 0; } } + if (wbuf_n) { fwrite(wbuf, 1, wbuf_n, fp); } + free(wbuf); + #undef WB_FMT + #undef WB_CH return 0; } diff --git a/tests/test_write_buffer.py b/tests/test_write_buffer.py new file mode 100644 index 0000000..cd8dac5 --- /dev/null +++ b/tests/test_write_buffer.py @@ -0,0 +1,65 @@ +"""The buffered C writer must produce byte-identical output and round-trip. + +The per-atom write loop builds lines in a memory buffer and `fwrite`s them in +blocks instead of one `fprintf` per value. That's purely an I/O change, so the +bytes must be exactly what the per-cell `fprintf` writer produced — this golden +locks the format, and the round-trip checks values survive a write→read. +""" +import numpy as np + +from extxyz import Frame, read_dicts, write_dicts + + +def _frame(): + return Frame( + natoms=2, cell=np.diag([10.0, 11.0, 12.0]), + pbc=np.array([True, False, True]), + info={"energy": -1.5}, + arrays={"species": np.array(["H", "Cu"]), + "pos": np.array([[1.123456789, -0.00623274, 0.0], + [-20.0, 3.5, 12.03120730]]), + "z": np.array([1, 29], dtype=np.int32)}) + + +GOLDEN = ( + '2\n' + 'energy=-1.50000000 ' + 'Lattice="10.00000000 0.00000000 0.00000000 0.00000000 11.00000000 ' + '0.00000000 0.00000000 0.00000000 12.00000000" pbc=[T, F, T] ' + 'Properties=species:S:1:pos:R:3:z:I:1\n' + 'H 1.12345679 -0.00623274 0.00000000 1\n' + 'Cu -20.00000000 3.50000000 12.03120730 29\n' +) + + +def test_write_bytes_are_golden(tmp_path): + p = tmp_path / "g.xyz" + write_dicts(p, [_frame()], use_cextxyz=True) + assert p.read_text() == GOLDEN + + +def test_write_roundtrips(tmp_path): + f = _frame() + p = tmp_path / "rt.xyz" + write_dicts(p, [f], use_cextxyz=True) + back = read_dicts(p, use_cextxyz=True) + assert list(back.arrays["species"]) == ["H", "Cu"] + assert list(back.arrays["z"]) == [1, 29] + # pos round-trips to the written 8-decimal precision + np.testing.assert_allclose(back.arrays["pos"], f.arrays["pos"], atol=1e-8) + assert (back.pbc == [True, False, True]).all() + + +def test_write_block_boundary(tmp_path): + """A frame large enough to cross the internal flush threshold many times + still round-trips exactly (exercises the buffer grow/flush path).""" + n = 20000 + rng = np.random.default_rng(0) + f = Frame(natoms=n, cell=np.eye(3) * 50, pbc=np.array([True] * 3), info={}, + arrays={"species": rng.choice(["H", "C", "N", "O"], size=n), + "pos": rng.random((n, 3)) * 40 - 20}) + p = tmp_path / "big.xyz" + write_dicts(p, [f], use_cextxyz=True) + back = read_dicts(p, use_cextxyz=True) + assert list(back.arrays["species"]) == list(f.arrays["species"]) + np.testing.assert_allclose(back.arrays["pos"], f.arrays["pos"], atol=1e-8)