Skip to content

Chunked dual transpose - #32

Merged
aarmey merged 5 commits into
mainfrom
perf/chunked-dual-transpose
Sep 4, 2026
Merged

Chunked dual transpose#32
aarmey merged 5 commits into
mainfrom
perf/chunked-dual-transpose

Conversation

@fishidaho

Copy link
Copy Markdown
Contributor

Regroup a chunk at a time for the misaligned matmul direction

What's wrong

_get_dual (src/vsparse/_vcs_matmul.py) built a full opposite-format copy of the array the first time a normalized-view matmul needed the direction the storage isn't aligned for:

def_get_dual(nview):
ifnview._dual_arrisNone:
nview._dual_arr=nview._arr._transpose_major()
returnnview._dual_arr

A second complete copy of the dataset, allocated silently inside a matmul.

What this changes

The regrouping doesn't have to happen all at once. Delta @ B splits over the contracted axis (sum of Delta[:, C] @ B[C, :] over column chunks), and B @ Delta splits the same way over row chunks. So a contiguous range of major slices is transposed on its own, fed to the same major-aligned kernel, accumulated into the shared output, and dropped. Peak memory becomes one chunk's regrouping — set by a byte budget — instead of the whole array's.

_VCSBase._major_range(start, stop) is the new chunking primitive. A contiguous major range is already contiguous in every stored array, so values/indices come back as views (no copy) and only the two small pointer arrays are rebuilt. That's what makes per-chunk work cheap enough to be the default.

Dead code removed: _matmul_vcsr/_rmatmul_vcsc were thin wrappers that only existed to allocate the output the entry points now own.

The trade-off here

Chunking is not a free win in every regime. Measured on 72M nonzeros (30000×3000, VCSC, self @ B):

first callpeak RSSsteady-state per call
chunked (new default)10.30s+0 MB9.83s
pinned full dual (old behavior)14.83s+851 MB0.12s

Max absolute difference between the two: 1.9e-14.

So the new default is faster to a first result and costs no extra memory, but repeated misaligned products against a cached dual are ~55× faster, because chunking redoes the regrouping every call. Power-iteration workloads (PCA, parafac2) hit exactly that pattern.

Two things address it:

  1. An array whose entire regrouping already fits inside one chunk's budget is still cached automatically. Transposing it per call would allocate exactly that much anyway, so keeping it costs no additional peak memory.

  2. view.pin_dual() is the explicit opt-in past that point, for callers who will run many misaligned products and know the memory is affordable.

Net effect, verified at both ends:

small (fits one chunk) nnz= 0.5M dual cached: True +RSS 0 MB
large (multi-chunk) nnz= 72.0M dual cached: False +RSS 0 MB

The budget is one module constant (_CHUNK_BUDGET_BYTES, 128 MiB) plus a deliberately generous bytes-per-nonzero estimate for transpose_major's transient sort arrays. Both are judgement calls that would be best introduced by creating a
user-facing memory budget.

fishidahoand others added 2 commits September 2, 2026 17:22
`_get_dual` built a full opposite-format copy of the array the first time a
normalized-view matmul needed the direction the storage isn't aligned for.
No size gate, no opt-in, no warning: a second complete copy of the dataset,
allocated silently inside a matmul. On a 72M-nonzero array that's +851 MB;
at atlas scale it's simply fatal on a machine the original already fits.
The regrouping doesn't have to happen all at once. `Delta @ B` splits over
the contracted axis (sum of `Delta[:, C] @ B[C, :]` over column chunks) and
`B @ Delta` splits the same way over row chunks, so a contiguous range of
major slices can be transposed alone, fed to the same major-aligned kernel,
accumulated into the shared output, and dropped. Peak memory is then one
chunk's regrouping -- set by a byte budget -- instead of the whole array's.
`_VCSBase._major_range` is the chunking primitive: a contiguous major range
is already contiguous in every stored array, so `values`/`indices` come back
as views and only the two small pointer arrays are rebuilt.
One case still caches without being asked: an array whose *entire*
regrouping already fits inside a single chunk's budget. Transposing it per
call would allocate exactly that much anyway, so keeping it costs no extra
peak memory -- and it matters, because repeated misaligned products run ~55x
faster against a cached dual than re-regrouping every call. Past one chunk
nothing is cached, and `pin_dual()` is the explicit opt-in for callers who
want the cached dual anyway and know the memory is affordable.
Measured on 72M nonzeros (30000x3000, VCSC, `self @ B`):
chunked first call 10.30s, +0 MB peak RSS
pinned first call 14.83s, +851 MB peak RSS
max abs difference between the two: 1.9e-14
so the default is now both faster to a first result and free of the extra
copy; steady-state repeated calls are where pinning wins, hence the opt-in.
`test_dual_array_is_built_lazily_and_cached` becomes two tests, since the
behavior it pinned is exactly what changed: small arrays still cache lazily,
and multi-chunk arrays cache nothing.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The one conflict in the wave, and a benign one: this branch added
`_major_range` immediately before `__getitem__`, #23 added `_select_minor`
in the same place and rewrote `__getitem__`. The two methods are independent
-- a contiguous major range that returns views, versus an arbitrary
minor-axis selection that filters and remaps indices -- so both are kept,
and `__getitem__` is taken from main unchanged.
They now overlap in purpose for one case: `__getitem__` has a native path
for a contiguous slice, which `_major_range` also covers. `_major_range`
stays, because it is the reason the chunked matmul is affordable --
`values`/`indices` come back as views into the parent's buffers, where
`_select_major` gathers. Verified still true after the merge
(`np.shares_memory` on both), and a test now pins the two against each
other: if they ever disagree, the chunked matmul is computing against a
different sub-array than a caller would get.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@fishidahofishidaho mentioned this pull request Sep 3, 2026
The public pin_dual had no caller outside its own test. The internal
_build_dual it wrapped stays, since that is how a whole-array regrouping
gets cached when it fits one chunk's budget.
Chunking tests fold into the properties that matter: chunks tile the axis
and respect the budget, _major_range shares buffers and matches ordinary
slicing, results are invariant to chunk size, and peak memory tracks the
budget rather than the array.
@fishidahofishidaho self-assigned this Sep 4, 2026
@fishidaho
fishidaho marked this pull request as ready for review September 4, 2026 00:26
@aarmey
aarmey merged commit e8f010f into mainSep 4, 2026
3 checks passed
@aarmey
aarmey deleted the perf/chunked-dual-transpose branch September 4, 2026 01:52
Sign up for freeto 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.

2 participants

@fishidaho@aarmey