Build permuted_dense inverse arrays lazily - #108
Open
Transurgeon wants to merge 3 commits into
Open
Transurgeon wants to merge 3 commits into
Transurgeon wants to merge 3 commits into
Conversation
…blocks Every permuted_dense allocated col_inv (global n ints) and row_inv (global m ints) at construction. On the dense left_matmul kron path this metadata dominates: p independent blocks each carry a full n_vars-sized col_inv, O(p*n_vars + p^2*m0) ints per node against exact-nnz values (5.9 GB engine peak on the OptimalAdvertising benchmark where true nnz needs ~5 MB). This introduces a compact PD variant whose storage stays proportional to the block rather than the global shape, scoped to the kron-csc family: - new_permuted_dense_compact leaves both inverse arrays NULL; the default constructor is unchanged, so every other PD keeps today's layout and O(1) code paths. - BA_pd_csc_alloc builds its output compact, and compactness propagates through copy_sparsity/transpose/index chains of compact sources. - Init-path membership tests (idxs_hits_set callers, index_pd_alloc) gate on a NULL inv and fall back to binary-search scans of the sorted perms (sorted_pos / sorted_hits in utils). - Eval-path consumers ensure the operand-side array at alloc time only when the product is non-empty (permuted_dense_ensure_col_inv / _row_inv), keeping every fill kernel untouched (asserted in BA_pd_csc_fill_values); index_pd_fill_values of a compact source reads a source-row map precomputed into kernel_iwork at alloc. - The mutable kron scratch keeps its arrays (its kernels write col_inv). New peak-memory regression test builds the row-sum + col-sum Jacobian shape at 128x128: init peak drops from 1324 to 292 bytes/var (21.7 MB to 4.8 MB); the test asserts < 400 bytes/var. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Transurgeon
force-pushed
the
lazy-inv-arrays
branch
from
August 2, 2026 23:16
12d8c9d to
960e02b
Compare
Transurgeon
marked this pull request as draft
August 3, 2026 00:39
# Conflicts: # src/utils/permuted_dense.c # tests/all_tests.c # tests/utils/test_permuted_dense.h
Replace the two-mode design (full vs compact PDs, with compactness propagating through copies and NULL-guarded fallbacks at every reader) with one rule: no constructor builds col_inv / row_inv; an alloc whose paired fill reads one calls permuted_dense_ensure_col_inv / _row_inv on that operand, and the fills assert it is present. row_gather is the one reader that only needs membership at alloc (its fill copies through bound_iwork), so it binary-searches row_perm via sorted_pos instead. Deletes new_permuted_dense_compact, pd_is_compact, sorted_hits, and the unreachable sorted_pos branches in coalesce_spd_scatter and BTA_pd_spd_fill_values. Same Jacobian-init saving on the kron-csc path (1324 -> 292 B/var at 128x128), and it also covers PDs no fill ever reads inverse arrays of (kron-pd/spd outputs, ATA raw blocks): an exp(A @ X) Hessian init drops 94.9 MB -> 62.8 MB (PR-as-written: 85.3). CVXPY OptimalAdvertising (m=250, n=1000) through the DIFFENGINE canon backend: engine peak 5973 MB -> 628 MB, get_problem_data 1.06 s -> 0.62 s (median of 3), and P/A/c/b bit-identical to main. CVaRBenchmark and SemidefiniteProgramming peaks are unchanged (a different cause). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Transurgeon
marked this pull request as ready for review
October 6, 2026 01:27
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.
Motivation
Every
permuted_dense(PD) used to build two dense inverse permutations at construction:col_inv(global n ints) androw_inv(global m ints). On the denseleft_matmulkron path, the Jacobian is astacked_pdof p blocks, and every block's column space is the full variable space. That made the index metadata O(p·n_vars + p²·m0), compared with O(nnz) for the values. On CVXPY's OptimalAdvertising benchmark this put the engine's peak at about 6 GB, while the actual nonzeros need about 5 MB.Change
The inverse arrays are now built lazily, under one rule:
col_inv/row_inv; they stayNULL.unchanged, and the fillsassertthe array is present.row_gather_pd_allocis the one exception. Its fill copies rows throughbound_iwork, so at alloc it only needs a membership test, which it answers with a binary search (sorted_pos) overrow_perminstead of building a global-size table.Ensure calls were added in
BA_pd_csc_alloc,BTA_pd_csc_alloc,BTA_csc_pd_alloc, the coalesce output blocks,BTA_pd_spd_alloc, the kron scratch / BT, and old-code.An earlier revision added a second "compact" PD mode instead. That needed NULL fallbacks in about six places, and only the kron-csc outputs were made compact. The single lazy rule replaces it with a smaller diff (src+include: +88/−23) and also covers the kron-pd/spd outputs and the ATA raw blocks.
Results
g_peak_bytes)get_problem_data(median of 3)sum(exp(A @ X))Hessian init, 96×96Tests
test_peak_memory_kron_jacobian, which asserts < 400 B/var. It is guarded bySP_TRACK_MEMORY.test_permuted_dense_lazy_invandtest_permuted_dense_times_csc_lazy_inv.-DSP_TRACK_MEMORY=ON, and under ASan+UBSan.transpose(A @ X)problem-Jacobian test covers therow_gatherpath: it segfaults ifrow_gatherreadsrow_invdirectly.🤖 Generated with Claude Code