start looking at memory

Author: dance858

include/utils/matrix.h

@@ -106,6 +106,17 @@ typedef void (*matrix_broadcast_fill_values_fn)(matrix *A, broadcast_type type,
 typedef matrix *(*matrix_diag_vec_alloc_fn)(matrix *A);
 typedef void (*matrix_diag_vec_fill_values_fn)(matrix *A, matrix *out);
 
+/* sum_all_rows: C = sum over all rows of A. Output shape is (1, A->n).
+   Implementations choose the output's matrix type to preserve structure
+   (sparse → sparse, permuted_dense → permuted_dense, stacked_pd →
+   permuted_dense over the union of block col_perms). idx_map (sized at
+   A->nnz, allocated by caller) is filled such that idx_map[k] is the
+   position in C->x where A's k-th cell (in A->base.x ordering)
+   accumulates — eval consumers run `accumulator(A->x, A->nnz, idx_map,
+   C->x)` unchanged. No paired _fill_values is needed: sum's existing
+   eval_jacobian does the value fill polymorphically. */
+typedef matrix *(*matrix_sum_all_rows_alloc_fn)(matrix *A, int *idx_map);
+
 typedef void (*matrix_free_fn)(matrix *self);
 
 struct matrix
@@ -141,6 +152,7 @@ struct matrix
     matrix_broadcast_fill_values_fn broadcast_fill_values;
     matrix_diag_vec_alloc_fn diag_vec_alloc;
     matrix_diag_vec_fill_values_fn diag_vec_fill_values;
+    matrix_sum_all_rows_alloc_fn sum_all_rows_alloc;
 
     /* Lifecycle */
     matrix_free_fn free_fn;

include/utils/permuted_dense.h

@@ -81,6 +81,15 @@ 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);
 
+/* Like new_permuted_dense, but the X buffer is borrowed from the caller
+   (owns_X = false). The caller is responsible for keeping borrowed_X alive
+   for at least the lifetime of the returned PD and for freeing it
+   separately. row_perm / col_perm / row_inv / col_inv are still owned by
+   the PD (copied / computed internally). Used by stacked_pd output blocks
+   that share one shared values buffer. */
+matrix *new_permuted_dense_view(int m, int n, int m0, int n0, const int *row_perm,
+                                const int *col_perm, double *borrowed_X);
+
 /* 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/stacked_pd.h

@@ -22,6 +22,35 @@
 #include "matrix.h"
 #include "permuted_dense.h"
 
+/* Shape metadata for one would-be source block of a coalesce. Used by
+   spd_blockwise_alloc_coalesce to drive symbolic coalesce without
+   allocating per-block Gram PDs, and by spd_blockwise_fill_coalesce to
+   re-aim a shared scratch PD at each input block's metadata.
+   - (m0, n0): block shape.
+   - row_perm, col_perm: aliased pointers into operand arrays; not owned.
+   - scratch_*_needed: sizes the matching fill op will read from its
+     destination PD's kernel_dwork / kernel_iwork. Coalesce ignores these;
+     spd_blockwise uses them to size the shared scratch PD once at the
+     max-across-blocks. */
+typedef struct
+{
+    int m0;
+    int n0;
+    /* Both arrays are SP_MALLOC'd by the shape callback and owned by the
+       per_block_shapes array (freed by stacked_pd_free). */
+    int *row_perm;
+    int *col_perm;
+    /* Sizes the matching fill op will read from its destination PD's
+       kernel_dwork / kernel_iwork. */
+    size_t scratch_dwork_needed;
+    size_t scratch_iwork_needed;
+    /* Optional metadata to seed scratch->kernel_iwork[0..init_len) with
+       before invoking the per-block fill op. Used by BTA_pd_spd's fill
+       which reads s_max / max_n0_A from iwork[0..1]. */
+    int iwork_init[2];
+    int iwork_init_len;
+} spd_per_block_shape;
+
 /* stacked_pd represents a matrix that is the (vertical) union of 'n_blocks'
    permuted_dense blocks. Two different blocks have disjoint row permutations,
    but the column permutations may overlap across blocks.
@@ -56,6 +85,44 @@ typedef struct stacked_pd
        the kron-expanded transpose-of-A buffer across fills. NULL when
        not used. */
     permuted_dense *kernel_pd_scratch;
+
+    /* Streaming-fill state for stacked_pds produced by
+       spd_blockwise_alloc_coalesce. NULL when not used. All buffers are
+       sized at alloc time so the matching fill function does zero
+       SP_MALLOC.
+       - scratch_X: per-iteration buffer holding one source block's Gram.
+         Sized at max over input blocks of (m0_partial * n0_partial).
+       - scratch_dwork / scratch_iwork: workspaces the per-block fill op
+         reads from its (scratch) destination PD's kernel_dwork /
+         kernel_iwork. Sized at max-across-blocks.
+       - per_block_shapes: shape (and scratch sizes) for each input block's
+         partial. Indexed [0, n_input_blocks). The scratch PD's metadata
+         (m0, n0, row_perm, col_perm) is re-aimed from this array each
+         iteration.
+       - n_input_blocks: number of input blocks (length of per_block_shapes
+         and of src_to_outs_p[]-1). Distinct from n_blocks (which counts
+         OUTPUT blocks after coalesce).
+       - src_to_outs_p, src_to_outs_data: CSR-style inverse of
+         src_block_idx_p/_data. src_to_outs_data[src_to_outs_p[k] ..
+         src_to_outs_p[k+1]) lists the output block indices that block k
+         contributes to. Drives the per-input-block scatter loop in fill. */
+    double *scratch_X;
+    size_t scratch_X_capacity;
+    double *scratch_dwork;
+    size_t scratch_dwork_capacity;
+    int *scratch_iwork;
+    size_t scratch_iwork_capacity;
+    /* Re-initialized per iteration from per_block_shapes[k].row_perm /
+       col_perm. Sized at (m, n) — the global dimensions. Some fill ops
+       (e.g. BTA_pd_spd_fill_values at stacked_pd_linalg.c:477) read
+       C->col_inv to look up scatter positions; the scratch must have
+       it populated. */
+    int *scratch_row_inv;
+    int *scratch_col_inv;
+    spd_per_block_shape *per_block_shapes;
+    int n_input_blocks;
+    int *src_to_outs_p;
+    int *src_to_outs_data;
 } stacked_pd;
 
 /* Constructor for stacked_pd. Takes ownership of every block in 'blocks'. The
@@ -75,6 +142,26 @@ matrix *new_stacked_pd_unchecked(int m, int n, int n_blocks, permuted_dense **bl
                                  const int *src_block_idx_p,
                                  const int *src_block_idx);
 
+/* Like new_stacked_pd_unchecked, but the blocks' X pointers already point
+   into the caller-supplied `shared_x` buffer at the correct offsets
+   (typically constructed via new_permuted_dense_view). No absorb / memcpy
+   is performed. The stacked_pd takes ownership of `shared_x` and frees it
+   on destruction; each block must have owns_X == false. */
+matrix *new_stacked_pd_borrowed_x(int m, int n, int n_blocks,
+                                  permuted_dense **blocks,
+                                  const int *src_block_idx_p,
+                                  const int *src_block_idx, double *shared_x);
+
+/* Build a stacked_pd from precomputed per-block shapes. Sums total_nnz,
+   SP_MALLOCs one shared X buffer, constructs n_blocks view PDs pointing
+   into that buffer at sequential offsets, and wraps via
+   new_stacked_pd_borrowed_x. No absorb / double-count. Identity
+   src_block_idx (one source per output block). Reads only the shape's
+   m0 / n0 / row_perm / col_perm fields; per-block kernel_iwork /
+   kernel_dwork are the caller's responsibility. */
+matrix *new_stacked_pd_from_shapes_unchecked(int m, int n, int n_blocks,
+                                             const spd_per_block_shape *shapes);
+
 /* Filter-map over a stacked_pd's blocks. For each block of B, calls op(Bk,
    ctx). Drop blocks whose result has nnz == 0 and assembles the survivors into
    a new stacked_pd (dimensions Cm x Cn) with one source per output block. */
@@ -93,13 +180,31 @@ matrix *coalesce_spd_alloc(const stacked_pd *A);
    overlapping row permutations. */
 matrix *coalesce_spd_alloc_unchecked(const stacked_pd *A);
 
+/* Symbolic coalesce driven by a shapes array (no per-block X / PDs needed).
+   Returns an output stacked_pd whose blocks are views into one shared X
+   buffer that the output owns. The output's src_block_idx_p / src_block_idx
+   hold the forward map; callers that need the inverse read those fields
+   directly. Equivalent to coalesce_spd_alloc_unchecked but works from shape
+   metadata only — used by spd_blockwise_alloc_coalesce to skip the
+   per-block Gram materialization. */
+matrix *coalesce_spd_alloc_from_shapes_unchecked(const spd_per_block_shape *shapes,
+                                                 int n_input_blocks, int m, int n);
+
 /* Fill values of C = coalesce(A). */
 void coalesce_spd_fill_values(const stacked_pd *A, stacked_pd *C);
 
 /* Same scatter as coalesce_spd_fill_values but with += instead of =. Caller is
    responsible for zeroing C->base.x first. */
 void coalesce_spd_fill_values_accumulate(const stacked_pd *A, stacked_pd *C);
 
+/* Scatter the cells of one source PD into one output PD (using out_k's
+   row_inv / col_inv to map global row / col indices to local positions),
+   accumulating with += (caller must zero out_k->X before the first call
+   in a scatter cycle). Used by the per-input-block streaming fill in
+   spd_blockwise_fill_coalesce_accumulate. */
+void scatter_one_source_into_one_output_accumulate(const permuted_dense *src_k,
+                                                   permuted_dense *out_k);
+
 /* Re-index `idx_map` in place from CSR ordering to spd-native ordering
    (block-major). Useful for atoms that build idx_maps from a `to_csr`
    view at init time and want to read spd values directly at eval time,

include/utils/tracked_alloc.h

@@ -21,20 +21,61 @@
 #include <stddef.h>
 #include <stdlib.h>
 
-extern size_t g_allocated_bytes;
+/* Platform shim for "how many usable bytes are at this malloc'd pointer".
+   Apple's malloc_size, glibc's malloc_usable_size, MSVC's _msize. The
+   returned size is the *usable* size (may exceed the requested size by a
+   few bytes of allocator rounding), but alloc and free use the same
+   function so the running totals stay symmetric. */
+#if defined(__APPLE__)
+#include <malloc/malloc.h>
+#define TRACKED_BLOCK_SIZE(p) malloc_size(p)
+#elif defined(_WIN32) || defined(_WIN64)
+#include <malloc.h>
+#define TRACKED_BLOCK_SIZE(p) _msize(p)
+#else
+#include <malloc.h>
+#define TRACKED_BLOCK_SIZE(p) malloc_usable_size(p)
+#endif
+
+extern size_t g_allocated_bytes; /* current live bytes */
+extern size_t g_peak_bytes;      /* high-water mark since last reset */
 
 static inline void *SP_MALLOC(size_t size)
 {
     void *ptr = malloc(size);
-    if (ptr) g_allocated_bytes += size;
+    if (ptr)
+    {
+        g_allocated_bytes += TRACKED_BLOCK_SIZE(ptr);
+        if (g_allocated_bytes > g_peak_bytes) g_peak_bytes = g_allocated_bytes;
+    }
     return ptr;
 }
 
 static inline void *SP_CALLOC(size_t count, size_t size)
 {
     void *ptr = calloc(count, size);
-    if (ptr) g_allocated_bytes += count * size;
+    if (ptr)
+    {
+        g_allocated_bytes += TRACKED_BLOCK_SIZE(ptr);
+        if (g_allocated_bytes > g_peak_bytes) g_peak_bytes = g_allocated_bytes;
+    }
     return ptr;
 }
 
+static inline void SP_FREE(void *ptr)
+{
+    if (ptr)
+    {
+        g_allocated_bytes -= TRACKED_BLOCK_SIZE(ptr);
+        free(ptr);
+    }
+}
+
+/* Auto-route plain malloc/calloc/free in caller translation units through
+   the tracked wrappers. Defined AFTER the wrapper bodies so SP_MALLOC /
+   SP_CALLOC / SP_FREE themselves still see the real stdlib symbols. */
+#define malloc SP_MALLOC
+#define calloc SP_CALLOC
+#define free SP_FREE
+
 #endif /* TRACKED_ALLOC_H */

src/atoms/affine/sum.c

@@ -89,43 +89,57 @@ static void jacobian_init_impl(expr *node)
 
     /* initialize child's jacobian */
     jacobian_init(x);
-    CSR_matrix *Jx = x->jacobian->to_csr(x->jacobian);
 
-    /* we never have to store more than the child's nnz */
-    CSR_matrix *jac = new_CSR_matrix(node->size, node->n_vars, Jx->nnz);
-    node->work->iwork = SP_MALLOC(MAX(jac->n, Jx->nnz) * sizeof(int));
-    snode->idx_map = SP_MALLOC(Jx->nnz * sizeof(int));
-
-    /* the idx_map array maps each nonzero entry j in x->jacobian
-       to the corresponding index in the output row matrix C. Specifically, for
-       each nonzero entry j in A, idx_map[j] gives the position in C->x where
-       the value from x->jacobian->x[j] should be accumulated. */
+    /* idx_map maps each nonzero entry j of x->jacobian (in its native
+       base.x order) to the corresponding index in this node's jacobian
+       value array. eval_jacobian then runs
+           accumulator(x->jacobian->x, x->jacobian->nnz, idx_map, ...). */
+    snode->idx_map = SP_MALLOC(x->jacobian->nnz * sizeof(int));
 
     if (axis == -1)
     {
-        sum_all_rows_csr_alloc(Jx, jac, node->work->iwork, snode->idx_map);
-    }
-    else if (axis == 0)
-    {
-        sum_block_of_rows_csr_alloc(Jx, jac, x->d1, node->work->iwork,
-                                    snode->idx_map);
+        /* Polymorphic native-layout sum: walks x->jacobian without going
+           through CSR. The output's matrix type follows the input's
+           (sparse → sparse, PD → PD, stacked_pd → PD over col union).
+           idx_map comes back in x->jacobian->base.x ordering, so no CSR
+           re-permutation is needed. */
+        node->jacobian =
+            x->jacobian->sum_all_rows_alloc(x->jacobian, snode->idx_map);
     }
-    else if (axis == 1)
+    else
     {
-        sum_evenly_spaced_rows_csr_alloc(Jx, jac, node->size, node->work->iwork,
-                                         snode->idx_map);
-    }
+        /* axis=0 / axis=1: still go through CSR. The output has multiple
+           rows and the natural per-type output is less clean; the CSR
+           path is fine since these aren't hot in current benchmarks. */
+        CSR_matrix *Jx = x->jacobian->to_csr(x->jacobian);
+        size_t max_out_nnz =
+            MIN((size_t) Jx->nnz, (size_t) node->size * (size_t) node->n_vars);
+        CSR_matrix *jac =
+            new_CSR_matrix(node->size, node->n_vars, (int) max_out_nnz);
+        node->work->iwork = SP_MALLOC(MAX(jac->n, Jx->nnz) * sizeof(int));
+
+        if (axis == 0)
+        {
+            sum_block_of_rows_csr_alloc(Jx, jac, x->d1, node->work->iwork,
+                                        snode->idx_map);
+        }
+        else
+        {
+            sum_evenly_spaced_rows_csr_alloc(Jx, jac, node->size, node->work->iwork,
+                                             snode->idx_map);
+        }
 
-    /* For stacked_pd children, child->jacobian->base.x is block-major while
-       csr->x is row-major sorted. Re-index idx_map so it can be applied
-       directly to base.x in eval_jacobian. */
-    if (x->jacobian->is_stacked_pd)
-    {
-        compose_csr_idx_map_for_spd((const stacked_pd *) x->jacobian, Jx,
-                                    snode->idx_map);
-    }
+        /* For stacked_pd children, child->jacobian->base.x is block-major
+           while csr->x is row-major sorted. Re-index idx_map so it can be
+           applied directly to base.x in eval_jacobian. */
+        if (x->jacobian->is_stacked_pd)
+        {
+            compose_csr_idx_map_for_spd((const stacked_pd *) x->jacobian, Jx,
+                                        snode->idx_map);
+        }
 
-    node->jacobian = new_sparse_matrix(jac);
+        node->jacobian = new_sparse_matrix(jac);
+    }
 }
 
 static void eval_jacobian(expr *node)

src/problem.c

@@ -32,7 +32,11 @@ static void problem_lagrange_hess_fill_sparsity(problem *prob, int *iwork);
 problem *new_problem(expr *objective, expr **constraints, int n_constraints,
                      bool verbose)
 {
-    g_allocated_bytes = 0;
+    /* Baseline the peak watermark to the current live bytes — peak then
+       captures any growth during this problem's lifetime. Don't reset
+       g_allocated_bytes itself: allocations made before new_problem are
+       still alive, and their frees would subtract from this counter. */
+    g_peak_bytes = g_allocated_bytes;
     problem *prob = (problem *) SP_CALLOC(1, sizeof(problem));
     if (!prob) return NULL;
 
@@ -321,7 +325,7 @@ static inline void print_end_message(const Diff_engine_stats *stats)
     printf("  Lagrange Hessian (nnz):                 %d\n", stats->nnz_hessian);
     char mem_buf[64];
     format_memory(stats->memory_bytes, mem_buf, sizeof(mem_buf));
-    printf("  Allocated memory:                       %s\n", mem_buf);
+    printf("  Peak memory:                            %s\n", mem_buf);
 
     printf("\nTiming (seconds):\n");
     printf("  Derivative structure (sparsity):     %8.3f\n",
@@ -349,7 +353,7 @@ void free_problem(problem *prob)
 
     if (prob->verbose)
     {
-        prob->stats.memory_bytes = g_allocated_bytes;
+        prob->stats.memory_bytes = g_peak_bytes;
         print_end_message(&prob->stats);
     }