Build permuted_dense inverse arrays lazily

.github/workflows/cmake.yml

@@ -31,6 +31,7 @@ jobs:
         -DCMAKE_BUILD_TYPE=${{ matrix.build_type }}
         -DCMAKE_VERBOSE_MAKEFILE=ON
         ${{ runner.os == 'Windows' && '-DCMAKE_TOOLCHAIN_FILE=C:/vcpkg/scripts/buildsystems/vcpkg.cmake' || '' }}
+        ${{ matrix.os == 'macos-latest' && matrix.build_type == 'Debug' && '-DSP_TRACK_MEMORY=ON' || '' }}
 
     - name: Build
       run: cmake --build build --config ${{ matrix.build_type }}

include/utils/permuted_dense.h

@@ -38,8 +38,16 @@ 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. NULL until built by permuted_dense_ensure_col_inv /
+       _row_inv: a kernel whose fill reads one ensures it on its operand in
+       the paired alloc. Built lazily because many PDs never need them and
+       their size is global, not per-block (e.g. the p blocks of a kron
+       Jacobian each have the full variable space as columns). */
+    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
@@ -92,6 +100,11 @@ 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);
 
+/* Build 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);

include/utils/utils.h

@@ -60,6 +60,10 @@ 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. */
+int sorted_pos(const int *perm, int n0, int g);
+
 /* in-place cumulative sum */
 void cumsum(int *p, int n);
 

src/old-code/old_permuted_dense.c

@@ -58,6 +58,7 @@ matrix *BTA_pd_csr_alloc(const permuted_dense *B, const CSR_matrix *A)
 
     matrix *C =
         new_permuted_dense(B->base.n, p, B->n0, s_A, B->col_perm, col_active, NULL);
+    permuted_dense_ensure_col_inv((permuted_dense *) C);
     sp_free(col_active);
     sp_free(seen);
 
@@ -89,7 +90,7 @@ void BTA_pd_csr_fill_values(const permuted_dense *B, const CSR_matrix *A_csr,
         return;
     }
 
-    /* Use C->col_inv (pre-built by new_permuted_dense) as col_inv_out and
+    /* Use C->col_inv (built at alloc) as col_inv_out and
        C->kernel_dwork as A_sub_dense; both are owned by C. dwork is sized at alloc
        time to cover m0 * s_A; only that prefix is touched. */
     double *A_sub_dense = C->kernel_dwork;
@@ -191,6 +192,7 @@ matrix *BTA_csr_pd_alloc(const CSR_matrix *B_csr, const permuted_dense *A)
 
     matrix *C =
         new_permuted_dense(q, A->base.n, r_B, A->n0, row_active, A->col_perm, NULL);
+    permuted_dense_ensure_row_inv((permuted_dense *) C);
     sp_free(row_active);
     sp_free(seen);
 
@@ -222,7 +224,7 @@ void BTA_csr_pd_fill_values(const CSR_matrix *B_csr, const permuted_dense *A,
         return;
     }
 
-    /* Use C->row_inv (pre-built by new_permuted_dense) as row_inv_out and
+    /* Use C->row_inv (built at alloc) as row_inv_out and
        C->kernel_dwork as B_sub_dense; both are owned by C. dwork is sized at alloc
        time to cover m0 * r_B; only that prefix is touched. */
     double *B_sub_dense = C->kernel_dwork;

src/utils/permuted_dense.c

@@ -114,13 +114,24 @@ matrix *row_gather_pd_alloc(const permuted_dense *A, const int *map, int m_out)
 {
     /* Output position i is dense iff map[i] hits a row in A->row_perm. The kept
        positions form C's row_perm (strictly increasing by construction, repeats
-       in map included); src[k] is the dense row of A that C's row k copies. */
+       in map included); src[k] is the dense row of A that C's row k copies.
+       The fill reads only src, so membership is a binary search in row_perm
+       rather than materializing A->row_inv. */
     int *new_row_perm = (int *) sp_malloc(m_out * sizeof(int));
     int *src = (int *) sp_malloc(m_out * sizeof(int));
     int new_m0 = 0;
+
+    /* map entries outside [row_lo, row_hi] miss without a search: the common
+       case when A is one block of a stacked_pd and map spans all blocks */
+    int row_lo = A->m0 > 0 ? A->row_perm[0] : 0;
+    int row_hi = A->m0 > 0 ? A->row_perm[A->m0 - 1] : -1;
     for (int i = 0; i < m_out; i++)
     {
-        int ii = A->row_inv[map[i]];
+        if (map[i] < row_lo || map[i] > row_hi)
+        {
+            continue;
+        }
+        int ii = sorted_pos(A->row_perm, A->m0, map[i]);
         if (ii >= 0)
         {
             new_row_perm[new_m0] = i;
@@ -390,8 +401,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)
     {
@@ -402,30 +411,42 @@ 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));
     }
 
-    for (int j = 0; j < n; j++)
+    if (X_data != NULL && sz > 0)
+    {
+        memcpy(pd->X, X_data, sz * sizeof(double));
+    }
+
+    return &pd->base;
+}
+
+void permuted_dense_ensure_col_inv(const permuted_dense *pd_const)
+{
+    permuted_dense *pd = (permuted_dense *) pd_const;
+    if (pd->col_inv != NULL) return;
+    pd->col_inv = (int *) sp_malloc(pd->base.n * sizeof(int));
+    for (int j = 0; j < pd->base.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;
     }
+}
 
-    for (int i = 0; i < m; i++)
+void permuted_dense_ensure_row_inv(const permuted_dense *pd_const)
+{
+    permuted_dense *pd = (permuted_dense *) pd_const;
+    if (pd->row_inv != NULL) return;
+    pd->row_inv = (int *) sp_malloc(pd->base.m * sizeof(int));
+    for (int i = 0; i < pd->base.m; i++)
     {
         pd->row_inv[i] = -1;
     }
-    for (int ii = 0; ii < m0; ii++)
-    {
-        pd->row_inv[row_perm[ii]] = ii;
-    }
-
-    if (X_data != NULL && sz > 0)
+    for (int ii = 0; ii < pd->m0; ii++)
     {
-        memcpy(pd->X, X_data, sz * sizeof(double));
+        pd->row_inv[pd->row_perm[ii]] = ii;
     }
-
-    return &pd->base;
 }
 
 matrix *new_permuted_dense_full(int m, int n, const double *data)

src/utils/permuted_dense_linalg.c

@@ -274,6 +274,7 @@ matrix *BA_pd_csc_alloc(const permuted_dense *B, const CSC_matrix *A)
     the columns of A. For each column of A, we check if it has any nonzeros in
     rows that are in B's col_perm. If yes, column j of C will have a nonzero
     block corresponding to the rows of B */
+    permuted_dense_ensure_col_inv(B);
     iVec *col_perm_C = iVec_new(10);
     for (int j = 0; j < A->n; j++)
     {
@@ -298,6 +299,7 @@ void BA_pd_csc_fill_values(const double *B, int n0_B, const int *inv,
        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). */
+    assert(inv != NULL);
 
     /* row i of C */
     for (int i = 0; i < C->m0; i++)
@@ -408,6 +410,7 @@ matrix *BTA_pd_csc_alloc(const permuted_dense *B, const CSC_matrix *A)
     through the columns of A. For each column of A, we check if it has any
     nonzeros in rows that are in B's row_perm. If yes, column j of C will
     have a nonzero block corresponding to the columns of B */
+    permuted_dense_ensure_row_inv(B);
     iVec *col_active = iVec_new(8);
     for (int j = 0; j < A->n; j++)
     {
@@ -476,7 +479,7 @@ matrix *BTA_csc_pd_alloc(const CSC_matrix *B, const permuted_dense *A)
        columns of B. For each column of B, we check if it has any nonzeros in rows
        that are in A->row_perm. If yes, column i of C will have a nonzero block
        corresponding to the columns of A */
-
+    permuted_dense_ensure_row_inv(A);
     iVec *row_active = iVec_new(10);
     for (int i = 0; i < B->n; i++)
     {
@@ -506,6 +509,7 @@ static void BTA_csc_denseT_fill_values(const CSC_matrix *B, const double *A_T,
                                        int m0_A, const int *inv, permuted_dense *C)
 {
     /* C[i_C, j_C] = dot(col C->row_perm[i_C] of B, row j_C of A_T). */
+    assert(inv != NULL);
     for (int i_C = 0; i_C < C->m0; i_C++)
     {
         int B_col = C->row_perm[i_C];

src/utils/stacked_pd_coalesce.c

@@ -234,6 +234,10 @@ static permuted_dense **build_output_blocks(const stacked_pd *A,
         const int *row_perm = int_csr_get(sig_to_rows, s, &m0);
         out_blocks[s] = (permuted_dense *) new_permuted_dense(
             m, n, m0, col_union->len, row_perm, col_union->data, NULL);
+
+        /* coalesce_spd_scatter reads both inverse arrays every fill. */
+        permuted_dense_ensure_row_inv(out_blocks[s]);
+        permuted_dense_ensure_col_inv(out_blocks[s]);
     }
 
     iVec_free(col_union);
@@ -331,6 +335,7 @@ static inline void coalesce_spd_scatter(const stacked_pd *src, stacked_pd *out,
     for (int k = 0; k < out->n_blocks; k++)
     {
         permuted_dense *out_k = out->blocks[k];
+        assert(out_k->row_inv != NULL && out_k->col_inv != NULL);
         int s_lo = out->src_block_idx_p[k];
         int s_hi = out->src_block_idx_p[k + 1];
 

src/utils/stacked_pd_kron_linalg.c

@@ -22,6 +22,7 @@
 #include "utils/stacked_pd.h"
 #include "utils/stacked_pd_linalg.h"
 #include "utils/tracked_alloc.h"
+#include <assert.h>
 #include <stdlib.h>
 
 // ====================================================================================
@@ -65,6 +66,7 @@ static permuted_dense *kron_scratch_init(const permuted_dense *A, int p)
     matrix *pd_m = new_permuted_dense(m * p, n * p, m, n, row_perm, col_perm, NULL);
     sp_free(row_perm);
     sp_free(col_perm);
+    permuted_dense_ensure_col_inv((permuted_dense *) pd_m);
 
     permuted_dense *pd = (permuted_dense *) pd_m;
     sp_free(pd->X);
@@ -89,16 +91,18 @@ static permuted_dense *kron_BT_alloc(const permuted_dense *A, int p)
     matrix *BT_m = new_permuted_dense(n * p, m * p, n, m, row_perm, col_perm, NULL);
     sp_free(row_perm);
     sp_free(col_perm);
+    permuted_dense_ensure_col_inv((permuted_dense *) BT_m);
     return (permuted_dense *) BT_m;
 }
 
 /* Mutate the state to represent kron block k. last_k is the block currently
    encoded (0 immediately after _init); the call clears last_k's col_inv
    entries and sets block-k's, plus refreshes row_perm and col_perm. O(m + n)
-   work. row_inv is left stale — no per-block kernel reads it. Works on both
+   work. row_inv is never built — no per-block kernel reads it. Works on both
    scratch (pd/csc) and BT (spd) under the "full A" contract. */
 static void kron_scratch_set_block(permuted_dense *state, int k, int last_k)
 {
+    assert(state->col_inv != NULL);
     int m = state->m0;
     int n = state->n0;
     for (int i = 0; i < m; i++)

src/utils/stacked_pd_linalg.c

@@ -278,6 +278,7 @@ matrix *BTA_pd_spd_alloc(const permuted_dense *B, const stacked_pd *A)
     // -------------------------------------------------------------------------------
     matrix *C = new_permuted_dense(B->base.n, A->base.n, B->n0, col_union->len,
                                    B->col_perm, col_union->data, NULL);
+    permuted_dense_ensure_col_inv((permuted_dense *) C);
 
     iVec_free(col_union);
     sp_free(col_perms);
@@ -327,6 +328,7 @@ static void BTA_pd_spd_core(const permuted_dense *B, const double *d,
     {
         return;
     }
+    assert(C->col_inv != NULL);
 
     /* reset values of C */
     memset(C->X, 0, (size_t) C->m0 * C->n0 * sizeof(double));

src/utils/utils.c

@@ -124,3 +124,26 @@ 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;
+}

tests/all_tests.c

@@ -62,6 +62,9 @@
 #include "problem/test_param_broadcast.h"
 #include "problem/test_param_prob.h"
 #include "problem/test_param_source_refresh.h"
+#ifdef SP_TRACK_MEMORY
+#include "problem/test_peak_memory.h"
+#endif
 #include "problem/test_problem.h"
 #include "utils/test_COO_matrix.h"
 #include "utils/test_alloc_overflow.h"
@@ -448,7 +451,7 @@ int main(void)
     mu_run_test(test_permuted_dense_times_csc, tests_run);
     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_lazy_inv, tests_run);
     mu_run_test(test_permuted_dense_row_gather, tests_run);
     mu_run_test(test_row_gather_sparse, tests_run);
     mu_run_test(test_row_gather_pd_vs_sparse_twin, tests_run);
@@ -464,6 +467,7 @@ int main(void)
 #ifdef SP_TRACK_MEMORY
     mu_run_test(test_row_reduce_spd_fill_no_transient_alloc, tests_run);
 #endif
+    mu_run_test(test_permuted_dense_times_csc_lazy_inv, tests_run);
     mu_run_test(test_permuted_dense_diag_vec, tests_run);
     mu_run_test(test_permuted_dense_BTA_matching_row_perm, tests_run);
     mu_run_test(test_permuted_dense_BTA_empty_overlap, tests_run);
@@ -593,6 +597,9 @@ 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);
+#ifdef SP_TRACK_MEMORY
+    mu_run_test(test_peak_memory_kron_jacobian, tests_run);
+#endif
 
     printf("\n--- Parameter Tests ---\n");
     mu_run_test(test_param_scalar_mult_problem, tests_run);