Stack constraint and hstack Jacobians with vstack_matrices_*

include/problem.h

@@ -59,17 +59,14 @@ typedef struct problem
     double *gradient_values;
 
     /* allocated by problem_init_derivatives */
-    CSR_matrix *jacobian;
+    CSR_matrix *jacobian; /* CSR view of jac_stack; do not free */
+    matrix *jac_stack;    /* vertical stack of the constraint jacobians */
+    matrix **jac_ops;     /* the constraints' jacobians: jac_stack's operands */
     CSR_matrix *lagrange_hessian; /* CSR view of hess_sum; do not free */
     matrix *hess_sum;             /* sum of all wsum_hess */
     COO_matrix *jacobian_coo;
     COO_matrix *lagrange_hessian_coo; /* lower triangular part stored in COO */
 
-    /* per-constraint jacobian values_version last copied into the aggregated
-     * jacobian; the copy is skipped when the version is unchanged (affine
-     * constraints after their first eval) */
-    uint64_t *constraint_jac_seen;
-
     /* statistics for performance measurement */
     Diff_engine_stats stats;
     bool verbose;

include/subexpr.h

@@ -19,7 +19,6 @@
 #define SUBEXPR_H
 
 #include "expr.h"
-#include "sparsediff_linalg/CSC_matrix.h"
 #include "sparsediff_linalg/CSR_matrix.h"
 #include "sparsediff_linalg/matrix.h"
 
@@ -113,6 +112,7 @@ typedef struct hstack_expr
     expr base;
     expr **args;
     int n_args;
+    matrix **child_jac;  /* args' Jacobians: the operands of the vstack */
     matrix **child_hess; /* args' Hessians: the operands of the sum */
 } hstack_expr;
 
@@ -132,9 +132,6 @@ typedef struct left_matmul_expr
     matrix *A;
     matrix *AT;
     int n_blocks;
-    CSC_matrix *Jchild_CSC;
-    CSC_matrix *J_CSC;
-    int *csc_to_csr_work;
     expr *param_source;
     void (*refresh_param_values)(struct left_matmul_expr *);
 } left_matmul_expr;

src/atoms/affine/hstack.c

@@ -16,9 +16,9 @@
  * limitations under the License.
  */
 #include "atoms/affine.h"
-#include "sparsediff_linalg/CSR_matrix.h"
 #include "sparsediff_linalg/matrix.h"
 #include "sparsediff_linalg/matrix_sum.h"
+#include "sparsediff_linalg/matrix_vstack.h"
 #include "sparsediff_linalg/tracked_alloc.h"
 #include "subexpr.h"
 #include <assert.h>
@@ -50,58 +50,27 @@ static void jacobian_init_impl(expr *node)
 {
     hstack_expr *hnode = (hstack_expr *) node;
 
-    /* initialize children's jacobians */
-    int nnz = 0;
+    /* initialize children's jacobians; this node's is their vertical stack */
+    hnode->child_jac = (matrix **) sp_malloc(hnode->n_args * sizeof(matrix *));
     for (int i = 0; i < hnode->n_args; i++)
     {
         assert(hnode->args[i] != NULL);
         jacobian_init(hnode->args[i]);
-        nnz += hnode->args[i]->jacobian->nnz;
+        hnode->child_jac[i] = hnode->args[i]->jacobian;
     }
-
-    CSR_matrix *A = new_CSR_matrix(node->size, node->n_vars, nnz);
-
-    /* precompute sparsity pattern of this node's jacobian */
-    int row_offset = 0;
-    int cursor = 0;
-
-    for (int i = 0; i < hnode->n_args; i++)
-    {
-        expr *child = hnode->args[i];
-        CSR_matrix *B = child->jacobian->to_csr(child->jacobian);
-
-        /* copy columns */
-        memcpy(A->i + cursor, B->i, B->nnz * sizeof(int));
-
-        /* set row pointers */
-        for (int r = 0; r < child->size; r++)
-        {
-            A->p[row_offset + r] = cursor + B->p[r];
-        }
-
-        cursor += B->nnz;
-        row_offset += child->size;
-    }
-    A->p[node->size] = cursor;
-    assert(cursor == A->nnz);
-    node->jacobian = new_sparse_matrix(A);
+    node->jacobian = vstack_matrices_alloc(hnode->child_jac, hnode->n_args);
 }
 
 static void eval_jacobian_impl(expr *node)
 {
     hstack_expr *hnode = (hstack_expr *) node;
-    int cursor = 0;
-
     for (int i = 0; i < hnode->n_args; i++)
     {
-        expr *child = hnode->args[i];
-        eval_jacobian(child);
-        /* to_csr needed for stacked_pd */
-        CSR_matrix *child_csr = child->jacobian->to_csr(child->jacobian);
-        memcpy(node->jacobian->x + cursor, child_csr->x,
-               child_csr->nnz * sizeof(double));
-        cursor += child_csr->nnz;
+        eval_jacobian(hnode->args[i]);
     }
+
+    /* copies only the children whose jacobian values changed */
+    vstack_matrices_fill_values(hnode->child_jac, hnode->n_args, node->jacobian);
 }
 
 static void wsum_hess_init_impl(expr *node)
@@ -172,6 +141,8 @@ static void free_type_data(expr *node)
 
     sp_free(hnode->child_hess);
     hnode->child_hess = NULL;
+    sp_free(hnode->child_jac);
+    hnode->child_jac = NULL;
     sp_free(hnode->args);
     hnode->args = NULL;
 }

src/atoms/affine/left_matmul.c

@@ -102,16 +102,10 @@ static void free_type_data(expr *node)
     left_matmul_expr *lnode = (left_matmul_expr *) node;
     free_matrix(lnode->A);
     free_matrix(lnode->AT);
-    free_CSC_matrix(lnode->Jchild_CSC);
-    free_CSC_matrix(lnode->J_CSC);
-    sp_free(lnode->csc_to_csr_work);
     free_expr(lnode->param_source);
 
     lnode->A = NULL;
     lnode->AT = NULL;
-    lnode->Jchild_CSC = NULL;
-    lnode->J_CSC = NULL;
-    lnode->csc_to_csr_work = NULL;
     lnode->param_source = NULL;
 }
 
@@ -270,22 +264,14 @@ expr *new_left_matmul(expr *param_node, expr *u, const CSR_matrix *A)
     left_matmul_expr *lnode =
         (left_matmul_expr *) sp_calloc(1, sizeof(left_matmul_expr));
     expr *node = &lnode->base;
-    /* Sparse A: a single block goes through the dispatchers; the kron form
-       keeps the CSC-mirror path. */
+    /* Sparse A: a single block goes through BA_matrices_*, the kron form
+       through BA_sparse_kron_matrices_*. */
     init_expr(node, d1, d2, u->n_vars, forward, jacobian_init_sparse,
               eval_jacobian_sparse, is_affine, wsum_hess_init_impl,
               eval_wsum_hess_impl, free_type_data);
     node->left = u;
     expr_retain(u);
 
-    /* Workspace of the kron form only. iwork is used for converting J_child
-       csr to csc (requiring size node->n_vars). csc_to_csr_work is used for
-       converting J_CSC to CSR_matrix (requiring node->size) */
-    if (n_blocks > 1)
-    {
-        node->work->iwork = (int *) sp_malloc(node->n_vars * sizeof(int));
-        lnode->csc_to_csr_work = (int *) sp_malloc(node->size * sizeof(int));
-    }
     lnode->n_blocks = n_blocks;
 
     /* store A and AT. new_sparse_matrix takes ownership, so clone first. */

src/problem.c

@@ -18,6 +18,7 @@
 #include "problem.h"
 #include "sparsediff_linalg/matrix.h"
 #include "sparsediff_linalg/matrix_sum.h"
+#include "sparsediff_linalg/matrix_vstack.h"
 #include "sparsediff_linalg/tracked_alloc.h"
 #include "subexpr.h"
 #include <assert.h>
@@ -92,51 +93,35 @@ void problem_init_jacobian(problem *prob)
     jacobian_init(prob->objective);
     if (prob->n_constraints > 0)
     {
-        prob->constraint_jac_seen =
-            (uint64_t *) sp_malloc(prob->n_constraints * sizeof(uint64_t));
+        prob->jac_ops =
+            (matrix **) sp_malloc(prob->n_constraints * sizeof(matrix *));
     }
-    int nnz = 0;
     for (int i = 0; i < prob->n_constraints; i++)
     {
         expr *c = prob->constraints[i];
         jacobian_init(c);
-        /* deliberately stale so the first problem_jacobian call copies */
-        prob->constraint_jac_seen[i] = c->jacobian->values_version - 1;
-        CSR_matrix *Jc = c->jacobian->to_csr(c->jacobian);
-        nnz += Jc->nnz;
+        prob->jac_ops[i] = c->jacobian;
 
         if (c->is_affine(c))
         {
-            prob->stats.nnz_affine += Jc->nnz;
+            prob->stats.nnz_affine += c->jacobian->nnz;
         }
         else
         {
-            prob->stats.nnz_nonlinear += Jc->nnz;
+            prob->stats.nnz_nonlinear += c->jacobian->nnz;
         }
     }
 
-    prob->jacobian = new_CSR_matrix(prob->total_constraint_size, prob->n_vars, nnz);
-
-    /* set sparsity pattern of jacobian */
-    CSR_matrix *H = prob->jacobian;
-    H->p[0] = 0;
-    int row_offset = 0;
-    int nnz_offset = 0;
-    for (int i = 0; i < prob->n_constraints; i++)
+    /* the constraint jacobians stacked row-wise; vstack needs an operand */
+    if (prob->n_constraints > 0)
     {
-        expr *c = prob->constraints[i];
-        CSR_matrix *Jc = c->jacobian->to_csr(c->jacobian);
-
-        for (int r = 1; r <= Jc->m; r++)
-        {
-            H->p[row_offset + r] = nnz_offset + Jc->p[r];
-        }
-
-        memcpy(H->i + nnz_offset, Jc->i, Jc->nnz * sizeof(int));
-        row_offset += Jc->m;
-        nnz_offset += Jc->nnz;
+        prob->jac_stack = vstack_matrices_alloc(prob->jac_ops, prob->n_constraints);
+    }
+    else
+    {
+        prob->jac_stack = new_sparse_matrix_alloc(0, prob->n_vars, 0);
     }
-    assert(nnz_offset == nnz);
+    prob->jacobian = prob->jac_stack->to_csr(prob->jac_stack);
 
     clock_gettime(CLOCK_MONOTONIC, &timer.end);
     prob->stats.time_init_derivatives += GET_ELAPSED_SECONDS(timer);
@@ -284,11 +269,11 @@ void free_problem(problem *prob)
     /* Free allocated arrays */
     sp_free(prob->constraint_values);
     sp_free(prob->gradient_values);
-    free_CSR_matrix(prob->jacobian);
+    free_matrix(prob->jac_stack);
+    sp_free(prob->jac_ops);
     free_matrix(prob->hess_sum);
     free_COO_matrix(prob->jacobian_coo);
     free_COO_matrix(prob->lagrange_hessian_coo);
-    sp_free(prob->constraint_jac_seen);
 
     /* Release expression references (decrements refcount) */
     free_expr(prob->objective);
@@ -415,31 +400,19 @@ void problem_jacobian(problem *prob)
 {
     Timer timer;
     clock_gettime(CLOCK_MONOTONIC, &timer.start);
-    CSR_matrix *J = prob->jacobian;
-    int nnz_offset = 0;
-
     for (int i = 0; i < prob->n_constraints; i++)
     {
-        expr *c = prob->constraints[i];
-        eval_jacobian(c);
-
-        /* copy only when the constraint's jacobian values actually changed
-           (an affine constraint's eval is a no-op after its first call and
-           leaves the version untouched until the next parameter update).
-           Copy through the CSR view: J's rows were laid out from to_csr in
-           problem_init_jacobian, and a stacked_pd's own value buffer is
-           block-major, which is not row order when its blocks interleave
-           rows. For sparse and permuted_dense the view aliases the buffer. */
-        if (prob->constraint_jac_seen[i] != c->jacobian->values_version)
-        {
-            const CSR_matrix *Jc = c->jacobian->to_csr(c->jacobian);
-            memcpy(J->x + nnz_offset, Jc->x, Jc->nnz * sizeof(double));
-            prob->constraint_jac_seen[i] = c->jacobian->values_version;
-        }
-        nnz_offset += c->jacobian->nnz;
+        eval_jacobian(prob->constraints[i]);
     }
 
-    assert(nnz_offset == J->nnz);
+    /* copies only the constraints whose jacobian values changed (an affine
+       constraint's eval is a no-op after its first call and leaves the version
+       untouched until the next parameter update) */
+    if (prob->n_constraints > 0)
+    {
+        vstack_matrices_fill_values(prob->jac_ops, prob->n_constraints,
+                                    prob->jac_stack);
+    }
     clock_gettime(CLOCK_MONOTONIC, &timer.end);
     prob->stats.time_eval_jacobian += GET_ELAPSED_SECONDS(timer);
 }

tests/all_tests.c

@@ -233,6 +233,7 @@ int main(void)
     mu_run_test(test_jacobian_sum_spd_child_axis_1, tests_run);
     mu_run_test(test_jacobian_hstack_vectors, tests_run);
     mu_run_test(test_jacobian_hstack_matrix, tests_run);
+    mu_run_test(test_jacobian_hstack_stacked_pd_child, tests_run);
     mu_run_test(test_jacobian_vstack_vectors, tests_run);
     mu_run_test(test_jacobian_vstack_matrix, tests_run);
     mu_run_test(test_index_forward_simple, tests_run);

tests/jacobian_tests/affine/test_hstack.h

@@ -6,6 +6,7 @@
 #include "atoms/elementwise_restricted_dom.h"
 #include "expr.h"
 #include "minunit.h"
+#include "numerical_diff.h"
 #include "test_helpers.h"
 
 const char *test_jacobian_hstack_vectors(void)
@@ -93,3 +94,29 @@ const char *test_jacobian_hstack_matrix(void)
     free_expr(stack);
     return 0;
 }
+
+/* hstack(A @ exp(X), sin(X)) for a 2x2 variable X: the left_matmul child's
+   Jacobian is kron(I_2, A) diag(exp(x)), a stacked_pd, so the stack reads it
+   blockwise rather than through a CSR view. Checked at two points so the
+   second evaluation goes through the refill. */
+const char *test_jacobian_hstack_stacked_pd_child(void)
+{
+    double A[4] = {1.0, 2.0, -0.5, 3.0};
+    double u1[4] = {0.1, 0.2, 0.3, 0.4};
+    double u2[4] = {-0.7, 0.5, 1.1, -0.2};
+    expr *X = new_variable(2, 2, 0, 4);
+    expr *f = new_left_matmul_dense(NULL, new_exp(X), 2, 2, A);
+    expr *g = new_sin(X);
+    expr *args[2] = {f, g};
+    expr *stack = new_hstack(args, 2, 4);
+
+    mu_assert("check_jacobian at u1 failed",
+              check_jacobian_num(stack, u1, NUMERICAL_DIFF_DEFAULT_H));
+    mu_assert("child jacobian is stacked_pd",
+              f->jacobian->kind == MATRIX_STACKED_PD);
+    mu_assert("check_jacobian at u2 failed",
+              check_jacobian_num(stack, u2, NUMERICAL_DIFF_DEFAULT_H));
+
+    free_expr(stack);
+    return 0;
+}