unify code in stacked_pd

Author: dance858

include/utils/permuted_dense_linalg.h

@@ -81,4 +81,14 @@ matrix *BTA_csc_pd_alloc(const CSC_matrix *B, const permuted_dense *A);
 void BTDA_csc_pd_fill_values(const CSC_matrix *B, const double *d,
                              const permuted_dense *A, permuted_dense *C);
 
+/* Allocate a new PD C with the same sparsity (m, n, m0, n0, row_perm,
+   col_perm) as A; values uninitialized. */
+matrix *copy_sparsity_pd_alloc(const permuted_dense *A);
+
+/* Allocate C = A^T (transpose). */
+matrix *transpose_pd_alloc(const permuted_dense *A);
+
+/* Fill values of C = A^T. */
+void transpose_pd_fill_values(const permuted_dense *A, permuted_dense *C);
+
 #endif /* PERMUTED_DENSE_LINALG_H */

include/utils/stacked_pd.h

@@ -114,4 +114,10 @@ matrix *coalesce_spd_alloc_unchecked(const stacked_pd *src);
    (block count, row_perms, col_perms) as the one passed now. */
 void coalesce_spd_fill_values(const stacked_pd *src, stacked_pd *out);
 
+/* Same scatter as coalesce_spd_fill_values but with += instead of =.
+   Use when src blocks share cells targeting the same out block (e.g.
+   the B_k^T B_k summands aggregated by ATDA_spd_fill_values). Caller is
+   responsible for zeroing out->base.x first. */
+void coalesce_spd_fill_values_accumulate(const stacked_pd *src, stacked_pd *out);
+
 #endif /* STACKED_PD_H */

src/utils/permuted_dense.c

@@ -59,9 +59,7 @@ static void permuted_dense_refresh_csc_values(matrix *self)
 /* Vtable adapters — each delegates to the existing permuted_dense_* kernel. */
 static matrix *permuted_dense_vtable_copy_sparsity(const matrix *self)
 {
-    const permuted_dense *pd = (const permuted_dense *) self;
-    return new_permuted_dense(pd->base.m, pd->base.n, pd->m0, pd->n0, pd->row_perm,
-                              pd->col_perm, NULL);
+    return copy_sparsity_pd_alloc((const permuted_dense *) self);
 }
 
 static void permuted_dense_vtable_DA_fill_values(const double *d, const matrix *self,
@@ -99,28 +97,13 @@ static CSR_matrix *permuted_dense_to_csr(matrix *self)
 
 static matrix *permuted_dense_vtable_transpose_alloc(const matrix *self)
 {
-    const permuted_dense *pd = (const permuted_dense *) self;
-    /* Swap (m, n), (m0, n0), and (row_perm, col_perm). The constructor
-       asserts strict increase of both perms, which holds by construction. */
-    return new_permuted_dense(pd->base.n, pd->base.m, pd->n0, pd->m0, pd->col_perm,
-                              pd->row_perm, NULL);
+    return transpose_pd_alloc((const permuted_dense *) self);
 }
 
 static void permuted_dense_vtable_transpose_fill_values(const matrix *self,
                                                         matrix *out)
 {
-    const permuted_dense *pd_in = (const permuted_dense *) self;
-    permuted_dense *pd_out = (permuted_dense *) out;
-    int m0 = pd_in->m0;
-    int n0 = pd_in->n0;
-    /* pd_out has shape (n0, m0); transpose pd_in->X into pd_out->X. */
-    for (int ii = 0; ii < m0; ii++)
-    {
-        for (int jj = 0; jj < n0; jj++)
-        {
-            pd_out->X[jj * m0 + ii] = pd_in->X[ii * n0 + jj];
-        }
-    }
+    transpose_pd_fill_values((const permuted_dense *) self, (permuted_dense *) out);
 }
 
 static matrix *permuted_dense_vtable_index_alloc(matrix *self, const int *indices,

src/utils/permuted_dense_linalg.c

@@ -26,6 +26,34 @@
 #include <stdlib.h>
 #include <string.h>
 
+matrix *copy_sparsity_pd_alloc(const permuted_dense *A)
+{
+    return new_permuted_dense(A->base.m, A->base.n, A->m0, A->n0, A->row_perm,
+                              A->col_perm, NULL);
+}
+
+matrix *transpose_pd_alloc(const permuted_dense *A)
+{
+    /* Swap (m, n), (m0, n0), and (row_perm, col_perm). The constructor
+       asserts strict increase of both perms, which holds by construction. */
+    return new_permuted_dense(A->base.n, A->base.m, A->n0, A->m0, A->col_perm,
+                              A->row_perm, NULL);
+}
+
+void transpose_pd_fill_values(const permuted_dense *A, permuted_dense *out)
+{
+    int m0 = A->m0;
+    int n0 = A->n0;
+    /* out has shape (n0, m0); transpose A->X into out->X. */
+    for (int ii = 0; ii < m0; ii++)
+    {
+        for (int jj = 0; jj < n0; jj++)
+        {
+            out->X[jj * m0 + ii] = A->X[ii * n0 + jj];
+        }
+    }
+}
+
 void DA_pd_fill_values(const double *d, const permuted_dense *A, permuted_dense *C)
 {
     int m0 = A->m0;

src/utils/stacked_pd.c

@@ -600,34 +600,69 @@ matrix *coalesce_spd_alloc_unchecked(const stacked_pd *src)
     return coalesce_spd_alloc_impl(src, false);
 }
 
-void coalesce_spd_fill_values(const stacked_pd *src, stacked_pd *out)
+/* Scatter src's blocks into out's blocks, following out's src_block_idx_*
+   map. If accumulate is true, += into out (caller must zero out->base.x
+   first); otherwise = (precondition: source cells are pairwise disjoint
+   across blocks targeting the same out block). `accumulate` is a
+   compile-time constant at every call site, so the inner-loop branch is
+   folded away when this is inlined into the public wrappers. */
+static inline void coalesce_spd_scatter(const stacked_pd *src, stacked_pd *out,
+                                        bool accumulate)
 {
+    /* for each block out_k of out */
     for (int k = 0; k < out->n_blocks; k++)
     {
         permuted_dense *out_k = out->blocks[k];
         int s_lo = out->src_block_idx_p[k];
         int s_hi = out->src_block_idx_p[k + 1];
+
+        /* for each block src_k in src contributing to out_k */
         for (int t = s_lo; t < s_hi; t++)
         {
             const permuted_dense *src_k = src->blocks[out->src_block_idx[t]];
+
+            /* for each row in src_k */
             for (int i = 0; i < src_k->m0; i++)
             {
                 int out_i = out_k->row_inv[src_k->row_perm[i]];
                 if (out_i < 0)
                 {
-                    continue; /* row not in this output PD */
+                    /* this row in src_k does not contribute to out_k */
+                    continue;
                 }
+
+                double *out_row = out_k->X + out_i * out_k->n0;
+                const double *src_row = src_k->X + i * src_k->n0;
+
+                /* place each col in src_row at its correct position in out_row */
                 for (int j = 0; j < src_k->n0; j++)
                 {
                     int out_j = out_k->col_inv[src_k->col_perm[j]];
-                    out_k->X[out_i * out_k->n0 + out_j] =
-                        src_k->X[i * src_k->n0 + j];
+                    assert(out_j >= 0);
+                    if (accumulate)
+                    {
+                        out_row[out_j] += src_row[j];
+                    }
+                    else
+                    {
+                        out_row[out_j] = src_row[j];
+                    }
                 }
             }
         }
     }
 }
 
+void coalesce_spd_fill_values(const stacked_pd *src, stacked_pd *out)
+{
+    coalesce_spd_scatter(src, out, false);
+}
+
+void coalesce_spd_fill_values_accumulate(const stacked_pd *src, stacked_pd *out)
+{
+    coalesce_spd_scatter(src, out, true);
+}
+
 #ifndef NDEBUG
 /* Pairwise verify that block row permutations are disjoint via a
    two-pointer merge over the sorted arrays. */

src/utils/stacked_pd_linalg.c

@@ -120,8 +120,7 @@ matrix *copy_sparsity_spd_alloc(const stacked_pd *A)
         (permuted_dense **) SP_MALLOC(n_blocks * sizeof(permuted_dense *));
     for (int k = 0; k < n_blocks; k++)
     {
-        const matrix *Ak = (matrix *) A->blocks[k];
-        C_blocks[k] = (permuted_dense *) Ak->copy_sparsity(Ak);
+        C_blocks[k] = (permuted_dense *) copy_sparsity_pd_alloc(A->blocks[k]);
     }
 
     matrix *C = new_stacked_pd(A->base.m, A->base.n, n_blocks, C_blocks, NULL, NULL);
@@ -137,9 +136,14 @@ void DA_spd_fill_values(const double *d, const stacked_pd *A, stacked_pd *C)
     }
 }
 
-/* ------------------------------------------------------------------ */
-/* Transpose                                                           */
-/* ------------------------------------------------------------------ */
+// ----------------------------------------------------------------------------------
+// Transpose C = A^T where A is stacked_pd. Blockwise logic followed by coalescing:
+// we first transpose each block independently, yielding a raw spd with the same
+// block count and swapped row/col perms as the final output but potentially
+// overlapping row perms. We then coalesce that raw spd to get the final output with
+// disjoint row perms. We keep the raw spd around on `out->work` for reuse by
+// `transpose_spd_fill_values` since its per-block transposes are also needed there.
+// ----------------------------------------------------------------------------------
 
 matrix *transpose_spd_alloc(const stacked_pd *A)
 {
@@ -148,8 +152,7 @@ matrix *transpose_spd_alloc(const stacked_pd *A)
         (permuted_dense **) SP_MALLOC(n_blocks * sizeof(permuted_dense *));
     for (int k = 0; k < n_blocks; k++)
     {
-        matrix *blk = (matrix *) A->blocks[k];
-        AT_blocks[k] = (permuted_dense *) blk->transpose_alloc(blk);
+        AT_blocks[k] = (permuted_dense *) transpose_pd_alloc(A->blocks[k]);
     }
 
     /* Raw spd: dimensions swapped, identity src_block_idx_*; rows may
@@ -168,16 +171,17 @@ void transpose_spd_fill_values(const stacked_pd *A, stacked_pd *C)
     stacked_pd *raw = C->work;
     for (int k = 0; k < A->n_blocks; k++)
     {
-        matrix *A_blk = (matrix *) A->blocks[k];
-        matrix *raw_blk = (matrix *) raw->blocks[k];
-        A_blk->transpose_fill_values(A_blk, raw_blk);
+        transpose_pd_fill_values(A->blocks[k], raw->blocks[k]);
     }
     coalesce_spd_fill_values(raw, C);
 }
 
-/* ------------------------------------------------------------------ */
-/* ATA / ATDA                                                          */
-/* ------------------------------------------------------------------ */
+// ------------------------------------------------------------------------------------
+// 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
+// overlap if Ai and Aj have overlapping column entries. We therefore first compute
+// Ai^T Di Ai and then take care of the accumulation.
+// ------------------------------------------------------------------------------------
 
 matrix *ATA_spd_alloc(const stacked_pd *A)
 {
@@ -207,49 +211,6 @@ matrix *ATA_spd_alloc(const stacked_pd *A)
     return out;
 }
 
-/* Scatter + add each scratch PD into its destination block in C. Multiple scratch
-   PDs may target the same Ck (their cells overlap by design — that's why we
-   accumulate). Caller is responsible for zeroing C->base.x first. */
-static inline void accumulate_scratch_blocks(const stacked_pd *scratch,
-                                             stacked_pd *C)
-{
-    /* for each block Ck of C */
-    for (int k = 0; k < C->n_blocks; k++)
-    {
-        permuted_dense *Ck = C->blocks[k];
-        int s_lo = C->src_block_idx_p[k];
-        int s_hi = C->src_block_idx_p[k + 1];
-
-        /* for each block Eq = Ai^T D Ai in scratch contributing to block Ck in C */
-        for (int qq = s_lo; qq < s_hi; qq++)
-        {
-            const permuted_dense *Eq = scratch->blocks[C->src_block_idx[qq]];
-
-            /* for each row in Eq */
-            for (int i = 0; i < Eq->m0; i++)
-            {
-                int Ck_row_idx = Ck->row_inv[Eq->row_perm[i]];
-                if (Ck_row_idx < 0)
-                {
-                    /* this row in Eq does not contribute to Ck */
-                    continue;
-                }
-
-                double *Ck_row = Ck->X + Ck_row_idx * Ck->n0;
-                double *Eq_row = Eq->X + i * Eq->n0;
-
-                /* place each col in Eq_row at its correct position in Ck_row */
-                for (int j = 0; j < Eq->n0; j++)
-                {
-                    int out_j = Ck->col_inv[Eq->col_perm[j]];
-                    assert(out_j >= 0);
-                    Ck_row[out_j] += Eq_row[j];
-                }
-            }
-        }
-    }
-}
-
 void ATDA_spd_fill_values(const stacked_pd *A, const double *d, stacked_pd *C)
 {
     /* Let A = [A1; A2; A3] where Ai has n columns (the same number as A). Then
@@ -270,7 +231,7 @@ void ATDA_spd_fill_values(const stacked_pd *A, const double *d, stacked_pd *C)
     memset(C->base.x, 0, C->base.nnz * sizeof(double));
 
     /* scatter + add each Ai^T Di Ai into its destination C block. */
-    accumulate_scratch_blocks(scratch, C);
+    coalesce_spd_fill_values_accumulate(scratch, C);
 }
 
 /* ------------------------------------------------------------------ */