Make strided AxisArray data contiguous before serializing - #269
Draft
cboulay wants to merge 1 commit into
Draft
Conversation
numpy hands a C- or F-contiguous array to a protocol-5 pickler out-of-band,
and ezmsg's marshal writes those buffers straight into shared memory without
copying them. An array that is neither falls back to an in-band copy through
the pickle stream. Unlike the rest of serialization -- which is a flat ~7us
regardless of payload -- that copy scales with size, and any transformer
emitting a decimated or channel-sliced view hits it silently.
ArrayWithNamedDims.__getstate__ now materializes a contiguous copy in that
case, which puts the payload back on the out-of-band path:
layout before after out-of-band buffers
C-contiguous (4 MB) 6.5us 7.1us 1 -> 1
F-order transposed 6.6us 6.8us 1 -> 1
strided [:, ::2] (2 MB) 311.9us 141.5us 0 -> 1
strided [::2, :] (2 MB) 253.5us 41.1us 0 -> 1
In-band bytes for the strided cases drop from 2,097,501 to 307.
This is transparent rather than a behaviour change: unpickling a strided
array already yields a C-contiguous one, so the receiver sees exactly the
same array either way.
F-contiguous arrays are deliberately excluded. They are already out-of-band,
and forcing C order on them measured ~80x worse -- a blanket
ascontiguousarray would have been a serious pessimization. The guard costs
two flag reads on the common path.
Only numpy arrays are touched; torch/cupy arrays have their own
serialization and no `flags` to consult. The fix lives on the shared base
class, so CoordinateAxis gets it too.
Co-Authored-By: Claude Opus 5 (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.
Draft, same as #268 — parking it for review rather than merge. Independent of #268; they touch different files and can land in either order.
The problem
numpy hands a C- or F-contiguous array to a protocol-5 pickler out-of-band, and ezmsg's marshal writes those buffers straight into shared memory without copying. An array that is neither falls back to an in-band copy through the pickle stream.
That matters because it is the one size-dependent cost in serialization. Everything else is a flat ~7 us regardless of payload. Any transformer emitting a decimated or channel-sliced view (
data[:, ::2],data[::2, :]) hits this silently, and it scales with the array.The fix
ArrayWithNamedDims.__getstate__makesdatacontiguous when it is neither C- nor F-contiguous, which puts the payload back on the out-of-band path.[:, ::2](2 MB)[::2, :](2 MB)In-band bytes for the strided cases drop from 2,097,501 to 307.
Why this is transparent, not a behaviour change
Unpickling a strided array already yields a C-contiguous one — numpy's reduce does
tobytes()on the way out. So the receiver sees exactly the same array either way; it just gets there for a lot less. Verified explicitly in the tests.Why F-order is excluded
This is the part worth a second pair of eyes. F-contiguous arrays are already on the out-of-band path, and forcing C order on them measured ~80x worse (7 us → 582 us on a 4 MB transposed array). A blanket
np.ascontiguousarraywould have been a serious pessimization on a common case — a transposedAxisArrayis not unusual. The guard isnot (c_contiguous or f_contiguous), costing two flag reads on the common path.Scope
flagsattribute to consult.ArrayWithNamedDims, the shared base, soCoordinateAxisgets it too — achaxis with a strideddatahit the same cliff.Testing
7 new tests: all four layouts round-trip out-of-band with correct values, F-order is confirmed not reordered, the source message is confirmed unmutated, and
CoordinateAxisis covered.Full suite green: 431 passed, 1 skipped. The 2 remaining ruff errors in these files (
E731ataxisarray.py,E402in the test) are pre-existing ondev— verified against the pristine versions.Measurements are single-machine (Darwin, arm64, Python 3.13).
🤖 Generated with Claude Code