diff --git a/include/expr.h b/include/expr.h index 85e91c8..a6f0339 100644 --- a/include/expr.h +++ b/include/expr.h @@ -40,7 +40,7 @@ typedef void (*local_jacobian_fn)(struct expr *node, double *out); typedef void (*local_wsum_hess_fn)(struct expr *node, double *out, const double *w); typedef bool (*is_affine_fn)(const struct expr *node); typedef void (*free_type_data_fn)(struct expr *node); -typedef void (*set_needs_refresh_children_fn)(struct expr *node); +typedef bool (*set_needs_refresh_children_fn)(struct expr *node); /* Workspace for derivative computation */ typedef struct @@ -97,13 +97,27 @@ typedef struct expr local_jacobian_fn local_jacobian; /* used by elementwise univariate atoms*/ local_wsum_hess_fn local_wsum_hess; /* used by elementwise univariate atoms*/ free_type_data_fn free_type_data; /* Cleanup for type-specific fields */ - /* Recursion hook for expr_set_needs_refresh: atoms holding children - outside left/right (hstack's args[]) set this so the parameter-refresh - walk reaches them. NULL for binary/unary atoms. */ + /* Recursion hook for expr_set_needs_refresh, which on its own only + descends left/right. Nodes that hold parameter-bearing subtrees + elsewhere install this so the walk reaches them: the parameter leaf + (reports itself), hstack (args[]) and the coefficient atoms + scalar_mult, vector_mult, kron, convolve, left_matmul and quad_form + (param_source). NULL everywhere else. + + Contract: the hook MUST return whether the nodes it reaches contain + an updatable parameter (param_id >= 0). The walk ORs the result into + has_params; a hook that returns false lets the subtree be memoized + parameter-free and pruned from the second update on. */ set_needs_refresh_children_fn set_needs_refresh_children; Expr_Work *work; /* derivative workspace */ - /* Set to true on all nodes by problem_update_params() via - expr_set_needs_refresh(). Atoms that cache parameter data + /* Does this subtree contain an updatable parameter? Starts true so + nothing is pruned before the first walk; expr_set_needs_refresh + refines and memoizes it from the left/right walk and the + set_needs_refresh_children hook. */ + bool has_params; + /* Set to true by problem_update_params() via expr_set_needs_refresh() + on every node whose subtree contains an updatable parameter; + parameter-free subtrees are skipped. Atoms that cache parameter data (e.g. left_matmul_dense) check this flag before their forward pass: if true, they refresh their cached matrices from param_source->value and clear the flag to false. */ @@ -139,8 +153,11 @@ void expr_refresh_jacobian_csc(expr *node); * Must be called after jacobian_init. */ void jacobian_csc_init(expr *node); -/* Recursively set needs_parameter_refresh on node and all children */ -void expr_set_needs_refresh(expr *node); +/* Mark the subtree dirty and report whether it contains an updatable + * parameter. A subtree known to be parameter-free is skipped entirely: its + * values and derivatives cannot have changed, so re-arming it would only + * force a recompute of a result we already hold. */ +bool expr_set_needs_refresh(expr *node); /* Reference counting helpers */ void expr_retain(expr *node); diff --git a/src/atoms/affine/convolve.c b/src/atoms/affine/convolve.c index c9ed047..39c69bd 100644 --- a/src/atoms/affine/convolve.c +++ b/src/atoms/affine/convolve.c @@ -39,9 +39,6 @@ static void forward(expr *node, const double *u) if (cnode->base.needs_parameter_refresh) { - /* Composite sources hold gated nodes of their own (promote, nested - mults): mark the whole side subtree before re-evaluating it. */ - expr_set_needs_refresh(cnode->param_source); cnode->param_source->forward(cnode->param_source, NULL); /* refresh the convolution matrix values if it exists (necessary to check for null in case someone calls forward before initializing the jacobian, @@ -143,6 +140,13 @@ static bool is_affine(const expr *node) return node->left->is_affine(node->left); } +/* param_source lives outside left/right, so the refresh walk reaches it + here -- and reports whether it actually holds an updatable parameter. */ +static bool set_needs_refresh_param_source(expr *node) +{ + return expr_set_needs_refresh(((convolve_expr *) node)->param_source); +} + static void free_type_data(expr *node) { convolve_expr *cnode = (convolve_expr *) node; @@ -197,6 +201,7 @@ expr *new_convolve(expr *param_node, expr *child) /* Ensure first forward() pulls current param values through any broadcast/promote wrappers and reflects them in T (once T is built). */ cnode->base.needs_parameter_refresh = true; + cnode->base.set_needs_refresh_children = set_needs_refresh_param_source; return node; } diff --git a/src/atoms/affine/hstack.c b/src/atoms/affine/hstack.c index 71554cf..5777f67 100644 --- a/src/atoms/affine/hstack.c +++ b/src/atoms/affine/hstack.c @@ -172,13 +172,15 @@ static bool is_affine(const expr *node) /* Children live in args[], not left/right, so the parameter-refresh walk needs this hook to reach them. */ -static void set_needs_refresh_children(expr *node) +static bool set_needs_refresh_children(expr *node) { hstack_expr *hnode = (hstack_expr *) node; + bool child_has_params = false; for (int i = 0; i < hnode->n_args; i++) { - expr_set_needs_refresh(hnode->args[i]); + child_has_params |= expr_set_needs_refresh(hnode->args[i]); } + return child_has_params; } static void free_type_data(expr *node) diff --git a/src/atoms/affine/kron.c b/src/atoms/affine/kron.c index b0dc5e1..6fd149a 100644 --- a/src/atoms/affine/kron.c +++ b/src/atoms/affine/kron.c @@ -44,9 +44,6 @@ static void refresh_param_values(kron_expr *knode) return; } - /* Composite sources hold gated nodes of their own (promote, nested - mults): mark the whole side subtree before re-evaluating it. */ - expr_set_needs_refresh(knode->param_source); knode->param_source->forward(knode->param_source, NULL); knode->base.needs_parameter_refresh = false; } @@ -176,6 +173,13 @@ static bool is_affine(const expr *node) return node->left->is_affine(node->left); } +/* param_source lives outside left/right, so the refresh walk reaches it + here -- and reports whether it actually holds an updatable parameter. */ +static bool set_needs_refresh_param_source(expr *node) +{ + return expr_set_needs_refresh(((kron_expr *) node)->param_source); +} + static void free_type_data(expr *node) { kron_expr *knode = (kron_expr *) node; @@ -214,6 +218,7 @@ static kron_expr *new_kron_common(expr *param_node, expr *child, int p, int q, i } knode->base.needs_parameter_refresh = true; + knode->base.set_needs_refresh_children = set_needs_refresh_param_source; return knode; } diff --git a/src/atoms/affine/left_matmul.c b/src/atoms/affine/left_matmul.c index 2cc4171..b3b9f33 100644 --- a/src/atoms/affine/left_matmul.c +++ b/src/atoms/affine/left_matmul.c @@ -70,9 +70,6 @@ static void forward(expr *node, const double *u) /* call forward on param_source if it exists and needs refresh */ if (lnode->param_source != NULL && lnode->base.needs_parameter_refresh) { - /* Composite sources hold gated nodes of their own (promote, nested - mults): mark the whole side subtree before re-evaluating it. */ - expr_set_needs_refresh(lnode->param_source); lnode->param_source->forward(lnode->param_source, NULL); } @@ -94,6 +91,13 @@ static bool is_affine(const expr *node) return node->left->is_affine(node->left); } +/* param_source lives outside left/right, so the refresh walk reaches it + here -- and reports whether it actually holds an updatable parameter. */ +static bool set_needs_refresh_param_source(expr *node) +{ + return expr_set_needs_refresh(((left_matmul_expr *) node)->param_source); +} + static void free_type_data(expr *node) { left_matmul_expr *lnode = (left_matmul_expr *) node; @@ -336,6 +340,7 @@ expr *new_left_matmul_dense(expr *param_node, expr *u, int m, int n, lnode->A = new_permuted_dense_full(m, n, NULL); lnode->AT = new_permuted_dense_full(n, m, NULL); node->needs_parameter_refresh = true; + node->set_needs_refresh_children = set_needs_refresh_param_source; } /* constant matrix case */ else diff --git a/src/atoms/affine/parameter.c b/src/atoms/affine/parameter.c index 39b010e..ccca48e 100644 --- a/src/atoms/affine/parameter.c +++ b/src/atoms/affine/parameter.c @@ -53,6 +53,12 @@ static void eval_wsum_hess_impl(expr *node, const double *w) (void) w; } +/* A leaf has no children, so it reports its own parameter-ness here. */ +static bool set_needs_refresh_children(expr *node) +{ + return ((parameter_expr *) node)->param_id >= 0; +} + static bool is_affine(const expr *node) { (void) node; @@ -68,6 +74,7 @@ expr *new_parameter(int d1, int d2, int param_id, int n_vars, const double *valu // TODO we should assert that the values array has the correct size. pnode->param_id = param_id; + node->set_needs_refresh_children = set_needs_refresh_children; if (values == NULL) { diff --git a/src/atoms/affine/scalar_mult.c b/src/atoms/affine/scalar_mult.c index e2886bb..f9b73f5 100644 --- a/src/atoms/affine/scalar_mult.c +++ b/src/atoms/affine/scalar_mult.c @@ -35,9 +35,6 @@ static void forward(expr *node, const double *u) its values) */ if (snode->base.needs_parameter_refresh) { - /* Composite sources hold gated nodes of their own (promote, nested - mults): mark the whole side subtree before re-evaluating it. */ - expr_set_needs_refresh(snode->param_source); snode->param_source->forward(snode->param_source, NULL); snode->base.needs_parameter_refresh = false; } @@ -108,6 +105,13 @@ static bool is_affine(const expr *node) return node->left->is_affine(node->left); } +/* param_source lives outside left/right, so the refresh walk reaches it + here -- and reports whether it actually holds an updatable parameter. */ +static bool set_needs_refresh_param_source(expr *node) +{ + return expr_set_needs_refresh(((scalar_mult_expr *) node)->param_source); +} + static void free_type_data(expr *node) { scalar_mult_expr *snode = (scalar_mult_expr *) node; @@ -135,6 +139,7 @@ expr *new_scalar_mult(expr *param_node, expr *child) /* special case for handling broadcasting of constants correctly */ mult_node->base.needs_parameter_refresh = true; + mult_node->base.set_needs_refresh_children = set_needs_refresh_param_source; return node; } diff --git a/src/atoms/affine/vector_mult.c b/src/atoms/affine/vector_mult.c index c463574..282bd84 100644 --- a/src/atoms/affine/vector_mult.c +++ b/src/atoms/affine/vector_mult.c @@ -35,9 +35,6 @@ static void forward(expr *node, const double *u) its values) */ if (vnode->base.needs_parameter_refresh) { - /* Composite sources hold gated nodes of their own (promote, nested - mults): mark the whole side subtree before re-evaluating it. */ - expr_set_needs_refresh(vnode->param_source); vnode->param_source->forward(vnode->param_source, NULL); vnode->base.needs_parameter_refresh = false; } @@ -109,6 +106,13 @@ static void eval_wsum_hess_impl(expr *node, const double *w) node->wsum_hess->nnz * sizeof(double)); } +/* param_source lives outside left/right, so the refresh walk reaches it + here -- and reports whether it actually holds an updatable parameter. */ +static bool set_needs_refresh_param_source(expr *node) +{ + return expr_set_needs_refresh(((vector_mult_expr *) node)->param_source); +} + static void free_type_data(expr *node) { vector_mult_expr *vnode = (vector_mult_expr *) node; @@ -141,6 +145,7 @@ expr *new_vector_mult(expr *param_node, expr *child) /* special case for handling broadcasting of constants correctly */ vnode->base.needs_parameter_refresh = true; + vnode->base.set_needs_refresh_children = set_needs_refresh_param_source; return node; } diff --git a/src/atoms/other/quad_form.c b/src/atoms/other/quad_form.c index 0386f94..54db5bc 100644 --- a/src/atoms/other/quad_form.c +++ b/src/atoms/other/quad_form.c @@ -56,9 +56,6 @@ static void forward(expr *node, const double *u) /* refresh Q from the parameter if needed (no-op on the constant/sparse path) */ if (qnode->param_source != NULL && node->needs_parameter_refresh) { - /* Composite sources hold gated nodes of their own (promote, nested - mults): mark the whole side subtree before re-evaluating it. */ - expr_set_needs_refresh(qnode->param_source); qnode->param_source->forward(qnode->param_source, NULL); } refresh_param_values_qf(qnode); @@ -340,6 +337,13 @@ static void eval_wsum_hess_dense(expr *node, const double *w) } } +/* param_source lives outside left/right, so the refresh walk reaches it + here -- and reports whether it actually holds an updatable parameter. */ +static bool set_needs_refresh_param_source(expr *node) +{ + return expr_set_needs_refresh(((quad_form_expr *) node)->param_source); +} + static void free_type_data(expr *node) { quad_form_expr *qnode = (quad_form_expr *) node; @@ -421,6 +425,7 @@ expr *new_quad_form_dense(expr *child, int n, const double *P_data, /* Q is filled from the parameter on the first forward pass. */ qnode->Q = new_permuted_dense_full(n, n, NULL); node->needs_parameter_refresh = true; + node->set_needs_refresh_children = set_needs_refresh_param_source; } else { diff --git a/src/expr.c b/src/expr.c index 0990727..c61055b 100644 --- a/src/expr.c +++ b/src/expr.c @@ -42,6 +42,7 @@ void init_expr(expr *node, int d1, int d2, int n_vars, forward_fn forward, node->eval_wsum_hess_impl = eval_wsum_hess; node->free_type_data = free_type_data; node->work = (Expr_Work *) sp_calloc(1, sizeof(Expr_Work)); + node->has_params = true; /* assume dirty until the first walk refines it */ } void jacobian_csc_init(expr *node) @@ -151,20 +152,32 @@ void eval_wsum_hess(expr *node, const double *w) matrix_values_changed(node->wsum_hess); } -void expr_set_needs_refresh(expr *node) +bool expr_set_needs_refresh(expr *node) { - if (node == NULL) return; + if (node == NULL) return false; + + /* Known parameter-free: nothing below can have changed, so leave the + node's jacobian_evaluated latch set and skip the whole subtree. */ + if (!node->has_params) return false; + node->needs_parameter_refresh = true; /* Re-arm the eval_jacobian wrapper's values_version bump: the next eval after a parameter update may change even an affine node's values. */ node->work->jacobian_evaluated = false; - expr_set_needs_refresh(node->left); - expr_set_needs_refresh(node->right); + + bool child_has_params = expr_set_needs_refresh(node->left); + child_has_params |= expr_set_needs_refresh(node->right); if (node->set_needs_refresh_children != NULL) { - node->set_needs_refresh_children(node); + child_has_params |= node->set_needs_refresh_children(node); } + + /* The hook reports nodes the left/right walk cannot see: a parameter + leaf reports itself, hstack its args[], the coefficient atoms their + param_source. So this assignment is the whole answer. */ + node->has_params = child_has_params; + return node->has_params; } void expr_retain(expr *node) diff --git a/tests/all_tests.c b/tests/all_tests.c index 98eb767..2377405 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -116,6 +116,7 @@ #ifdef PROFILE_ONLY #include "profiling/profile_BTA_pd_csr_vs_csc.h" #include "profiling/profile_hessian_exp_AX.h" +#include "profiling/profile_lasso.h" #include "profiling/profile_left_matmul.h" #include "profiling/profile_log_reg.h" #include "profiling/profile_memory.h" @@ -217,6 +218,10 @@ int main(void) mu_run_test(test_values_version_csc_mirror_dedup, tests_run); mu_run_test(test_values_version_stacked_pd_to_csr, tests_run); mu_run_test(test_values_version_param_under_hstack, tests_run); + mu_run_test(test_refresh_prunes_param_free, tests_run); + mu_run_test(test_refresh_rearms_param_dependent, tests_run); + mu_run_test(test_refresh_prunes_fixed_constant, tests_run); + mu_run_test(test_refresh_rearms_updatable_constant, tests_run); mu_run_test(test_values_version_spd_hess_terms, tests_run); /* commented out - see test_quad_form.h */ // mu_run_test(test_quad_form2, tests_run); @@ -612,6 +617,7 @@ int main(void) #ifdef PROFILE_ONLY printf("\n--- Profiling Tests ---\n"); + mu_run_test(profile_lasso, tests_run); mu_run_test(profile_left_matmul, tests_run); mu_run_test(profile_log_reg, tests_run); mu_run_test(profile_trimmed_log_reg, tests_run); diff --git a/tests/jacobian_tests/test_values_version.h b/tests/jacobian_tests/test_values_version.h index 3b352b2..9857e08 100644 --- a/tests/jacobian_tests/test_values_version.h +++ b/tests/jacobian_tests/test_values_version.h @@ -231,6 +231,27 @@ const char *test_values_version_param_under_hstack(void) mu_assert("hstack jacobian must pick up the new parameter value", cmp_double_array(h->jacobian->x, expected, nnz1)); + /* Second update. The hook must REPORT its args' parameter-dependence, not + merely visit them: a hook that walks args[] but returns false marks + everything correctly on the first walk and then memoizes the hstack + parameter-free, pruning it from the second update on. One walk cannot + see that, so this round is what pins the return value. + p = 3.0 -> 7.0 */ + double p2 = 7.0; + memcpy(p->value, &p2, sizeof(double)); + mu_assert("hstack must report itself parameter-dependent", + expr_set_needs_refresh(h) == true); + + h->forward(h, u); + eval_jacobian(h); + + for (int k = 0; k < nnz1; k++) + { + expected[k] *= p2 / p1; + } + mu_assert("hstack jacobian must track the parameter on the second update", + cmp_double_array(h->jacobian->x, expected, nnz1)); + free(expected); free_expr(h); return 0; @@ -272,3 +293,183 @@ const char *test_values_version_spd_hess_terms(void) free_expr(outer); return 0; } + +/* Parameter-free subtree: expr_set_needs_refresh resolves the dependency on + * its first walk and prunes every walk after that, so an affine node that no + * parameter can reach keeps its bump-skip armed for the life of the problem. + * Proven by poisoning the values and observing they survive a refresh+eval. */ +const char *test_refresh_prunes_param_free(void) +{ + double u[3] = {0.1, 0.2, 0.3}; + expr *x = new_variable(3, 1, 0, 3); + expr *m = new_neg(x); + + jacobian_init(m); + m->forward(m, u); + eval_jacobian(m); + + /* First walk still re-arms (it is the walk that discovers the subtree is + parameter-free), so this eval runs and restores the true values. */ + mu_assert("must start assumed-dirty", m->has_params == true); + mu_assert("first walk must report parameter-free", + expr_set_needs_refresh(m) == false); + mu_assert("first walk must memoize parameter-free", m->has_params == false); + eval_jacobian(m); + + uint64_t v1 = m->jacobian->values_version; + int nnz = m->jacobian->nnz; + mu_assert("neg jacobian must have entries", nnz == 3); + for (int ii = 0; ii < nnz; ii++) + { + m->jacobian->x[ii] = 42.0; /* poison */ + } + + /* Every later walk prunes: the latch is never cleared, so the impl does + not run and the poison survives. */ + mu_assert("later walk must still report parameter-free", + expr_set_needs_refresh(m) == false); + mu_assert("pruned node must keep its eval latch armed", + m->work->jacobian_evaluated == true); + eval_jacobian(m); + mu_assert("pruned re-eval must not bump", m->jacobian->values_version == v1); + for (int ii = 0; ii < nnz; ii++) + { + mu_assert("impl must not have run (poison must survive)", + m->jacobian->x[ii] == 42.0); + } + + free_expr(m); + return 0; +} + +/* A subtree a parameter can reach is re-armed by every walk, not just the + * first one -- the prune must not swallow real invalidation. */ +const char *test_refresh_rearms_param_dependent(void) +{ + int n_vars = 3; + double u[3] = {1.0, 2.0, 3.0}; + double p0[1] = {2.0}; + + expr *x = new_variable(3, 1, 0, n_vars); + expr *p = new_parameter(1, 1, 0, n_vars, p0); + expr *m = new_scalar_mult(p, x); /* p * x: affine, parameter-dependent */ + + jacobian_init(m); + m->forward(m, u); + eval_jacobian(m); + + int nnz = m->jacobian->nnz; + mu_assert("scalar_mult jacobian must have entries", nnz == 3); + for (int ii = 0; ii < nnz; ii++) + { + mu_assert("jacobian must be p", m->jacobian->x[ii] == 2.0); + } + + for (int round = 0; round < 3; round++) + { + uint64_t v = m->jacobian->values_version; + mu_assert("walk must report parameter-dependent", + expr_set_needs_refresh(m) == true); + mu_assert("parameter-dependent node must be re-armed", + m->work->jacobian_evaluated == false); + eval_jacobian(m); + mu_assert("re-armed eval must bump", m->jacobian->values_version == v + 1); + } + + /* and a real parameter change must show through */ + p->value[0] = -5.0; + expr_set_needs_refresh(m); + m->forward(m, u); + eval_jacobian(m); + for (int ii = 0; ii < nnz; ii++) + { + mu_assert("jacobian must track the new parameter", + m->jacobian->x[ii] == -5.0); + } + + free_expr(m); + return 0; +} + +/* A PARAM_FIXED constant is not an updatable parameter. add(x, c) contains a + * parameter NODE but nothing theta can move, so the walk must still memoize + * it parameter-free and prune it. Guards the parameter leaf's + * set_needs_refresh_children against reporting every parameter_expr as + * parametric regardless of param_id. */ +const char *test_refresh_prunes_fixed_constant(void) +{ + double u[3] = {0.1, 0.2, 0.3}; + double c_vals[3] = {1.0, 2.0, 3.0}; + + expr *x = new_variable(3, 1, 0, 3); + expr *c = new_parameter(3, 1, PARAM_FIXED, 3, c_vals); + expr *m = new_add(x, c); + + jacobian_init(m); + m->forward(m, u); + eval_jacobian(m); + + mu_assert("first walk must report a fixed constant as parameter-free", + expr_set_needs_refresh(m) == false); + mu_assert("constant subtree must memoize parameter-free", + m->has_params == false); + eval_jacobian(m); + + uint64_t v1 = m->jacobian->values_version; + int nnz = m->jacobian->nnz; + for (int ii = 0; ii < nnz; ii++) + { + m->jacobian->x[ii] = 42.0; /* poison */ + } + + mu_assert("later walk must still report parameter-free", + expr_set_needs_refresh(m) == false); + eval_jacobian(m); + mu_assert("pruned re-eval must not bump", m->jacobian->values_version == v1); + for (int ii = 0; ii < nnz; ii++) + { + mu_assert("impl must not have run (poison must survive)", + m->jacobian->x[ii] == 42.0); + } + + free_expr(m); + return 0; +} + +/* The same shape with an UPDATABLE parameter must behave the opposite way -- + * every walk re-arms it and the poison is overwritten. Together with + * test_refresh_prunes_fixed_constant this pins the param_id >= 0 test itself: + * a hook that ignores param_id passes one of these two and fails the other. */ +const char *test_refresh_rearms_updatable_constant(void) +{ + double u[3] = {0.1, 0.2, 0.3}; + double p_vals[3] = {1.0, 2.0, 3.0}; + + expr *x = new_variable(3, 1, 0, 3); + expr *p = new_parameter(3, 1, 0, 3, p_vals); + expr *m = new_add(x, p); + + jacobian_init(m); + m->forward(m, u); + eval_jacobian(m); + + mu_assert("walk must report an updatable parameter as parametric", + expr_set_needs_refresh(m) == true); + mu_assert("parametric subtree must memoize parametric", m->has_params == true); + + int nnz = m->jacobian->nnz; + for (int ii = 0; ii < nnz; ii++) + { + m->jacobian->x[ii] = 42.0; /* poison */ + } + + eval_jacobian(m); + for (int ii = 0; ii < nnz; ii++) + { + mu_assert("re-armed eval must overwrite the poison", + m->jacobian->x[ii] != 42.0); + } + + free_expr(m); + return 0; +} diff --git a/tests/problem/test_param_broadcast.h b/tests/problem/test_param_broadcast.h index c64d64b..bd31842 100644 --- a/tests/problem/test_param_broadcast.h +++ b/tests/problem/test_param_broadcast.h @@ -395,6 +395,19 @@ const char *test_param_scalar_mult_convolve(void) mu_assert("check_jacobian failed", check_jacobian_num(constraint, x_vals, NUMERICAL_DIFF_DEFAULT_H)); + /* test 3: s back to 1. The refresh walk memoizes parameter-dependence on + its first pass, so a kernel source the walk fails to reach is only + pruned -- and only serves a stale convolution matrix T -- from the + SECOND update on. One update cannot see it; this round pins the hook. */ + theta[0] = 1.0; + problem_update_params(prob, theta); + problem_constraint_forward(prob, x_vals); + problem_jacobian(prob); + mu_assert("stale constraint values on the second update", + cmp_double_array(prob->constraint_values, constrs, 6)); + mu_assert("stale jacobian values on the second update", + cmp_double_array(prob->jacobian->x, Ax, 12)); + free_problem(prob); return 0; } diff --git a/tests/problem/test_param_source_refresh.h b/tests/problem/test_param_source_refresh.h index 47550e8..07f52f7 100644 --- a/tests/problem/test_param_source_refresh.h +++ b/tests/problem/test_param_source_refresh.h @@ -84,6 +84,21 @@ const char *test_composite_source_left_matmul(void) mu_assert("stale jacobian values after update", cmp_double_array(prob->jacobian->x, Ax, 4)); + /* p back to 2. The refresh walk memoizes parameter-dependence on its + first pass, so a coefficient subtree the walk fails to reach is only + pruned -- and only serves stale values -- from the SECOND update on. + One update cannot see it; this round is what pins the hook. */ + theta[0] = 2.0; + problem_update_params(prob, theta); + problem_constraint_forward(prob, x_vals); + problem_jacobian(prob); + double constrs_2nd[2] = {10.0, 22.0}; + double Ax_2nd[4] = {2.0, 4.0, 6.0, 8.0}; + mu_assert("stale constraint values on the second update", + cmp_double_array(prob->constraint_values, constrs_2nd, 2)); + mu_assert("stale jacobian values on the second update", + cmp_double_array(prob->jacobian->x, Ax_2nd, 4)); + free_problem(prob); return 0; @@ -131,6 +146,19 @@ const char *test_composite_source_quad_form(void) mu_assert("stale gradient after update", cmp_double_array(prob->gradient_values, grad, 2)); + /* g back to 1. Parameter-dependence is memoized on the first walk, so a + source the walk fails to reach only goes stale from the SECOND update + on -- this round is what pins the hook. */ + theta[0] = 1.0; + problem_update_params(prob, theta); + obj_val = problem_objective_forward(prob, x_vals); + problem_gradient(prob); + grad[0] = 4.0; + grad[1] = 12.0; + mu_assert("stale objective on the second update", fabs(obj_val - 14.0) < 1e-10); + mu_assert("stale gradient on the second update", + cmp_double_array(prob->gradient_values, grad, 2)); + free_problem(prob); return 0; @@ -190,6 +218,21 @@ const char *test_composite_source_nested_gates(void) mu_assert("stale jacobian values after update", cmp_double_array(prob->jacobian->x, Ax, 4)); + /* p back to 1. The refresh walk memoizes parameter-dependence on its + first pass, so a coefficient subtree the walk fails to reach is only + pruned -- and only serves stale values -- from the SECOND update on. + One update cannot see it; this round is what pins the hook. */ + theta[0] = 1.0; + problem_update_params(prob, theta); + problem_constraint_forward(prob, x_vals); + problem_jacobian(prob); + double constrs_2nd[2] = {10.0, 22.0}; + double Ax_2nd[4] = {2.0, 4.0, 6.0, 8.0}; + mu_assert("stale constraint values on the second update", + cmp_double_array(prob->constraint_values, constrs_2nd, 2)); + mu_assert("stale jacobian values on the second update", + cmp_double_array(prob->jacobian->x, Ax_2nd, 4)); + free_problem(prob); return 0; @@ -239,6 +282,18 @@ const char *test_composite_source_kron(void) mu_assert("stale gradient after update", cmp_double_array(prob->gradient_values, grad, 4)); + /* p back to 2. Parameter-dependence is memoized on the first walk, so an + operand the walk fails to reach only goes stale from the SECOND update + on -- this round is what pins the hook. */ + theta[0] = 2.0; + problem_update_params(prob, theta); + obj_val = problem_objective_forward(prob, x_vals); + problem_gradient(prob); + grad[0] = grad[1] = grad[2] = grad[3] = 20.0; + mu_assert("stale objective on the second update", fabs(obj_val - 80.0) < 1e-10); + mu_assert("stale gradient on the second update", + cmp_double_array(prob->gradient_values, grad, 4)); + free_problem(prob); return 0; diff --git a/tests/profiling/profile_lasso.h b/tests/profiling/profile_lasso.h new file mode 100644 index 0000000..3786ade --- /dev/null +++ b/tests/profiling/profile_lasso.h @@ -0,0 +1,132 @@ +#ifndef PROFILE_LASSO_H +#define PROFILE_LASSO_H + +#include +#include +#include +#include + +#include "atoms/affine.h" +#include "atoms/elementwise_full_dom.h" +#include "expr.h" +#include "minunit.h" +#include "problem.h" +#include "subexpr.h" +#include "utils/Timer.h" + +/* Dense lasso over a lambda path. + * + * variables: x (n), t (n) -> n_vars = 2n + * objective: sum((A x - b)^2) + lam * sum(t) + * constraints: x - t <= 0, -x - t <= 0 + * + * lam is the ONLY registered parameter. A is a fixed dense buffer, so the + * whole A@x subtree is parameter-free -- the case the has_params prune is + * meant to exploit. Reports forward / gradient / jacobian time separately so + * the dense-Jacobian cost is not hidden behind the dgemv. + */ +const char *profile_lasso(void) +{ + const int m = 2000; + const int n = 785; + const int nv = 2 * n; + const int n_sweep = 50; + + double *A_data = (double *) malloc((size_t) m * n * sizeof(double)); + double *negb = (double *) malloc((size_t) m * sizeof(double)); + double *u = (double *) malloc((size_t) nv * sizeof(double)); + srand(42); + for (int i = 0; i < m * n; i++) A_data[i] = (double) rand() / RAND_MAX - 0.5; + for (int i = 0; i < m; i++) negb[i] = (double) rand() / RAND_MAX - 0.5; + for (int i = 0; i < nv; i++) u[i] = (double) rand() / RAND_MAX - 0.5; + + /* ---- objective: sum((Ax - b)^2) + lam*sum(t) ---- */ + expr *x = new_variable(n, 1, 0, nv); + expr *t = new_variable(n, 1, n, nv); + + expr *Ax = new_left_matmul_dense(NULL, x, m, n, A_data); + expr *b_const = new_parameter(m, 1, PARAM_FIXED, nv, negb); + expr *resid = new_add(Ax, b_const); + expr *sq = new_power(resid, 2.0); + expr *ssq = new_sum(sq, -1); + + double lam0 = 1.0; + expr *lam = new_parameter(1, 1, 0, nv, &lam0); + expr *sum_t = new_sum(t, -1); + expr *pen = new_scalar_mult(lam, sum_t); + + expr *objective = new_add(ssq, pen); + + /* ---- constraints: x - t <= 0, -x - t <= 0 ---- */ + expr *c1 = new_add(x, new_neg(t)); + expr *c2 = new_add(new_neg(x), new_neg(t)); + expr *constraints[2] = {c1, c2}; + + problem *prob = new_problem(objective, constraints, 2, false); + expr *param_nodes[1] = {lam}; + problem_register_params(prob, param_nodes, 1); + + Timer t_init; + clock_gettime(CLOCK_MONOTONIC, &t_init.start); + problem_init_derivatives(prob); + clock_gettime(CLOCK_MONOTONIC, &t_init.end); + double sec_init = GET_ELAPSED_SECONDS(t_init); + + /* ---- lambda path ---- */ + Timer t_upd, t_fwd, t_grad, t_jac; + double sec_upd = 0, sec_fwd = 0, sec_grad = 0, sec_jac = 0; + double checksum = 0.0; + + for (int k = 0; k < n_sweep; k++) + { + double theta[1] = {0.01 + 0.02 * k}; + + clock_gettime(CLOCK_MONOTONIC, &t_upd.start); + problem_update_params(prob, theta); + clock_gettime(CLOCK_MONOTONIC, &t_upd.end); + sec_upd += GET_ELAPSED_SECONDS(t_upd); + + clock_gettime(CLOCK_MONOTONIC, &t_fwd.start); + checksum += problem_objective_forward(prob, u); + problem_constraint_forward(prob, u); + clock_gettime(CLOCK_MONOTONIC, &t_fwd.end); + sec_fwd += GET_ELAPSED_SECONDS(t_fwd); + + clock_gettime(CLOCK_MONOTONIC, &t_grad.start); + problem_gradient(prob); + clock_gettime(CLOCK_MONOTONIC, &t_grad.end); + sec_grad += GET_ELAPSED_SECONDS(t_grad); + + clock_gettime(CLOCK_MONOTONIC, &t_jac.start); + problem_jacobian(prob); + clock_gettime(CLOCK_MONOTONIC, &t_jac.end); + sec_jac += GET_ELAPSED_SECONDS(t_jac); + + checksum += prob->gradient_values[0] + prob->jacobian->x[0]; + } + + printf("\n [lasso] m=%d n=%d n_vars=%d sweep=%d\n", m, n, nv, n_sweep); + printf(" [lasso] jacobian nnz = %d\n", prob->jacobian->nnz); + printf(" [lasso] objective jac nnz = %d\n", objective->jacobian->nnz); + printf(" [lasso] Ax jac nnz = %d\n", Ax->jacobian->nnz); + printf(" [lasso] init_derivatives = %.4f s\n", sec_init); + printf(" [lasso] update_params (total) = %.4f s (%.3f ms/iter)\n", sec_upd, + 1e3 * sec_upd / n_sweep); + printf(" [lasso] forward (total) = %.4f s (%.3f ms/iter)\n", sec_fwd, + 1e3 * sec_fwd / n_sweep); + printf(" [lasso] gradient (total) = %.4f s (%.3f ms/iter)\n", sec_grad, + 1e3 * sec_grad / n_sweep); + printf(" [lasso] jacobian (total) = %.4f s (%.3f ms/iter)\n", sec_jac, + 1e3 * sec_jac / n_sweep); + printf(" [lasso] checksum = %.6f\n", checksum); + + /* free_problem releases objective + constraints (it retained them). */ + free_problem(prob); + free(A_data); + free(negb); + free(u); + + return 0; +} + +#endif /* PROFILE_LASSO_H */