[WIP] values_version cache coherence for Jacobian/Hessian value caches

include/expr.h

@@ -23,6 +23,7 @@
 #include "utils/matrix.h"
 #include <stdbool.h>
 #include <stddef.h>
+#include <stdint.h>
 #include <string.h>
 
 #define JAC_IDXS_NOT_SET -1
@@ -39,6 +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);
 
 /* Workspace for derivative computation */
 typedef struct
@@ -48,10 +50,19 @@ typedef struct
     CSC_matrix *jacobian_csc;
     int *csc_work; /* for CSR_matrix-CSC_matrix conversion */
 
-    /* jacobian_csc_filled is only used for affine functions to avoid redundant
-       conversions. Could become relevant for non-affine functions if we start
-       supporting common subexpressions on the Python side. */
-    bool jacobian_csc_filled;
+    /* jacobian->values_version that the jacobian_csc mirror reflects;
+       expr_refresh_jacobian_csc refills iff it differs. */
+    uint64_t jacobian_csc_seen;
+
+    /* node->is_affine(node), computed once by jacobian_init (affinity is
+       structural, and the recursive is_affine is too costly per eval). */
+    bool is_affine_cached;
+
+    /* True once eval_jacobian has run this parameter epoch; cleared by
+       expr_set_needs_refresh. Only consulted for affine nodes, where it
+       lets the eval_jacobian wrapper skip the values_version bump (same
+       role the old jacobian_csc_filled latch played). */
+    bool jacobian_evaluated;
     double *local_jac_diag; /* cached f'(g(x)) diagonal */
     matrix *hess_term1;     /* Jg^T D Jg workspace */
     matrix *hess_term2;     /* child wsum_hess workspace */
@@ -76,8 +87,8 @@ typedef struct expr
     forward_fn forward;
     jacobian_init_fn jacobian_init_impl;
     wsum_hess_init_fn wsum_hess_init_impl;
-    eval_jacobian_fn eval_jacobian;
-    wsum_hess_fn eval_wsum_hess;
+    eval_jacobian_fn eval_jacobian_impl;
+    wsum_hess_fn eval_wsum_hess_impl;
 
     // ------------------------------------------------------------------------
     //                      other things
@@ -86,7 +97,11 @@ 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 */
-    Expr_Work *work;                    /* derivative workspace */
+    /* 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. */
+    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
        (e.g. left_matmul_dense) check this flag before their forward
@@ -111,6 +126,15 @@ void free_expr(expr *node);
 void jacobian_init(expr *node);
 void wsum_hess_init(expr *node);
 
+/* Eval wrappers: run the atom's eval_*_impl and bump the output matrix's
+ * values_version so version-guarded caches (CSC mirrors, spd CSR views)
+ * refresh. Always call these instead of the impl slots. */
+void eval_jacobian(expr *node);
+void eval_wsum_hess(expr *node, const double *w);
+
+/* Refresh work->jacobian_csc from node->jacobian iff its values changed. */
+void expr_refresh_jacobian_csc(expr *node);
+
 /* Initialize CSC_matrix form of the Jacobian from the CSR_matrix Jacobian.
  * Must be called after jacobian_init. */
 void jacobian_csc_init(expr *node);

include/utils/matrix.h

@@ -21,6 +21,7 @@
 #include "CSC_matrix.h"
 #include "CSR_matrix.h"
 #include <stdbool.h>
+#include <stdint.h>
 
 /* Broadcast shape used by the broadcast atom and its vtable methods. */
 typedef enum
@@ -77,7 +78,8 @@ typedef void (*matrix_transpose_fill_values_fn)(const matrix *A, matrix *AT);
 typedef CSR_matrix *(*matrix_to_csr_fn)(matrix *A);
 
 /* Refresh any internal caches (e.g. a CSC_matrix mirror) so subsequent ATA /
-   ATDA calls reflect the current values. */
+   ATDA calls reflect the current values. Version-guarded: a no-op when the
+   cache already matches values_version, so it is cheap to call when fresh. */
 typedef void (*matrix_refresh_csc_values_fn)(matrix *A);
 
 /* Allocate C = A[indices, :] */
@@ -128,6 +130,14 @@ struct matrix
     bool is_permuted_dense;
     bool is_stacked_pd;
 
+    /* Monotone counter bumped whenever the matrix's values change. Consumers
+       that mirror the values into a cache (CSC mirror, CSR view, ...) record
+       the version they last saw and refresh iff it differs. Code that writes
+       x directly must call matrix_values_changed on the OWNER of the buffer —
+       aliased children (spd blocks, cache views) have no version of their
+       own. */
+    uint64_t values_version;
+
     /* Operator ops */
     matrix_block_left_mult_vec_fn block_left_mult_vec;
     matrix_block_left_mult_sparsity_fn block_left_mult_sparsity;
@@ -160,6 +170,12 @@ struct matrix
     matrix_free_fn free_fn;
 };
 
+/* Notify the library after writing A->x directly. */
+static inline void matrix_values_changed(matrix *A)
+{
+    A->values_version++;
+}
+
 /* Free helper */
 static inline void free_matrix(matrix *m)
 {

include/utils/sparse_matrix.h

@@ -28,6 +28,7 @@ typedef struct sparse_matrix
     matrix base;
     CSR_matrix *csr;
     CSC_matrix *csc_cache;
+    uint64_t csc_seen; /* base.values_version the csc_cache values reflect */
     int *csc_iwork;
     int *transpose_iwork; /* sized csr->n; allocated by sparse_transpose_alloc
                              on the output sm and reused by

include/utils/stacked_pd.h

@@ -48,6 +48,7 @@ typedef struct stacked_pd
 
     /* lazily built CSR view */
     CSR_matrix *csr_cache;
+    uint64_t csr_seen; /* base.values_version the csr_cache values reflect */
 
     /* Private permuted_dense scratch owned by the kernel that produced
        this spd. Allocated by the producing _alloc, used (without

src/atoms/affine/add.c

@@ -50,11 +50,11 @@ static void jacobian_init_impl(expr *node)
     sum_matrices_alloc(node->left->jacobian, node->right->jacobian, node->jacobian);
 }
 
-static void eval_jacobian(expr *node)
+static void eval_jacobian_impl(expr *node)
 {
     /* evaluate children's jacobians */
-    node->left->eval_jacobian(node->left);
-    node->right->eval_jacobian(node->right);
+    eval_jacobian(node->left);
+    eval_jacobian(node->right);
 
     /* sum children's jacobians */
     sum_matrices_fill_values(node->left->jacobian, node->right->jacobian,
@@ -76,11 +76,11 @@ static void wsum_hess_init_impl(expr *node)
                        node->wsum_hess);
 }
 
-static void eval_wsum_hess(expr *node, const double *w)
+static void eval_wsum_hess_impl(expr *node, const double *w)
 {
     /* evaluate children's wsum_hess */
-    node->left->eval_wsum_hess(node->left, w);
-    node->right->eval_wsum_hess(node->right, w);
+    eval_wsum_hess(node->left, w);
+    eval_wsum_hess(node->right, w);
 
     /* sum children's wsum_hess */
     sum_matrices_fill_values(node->left->wsum_hess, node->right->wsum_hess,
@@ -97,7 +97,8 @@ expr *new_add(expr *left, expr *right)
     assert(left->d1 == right->d1 && left->d2 == right->d2);
     expr *node = (expr *) sp_calloc(1, sizeof(expr));
     init_expr(node, left->d1, left->d2, left->n_vars, forward, jacobian_init_impl,
-              eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
+              eval_jacobian_impl, is_affine, wsum_hess_init_impl,
+              eval_wsum_hess_impl, NULL);
     node->left = left;
     node->right = right;
     expr_retain(left);

src/atoms/affine/broadcast.c

@@ -77,9 +77,9 @@ static void jacobian_init_impl(expr *node)
         x->jacobian->broadcast_alloc(x->jacobian, bcast->type, node->d1, node->d2);
 }
 
-static void eval_jacobian(expr *node)
+static void eval_jacobian_impl(expr *node)
 {
-    node->left->eval_jacobian(node->left);
+    eval_jacobian(node->left);
 
     /* fill values into the preallocated output. */
     broadcast_expr *bcast = (broadcast_expr *) node;
@@ -99,7 +99,7 @@ static void wsum_hess_init_impl(expr *node)
     node->work->dwork = sp_malloc(node->size * sizeof(double));
 }
 
-static void eval_wsum_hess(expr *node, const double *w)
+static void eval_wsum_hess_impl(expr *node, const double *w)
 {
     broadcast_expr *bcast = (broadcast_expr *) node;
     expr *x = node->left;
@@ -139,7 +139,7 @@ static void eval_wsum_hess(expr *node, const double *w)
         }
     }
 
-    x->eval_wsum_hess(x, node->work->dwork);
+    eval_wsum_hess(x, node->work->dwork);
     memcpy(node->wsum_hess->x, x->wsum_hess->x,
            node->wsum_hess->nnz * sizeof(double));
 }
@@ -183,7 +183,8 @@ expr *new_broadcast(expr *child, int d1, int d2)
     //                  initialize the rest of the expression
     // --------------------------------------------------------------------------
     init_expr(node, d1, d2, child->n_vars, forward, jacobian_init_impl,
-              eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
+              eval_jacobian_impl, is_affine, wsum_hess_init_impl,
+              eval_wsum_hess_impl, NULL);
     node->left = child;
     expr_retain(child);
     bcast->type = type;

src/atoms/affine/convolve.c

@@ -91,12 +91,12 @@ static void jacobian_init_impl(expr *node)
         new_sparse_matrix(csr_csc_matmul_alloc(cnode->T, cnode->Jchild_CSC));
 }
 
-static void eval_jacobian(expr *node)
+static void eval_jacobian_impl(expr *node)
 {
     expr *child = node->left;
     convolve_expr *cnode = (convolve_expr *) node;
 
-    child->eval_jacobian(child);
+    eval_jacobian(child);
 
     /* J = T @ J_child */
     csr_to_csc_fill_values(child->jacobian->to_csr(child->jacobian),
@@ -115,7 +115,7 @@ static void wsum_hess_init_impl(expr *node)
     node->work->dwork = (double *) sp_malloc(cnode->n * sizeof(double));
 }
 
-static void eval_wsum_hess(expr *node, const double *w)
+static void eval_wsum_hess_impl(expr *node, const double *w)
 {
     expr *child = node->left;
     convolve_expr *cnode = (convolve_expr *) node;
@@ -133,7 +133,7 @@ static void eval_wsum_hess(expr *node, const double *w)
         w_prime[j] = sum;
     }
 
-    child->eval_wsum_hess(child, w_prime);
+    eval_wsum_hess(child, w_prime);
     memcpy(node->wsum_hess->x, child->wsum_hess->x,
            node->wsum_hess->nnz * sizeof(double));
 }
@@ -181,8 +181,8 @@ expr *new_convolve(expr *param_node, expr *child)
     convolve_expr *cnode = (convolve_expr *) sp_calloc(1, sizeof(convolve_expr));
     expr *node = &cnode->base;
     init_expr(node, d1, d2, child->n_vars, forward, jacobian_init_impl,
-              eval_jacobian, is_affine, wsum_hess_init_impl, eval_wsum_hess,
-              free_type_data);
+              eval_jacobian_impl, is_affine, wsum_hess_init_impl,
+              eval_wsum_hess_impl, free_type_data);
     node->left = child;
     expr_retain(child);
 

src/atoms/affine/diag_vec.c

@@ -54,9 +54,9 @@ static void jacobian_init_impl(expr *node)
     node->jacobian = x->jacobian->diag_vec_alloc(x->jacobian);
 }
 
-static void eval_jacobian(expr *node)
+static void eval_jacobian_impl(expr *node)
 {
-    node->left->eval_jacobian(node->left);
+    eval_jacobian(node->left);
 
     /* fill the diagonal rows of the preallocated output. */
     node->left->jacobian->diag_vec_fill_values(node->left->jacobian, node->jacobian);
@@ -77,7 +77,7 @@ static void wsum_hess_init_impl(expr *node)
     node->wsum_hess = x->wsum_hess->copy_sparsity(x->wsum_hess);
 }
 
-static void eval_wsum_hess(expr *node, const double *w)
+static void eval_wsum_hess_impl(expr *node, const double *w)
 {
     expr *x = node->left;
     int n = x->size;
@@ -89,7 +89,7 @@ static void eval_wsum_hess(expr *node, const double *w)
     }
 
     /* Evaluate child's Hessian with extracted weights */
-    x->eval_wsum_hess(x, node->work->dwork);
+    eval_wsum_hess(x, node->work->dwork);
     memcpy(node->wsum_hess->x, x->wsum_hess->x,
            node->wsum_hess->nnz * sizeof(double));
 }
@@ -107,8 +107,9 @@ expr *new_diag_vec(expr *child)
     /* n is the number of elements (works for both row and column vectors) */
     int n = child->size;
     expr *node = (expr *) sp_calloc(1, sizeof(expr));
-    init_expr(node, n, n, child->n_vars, forward, jacobian_init_impl, eval_jacobian,
-              is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
+    init_expr(node, n, n, child->n_vars, forward, jacobian_init_impl,
+              eval_jacobian_impl, is_affine, wsum_hess_init_impl,
+              eval_wsum_hess_impl, NULL);
     node->left = child;
     expr_retain(child);
 

src/atoms/affine/hstack.c

@@ -87,15 +87,15 @@ static void jacobian_init_impl(expr *node)
     node->jacobian = new_sparse_matrix(A);
 }
 
-static void eval_jacobian(expr *node)
+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];
-        child->eval_jacobian(child);
+        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,
@@ -148,7 +148,7 @@ static void wsum_hess_eval(expr *node, const double *w)
     for (int i = 0; i < hnode->n_args; i++)
     {
         expr *child = hnode->args[i];
-        child->eval_wsum_hess(child, w + row_offset);
+        eval_wsum_hess(child, w + row_offset);
         copy_CSR_matrix(H, hnode->CSR_work);
         sum_csr_fill_values(hnode->CSR_work,
                             child->wsum_hess->to_csr(child->wsum_hess), H);
@@ -170,6 +170,17 @@ static bool is_affine(const expr *node)
     return true;
 }
 
+/* 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)
+{
+    hstack_expr *hnode = (hstack_expr *) node;
+    for (int i = 0; i < hnode->n_args; i++)
+    {
+        expr_set_needs_refresh(hnode->args[i]);
+    }
+}
+
 static void free_type_data(expr *node)
 {
     hstack_expr *hnode = (hstack_expr *) node;
@@ -199,9 +210,11 @@ expr *new_hstack(expr **args, int n_args, int n_vars)
     hstack_expr *hnode = (hstack_expr *) sp_calloc(1, sizeof(hstack_expr));
     expr *node = &hnode->base;
     init_expr(node, args[0]->d1, d2, n_vars, forward, jacobian_init_impl,
-              eval_jacobian, is_affine, wsum_hess_init_impl, wsum_hess_eval,
+              eval_jacobian_impl, is_affine, wsum_hess_init_impl, wsum_hess_eval,
               free_type_data);
 
+    node->set_needs_refresh_children = set_needs_refresh_children;
+
     /* Set type-specific fields (deep copy args array) */
     hnode->args = (expr **) sp_calloc(n_args, sizeof(expr *));
     hnode->n_args = n_args;