add memory tracking

Author: dance858

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 memory_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 memory_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_memory_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
  *
@@ -60,4 +61,7 @@ void csr_to_csc_fill_values(const CSR_Matrix *A, CSC_Matrix *C, int *iwork);
 CSR_Matrix *csc_to_csr_alloc(const CSC_Matrix *A, int *iwork);
 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_memory_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
  *
@@ -49,6 +50,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_memory_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);

src/atoms/affine/add.c

@@ -43,6 +43,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->memory_bytes += csr_memory_bytes(node->jacobian);
 
     /* fill sparsity pattern  */
     sum_csr_alloc(node->left->jacobian, node->right->jacobian, node->jacobian);
@@ -67,6 +68,7 @@ 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->memory_bytes += csr_memory_bytes(node->wsum_hess);
 
     /* fill sparsity pattern of hessian */
     sum_csr_alloc(node->left->wsum_hess, node->right->wsum_hess, node->wsum_hess);

src/atoms/affine/broadcast.c

@@ -93,6 +93,7 @@ static void jacobian_init_impl(expr *node)
     }
 
     node->jacobian = new_csr_matrix(node->size, node->n_vars, total_nnz);
+    node->memory_bytes += csr_memory_bytes(node->jacobian);
 
     // ---------------------------------------------------------------------
     //                 fill sparsity pattern
@@ -192,9 +193,11 @@ static void wsum_hess_init_impl(expr *node)
 
     /* Same sparsity as child - weights get summed */
     node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
+    node->memory_bytes += csr_memory_bytes(node->wsum_hess);
 
     /* allocate space for weight vector */
     node->work->dwork = malloc(node->size * sizeof(double));
+    node->memory_bytes += node->size * sizeof(double);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)

src/atoms/affine/const_scalar_mult.c

@@ -48,6 +48,7 @@ static void jacobian_init_impl(expr *node)
 
     /* same sparsity as child */
     node->jacobian = new_csr_copy_sparsity(x->jacobian);
+    node->memory_bytes += csr_memory_bytes(node->jacobian);
 }
 
 static void eval_jacobian(expr *node)
@@ -74,6 +75,7 @@ static void wsum_hess_init_impl(expr *node)
 
     /* same sparsity as child */
     node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
+    node->memory_bytes += csr_memory_bytes(node->wsum_hess);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)

src/atoms/affine/const_vector_mult.c

@@ -47,6 +47,7 @@ static void jacobian_init_impl(expr *node)
 
     /* same sparsity as child */
     node->jacobian = new_csr_copy_sparsity(x->jacobian);
+    node->memory_bytes += csr_memory_bytes(node->jacobian);
 }
 
 static void eval_jacobian(expr *node)
@@ -76,8 +77,10 @@ static void wsum_hess_init_impl(expr *node)
 
     /* same sparsity as child */
     node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess);
+    node->memory_bytes += csr_memory_bytes(node->wsum_hess);
 
     node->work->dwork = (double *) malloc(node->size * sizeof(double));
+    node->memory_bytes += node->size * sizeof(double);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -123,6 +126,7 @@ expr *new_const_vector_mult(const double *a, expr *child)
 
     /* copy a vector */
     vnode->a = (double *) malloc(child->size * sizeof(double));
+    node->memory_bytes += child->size * sizeof(double);
     memcpy(vnode->a, a, child->size * sizeof(double));
 
     return node;

src/atoms/affine/constant.c

@@ -31,6 +31,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->memory_bytes += csr_memory_bytes(node->jacobian);
 }
 
 static void eval_jacobian(expr *node)
@@ -43,6 +44,7 @@ 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->memory_bytes += csr_memory_bytes(node->wsum_hess);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)