From 960e02be2adcdec59a796b1b5362a8636ceb3f82 Mon Sep 17 00:00:00 2001 From: Transurgeon Date: Sun, 2 Aug 2026 17:38:41 -0400 Subject: [PATCH] Compact permuted_dense (no inverse-permutation arrays) for kron-path 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 --- include/utils/permuted_dense.h | 24 ++++++- include/utils/utils.h | 10 +++ src/utils/permuted_dense.c | 114 ++++++++++++++++++++++++------ src/utils/permuted_dense_linalg.c | 69 ++++++++++++++++-- src/utils/stacked_pd_coalesce.c | 10 ++- src/utils/stacked_pd_linalg.c | 16 ++++- src/utils/utils.c | 35 +++++++++ tests/all_tests.c | 4 ++ tests/problem/test_peak_memory.h | 66 +++++++++++++++++ tests/utils/test_permuted_dense.h | 98 +++++++++++++++++++++++++ 10 files changed, 413 insertions(+), 33 deletions(-) create mode 100644 tests/problem/test_peak_memory.h diff --git a/include/utils/permuted_dense.h b/include/utils/permuted_dense.h index 47e4ab8d..1b1ffa5e 100644 --- a/include/utils/permuted_dense.h +++ b/include/utils/permuted_dense.h @@ -38,8 +38,14 @@ typedef struct permuted_dense int n0; int *row_perm; int *col_perm; - int *col_inv; /* col_inv[col_perm[jj]] = jj, otherwise -1 */ - int *row_inv; /* row_inv[row_perm[ii]] = ii, otherwise -1 */ + + /* Dense inverse permutations, sized by the GLOBAL dims (n and m): + col_inv[col_perm[jj]] = jj and row_inv[row_perm[ii]] = ii, -1 + elsewhere. PDs built by new_permuted_dense_compact leave both NULL; + consumers either fall back to sorted scans of the perms or call + permuted_dense_ensure_col_inv / _row_inv to materialize on demand. */ + int *col_inv; + int *row_inv; /* Row-major block of size m0 x n0. Owned by this PD when owns_X == true, and otherwise X is a view into a buffer (eg., stacked_pd's shared values @@ -81,6 +87,20 @@ matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm, col_perm = [0..n-1], dense block fills the full (m, n) shape. */ matrix *new_permuted_dense_full(int m, int n, const double *data); +/* Compact constructor: like new_permuted_dense, but leaves col_inv and + row_inv NULL instead of allocating the two global-dimension inverse + arrays — storage stays proportional to the block, not the global shape. + Used where those arrays dominate memory (the kron-path blocks, whose + global dims are the full variable space); consumers materialize them on + demand via the ensure helpers below, or scan the sorted perms instead. */ +matrix *new_permuted_dense_compact(int m, int n, int m0, int n0, const int *row_perm, + const int *col_perm, const double *X_data); + +/* Materialize A->col_inv / A->row_inv if NULL. Same const-cast on-demand + slot convention as permuted_dense_ensure_kernel_dwork. */ +void permuted_dense_ensure_col_inv(const permuted_dense *A); +void permuted_dense_ensure_row_inv(const permuted_dense *A); + /* Ensure A->kernel_dwork is sized at least 'size' doubles. Grows in place; contents are NOT preserved. */ void permuted_dense_ensure_kernel_dwork(const permuted_dense *A, size_t size); diff --git a/include/utils/utils.h b/include/utils/utils.h index c2a4b972..0588963f 100644 --- a/include/utils/utils.h +++ b/include/utils/utils.h @@ -60,6 +60,16 @@ int sorted_intersect_indices(const int *a, int a_len, const int *b, int b_len, void sorted_union_int_arrays(const int *const *arrs, const int *lens, int n_arrs, iVec *out); +/* Position of value g in the sorted, strictly-increasing array 'perm' of + length n0 (binary search), or -1 if absent. The scan replacement for a + dense inverse-permutation lookup inv[g]. */ +int sorted_pos(const int *perm, int n0, int g); + +/* Return true if any of the 'len' integers in 'idxs' (arbitrary order) is a + member of the sorted, strictly-increasing array 'perm' of length n0. The + scan replacement for idxs_hits_set when no dense inverse array exists. */ +bool sorted_hits(const int *idxs, int len, const int *perm, int n0); + /* in-place cumulative sum */ void cumsum(int *p, int n); diff --git a/src/utils/permuted_dense.c b/src/utils/permuted_dense.c index ed7dca50..e83b3710 100644 --- a/src/utils/permuted_dense.c +++ b/src/utils/permuted_dense.c @@ -113,16 +113,44 @@ matrix *index_pd_alloc(const permuted_dense *A, const int *indices, int n_idxs) /* Scan indices: which output positions i hit a row in A->row_perm? */ int *new_row_perm = (int *) sp_malloc(n_idxs * sizeof(int)); int new_m0 = 0; + if (A->row_inv != NULL) + { + for (int i = 0; i < n_idxs; i++) + { + if (A->row_inv[indices[i]] >= 0) + { + new_row_perm[new_m0++] = i; + } + } + + matrix *out = new_permuted_dense(n_idxs, A->base.n, new_m0, A->n0, + new_row_perm, A->col_perm, NULL); + sp_free(new_row_perm); + return out; + } + + /* Compact source (row_inv == NULL): membership via binary search in the sorted + row_perm. Record each hit's source dense row so index_pd_fill_values needs no + inverse array either — indices are owned by the atom and fixed + between alloc and fill, so the mapping stays valid. */ + int *src_rows = (int *) sp_malloc(n_idxs * sizeof(int)); for (int i = 0; i < n_idxs; i++) { - if (A->row_inv[indices[i]] >= 0) + int old_ii = sorted_pos(A->row_perm, A->m0, indices[i]); + if (old_ii >= 0) { + src_rows[new_m0] = old_ii; new_row_perm[new_m0++] = i; } } - matrix *out = new_permuted_dense(n_idxs, A->base.n, new_m0, A->n0, new_row_perm, - A->col_perm, NULL); + matrix *out = new_permuted_dense_compact(n_idxs, A->base.n, new_m0, A->n0, + new_row_perm, A->col_perm, NULL); + permuted_dense *C = (permuted_dense *) out; + C->kernel_iwork_size = (size_t) new_m0; + C->kernel_iwork = (int *) sp_malloc(new_m0 * sizeof(int)); + memcpy(C->kernel_iwork, src_rows, new_m0 * sizeof(int)); + sp_free(src_rows); sp_free(new_row_perm); return out; } @@ -132,11 +160,22 @@ void index_pd_fill_values(const permuted_dense *A, const int *indices, int n_idx { (void) n_idxs; int n0 = A->n0; + if (A->row_inv != NULL) + { + for (int k = 0; k < C->m0; k++) + { + int i = C->row_perm[k]; + int old_ii = A->row_inv[indices[i]]; + memcpy(C->X + k * n0, A->X + old_ii * n0, n0 * sizeof(double)); + } + return; + } + + /* Compact source (row_inv == NULL): C->kernel_iwork[k] holds the source dense + row for output dense row k, precomputed by index_pd_alloc. */ for (int k = 0; k < C->m0; k++) { - int i = C->row_perm[k]; - int old_ii = A->row_inv[indices[i]]; - memcpy(C->X + k * n0, A->X + old_ii * n0, n0 * sizeof(double)); + memcpy(C->X + k * n0, A->X + C->kernel_iwork[k] * n0, n0 * sizeof(double)); } } @@ -566,8 +605,9 @@ static void wire_vtable(permuted_dense *pd) pd->base.refresh_csc_values = permuted_dense_refresh_csc_values; } -matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm, - const int *col_perm, const double *X_data) +static matrix *new_permuted_dense_impl(int m, int n, int m0, int n0, + const int *row_perm, const int *col_perm, + const double *X_data, bool with_inv) { /* Validate sorted invariants. */ for (int ii = 1; ii < m0; ii++) @@ -602,8 +642,6 @@ matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm, pd->X = (double *) sp_malloc(sz * sizeof(double)); pd->base.x = pd->X; pd->owns_X = true; - pd->col_inv = (int *) sp_malloc(n * sizeof(int)); - pd->row_inv = (int *) sp_malloc(m * sizeof(int)); if (m0 > 0) { @@ -614,30 +652,64 @@ matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm, memcpy(pd->col_perm, col_perm, n0 * sizeof(int)); } + if (with_inv) + { + permuted_dense_ensure_col_inv(pd); + permuted_dense_ensure_row_inv(pd); + } + + if (X_data != NULL && sz > 0) + { + memcpy(pd->X, X_data, sz * sizeof(double)); + } + + return &pd->base; +} + +matrix *new_permuted_dense(int m, int n, int m0, int n0, const int *row_perm, + const int *col_perm, const double *X_data) +{ + return new_permuted_dense_impl(m, n, m0, n0, row_perm, col_perm, X_data, true); +} + +matrix *new_permuted_dense_compact(int m, int n, int m0, int n0, const int *row_perm, + const int *col_perm, const double *X_data) +{ + return new_permuted_dense_impl(m, n, m0, n0, row_perm, col_perm, X_data, false); +} + +void permuted_dense_ensure_col_inv(const permuted_dense *A) +{ + /* const-cast: lazily populating a cache slot, same convention as + permuted_dense_ensure_kernel_dwork. */ + permuted_dense *pd = (permuted_dense *) A; + if (pd->col_inv != NULL) return; + int n = pd->base.n; + pd->col_inv = (int *) sp_malloc(n * sizeof(int)); for (int j = 0; j < n; j++) { pd->col_inv[j] = -1; } - for (int jj = 0; jj < n0; jj++) + for (int jj = 0; jj < pd->n0; jj++) { - pd->col_inv[col_perm[jj]] = jj; + pd->col_inv[pd->col_perm[jj]] = jj; } +} +void permuted_dense_ensure_row_inv(const permuted_dense *A) +{ + permuted_dense *pd = (permuted_dense *) A; + if (pd->row_inv != NULL) return; + int m = pd->base.m; + pd->row_inv = (int *) sp_malloc(m * sizeof(int)); for (int i = 0; i < m; i++) { pd->row_inv[i] = -1; } - for (int ii = 0; ii < m0; ii++) + for (int ii = 0; ii < pd->m0; ii++) { - pd->row_inv[row_perm[ii]] = ii; + pd->row_inv[pd->row_perm[ii]] = ii; } - - if (X_data != NULL && sz > 0) - { - memcpy(pd->X, X_data, sz * sizeof(double)); - } - - return &pd->base; } matrix *new_permuted_dense_full(int m, int n, const double *data) diff --git a/src/utils/permuted_dense_linalg.c b/src/utils/permuted_dense_linalg.c index dc493e7b..277c156b 100644 --- a/src/utils/permuted_dense_linalg.c +++ b/src/utils/permuted_dense_linalg.c @@ -27,8 +27,23 @@ #include #include +/* Compactness propagates through structural copies: a compact PD (no + inverse arrays — the kron-path blocks) produces compact copies and + transposes, so the global-dimension inv arrays never materialize anywhere + along its transpose/reshape/copy_sparsity chain. Sources with prebuilt + inv arrays produce full copies. */ +static bool pd_is_compact(const permuted_dense *A) +{ + return A->col_inv == NULL || A->row_inv == NULL; +} + matrix *copy_sparsity_pd_alloc(const permuted_dense *A) { + if (pd_is_compact(A)) + { + return new_permuted_dense_compact(A->base.m, A->base.n, A->m0, A->n0, + A->row_perm, A->col_perm, NULL); + } return new_permuted_dense(A->base.m, A->base.n, A->m0, A->n0, A->row_perm, A->col_perm, NULL); } @@ -36,6 +51,11 @@ matrix *copy_sparsity_pd_alloc(const permuted_dense *A) matrix *transpose_pd_alloc(const permuted_dense *A) { /* swap (m, n), (m0, n0), and (row_perm, col_perm) */ + if (pd_is_compact(A)) + { + return new_permuted_dense_compact(A->base.n, A->base.m, A->n0, A->m0, + A->col_perm, A->row_perm, NULL); + } return new_permuted_dense(A->base.n, A->base.m, A->n0, A->m0, A->col_perm, A->row_perm, NULL); } @@ -265,14 +285,29 @@ matrix *BA_pd_csc_alloc(const permuted_dense *B, const CSC_matrix *A) { int start = A->p[j]; int len = A->p[j + 1] - start; - if (idxs_hits_set(A->i + start, len, B->col_inv)) + bool hit = B->col_inv != NULL + ? idxs_hits_set(A->i + start, len, B->col_inv) + : sorted_hits(A->i + start, len, B->col_perm, B->n0); + if (hit) { iVec_append(col_perm_C, j); } } - matrix *C = new_permuted_dense(B->base.m, A->n, B->m0, col_perm_C->len, - B->row_perm, col_perm_C->data, NULL); + /* The output is built compact: its inverse arrays would be sized by the + GLOBAL dims (m = B's global rows, n = A's full column count) and are + the dominant memory cost on the kron path, where this allocator runs + once per kron block. Every consumer of C either scans its sorted perms + or ensures the arrays on demand. */ + matrix *C = new_permuted_dense_compact(B->base.m, A->n, B->m0, col_perm_C->len, + B->row_perm, col_perm_C->data, NULL); + + /* BA_pd_csc_fill_values reads B->col_inv whenever C is non-empty (a + no-op when B already carries it, e.g. the kron scratch). */ + if (col_perm_C->len > 0) + { + permuted_dense_ensure_col_inv(B); + } iVec_free(col_perm_C); return C; } @@ -283,7 +318,9 @@ void BA_pd_csc_fill_values(const double *B, int n0_B, const int *inv, /* C[i, j] = bi^T @ ajj, where bi is the ith row of B_X (length n0_B, row stride n0_B) and ajj is the jjth column of A's sparse block (column jj = C->col_perm[j]). inv maps A's row indices to positions - in B_X (entries with inv[r] == -1 are skipped). */ + in B_X (entries with inv[r] == -1 are skipped). inv may be NULL only + for an empty C (the alloc twin ensures it whenever C is non-empty). */ + assert(inv != NULL || C->n0 == 0); /* row i of C */ for (int i = 0; i < C->m0; i++) @@ -399,7 +436,10 @@ matrix *BTA_pd_csc_alloc(const permuted_dense *B, const CSC_matrix *A) { int start = A->p[j]; int len = A->p[j + 1] - start; - if (idxs_hits_set(A->i + start, len, B->row_inv)) + bool hit = B->row_inv != NULL + ? idxs_hits_set(A->i + start, len, B->row_inv) + : sorted_hits(A->i + start, len, B->row_perm, B->m0); + if (hit) { iVec_append(col_active, j); } @@ -407,6 +447,13 @@ matrix *BTA_pd_csc_alloc(const permuted_dense *B, const CSC_matrix *A) matrix *C = new_permuted_dense(B->base.n, A->n, B->n0, col_active->len, B->col_perm, col_active->data, NULL); + + /* BTDA_pd_csc_fill_values reads B->row_inv whenever C is non-empty; a + disjoint (empty) product never materializes it. */ + if (col_active->len > 0) + { + permuted_dense_ensure_row_inv(B); + } iVec_free(col_active); /* Pre-size B's dwork for the BTDA fill (holds (diag(d) B)^T). */ @@ -453,7 +500,10 @@ matrix *BTA_csc_pd_alloc(const CSC_matrix *B, const permuted_dense *A) { int start = B->p[i]; int len = B->p[i + 1] - start; - if (idxs_hits_set(B->i + start, len, A->row_inv)) + bool hit = A->row_inv != NULL + ? idxs_hits_set(B->i + start, len, A->row_inv) + : sorted_hits(B->i + start, len, A->row_perm, A->m0); + if (hit) { iVec_append(row_active, i); } @@ -461,6 +511,13 @@ matrix *BTA_csc_pd_alloc(const CSC_matrix *B, const permuted_dense *A) matrix *C = new_permuted_dense(B->n, A->base.n, row_active->len, A->n0, row_active->data, A->col_perm, NULL); + + /* BTDA_csc_pd_fill_values reads A->row_inv whenever C is non-empty; a + disjoint (empty) product never materializes it. */ + if (row_active->len > 0) + { + permuted_dense_ensure_row_inv(A); + } iVec_free(row_active); /* Pre-size A's dwork for the BTDA fill (holds (diag(d_perm) X_A)^T). */ diff --git a/src/utils/stacked_pd_coalesce.c b/src/utils/stacked_pd_coalesce.c index 33a7edf3..4d69ebcb 100644 --- a/src/utils/stacked_pd_coalesce.c +++ b/src/utils/stacked_pd_coalesce.c @@ -342,7 +342,10 @@ static inline void coalesce_spd_scatter(const stacked_pd *src, stacked_pd *out, /* for each row in src_k */ for (int i = 0; i < src_k->m0; i++) { - int out_i = out_k->row_inv[src_k->row_perm[i]]; + int out_i = + out_k->row_inv != NULL + ? out_k->row_inv[src_k->row_perm[i]] + : sorted_pos(out_k->row_perm, out_k->m0, src_k->row_perm[i]); if (out_i < 0) { /* this row in src_k does not contribute to out_k */ @@ -355,7 +358,10 @@ static inline void coalesce_spd_scatter(const stacked_pd *src, stacked_pd *out, /* place each col in src_row at its correct position in out_row */ for (int j = 0; j < src_k->n0; j++) { - int out_j = out_k->col_inv[src_k->col_perm[j]]; + int out_j = out_k->col_inv != NULL + ? out_k->col_inv[src_k->col_perm[j]] + : sorted_pos(out_k->col_perm, out_k->n0, + src_k->col_perm[j]); assert(out_j >= 0); if (accumulate) { diff --git a/src/utils/stacked_pd_linalg.c b/src/utils/stacked_pd_linalg.c index 90a098f8..9e9c8c00 100644 --- a/src/utils/stacked_pd_linalg.c +++ b/src/utils/stacked_pd_linalg.c @@ -332,9 +332,21 @@ void BTA_pd_spd_fill_values(const permuted_dense *B, const stacked_pd *A, Bg, B->n0, Ag, Ak->n0, 0.0, Cg, Ak->n0); /* precompute scatter positions once per block */ - for (int j = 0; j < Ak->n0; j++) + if (C->col_inv != NULL) { - out_cols[j] = C->col_inv[Ak->col_perm[j]]; + for (int j = 0; j < Ak->n0; j++) + { + out_cols[j] = C->col_inv[Ak->col_perm[j]]; + } + } + else + { + /* compact C (col_inv == NULL): Ak's columns are a subset of C's, both + * sorted */ + for (int j = 0; j < Ak->n0; j++) + { + out_cols[j] = sorted_pos(C->col_perm, C->n0, Ak->col_perm[j]); + } } /* C += Cg (scatter + add into C) */ diff --git a/src/utils/utils.c b/src/utils/utils.c index 179f9435..c4e240dc 100644 --- a/src/utils/utils.c +++ b/src/utils/utils.c @@ -124,3 +124,38 @@ void cumsum(int *p, int n) p[i + 1] += p[i]; } } + +int sorted_pos(const int *perm, int n0, int g) +{ + int lo = 0; + int hi = n0 - 1; + while (lo <= hi) + { + int mid = lo + (hi - lo) / 2; + if (perm[mid] == g) + { + return mid; + } + if (perm[mid] < g) + { + lo = mid + 1; + } + else + { + hi = mid - 1; + } + } + return -1; +} + +bool sorted_hits(const int *idxs, int len, const int *perm, int n0) +{ + for (int ii = 0; ii < len; ii++) + { + if (sorted_pos(perm, n0, idxs[ii]) >= 0) + { + return true; + } + } + return false; +} diff --git a/tests/all_tests.c b/tests/all_tests.c index f5ddbed6..bd376cb1 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -61,6 +61,7 @@ #include "problem/test_param_broadcast.h" #include "problem/test_param_prob.h" #include "problem/test_param_source_refresh.h" +#include "problem/test_peak_memory.h" #include "problem/test_problem.h" #include "utils/test_COO_matrix.h" #include "utils/test_alloc_overflow.h" @@ -425,6 +426,8 @@ int main(void) mu_run_test(test_permuted_dense_times_csc_no_active, tests_run); mu_run_test(test_permuted_dense_to_csr_lazy, tests_run); mu_run_test(test_permuted_dense_col_inv, tests_run); + mu_run_test(test_permuted_dense_compact_inv, tests_run); + mu_run_test(test_permuted_dense_times_csc_compact_output, tests_run); mu_run_test(test_permuted_dense_index, tests_run); mu_run_test(test_permuted_dense_promote, tests_run); mu_run_test(test_permuted_dense_broadcast_scalar, tests_run); @@ -538,6 +541,7 @@ int main(void) mu_run_test(test_problem_constraint_forward, tests_run); mu_run_test(test_problem_hessian, tests_run); mu_run_test(test_problem_hessian_sum_exp_left_matmul_dense_transpose, tests_run); + mu_run_test(test_peak_memory_kron_jacobian, tests_run); printf("\n--- Parameter Tests ---\n"); mu_run_test(test_param_scalar_mult_problem, tests_run); diff --git a/tests/problem/test_peak_memory.h b/tests/problem/test_peak_memory.h new file mode 100644 index 00000000..1adefac5 --- /dev/null +++ b/tests/problem/test_peak_memory.h @@ -0,0 +1,66 @@ +#ifndef TEST_PEAK_MEMORY_H +#define TEST_PEAK_MEMORY_H + +#include + +#include "atoms/affine.h" +#include "expr.h" +#include "minunit.h" +#include "problem.h" +#include "utils/tracked_alloc.h" + +/* Peak-memory regression for the kron-path Jacobians (compact blocks, no + inverse-permutation arrays). + + Row-sum + col-sum constraints on a matrix variable X (a x b) both take the + dense left_matmul kron path: each Jacobian is a stacked_pd of p blocks + (p = b and p = a) whose global column space is all of n_vars = a * b. + Before the kron blocks were compact, every block eagerly carried col_inv + (n_vars ints) plus row_inv, so init peaked at O((a + b) * a * b) ints of + pure index metadata against O(a * b) true nnz: measured 1324 bytes/var + (21.7 MB) at a = b = 128. With compact blocks the same init peaks at 292 + bytes/var (4.8 MB), dominated by inherent O(n_vars) structure (the + children's identity Jacobians and their CSC caches). The 400 bytes/var + bound sits between the two with wide margins on both sides. */ +const char *test_peak_memory_kron_jacobian(void) +{ + const int a = 128; + const int b = 128; + const int n_vars = a * b; + + double *ones_a = (double *) malloc(a * sizeof(double)); + double *ones_b = (double *) malloc(b * sizeof(double)); + for (int i = 0; i < a; i++) ones_a[i] = 1.0; + for (int j = 0; j < b; j++) ones_b[j] = 1.0; + + expr *x_obj = new_variable(a, b, 0, n_vars); + expr *objective = new_sum(x_obj, -1); + + /* col sums: ones(1, a) @ X, shape (1, b) -> kron path with p = b blocks */ + expr *X1 = new_variable(a, b, 0, n_vars); + expr *colsum = new_left_matmul_dense(NULL, X1, 1, a, ones_a); + + /* row sums: ones(1, b) @ X^T, shape (1, a) -> kron path with p = a blocks */ + expr *X2 = new_variable(a, b, 0, n_vars); + expr *rowsum = new_left_matmul_dense(NULL, new_transpose(X2), 1, b, ones_b); + + expr *constraints[2] = {colsum, rowsum}; + problem *prob = new_problem(objective, constraints, 2, false); + mu_assert("new_problem failed", prob != NULL); + + /* new_problem rebaselined g_peak_bytes; measure the growth that + derivative initialization adds on top of the live bytes here. */ + size_t base = g_allocated_bytes; + problem_init_derivatives(prob); + size_t init_peak_growth = g_peak_bytes > base ? g_peak_bytes - base : 0; + + mu_assert("kron Jacobian init peak exceeds 400 bytes per variable", + init_peak_growth < (size_t) 400 * (size_t) n_vars); + + free_problem(prob); + free(ones_a); + free(ones_b); + return 0; +} + +#endif /* TEST_PEAK_MEMORY_H */ diff --git a/tests/utils/test_permuted_dense.h b/tests/utils/test_permuted_dense.h index 53a49bac..8d1ad40a 100644 --- a/tests/utils/test_permuted_dense.h +++ b/tests/utils/test_permuted_dense.h @@ -341,6 +341,104 @@ const char *test_permuted_dense_col_inv(void) return 0; } +/* Compact construction: inv arrays start NULL, the ensure helpers materialize + them with the same contents the full constructor would have produced, + and structural copies of a compact PD are compact themselves. */ +const char *test_permuted_dense_compact_inv(void) +{ + int row_perm[2] = {1, 4}; + int col_perm[2] = {0, 3}; + double X[4] = {1.0, 2.0, 3.0, 4.0}; + + matrix *M = new_permuted_dense_compact(5, 6, 2, 2, row_perm, col_perm, X); + permuted_dense *pd = (permuted_dense *) M; + + mu_assert("compact col_inv NULL", pd->col_inv == NULL); + mu_assert("compact row_inv NULL", pd->row_inv == NULL); + + /* Compactness propagates through copy_sparsity / transpose / index. */ + matrix *copy = copy_sparsity_pd_alloc(pd); + matrix *trans = transpose_pd_alloc(pd); + int indices[3] = {4, 0, 1}; + matrix *indexed = index_pd_alloc(pd, indices, 3); + mu_assert("copy is compact", ((permuted_dense *) copy)->col_inv == NULL); + mu_assert("transpose is compact", ((permuted_dense *) trans)->row_inv == NULL); + mu_assert("indexed is compact", ((permuted_dense *) indexed)->col_inv == NULL); + + /* index of a compact PD: rows {4, 0, 1} of M hit source rows {1, -, 0}. */ + permuted_dense *idx_pd = (permuted_dense *) indexed; + int idx_row_perm_expected[2] = {0, 2}; + mu_assert("indexed m0", idx_pd->m0 == 2); + mu_assert("indexed row_perm", + cmp_int_array(idx_pd->row_perm, idx_row_perm_expected, 2)); + index_pd_fill_values(pd, indices, 3, idx_pd); + double idx_X_expected[4] = {3.0, 4.0, 1.0, 2.0}; + mu_assert("indexed values", cmp_double_array(idx_pd->X, idx_X_expected, 4)); + + /* Ensure materializes the same arrays new_permuted_dense builds. */ + permuted_dense_ensure_col_inv(pd); + permuted_dense_ensure_row_inv(pd); + int col_inv_expected[6] = {0, -1, -1, 1, -1, -1}; + int row_inv_expected[5] = {-1, 0, -1, -1, 1}; + mu_assert("ensured col_inv", cmp_int_array(pd->col_inv, col_inv_expected, 6)); + mu_assert("ensured row_inv", cmp_int_array(pd->row_inv, row_inv_expected, 5)); + + /* Copies of a now-ensured PD carry prebuilt inv arrays again. */ + matrix *copy2 = copy_sparsity_pd_alloc(pd); + mu_assert("copy of ensured PD has col_inv", + ((permuted_dense *) copy2)->col_inv != NULL); + + free_matrix(copy2); + free_matrix(indexed); + free_matrix(trans); + free_matrix(copy); + free_matrix(M); + return 0; +} + +/* BA_pd_csc_alloc builds its output compact: same sparsity and values as + before, but no inverse arrays on the result. The left operand (with prebuilt inv + arrays) keeps working as the fill kernel's inv provider. */ +const char *test_permuted_dense_times_csc_compact_output(void) +{ + /* B: 3x4 PD, dense block rows {0, 2} x cols {1, 3}. */ + int row_perm[2] = {0, 2}; + int col_perm[2] = {1, 3}; + double BX[4] = {1.0, 2.0, 3.0, 4.0}; + matrix *B = new_permuted_dense(3, 4, 2, 2, row_perm, col_perm, BX); + permuted_dense *B_pd = (permuted_dense *) B; + + /* A: 4x3 CSC. Column 0 hits row 1 (in B's col_perm), column 1 hits + row 2 (not in col_perm), column 2 hits rows 1 and 3. */ + CSC_matrix *A = new_CSC_matrix(4, 3, 4); + int Ap[4] = {0, 1, 2, 4}; + int Ai[4] = {1, 2, 1, 3}; + double Ax[4] = {10.0, 5.0, 2.0, 3.0}; + memcpy(A->p, Ap, 4 * sizeof(int)); + memcpy(A->i, Ai, 4 * sizeof(int)); + memcpy(A->x, Ax, 4 * sizeof(double)); + + matrix *C = BA_pd_csc_alloc(B_pd, A); + permuted_dense *C_pd = (permuted_dense *) C; + + mu_assert("C col_inv NULL", C_pd->col_inv == NULL); + mu_assert("C row_inv NULL", C_pd->row_inv == NULL); + int C_col_perm_expected[2] = {0, 2}; + mu_assert("C n0", C_pd->n0 == 2); + mu_assert("C col_perm", cmp_int_array(C_pd->col_perm, C_col_perm_expected, 2)); + mu_assert("C row_perm", cmp_int_array(C_pd->row_perm, row_perm, 2)); + + BA_pd_csc_fill_values(B_pd->X, B_pd->n0, B_pd->col_inv, A, C_pd); + /* C[:, 0] = 10 * B[:, 1]; C[:, 1] = 2 * B[:, 1] + 3 * B[:, 3]. */ + double CX_expected[4] = {10.0, 8.0, 30.0, 18.0}; + mu_assert("C values", cmp_double_array(C_pd->X, CX_expected, 4)); + + free_matrix(C); + free_matrix(B); + free_CSC_matrix(A); + return 0; +} + /* PD index_alloc / index_fill_values: select rows from a PD; output must be another PD with row_perm equal to the output positions where indices[i] hit the source row_perm. */