permuted dense sum different rows

Author: dance858

include/utils/matrix.h

@@ -109,20 +109,15 @@ typedef void (*matrix_diag_vec_fill_values_fn)(matrix *A, matrix *out);
 /* Allocate C as a row-wise reduction of A. The reduction pattern is chosen by
    axis:
      - axis = -1: sum all rows of A. C has shape (1, A->n).
-                  C[0, k] = sum_{i in [0, A->m)} A[i, k].
-     - axis = 0:  block-sum rows in consecutive groups of d1 (requires
-                  A->m % d1 == 0). C has shape (A->m / d1, A->n).
-                  C[j, k] = sum_{i in [j*d1, (j+1)*d1)} A[i, k].
-     - axis = 1:  stride-sum rows at modular spacing d1 (requires
-                  A->m % d1 == 0). C has shape (d1, A->n).
-                  C[j, k] = sum_{i : i % d1 == j} A[i, k].
-   d1 is ignored when axis == -1.
-
-   Caller pre-allocates idx_map of size A->nnz. On return idx_map[ii] holds the
-   position in C's base.x where A's base.x[ii] scatter-adds. A values-fill pass
-   can compute C->x by zeroing it and accumulating A->x[ii] into
-   C->x[idx_map[ii]] for ii in [0, A->nnz). */
-typedef matrix *(*matrix_sum_alloc_fn)(matrix *A, int axis, int d1, int *idx_map);
+     - axis = 0:  block-sum rows in consecutive groups of d1. C has shape (A->m
+                  / d1, A->n). C[j, :] = sum_{i in [j*d1, (j+1)*d1)} A[i, :].
+     - axis = 1:  stride-sum rows at modular spacing d1. C has shape (d1, A->n).
+                  C[j, :] = sum_{i : i % d1 == j} A[i, :].
+
+   Caller pre-allocates idx_map of size A->nnz that can be used to compute the
+   numerical result of the operation using via accumulation. */
+typedef matrix *(*matrix_sum_row_partition_alloc_fn)(matrix *A, int axis, int d1,
+                                                     int *idx_map);
 
 typedef void (*matrix_free_fn)(matrix *self);
 
@@ -159,7 +154,7 @@ struct matrix
     matrix_broadcast_fill_values_fn broadcast_fill_values;
     matrix_diag_vec_alloc_fn diag_vec_alloc;
     matrix_diag_vec_fill_values_fn diag_vec_fill_values;
-    matrix_sum_alloc_fn sum_alloc;
+    matrix_sum_row_partition_alloc_fn sum_row_partition_alloc;
 
     /* Lifecycle */
     matrix_free_fn free_fn;

src/atoms/affine/sum.c

@@ -87,11 +87,11 @@ static void jacobian_init_impl(expr *node)
     sum_expr *snode = (sum_expr *) node;
     jacobian_init(x);
 
-    /* sum_alloc fills idx_map so eval_jacobian can accumulate from
+    /* sum_row_partition_alloc fills idx_map so eval_jacobian can accumulate from
        child->jacobian->x. */
     snode->idx_map = sp_malloc(x->jacobian->nnz * sizeof(int));
-    node->jacobian =
-        x->jacobian->sum_alloc(x->jacobian, snode->axis, x->d1, snode->idx_map);
+    node->jacobian = x->jacobian->sum_row_partition_alloc(x->jacobian, snode->axis,
+                                                          x->d1, snode->idx_map);
 }
 
 static void eval_jacobian(expr *node)

src/utils/permuted_dense.c

@@ -523,8 +523,8 @@ static matrix *sum_evenly_spaced_rows_pd_alloc(matrix *self, int d1, int *idx_ma
     return C;
 }
 
-static matrix *permuted_dense_vtable_sum_alloc(matrix *self, int axis, int d1,
-                                               int *idx_map)
+static matrix *permuted_dense_vtable_sum_row_partition_alloc(matrix *self, int axis,
+                                                             int d1, int *idx_map)
 {
     if (axis == -1)
     {
@@ -562,7 +562,7 @@ static void wire_vtable(permuted_dense *pd)
     pd->base.broadcast_fill_values = permuted_dense_vtable_broadcast_fill_values;
     pd->base.diag_vec_alloc = permuted_dense_vtable_diag_vec_alloc;
     pd->base.diag_vec_fill_values = permuted_dense_vtable_diag_vec_fill_values;
-    pd->base.sum_alloc = permuted_dense_vtable_sum_alloc;
+    pd->base.sum_row_partition_alloc = permuted_dense_vtable_sum_row_partition_alloc;
     pd->base.refresh_csc_values = permuted_dense_refresh_csc_values;
 }
 

src/utils/sparse_matrix.c

@@ -312,7 +312,8 @@ static void sparse_refresh_csc_values(matrix *self)
     csr_to_csc_fill_values(sm->csr, sm->csc_cache, sm->csc_iwork);
 }
 
-static matrix *sparse_sum_alloc(matrix *self, int axis, int d1, int *idx_map)
+static matrix *sparse_sum_row_partition_alloc(matrix *self, int axis, int d1,
+                                              int *idx_map)
 {
     CSR_matrix *A = ((sparse_matrix *) self)->csr;
     int m;
@@ -369,7 +370,7 @@ static void wire_vtable(sparse_matrix *sm)
     sm->base.broadcast_fill_values = sparse_broadcast_fill_values;
     sm->base.diag_vec_alloc = sparse_diag_vec_alloc;
     sm->base.diag_vec_fill_values = sparse_diag_vec_fill_values;
-    sm->base.sum_alloc = sparse_sum_alloc;
+    sm->base.sum_row_partition_alloc = sparse_sum_row_partition_alloc;
     sm->base.refresh_csc_values = sparse_refresh_csc_values;
     sm->base.free_fn = sparse_free;
 }

src/utils/stacked_pd.c

@@ -394,8 +394,8 @@ static void assert_disjoint_row_perms(int n_blocks, permuted_dense *const *block
    existing CSR helper, then re-index idx_map from CSR ordering into A's
    block-major base.x ordering so a downstream values-fill pass can read
    directly from A->base.x. */
-static matrix *stacked_pd_vtable_sum_alloc(matrix *self, int axis, int d1,
-                                           int *idx_map)
+static matrix *stacked_pd_vtable_sum_row_partition_alloc(matrix *self, int axis,
+                                                         int d1, int *idx_map)
 {
     stacked_pd *spd = (stacked_pd *) self;
 
@@ -488,7 +488,7 @@ static void wire_vtable(stacked_pd *spd)
     spd->base.diag_vec_fill_values = stacked_pd_vtable_diag_vec_fill_values;
     spd->base.broadcast_alloc = stacked_pd_vtable_broadcast_alloc;
     spd->base.broadcast_fill_values = stacked_pd_vtable_broadcast_fill_values;
-    spd->base.sum_alloc = stacked_pd_vtable_sum_alloc;
+    spd->base.sum_row_partition_alloc = stacked_pd_vtable_sum_row_partition_alloc;
 }
 
 matrix *new_stacked_pd_unchecked(int m, int n, int n_blocks, permuted_dense **blocks,

tests/utils/test_permuted_dense.h

@@ -1047,8 +1047,8 @@ const char *test_BA_pd_matrices_fast_path(void)
     return 0;
 }
 
-/* Direct vtable tests for sum_alloc. The test PD represents a (6, 4) matrix
-   with a (3, 2) dense block at rows {0, 3, 4}, cols {1, 3}. We exercise all
+/* Direct vtable tests for sum_row_partition_alloc. The test PD represents a (6, 4)
+   matrix with a (3, 2) dense block at rows {0, 3, 4}, cols {1, 3}. We exercise all
    three axes; for axis=0 (d1=2) the buckets {0/2, 3/2, 4/2} = {0, 1, 2}
    are all distinct and non-decreasing (linear-scan dedupe); for axis=1
    (d1=3) the buckets {0%3, 3%3, 4%3} = {0, 0, 1} collapse two input rows
@@ -1062,7 +1062,7 @@ const char *test_permuted_dense_sum_all_rows(void)
     matrix *M = new_permuted_dense(6, 4, 3, 2, row_perm, col_perm, X);
 
     int idx_map[6];
-    matrix *out = M->sum_alloc(M, -1, /*d1 unused*/ 0, idx_map);
+    matrix *out = M->sum_row_partition_alloc(M, -1, /*d1 unused*/ 0, idx_map);
 
     mu_assert("output is PD", out->is_permuted_dense);
     permuted_dense *opd = (permuted_dense *) out;
@@ -1092,7 +1092,7 @@ const char *test_permuted_dense_sum_block_of_rows(void)
 
     int d1 = 2; /* child shape (d1, d2) = (2, 3); output rows = d2 = 3 */
     int idx_map[6];
-    matrix *out = M->sum_alloc(M, 0, d1, idx_map);
+    matrix *out = M->sum_row_partition_alloc(M, 0, d1, idx_map);
 
     mu_assert("output is PD", out->is_permuted_dense);
     permuted_dense *opd = (permuted_dense *) out;
@@ -1121,7 +1121,7 @@ const char *test_permuted_dense_sum_evenly_spaced_rows(void)
 
     int d1 = 3; /* output rows = d1 = 3, buckets = {0%3, 3%3, 4%3} = {0, 0, 1} */
     int idx_map[6];
-    matrix *out = M->sum_alloc(M, 1, d1, idx_map);
+    matrix *out = M->sum_row_partition_alloc(M, 1, d1, idx_map);
 
     mu_assert("output is PD", out->is_permuted_dense);
     permuted_dense *opd = (permuted_dense *) out;