From 575ff784066155f042e80d2d93662aeae2c4ff77 Mon Sep 17 00:00:00 2001 From: dance858 Date: Tue, 13 Oct 2026 14:48:37 +0200 Subject: [PATCH] Stack constraint and hstack Jacobians with vstack_matrices_* - problem.c: the global constraint Jacobian is a vstack of the constraint Jacobians (jac_stack over jac_ops); prob->jacobian becomes its CSR view. The fill copies only constraints whose values_version changed, which replaces constraint_jac_seen and both hand-written to_csr loops. With no constraints, jac_stack is an empty 0 x n_vars sparse matrix. - hstack.c: the Jacobian is a vstack of the children's Jacobians (child_jac), so unchanged children are no longer recopied on each eval. - Both read each operand's own storage: a stacked_pd Jacobian is copied once per change instead of twice through a to_csr view, and no CSR cache is created for it. - left_matmul: drop Jchild_CSC / J_CSC / csc_to_csr_work and the kron-only workspace, dead since the kron branch moved to BA_sparse_kron_matrices_*. subexpr.h no longer needs CSC_matrix.h. - New test test_jacobian_hstack_stacked_pd_child: hstack with a stacked_pd child, checked against finite differences at two points (refill). Co-Authored-By: Claude Opus 5.5 --- include/problem.h | 9 +-- include/subexpr.h | 5 +- src/atoms/affine/hstack.c | 51 ++++----------- src/atoms/affine/left_matmul.c | 18 +----- src/problem.c | 77 ++++++++--------------- tests/all_tests.c | 1 + tests/jacobian_tests/affine/test_hstack.h | 27 ++++++++ 7 files changed, 70 insertions(+), 118 deletions(-) diff --git a/include/problem.h b/include/problem.h index b9f28777..b7645860 100644 --- a/include/problem.h +++ b/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; diff --git a/include/subexpr.h b/include/subexpr.h index cdd29308..fedf180d 100644 --- a/include/subexpr.h +++ b/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; diff --git a/src/atoms/affine/hstack.c b/src/atoms/affine/hstack.c index be45c4f9..c93a6a61 100644 --- a/src/atoms/affine/hstack.c +++ b/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 @@ -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; } diff --git a/src/atoms/affine/left_matmul.c b/src/atoms/affine/left_matmul.c index 8640e153..322e3213 100644 --- a/src/atoms/affine/left_matmul.c +++ b/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. */ diff --git a/src/problem.c b/src/problem.c index fe5b859b..66a9278f 100644 --- a/src/problem.c +++ b/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 @@ -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); } diff --git a/tests/all_tests.c b/tests/all_tests.c index 768d8b56..55bd5879 100644 --- a/tests/all_tests.c +++ b/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); diff --git a/tests/jacobian_tests/affine/test_hstack.h b/tests/jacobian_tests/affine/test_hstack.h index bd61df82..98551705 100644 --- a/tests/jacobian_tests/affine/test_hstack.h +++ b/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; +}