Chunked dual transpose - #32
Merged
Merged
Conversation
`_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>
Merged
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.
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.
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: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 @ Bsplits over the contracted axis (sum ofDelta[:, C] @ B[C, :]over column chunks), andB @ Deltasplits 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, sovalues/indicescome 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_vcscwere 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):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:
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.
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:
The budget is one module constant (
_CHUNK_BUDGET_BYTES, 128 MiB) plus a deliberately generous bytes-per-nonzero estimate fortranspose_major's transient sort arrays. Both are judgement calls that would be best introduced by creating auser-facing memory budget.