Skip to content

Make strided AxisArray data contiguous before serializing - #269

Draft
cboulay wants to merge 1 commit into
devfrom
cboulay/contiguous-before-serialize
Draft

Make strided AxisArray data contiguous before serializing#269
cboulay wants to merge 1 commit into
devfrom
cboulay/contiguous-before-serialize

Conversation

@cboulay

@cboulay cboulay commented Sep 4, 2026

Copy link
Copy Markdown
Member

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__ makes data contiguous when it is neither C- nor F-contiguous, which puts the payload back on the out-of-band path.

layout before after out-of-band buffers
C-contiguous (4 MB) 6.5 us 7.1 us 1 → 1
F-order transposed 6.6 us 6.8 us 1 → 1
strided [:, ::2] (2 MB) 311.9 us 141.5 us 0 → 1
strided [::2, :] (2 MB) 253.5 us 41.1 us 0 → 1

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.ascontiguousarray would have been a serious pessimization on a common case — a transposed AxisArray is not unusual. The guard is not (c_contiguous or f_contiguous), costing two flag reads on the common path.

Scope

  • Only numpy arrays are touched. torch/cupy have their own serialization and no flags attribute to consult.
  • Lives on ArrayWithNamedDims, the shared base, so CoordinateAxis gets it too — a ch axis with a strided data hit the same cliff.
  • The caller's array is never mutated; the copy goes into the returned state dict only.

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 CoordinateAxis is covered.

Full suite green: 431 passed, 1 skipped. The 2 remaining ruff errors in these files (E731 at axisarray.py, E402 in the test) are pre-existing on dev — verified against the pristine versions.

Measurements are single-machine (Darwin, arm64, Python 3.13).

🤖 Generated with Claude Code

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant