Pr3 peak live memory tracking

include/utils/tracked_alloc.h

@@ -21,20 +21,60 @@
 #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".
+    Used to track total live bytes */
+#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
 
-static inline void *SP_MALLOC(size_t size)
+extern size_t g_allocated_bytes; /* current live bytes */
+extern size_t g_peak_bytes;      /* high-water mark since last reset */
+
+/* All allocations in src/ must go through these wrappers (and pair sp_free
+   with sp_malloc / sp_calloc). Tests may use plain malloc/free — those
+   bytes are simply not tracked. */
+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)
+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);
+    }
+}
+
 #endif /* TRACKED_ALLOC_H */

src/atoms/affine/add.c

@@ -95,7 +95,7 @@ static bool is_affine(const expr *node)
 expr *new_add(expr *left, expr *right)
 {
     assert(left->d1 == right->d1 && left->d2 == right->d2);
-    expr *node = (expr *) SP_CALLOC(1, sizeof(expr));
+    expr *node = (expr *) sp_calloc(1, sizeof(expr));
     init_expr(node, left->d1, left->d2, left->n_vars, forward, jacobian_init_impl,
               eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     node->left = left;

src/atoms/affine/broadcast.c

@@ -96,7 +96,7 @@ static void wsum_hess_init_impl(expr *node)
     node->wsum_hess = x->wsum_hess->copy_sparsity(x->wsum_hess);
 
     /* allocate space for weight vector */
-    node->work->dwork = SP_MALLOC(node->size * sizeof(double));
+    node->work->dwork = sp_malloc(node->size * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -176,7 +176,7 @@ expr *new_broadcast(expr *child, int d1, int d2)
         exit(1);
     }
 
-    broadcast_expr *bcast = (broadcast_expr *) SP_CALLOC(1, sizeof(broadcast_expr));
+    broadcast_expr *bcast = (broadcast_expr *) sp_calloc(1, sizeof(broadcast_expr));
     expr *node = (expr *) bcast;
 
     // --------------------------------------------------------------------------

src/atoms/affine/convolve.c

@@ -82,7 +82,8 @@ static void jacobian_init_impl(expr *node)
     conv_matrix_fill_values(cnode->T, a);
 
     /* J = T @ J_child */
-    cnode->Jchild_CSC = csr_to_csc_alloc(child->jacobian->to_csr(child->jacobian), node->work->iwork);
+    cnode->Jchild_CSC = csr_to_csc_alloc(child->jacobian->to_csr(child->jacobian),
+                                         node->work->iwork);
     node->jacobian =
         new_sparse_matrix(csr_csc_matmul_alloc(cnode->T, cnode->Jchild_CSC));
 }
@@ -95,8 +96,8 @@ static void eval_jacobian(expr *node)
     child->eval_jacobian(child);
 
     /* J = T @ J_child */
-    csr_to_csc_fill_values(child->jacobian->to_csr(child->jacobian), cnode->Jchild_CSC,
-                           node->work->iwork);
+    csr_to_csc_fill_values(child->jacobian->to_csr(child->jacobian),
+                           cnode->Jchild_CSC, node->work->iwork);
     csr_csc_matmul_fill_values(cnode->T, cnode->Jchild_CSC,
                                node->jacobian->to_csr(node->jacobian));
 }
@@ -108,7 +109,7 @@ static void wsum_hess_init_impl(expr *node)
 
     wsum_hess_init(child);
     node->wsum_hess = child->wsum_hess->copy_sparsity(child->wsum_hess);
-    node->work->dwork = (double *) SP_MALLOC(cnode->n * sizeof(double));
+    node->work->dwork = (double *) sp_malloc(cnode->n * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -174,7 +175,7 @@ expr *new_convolve(expr *param_node, expr *child)
         exit(1);
     }
 
-    convolve_expr *cnode = (convolve_expr *) SP_CALLOC(1, sizeof(convolve_expr));
+    convolve_expr *cnode = (convolve_expr *) sp_calloc(1, sizeof(convolve_expr));
     expr *node = &cnode->base;
     init_expr(node, d1, d2, child->n_vars, forward, jacobian_init_impl,
               eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess,
@@ -188,7 +189,7 @@ expr *new_convolve(expr *param_node, expr *child)
     cnode->n = n;
 
     /* iwork is used for csr_to_csc of child's Jacobian (size n_vars). */
-    node->work->iwork = (int *) SP_MALLOC(node->n_vars * sizeof(int));
+    node->work->iwork = (int *) sp_malloc(node->n_vars * sizeof(int));
 
     /* Ensure first forward() pulls current param values through any
        broadcast/promote wrappers and reflects them in T (once T is built). */

src/atoms/affine/diag_vec.c

@@ -70,7 +70,7 @@ static void wsum_hess_init_impl(expr *node)
     wsum_hess_init(x);
 
     /* workspace for extracting diagonal weights */
-    node->work->dwork = (double *) SP_CALLOC(x->size, sizeof(double));
+    node->work->dwork = (double *) sp_calloc(x->size, sizeof(double));
 
     /* Copy child's Hessian structure (diag_vec is linear, so its own Hessian is
      * zero) */
@@ -106,7 +106,7 @@ expr *new_diag_vec(expr *child)
 
     /* n is the number of elements (works for both row and column vectors) */
     int n = child->size;
-    expr *node = (expr *) SP_CALLOC(1, sizeof(expr));
+    expr *node = (expr *) sp_calloc(1, sizeof(expr));
     init_expr(node, n, n, child->n_vars, forward, jacobian_init_impl, eval_jacobian,
               is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     node->left = child;

src/atoms/affine/hstack.c

@@ -187,14 +187,14 @@ expr *new_hstack(expr **args, int n_args, int n_vars)
     }
 
     /* Allocate the type-specific struct */
-    hstack_expr *hnode = (hstack_expr *) SP_CALLOC(1, sizeof(hstack_expr));
+    hstack_expr *hnode = (hstack_expr *) sp_calloc(1, sizeof(hstack_expr));
     expr *node = &hnode->base;
     init_expr(node, args[0]->d1, d2, n_vars, forward, jacobian_init_impl,
               eval_jacobian, is_affine, wsum_hess_init_impl, wsum_hess_eval,
               free_type_data);
 
     /* Set type-specific fields (deep copy args array) */
-    hnode->args = (expr **) SP_CALLOC(n_args, sizeof(expr *));
+    hnode->args = (expr **) sp_calloc(n_args, sizeof(expr *));
     hnode->n_args = n_args;
     for (int i = 0; i < n_args; i++)
     {

src/atoms/affine/index.c

@@ -29,7 +29,7 @@
  * Returns true if duplicates exist, false otherwise. */
 static bool check_for_duplicates(const int *indices, int n_idxs, int max_idx)
 {
-    bool *seen = (bool *) SP_CALLOC(max_idx, sizeof(bool));
+    bool *seen = (bool *) sp_calloc(max_idx, sizeof(bool));
     bool has_dup = false;
     for (int i = 0; i < n_idxs && !has_dup; i++)
     {
@@ -89,7 +89,7 @@ static void wsum_hess_init_impl(expr *node)
     wsum_hess_init(x);
 
     /* for setting weight vector to evaluate hessian of child */
-    node->work->dwork = (double *) SP_CALLOC(x->size, sizeof(double));
+    node->work->dwork = (double *) sp_calloc(x->size, sizeof(double));
 
     /* in the implementation of eval_wsum_hess we evaluate the
        child's hessian with a weight vector that has w[i] = 0
@@ -148,7 +148,7 @@ expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs)
 {
     assert(d1 * d2 == n_idxs);
     /* allocate type-specific struct */
-    index_expr *idx = (index_expr *) SP_CALLOC(1, sizeof(index_expr));
+    index_expr *idx = (index_expr *) sp_calloc(1, sizeof(index_expr));
     expr *node = &idx->base;
 
     init_expr(node, d1, d2, child->n_vars, forward, jacobian_init_impl,
@@ -159,7 +159,7 @@ expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs)
     expr_retain(child);
 
     /* copy indices */
-    idx->indices = (int *) SP_MALLOC(n_idxs * sizeof(int));
+    idx->indices = (int *) sp_malloc(n_idxs * sizeof(int));
     memcpy(idx->indices, indices, n_idxs * sizeof(int));
     idx->n_idxs = n_idxs;
 

src/atoms/affine/left_matmul.c

@@ -201,7 +201,7 @@ static void wsum_hess_init_impl(expr *node)
     /* work for computing A^T w*/
     int n_blocks = ((left_matmul_expr *) node)->n_blocks;
     int dim = ((left_matmul_expr *) node)->AT->m * n_blocks;
-    node->work->dwork = (double *) SP_MALLOC(dim * sizeof(double));
+    node->work->dwork = (double *) sp_malloc(dim * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -263,7 +263,7 @@ expr *new_left_matmul(expr *param_node, expr *u, const CSR_matrix *A)
 
     /* Allocate the type-specific struct */
     left_matmul_expr *lnode =
-        (left_matmul_expr *) SP_CALLOC(1, sizeof(left_matmul_expr));
+        (left_matmul_expr *) sp_calloc(1, sizeof(left_matmul_expr));
     expr *node = &lnode->base;
     /* Sparse A — always the general CSC-mirror path. */
     init_expr(node, d1, d2, u->n_vars, forward, jacobian_init_sparse,
@@ -276,8 +276,8 @@ expr *new_left_matmul(expr *param_node, expr *u, const CSR_matrix *A)
        (requiring size node->n_vars).
        csc_to_csr_work is used for converting J_CSC to CSR_matrix (requiring
        node->size) */
-    node->work->iwork = (int *) SP_MALLOC(node->n_vars * sizeof(int));
-    lnode->csc_to_csr_work = (int *) SP_MALLOC(node->size * sizeof(int));
+    node->work->iwork = (int *) sp_malloc(node->n_vars * sizeof(int));
+    lnode->csc_to_csr_work = (int *) sp_malloc(node->size * sizeof(int));
     lnode->n_blocks = n_blocks;
 
     /* store A and AT. new_sparse_matrix takes ownership, so clone first. */
@@ -304,7 +304,7 @@ expr *new_left_matmul_dense(expr *param_node, expr *u, int m, int n,
     compute_dims(u, m, n, &d1, &d2, &n_blocks);
 
     left_matmul_expr *lnode =
-        (left_matmul_expr *) SP_CALLOC(1, sizeof(left_matmul_expr));
+        (left_matmul_expr *) sp_calloc(1, sizeof(left_matmul_expr));
     expr *node = &lnode->base;
     init_expr(node, d1, d2, u->n_vars, forward, jacobian_init_dense,
               eval_jacobian_dense, is_affine, wsum_hess_init_impl, eval_wsum_hess,

src/atoms/affine/neg.c

@@ -85,7 +85,7 @@ static bool is_affine(const expr *node)
 
 expr *new_neg(expr *child)
 {
-    expr *node = (expr *) SP_CALLOC(1, sizeof(expr));
+    expr *node = (expr *) sp_calloc(1, sizeof(expr));
     init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init_impl,
               eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     node->left = child;

src/atoms/affine/parameter.c

@@ -61,7 +61,7 @@ static bool is_affine(const expr *node)
 
 expr *new_parameter(int d1, int d2, int param_id, int n_vars, const double *values)
 {
-    parameter_expr *pnode = (parameter_expr *) SP_CALLOC(1, sizeof(parameter_expr));
+    parameter_expr *pnode = (parameter_expr *) sp_calloc(1, sizeof(parameter_expr));
     expr *node = &pnode->base;
     init_expr(node, d1, d2, n_vars, forward, jacobian_init_impl, eval_jacobian,
               is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);

src/atoms/affine/promote.c

@@ -84,7 +84,7 @@ static bool is_affine(const expr *node)
 expr *new_promote(expr *child, int d1, int d2)
 {
     assert(child->size == 1);
-    expr *node = (expr *) SP_CALLOC(1, sizeof(expr));
+    expr *node = (expr *) sp_calloc(1, sizeof(expr));
     init_expr(node, d1, d2, child->n_vars, forward, jacobian_init_impl,
               eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     node->left = child;

src/atoms/affine/reshape.c

@@ -69,7 +69,7 @@ static bool is_affine(const expr *node)
 expr *new_reshape(expr *child, int d1, int d2)
 {
     assert(d1 * d2 == child->size);
-    expr *node = (expr *) SP_CALLOC(1, sizeof(expr));
+    expr *node = (expr *) sp_calloc(1, sizeof(expr));
     init_expr(node, d1, d2, child->n_vars, forward, jacobian_init_impl,
               eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     node->left = child;

src/atoms/affine/right_matmul.c

@@ -58,7 +58,7 @@ expr *new_right_matmul(expr *param_node, expr *u, const CSR_matrix *A)
 {
     /* We can express right matmul using left matmul and transpose:
        u @ A = (A^T @ u^T)^T. */
-    int *work_transpose = (int *) SP_MALLOC(A->n * sizeof(int));
+    int *work_transpose = (int *) sp_malloc(A->n * sizeof(int));
     CSR_matrix *AT = transpose(A, work_transpose);
 
     expr *u_transpose = new_transpose(u);
@@ -104,7 +104,7 @@ expr *new_right_matmul_dense(expr *param_node, expr *u, int m, int n,
     }
     else
     {
-        double *AT = (double *) SP_MALLOC(n * m * sizeof(double));
+        double *AT = (double *) sp_malloc(n * m * sizeof(double));
         A_transpose(AT, data, m, n);
         left_matmul_node = new_left_matmul_dense(NULL, u_transpose, n, m, AT);
         free(AT);

src/atoms/affine/scalar_mult.c

@@ -118,7 +118,7 @@ static void free_type_data(expr *node)
 expr *new_scalar_mult(expr *param_node, expr *child)
 {
     scalar_mult_expr *mult_node =
-        (scalar_mult_expr *) SP_CALLOC(1, sizeof(scalar_mult_expr));
+        (scalar_mult_expr *) sp_calloc(1, sizeof(scalar_mult_expr));
     expr *node = &mult_node->base;
 
     init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init_impl,

src/atoms/affine/sum.c

@@ -93,8 +93,8 @@ static void jacobian_init_impl(expr *node)
 
     /* 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));
+    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
@@ -150,7 +150,7 @@ static void wsum_hess_init_impl(expr *node)
 
     /* we never have to store more than the child's nnz */
     node->wsum_hess = child->wsum_hess->copy_sparsity(child->wsum_hess);
-    node->work->dwork = SP_MALLOC(child->size * sizeof(double));
+    node->work->dwork = sp_malloc(child->size * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -210,7 +210,7 @@ expr *new_sum(expr *child, int axis)
     }
 
     /* Allocate the type-specific struct */
-    sum_expr *snode = (sum_expr *) SP_CALLOC(1, sizeof(sum_expr));
+    sum_expr *snode = (sum_expr *) sp_calloc(1, sizeof(sum_expr));
     expr *node = &snode->base;
 
     /* to be consistent with CVXPY and NumPy we treat the result from

src/atoms/affine/trace.c

@@ -72,14 +72,14 @@ static void jacobian_init_impl(expr *node)
     // fill sparsity pattern and idx_map
     // ---------------------------------------------------------------
     trace_expr *tnode = (trace_expr *) node;
-    node->work->iwork = SP_MALLOC(MAX(jac->n, total_nnz) * sizeof(int));
+    node->work->iwork = sp_malloc(MAX(jac->n, total_nnz) * sizeof(int));
 
     /* the idx_map array maps each nonzero entry j in the original matrix A (from the
        selected, evenly spaced rows) to the corresponding index in the output row
        matrix C. Specifically, for each nonzero entry j in A (from the selected
        rows), idx_map[j] gives the position in C->x where the value from A->x[j]
        should be accumulated. */
-    tnode->idx_map = SP_MALLOC(A->nnz * sizeof(int));
+    tnode->idx_map = sp_malloc(A->nnz * sizeof(int));
     sum_spaced_rows_into_row_csr_alloc(A, jac, row_spacing, node->work->iwork,
                                        tnode->idx_map);
     node->jacobian = new_sparse_matrix(jac);
@@ -107,7 +107,7 @@ static void wsum_hess_init_impl(expr *node)
     /* initialize child's hessian */
     wsum_hess_init(x);
 
-    node->work->dwork = (double *) SP_CALLOC(x->size, sizeof(double));
+    node->work->dwork = (double *) sp_calloc(x->size, sizeof(double));
 
     /* We copy over the sparsity pattern from the child. This also includes the
        contribution to wsum_hess of entries of the child that will always have
@@ -148,7 +148,7 @@ static void free_type_data(expr *node)
 
 expr *new_trace(expr *child)
 {
-    trace_expr *tnode = (trace_expr *) SP_CALLOC(1, sizeof(trace_expr));
+    trace_expr *tnode = (trace_expr *) sp_calloc(1, sizeof(trace_expr));
     expr *node = &tnode->base;
     init_expr(node, 1, 1, child->n_vars, forward, jacobian_init_impl, eval_jacobian,
               is_affine, wsum_hess_init_impl, eval_wsum_hess, free_type_data);