diff --git a/include/utils/linalg_sparse_matmuls.h b/include/utils/linalg_sparse_matmuls.h index 516ff18..f0485ea 100644 --- a/include/utils/linalg_sparse_matmuls.h +++ b/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 diff --git a/include/utils/sparse_matrix.h b/include/utils/sparse_matrix.h index 039ca83..d6e390c 100644 --- a/include/utils/sparse_matrix.h +++ b/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 diff --git a/src/utils/linalg_sparse_matmuls.c b/src/utils/linalg_sparse_matmuls.c index f6ba8b8..8b94716 100644 --- a/src/utils/linalg_sparse_matmuls.c +++ b/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) diff --git a/src/utils/sparse_matrix.c b/src/utils/sparse_matrix.c index 7459dd1..435c2fc 100644 --- a/src/utils/sparse_matrix.c +++ b/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); } diff --git a/tests/all_tests.c b/tests/all_tests.c index a4486f7..7081d5c 100644 --- a/tests/all_tests.c +++ b/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); diff --git a/tests/utils/test_linalg_sparse_matmuls.h b/tests/utils/test_linalg_sparse_matmuls.h index 0e74ed7..5f36b7e 100644 --- a/tests/utils/test_linalg_sparse_matmuls.h +++ b/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; +}