From d519febc61952dac564904ac864fac9e21135318 Mon Sep 17 00:00:00 2001 From: Daniel Date: Sat, 17 Jan 2026 05:20:48 -0800 Subject: [PATCH 1/2] affine shortcut for jacobian and hessian evaluation --- include/problem.h | 4 ++++ src/bivariate/const_scalar_mult.c | 8 +++++++- src/bivariate/const_vector_mult.c | 8 +++++++- src/bivariate/multiply.c | 8 +++++++- src/bivariate/quad_over_lin.c | 7 +++++++ src/bivariate/rel_entr.c | 7 +++++++ src/bivariate/rel_entr_scalar_vector.c | 7 +++++++ src/bivariate/rel_entr_vector_scalar.c | 9 +++++++++ src/other/quad_form.c | 10 ++++++++-- src/problem.c | 18 ++++++++++++++++++ 10 files changed, 81 insertions(+), 5 deletions(-) diff --git a/include/problem.h b/include/problem.h index 63bd408..31f5e84 100644 --- a/include/problem.h +++ b/include/problem.h @@ -33,6 +33,10 @@ typedef struct problem int *hess_idx_map; /* Maps all wsum_hess nnz to lagrange_hessian (obj + constraints) */ + /* for the affine shortcut we keep track of the first time the jacobian and + * hessian are called */ + bool jacobian_called; + /* Statistics for performance measurement */ stats stats; } problem; diff --git a/src/bivariate/const_scalar_mult.c b/src/bivariate/const_scalar_mult.c index dc08ec6..d7f6fb4 100644 --- a/src/bivariate/const_scalar_mult.c +++ b/src/bivariate/const_scalar_mult.c @@ -75,6 +75,12 @@ static void eval_wsum_hess(expr *node, const double *w) } } +static bool is_affine(const expr *node) +{ + /* Affine iff the child is affine */ + return node->left->is_affine(node->left); +} + expr *new_const_scalar_mult(double a, expr *child) { const_scalar_mult_expr *mult_node = @@ -82,7 +88,7 @@ expr *new_const_scalar_mult(double a, expr *child) expr *node = &mult_node->base; init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init, - eval_jacobian, NULL, NULL); + eval_jacobian, is_affine, NULL); node->wsum_hess_init = wsum_hess_init; node->eval_wsum_hess = eval_wsum_hess; node->left = child; diff --git a/src/bivariate/const_vector_mult.c b/src/bivariate/const_vector_mult.c index 4846a7c..efee54d 100644 --- a/src/bivariate/const_vector_mult.c +++ b/src/bivariate/const_vector_mult.c @@ -90,6 +90,12 @@ static void free_type_data(expr *node) free(vnode->a); } +static bool is_affine(const expr *node) +{ + /* Affine iff the child is affine */ + return node->left->is_affine(node->left); +} + expr *new_const_vector_mult(const double *a, expr *child) { const_vector_mult_expr *vnode = @@ -97,7 +103,7 @@ expr *new_const_vector_mult(const double *a, expr *child) expr *node = &vnode->base; init_expr(node, child->d1, child->d2, child->n_vars, forward, jacobian_init, - eval_jacobian, NULL, free_type_data); + eval_jacobian, is_affine, free_type_data); node->wsum_hess_init = wsum_hess_init; node->eval_wsum_hess = eval_wsum_hess; node->left = child; diff --git a/src/bivariate/multiply.c b/src/bivariate/multiply.c index f30e916..770f34d 100644 --- a/src/bivariate/multiply.c +++ b/src/bivariate/multiply.c @@ -184,6 +184,12 @@ static void free_type_data(expr *node) free_csr_matrix(((elementwise_mult_expr *) node)->CSR_work2); } +static bool is_affine(const expr *node) +{ + (void) node; + return false; +} + expr *new_elementwise_mult(expr *left, expr *right) { elementwise_mult_expr *mul_node = @@ -191,7 +197,7 @@ expr *new_elementwise_mult(expr *left, expr *right) expr *node = &mul_node->base; init_expr(node, left->d1, left->d2, left->n_vars, forward, jacobian_init, - eval_jacobian, NULL, free_type_data); + eval_jacobian, is_affine, free_type_data); node->wsum_hess_init = wsum_hess_init; node->eval_wsum_hess = eval_wsum_hess; node->left = left; diff --git a/src/bivariate/quad_over_lin.c b/src/bivariate/quad_over_lin.c index ce28685..1ceff37 100644 --- a/src/bivariate/quad_over_lin.c +++ b/src/bivariate/quad_over_lin.c @@ -292,6 +292,12 @@ static void eval_wsum_hess(expr *node, const double *w) } } +static bool is_affine(const expr *node) +{ + (void) node; + return false; +} + expr *new_quad_over_lin(expr *left, expr *right) { assert((right->d2 == 1 && right->d2 == 1)); /* right must be scalar*/ @@ -305,5 +311,6 @@ expr *new_quad_over_lin(expr *left, expr *right) node->eval_jacobian = eval_jacobian; node->wsum_hess_init = wsum_hess_init; node->eval_wsum_hess = eval_wsum_hess; + node->is_affine = is_affine; return node; } diff --git a/src/bivariate/rel_entr.c b/src/bivariate/rel_entr.c index f0496df..5f9f9e9 100644 --- a/src/bivariate/rel_entr.c +++ b/src/bivariate/rel_entr.c @@ -169,6 +169,12 @@ static void eval_wsum_hess_vector_args(expr *node, const double *w) } } +static bool is_affine(const expr *node) +{ + (void) node; + return false; +} + expr *new_rel_entr_vector_args(expr *left, expr *right) { expr *node = new_expr(left->d1, left->d2, left->n_vars); @@ -181,5 +187,6 @@ expr *new_rel_entr_vector_args(expr *left, expr *right) node->eval_jacobian = eval_jacobian_vector_args; node->wsum_hess_init = wsum_hess_init_vector_args; node->eval_wsum_hess = eval_wsum_hess_vector_args; + node->is_affine = is_affine; return node; } diff --git a/src/bivariate/rel_entr_scalar_vector.c b/src/bivariate/rel_entr_scalar_vector.c index 1d0758e..a1b4aed 100644 --- a/src/bivariate/rel_entr_scalar_vector.c +++ b/src/bivariate/rel_entr_scalar_vector.c @@ -191,6 +191,12 @@ static void eval_wsum_hess_scalar_vector(expr *node, const double *w) } } +static bool is_affine(const expr *node) +{ + (void) node; + return false; +} + expr *new_rel_entr_first_arg_scalar(expr *left, expr *right) { assert(left->d1 == 1 && left->d2 == 1); @@ -204,5 +210,6 @@ expr *new_rel_entr_first_arg_scalar(expr *left, expr *right) node->eval_jacobian = eval_jacobian_scalar_vector; node->wsum_hess_init = wsum_hess_init_scalar_vector; node->eval_wsum_hess = eval_wsum_hess_scalar_vector; + node->is_affine = is_affine; return node; } diff --git a/src/bivariate/rel_entr_vector_scalar.c b/src/bivariate/rel_entr_vector_scalar.c index 0b36ca6..b3b7fc1 100644 --- a/src/bivariate/rel_entr_vector_scalar.c +++ b/src/bivariate/rel_entr_vector_scalar.c @@ -192,6 +192,14 @@ static void eval_wsum_hess_vector_scalar(expr *node, const double *w) } } +static bool is_affine(const expr *node) +{ + (void) node; + return false; +} + +/* TODO: modify relative entropy atoms to use init_expr initialization? Also make + * init_expr take in the hessian functions?*/ expr *new_rel_entr_second_arg_scalar(expr *left, expr *right) { assert(right->d1 == 1 && right->d2 == 1); @@ -205,5 +213,6 @@ expr *new_rel_entr_second_arg_scalar(expr *left, expr *right) node->eval_jacobian = eval_jacobian_vector_scalar; node->wsum_hess_init = wsum_hess_init_vector_scalar; node->eval_wsum_hess = eval_wsum_hess_vector_scalar; + node->is_affine = is_affine; return node; } diff --git a/src/other/quad_form.c b/src/other/quad_form.c index 42fb989..2cfbf10 100644 --- a/src/other/quad_form.c +++ b/src/other/quad_form.c @@ -177,14 +177,20 @@ static void free_type_data(expr *node) qnode->Q = NULL; } +static bool is_affine(const expr *node) +{ + (void) node; + return false; +} + expr *new_quad_form(expr *left, CSR_Matrix *Q) { assert(left->d1 == 1 || left->d2 == 1); /* left must be a vector */ quad_form_expr *qnode = (quad_form_expr *) calloc(1, sizeof(quad_form_expr)); expr *node = &qnode->base; - init_expr(node, 1, 1, left->n_vars, forward, jacobian_init, eval_jacobian, NULL, - free_type_data); + init_expr(node, 1, 1, left->n_vars, forward, jacobian_init, eval_jacobian, + is_affine, free_type_data); node->left = left; expr_retain(left); node->wsum_hess_init = wsum_hess_init; diff --git a/src/problem.c b/src/problem.c index 0372494..e5b4764 100644 --- a/src/problem.c +++ b/src/problem.c @@ -17,6 +17,7 @@ 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; @@ -99,6 +100,7 @@ void problem_init_derivatives(problem *prob) } prob->lagrange_hessian = new_csr_matrix(prob->n_vars, prob->n_vars, nnz); + memset(prob->lagrange_hessian->x, 0, nnz * sizeof(double)); /* affine shortcut */ prob->hess_idx_map = (int *) malloc(nnz * sizeof(int)); int *iwork = (int *) malloc(MAX(nnz, prob->n_vars) * sizeof(int)); problem_lagrange_hess_fill_sparsity(prob, iwork); @@ -274,6 +276,7 @@ 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; @@ -281,6 +284,14 @@ void problem_jacobian(problem *prob) for (int i = 0; i < prob->n_constraints; i++) { expr *c = prob->constraints[i]; + + if (!first_call && c->is_affine(c)) + { + /* skip evaluation for affine constraints after first call */ + nnz_offset += c->jacobian->nnz; + continue; + } + c->eval_jacobian(c); memcpy(J->x + nnz_offset, c->jacobian->x, c->jacobian->nnz * sizeof(double)); nnz_offset += c->jacobian->nnz; @@ -289,6 +300,7 @@ void problem_jacobian(problem *prob) /* update actual nnz (may be less than allocated) */ J->nnz = nnz_offset; + prob->jacobian_called = true; clock_gettime(CLOCK_MONOTONIC, &timer.end); prob->stats.time_eval_jacobian += GET_ELAPSED_SECONDS(timer); } @@ -308,6 +320,12 @@ void problem_hessian(problem *prob, double obj_w, const double *w) expr **constrs = prob->constraints; for (int i = 0; i < prob->n_constraints; i++) { + if (constrs[i]->is_affine(constrs[i])) + { + /* skip evaluation for affine constraints */ + offset += constrs[i]->size; + continue; + } constrs[i]->eval_wsum_hess(constrs[i], w + offset); offset += constrs[i]->size; } From 6e6915e860d5a619082339aa1cc4fdf5667dbf92 Mon Sep 17 00:00:00 2001 From: Daniel Date: Sat, 17 Jan 2026 05:23:15 -0800 Subject: [PATCH 2/2] cleaned up todo a bit --- TODO.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/TODO.md b/TODO.md index 64dc697..8854918 100644 --- a/TODO.md +++ b/TODO.md @@ -1,8 +1,4 @@ -3. more tests for chain rule elementwise univariate hessian 4. in the refactor, add consts -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. -8. Shortcut hessians of affine stuff? 10. For performance reasons, is it useful to have a dense matmul with A and B as dense matrices? 11. right matmul, add broadcasting logic as in left matmul