[WIP] add memory tracking

include/expr.h

@@ -21,7 +21,7 @@
 #include "utils/CSC_Matrix.h"
 #include "utils/CSR_Matrix.h"
 #include <stdbool.h>
-#include <stddef.h>
+#include <stddef.h> /* size_t */
 #include <string.h>
 
 #define JAC_IDXS_NOT_SET -1
@@ -63,6 +63,8 @@ typedef struct expr
     //                         general quantities
     // ------------------------------------------------------------------------
     int d1, d2, size, n_vars, refcount, var_id;
+    size_t bytes;
+    bool visited;
     struct expr *left;
     struct expr *right;
 

include/problem.h

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

include/utils/COO_Matrix.h

@@ -2,6 +2,7 @@
 #define COO_MATRIX_H
 
 #include "CSR_Matrix.h"
+#include <stddef.h>
 
 /* COO (Coordinate) Sparse Matrix Format
  *
@@ -40,4 +41,8 @@ void refresh_lower_triangular_coo(COO_Matrix *coo, const double *vals);
 
 void free_coo_matrix(COO_Matrix *matrix);
 
+/* Returns total bytes used by rows, cols, x, value_map arrays
+   (0 if A is NULL) */
+size_t coo_bytes(const COO_Matrix *A);
+
 #endif /* COO_MATRIX_H */

include/utils/CSC_Matrix.h

@@ -2,6 +2,7 @@
 #define CSC_MATRIX_H
 
 #include "CSR_Matrix.h"
+#include <stddef.h>
 
 /* CSC (Compressed Sparse Column) Matrix Format
  *
@@ -23,18 +24,19 @@ typedef struct CSC_Matrix
     int nnz;
 } CSC_Matrix;
 
-/* constructor and destructor */
-CSC_Matrix *new_csc_matrix(int m, int n, int nnz);
+/* constructor and destructor.
+   If mem is non-NULL, *mem is incremented by the bytes allocated. */
+CSC_Matrix *new_csc_matrix(int m, int n, int nnz, size_t *mem);
 void free_csc_matrix(CSC_Matrix *matrix);
 
 /* Fill sparsity of C = A^T D A for diagonal D */
-CSR_Matrix *ATA_alloc(const CSC_Matrix *A);
+CSR_Matrix *ATA_alloc(const CSC_Matrix *A, size_t *mem);
 
 /* Fill sparsity of C = B^T D A for diagonal D */
-CSR_Matrix *BTA_alloc(const CSC_Matrix *A, const CSC_Matrix *B);
+CSR_Matrix *BTA_alloc(const CSC_Matrix *A, const CSC_Matrix *B, size_t *mem);
 
 /* Fill sparsity of C = BA, where B is symmetric. */
-CSC_Matrix *symBA_alloc(const CSR_Matrix *B, const CSC_Matrix *A);
+CSC_Matrix *symBA_alloc(const CSR_Matrix *B, const CSC_Matrix *A, size_t *mem);
 
 /* Compute values for C = A^T D A (null d corresponds to D as identity) */
 void ATDA_fill_values(const CSC_Matrix *A, const double *d, CSR_Matrix *C);
@@ -53,11 +55,14 @@ void yTA_fill_values(const CSC_Matrix *A, const double *x, CSR_Matrix *C);
 int count_nonzero_cols_csc(const CSC_Matrix *A);
 
 /* convert from CSR to CSC format */
-CSC_Matrix *csr_to_csc_alloc(const CSR_Matrix *A, int *iwork);
+CSC_Matrix *csr_to_csc_alloc(const CSR_Matrix *A, int *iwork, size_t *mem);
 void csr_to_csc_fill_values(const CSR_Matrix *A, CSC_Matrix *C, int *iwork);
 
 /* convert from CSC to CSR format */
-CSR_Matrix *csc_to_csr_alloc(const CSC_Matrix *A, int *iwork);
+CSR_Matrix *csc_to_csr_alloc(const CSC_Matrix *A, int *iwork, size_t *mem);
 void csc_to_csr_fill_values(const CSC_Matrix *A, CSR_Matrix *C, int *iwork);
 
+/* Returns total bytes used by p, i, x arrays (0 if A is NULL) */
+size_t csc_bytes(const CSC_Matrix *A);
+
 #endif /* CSC_MATRIX_H */

include/utils/CSR_Matrix.h

@@ -1,6 +1,7 @@
 #ifndef CSR_MATRIX_H
 #define CSR_MATRIX_H
 #include <stdbool.h>
+#include <stddef.h>
 
 /* CSR (Compressed Sparse Row) Matrix Format
  *
@@ -22,16 +23,17 @@ typedef struct CSR_Matrix
     int nnz;
 } CSR_Matrix;
 
-/* constructors and destructors */
-CSR_Matrix *new_csr_matrix(int m, int n, int nnz);
-CSR_Matrix *new_csr(const CSR_Matrix *A);
-CSR_Matrix *new_csr_copy_sparsity(const CSR_Matrix *A);
+/* constructors and destructors.
+   If mem is non-NULL, *mem is incremented by the bytes allocated. */
+CSR_Matrix *new_csr_matrix(int m, int n, int nnz, size_t *mem);
+CSR_Matrix *new_csr(const CSR_Matrix *A, size_t *mem);
+CSR_Matrix *new_csr_copy_sparsity(const CSR_Matrix *A, size_t *mem);
 void free_csr_matrix(CSR_Matrix *matrix);
 void copy_csr_matrix(const CSR_Matrix *A, CSR_Matrix *C);
 
 /* transpose functionality (iwork must be of size A->n) */
 CSR_Matrix *transpose(const CSR_Matrix *A, int *iwork);
-CSR_Matrix *AT_alloc(const CSR_Matrix *A, int *iwork);
+CSR_Matrix *AT_alloc(const CSR_Matrix *A, int *iwork, size_t *mem);
 void AT_fill_values(const CSR_Matrix *A, CSR_Matrix *AT, int *iwork);
 
 /* computes dense y = Ax */
@@ -49,6 +51,9 @@ void insert_idx(int idx, int *arr, int len);
 /* get value at position (row, col) in A */
 double csr_get_value(const CSR_Matrix *A, int row, int col);
 
+/* Returns total bytes used by p, i, x arrays (0 if A is NULL) */
+size_t csr_bytes(const CSR_Matrix *A);
+
 /* Expand symmetric CSR matrix A to full matrix C. A is assumed to store
    only upper triangle. C must be pre-allocated with sufficient nnz */
 void symmetrize_csr(const int *Ap, const int *Ai, int m, CSR_Matrix *C);

include/utils/CSR_sum.h

@@ -46,7 +46,7 @@ void sum_spaced_rows_into_row_csr_alloc(const CSR_Matrix *A, CSR_Matrix *C,
 /* Compute sparsity pattern of out = A + B + C + D */
 CSR_Matrix *sum_4_csr_alloc(const CSR_Matrix *A, const CSR_Matrix *B,
                             const CSR_Matrix *C, const CSR_Matrix *D,
-                            int *idx_maps[4]);
+                            int *idx_maps[4], size_t *mem);
 // ------------------------------------------------------------------------------------
 
 /* Accumulates values from A according to map. Must memset to zero before calling. */

include/utils/linalg_dense_sparse_matmuls.h

@@ -8,18 +8,18 @@
 /* C = (I_p kron A) @ J via the polymorphic Matrix interface.
  * A is dense m x n, J is (n*p) x k in CSC, C is (m*p) x k in CSC. */
 // TODO: maybe we can replace these with I_kron_X functionality?
-CSC_Matrix *I_kron_A_alloc(const Matrix *A, const CSC_Matrix *J, int p);
+CSC_Matrix *I_kron_A_alloc(const Matrix *A, const CSC_Matrix *J, int p, size_t *mem);
 void I_kron_A_fill_values(const Matrix *A, const CSC_Matrix *J, CSC_Matrix *C);
 
 /* Sparsity and values of C = (Y^T kron I_m) @ J where Y is k x n, J is (m*k) x p,
    and C is (m*n) x p. Y is given in column-major dense format. */
-CSR_Matrix *YT_kron_I_alloc(int m, int k, int n, const CSC_Matrix *J);
+CSR_Matrix *YT_kron_I_alloc(int m, int k, int n, const CSC_Matrix *J, size_t *mem);
 void YT_kron_I_fill_values(int m, int k, int n, const double *Y, const CSC_Matrix *J,
                            CSR_Matrix *C);
 
 /* Sparsity and values of C = (I_n kron X) @ J where X is m x k (col-major dense),
    J is (k*n) x p, and C is (m*n) x p. */
-CSR_Matrix *I_kron_X_alloc(int m, int k, int n, const CSC_Matrix *J);
+CSR_Matrix *I_kron_X_alloc(int m, int k, int n, const CSC_Matrix *J, size_t *mem);
 void I_kron_X_fill_values(int m, int k, int n, const double *X, const CSC_Matrix *J,
                           CSR_Matrix *C);
 

include/utils/linalg_sparse_matmuls.h

@@ -14,8 +14,8 @@
     * Mathematically it corresponds to  C = [A @ J1; A @ J2; ...; A @ Jp],
       where J = [J1; J2; ...; Jp]
 */
-CSC_Matrix *block_left_multiply_fill_sparsity(const CSR_Matrix *A,
-                                              const CSC_Matrix *J, int p);
+CSC_Matrix *block_left_multiply_alloc(const CSR_Matrix *A, const CSC_Matrix *J,
+                                      int p, size_t *mem);
 
 void block_left_multiply_fill_values(const CSR_Matrix *A, const CSC_Matrix *J,
                                      CSC_Matrix *C);
@@ -36,6 +36,7 @@ void csr_csc_matmul_fill_values(const CSR_Matrix *A, const CSC_Matrix *B,
 /* C = A @ B where A is CSR, B is CSC. Result C is CSR.
  * Allocates and precomputes sparsity pattern. No workspace required.
  */
-CSR_Matrix *csr_csc_matmul_alloc(const CSR_Matrix *A, const CSC_Matrix *B);
+CSR_Matrix *csr_csc_matmul_alloc(const CSR_Matrix *A, const CSC_Matrix *B,
+                                 size_t *mem);
 
 #endif /* LINALG_H */

include/utils/matrix.h

@@ -28,7 +28,7 @@ typedef struct Matrix
     void (*block_left_mult_vec)(const struct Matrix *self, const double *x,
                                 double *y, int p);
     CSC_Matrix *(*block_left_mult_sparsity)(const struct Matrix *self,
-                                            const CSC_Matrix *J, int p);
+                                            const CSC_Matrix *J, int p, size_t *mem);
     void (*block_left_mult_values)(const struct Matrix *self, const CSC_Matrix *J,
                                    CSC_Matrix *C);
     void (*free_fn)(struct Matrix *self);

src/atoms/affine/add.c

@@ -42,7 +42,7 @@ static void jacobian_init_impl(expr *node)
 
     /* we never have to store more than the sum of children's nnz */
     int nnz_max = node->left->jacobian->nnz + node->right->jacobian->nnz;
-    node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz_max);
+    node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz_max, &node->bytes);
 
     /* fill sparsity pattern  */
     sum_csr_alloc(node->left->jacobian, node->right->jacobian, node->jacobian);
@@ -66,7 +66,8 @@ static void wsum_hess_init_impl(expr *node)
 
     /* we never have to store more than the sum of children's nnz */
     int nnz_max = node->left->wsum_hess->nnz + node->right->wsum_hess->nnz;
-    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, nnz_max);
+    node->wsum_hess =
+        new_csr_matrix(node->n_vars, node->n_vars, nnz_max, &node->bytes);
 
     /* fill sparsity pattern of hessian */
     sum_csr_alloc(node->left->wsum_hess, node->right->wsum_hess, node->wsum_hess);

src/atoms/affine/broadcast.c

@@ -92,7 +92,8 @@ static void jacobian_init_impl(expr *node)
         total_nnz = x->jacobian->nnz * node->size;
     }
 
-    node->jacobian = new_csr_matrix(node->size, node->n_vars, total_nnz);
+    node->jacobian =
+        new_csr_matrix(node->size, node->n_vars, total_nnz, &node->bytes);
 
     // ---------------------------------------------------------------------
     //                 fill sparsity pattern
@@ -191,10 +192,11 @@ static void wsum_hess_init_impl(expr *node)
     wsum_hess_init(x);
 
     /* Same sparsity as child - weights get summed */
-    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
+    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess, &node->bytes);
 
     /* allocate space for weight vector */
     node->work->dwork = malloc(node->size * sizeof(double));
+    node->bytes += node->size * sizeof(double);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)

src/atoms/affine/const_scalar_mult.c

@@ -47,7 +47,7 @@ static void jacobian_init_impl(expr *node)
     jacobian_init(x);
 
     /* same sparsity as child */
-    node->jacobian = new_csr_copy_sparsity(x->jacobian);
+    node->jacobian = new_csr_copy_sparsity(x->jacobian, &node->bytes);
 }
 
 static void eval_jacobian(expr *node)
@@ -73,7 +73,7 @@ static void wsum_hess_init_impl(expr *node)
     wsum_hess_init(x);
 
     /* same sparsity as child */
-    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
+    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess, &node->bytes);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)

src/atoms/affine/const_vector_mult.c

@@ -46,7 +46,7 @@ static void jacobian_init_impl(expr *node)
     jacobian_init(x);
 
     /* same sparsity as child */
-    node->jacobian = new_csr_copy_sparsity(x->jacobian);
+    node->jacobian = new_csr_copy_sparsity(x->jacobian, &node->bytes);
 }
 
 static void eval_jacobian(expr *node)
@@ -75,9 +75,10 @@ static void wsum_hess_init_impl(expr *node)
     wsum_hess_init(x);
 
     /* same sparsity as child */
-    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
+    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess, &node->bytes);
 
     node->work->dwork = (double *) malloc(node->size * sizeof(double));
+    node->bytes += node->size * sizeof(double);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -123,6 +124,7 @@ expr *new_const_vector_mult(const double *a, expr *child)
 
     /* copy a vector */
     vnode->a = (double *) malloc(child->size * sizeof(double));
+    node->bytes += child->size * sizeof(double);
     memcpy(vnode->a, a, child->size * sizeof(double));
 
     return node;

src/atoms/affine/constant.c

@@ -30,7 +30,7 @@ static void jacobian_init_impl(expr *node)
 {
     /* Constant jacobian is all zeros: size x n_vars with 0 nonzeros.
      * new_csr_matrix uses calloc for row pointers, so they're already 0. */
-    node->jacobian = new_csr_matrix(node->size, node->n_vars, 0);
+    node->jacobian = new_csr_matrix(node->size, node->n_vars, 0, &node->bytes);
 }
 
 static void eval_jacobian(expr *node)
@@ -42,7 +42,7 @@ static void eval_jacobian(expr *node)
 static void wsum_hess_init_impl(expr *node)
 {
     /* Constant Hessian is all zeros: n_vars x n_vars with 0 nonzeros. */
-    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 0);
+    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, 0, &node->bytes);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)

src/atoms/affine/diag_vec.c

@@ -51,7 +51,7 @@ static void jacobian_init_impl(expr *node)
     jacobian_init(x);
 
     CSR_Matrix *Jx = x->jacobian;
-    CSR_Matrix *J = new_csr_matrix(node->size, node->n_vars, Jx->nnz);
+    CSR_Matrix *J = new_csr_matrix(node->size, node->n_vars, Jx->nnz, &node->bytes);
 
     /* Output has n^2 rows but only n diagonal positions are non-empty.
      * Diagonal position i is at row i*(n+1) in Fortran order. */
@@ -101,11 +101,12 @@ static void wsum_hess_init_impl(expr *node)
 
     /* workspace for extracting diagonal weights */
     node->work->dwork = (double *) calloc(x->size, sizeof(double));
+    node->bytes += x->size * sizeof(double);
 
     /* Copy child's Hessian structure (diag_vec is linear, so its own Hessian is
      * zero) */
     CSR_Matrix *Hx = x->wsum_hess;
-    node->wsum_hess = new_csr_copy_sparsity(Hx);
+    node->wsum_hess = new_csr_copy_sparsity(Hx, &node->bytes);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)

src/atoms/affine/hstack.c

@@ -55,7 +55,7 @@ static void jacobian_init_impl(expr *node)
         nnz += hnode->args[i]->jacobian->nnz;
     }
 
-    node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz);
+    node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz, &node->bytes);
 
     /* precompute sparsity pattern of this node's jacobian */
     int row_offset = 0;
@@ -113,8 +113,8 @@ static void wsum_hess_init_impl(expr *node)
 
     /* worst-case scenario the nnz of node->wsum_hess is the sum of children's
        nnz */
-    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, nnz);
-    hnode->CSR_work = new_csr_matrix(node->n_vars, node->n_vars, nnz);
+    node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, nnz, &node->bytes);
+    hnode->CSR_work = new_csr_matrix(node->n_vars, node->n_vars, nnz, &node->bytes);
 
     /* fill sparsity pattern */
     CSR_Matrix *H = node->wsum_hess;
@@ -193,6 +193,7 @@ expr *new_hstack(expr **args, int n_args, int n_vars)
 
     /* Set type-specific fields (deep copy args array) */
     hnode->args = (expr **) calloc(n_args, sizeof(expr *));
+    node->bytes += n_args * sizeof(expr *);
     hnode->n_args = n_args;
     for (int i = 0; i < n_args; i++)
     {

src/atoms/affine/index.c

@@ -64,7 +64,7 @@ static void jacobian_init_impl(expr *node)
     jacobian_init(x);
 
     CSR_Matrix *Jx = x->jacobian;
-    CSR_Matrix *J = new_csr_matrix(node->size, node->n_vars, Jx->nnz);
+    CSR_Matrix *J = new_csr_matrix(node->size, node->n_vars, Jx->nnz, &node->bytes);
 
     /* set sparsity pattern */
     J->p[0] = 0;
@@ -105,6 +105,7 @@ static void wsum_hess_init_impl(expr *node)
 
     /* for setting weight vector to evaluate hessian of child */
     node->work->dwork = (double *) calloc(x->size, sizeof(double));
+    node->bytes += 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
@@ -113,7 +114,7 @@ static void wsum_hess_init_impl(expr *node)
        structural zeros, but we do not try to exploit that sparsity
        right now. */
     CSR_Matrix *Hx = x->wsum_hess;
-    node->wsum_hess = new_csr_copy_sparsity(Hx);
+    node->wsum_hess = new_csr_copy_sparsity(Hx, &node->bytes);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -175,6 +176,7 @@ expr *new_index(expr *child, int d1, int d2, const int *indices, int n_idxs)
 
     /* copy indices */
     idx->indices = (int *) malloc(n_idxs * sizeof(int));
+    node->bytes += n_idxs * sizeof(int);
     memcpy(idx->indices, indices, n_idxs * sizeof(int));
     idx->n_idxs = n_idxs;