Better BTA without allocation

src/utils/permuted_dense_linalg.c

@@ -144,8 +144,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)
@@ -156,8 +162,21 @@ 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. */
+            memcpy(A->kernel_dwork, A->X, (size_t) A->m0 * A->n0 * sizeof(double));
+            for (int ii = 0; ii < A->m0; ii++)
+            {
+                cblas_dscal(A->n0, d[A->row_perm[ii]], A->kernel_dwork + ii * A->n0,
+                            1);
+            }
+            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;
     }
 
@@ -174,9 +193,10 @@ 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++)
     {
@@ -185,38 +205,30 @@ void BTA_pd_pd_fill_values(const permuted_dense *B, const permuted_dense *A,
         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++)
+        {
+            cblas_dscal(A->n0, d[A->row_perm[idx_A[k]]], A->kernel_dwork + k * A->n0,
+                        1);
+        }
+    }
 
     /* 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

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,12 +323,20 @@ 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 */
+        /* Ag = (diag(d)A)[idx_A, :] where idx_A contains the overlapping row
+           indices */
         for (int p = 0; p < s; p++)
         {
             memcpy(Ag + p * Ak->n0, Ak->X + idx_A[p] * Ak->n0,
                    Ak->n0 * sizeof(double));
         }
+        if (d != NULL)
+        {
+            for (int p = 0; p < s; p++)
+            {
+                cblas_dscal(Ak->n0, d[Ak->row_perm[idx_A[p]]], Ag + p * Ak->n0, 1);
+            }
+        }
 
         /* Cg = Bg^T @ Ag. Bg is (s, B->n0) row-major (lda = B->n0); we want
            output Cg = (B->n0, Ak->n0). */
@@ -350,6 +362,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 +377,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 +389,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)
 {

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,9 @@ 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);
+#ifdef SP_TRACK_MEMORY
+    mu_run_test(test_BTDA_fill_no_transient_alloc, tests_run);
+#endif
     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);

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 <stdlib.h>
 #include <string.h>
@@ -2436,6 +2437,163 @@ const char *test_BA_pd_kron_spd_no_cache_staleness(void)
     return 0;
 }
 
+/* Only compiled with -DSP_TRACK_MEMORY=ON: the test reads the tracked
+   allocator counters, which do not exist in a default build. */
+#ifdef SP_TRACK_MEMORY
+
+/* 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 /* SP_TRACK_MEMORY */
+
 /* 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