Make block_left_multiply_fill_values linear in the multiply-adds

include/utils/linalg_sparse_matmuls.h

@@ -37,6 +37,12 @@ CSC_matrix *block_left_multiply_fill_sparsity(const CSR_matrix *A,
 void block_left_multiply_fill_values(const CSR_matrix *A, const CSC_matrix *J,
                                      CSC_matrix *C);
 
+/* Same values from A's CSC mirror (A_csc, values current) with a caller-owned
+   accumulator of A_csc->m doubles; linear in the multiply-adds. */
+void block_left_multiply_fill_values_csc(const CSC_matrix *A_csc,
+                                         const CSC_matrix *J, CSC_matrix *C,
+                                         double *acc);
+
 /* Compute y = kron(I_p, A) @ x where A is m x n and x is(n*p)-length vector.
    The output y is m*p-length vector corresponding to
    y = [A @ x1; A @ x2; ...; A @ xp] where x is divided into p blocks of n

include/utils/sparse_matrix.h

@@ -37,6 +37,8 @@ typedef struct sparse_matrix
        position map for row_reduce_fill_values. NULL otherwise; never touched
        by any other kernel. */
     int *bound_iwork;
+    /* Accumulator (csr->m doubles) of block_left_mult_values; lazily allocated. */
+    double *bl_acc;
 } sparse_matrix;
 
 /* Constructor. Takes ownership of A; the caller must not free A separately

src/utils/linalg_sparse_matmuls.c

@@ -207,68 +207,82 @@ CSC_matrix *block_left_multiply_fill_sparsity(const CSR_matrix *A,
     return C;
 }
 
-void block_left_multiply_fill_values(const CSR_matrix *A, const CSC_matrix *J,
-                                     CSC_matrix *C)
+/* Numeric phase of Gustavson's matmul, column by column of C: for each block of
+   column j, scatter A[:, c] * J[c, j] into a dense accumulator over the rows of
+   A for every entry of J in the block, then gather the accumulator into the
+   block's entries of C. Cost is the number of multiply-adds, independent of
+   the row lengths of A. The previous version took a merge-based sparse dot of
+   a whole row of A per entry of C, which is O(m_out * nnz(row)) and quadratic
+   for a long dense row (c @ x with c of length n: n^2).
+
+   Each row's terms are added in increasing column order of A (the order of J's
+   row indices within the block), the same order as the merge-based dot, so
+   the values are bit-identical to it. acc must hold A_csc->m doubles; it need
+   not be initialized. */
+void block_left_multiply_fill_values_csc(const CSC_matrix *A_csc,
+                                         const CSC_matrix *J, CSC_matrix *C,
+                                         double *acc)
 {
     /* A is m x n, J is (n*p) x k, C is (m*p) x k */
-    int m = A->m;
-    int n = A->n;
-    int k = J->n;
+    int m = A_csc->m;
+    int n = A_csc->n;
 
-    int i, j, row_a, block, block_start, block_end, start, end;
-
-    /* to get rid of unitialized warnings */
-    block = 0;
-    block_start = 0;
-    block_end = 0;
-    start = 0;
-    end = 0;
-
-    /* for each column of J (and C) */
-    for (j = 0; j < k; j++)
+    for (int j = 0; j < J->n; j++)
     {
-        int previous_block = -1;
-
-        for (i = C->p[j]; i < C->p[j + 1]; i++)
+        int jj = J->p[j];
+        int i = C->p[j];
+        while (i < C->p[j + 1])
         {
-            /* choose row of A and block of column of J */
-            row_a = C->i[i] % m;
-            block = C->i[i] / m;
-
-            // -------------------------------------------------------------------------
-            //          find the part of the column of J in the current block
-            // -------------------------------------------------------------------------
-            if (block != previous_block)
+            /* C's row indices are sorted, so one block's entries are contiguous */
+            int block = C->i[i] / m;
+            int row_offset = block * m;
+            int block_start = block * n;
+            int block_end = block_start + n;
+            int i_end = i;
+            while (i_end < C->p[j + 1] && C->i[i_end] < row_offset + m)
             {
-                previous_block = block;
-                block_start = block * n;
-                block_end = block_start + n;
-                start = J->p[j];
-                end = J->p[j + 1];
-
-                while (start < J->p[j + 1] && J->i[start] < block_start)
-                {
-                    start++;
-                }
+                acc[C->i[i_end] - row_offset] = 0.0;
+                i_end++;
+            }
 
-                while (end > start && J->i[end - 1] >= block_end)
+            /* J's entries of this block (blocks are visited in increasing order) */
+            while (jj < J->p[j + 1] && J->i[jj] < block_start)
+            {
+                jj++;
+            }
+            for (; jj < J->p[j + 1] && J->i[jj] < block_end; jj++)
+            {
+                int c = J->i[jj] - block_start;
+                double v = J->x[jj];
+                for (int q = A_csc->p[c]; q < A_csc->p[c + 1]; q++)
                 {
-                    end--;
+                    acc[A_csc->i[q]] += A_csc->x[q] * v;
                 }
             }
 
-            // ------------------------------------------------------------------------------
-            // compute value as sparse dot product of row of A and column of J in
-            // this block
-            // ------------------------------------------------------------------------------
-            int a_len = A->p[row_a + 1] - A->p[row_a];
-            C->x[i] =
-                sparse_dot(A->x + A->p[row_a], A->i + A->p[row_a], a_len,
-                           J->x + start, J->i + start, end - start, block_start);
+            for (; i < i_end; i++)
+            {
+                C->x[i] = acc[C->i[i] - row_offset];
+            }
         }
     }
 }
 
+void block_left_multiply_fill_values(const CSR_matrix *A, const CSC_matrix *J,
+                                     CSC_matrix *C)
+{
+    /* One-off convenience form: builds A's CSC mirror per call. Callers that
+       fill repeatedly (sparse_matrix) keep the mirror and the accumulator. */
+    int *iwork = (int *) sp_malloc((A->n > 0 ? A->n : 1) * sizeof(int));
+    CSC_matrix *A_csc = csr_to_csc_alloc(A, iwork);
+    csr_to_csc_fill_values(A, A_csc, iwork);
+    double *acc = (double *) sp_malloc((A->m > 0 ? A->m : 1) * sizeof(double));
+    block_left_multiply_fill_values_csc(A_csc, J, C, acc);
+    sp_free(acc);
+    sp_free(iwork);
+    free_CSC_matrix(A_csc);
+}
+
 /* Fill values of C = A @ B where A is CSR_matrix, B is CSC_matrix. */
 void csr_csc_matmul_fill_values(const CSR_matrix *A, const CSC_matrix *B,
                                 CSR_matrix *C)

src/utils/sparse_matrix.c

@@ -43,11 +43,22 @@ static CSC_matrix *sparse_block_left_mult_sparsity(const matrix *self,
     return block_left_multiply_fill_sparsity(sm->csr, J, p);
 }
 
+static void sparse_refresh_csc_values(matrix *self);
+
 static void sparse_block_left_mult_values(const matrix *self, const CSC_matrix *J,
                                           CSC_matrix *C)
 {
-    const sparse_matrix *sm = (const sparse_matrix *) self;
-    block_left_multiply_fill_values(sm->csr, J, C);
+    /* The CSC mirror and the accumulator are caches: logically const, so cast
+       like ATA_alloc does. The mirror is version-guarded, so a constant A
+       converts once and a refreshed one reconverts. */
+    sparse_matrix *sm = (sparse_matrix *) self;
+    sparse_refresh_csc_values(&sm->base);
+    if (sm->bl_acc == NULL)
+    {
+        sm->bl_acc =
+            (double *) sp_malloc((sm->csr->m > 0 ? sm->csr->m : 1) * sizeof(double));
+    }
+    block_left_multiply_fill_values_csc(sm->csc_cache, J, C, sm->bl_acc);
 }
 
 static void sparse_free(matrix *self)
@@ -57,6 +68,7 @@ static void sparse_free(matrix *self)
     free_CSC_matrix(sm->csc_cache);
     sp_free(sm->csc_iwork);
     sp_free(sm->bound_iwork);
+    sp_free(sm->bl_acc);
     sp_free(sm);
 }
 

tests/all_tests.c

@@ -411,6 +411,9 @@ int main(void)
     mu_run_test(test_block_left_multiply_zero_column, tests_run);
     mu_run_test(test_block_left_multiply_dedup_order, tests_run);
     mu_run_test(test_block_left_multiply_matches_reference_random, tests_run);
+    mu_run_test(test_block_left_multiply_values_two_blocks, tests_run);
+    mu_run_test(test_block_left_multiply_values_dense_row, tests_run);
+    mu_run_test(test_sparse_matrix_block_left_mult_values_refill, tests_run);
     mu_run_test(test_csr_csc_matmul_alloc_basic, tests_run);
     mu_run_test(test_csr_csc_matmul_alloc_sparse, tests_run);
     mu_run_test(test_csr_csc_matmul_alloc_dedup_order, tests_run);

tests/utils/test_linalg_sparse_matmuls.h

@@ -9,6 +9,7 @@
 #include "utils/CSR_matrix.h"
 #include "utils/iVec.h"
 #include "utils/linalg_sparse_matmuls.h"
+#include "utils/sparse_matrix.h"
 #include "utils/tracked_alloc.h"
 #include "utils/utils.h"
 
@@ -647,3 +648,164 @@ const char *test_block_left_multiply_vec_three_blocks(void)
     free_CSR_matrix(A);
     return NULL;
 }
+
+/* Test block_left_multiply_fill_values on two blocks, where an entry of C sums
+ * several products and A's rows meet each block in different columns */
+const char *test_block_left_multiply_values_two_blocks(void)
+{
+    /* A is 2x3 CSR_matrix:
+     * [1.0  2.0  0.0]
+     * [0.0  3.0  4.0]
+     */
+    CSR_matrix *A = new_CSR_matrix(2, 3, 4);
+    double Ax[4] = {1.0, 2.0, 3.0, 4.0};
+    int Ai[4] = {0, 1, 1, 2};
+    int Ap[3] = {0, 2, 4};
+    memcpy(A->x, Ax, 4 * sizeof(double));
+    memcpy(A->i, Ai, 4 * sizeof(int));
+    memcpy(A->p, Ap, 3 * sizeof(int));
+
+    /* J is 6x2 CSC_matrix (two blocks of 3 rows each):
+     * Block 1 rows [0,1,2]:
+     * [5.0  0.0]
+     * [6.0  0.0]
+     * [0.0  8.0]
+     * Block 2 rows [3,4,5]:
+     * [0.0  9.0]
+     * [0.0  0.0]
+     * [7.0  0.0]
+     */
+    CSC_matrix *J = new_CSC_matrix(6, 2, 5);
+    double Jx[5] = {5.0, 6.0, 7.0, 8.0, 9.0};
+    int Ji[5] = {0, 1, 5, 2, 3};
+    int Jp[3] = {0, 3, 5};
+    memcpy(J->x, Jx, 5 * sizeof(double));
+    memcpy(J->i, Ji, 5 * sizeof(int));
+    memcpy(J->p, Jp, 3 * sizeof(int));
+
+    /* C = [A @ J1; A @ J2] is 4x2:
+     * A @ J1 = [[1*5 + 2*6, 0], [3*6, 4*8]] = [[17, 0], [18, 32]]
+     * A @ J2 = [[0, 1*9], [4*7, 0]]          = [[0, 9], [28, 0]]
+     * So C is:
+     * [17.0   0.0]
+     * [18.0  32.0]
+     * [ 0.0   9.0]
+     * [28.0   0.0]
+     */
+    CSC_matrix *C = block_left_multiply_fill_sparsity(A, J, 2);
+    block_left_multiply_fill_values(A, J, C);
+
+    int expected_p[3] = {0, 3, 5};
+    int expected_i[5] = {0, 1, 3, 1, 2};
+    double expected_x[5] = {17.0, 18.0, 28.0, 32.0, 9.0};
+
+    mu_assert("C dims incorrect", C->m == 4 && C->n == 2 && C->nnz == 5);
+    mu_assert("C col pointers incorrect", cmp_int_array(C->p, expected_p, 3));
+    mu_assert("C row indices incorrect", cmp_int_array(C->i, expected_i, 5));
+    mu_assert("C values incorrect", cmp_double_array(C->x, expected_x, 5));
+
+    free_CSC_matrix(C);
+    free_CSR_matrix(A);
+    free_CSC_matrix(J);
+    return NULL;
+}
+
+/* Test block_left_multiply_fill_values with a dense row, as in the gradient of
+ * c @ x: every entry of C needs one entry of a long row of A */
+const char *test_block_left_multiply_values_dense_row(void)
+{
+    /* A is 1x4 CSR_matrix:
+     * [1.0  2.0  3.0  4.0]
+     */
+    CSR_matrix *A = new_CSR_matrix(1, 4, 4);
+    double Ax[4] = {1.0, 2.0, 3.0, 4.0};
+    int Ai[4] = {0, 1, 2, 3};
+    int Ap[2] = {0, 4};
+    memcpy(A->x, Ax, 4 * sizeof(double));
+    memcpy(A->i, Ai, 4 * sizeof(int));
+    memcpy(A->p, Ap, 2 * sizeof(int));
+
+    /* J is 4x4 CSC_matrix (single block), diagonal:
+     * diag(1.0, -1.0, 2.0, 0.5)
+     */
+    CSC_matrix *J = new_CSC_matrix(4, 4, 4);
+    double Jx[4] = {1.0, -1.0, 2.0, 0.5};
+    int Ji[4] = {0, 1, 2, 3};
+    int Jp[5] = {0, 1, 2, 3, 4};
+    memcpy(J->x, Jx, 4 * sizeof(double));
+    memcpy(J->i, Ji, 4 * sizeof(int));
+    memcpy(J->p, Jp, 5 * sizeof(int));
+
+    /* C = A @ J is 1x4: [1*1, 2*(-1), 3*2, 4*0.5] = [1.0, -2.0, 6.0, 2.0] */
+    CSC_matrix *C = block_left_multiply_fill_sparsity(A, J, 1);
+    block_left_multiply_fill_values(A, J, C);
+
+    int expected_p[5] = {0, 1, 2, 3, 4};
+    int expected_i[4] = {0, 0, 0, 0};
+    double expected_x[4] = {1.0, -2.0, 6.0, 2.0};
+
+    mu_assert("C dims incorrect", C->m == 1 && C->n == 4 && C->nnz == 4);
+    mu_assert("C col pointers incorrect", cmp_int_array(C->p, expected_p, 5));
+    mu_assert("C row indices incorrect", cmp_int_array(C->i, expected_i, 4));
+    mu_assert("C values incorrect", cmp_double_array(C->x, expected_x, 4));
+
+    free_CSC_matrix(C);
+    free_CSR_matrix(A);
+    free_CSC_matrix(J);
+    return NULL;
+}
+
+/* Test the sparse_matrix fill, which keeps A's CSC mirror between calls: new J
+ * values are picked up, and new A values are picked up after
+ * matrix_values_changed */
+const char *test_sparse_matrix_block_left_mult_values_refill(void)
+{
+    /* A is 2x2 CSR_matrix:
+     * [1.0  2.0]
+     * [0.0  3.0]
+     */
+    CSR_matrix *A_csr = new_CSR_matrix(2, 2, 3);
+    double Ax[3] = {1.0, 2.0, 3.0};
+    int Ai[3] = {0, 1, 1};
+    int Ap[3] = {0, 2, 3};
+    memcpy(A_csr->x, Ax, 3 * sizeof(double));
+    memcpy(A_csr->i, Ai, 3 * sizeof(int));
+    memcpy(A_csr->p, Ap, 3 * sizeof(int));
+    matrix *A = new_sparse_matrix(A_csr);
+
+    /* J is 2x1 CSC_matrix: [1.0; 1.0] */
+    CSC_matrix *J = new_CSC_matrix(2, 1, 2);
+    double Jx[2] = {1.0, 1.0};
+    int Ji[2] = {0, 1};
+    int Jp[2] = {0, 2};
+    memcpy(J->x, Jx, 2 * sizeof(double));
+    memcpy(J->i, Ji, 2 * sizeof(int));
+    memcpy(J->p, Jp, 2 * sizeof(int));
+
+    /* C = A @ J = [1 + 2, 3] = [3.0, 3.0] */
+    CSC_matrix *C = A->block_left_mult_sparsity(A, J, 1);
+    A->block_left_mult_values(A, J, C);
+    double expected_1[2] = {3.0, 3.0};
+    mu_assert("C values incorrect", cmp_double_array(C->x, expected_1, 2));
+
+    /* J = [2.0; -1.0]: C = [1*2 + 2*(-1), 3*(-1)] = [0.0, -3.0] */
+    J->x[0] = 2.0;
+    J->x[1] = -1.0;
+    A->block_left_mult_values(A, J, C);
+    double expected_2[2] = {0.0, -3.0};
+    mu_assert("C values after new J incorrect",
+              cmp_double_array(C->x, expected_2, 2));
+
+    /* A = [[1, 2], [0, 5]]: C = [0.0, 5*(-1)] = [0.0, -5.0] */
+    A_csr->x[2] = 5.0;
+    matrix_values_changed(A);
+    A->block_left_mult_values(A, J, C);
+    double expected_3[2] = {0.0, -5.0};
+    mu_assert("C values after new A incorrect",
+              cmp_double_array(C->x, expected_3, 2));
+
+    free_CSC_matrix(C);
+    free_CSC_matrix(J);
+    free_matrix(A);
+    return NULL;
+}