new approach

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>  /* size_t */
+#include <stddef.h> /* size_t */
 #include <string.h>
 
 #define JAC_IDXS_NOT_SET -1

include/utils/CSC_Matrix.h

@@ -24,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);
@@ -54,11 +55,11 @@ 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) */

include/utils/CSR_Matrix.h

@@ -23,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 */

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

@@ -13,13 +13,13 @@ 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

@@ -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 */

src/atoms/affine/add.c

@@ -42,8 +42,8 @@ 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);
+    node->jacobian =
+        new_csr_matrix(node->size, node->n_vars, nnz_max, &node->memory_bytes);
 
     /* fill sparsity pattern  */
     sum_csr_alloc(node->left->jacobian, node->right->jacobian, node->jacobian);
@@ -67,8 +67,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->memory_bytes += csr_memory_bytes(node->wsum_hess);
+    node->wsum_hess =
+        new_csr_matrix(node->n_vars, node->n_vars, nnz_max, &node->memory_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,8 +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->memory_bytes += csr_memory_bytes(node->jacobian);
+    node->jacobian =
+        new_csr_matrix(node->size, node->n_vars, total_nnz, &node->memory_bytes);
 
     // ---------------------------------------------------------------------
     //                 fill sparsity pattern
@@ -192,8 +192,7 @@ 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->memory_bytes += csr_memory_bytes(node->wsum_hess);
+    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess, &node->memory_bytes);
 
     /* allocate space for weight vector */
     node->work->dwork = malloc(node->size * sizeof(double));

src/atoms/affine/const_scalar_mult.c

@@ -47,8 +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->memory_bytes += csr_memory_bytes(node->jacobian);
+    node->jacobian = new_csr_copy_sparsity(x->jacobian, &node->memory_bytes);
 }
 
 static void eval_jacobian(expr *node)
@@ -74,8 +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->memory_bytes += csr_memory_bytes(node->wsum_hess);
+    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess, &node->memory_bytes);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)

src/atoms/affine/const_vector_mult.c

@@ -46,8 +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->memory_bytes += csr_memory_bytes(node->jacobian);
+    node->jacobian = new_csr_copy_sparsity(x->jacobian, &node->memory_bytes);
 }
 
 static void eval_jacobian(expr *node)
@@ -76,8 +75,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->memory_bytes += csr_memory_bytes(node->wsum_hess);
+    node->wsum_hess = new_csr_copy_sparsity(x->wsum_hess, &node->memory_bytes);
 
     node->work->dwork = (double *) malloc(node->size * sizeof(double));
     node->memory_bytes += node->size * sizeof(double);