From 065fa0a2f224ef9e80ece8401d0c63c3bed2948d Mon Sep 17 00:00:00 2001 From: dance858 Date: Mon, 31 Aug 2026 20:12:50 -0700 Subject: [PATCH 1/3] Fold diag(d) into the BTA gathers, delete the per-eval DA temps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BTDA_pd_pd_fill_values and BTDA_pd_spd_fill_values allocated, filled, and freed a diag(d) @ A intermediate on every Hessian evaluation (the "must remove this allocation. Very important" TODOs), and the blockwise BTDA_spd_pd/spd_spd variants compounded that to n_blocks temps per fill. Restructure each BTA kernel into a shared core taking an optional d (NULL = identity) that scales A's rows by d[global_row] while they are being gathered into the pre-sized kernel_dwork scratch — the same pattern BTDA_pd_csc/BTDA_csc_pd already use — so the BTDA entry points become thin wrappers and no fill allocates. The pd_pd fast path (identical row_perms) scales into A->kernel_dwork, which BTA_pd_pd_alloc already sizes to exactly m0*n0 on that path. New tests: direct BTDA_pd_pd correctness on both the matching-row_perm and gather paths against the diag-decomposition oracle; a no-transient-allocation regression test over pd_pd (both paths), pd_spd, and blockwise spd_pd using the tracked allocator's peak-bytes counter (fails on the previous kernels); and an end-to-end multiply Hessian test with dense operators — the first wsum_hess coverage of multiply's pd/pd dispatch. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HRnTPZsCXK6x2PSGUzanmg --- src/utils/permuted_dense_linalg.c | 89 ++++++---- src/utils/stacked_pd_linalg.c | 64 ++++---- tests/all_tests.c | 4 + tests/utils/test_matmul_dispatchers.h | 152 ++++++++++++++++++ tests/utils/test_permuted_dense.h | 77 +++++++++ .../bivariate_full_dom/test_multiply.h | 24 +++ 6 files changed, 349 insertions(+), 61 deletions(-) diff --git a/src/utils/permuted_dense_linalg.c b/src/utils/permuted_dense_linalg.c index f6a9cd26..0f074275 100644 --- a/src/utils/permuted_dense_linalg.c +++ b/src/utils/permuted_dense_linalg.c @@ -142,8 +142,14 @@ static int int_arrays_equal(const int *a, const int *b, int n) return 1; } -void BTA_pd_pd_fill_values(const permuted_dense *B, const permuted_dense *A, - permuted_dense *C) +/* Shared core for C = B^T @ diag(d) @ A with d == NULL meaning identity. + diag(d) is folded into A's copy into kernel_dwork (d indexed by global row), + so no intermediate is ever allocated: all buffers are pre-sized by + BTA_pd_pd_alloc. Aliasing invariant: B == A implies identical row_perms and + therefore the single-matmul path (a dwork write vs an X read — safe); the + gather path must never run with B == A sharing one kernel_dwork. */ +static void BTA_pd_pd_core(const permuted_dense *B, const double *d, + const permuted_dense *A, permuted_dense *C) { /* C may be empty if there is no overlap in row permutations */ if (C->base.nnz == 0) @@ -154,8 +160,25 @@ void BTA_pd_pd_fill_values(const permuted_dense *B, const permuted_dense *A, /* if B and A have identical row_perms, one matmul suffices */ if (A->m0 == B->m0 && int_arrays_equal(A->row_perm, B->row_perm, A->m0)) { + const double *A_rows = A->X; + if (d != NULL) + { + /* A->kernel_dwork = diag(d) X_A. Pre-sized by BTA_pd_pd_alloc to + MIN(A->m0, B->m0) * A->n0 = A->m0 * A->n0 on this path. */ + for (int ii = 0; ii < A->m0; ii++) + { + double dk = d[A->row_perm[ii]]; + const double *src = A->X + ii * A->n0; + double *dst = A->kernel_dwork + ii * A->n0; + for (int jj = 0; jj < A->n0; jj++) + { + dst[jj] = dk * src[jj]; + } + } + A_rows = A->kernel_dwork; + } cblas_dgemm(CblasRowMajor, CblasTrans, CblasNoTrans, B->n0, A->n0, A->m0, - 1.0, B->X, B->n0, A->X, A->n0, 0.0, C->X, A->n0); + 1.0, B->X, B->n0, A_rows, A->n0, 0.0, C->X, A->n0); return; } @@ -172,49 +195,53 @@ void BTA_pd_pd_fill_values(const permuted_dense *B, const permuted_dense *A, assert(s > 0); // ------------------------------------------------------------------------ - // Gather the matching rows into A->kernel_dwork and B->kernel_dwork. dwork is - // pre-sized by BTA_pd_pd_alloc (one ensure_dwork call per operand at alloc - // time). + // Gather the matching rows into A->kernel_dwork and B->kernel_dwork, + // scaling A's rows by diag(d) on the way when d is given. dwork is + // pre-sized by BTA_pd_pd_alloc (one ensure_dwork call per operand at + // alloc time). // ------------------------------------------------------------------------ for (int k = 0; k < s; k++) { - memcpy(A->kernel_dwork + k * A->n0, A->X + idx_A[k] * A->n0, - A->n0 * sizeof(double)); memcpy(B->kernel_dwork + k * B->n0, B->X + idx_B[k] * B->n0, B->n0 * sizeof(double)); } + if (d == NULL) + { + for (int k = 0; k < s; k++) + { + memcpy(A->kernel_dwork + k * A->n0, A->X + idx_A[k] * A->n0, + A->n0 * sizeof(double)); + } + } + else + { + for (int k = 0; k < s; k++) + { + double dk = d[A->row_perm[idx_A[k]]]; + const double *src = A->X + idx_A[k] * A->n0; + double *dst = A->kernel_dwork + k * A->n0; + for (int jj = 0; jj < A->n0; jj++) + { + dst[jj] = dk * src[jj]; + } + } + } /* matmul on the gathered rows */ cblas_dgemm(CblasRowMajor, CblasTrans, CblasNoTrans, B->n0, A->n0, s, 1.0, B->kernel_dwork, B->n0, A->kernel_dwork, A->n0, 0.0, C->X, A->n0); } +void BTA_pd_pd_fill_values(const permuted_dense *B, const permuted_dense *A, + permuted_dense *C) +{ + BTA_pd_pd_core(B, NULL, A, C); +} + void BTDA_pd_pd_fill_values(const permuted_dense *B, const double *d, const permuted_dense *A, permuted_dense *C) { - /* C may be empty if there is no overlap in row permutations of A and B */ - if (C->base.nnz == 0) - { - return; - } - - /* TODO: must remove this allocation. Very important. The DA - intermediate PD is allocated and freed on every Hessian iteration - — violates the no-alloc-in-fill policy. Fix is to fold diag(d) - directly into BTA_pd_pd_fill_values's gather/dgemm (either via a - shared internal helper that takes an optional d, or by rewriting - this kernel inline using pre-sized A->kernel_dwork). */ - /* C = BT @ (DA) */ - permuted_dense *DA = (permuted_dense *) A->base.copy_sparsity(&A->base); - DA_pd_fill_values(d, A, DA); - /* DA is freshly created via copy_sparsity (no kernel_dwork sized). - BTA_pd_pd_fill_values' slow path (non-identical row_perms) gathers - rows into DA->kernel_dwork — size it to match what - BTA_pd_pd_alloc would have done. */ - int s_max = MIN(DA->m0, B->m0); - permuted_dense_ensure_kernel_dwork(DA, (size_t) s_max * DA->n0); - BTA_pd_pd_fill_values(B, DA, C); - free_matrix(&DA->base); + BTA_pd_pd_core(B, d, A, C); } /* The CSR-flavored kernels for (B=Sparse, A=PD) live in src/old-code; the diff --git a/src/utils/stacked_pd_linalg.c b/src/utils/stacked_pd_linalg.c index dc46743f..3b6cc542 100644 --- a/src/utils/stacked_pd_linalg.c +++ b/src/utils/stacked_pd_linalg.c @@ -277,8 +277,12 @@ matrix *BTA_pd_spd_alloc(const permuted_dense *B, const stacked_pd *A) return C; } -void BTA_pd_spd_fill_values(const permuted_dense *B, const stacked_pd *A, - permuted_dense *C) +/* Shared core for C = B^T @ diag(d) @ A with d == NULL meaning identity. + diag(d) is folded into the per-block gather of A_k's rows (d indexed by + global row), so no intermediate is ever allocated: Bg | Ag | Cg live in + C->kernel_dwork, pre-sized by BTA_pd_spd_alloc. */ +static void BTA_pd_spd_core(const permuted_dense *B, const double *d, + const stacked_pd *A, permuted_dense *C) { /* return if C is empty */ if (C->base.nnz == 0) @@ -319,11 +323,28 @@ void BTA_pd_spd_fill_values(const permuted_dense *B, const stacked_pd *A, memcpy(Bg + p * B->n0, B->X + idx_B[p] * B->n0, B->n0 * sizeof(double)); } - /* Ag = A[idx_A, :] where idx_A contains the overlapping row indices */ - for (int p = 0; p < s; p++) + /* Ag = (diag(d)A)[idx_A, :] where idx_A contains the overlapping row + indices */ + if (d == NULL) + { + for (int p = 0; p < s; p++) + { + memcpy(Ag + p * Ak->n0, Ak->X + idx_A[p] * Ak->n0, + Ak->n0 * sizeof(double)); + } + } + else { - memcpy(Ag + p * Ak->n0, Ak->X + idx_A[p] * Ak->n0, - Ak->n0 * sizeof(double)); + for (int p = 0; p < s; p++) + { + double dk = d[Ak->row_perm[idx_A[p]]]; + const double *src = Ak->X + idx_A[p] * Ak->n0; + double *dst = Ag + p * Ak->n0; + for (int j = 0; j < Ak->n0; j++) + { + dst[j] = dk * src[j]; + } + } } /* Cg = Bg^T @ Ag. Bg is (s, B->n0) row-major (lda = B->n0); we want @@ -350,6 +371,12 @@ void BTA_pd_spd_fill_values(const permuted_dense *B, const stacked_pd *A, } } +void BTA_pd_spd_fill_values(const permuted_dense *B, const stacked_pd *A, + permuted_dense *C) +{ + BTA_pd_spd_core(B, NULL, A, C); +} + // --------------------------------------------------------------------------------- // BTDA_pd_spd: C = B^T @ diag(d) @ A. No separate alloc — output sparsity // is identical to BTA_pd_spd (D doesn't add/remove nonzeros), so callers @@ -359,24 +386,7 @@ void BTA_pd_spd_fill_values(const permuted_dense *B, const stacked_pd *A, void BTDA_pd_spd_fill_values(const permuted_dense *B, const double *d, const stacked_pd *A, permuted_dense *C) { - /* skip if C is empty (no contributing A-blocks) */ - if (C->base.nnz == 0) - { - return; - } - - /* TODO: must remove this allocation. Very important. The DA - intermediate spd is allocated and freed on every Hessian - iteration — violates the no-alloc-in-fill policy. Fix is to - fold diag(d) directly into BTA_pd_spd_fill_values (either via a - shared internal helper that takes an optional d, or by stashing - a persistent DA scratch on C via a new aux slot — mirror of the - transpose_cache pattern). */ - /* C = BT @ (DA) */ - stacked_pd *DA = (stacked_pd *) copy_sparsity_spd_alloc(A); - DA_spd_fill_values(d, A, DA); - BTA_pd_spd_fill_values(B, DA, C); - free_matrix(&DA->base); + BTA_pd_spd_core(B, d, A, C); } // --------------------------------------------------------------------------------- @@ -388,12 +398,6 @@ void BTDA_pd_spd_fill_values(const permuted_dense *B, const double *d, // // B = B1 + B2 + B3 where each Bi is a global permuted dense. Then // C = B^T D A = C1 + C2 + C3 where Ci = Bi^T D A. -// -// TODO: each BTDA_pd_pd_fill_values call internally allocates a DA intermediate -// (see permuted_dense_linalg.c BTDA_pd_pd_fill_values). That means the BTDA -// variant here allocates n_blocks DA temps per fill, all on the hot Hessian -// path. Must be fixed — same future remedy as the per-block BTDA: fold -// diag(d) directly into BTA_pd_pd's gather step. // --------------------------------------------------------------------------------- static matrix *wrapper_BTA_pd_pd(const permuted_dense *Bk, const void *ctx) { diff --git a/tests/all_tests.c b/tests/all_tests.c index 7ec11799..53c76d48 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -327,6 +327,7 @@ int main(void) mu_run_test(test_wsum_hess_vector_mult_log_vector, tests_run); mu_run_test(test_wsum_hess_vector_mult_log_matrix, tests_run); mu_run_test(test_wsum_hess_multiply_linear_ops, tests_run); + mu_run_test(test_wsum_hess_multiply_dense_ops, tests_run); mu_run_test(test_wsum_hess_multiply_sparse_random, tests_run); mu_run_test(test_wsum_hess_multiply_1, tests_run); mu_run_test(test_wsum_hess_multiply_2, tests_run); @@ -444,6 +445,8 @@ int main(void) mu_run_test(test_permuted_dense_BTA_empty_overlap, tests_run); mu_run_test(test_permuted_dense_BTA_partial_overlap, tests_run); mu_run_test(test_permuted_dense_BTDA_decomposition, tests_run); + mu_run_test(test_permuted_dense_BTDA_matching_row_perm, tests_run); + mu_run_test(test_permuted_dense_BTDA_partial_overlap, tests_run); mu_run_test(test_permuted_dense_sum_all_rows, tests_run); mu_run_test(test_permuted_dense_sum_block_of_rows, tests_run); mu_run_test(test_permuted_dense_sum_evenly_spaced_rows, tests_run); @@ -485,6 +488,7 @@ int main(void) mu_run_test(test_BTA_pd_spd_two_blocks_both_kept, tests_run); mu_run_test(test_BTDA_pd_spd_two_blocks_both_kept, tests_run); mu_run_test(test_BTDA_spd_pd_overlapping_cp, tests_run); + mu_run_test(test_BTDA_fill_no_transient_alloc, tests_run); mu_run_test(test_BTA_spd_pd_overlapping_cp, tests_run); mu_run_test(test_BTDA_spd_csc_overlapping_cp, tests_run); mu_run_test(test_BTA_spd_csc_overlapping, tests_run); diff --git a/tests/utils/test_matmul_dispatchers.h b/tests/utils/test_matmul_dispatchers.h index 514ca020..68db075e 100644 --- a/tests/utils/test_matmul_dispatchers.h +++ b/tests/utils/test_matmul_dispatchers.h @@ -12,6 +12,7 @@ #include "utils/sparse_matrix.h" #include "utils/stacked_pd.h" #include "utils/stacked_pd_linalg.h" +#include "utils/tracked_alloc.h" #include "utils/utils.h" #include #include @@ -2434,4 +2435,155 @@ const char *test_BA_pd_kron_spd_no_cache_staleness(void) return 0; } +/* No-alloc-in-fill contract for the BTDA kernels: after alloc and one warm-up + fill, a second fill must not touch the tracked allocator at all. Any + transient sp_malloc inside the fill raises g_peak_bytes above the baseline + even if freed before returning; a permanent one raises g_allocated_bytes. + Covers BTDA_pd_pd (both the matching-row_perm and gather paths), + BTDA_pd_spd, and the blockwise BTDA_spd_pd. */ +static int fill_is_alloc_free(void (*fill)(const void *ctx), const void *ctx) +{ + size_t base = g_allocated_bytes; + g_peak_bytes = base; + fill(ctx); + return g_allocated_bytes == base && g_peak_bytes == base; +} + +typedef struct +{ + const permuted_dense *B; + const double *d; + const void *A; + matrix *C; +} btda_fill_args; + +static void run_BTDA_pd_pd(const void *ctx) +{ + const btda_fill_args *a = (const btda_fill_args *) ctx; + BTDA_pd_pd_fill_values(a->B, a->d, (const permuted_dense *) a->A, + (permuted_dense *) a->C); +} + +static void run_BTDA_pd_spd(const void *ctx) +{ + const btda_fill_args *a = (const btda_fill_args *) ctx; + BTDA_pd_spd_fill_values(a->B, a->d, (const stacked_pd *) a->A, + (permuted_dense *) a->C); +} + +typedef struct +{ + const stacked_pd *B; + const double *d; + const permuted_dense *A; + stacked_pd *C; +} btda_spd_fill_args; + +static void run_BTDA_spd_pd(const void *ctx) +{ + const btda_spd_fill_args *a = (const btda_spd_fill_args *) ctx; + BTDA_spd_pd_fill_values(a->B, a->d, a->A, a->C); +} + +const char *test_BTDA_fill_no_transient_alloc(void) +{ + double d[8] = {2.0, -1.5, 0.5, 1.25, 3.0, -0.5, 1.0, 2.5}; + + /* --- pd_pd, matching row_perms (fast path) --- */ + { + int row_perm[2] = {1, 3}; + int cp_A[2] = {0, 2}; + int cp_B[2] = {1, 3}; + double XA[4] = {1.0, 2.0, 3.0, 4.0}; + double XB[4] = {5.0, 6.0, 7.0, 8.0}; + matrix *A = new_permuted_dense(4, 4, 2, 2, row_perm, cp_A, XA); + matrix *B = new_permuted_dense(4, 4, 2, 2, row_perm, cp_B, XB); + matrix *C = BTA_pd_pd_alloc((permuted_dense *) B, (permuted_dense *) A); + btda_fill_args args = {(permuted_dense *) B, d, A, C}; + run_BTDA_pd_pd(&args); /* warm-up */ + mu_assert("pd_pd fast path allocates in fill", + fill_is_alloc_free(run_BTDA_pd_pd, &args)); + free_matrix(C); + free_matrix(B); + free_matrix(A); + } + + /* --- pd_pd, partial overlap (gather path) --- */ + { + int rp_A[3] = {1, 3, 5}; + int rp_B[3] = {3, 5, 7}; + int cp_A[2] = {0, 2}; + int cp_B[2] = {1, 3}; + double XA[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + double XB[6] = {10.0, 20.0, 30.0, 40.0, 50.0, 60.0}; + matrix *A = new_permuted_dense(8, 4, 3, 2, rp_A, cp_A, XA); + matrix *B = new_permuted_dense(8, 4, 3, 2, rp_B, cp_B, XB); + matrix *C = BTA_pd_pd_alloc((permuted_dense *) B, (permuted_dense *) A); + btda_fill_args args = {(permuted_dense *) B, d, A, C}; + run_BTDA_pd_pd(&args); + mu_assert("pd_pd gather path allocates in fill", + fill_is_alloc_free(run_BTDA_pd_pd, &args)); + free_matrix(C); + free_matrix(B); + free_matrix(A); + } + + /* --- pd_spd (two-block A, both kept) --- */ + { + int A0_rp[2] = {0, 1}; + int A0_cp[2] = {0, 2}; + double A0X[4] = {1, 2, 3, 4}; + matrix *blk0 = new_permuted_dense(4, 3, 2, 2, A0_rp, A0_cp, A0X); + int A1_rp[2] = {2, 3}; + int A1_cp[2] = {1, 2}; + double A1X[4] = {5, 6, 7, 8}; + matrix *blk1 = new_permuted_dense(4, 3, 2, 2, A1_rp, A1_cp, A1X); + permuted_dense *A_blocks[2] = {(permuted_dense *) blk0, + (permuted_dense *) blk1}; + matrix *A_spd = new_stacked_pd(4, 3, 2, A_blocks, NULL, NULL); + int B_rp[3] = {0, 1, 2}; + int B_cp[3] = {1, 3, 4}; + double BX[9] = {1, 2, 3, 4, 5, 6, 7, 8, 9}; + matrix *B = new_permuted_dense(4, 5, 3, 3, B_rp, B_cp, BX); + matrix *C = BTA_pd_spd_alloc((permuted_dense *) B, (stacked_pd *) A_spd); + btda_fill_args args = {(permuted_dense *) B, d, A_spd, C}; + run_BTDA_pd_spd(&args); + mu_assert("pd_spd allocates in fill", + fill_is_alloc_free(run_BTDA_pd_spd, &args)); + free_matrix(C); + free_matrix(B); + free_matrix(A_spd); + } + + /* --- spd_pd (blockwise, overlapping col_perms) --- */ + { + int B0_rp[2] = {0, 1}; + int B0_cp[2] = {0, 2}; + double B0X[4] = {1, 2, 3, 4}; + matrix *blk0 = new_permuted_dense(4, 3, 2, 2, B0_rp, B0_cp, B0X); + int B1_rp[2] = {2, 3}; + int B1_cp[2] = {1, 2}; + double B1X[4] = {5, 6, 7, 8}; + matrix *blk1 = new_permuted_dense(4, 3, 2, 2, B1_rp, B1_cp, B1X); + permuted_dense *B_blocks[2] = {(permuted_dense *) blk0, + (permuted_dense *) blk1}; + matrix *B_spd = new_stacked_pd(4, 3, 2, B_blocks, NULL, NULL); + int A_rp[4] = {0, 1, 2, 3}; + int A_cp[3] = {0, 1, 2}; + double AX[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + matrix *A = new_permuted_dense(4, 3, 4, 3, A_rp, A_cp, AX); + matrix *C = BTA_spd_pd_alloc((stacked_pd *) B_spd, (permuted_dense *) A); + btda_spd_fill_args args = {(stacked_pd *) B_spd, d, (permuted_dense *) A, + (stacked_pd *) C}; + run_BTDA_spd_pd(&args); + mu_assert("spd_pd blockwise allocates in fill", + fill_is_alloc_free(run_BTDA_spd_pd, &args)); + free_matrix(C); + free_matrix(B_spd); + free_matrix(A); + } + + return 0; +} + #endif /* TEST_MATMUL_DISPATCHERS_H */ diff --git a/tests/utils/test_permuted_dense.h b/tests/utils/test_permuted_dense.h index 53a49bac..761708b7 100644 --- a/tests/utils/test_permuted_dense.h +++ b/tests/utils/test_permuted_dense.h @@ -777,6 +777,83 @@ const char *test_permuted_dense_BTDA_decomposition(void) return 0; } +/* Direct BTDA_pd_pd on the matching-row_perm (fast) path: A and B share + row_perm, so the kernel takes the single-dgemm route with diag(d) folded + in. Oracle is the decomposition tmp = diag(d) A; C_ref = B^T tmp. */ +const char *test_permuted_dense_BTDA_matching_row_perm(void) +{ + int row_perm[2] = {1, 3}; + int col_perm_A[2] = {0, 2}; + int col_perm_B[2] = {1, 3}; + double XA[4] = {1.0, 2.0, 3.0, 4.0}; + double XB[4] = {5.0, 6.0, 7.0, 8.0}; + /* d indexed by global row; only d[1], d[3] are read. */ + double d[4] = {-2.0, 0.5, 100.0, 3.0}; + matrix *A_m = new_permuted_dense(4, 4, 2, 2, row_perm, col_perm_A, XA); + matrix *B_m = new_permuted_dense(4, 4, 2, 2, row_perm, col_perm_B, XB); + permuted_dense *A = (permuted_dense *) A_m; + permuted_dense *B = (permuted_dense *) B_m; + + matrix *C_m = BTA_pd_pd_alloc(B, A); + permuted_dense *C = (permuted_dense *) C_m; + BTDA_pd_pd_fill_values(B, d, A, C); + + matrix *tmp_m = A_m->copy_sparsity(A_m); + permuted_dense *tmp = (permuted_dense *) tmp_m; + DA_pd_fill_values(d, A, tmp); + matrix *C_ref_m = BTA_pd_pd_alloc(B, tmp); + permuted_dense *C_ref = (permuted_dense *) C_ref_m; + BTA_pd_pd_fill_values(B, tmp, C_ref); + + mu_assert("values", cmp_double_array(C->X, C_ref->X, 4)); + + free_matrix(C_ref_m); + free_matrix(tmp_m); + free_matrix(C_m); + free_matrix(B_m); + free_matrix(A_m); + return 0; +} + +/* Direct BTDA_pd_pd on the partial-overlap (gather) path: row_perm_A = + [1, 3, 5], row_perm_B = [3, 5, 7], intersection {3, 5}. Same + decomposition oracle as above. */ +const char *test_permuted_dense_BTDA_partial_overlap(void) +{ + int row_perm_A[3] = {1, 3, 5}; + int row_perm_B[3] = {3, 5, 7}; + int col_perm_A[2] = {0, 2}; + int col_perm_B[2] = {1, 3}; + double XA[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + double XB[6] = {10.0, 20.0, 30.0, 40.0, 50.0, 60.0}; + /* d indexed by global row; only d[3], d[5] are read. */ + double d[8] = {9.0, 9.0, 9.0, 2.0, 9.0, -1.5, 9.0, 9.0}; + matrix *A_m = new_permuted_dense(8, 4, 3, 2, row_perm_A, col_perm_A, XA); + matrix *B_m = new_permuted_dense(8, 4, 3, 2, row_perm_B, col_perm_B, XB); + permuted_dense *A = (permuted_dense *) A_m; + permuted_dense *B = (permuted_dense *) B_m; + + matrix *C_m = BTA_pd_pd_alloc(B, A); + permuted_dense *C = (permuted_dense *) C_m; + BTDA_pd_pd_fill_values(B, d, A, C); + + matrix *tmp_m = A_m->copy_sparsity(A_m); + permuted_dense *tmp = (permuted_dense *) tmp_m; + DA_pd_fill_values(d, A, tmp); + matrix *C_ref_m = BTA_pd_pd_alloc(B, tmp); + permuted_dense *C_ref = (permuted_dense *) C_ref_m; + BTA_pd_pd_fill_values(B, tmp, C_ref); + + mu_assert("values", cmp_double_array(C->X, C_ref->X, 4)); + + free_matrix(C_ref_m); + free_matrix(tmp_m); + free_matrix(C_m); + free_matrix(B_m); + free_matrix(A_m); + return 0; +} + /* BTA(CSR_matrix A, PD B): basic correctness against a dense reference. A is (4, 5) CSR_matrix with mixed sparsity; B is (4, 4) PD with row_perm = [1, 3], col_perm = [0, 2], dense block (2, 2). */ diff --git a/tests/wsum_hess/bivariate_full_dom/test_multiply.h b/tests/wsum_hess/bivariate_full_dom/test_multiply.h index b180c3d7..5ac93ae9 100644 --- a/tests/wsum_hess/bivariate_full_dom/test_multiply.h +++ b/tests/wsum_hess/bivariate_full_dom/test_multiply.h @@ -2,6 +2,7 @@ #include "atoms/bivariate_full_dom.h" #include "expr.h" #include "minunit.h" +#include "numerical_diff.h" #include "test_helpers.h" #include #include @@ -108,6 +109,29 @@ const char *test_wsum_hess_multiply_sparse_random(void) return 0; } +/* Hessian for mult(A u, B u) with DENSE 4x3 operators: both child jacobians + are permuted_dense, so multiply's Hessian fill goes through the pd/pd BTDA + dispatch (every other multiply test here uses sparse operators). Verified + against numerical differentiation. */ +const char *test_wsum_hess_multiply_dense_ops(void) +{ + double A[12] = {1.0, 0.0, 2.0, 0.0, 3.0, 0.0, 4.0, 0.0, 5.0, 0.0, 6.0, 0.0}; + double B[12] = {1.0, 0.0, 4.0, 0.0, 2.0, 7.0, 3.0, 0.0, 2.0, 0.0, 4.0, -1.0}; + double u_vals[3] = {0.5, -1.0, 2.0}; + double w[4] = {1.0, -2.0, 0.5, 3.0}; + + expr *x = new_variable(3, 1, 0, 3); + expr *Ax_node = new_left_matmul_dense(NULL, x, 4, 3, A); + expr *Bx_node = new_left_matmul_dense(NULL, x, 4, 3, B); + expr *mult_node = new_elementwise_mult(Ax_node, Bx_node); + + mu_assert("multiply dense-operator wsum_hess failed", + check_wsum_hess(mult_node, u_vals, w, NUMERICAL_DIFF_DEFAULT_H)); + + free_expr(mult_node); + return 0; +} + const char *test_wsum_hess_multiply_linear_ops(void) { /* Test Hessian for mult(Ax, Bx) where A, B are 4x3 linear operators From 92ee9e45fb0bef3721eaf9112f63a20052d0db57 Mon Sep 17 00:00:00 2001 From: dance858 Date: Mon, 31 Aug 2026 20:19:03 -0700 Subject: [PATCH 2/3] Use memcpy + cblas_dscal for the diag-fold scale steps Replace the hand-written scale-copy loops with an unconditional memcpy gather followed by a per-row cblas_dscal pass when d is given, matching the existing DA_pd_fill_values / ATDA_pd_fill_values idiom. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HRnTPZsCXK6x2PSGUzanmg --- src/utils/permuted_dense_linalg.c | 31 ++++++++----------------------- src/utils/stacked_pd_linalg.c | 19 +++++-------------- 2 files changed, 13 insertions(+), 37 deletions(-) diff --git a/src/utils/permuted_dense_linalg.c b/src/utils/permuted_dense_linalg.c index 0f074275..773305ce 100644 --- a/src/utils/permuted_dense_linalg.c +++ b/src/utils/permuted_dense_linalg.c @@ -165,15 +165,11 @@ static void BTA_pd_pd_core(const permuted_dense *B, const double *d, { /* A->kernel_dwork = diag(d) X_A. Pre-sized by BTA_pd_pd_alloc to MIN(A->m0, B->m0) * A->n0 = A->m0 * A->n0 on this path. */ + memcpy(A->kernel_dwork, A->X, (size_t) A->m0 * A->n0 * sizeof(double)); for (int ii = 0; ii < A->m0; ii++) { - double dk = d[A->row_perm[ii]]; - const double *src = A->X + ii * A->n0; - double *dst = A->kernel_dwork + ii * A->n0; - for (int jj = 0; jj < A->n0; jj++) - { - dst[jj] = dk * src[jj]; - } + cblas_dscal(A->n0, d[A->row_perm[ii]], A->kernel_dwork + ii * A->n0, + 1); } A_rows = A->kernel_dwork; } @@ -202,28 +198,17 @@ static void BTA_pd_pd_core(const permuted_dense *B, const double *d, // ------------------------------------------------------------------------ for (int k = 0; k < s; k++) { + memcpy(A->kernel_dwork + k * A->n0, A->X + idx_A[k] * A->n0, + A->n0 * sizeof(double)); memcpy(B->kernel_dwork + k * B->n0, B->X + idx_B[k] * B->n0, B->n0 * sizeof(double)); } - if (d == NULL) - { - for (int k = 0; k < s; k++) - { - memcpy(A->kernel_dwork + k * A->n0, A->X + idx_A[k] * A->n0, - A->n0 * sizeof(double)); - } - } - else + if (d != NULL) { for (int k = 0; k < s; k++) { - double dk = d[A->row_perm[idx_A[k]]]; - const double *src = A->X + idx_A[k] * A->n0; - double *dst = A->kernel_dwork + k * A->n0; - for (int jj = 0; jj < A->n0; jj++) - { - dst[jj] = dk * src[jj]; - } + cblas_dscal(A->n0, d[A->row_perm[idx_A[k]]], A->kernel_dwork + k * A->n0, + 1); } } diff --git a/src/utils/stacked_pd_linalg.c b/src/utils/stacked_pd_linalg.c index 3b6cc542..4efadf87 100644 --- a/src/utils/stacked_pd_linalg.c +++ b/src/utils/stacked_pd_linalg.c @@ -325,25 +325,16 @@ static void BTA_pd_spd_core(const permuted_dense *B, const double *d, /* Ag = (diag(d)A)[idx_A, :] where idx_A contains the overlapping row indices */ - if (d == NULL) + for (int p = 0; p < s; p++) { - for (int p = 0; p < s; p++) - { - memcpy(Ag + p * Ak->n0, Ak->X + idx_A[p] * Ak->n0, - Ak->n0 * sizeof(double)); - } + memcpy(Ag + p * Ak->n0, Ak->X + idx_A[p] * Ak->n0, + Ak->n0 * sizeof(double)); } - else + if (d != NULL) { for (int p = 0; p < s; p++) { - double dk = d[Ak->row_perm[idx_A[p]]]; - const double *src = Ak->X + idx_A[p] * Ak->n0; - double *dst = Ag + p * Ak->n0; - for (int j = 0; j < Ak->n0; j++) - { - dst[j] = dk * src[j]; - } + cblas_dscal(Ak->n0, d[Ak->row_perm[idx_A[p]]], Ag + p * Ak->n0, 1); } } From d8acd55a932c27ec4c50be1a18fe4b5689af3d33 Mon Sep 17 00:00:00 2001 From: dance858 Date: Mon, 31 Aug 2026 20:20:21 -0700 Subject: [PATCH 3/3] clang-format permuted_dense_linalg.c Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HRnTPZsCXK6x2PSGUzanmg --- src/utils/permuted_dense_linalg.c | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/utils/permuted_dense_linalg.c b/src/utils/permuted_dense_linalg.c index 773305ce..6ea8cac7 100644 --- a/src/utils/permuted_dense_linalg.c +++ b/src/utils/permuted_dense_linalg.c @@ -571,10 +571,10 @@ void BTDA_csc_pd_fill_values(const CSC_matrix *B, const double *d, #if defined(__GNUC__) || defined(__clang__) __attribute__((unused)) #endif -static void -BTDA_csc_pd_fill_values_via_transpose_dead(const CSC_matrix *B, const double *d, - const permuted_dense *A, - permuted_dense *C) +static void BTDA_csc_pd_fill_values_via_transpose_dead(const CSC_matrix *B, + const double *d, + const permuted_dense *A, + permuted_dense *C) { if (C->base.nnz == 0) {