diff --git a/include/utils/permuted_dense.h b/include/utils/permuted_dense.h index 47e4ab8..671a0fc 100644 --- a/include/utils/permuted_dense.h +++ b/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 diff --git a/src/atoms/affine/left_matmul.c b/src/atoms/affine/left_matmul.c index e4680f6..2cc4171 100644 --- a/src/atoms/affine/left_matmul.c +++ b/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 diff --git a/src/atoms/other/quad_form.c b/src/atoms/other/quad_form.c index 5ba79ae..f7e98b8 100644 --- a/src/atoms/other/quad_form.c +++ b/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) diff --git a/src/utils/permuted_dense_linalg.c b/src/utils/permuted_dense_linalg.c index dc493e7..bf006b8 100644 --- a/src/utils/permuted_dense_linalg.c +++ b/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; } diff --git a/src/utils/stacked_pd_linalg.c b/src/utils/stacked_pd_linalg.c index 90a098f..522d898 100644 --- a/src/utils/stacked_pd_linalg.c +++ b/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]); } } diff --git a/tests/all_tests.c b/tests/all_tests.c index c6f137d..c0fcfd9 100644 --- a/tests/all_tests.c +++ b/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); diff --git a/tests/utils/test_matmul_dispatchers.h b/tests/utils/test_matmul_dispatchers.h index 5620574..8d2ad81 100644 --- a/tests/utils/test_matmul_dispatchers.h +++ b/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 */ diff --git a/tests/utils/test_stacked_pd.h b/tests/utils/test_stacked_pd.h index f27a400..08e70f4 100644 --- a/tests/utils/test_stacked_pd.h +++ b/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;