Precompute sparsity pattern for a few atoms and start on trace

README.md

@@ -2,4 +2,13 @@
 4. in the refactor, add consts
 5. multiply with one constant vector/scalar argument
 6. AX where X is a matrix. Can that happen? How is that canonicalized? Maybe it can't happen.
-7. Must be able to compute jacobian and hessian of A @ phi(x), so linear operator needs other code! This requires new infrastructure, I think.
+7. Must be able to compute jacobian and hessian of A @ phi(x), so linear operator needs other code! This requires new infrastructure, I think.
+8. Shortcut hessians of affine stuff?
+
+Going through all atoms to see that sparsity pattern is computed in initialization of jacobian:
+2. trace - not ok
+
+Going through all atoms to see that sparsity pattern is computed in initialization of hessian:
+2. hstack - not ok
+3. trace - not ok
+

include/affine.h

@@ -11,6 +11,7 @@ expr *new_add(expr *left, expr *right);
 
 expr *new_sum(expr *child, int axis);
 expr *new_hstack(expr **args, int n_args, int n_vars);
+expr *new_trace(expr *child);
 
 expr *new_constant(int d1, int d2, int n_vars, const double *values);
 expr *new_variable(int d1, int d2, int var_id, int n_vars);

include/expr.h

@@ -5,6 +5,7 @@
 #include "utils/CSR_Matrix.h"
 #include <stdbool.h>
 #include <stddef.h>
+#include <string.h>
 
 #define JAC_IDXS_NOT_SET -1
 #define NOT_A_VARIABLE -1
@@ -53,6 +54,9 @@ typedef struct expr
     local_wsum_hess_fn local_wsum_hess; /* used by elementwise univariate atoms*/
     free_type_data_fn free_type_data;   /* Cleanup for type-specific fields */
 
+    // name of node just for debugging - should be removed later
+    char name[32];
+
 } expr;
 
 void init_expr(expr *node, int d1, int d2, int n_vars, forward_fn forward,

include/subexpr.h

@@ -37,9 +37,16 @@ typedef struct sum_expr
 {
     expr base;
     int axis;
-    struct int_double_pair *int_double_pairs; /* for sorting jacobian entries */
+    int *idx_map; /* maps child nnz to summed-row positions */
 } sum_expr;
 
+/* Trace-like reduction: sums entries spaced by child->d1 */
+typedef struct trace_expr
+{
+    expr base;
+    struct int_double_pair *int_double_pairs; /* for sorting jacobian entries */
+} trace_expr;
+
 /* Product of all entries */
 typedef struct prod_expr
 {

include/utils/CSR_Matrix.h

@@ -52,29 +52,86 @@ void csr_insert_value(CSR_Matrix *A, int col_idx, double value);
  * d must have length m
  * C must be pre-allocated with same dimensions as A */
 void diag_csr_mult(const double *d, const CSR_Matrix *A, CSR_Matrix *C);
+void diag_csr_mult_fill_values(const double *d, const CSR_Matrix *A, CSR_Matrix *C);
 
 /* Compute C = A + B where A, B, C are CSR matrices
  * A and B must have same dimensions
  * C must be pre-allocated with sufficient nnz capacity.
  * C must be different from A and B */
 void sum_csr_matrices(const CSR_Matrix *A, const CSR_Matrix *B, CSR_Matrix *C);
+/* Compute sparsity pattern of A + B where A, B, C are CSR matrices.
+ * Fills C->p, C->i, and C->nnz; does not touch C->x. */
+void sum_csr_matrices_fill_sparsity(const CSR_Matrix *A, const CSR_Matrix *B,
+                                    CSR_Matrix *C);
+
+/* Fill only the values of C = A + B, assuming C's sparsity pattern (p and i)
+ * is already filled and matches the union of A and B per row. Does not modify
+ * C->p, C->i, or C->nnz. */
+void sum_csr_matrices_fill_values(const CSR_Matrix *A, const CSR_Matrix *B,
+                                  CSR_Matrix *C);
 
 /* Compute C = diag(d1) * A + diag(d2) * B where A, B, C are CSR matrices */
 void sum_scaled_csr_matrices(const CSR_Matrix *A, const CSR_Matrix *B, CSR_Matrix *C,
                              const double *d1, const double *d2);
 
+/* Fill only the values of C = diag(d1) * A + diag(d2) * B, assuming C's sparsity
+ * pattern (p and i) is already filled and matches the union of A and B per row.
+ * Does not modify C->p, C->i, or C->nnz. */
+void sum_scaled_csr_matrices_fill_values(const CSR_Matrix *A, const CSR_Matrix *B,
+                                         CSR_Matrix *C, const double *d1,
+                                         const double *d2);
+
 /* Sum all rows of A into a single row matrix C */
 void sum_all_rows_csr(const CSR_Matrix *A, CSR_Matrix *C,
                       struct int_double_pair *pairs);
 
+/* iwork must have size max(C->n, A->nnz), and idx_map must have size A->nnz. */
+void sum_all_rows_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A, CSR_Matrix *C,
+                                                int *iwork, int *idx_map);
+
+/* Fill values of summed rows using precomputed idx_map and sparsity of C */
+// void sum_all_rows_csr_fill_values(const CSR_Matrix *A, CSR_Matrix *C,
+//                                  const int *idx_map);
+
+/* Fill accumulator for summing rows using precomputed idx_map for each nnz of A*/
+void idx_map_accumulator(const CSR_Matrix *A, const int *idx_map,
+                         double *accumulator);
+
 /* Sum blocks of rows of A into a matrix C */
 void sum_block_of_rows_csr(const CSR_Matrix *A, CSR_Matrix *C,
                            struct int_double_pair *pairs, int row_block_size);
 
+/* Build sparsity and index map for summing blocks of rows.
+ * iwork must have size max(A->n, A->nnz), and idx_map must have size A->nnz. */
+void sum_block_of_rows_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A,
+                                                     CSR_Matrix *C,
+                                                     int row_block_size, int *iwork,
+                                                     int *idx_map);
+
+/* Fill values for summing blocks of rows using precomputed idx_map */
+// void sum_block_of_rows_csr_fill_values(const CSR_Matrix *A, CSR_Matrix *C,
+//                                        const int *idx_map);
+
 /* Sum evenly spaced rows of A into a matrix C */
 void sum_evenly_spaced_rows_csr(const CSR_Matrix *A, CSR_Matrix *C,
                                 struct int_double_pair *pairs, int row_spacing);
 
+/* Build sparsity and index map for summing evenly spaced rows.
+ * iwork must have size max(A->n, A->nnz), and idx_map must have size A->nnz. */
+void sum_evenly_spaced_rows_csr_fill_sparsity_and_idx_map(const CSR_Matrix *A,
+                                                          CSR_Matrix *C,
+                                                          int row_spacing,
+                                                          int *iwork, int *idx_map);
+
+/* Fill values for summing evenly spaced rows using precomputed idx_map */
+// void sum_evenly_spaced_rows_csr_fill_values(const CSR_Matrix *A, CSR_Matrix *C,
+//                                             const int *idx_map);
+
+/* Sum evenly spaced rows of A starting at offset into a row matrix C */
+void sum_spaced_rows_into_row_csr(const CSR_Matrix *A, CSR_Matrix *C,
+                                  struct int_double_pair *pairs, int offset,
+                                  int spacing);
+
 /* Count number of columns with nonzero entries */
 int count_nonzero_cols(const CSR_Matrix *A, bool *col_nz);
 

src/affine/add.c

@@ -22,6 +22,10 @@ static void jacobian_init(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);
+
+    /* fill sparsity pattern  */
+    sum_csr_matrices_fill_sparsity(node->left->jacobian, node->right->jacobian,
+                                   node->jacobian);
 }
 
 static void eval_jacobian(expr *node)
@@ -31,7 +35,8 @@ static void eval_jacobian(expr *node)
     node->right->eval_jacobian(node->right);
 
     /* sum children's jacobians */
-    sum_csr_matrices(node->left->jacobian, node->right->jacobian, node->jacobian);
+    sum_csr_matrices_fill_values(node->left->jacobian, node->right->jacobian,
+                                 node->jacobian);
 }
 
 static void wsum_hess_init(expr *node)
@@ -43,6 +48,10 @@ static void wsum_hess_init(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);
+
+    /* TODO: we should fill sparsity pattern here for consistency */
+    sum_csr_matrices_fill_sparsity(node->left->wsum_hess, node->right->wsum_hess,
+                                   node->wsum_hess);
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -52,7 +61,8 @@ static void eval_wsum_hess(expr *node, const double *w)
     node->right->eval_wsum_hess(node->right, w);
 
     /* sum children's wsum_hess */
-    sum_csr_matrices(node->left->wsum_hess, node->right->wsum_hess, node->wsum_hess);
+    sum_csr_matrices_fill_values(node->left->wsum_hess, node->right->wsum_hess,
+                                 node->wsum_hess);
 }
 
 static bool is_affine(const expr *node)
@@ -74,5 +84,8 @@ expr *new_add(expr *left, expr *right)
     node->wsum_hess_init = wsum_hess_init;
     node->eval_wsum_hess = eval_wsum_hess;
 
+    // just for debugging, should be removed
+    strcpy(node->name, "add");
+
     return node;
 }

src/affine/sum.c

@@ -5,6 +5,8 @@
 #include <stdlib.h>
 #include <string.h>
 
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+
 static void forward(expr *node, const double *u)
 {
     int i, j, end;
@@ -62,42 +64,46 @@ static void jacobian_init(expr *node)
 {
     expr *x = node->left;
     sum_expr *snode = (sum_expr *) node;
+    int axis = snode->axis;
 
     /* initialize child's jacobian */
     x->jacobian_init(x);
 
     /* we never have to store more than the child's nnz */
     node->jacobian = new_csr_matrix(node->d1, node->n_vars, x->jacobian->nnz);
-    snode->int_double_pairs = new_int_double_pair_array(x->jacobian->nnz);
-}
-
-static void eval_jacobian(expr *node)
-{
-    expr *x = node->left;
-    sum_expr *snode = (sum_expr *) node;
-    int axis = snode->axis;
-
-    /* evaluate child's jacobian */
-    x->eval_jacobian(x);
+    node->iwork = malloc(MAX(node->jacobian->n, x->jacobian->nnz) * sizeof(int));
+    snode->idx_map = malloc(x->jacobian->nnz * sizeof(int));
 
-    /* sum rows or columns of child's jacobian */
     if (axis == -1)
     {
-        sum_all_rows_csr(x->jacobian, node->jacobian, snode->int_double_pairs);
+        sum_all_rows_csr_fill_sparsity_and_idx_map(x->jacobian, node->jacobian,
+                                                   node->iwork, snode->idx_map);
     }
     else if (axis == 0)
     {
-        sum_block_of_rows_csr(x->jacobian, node->jacobian, snode->int_double_pairs,
-                              x->d1);
+        sum_block_of_rows_csr_fill_sparsity_and_idx_map(
+            x->jacobian, node->jacobian, x->d1, node->iwork, snode->idx_map);
     }
-
     else if (axis == 1)
     {
-        sum_evenly_spaced_rows_csr(x->jacobian, node->jacobian,
-                                   snode->int_double_pairs, node->d1);
+        sum_evenly_spaced_rows_csr_fill_sparsity_and_idx_map(
+            x->jacobian, node->jacobian, node->d1, node->iwork, snode->idx_map);
     }
 }
 
+static void eval_jacobian(expr *node)
+{
+    expr *x = node->left;
+
+    /* evaluate child's jacobian */
+    x->eval_jacobian(x);
+
+    /* we have precomputed an idx map between the nonzeros of the child's jacobian
+       and this node's jacobian, so we just accumulate accordingly */
+    idx_map_accumulator(x->jacobian, ((sum_expr *) node)->idx_map,
+                        node->jacobian->x);
+}
+
 static void wsum_hess_init(expr *node)
 {
     expr *x = node->left;
@@ -107,6 +113,10 @@ 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));
+
+    /* copy sparsity pattern */
+    memcpy(node->wsum_hess->p, x->wsum_hess->p, (x->n_vars + 1) * sizeof(int));
+    memcpy(node->wsum_hess->i, x->wsum_hess->i, x->wsum_hess->nnz * sizeof(int));
 }
 
 static void eval_wsum_hess(expr *node, const double *w)
@@ -130,8 +140,8 @@ static void eval_wsum_hess(expr *node, const double *w)
 
     x->eval_wsum_hess(x, node->dwork);
 
-    /* todo: is this copy necessary or can we just change pointers? */
-    copy_csr_matrix(x->wsum_hess, node->wsum_hess);
+    /* copy values (TODO: is this necessary or can we just change pointers?)*/
+    memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));
 }
 
 static bool is_affine(const expr *node)
@@ -142,7 +152,7 @@ static bool is_affine(const expr *node)
 static void free_type_data(expr *node)
 {
     sum_expr *snode = (sum_expr *) node;
-    free_int_double_pair_array(snode->int_double_pairs);
+    free(snode->idx_map);
 }
 
 expr *new_sum(expr *child, int axis)
@@ -179,7 +189,6 @@ expr *new_sum(expr *child, int axis)
 
     /* Set type-specific fields */
     snode->axis = axis;
-    snode->int_double_pairs = NULL;
 
     return node;
 }