wrap malloc/alloc

include/problem.h

@@ -38,6 +38,7 @@ typedef struct
     int nnz_hessian;
     int n_vars;
     int total_constraint_size;
+    size_t memory_bytes;
 } Diff_engine_stats;
 
 typedef struct problem

include/utils/tracked_alloc.h

@@ -0,0 +1,23 @@
+#ifndef TRACKED_ALLOC_H
+#define TRACKED_ALLOC_H
+
+#include <stddef.h>
+#include <stdlib.h>
+
+extern size_t g_allocated_bytes;
+
+static inline void *SP_MALLOC(size_t size)
+{
+    void *ptr = malloc(size);
+    if (ptr) g_allocated_bytes += size;
+    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;
+    return ptr;
+}
+
+#endif /* TRACKED_ALLOC_H */

src/atoms/affine/add.c

@@ -17,6 +17,7 @@
  */
 #include "atoms/affine.h"
 #include "utils/CSR_sum.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -91,7 +92,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 *) 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

@@ -18,6 +18,7 @@
 #include "atoms/affine.h"
 #include "subexpr.h"
 #include "utils/mini_numpy.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -194,7 +195,7 @@ static void wsum_hess_init_impl(expr *node)
     node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
 
     /* allocate space for weight vector */
-    node->work->dwork = malloc(node->size * sizeof(double));
+    node->work->dwork = SP_MALLOC(node->size * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -273,7 +274,7 @@ expr *new_broadcast(expr *child, int d1, int d2)
         exit(1);
     }
 
-    broadcast_expr *bcast = (broadcast_expr *) calloc(1, sizeof(broadcast_expr));
+    broadcast_expr *bcast = (broadcast_expr *) SP_CALLOC(1, sizeof(broadcast_expr));
     expr *node = (expr *) bcast;
 
     // --------------------------------------------------------------------------

src/atoms/affine/const_scalar_mult.c

@@ -17,6 +17,7 @@
  */
 #include "atoms/affine.h"
 #include "subexpr.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -97,7 +98,7 @@ static bool is_affine(const expr *node)
 expr *new_const_scalar_mult(double a, expr *child)
 {
     const_scalar_mult_expr *mult_node =
-        (const_scalar_mult_expr *) calloc(1, sizeof(const_scalar_mult_expr));
+        (const_scalar_mult_expr *) SP_CALLOC(1, sizeof(const_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/const_vector_mult.c

@@ -17,6 +17,7 @@
  */
 #include "atoms/affine.h"
 #include "subexpr.h"
+#include "utils/tracked_alloc.h"
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -77,7 +78,7 @@ static void wsum_hess_init_impl(expr *node)
     /* same sparsity as child */
     node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
 
-    node->work->dwork = (double *) malloc(node->size * sizeof(double));
+    node->work->dwork = (double *) SP_MALLOC(node->size * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -112,7 +113,7 @@ static bool is_affine(const expr *node)
 expr *new_const_vector_mult(const double *a, expr *child)
 {
     const_vector_mult_expr *vnode =
-        (const_vector_mult_expr *) calloc(1, sizeof(const_vector_mult_expr));
+        (const_vector_mult_expr *) SP_CALLOC(1, sizeof(const_vector_mult_expr));
     expr *node = &vnode->base;
 
     init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init_impl,
@@ -122,7 +123,7 @@ expr *new_const_vector_mult(const double *a, expr *child)
     expr_retain(child);
 
     /* copy a vector */
-    vnode->a = (double *) malloc(child->size * sizeof(double));
+    vnode->a = (double *) SP_MALLOC(child->size * sizeof(double));
     memcpy(vnode->a, a, child->size * sizeof(double));
 
     return node;

src/atoms/affine/constant.c

@@ -16,6 +16,7 @@
  * limitations under the License.
  */
 #include "atoms/affine.h"
+#include "utils/tracked_alloc.h"
 #include <stdlib.h>
 #include <string.h>
 
@@ -60,7 +61,7 @@ static bool is_affine(const expr *node)
 
 expr *new_constant(int d1, int d2, int n_vars, const double *values)
 {
-    expr *node = (expr *) calloc(1, sizeof(expr));
+    expr *node = (expr *) SP_CALLOC(1, sizeof(expr));
     init_expr(node, d1, d2, n_vars, forward, jacobian_init_impl, eval_jacobian,
               is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
     memcpy(node->value, values, node->size * sizeof(double));

src/atoms/affine/diag_vec.c

@@ -18,6 +18,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #include "atoms/affine.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdlib.h>
 #include <string.h>
@@ -100,7 +101,7 @@ static void wsum_hess_init_impl(expr *node)
     wsum_hess_init(x);
 
     /* workspace for extracting diagonal weights */
-    node->work->dwork = (double *) 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) */
@@ -136,7 +137,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 *) 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

@@ -17,6 +17,7 @@
  */
 #include "atoms/affine.h"
 #include "utils/CSR_sum.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -185,14 +186,14 @@ expr *new_hstack(expr **args, int n_args, int n_vars)
     }
 
     /* Allocate the type-specific struct */
-    hstack_expr *hnode = (hstack_expr *) 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 **) 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

@@ -17,6 +17,7 @@
  */
 #include "atoms/affine.h"
 #include "subexpr.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -28,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 *) 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++)
     {
@@ -104,7 +105,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 *) 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
@@ -163,7 +164,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 *) 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,
@@ -174,7 +175,7 @@ expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs)
     expr_retain(child);
 
     /* copy indices */
-    idx->indices = (int *) 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

@@ -46,6 +46,7 @@
     vector-valued or matrix-valued.
 */
 
+#include "utils/tracked_alloc.h"
 #include "utils/utils.h"
 
 static void forward(expr *node, const double *u)
@@ -125,7 +126,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 *) malloc(dim * sizeof(double));
+    node->work->dwork = (double *) SP_MALLOC(dim * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -167,7 +168,7 @@ expr *new_left_matmul(expr *u, const CSR_Matrix *A)
 
     /* Allocate the type-specific struct */
     left_matmul_expr *lnode =
-        (left_matmul_expr *) 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_impl, eval_jacobian,
               is_affine, wsum_hess_init_impl, eval_wsum_hess, free_type_data);
@@ -178,8 +179,8 @@ expr *new_left_matmul(expr *u, const CSR_Matrix *A)
        (requiring size node->n_vars) and for transposing A (requiring size A->n).
        csc_to_csr_work is used for converting J_CSC to CSR (requiring
        node->size) */
-    node->work->iwork = (int *) malloc(MAX(A->n, node->n_vars) * sizeof(int));
-    lnode->csc_to_csr_work = (int *) malloc(node->size * sizeof(int));
+    node->work->iwork = (int *) SP_MALLOC(MAX(A->n, 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 */
@@ -212,15 +213,15 @@ expr *new_left_matmul_dense(expr *u, int m, int n, const double *data)
     }
 
     left_matmul_expr *lnode =
-        (left_matmul_expr *) 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_impl, eval_jacobian,
               is_affine, wsum_hess_init_impl, eval_wsum_hess, free_type_data);
     node->left = u;
     expr_retain(u);
 
-    node->work->iwork = (int *) malloc(MAX(n, node->n_vars) * sizeof(int));
-    lnode->csc_to_csr_work = (int *) malloc(node->size * sizeof(int));
+    node->work->iwork = (int *) SP_MALLOC(MAX(n, node->n_vars) * sizeof(int));
+    lnode->csc_to_csr_work = (int *) SP_MALLOC(node->size * sizeof(int));
     lnode->n_blocks = n_blocks;
 
     lnode->A = new_dense_matrix(m, n, data);

src/atoms/affine/neg.c

@@ -16,6 +16,7 @@
  * limitations under the License.
  */
 #include "atoms/affine.h"
+#include "utils/tracked_alloc.h"
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -87,7 +88,7 @@ static bool is_affine(const expr *node)
 
 expr *new_neg(expr *child)
 {
-    expr *node = (expr *) 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/promote.c

@@ -16,6 +16,7 @@
  * limitations under the License.
  */
 #include "atoms/affine.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdlib.h>
 #include <string.h>
@@ -106,7 +107,7 @@ static bool is_affine(const expr *node)
 expr *new_promote(expr *child, int d1, int d2)
 {
     assert(child->size == 1);
-    expr *node = (expr *) 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

@@ -16,6 +16,7 @@
  * limitations under the License.
  */
 #include "atoms/affine.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -67,7 +68,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 *) 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

@@ -18,6 +18,7 @@
 #include "atoms/affine.h"
 
 #include "utils/CSR_Matrix.h"
+#include "utils/tracked_alloc.h"
 #include <stdlib.h>
 
 /* This file implements the atom 'right_matmul' corresponding to the operation y =
@@ -28,7 +29,7 @@ expr *new_right_matmul(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 *) 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);
@@ -44,7 +45,7 @@ expr *new_right_matmul_dense(expr *u, int m, int n, const double *data)
 {
     /* We express: u @ A = (A^T @ u^T)^T
        A is m x n, so A^T is n x m. */
-    double *AT = (double *) malloc(n * m * sizeof(double));
+    double *AT = (double *) SP_MALLOC(n * m * sizeof(double));
     for (int i = 0; i < m; i++)
     {
         for (int j = 0; j < n; j++)