BA_pd_spd transpose cache

include/utils/permuted_dense.h

@@ -67,8 +67,11 @@ typedef struct permuted_dense
     size_t kernel_iwork_size;
 
     /* Cached transpose of this PD as another permuted_dense, allocated lazily
-       on first call to permuted_dense_ensure_transpose_cache. */
+       on first call to permuted_dense_ensure_transpose_cache. On the cache PD
+       itself, transpose_seen records the source's base.values_version whose
+       values the cache holds; consumers refill iff it is stale. */
     struct permuted_dense *transpose_cache;
+    uint64_t transpose_seen;
 } permuted_dense;
 
 /* Constructor. row_perm and col_perm must be strictly increasing in their

src/atoms/affine/left_matmul.c

@@ -231,6 +231,8 @@ static void refresh_dense_left(left_matmul_expr *lnode)
        actually corresponds to the transpose of A, and we transpose AT to get A. */
     memcpy(lnode->AT->x, lnode->param_source->value, m * n * sizeof(double));
     A_transpose(lnode->A->x, lnode->AT->x, n, m);
+    matrix_values_changed(lnode->AT);
+    matrix_values_changed(lnode->A);
 }
 
 /* We expect u->d1 == A->n. However, numpy's broadcasting rules allow users to

src/atoms/other/quad_form.c

@@ -45,6 +45,7 @@ static void refresh_param_values_qf(quad_form_expr *qnode)
     qnode->base.needs_parameter_refresh = false;
     memcpy(qnode->Q->x, qnode->param_source->value,
            (size_t) qnode->n * qnode->n * sizeof(double));
+    matrix_values_changed(qnode->Q);
 }
 
 static void forward(expr *node, const double *u)

src/utils/permuted_dense_linalg.c

@@ -57,6 +57,8 @@ permuted_dense *permuted_dense_ensure_transpose_cache(const permuted_dense *B_co
     }
     permuted_dense *BT = (permuted_dense *) transpose_pd_alloc(B);
     B->transpose_cache = BT;
+    /* Deliberately stale so the first fill always refreshes the cache. */
+    BT->transpose_seen = B->base.values_version - 1;
     return BT;
 }
 

src/utils/stacked_pd_linalg.c

@@ -515,14 +515,18 @@ void BTDA_csc_spd_fill_values(const CSC_matrix *B, const double *d,
 // BA_pd_spd: C = B @ A where B is permuted_dense and A is stacked_pd. Thin
 // wrapper over the canonical BTA_pd_spd_* kernel: use B's lazily-cached
 // transpose and call BTA. The cache is populated on first call (in alloc)
-// and reused across subsequent fills.
+// and reused across subsequent fills; its values are refreshed only when
+// B's values_version has moved since the last fill (transpose_seen).
 //
 // Contract: B's perms must be immutable between alloc and fill (the cache
 // records B's perms at alloc time and is not re-validated at fill). For
 // callers where B's perms change between calls — notably the kron-spd path
 // that reuses a mutating scratch — bypass this wrapper and call
 // BTA_pd_spd_* directly. BA_dense_kron_spd does exactly that
-// (stacked_pd_kron_linalg.c) and is the only such caller today.
+// (stacked_pd_kron_linalg.c). The values_version guard additionally
+// requires B to be the owner of its value buffer; BA_spd_spd_fill_values
+// passes spd blocks (no version of their own) and therefore also bypasses
+// the wrapper for its fills.
 // ---------------------------------------------------------------------------------
 matrix *BA_pd_spd_alloc(const permuted_dense *B, const stacked_pd *A)
 {
@@ -534,7 +538,11 @@ void BA_pd_spd_fill_values(const permuted_dense *B, const stacked_pd *A,
                            permuted_dense *C)
 {
     permuted_dense *BT = B->transpose_cache;
-    transpose_pd_fill_values(B, BT);
+    if (BT->transpose_seen != B->base.values_version)
+    {
+        transpose_pd_fill_values(B, BT);
+        BT->transpose_seen = B->base.values_version;
+    }
     BTA_pd_spd_fill_values(BT, A, C);
 }
 
@@ -609,6 +617,11 @@ void BA_spd_spd_fill_values(const stacked_pd *B, const stacked_pd *A, stacked_pd
     {
         int q = C->src_block_idx[C->src_block_idx_p[k]];
         const permuted_dense *Bq = B->blocks[q];
-        BA_pd_spd_fill_values(Bq, A, C->blocks[k]);
+        /* Bypass BA_pd_spd_fill_values' version guard: spd blocks have no
+           values_version of their own (writers bump the owning spd), so the
+           cached transpose must be refreshed unconditionally here. */
+        permuted_dense *BqT = Bq->transpose_cache;
+        transpose_pd_fill_values(Bq, BqT);
+        BTA_pd_spd_fill_values(BqT, A, C->blocks[k]);
     }
 }

tests/all_tests.c

@@ -490,6 +490,7 @@ int main(void)
     mu_run_test(test_BTA_sparse_matrices_csc_A, tests_run);
     mu_run_test(test_BTA_sparse_matrices_spd_A, tests_run);
     mu_run_test(test_BA_pd_kron_spd_no_cache_staleness, tests_run);
+    mu_run_test(test_BA_pd_spd_transpose_cache_refresh, tests_run);
     mu_run_test(test_stacked_pd_construct_and_free, tests_run);
     mu_run_test(test_coalesce_no_overlap, tests_run);
     mu_run_test(test_coalesce_three_signatures, tests_run);

tests/utils/test_matmul_dispatchers.h

@@ -1328,8 +1328,10 @@ const char *test_BTA_sparse_matrices_spd_A(void)
    its row_perm / col_perm in place. A prior implementation of
    BA_pd_spd_fill_values cached B's transpose on B->transpose_cache,
    which (a) was never populated for the fill-path scratch, and (b)
-   would have held stale perms even if populated. Now BA_pd_spd does
-   per-call transpose alloc/free so this scenario is correct. The
+   would have held stale perms even if populated. BA_pd_spd still uses
+   that cache, but the kron path stays correct because
+   BA_dense_kron_spd bypasses the wrapper and calls BTA_pd_spd_*
+   directly on its mutating scratch (see stacked_pd_linalg.c). The
    reference path runs the same dispatch with J flattened to
    sparse_matrix (which routes via BA_dense_kron_csc, no transpose-cache
    involvement); we compare via to_csr. */
@@ -1393,4 +1395,68 @@ const char *test_BA_pd_kron_spd_no_cache_staleness(void)
     return 0;
 }
 
+/* BA_pd_spd transpose cache: the fill refreshes B's cached transpose iff
+   B's values_version moved since the last fill. Fill once, mutate B's
+   values + bump, refill, and compare against a fresh computation with the
+   mutated values. */
+const char *test_BA_pd_spd_transpose_cache_refresh(void)
+{
+    /* B: 3x4 pd with a non-square 2x3 block. */
+    int B_rp[2] = {0, 2};
+    int B_cp[3] = {0, 1, 3};
+    double BX[6] = {1, 2, 3, 4, 5, 6};
+    matrix *B_m = new_permuted_dense(3, 4, 2, 3, B_rp, B_cp, BX);
+    permuted_dense *B = (permuted_dense *) B_m;
+
+    /* A: 4x5 spd with two disjoint-row blocks. */
+    int A0_rp[2] = {0, 1};
+    int A0_cp[2] = {0, 2};
+    double A0X[4] = {1, 2, 3, 4};
+    matrix *Ablk0 = new_permuted_dense(4, 5, 2, 2, A0_rp, A0_cp, A0X);
+    int A1_rp[2] = {2, 3};
+    int A1_cp[2] = {1, 4};
+    double A1X[4] = {5, 6, 7, 8};
+    matrix *Ablk1 = new_permuted_dense(4, 5, 2, 2, A1_rp, A1_cp, A1X);
+    permuted_dense *A_blocks[2] = {(permuted_dense *) Ablk0,
+                                   (permuted_dense *) Ablk1};
+    matrix *A_spd = new_stacked_pd(4, 5, 2, A_blocks, NULL, NULL);
+    stacked_pd *A = (stacked_pd *) A_spd;
+
+    matrix *C = BA_pd_spd_alloc(B, A);
+    BA_pd_spd_fill_values(B, A, (permuted_dense *) C);
+    mu_assert("seen must match after fill",
+              B->transpose_cache->transpose_seen == B->base.values_version);
+
+    /* Mutate B's values and bump; the guarded refill must refresh the
+       cached transpose. */
+    double BX2[6] = {-1, 7, 0.5, 2, -3, 6};
+    memcpy(B->X, BX2, 6 * sizeof(double));
+    matrix_values_changed(B_m);
+    BA_pd_spd_fill_values(B, A, (permuted_dense *) C);
+    mu_assert("seen must catch up after bump + refill",
+              B->transpose_cache->transpose_seen == B->base.values_version);
+
+    /* Reference: fresh B with the mutated values, fresh cache. */
+    matrix *B2_m = new_permuted_dense(3, 4, 2, 3, B_rp, B_cp, BX2);
+    permuted_dense *B2 = (permuted_dense *) B2_m;
+    matrix *C_ref = BA_pd_spd_alloc(B2, A);
+    BA_pd_spd_fill_values(B2, A, (permuted_dense *) C_ref);
+
+    CSR_matrix *csr_ours = C->to_csr(C);
+    CSR_matrix *csr_ref = C_ref->to_csr(C_ref);
+    mu_assert("m", csr_ours->m == csr_ref->m);
+    mu_assert("n", csr_ours->n == csr_ref->n);
+    mu_assert("nnz", csr_ours->nnz == csr_ref->nnz);
+    mu_assert("p", cmp_int_array(csr_ours->p, csr_ref->p, csr_ours->m + 1));
+    mu_assert("i", cmp_int_array(csr_ours->i, csr_ref->i, csr_ours->nnz));
+    mu_assert("x", cmp_double_array(csr_ours->x, csr_ref->x, csr_ours->nnz));
+
+    free_matrix(C_ref);
+    free_matrix(B2_m);
+    free_matrix(C);
+    free_matrix(A_spd);
+    free_matrix(B_m);
+    return 0;
+}
+
 #endif /* TEST_MATMUL_DISPATCHERS_H */

tests/utils/test_stacked_pd.h

@@ -1118,12 +1118,15 @@ const char *test_BA_pd_spd_alloc_then_fill_values(void)
     BA_pd_spd_fill_values((permuted_dense *) B, (stacked_pd *) A,
                           (permuted_dense *) C_m);
 
-    /* Mutate B and A_0 values. */
+    /* Mutate B and A_0 values. B is an owner pd, so the write must be
+       announced (matrix.h contract) — BA_pd_spd_fill_values refreshes its
+       cached transpose of B iff B's values_version moved. */
     permuted_dense *B_pd = (permuted_dense *) B;
     B_pd->X[0] = 5;
     B_pd->X[1] = 6;
     B_pd->X[2] = 7;
     B_pd->X[3] = 8;
+    matrix_values_changed(B);
     permuted_dense *A0_pd = (permuted_dense *) A0;
     A0_pd->X[0] = 100;
     A0_pd->X[1] = 200;