Repository navigation
perf: marshal the C read path's dicts in C (numpy C-API), ~1.5x reads - #44
Merged
Merged
Conversation
The C reader handed each parsed frame back to Python by walking the DictEntry linked list one field at a time through ctypes (c_to_py_dict): ~1.09M ctypes.cast + 737k copy.copy across a 76k-frame read, pure boundary overhead that dominated files with many small frames. Add a CPython C-API entry point (_extxyz.read_frame in libextxyz/pyext.c) that builds the info/arrays dicts of numpy arrays / scalars directly in C and frees the C dicts, replacing that per-node loop. It is bit-identical to the ctypes path (tests/test_marshal_parity.py checks both backends and both use_regex values). cextxyz.read_frame_dicts dispatches to it when available and falls back to read_frame_dicts_ctypes when the extension was built without numpy; EXTXYZ_LEGACY_MARSHAL=1 forces the legacy path. The shared C core is untouched (zero diff to extxyz.c / the extxyz.h ABI): pyext.c is compiled ONLY into the _extxyz Python extension, never the standalone libextxyz shared library or the C/Fortran/Julia consumers. numpy becomes a build-time dependency, guarded non-fatally so a numpy-less standalone build still configures and uses the ctypes path. On Windows a second .def (_extxyz_pyext.def) exports PyInit__extxyz so the module is importable; CI asserts every wheel shipped the C path. Measured on a 76,346-frame, ~27-atom/frame set: dict-level parse 3.6s -> 2.3s (~1.5x), full ASE read 4.8s -> 3.3s; 0 leaks over 12k frames across success/EOF/error paths. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Move the C read path's dict marshalling into the
_extxyzextension. After the C reader parses a frame, it previously handed the data back to Python by walking theDictEntrylinked list one field at a time throughctypes(c_to_py_dict) — ~1.09Mctypes.cast+ 737kcopy.copyacross a 76k-frame read, pure boundary overhead that dominates files with many small frames.This adds a CPython C-API entry point (
_extxyz.read_frame, inlibextxyz/pyext.c) that builds theinfo/arraysdicts of numpy arrays / scalars directly in C and frees the C dicts, replacing the per-node Python loop.Why it's safe
tests/test_marshal_parity.pyreads each fixture through both backends (and bothuse_regexvalues) and asserts equal keys, dtypes, shapes, types, and bit-identical floats. Verified additionally on all 76,346 frames of a real training set: 0 mismatches.cextxyz.read_frame_dictsdispatches to the C path only when present; otherwise it uses the preservedread_frame_dicts_ctypes.EXTXYZ_LEGACY_MARSHAL=1forces the legacy path.extxyz.c/ theextxyz.hABI.pyext.cis compiled only into the_extxyzPython extension — never the standalonelibextxyzshared library or the C/Fortran (fextxyz/QUIP)/Julia consumers.leaksover 12k frames across success/EOF/scattered-string/error paths: 0 leaks, 0 bytes.Build / packaging
meson.build: a numpy-less standalone C/Fortran/Julia build still configures and just uses the ctypes path. Verified both a numpy-present and a numpy-lessmeson setup/compile._extxyz_pyext.def) exportsPyInit__extxyzso the module is importable; the numpy-less build keeps the original.def(no unresolved-export link error).tools/assert_c_read_path.py), so a silent ctypes fallback fails the build instead of shipping quietly slow.Measured (76,346 frames, ~27 atoms/frame)
read_dicts)format='cextxyz')Most visible on many-small-frame files / rich comment lines, where per-frame overhead — not per-atom parsing — dominates; on the large single-frame Cu benchmark the effect is small.
Notes for reviewers
benchmarks/(bench_mad.py,scope_comment.py,verify_marshal.py).🤖 Generated with Claude Code