From 46cb698f79aa944f17f4336d56940c9f6dc48c26 Mon Sep 17 00:00:00 2001 From: dance858 Date: Sat, 12 Sep 2026 13:51:44 -0700 Subject: [PATCH] Migrate promote onto row_gather; delete the promote vtable slots (lesson 3, M2a) promote's Jacobian is a row gather with an all-zero map. The atom now builds that map in jacobian_init and calls row_gather_alloc/fill_values; the promote_alloc / promote_fill_values slots and their sparse, permuted_dense and stacked_pd kernels are removed. New atom-level test proves a pd child Jacobian stays pd through promote. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Azb8o1AgiFRU4t9aAY7oDr --- include/utils/matrix.h | 8 ---- include/utils/permuted_dense.h | 6 --- src/atoms/affine/promote.c | 11 ++++-- src/utils/permuted_dense.c | 44 ---------------------- src/utils/sparse_matrix.c | 28 -------------- src/utils/stacked_pd.c | 27 ------------- tests/all_tests.c | 3 +- tests/jacobian_tests/affine/test_promote.h | 31 +++++++++++++++ tests/utils/test_permuted_dense.h | 34 ----------------- tests/utils/test_stacked_pd.h | 37 ------------------ 10 files changed, 39 insertions(+), 190 deletions(-) diff --git a/include/utils/matrix.h b/include/utils/matrix.h index 0fc7287..041f7a8 100644 --- a/include/utils/matrix.h +++ b/include/utils/matrix.h @@ -89,12 +89,6 @@ typedef matrix *(*matrix_row_gather_alloc_fn)(const matrix *A, const int *map, /* Fill values of C = A[map, :] */ typedef void (*matrix_row_gather_fill_values_fn)(const matrix *A, matrix *C); -/* Row-tiling for the promote atom: A must be a 1-row matrix; returns - a new matrix of shape (size, A->n) where every row is a copy of A's - single row. */ -typedef matrix *(*matrix_promote_alloc_fn)(matrix *A, int size); -typedef void (*matrix_promote_fill_values_fn)(matrix *A, matrix *out); - /* Broadcast: lift the child Jacobian of a broadcast atom into the output Jacobian. `type` is the broadcast variant; (d1, d2) is the output shape. */ typedef matrix *(*matrix_broadcast_alloc_fn)(matrix *A, broadcast_type type, int d1, @@ -158,8 +152,6 @@ struct matrix /* Atom-specific ops */ matrix_row_gather_alloc_fn row_gather_alloc; matrix_row_gather_fill_values_fn row_gather_fill_values; - matrix_promote_alloc_fn promote_alloc; - matrix_promote_fill_values_fn promote_fill_values; matrix_broadcast_alloc_fn broadcast_alloc; matrix_broadcast_fill_values_fn broadcast_fill_values; matrix_diag_vec_alloc_fn diag_vec_alloc; diff --git a/include/utils/permuted_dense.h b/include/utils/permuted_dense.h index bb01b88..f8cc704 100644 --- a/include/utils/permuted_dense.h +++ b/include/utils/permuted_dense.h @@ -109,12 +109,6 @@ matrix *row_gather_pd_alloc(const permuted_dense *A, const int *map, int m_out); /* Fill values of C = A[map, :], where A and C are permuted dense. */ void row_gather_pd_fill_values(const permuted_dense *A, permuted_dense *C); -/* Allocate C = promote(A, size), where A and C are permuted dense. */ -matrix *promote_pd_alloc(const permuted_dense *A, int size); - -/* Fill values of C = promote(A, size). */ -void promote_pd_fill_values(const permuted_dense *A, permuted_dense *C); - /* Allocate C = diag_vec(A), where A and C are permuted dense. */ matrix *diag_vec_pd_alloc(const permuted_dense *A); diff --git a/src/atoms/affine/promote.c b/src/atoms/affine/promote.c index 7fea914..6750c18 100644 --- a/src/atoms/affine/promote.c +++ b/src/atoms/affine/promote.c @@ -40,9 +40,11 @@ static void jacobian_init_impl(expr *node) expr *x = node->left; jacobian_init(x); - /* allocate sparsity for an (node->size, n_vars) matrix whose rows are all - copies of the child's single row; output type matches child's type. */ - node->jacobian = x->jacobian->promote_alloc(x->jacobian, node->size); + /* every output row is the child's single row (row 0): a row gather with an + all-zero map. The map is bound to node->jacobian, so it is not kept. */ + int *map = (int *) sp_calloc(node->size, sizeof(int)); + node->jacobian = x->jacobian->row_gather_alloc(x->jacobian, map, node->size); + sp_free(map); } static void eval_jacobian_impl(expr *node) @@ -50,7 +52,8 @@ static void eval_jacobian_impl(expr *node) eval_jacobian(node->left); /* tile the child's single row into the preallocated output. */ - node->left->jacobian->promote_fill_values(node->left->jacobian, node->jacobian); + node->left->jacobian->row_gather_fill_values(node->left->jacobian, + node->jacobian); } static void wsum_hess_init_impl(expr *node) diff --git a/src/utils/permuted_dense.c b/src/utils/permuted_dense.c index f5ad9ea..17958e0 100644 --- a/src/utils/permuted_dense.c +++ b/src/utils/permuted_dense.c @@ -164,48 +164,6 @@ static void permuted_dense_vtable_row_gather_fill_values(const matrix *self, row_gather_pd_fill_values((const permuted_dense *) self, (permuted_dense *) out); } -matrix *promote_pd_alloc(const permuted_dense *A, int size) -{ - assert(A->m0 <= 1); - - if (A->m0 == 0) - { - /* source row is all-zero; output is also structurally all-zero. */ - return new_permuted_dense(size, A->base.n, 0, A->n0, NULL, A->col_perm, - NULL); - } - - int *new_row_perm = (int *) sp_malloc(size * sizeof(int)); - for (int i = 0; i < size; i++) - { - new_row_perm[i] = i; - } - matrix *out = new_permuted_dense(size, A->base.n, size, A->n0, new_row_perm, - A->col_perm, NULL); - sp_free(new_row_perm); - return out; -} - -void promote_pd_fill_values(const permuted_dense *A, permuted_dense *C) -{ - if (A->m0 == 0) return; - int n0 = A->n0; - for (int k = 0; k < C->m0; k++) - { - memcpy(C->X + k * n0, A->X, n0 * sizeof(double)); - } -} - -static matrix *permuted_dense_vtable_promote_alloc(matrix *self, int size) -{ - return promote_pd_alloc((const permuted_dense *) self, size); -} - -static void permuted_dense_vtable_promote_fill_values(matrix *self, matrix *out) -{ - promote_pd_fill_values((const permuted_dense *) self, (permuted_dense *) out); -} - matrix *broadcast_pd_alloc(const permuted_dense *A, broadcast_type type, int d1, int d2) { @@ -567,8 +525,6 @@ static void wire_vtable(permuted_dense *pd) pd->base.transpose_fill_values = permuted_dense_vtable_transpose_fill_values; pd->base.row_gather_alloc = permuted_dense_vtable_row_gather_alloc; pd->base.row_gather_fill_values = permuted_dense_vtable_row_gather_fill_values; - pd->base.promote_alloc = permuted_dense_vtable_promote_alloc; - pd->base.promote_fill_values = permuted_dense_vtable_promote_fill_values; pd->base.broadcast_alloc = permuted_dense_vtable_broadcast_alloc; pd->base.broadcast_fill_values = permuted_dense_vtable_broadcast_fill_values; pd->base.diag_vec_alloc = permuted_dense_vtable_diag_vec_alloc; diff --git a/src/utils/sparse_matrix.c b/src/utils/sparse_matrix.c index 762107b..b85ede9 100644 --- a/src/utils/sparse_matrix.c +++ b/src/utils/sparse_matrix.c @@ -185,32 +185,6 @@ static void sparse_row_gather_fill_values(const matrix *self, matrix *out) } } -static matrix *sparse_promote_alloc(matrix *self, int size) -{ - CSR_matrix *Jx = ((sparse_matrix *) self)->csr; - int row_nnz = Jx->nnz; - CSR_matrix *J = new_CSR_matrix(size, self->n, size * row_nnz); - - for (int row = 0; row < size; row++) - { - J->p[row] = row * row_nnz; - memcpy(J->i + row * row_nnz, Jx->i, row_nnz * sizeof(int)); - } - J->p[size] = size * row_nnz; - J->nnz = size * row_nnz; - return new_sparse_matrix(J); -} - -static void sparse_promote_fill_values(matrix *self, matrix *out) -{ - CSR_matrix *Jx = ((sparse_matrix *) self)->csr; - int row_nnz = Jx->nnz; - for (int row = 0; row < out->m; row++) - { - memcpy(out->x + row * row_nnz, Jx->x, row_nnz * sizeof(double)); - } -} - static matrix *sparse_broadcast_alloc(matrix *self, broadcast_type type, int d1, int d2) { @@ -402,8 +376,6 @@ static void wire_vtable(sparse_matrix *sm) sm->base.transpose_fill_values = sparse_transpose_fill_values; sm->base.row_gather_alloc = sparse_row_gather_alloc; sm->base.row_gather_fill_values = sparse_row_gather_fill_values; - sm->base.promote_alloc = sparse_promote_alloc; - sm->base.promote_fill_values = sparse_promote_fill_values; sm->base.broadcast_alloc = sparse_broadcast_alloc; sm->base.broadcast_fill_values = sparse_broadcast_fill_values; sm->base.diag_vec_alloc = sparse_diag_vec_alloc; diff --git a/src/utils/stacked_pd.c b/src/utils/stacked_pd.c index a178279..fca0cbe 100644 --- a/src/utils/stacked_pd.c +++ b/src/utils/stacked_pd.c @@ -267,31 +267,6 @@ static void stacked_pd_vtable_row_gather_fill_values(const matrix *self, matrix } } -// ----------------------------------------------------------------------------- -// promote: C = promote(A) where A is stacked_pd -// ----------------------------------------------------------------------------- -static matrix *wrapper_pd_promote(permuted_dense *Bk, const void *ctx) -{ - return promote_pd_alloc(Bk, *(const int *) ctx); -} - -static matrix *stacked_pd_vtable_promote_alloc(matrix *self, int size) -{ - stacked_pd *src = (stacked_pd *) self; - return spd_map_filter_blocks(src, size, src->base.n, wrapper_pd_promote, &size); -} - -static void stacked_pd_vtable_promote_fill_values(matrix *self, matrix *out) -{ - stacked_pd *src = (stacked_pd *) self; - stacked_pd *out_spd = (stacked_pd *) out; - for (int k = 0; k < out_spd->n_blocks; k++) - { - int sk = out_spd->src_block_idx[k]; - promote_pd_fill_values(src->blocks[sk], out_spd->blocks[k]); - } -} - // ----------------------------------------------------------------------------- // diag_vec: C = diag(vec(A)) where A is stacked_pd // ----------------------------------------------------------------------------- @@ -496,8 +471,6 @@ static void wire_vtable(stacked_pd *spd) spd->base.to_csr = stacked_pd_to_csr; spd->base.row_gather_alloc = stacked_pd_vtable_row_gather_alloc; spd->base.row_gather_fill_values = stacked_pd_vtable_row_gather_fill_values; - spd->base.promote_alloc = stacked_pd_vtable_promote_alloc; - spd->base.promote_fill_values = stacked_pd_vtable_promote_fill_values; spd->base.diag_vec_alloc = stacked_pd_vtable_diag_vec_alloc; spd->base.diag_vec_fill_values = stacked_pd_vtable_diag_vec_fill_values; spd->base.broadcast_alloc = stacked_pd_vtable_broadcast_alloc; diff --git a/tests/all_tests.c b/tests/all_tests.c index 4be944e..fda2287 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -245,6 +245,7 @@ int main(void) mu_run_test(test_sum_of_index, tests_run); mu_run_test(test_promote_scalar_jacobian, tests_run); mu_run_test(test_promote_scalar_to_matrix_jacobian, tests_run); + mu_run_test(test_promote_jacobian_pd_preserved, tests_run); mu_run_test(test_broadcast_row_jacobian, tests_run); mu_run_test(test_broadcast_col_jacobian, tests_run); mu_run_test(test_broadcast_scalar_to_matrix_jacobian, tests_run); @@ -444,7 +445,6 @@ int main(void) #ifdef SP_TRACK_MEMORY mu_run_test(test_row_gather_spd_fill_no_transient_alloc, tests_run); #endif - mu_run_test(test_permuted_dense_promote, tests_run); mu_run_test(test_permuted_dense_broadcast_scalar, tests_run); mu_run_test(test_permuted_dense_broadcast_row, tests_run); mu_run_test(test_permuted_dense_broadcast_col, tests_run); @@ -559,7 +559,6 @@ int main(void) mu_run_test(test_spd_vtable_transpose, tests_run); mu_run_test(test_spd_vtable_refresh_csc_values_noop, tests_run); mu_run_test(test_spd_vtable_row_gather, tests_run); - mu_run_test(test_spd_vtable_promote, tests_run); mu_run_test(test_spd_vtable_diag_vec, tests_run); mu_run_test(test_spd_vtable_broadcast_row, tests_run); mu_run_test(test_YT_kron_I, tests_run); diff --git a/tests/jacobian_tests/affine/test_promote.h b/tests/jacobian_tests/affine/test_promote.h index 4746645..46a8cec 100644 --- a/tests/jacobian_tests/affine/test_promote.h +++ b/tests/jacobian_tests/affine/test_promote.h @@ -6,6 +6,7 @@ #include "expr.h" #include "minunit.h" #include "test_helpers.h" +#include "utils/permuted_dense.h" const char *test_promote_scalar_jacobian(void) { @@ -58,3 +59,33 @@ const char *test_promote_scalar_to_matrix_jacobian(void) free_expr(promote_node); return 0; } + +/* A pd child Jacobian stays pd through promote: AU = A @ u with A a 1x2 dense + constant and u a 2x1 variable has a (1, 2) pd Jacobian with m0 = 1; + promote(AU, 3, 2) gathers that single row six times. */ +const char *test_promote_jacobian_pd_preserved(void) +{ + double A[2] = {1.5, -2.0}; + expr *u = new_variable(2, 1, 0, 2); + expr *AU = new_left_matmul_dense(NULL, u, 1, 2, A); + expr *P = new_promote(AU, 3, 2); + + double u_vals[2] = {0.5, -1.5}; + jacobian_init(P); + P->forward(P, u_vals); + eval_jacobian(P); + + mu_assert("promote Jacobian should be PD", P->jacobian->is_permuted_dense); + permuted_dense *pd = (permuted_dense *) P->jacobian; + mu_assert("shape", P->jacobian->m == 6 && P->jacobian->n == 2); + mu_assert("m0", pd->m0 == 6); + mu_assert("n0", pd->n0 == 2); + int expected_row_perm[6] = {0, 1, 2, 3, 4, 5}; + mu_assert("row_perm", cmp_int_array(pd->row_perm, expected_row_perm, 6)); + double expected_X[12] = {1.5, -2.0, 1.5, -2.0, 1.5, -2.0, + 1.5, -2.0, 1.5, -2.0, 1.5, -2.0}; + mu_assert("X values", cmp_double_array(pd->X, expected_X, 12)); + + free_expr(P); + return 0; +} diff --git a/tests/utils/test_permuted_dense.h b/tests/utils/test_permuted_dense.h index ffed992..db6d55b 100644 --- a/tests/utils/test_permuted_dense.h +++ b/tests/utils/test_permuted_dense.h @@ -387,40 +387,6 @@ const char *test_permuted_dense_row_gather(void) return 0; } -/* PD promote_alloc / promote_fill_values: tile a 1-row PD into a - `size`-row PD where every row is a copy of the source row. */ -const char *test_permuted_dense_promote(void) -{ - /* Source PD, shape (1, 5), single dense row at row 0, cols {1, 3}. */ - int row_perm[1] = {0}; - int col_perm[2] = {1, 3}; - double X[2] = {7.0, 9.0}; - matrix *M = new_permuted_dense(1, 5, 1, 2, row_perm, col_perm, X); - - matrix *out = M->promote_alloc(M, 4); - permuted_dense *out_pd = (permuted_dense *) out; - - mu_assert("out m", out->m == 4); - mu_assert("out n", out->n == 5); - mu_assert("out nnz", out->nnz == 8); /* m0=4 * n0=2 */ - mu_assert("m0", out_pd->m0 == 4); - mu_assert("n0", out_pd->n0 == 2); - - int expected_row_perm[4] = {0, 1, 2, 3}; - mu_assert("row_perm", cmp_int_array(out_pd->row_perm, expected_row_perm, 4)); - int expected_col_perm[2] = {1, 3}; - mu_assert("col_perm", cmp_int_array(out_pd->col_perm, expected_col_perm, 2)); - - M->promote_fill_values(M, out); - - double expected_X[8] = {7.0, 9.0, 7.0, 9.0, 7.0, 9.0, 7.0, 9.0}; - mu_assert("values", cmp_double_array(out_pd->X, expected_X, 8)); - - free_matrix(out); - free_matrix(M); - return 0; -} - /* PD broadcast_alloc / broadcast_fill_values, SCALAR variant. (1, 5) PD with single dense row -> (d1*d2, 5) PD with that row tiled. */ const char *test_permuted_dense_broadcast_scalar(void) diff --git a/tests/utils/test_stacked_pd.h b/tests/utils/test_stacked_pd.h index 495b62f..4777225 100644 --- a/tests/utils/test_stacked_pd.h +++ b/tests/utils/test_stacked_pd.h @@ -1395,43 +1395,6 @@ const char *test_spd_vtable_row_gather(void) return 0; } -/* promote_* on spd: replicate the single row across `size` rows; per-block - delegation drops empty blocks so output has at most one block. */ -const char *test_spd_vtable_promote(void) -{ - /* 1x4 spd, single block carries row 0 at cols {0, 2} with values - [9, 11]. Promote to size=3. */ - int row_perm[1] = {0}; - int col_perm[2] = {0, 2}; - double X[2] = {9.0, 11.0}; - matrix *blk = new_permuted_dense(1, 4, 1, 2, row_perm, col_perm, X); - - permuted_dense *blocks[1] = {(permuted_dense *) blk}; - matrix *M = new_stacked_pd(1, 4, 1, blocks, NULL, NULL); - - matrix *C_m = M->promote_alloc(M, 3); - M->promote_fill_values(M, C_m); - stacked_pd *C = (stacked_pd *) C_m; - - mu_assert("n_blocks", C->n_blocks == 1); - mu_assert("base.m", C_m->m == 3); - mu_assert("base.n", C_m->n == 4); - - permuted_dense *out0 = C->blocks[0]; - int expected_row_perm[3] = {0, 1, 2}; - int expected_col_perm[2] = {0, 2}; - double expected_X[6] = {9.0, 11.0, 9.0, 11.0, 9.0, 11.0}; - mu_assert("out0 m0", out0->m0 == 3); - mu_assert("out0 n0", out0->n0 == 2); - mu_assert("out0 row_perm", cmp_int_array(out0->row_perm, expected_row_perm, 3)); - mu_assert("out0 col_perm", cmp_int_array(out0->col_perm, expected_col_perm, 2)); - mu_assert("out0 X", cmp_double_array(out0->X, expected_X, 6)); - - free_matrix(C_m); - free_matrix(M); - return 0; -} - /* diag_vec_* on spd: per-block row_perm entries are rescaled r -> r*(n+1); X buffers are unchanged; structure is preserved (same n_blocks). */ const char *test_spd_vtable_diag_vec(void)