unify BTA_pd_spd like kernels

Author: dance858

include/utils/stacked_pd_linalg.h

@@ -59,6 +59,15 @@ matrix *BTA_spd_pd_alloc(const stacked_pd *B, const permuted_dense *A);
 void BTDA_spd_pd_fill_values(const stacked_pd *B, const double *d,
                              const permuted_dense *A, stacked_pd *C);
 
+/* Allocate sparsity for C = B^T @ A where B is stacked_pd and A is CSC.
+   The output C is stacked_pd. */
+matrix *BTA_spd_csc_alloc(const stacked_pd *B, const CSC_matrix *A);
+
+/* Fill values of C = B^T @ diag(d) @ A where B is stacked_pd, A is CSC,
+   d is a global vector of length B->m, and C is stacked_pd. */
+void BTDA_spd_csc_fill_values(const stacked_pd *B, const double *d,
+                              const CSC_matrix *A, stacked_pd *C);
+
 /* Allocate sparsity for a new stacked_pd C = A^T A. */
 matrix *ATA_spd_alloc(const stacked_pd *A);
 

src/utils/stacked_pd_linalg.c

@@ -94,6 +94,60 @@ void transpose_spd_fill_values(const stacked_pd *A, stacked_pd *C)
     coalesce_spd_fill_values(raw, C);
 }
 
+// ====================================================================================
+// Shared "spd blockwise + coalesce-accumulate" skeleton.
+//
+// Used by kernels of the shape C = f(B, X) where B is a stacked_pd and each
+// per-block partial f(B_k, X) is a PD. The partials' row_perms (and cells)
+// may overlap across k, so we collect them via new_stacked_pd_unchecked,
+// coalesce with the _unchecked variant, and use the accumulating scatter
+// at fill time. The raw spd is stashed on out->pre_coalesce so fill can
+// refresh per-block values without reallocating structure.
+//
+// Three kernel pairs build on this skeleton: BTA_spd_pd / BTDA_spd_pd,
+// BTA_spd_csc / BTDA_spd_csc, and ATA_spd / ATDA_spd. Each variant supplies
+// per-block alloc and fill kernels via callback + context.
+// ====================================================================================
+typedef matrix *(*spd_per_block_alloc_fn)(const permuted_dense *Bk, const void *ctx);
+
+typedef void (*spd_per_block_fill_fn)(const permuted_dense *Bk, const double *d,
+                                      const void *ctx, permuted_dense *Ck);
+
+static matrix *spd_blockwise_alloc_coalesce(const stacked_pd *B, int Cn,
+                                            spd_per_block_alloc_fn op,
+                                            const void *ctx)
+{
+    int n_blocks = B->n_blocks;
+    permuted_dense **partials =
+        (permuted_dense **) SP_MALLOC(n_blocks * sizeof(permuted_dense *));
+    for (int k = 0; k < n_blocks; k++)
+    {
+        partials[k] = (permuted_dense *) op(B->blocks[k], ctx);
+    }
+    matrix *raw =
+        new_stacked_pd_unchecked(B->base.n, Cn, n_blocks, partials, NULL, NULL);
+    free(partials);
+    matrix *C = coalesce_spd_alloc_unchecked((stacked_pd *) raw);
+    ((stacked_pd *) C)->pre_coalesce = (stacked_pd *) raw;
+    return C;
+}
+
+static void spd_blockwise_fill_coalesce_accumulate(const stacked_pd *B,
+                                                   const double *d, const void *ctx,
+                                                   stacked_pd *C,
+                                                   spd_per_block_fill_fn op)
+{
+    if (C->base.nnz == 0) return;
+
+    stacked_pd *raw = C->pre_coalesce;
+    for (int k = 0; k < B->n_blocks; k++)
+    {
+        op(B->blocks[k], d, ctx, raw->blocks[k]);
+    }
+    memset(C->base.x, 0, C->base.nnz * sizeof(double));
+    coalesce_spd_fill_values_accumulate(raw, C);
+}
+
 // ------------------------------------------------------------------------------------
 // C = ATDA for stacked_pd A. Let A = [A1; A2; A3] where Ai has n columns (the same
 // number as A). Then ATDA = A1^T D1 A1 + A2^T D2 A2 + A3^T D3 A3. Term i and j
@@ -108,50 +162,33 @@ share cells.The output groups cols of A by signature sig_C(c) = {k: c ∈ C_k};
    and col_perm = ⋃ C_k for k in the signature. No structural zeros.
    The output's `pre_coalesce` slot holds per-source scratch PDs (one symmetric
    PD per source block, row_perm = col_perm = C_k) that
-   ATDA_spd_fill_values writes into. */
-matrix *ATA_spd_alloc(const stacked_pd *A)
+   ATDA_spd_fill_values writes into.
+
+   Note on per-block kernel choice: we use ATA_pd_alloc (rather than
+   new_permuted_dense directly) because it also pre-sizes
+   A->blocks[k]->kernel_dwork — that scratch buffer is what
+   ATDA_pd_fill_values reads from during the fill phase. */
+static matrix *wrapper_ATA_pd(const permuted_dense *Ak, const void *ctx)
 {
-    int n = A->base.n;
-    int n_blocks = A->n_blocks;
+    (void) ctx;
+    return ATA_pd_alloc(Ak);
+}
 
-    /* Pseudo-spd: per source block k, a square PD with
-       row_perm = col_perm = C_k and uninitialized X. We use the
-       PD-level ATA_pd_alloc here (rather than new_permuted_dense
-       directly) because it also pre-sizes A->blocks[k]->kernel_dwork — that
-       scratch buffer is what ATDA_pd_fill_values reads from during the
-       fill phase. */
-    permuted_dense **scratch_blocks =
-        (permuted_dense **) SP_MALLOC(n_blocks * sizeof(permuted_dense *));
-    for (int k = 0; k < n_blocks; k++)
-    {
-        scratch_blocks[k] = (permuted_dense *) ATA_pd_alloc(A->blocks[k]);
-    }
-    matrix *scratch =
-        new_stacked_pd_unchecked(n, n, n_blocks, scratch_blocks, NULL, NULL);
-    free(scratch_blocks);
+static void wrapper_ATDA_pd(const permuted_dense *Ak, const double *d,
+                            const void *ctx, permuted_dense *Ck)
+{
+    (void) ctx;
+    ATDA_pd_fill_values(Ak, d, Ck);
+}
 
-    /* prepare coalescing of overlapping scratch blocks */
-    matrix *C = coalesce_spd_alloc_unchecked((stacked_pd *) scratch);
-    ((stacked_pd *) C)->pre_coalesce = (stacked_pd *) scratch;
-    return C;
+matrix *ATA_spd_alloc(const stacked_pd *A)
+{
+    return spd_blockwise_alloc_coalesce(A, A->base.n, wrapper_ATA_pd, NULL);
 }
 
 void ATDA_spd_fill_values(const stacked_pd *A, const double *d, stacked_pd *C)
 {
-    stacked_pd *scratch = C->pre_coalesce;
-
-    /* compute Ai^T @ Di @ Ai into the scratch PDs. */
-    for (int k = 0; k < A->n_blocks; k++)
-    {
-        ATDA_pd_fill_values(A->blocks[k], d, scratch->blocks[k]);
-    }
-
-    /* zero C (one memset is sufficient since, by design, the values of
-       different blocks are stored consecutively in memory). */
-    memset(C->base.x, 0, C->base.nnz * sizeof(double));
-
-    /* scatter + add each Ai^T Di Ai into its destination C block. */
-    coalesce_spd_fill_values_accumulate(scratch, C);
+    spd_blockwise_fill_coalesce_accumulate(A, d, NULL, C, wrapper_ATDA_pd);
 }
 
 // ---------------------------------------------------------------------------------
@@ -344,52 +381,62 @@ 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.
 // ---------------------------------------------------------------------------------
-matrix *BTA_spd_pd_alloc(const stacked_pd *B, const permuted_dense *A)
+static matrix *wrapper_BTA_pd_pd(const permuted_dense *Bk, const void *ctx)
 {
-    int n_blocks = B->n_blocks;
-
-    /* Ck = B_k^T @ A */
-    permuted_dense **Cks =
-        (permuted_dense **) SP_MALLOC(n_blocks * sizeof(permuted_dense *));
-    for (int k = 0; k < n_blocks; k++)
-    {
-        Cks[k] = (permuted_dense *) BTA_pd_pd_alloc(B->blocks[k], A);
-    }
+    return BTA_pd_pd_alloc(Bk, (const permuted_dense *) ctx);
+}
 
-    matrix *raw =
-        new_stacked_pd_unchecked(B->base.n, A->base.n, n_blocks, Cks, NULL, NULL);
-    free(Cks);
+static void wrapper_BTDA_pd_pd(const permuted_dense *Bk, const double *d,
+                               const void *ctx, permuted_dense *Ck)
+{
+    BTDA_pd_pd_fill_values(Bk, d, (const permuted_dense *) ctx, Ck);
+}
 
-    /* coalesce */
-    matrix *C = coalesce_spd_alloc_unchecked((stacked_pd *) raw);
-    ((stacked_pd *) C)->pre_coalesce = (stacked_pd *) raw;
-    return C;
+matrix *BTA_spd_pd_alloc(const stacked_pd *B, const permuted_dense *A)
+{
+    return spd_blockwise_alloc_coalesce(B, A->base.n, wrapper_BTA_pd_pd, A);
 }
 
 void BTDA_spd_pd_fill_values(const stacked_pd *B, const double *d,
                              const permuted_dense *A, stacked_pd *C)
 {
-    if (C->base.nnz == 0)
-    {
-        return;
-    }
+    spd_blockwise_fill_coalesce_accumulate(B, d, A, C, wrapper_BTDA_pd_pd);
+}
 
-    /* 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 this function 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. */
-    stacked_pd *raw = C->pre_coalesce;
-    for (int k = 0; k < B->n_blocks; k++)
-    {
-        BTDA_pd_pd_fill_values(B->blocks[k], d, A, raw->blocks[k]);
-    }
+// ---------------------------------------------------------------------------------
+// BTA_spd_csc / BTDA_spd_csc: C = B^T @ (diag(d) @) A where B is stacked_pd,
+// A is CSC. Output C is stacked_pd. Same two-stage skeleton as BTA_spd_pd /
+// BTDA_spd_pd, just with BTA_pd_csc / BTDA_pd_csc as the per-block kernel.
+// Unlike the pd-on-the-right path, BTDA_pd_csc reuses the per-block
+// B_k->kernel_dwork sized by BTA_pd_csc_alloc — no temp DA allocation.
+// ---------------------------------------------------------------------------------
+static matrix *wrapper_BTA_pd_csc(const permuted_dense *Bk, const void *ctx)
+{
+    return BTA_pd_csc_alloc(Bk, (const CSC_matrix *) ctx);
+}
 
-    /* zero C before accumulating (matches ATDA_spd_fill_values). */
-    memset(C->base.x, 0, C->base.nnz * sizeof(double));
-    coalesce_spd_fill_values_accumulate(raw, C);
+static void wrapper_BTDA_pd_csc(const permuted_dense *Bk, const double *d,
+                                const void *ctx, permuted_dense *Ck)
+{
+    BTDA_pd_csc_fill_values(Bk, d, (const CSC_matrix *) ctx, Ck);
+}
+
+matrix *BTA_spd_csc_alloc(const stacked_pd *B, const CSC_matrix *A)
+{
+    return spd_blockwise_alloc_coalesce(B, A->n, wrapper_BTA_pd_csc, A);
+}
+
+void BTDA_spd_csc_fill_values(const stacked_pd *B, const double *d,
+                              const CSC_matrix *A, stacked_pd *C)
+{
+    spd_blockwise_fill_coalesce_accumulate(B, d, A, C, wrapper_BTDA_pd_csc);
 }
 
 // ---------------------------------------------------------------------------------

tests/all_tests.c

@@ -439,6 +439,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_spd_csc_overlapping_cp, 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_matrix_BTA.h

@@ -560,4 +560,79 @@ const char *test_BTDA_spd_pd_overlapping_cp(void)
     return 0;
 }
 
+/* Primitive BTDA_spd_csc kernel (ATA-style direct: per-block BTDA_pd_csc
+   + accumulating coalesce). C = B^T @ diag(d) @ A with B a 2-block
+   stacked_pd whose col_perms share column 2, and A a sparse_matrix
+   wrapping a small CSR. Reference is the production dispatcher with B
+   flattened to a sparse_matrix; outputs differ in storage (stacked_pd
+   vs sparse_matrix), so compare via to_csr. */
+const char *test_BTDA_spd_csc_overlapping_cp(void)
+{
+    /* B: 4x3 spd, two blocks with overlapping col_perms (share col 2).
+       Same layout as test_BTDA_spd_pd_overlapping_cp.
+       blk0: rows {0,1}, cols {0,2}, X = [[1,2],[3,4]]
+       blk1: rows {2,3}, cols {1,2}, X = [[5,6],[7,8]]                     */
+    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);
+
+    /* A: 4x3 sparse_matrix wrapping a CSR with arbitrary nonzeros.
+       Row 0: col 0 = 1, col 2 = 2
+       Row 1: col 1 = 3
+       Row 2: col 0 = 4, col 1 = 5, col 2 = 6
+       Row 3: col 2 = 7
+       p = [0, 2, 3, 6, 7], i = [0,2, 1, 0,1,2, 2], x = [1,2, 3, 4,5,6, 7]. */
+    CSR_matrix *A_csr = new_CSR_matrix(4, 3, 7);
+    int Ap[5] = {0, 2, 3, 6, 7};
+    int Ai[7] = {0, 2, 1, 0, 1, 2, 2};
+    double Ax[7] = {1, 2, 3, 4, 5, 6, 7};
+    memcpy(A_csr->p, Ap, sizeof(Ap));
+    memcpy(A_csr->i, Ai, sizeof(Ai));
+    memcpy(A_csr->x, Ax, sizeof(Ax));
+    matrix *A_sm = new_sparse_matrix(A_csr);
+
+    /* Non-trivial d so BTDA != BTA. */
+    double d[4] = {2.0, -1.5, 0.5, 1.25};
+
+    /* Route 1: our new BTDA_spd_csc_fill_values. Need A's csc_cache. */
+    sparse_matrix_ensure_csc_cache((sparse_matrix *) A_sm);
+    A_sm->refresh_csc_values(A_sm);
+    matrix *C_ours =
+        BTA_spd_csc_alloc((stacked_pd *) B_spd, ((sparse_matrix *) A_sm)->csc_cache);
+    BTDA_spd_csc_fill_values((stacked_pd *) B_spd, d,
+                             ((sparse_matrix *) A_sm)->csc_cache,
+                             (stacked_pd *) C_ours);
+
+    /* Route 2: dispatcher with B flattened to sparse_matrix. Note
+       BTA_matrices_alloc(A, B) computes B^T @ A, so A goes first. */
+    matrix *B_sparse = spd_to_sparse_matrix_copy(B_spd);
+    matrix *C_ref = BTA_matrices_alloc(A_sm, B_sparse);
+    B_sparse->refresh_csc_values(B_sparse);
+    BTDA_matrices_fill_values(A_sm, d, B_sparse, C_ref);
+
+    /* C_ours is stacked_pd, C_ref is sparse_matrix — compare via to_csr. */
+    CSR_matrix *csr_ours = C_ours->to_csr(C_ours);
+    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(B_sparse);
+    free_matrix(C_ours);
+    free_matrix(A_sm);
+    free_matrix(B_spd);
+    return 0;
+}
+
 #endif /* TEST_MATRIX_BTA_H */