From d1fbf68817698f5b507f67a167cc81c1d06de550 Mon Sep 17 00:00:00 2001 From: dance858 Date: Sun, 30 Aug 2026 21:06:16 -0700 Subject: [PATCH] Replace jacobian_called flag with per-constraint version checks problem_jacobian used a problem-global bool: first call evaluates all constraints, later calls skip affine ones entirely, and problem_update_params resets the flag wholesale. Replace it with a per-constraint seen counter against the constraint jacobian's values_version: every constraint is eval'd (affine ones return instantly via the impl-skip from #114), and a block is memcpy'd into the aggregated jacobian only when its version moved. Re-copying after a parameter update happens automatically because expr_set_needs_refresh re-arms the evals, which bump the versions - so the copy skip is per-constraint and needs no reset. The seen array is seeded deliberately stale in problem_init_jacobian so the first call always copies. New poison-based regression test proves the affine slice is not rewritten on re-eval and is restored after a parameter update. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016rQ3LPdi31kNLmTy9F2oEa --- include/problem.h | 7 +++-- src/problem.c | 31 ++++++++++--------- tests/all_tests.c | 1 + tests/problem/test_param_prob.h | 55 +++++++++++++++++++++++++++++++++ 4 files changed, 77 insertions(+), 17 deletions(-) diff --git a/include/problem.h b/include/problem.h index 530907cc..e77794d1 100644 --- a/include/problem.h +++ b/include/problem.h @@ -65,9 +65,10 @@ typedef struct problem COO_matrix *jacobian_coo; COO_matrix *lagrange_hessian_coo; /* lower triangular part stored in COO */ - /* for the affine shortcut we keep track of the first time the jacobian and - * hessian are called */ - bool jacobian_called; + /* 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; diff --git a/src/problem.c b/src/problem.c index 11b49f00..083f2092 100644 --- a/src/problem.c +++ b/src/problem.c @@ -43,7 +43,6 @@ problem *new_problem(expr *objective, expr **constraints, int n_constraints, prob->objective = objective; expr_retain(objective); prob->n_vars = objective->n_vars; - prob->jacobian_called = false; /* constraints array */ prob->total_constraint_size = 0; @@ -192,11 +191,18 @@ void problem_init_jacobian(problem *prob) // Jacobian structure // ------------------------------------------------------------------------------- jacobian_init(prob->objective); + if (prob->n_constraints > 0) + { + prob->constraint_jac_seen = + (uint64_t *) sp_malloc(prob->n_constraints * sizeof(uint64_t)); + } 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; @@ -384,6 +390,7 @@ void free_problem(problem *prob) free_COO_matrix(prob->jacobian_coo); free_COO_matrix(prob->lagrange_hessian_coo); sp_free(prob->hess_idx_map); + sp_free(prob->constraint_jac_seen); /* Release expression references (decrements refcount) */ free_expr(prob->objective); @@ -450,9 +457,6 @@ void problem_update_params(problem *prob, const double *theta) { expr_set_needs_refresh(prob->constraints[i]); } - - /* Force re-evaluation of affine Jacobians on next call */ - prob->jacobian_called = false; } double problem_objective_forward(problem *prob, const double *u) @@ -513,28 +517,27 @@ void problem_jacobian(problem *prob) { Timer timer; clock_gettime(CLOCK_MONOTONIC, &timer.start); - bool first_call = !prob->jacobian_called; - CSR_matrix *J = prob->jacobian; int nnz_offset = 0; for (int i = 0; i < prob->n_constraints; i++) { expr *c = prob->constraints[i]; - if (!first_call && c->is_affine(c)) + 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) */ + if (prob->constraint_jac_seen[i] != c->jacobian->values_version) { - /* skip evaluation for affine constraints after first call */ - nnz_offset += c->jacobian->nnz; - continue; + memcpy(J->x + nnz_offset, c->jacobian->x, + c->jacobian->nnz * sizeof(double)); + prob->constraint_jac_seen[i] = c->jacobian->values_version; } - - eval_jacobian(c); - memcpy(J->x + nnz_offset, c->jacobian->x, c->jacobian->nnz * sizeof(double)); nnz_offset += c->jacobian->nnz; } assert(nnz_offset == J->nnz); - prob->jacobian_called = true; 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 e1bb5572..21483121 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -556,6 +556,7 @@ int main(void) mu_run_test(test_param_right_matmul_rectangular, tests_run); mu_run_test(test_param_shared_left_matmul_problem, tests_run); mu_run_test(test_param_fixed_skip_in_update, tests_run); + mu_run_test(test_problem_jacobian_memcpy_skip, tests_run); mu_run_test(test_param_scalar_mult_problem_with_constant, tests_run); mu_run_test(test_param_convolve_problem, tests_run); diff --git a/tests/problem/test_param_prob.h b/tests/problem/test_param_prob.h index a638b31b..3bbf3f15 100644 --- a/tests/problem/test_param_prob.h +++ b/tests/problem/test_param_prob.h @@ -618,4 +618,59 @@ const char *test_param_convolve_problem(void) return 0; } +/* problem_jacobian copies a constraint's block into the aggregated jacobian + only when the constraint's values_version moved: an affine constraint's + slice is not rewritten after its first copy (proven by poisoning it), and + a parameter update re-arms the copy. */ +const char *test_problem_jacobian_memcpy_skip(void) +{ + int n = 2; + expr *x = new_variable(2, 1, 0, n); + + /* objective: sum(log(x)) */ + expr *objective = new_sum(new_log(x), -1); + + /* constraint 1 (affine, parameterized): a * x */ + double theta[1] = {3.0}; + expr *a_param = new_parameter(1, 1, 0, n, theta); + expr *scaled = new_scalar_mult(a_param, x); + + /* constraint 2 (non-affine): exp(x) */ + expr *exp_c = new_exp(x); + + expr *constraints[2] = {scaled, exp_c}; + problem *prob = new_problem(objective, constraints, 2, false); + expr *param_nodes[1] = {a_param}; + problem_register_params(prob, param_nodes, 1); + problem_init_derivatives(prob); + + double u[2] = {1.0, 2.0}; + problem_constraint_forward(prob, u); + problem_jacobian(prob); + + double expected1[4] = {3.0, 3.0, exp(1.0), exp(2.0)}; + mu_assert("first jacobian wrong", + cmp_double_array(prob->jacobian->x, expected1, 4)); + + /* poison the affine slice; a re-eval must not rewrite it */ + prob->jacobian->x[0] = 42.0; + prob->jacobian->x[1] = 42.0; + problem_jacobian(prob); + mu_assert("affine slice must not be recopied", prob->jacobian->x[0] == 42.0); + mu_assert("affine slice must not be recopied", prob->jacobian->x[1] == 42.0); + mu_assert("non-affine slice must be recopied", prob->jacobian->x[2] == exp(1.0)); + + /* a parameter update re-arms the copy */ + theta[0] = 5.0; + problem_update_params(prob, theta); + problem_constraint_forward(prob, u); + problem_jacobian(prob); + double expected2[4] = {5.0, 5.0, exp(1.0), exp(2.0)}; + mu_assert("jacobian after update wrong", + cmp_double_array(prob->jacobian->x, expected2, 4)); + + free_problem(prob); + return 0; +} + #endif /* TEST_PARAM_PROB_H */