add work struct for expression (#60)

* add work struct for expression

* run formatter

Author: dance858

include/expr.h

@@ -39,7 +39,24 @@ typedef void (*local_wsum_hess_fn)(struct expr *node, double *out, const double
 typedef bool (*is_affine_fn)(const struct expr *node);
 typedef void (*free_type_data_fn)(struct expr *node);
 
-/* Base expression node structure - contains only common fields */
+/* Workspace for derivative computation */
+typedef struct
+{
+    double *dwork;
+    int *iwork;
+    CSC_Matrix *jacobian_csc;
+    int *csc_work; /* for CSR-CSC conversion */
+
+    /* jacobian_csc_filled is only used for affine functions to avoid redundant
+       conversions. Could become relevant for non-affine functions if we start
+       supporting common subexpressions on the Python side. */
+    bool jacobian_csc_filled;
+    double *local_jac_diag; /* cached f'(g(x)) diagonal */
+    CSR_Matrix *hess_term1; /* Jg^T D Jg workspace */
+    CSR_Matrix *hess_term2; /* child wsum_hess workspace */
+} Expr_Work;
+
+/* Base expression node structure */
 typedef struct expr
 {
     // ------------------------------------------------------------------------
@@ -48,24 +65,13 @@ typedef struct expr
     int d1, d2, size, n_vars, refcount, var_id;
     struct expr *left;
     struct expr *right;
-    double *dwork;
-    int *iwork;
 
     // ------------------------------------------------------------------------
     //                     oracle related quantities
     // ------------------------------------------------------------------------
     double *value;
     CSR_Matrix *jacobian;
-    CSC_Matrix *jacobian_csc;
-    int *csc_work; /* workspace for CSR-CSC conversion */
-
-    /* jacobian_csc_filled is only used for affine functions to avoid redundant
-       conversions. Could become relevant for non-affine functions if we start
-       supporting common subexpressions on the Python side. */
-    bool jacobian_csc_filled;
     CSR_Matrix *wsum_hess;
-    CSR_Matrix *hess_term1; /* Jg^T D Jg workspace */
-    CSR_Matrix *hess_term2; /* child wsum_hess workspace */
     forward_fn forward;
     jacobian_init_fn jacobian_init;
     wsum_hess_init_fn wsum_hess_init;
@@ -76,10 +82,10 @@ typedef struct expr
     //                      other things
     // ------------------------------------------------------------------------
     is_affine_fn is_affine;
-    double *local_jac_diag;             /* cached f'(g(x)) diagonal */
     local_jacobian_fn local_jacobian;   /* used by elementwise univariate atoms*/
     local_wsum_hess_fn local_wsum_hess; /* used by elementwise univariate atoms*/
     free_type_data_fn free_type_data;   /* Cleanup for type-specific fields */
+    Expr_Work *work;                    /* derivative workspace */
 
     // name of node just for debugging - should be removed later
     char name[32];

src/affine/broadcast.c

@@ -196,7 +196,7 @@ static void wsum_hess_init(expr *node)
     memcpy(node->wsum_hess->i, x->wsum_hess->i, x->wsum_hess->nnz * sizeof(int));
 
     /* allocate space for weight vector */
-    node->dwork = malloc(node->size * sizeof(double));
+    node->work->dwork = malloc(node->size * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -205,7 +205,7 @@ static void eval_wsum_hess(expr *node, const double *w)
     expr *x = node->left;
 
     /* Zero out the work array first */
-    memset(node->dwork, 0, x->size * sizeof(double));
+    memset(node->work->dwork, 0, x->size * sizeof(double));
 
     if (bcast->type == BROADCAST_ROW)
     {
@@ -214,7 +214,7 @@ static void eval_wsum_hess(expr *node, const double *w)
         {
             for (int i = 0; i < node->d1; i++)
             {
-                node->dwork[j] += w[i + j * node->d1];
+                node->work->dwork[j] += w[i + j * node->d1];
             }
         }
     }
@@ -225,21 +225,21 @@ static void eval_wsum_hess(expr *node, const double *w)
         {
             for (int i = 0; i < node->d1; i++)
             {
-                node->dwork[i] += w[i + j * node->d1];
+                node->work->dwork[i] += w[i + j * node->d1];
             }
         }
     }
     else
     {
         /* (1, 1) -> (m, n): scalar has m*n weights to sum */
-        node->dwork[0] = 0.0;
+        node->work->dwork[0] = 0.0;
         for (int k = 0; k < node->size; k++)
         {
-            node->dwork[0] += w[k];
+            node->work->dwork[0] += w[k];
         }
     }
 
-    x->eval_wsum_hess(x, node->dwork);
+    x->eval_wsum_hess(x, node->work->dwork);
     memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));
 }
 

src/affine/diag_vec.c

@@ -100,7 +100,7 @@ static void wsum_hess_init(expr *node)
     x->wsum_hess_init(x);
 
     /* workspace for extracting diagonal weights */
-    node->dwork = (double *) calloc(x->size, sizeof(double));
+    node->work->dwork = (double *) calloc(x->size, sizeof(double));
 
     /* Copy child's Hessian structure (diag_vec is linear, so its own Hessian is
      * zero) */
@@ -118,11 +118,11 @@ static void eval_wsum_hess(expr *node, const double *w)
     /* Extract weights from diagonal positions of w (which has n^2 elements) */
     for (int i = 0; i < n; i++)
     {
-        node->dwork[i] = w[i * (n + 1)];
+        node->work->dwork[i] = w[i * (n + 1)];
     }
 
     /* Evaluate child's Hessian with extracted weights */
-    x->eval_wsum_hess(x, node->dwork);
+    x->eval_wsum_hess(x, node->work->dwork);
     memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));
 }
 

src/affine/index.c

@@ -104,7 +104,7 @@ static void wsum_hess_init(expr *node)
     x->wsum_hess_init(x);
 
     /* for setting weight vector to evaluate hessian of child */
-    node->dwork = (double *) calloc(x->size, sizeof(double));
+    node->work->dwork = (double *) 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
@@ -126,23 +126,23 @@ static void eval_wsum_hess(expr *node, const double *w)
     if (idx->has_duplicates)
     {
         /* zero and accumulate for repeated indices */
-        memset(node->dwork, 0, x->size * sizeof(double));
+        memset(node->work->dwork, 0, x->size * sizeof(double));
         for (int i = 0; i < idx->n_idxs; i++)
         {
-            node->dwork[idx->indices[i]] += w[i];
+            node->work->dwork[idx->indices[i]] += w[i];
         }
     }
     else
     {
         /* direct write (no memset needed, no accumulation) */
         for (int i = 0; i < idx->n_idxs; i++)
         {
-            node->dwork[idx->indices[i]] = w[i];
+            node->work->dwork[idx->indices[i]] = w[i];
         }
     }
 
     /* evalute hessian of child */
-    x->eval_wsum_hess(x, node->dwork);
+    x->eval_wsum_hess(x, node->work->dwork);
     memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));
 }
 

src/affine/sum.c

@@ -88,7 +88,8 @@ static void jacobian_init(expr *node)
 
     /* we never have to store more than the child's nnz */
     node->jacobian = new_csr_matrix(node->size, node->n_vars, x->jacobian->nnz);
-    node->iwork = malloc(MAX(node->jacobian->n, x->jacobian->nnz) * sizeof(int));
+    node->work->iwork =
+        malloc(MAX(node->jacobian->n, x->jacobian->nnz) * sizeof(int));
     snode->idx_map = malloc(x->jacobian->nnz * sizeof(int));
 
     /* the idx_map array maps each nonzero entry j in x->jacobian
@@ -98,18 +99,19 @@ static void jacobian_init(expr *node)
 
     if (axis == -1)
     {
-        sum_all_rows_csr_fill_sparsity_and_idx_map(x->jacobian, node->jacobian,
-                                                   node->iwork, snode->idx_map);
+        sum_all_rows_csr_fill_sparsity_and_idx_map(
+            x->jacobian, node->jacobian, node->work->iwork, snode->idx_map);
     }
     else if (axis == 0)
     {
         sum_block_of_rows_csr_fill_sparsity_and_idx_map(
-            x->jacobian, node->jacobian, x->d1, node->iwork, snode->idx_map);
+            x->jacobian, node->jacobian, x->d1, node->work->iwork, snode->idx_map);
     }
     else if (axis == 1)
     {
         sum_evenly_spaced_rows_csr_fill_sparsity_and_idx_map(
-            x->jacobian, node->jacobian, node->size, node->iwork, snode->idx_map);
+            x->jacobian, node->jacobian, node->size, node->work->iwork,
+            snode->idx_map);
     }
 }
 
@@ -135,7 +137,7 @@ static void wsum_hess_init(expr *node)
 
     /* we never have to store more than the child's nnz */
     node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, x->wsum_hess->nnz);
-    node->dwork = malloc(x->size * sizeof(double));
+    node->work->dwork = malloc(x->size * sizeof(double));
 
     /* copy sparsity pattern */
     memcpy(node->wsum_hess->p, x->wsum_hess->p, (x->n_vars + 1) * sizeof(int));
@@ -150,18 +152,18 @@ static void eval_wsum_hess(expr *node, const double *w)
 
     if (axis == -1)
     {
-        scaled_ones(node->dwork, x->size, *w);
+        scaled_ones(node->work->dwork, x->size, *w);
     }
     else if (axis == 0)
     {
-        repeat(node->dwork, w, x->d2, x->d1);
+        repeat(node->work->dwork, w, x->d2, x->d1);
     }
     else if (axis == 1)
     {
-        tile_double(node->dwork, w, x->d1, x->d2);
+        tile_double(node->work->dwork, w, x->d1, x->d2);
     }
 
-    x->eval_wsum_hess(x, node->dwork);
+    x->eval_wsum_hess(x, node->work->dwork);
 
     /* copy values */
     memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));

src/affine/trace.c

@@ -69,7 +69,7 @@ static void jacobian_init(expr *node)
     // fill sparsity pattern and idx_map
     // ---------------------------------------------------------------
     trace_expr *tnode = (trace_expr *) node;
-    node->iwork = malloc(MAX(node->jacobian->n, total_nnz) * sizeof(int));
+    node->work->iwork = malloc(MAX(node->jacobian->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
@@ -78,7 +78,7 @@ static void jacobian_init(expr *node)
        should be accumulated. */
     tnode->idx_map = malloc(x->jacobian->nnz * sizeof(int));
     sum_spaced_rows_into_row_csr_fill_sparsity_and_idx_map(
-        A, node->jacobian, row_spacing, node->iwork, tnode->idx_map);
+        A, node->jacobian, row_spacing, node->work->iwork, tnode->idx_map);
 }
 
 static void eval_jacobian(expr *node)
@@ -104,7 +104,7 @@ static void wsum_hess_init(expr *node)
     x->wsum_hess_init(x);
 
     node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, x->wsum_hess->nnz);
-    node->dwork = (double *) calloc(x->size, sizeof(double));
+    node->work->dwork = (double *) 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
@@ -121,10 +121,10 @@ static void eval_wsum_hess(expr *node, const double *w)
     int row_spacing = x->d1 + 1;
     for (int i = 0; i < x->size; i += row_spacing)
     {
-        node->dwork[i] = w[0];
+        node->work->dwork[i] = w[0];
     }
 
-    x->eval_wsum_hess(x, node->dwork);
+    x->eval_wsum_hess(x, node->work->dwork);
 
     memcpy(node->wsum_hess->x, x->wsum_hess->x, sizeof(double) * x->wsum_hess->nnz);
 }

src/affine/transpose.c

@@ -99,7 +99,7 @@ static void wsum_hess_init(expr *node)
     node->wsum_hess = H;
 
     /* for computing Kw where K is the commutation matrix */
-    node->dwork = (double *) malloc(node->size * sizeof(double));
+    node->work->dwork = (double *) malloc(node->size * sizeof(double));
 }
 static void eval_wsum_hess(expr *node, const double *w)
 {
@@ -112,11 +112,11 @@ static void eval_wsum_hess(expr *node, const double *w)
     {
         for (int j = 0; j < d1; ++j)
         {
-            node->dwork[j * d2 + i] = w[i * d1 + j];
+            node->work->dwork[j * d2 + i] = w[i * d1 + j];
         }
     }
 
-    node->left->eval_wsum_hess(node->left, node->dwork);
+    node->left->eval_wsum_hess(node->left, node->work->dwork);
 
     /* copy to this node's hessian */
     memcpy(node->wsum_hess->x, node->left->wsum_hess->x,

src/bivariate/const_vector_mult.c

@@ -81,7 +81,7 @@ static void wsum_hess_init(expr *node)
     memcpy(node->wsum_hess->p, x->wsum_hess->p, (node->n_vars + 1) * sizeof(int));
     memcpy(node->wsum_hess->i, x->wsum_hess->i, x->wsum_hess->nnz * sizeof(int));
 
-    node->dwork = (double *) malloc(node->size * sizeof(double));
+    node->work->dwork = (double *) malloc(node->size * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -92,10 +92,10 @@ static void eval_wsum_hess(expr *node, const double *w)
     /* scale weights w by a */
     for (int i = 0; i < node->size; i++)
     {
-        node->dwork[i] = a[i] * w[i];
+        node->work->dwork[i] = a[i] * w[i];
     }
 
-    x->eval_wsum_hess(x, node->dwork);
+    x->eval_wsum_hess(x, node->work->dwork);
 
     /* copy values from child to this node */
     memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));

src/bivariate/left_matmul.c

@@ -88,7 +88,7 @@ static void jacobian_init(expr *node)
 
     /* initialize child's jacobian and precompute sparsity of its CSC */
     x->jacobian_init(x);
-    lnode->Jchild_CSC = csr_to_csc_fill_sparsity(x->jacobian, node->iwork);
+    lnode->Jchild_CSC = csr_to_csc_fill_sparsity(x->jacobian, node->work->iwork);
 
     /* precompute sparsity of this node's jacobian in CSC and CSR */
     lnode->J_CSC = lnode->A->block_left_mult_sparsity(lnode->A, lnode->Jchild_CSC,
@@ -106,7 +106,7 @@ static void eval_jacobian(expr *node)
 
     /* evaluate child's jacobian and convert to CSC */
     x->eval_jacobian(x);
-    csr_to_csc_fill_values(x->jacobian, Jchild_CSC, node->iwork);
+    csr_to_csc_fill_values(x->jacobian, Jchild_CSC, node->work->iwork);
 
     /* compute this node's jacobian: */
     lnode->A->block_left_mult_values(lnode->A, Jchild_CSC, J_CSC);
@@ -127,17 +127,17 @@ static void wsum_hess_init(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->dwork = (double *) malloc(dim * sizeof(double));
+    node->work->dwork = (double *) malloc(dim * sizeof(double));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
 {
     /* compute A^T w*/
     Matrix *AT = ((left_matmul_expr *) node)->AT;
     int n_blocks = ((left_matmul_expr *) node)->n_blocks;
-    AT->block_left_mult_vec(AT, w, node->dwork, n_blocks);
+    AT->block_left_mult_vec(AT, w, node->work->dwork, n_blocks);
 
-    node->left->eval_wsum_hess(node->left, node->dwork);
+    node->left->eval_wsum_hess(node->left, node->work->dwork);
     memcpy(node->wsum_hess->x, node->left->wsum_hess->x,
            node->wsum_hess->nnz * sizeof(double));
 }
@@ -180,13 +180,14 @@ 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->iwork = (int *) malloc(MAX(A->n, node->n_vars) * sizeof(int));
+    node->work->iwork = (int *) malloc(MAX(A->n, node->n_vars) * sizeof(int));
     lnode->csc_to_csr_work = (int *) malloc(node->size * sizeof(int));
     lnode->n_blocks = n_blocks;
 
     /* store A and AT */
     lnode->A = new_sparse_matrix(A);
-    lnode->AT = sparse_matrix_trans((const Sparse_Matrix *) lnode->A, node->iwork);
+    lnode->AT =
+        sparse_matrix_trans((const Sparse_Matrix *) lnode->A, node->work->iwork);
 
     return node;
 }
@@ -220,7 +221,7 @@ expr *new_left_matmul_dense(expr *u, int m, int n, const double *data)
     node->left = u;
     expr_retain(u);
 
-    node->iwork = (int *) malloc(MAX(n, node->n_vars) * sizeof(int));
+    node->work->iwork = (int *) malloc(MAX(n, node->n_vars) * sizeof(int));
     lnode->csc_to_csr_work = (int *) malloc(node->size * sizeof(int));
     lnode->n_blocks = n_blocks;