Add a row_reduce primitive and migrate sum onto it (lesson 3, M3)

include/old-code/old_CSR_sum.h

@@ -58,4 +58,23 @@ void sum_all_rows_csr_fill_values(const CSR_matrix *A, CSR_matrix *C,
 void sum_block_of_rows_csr_fill_values(const CSR_matrix *A, CSR_matrix *C,
                                        const int *idx_map);
 
+/* Row-sum kernels that also produce an idx_map (input nnz -> position in C->x),
+   so values can be filled by zeroing C->x and calling accumulator() from
+   utils/CSR_sum.h. C must be pre-allocated with capacity >= A->nnz; iwork must
+   have size max(A->n, A->nnz); idx_map must have size A->nnz. C->nnz is set. */
+
+/* All rows of A into the single row of C (C->m == 1). */
+void sum_all_rows_csr_alloc(const CSR_matrix *A, CSR_matrix *C, int *iwork,
+                            int *idx_map);
+
+/* Consecutive blocks of row_block_size rows of A into one row of C each
+   (C->m == A->m / row_block_size). */
+void sum_block_of_rows_csr_alloc(const CSR_matrix *A, CSR_matrix *C,
+                                 int row_block_size, int *iwork, int *idx_map);
+
+/* Rows of A at stride row_spacing into one row of C each (C->m == row_spacing):
+   C[j, :] = sum_{i : i % row_spacing == j} A[i, :]. */
+void sum_evenly_spaced_rows_csr_alloc(const CSR_matrix *A, CSR_matrix *C,
+                                      int row_spacing, int *iwork, int *idx_map);
+
 #endif /* OLD_CSR_SUM_H */

include/old-code/old_mini_numpy.h

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2026 Daniel Cederberg and William Zhang
+ *
+ * This file is part of the SparseDiffEngine project.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef OLD_MINI_NUMPY_H
+#define OLD_MINI_NUMPY_H
+
+/* Example: a = [1, 2], len = 2, tiles = 3, result = [1, 2, 1, 2, 1, 2].
+   Retired from utils/mini_numpy.h with the CSR broadcast kernel, its last
+   caller. */
+void tile_int(int *result, const int *a, int len, int tiles);
+
+#endif /* OLD_MINI_NUMPY_H */

include/subexpr.h

@@ -76,7 +76,6 @@ typedef struct sum_expr
 {
     expr base;
     int axis;
-    int *idx_map; /* maps child nnz to summed-row positions */
 } sum_expr;
 
 /* trace */

include/utils/CSR_sum.h

@@ -34,7 +34,7 @@ void sum_scaled_csr_matrices_fill_values(const CSR_matrix *A, const CSR_matrix *
                                          CSR_matrix *C, const double *d1,
                                          const double *d2);
 
-/* The following five functions are used for summing either more than two CSR_matrix
+/* The following two functions are used for summing either more than two CSR_matrix
    matrices or rows of CSR_matrix matrices. To implement the filling of values
    efficiently, we compute an idx_map when we fill the sparsity pattern of the output
    matrix, which maps each nonzero entry in the input matrix to its position in the
@@ -46,17 +46,8 @@ void sum_scaled_csr_matrices_fill_values(const CSR_matrix *A, const CSR_matrix *
    matrix) corresponding to the j-th nonzero in the input matrix.
 
    Output matrix C, input matrix A, iwork->size = max(A->n, A->nnz) for the first
-   four functions. The last function allocates the output matrix and returns it. */
+   function. The last function allocates the output matrix and returns it. */
 // ------------------------------------------------------------------------------------
-void sum_all_rows_csr_alloc(const CSR_matrix *A, CSR_matrix *C, int *iwork,
-                            int *idx_map);
-
-void sum_block_of_rows_csr_alloc(const CSR_matrix *A, CSR_matrix *C,
-                                 int row_block_size, int *iwork, int *idx_map);
-
-void sum_evenly_spaced_rows_csr_alloc(const CSR_matrix *A, CSR_matrix *C,
-                                      int row_spacing, int *iwork, int *idx_map);
-
 void sum_spaced_rows_into_row_csr_alloc(const CSR_matrix *A, CSR_matrix *C,
                                         int spacing, int *iwork, int *idx_map);
 

include/utils/matrix.h

@@ -87,18 +87,14 @@ typedef void (*matrix_row_gather_fill_values_fn)(const matrix *A, matrix *C);
 typedef matrix *(*matrix_diag_vec_alloc_fn)(matrix *A);
 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).
-     - 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 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);
+/* Allocate C = row-reduce of A: C[j, :] = sum of rows i with group[i] == j,
+   group[i] in [0, m_out). C stores the reduction map internally, so the fill
+   takes none. */
+typedef matrix *(*matrix_row_reduce_alloc_fn)(const matrix *A, const int *group,
+                                              int m_out);
+
+/* Fill values of C = row-reduce of A. */
+typedef void (*matrix_row_reduce_fill_values_fn)(const matrix *A, matrix *C);
 
 typedef void (*matrix_free_fn)(matrix *self);
 
@@ -139,7 +135,8 @@ struct matrix
     matrix_row_gather_fill_values_fn row_gather_fill_values;
     matrix_diag_vec_alloc_fn diag_vec_alloc;
     matrix_diag_vec_fill_values_fn diag_vec_fill_values;
-    matrix_sum_row_partition_alloc_fn sum_row_partition_alloc;
+    matrix_row_reduce_alloc_fn row_reduce_alloc;
+    matrix_row_reduce_fill_values_fn row_reduce_fill_values;
 
     /* Lifecycle */
     matrix_free_fn free_fn;

include/utils/permuted_dense.h

@@ -68,8 +68,10 @@ typedef struct permuted_dense
 
     /* Int state bound by the alloc that produced this PD and read by the
        matching fill: row_gather_alloc stores, per dense row of this PD, the
-       dense row of the source it copies (length m0). NULL otherwise; never
-       touched by any other kernel (unlike kernel_iwork). */
+       dense row of the source it copies (length m0); row_reduce_alloc stores,
+       per dense row of the source, the dense row of this PD it adds into
+       (length source m0). NULL otherwise; never touched by any other kernel
+       (unlike kernel_iwork). */
     int *bound_iwork;
 
     /* Cached transpose of this PD as another permuted_dense, allocated lazily
@@ -101,6 +103,14 @@ matrix *row_gather_pd_alloc(const permuted_dense *A, const int *map, int m_out);
 /* Fill values of C = A[map, :], where A and C are permuted dense. */
 void row_gather_pd_fill_values(const permuted_dense *A, permuted_dense *C);
 
+/* Allocate C = row-reduce of A (C[j, :] = sum of rows i with group[i] == j,
+   group[i] in [0, m_out)), where A and C are permuted dense. C stores the
+   reduction map internally, so the fill takes none. */
+matrix *row_reduce_pd_alloc(const permuted_dense *A, const int *group, int m_out);
+
+/* Fill values of C = row-reduce of A, where A and C are permuted dense. */
+void row_reduce_pd_fill_values(const permuted_dense *A, permuted_dense *C);
+
 /* Allocate C = diag_vec(A), where A and C are permuted dense. */
 matrix *diag_vec_pd_alloc(const permuted_dense *A);
 

include/utils/sparse_matrix.h

@@ -33,8 +33,9 @@ typedef struct sparse_matrix
     /* Int state bound by the alloc that produced this matrix and read by the
        matching fill: transpose_alloc stores csr->n scratch for
        transpose_fill_values; row_gather_alloc stores the base.m-long row map
-       for row_gather_fill_values. NULL otherwise; never touched by any other
-       kernel. */
+       for row_gather_fill_values; row_reduce_alloc stores the source-nnz-long
+       position map for row_reduce_fill_values. NULL otherwise; never touched
+       by any other kernel. */
     int *bound_iwork;
 } sparse_matrix;
 

include/utils/stacked_pd_linalg.h

@@ -32,6 +32,14 @@ matrix *transpose_spd_alloc(const stacked_pd *A);
 /* Fill values of C = transpose(A). */
 void transpose_spd_fill_values(const stacked_pd *A, stacked_pd *C);
 
+/* Allocate C = row-reduce of A (C[j, :] = sum of rows i with group[i] == j,
+   group[i] in [0, m_out)); C is a stacked_pd. C stores the reduction map
+   internally, so the fill takes none. */
+matrix *row_reduce_spd_alloc(const stacked_pd *A, const int *group, int m_out);
+
+/* Fill values of C = row-reduce of A. */
+void row_reduce_spd_fill_values(const stacked_pd *A, stacked_pd *C);
+
 /* Fill values of C = diag(d) @ A, where 'd' is 'global' with length A->m. */
 void DA_spd_fill_values(const double *d, const stacked_pd *A, stacked_pd *C);
 

src/atoms/affine/sum.c

@@ -17,15 +17,9 @@
  */
 #include "atoms/affine.h"
 #include "subexpr.h"
-#include "utils/CSR_sum.h"
-#include "utils/int_double_pair.h"
 #include "utils/mini_numpy.h"
-#include "utils/sparse_matrix.h"
-#include "utils/stacked_pd.h"
 #include "utils/tracked_alloc.h"
-#include "utils/utils.h"
 #include <assert.h>
-#include <stdlib.h>
 #include <string.h>
 
 static void forward(expr *node, const double *u)
@@ -86,26 +80,32 @@ static void jacobian_init_impl(expr *node)
     expr *x = node->left;
     sum_expr *snode = (sum_expr *) node;
     jacobian_init(x);
-
-    /* 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_row_partition_alloc(x->jacobian, snode->axis,
-                                                          x->d1, snode->idx_map);
+    assert(x->jacobian->m == x->size);
+
+    /* Child rows are column-major: row i = r + c * x->d1. Row i of the child
+       Jacobian is summed into output row 0 (axis -1), c (axis 0) or r (axis 1).
+       The reduction map is bound to node->jacobian, so group is not kept. */
+    int d1 = x->d1;
+    int m_out = 1;
+    if (snode->axis == 0) m_out = x->d2;
+    if (snode->axis == 1) m_out = d1;
+    int *group = (int *) sp_malloc(x->size * sizeof(int));
+    for (int i = 0; i < x->size; i++)
+    {
+        int g = 0;
+        if (snode->axis == 0) g = i / d1;
+        if (snode->axis == 1) g = i % d1;
+        group[i] = g;
+    }
+    node->jacobian = x->jacobian->row_reduce_alloc(x->jacobian, group, m_out);
+    sp_free(group);
 }
 
 static void eval_jacobian_impl(expr *node)
 {
     expr *child = node->left;
-
-    /* evaluate child's jacobian */
     eval_jacobian(child);
-
-    /* we have precomputed an idx map between the nonzeros of the child's jacobian
-       and this node's jacobian, so we just accumulate accordingly */
-    memset(node->jacobian->x, 0, node->jacobian->nnz * sizeof(double));
-    accumulator(child->jacobian->x, child->jacobian->nnz,
-                ((sum_expr *) node)->idx_map, node->jacobian->x);
+    child->jacobian->row_reduce_fill_values(child->jacobian, node->jacobian);
 }
 
 static void wsum_hess_init_impl(expr *node)
@@ -149,12 +149,6 @@ static bool is_affine(const expr *node)
     return node->left->is_affine(node->left);
 }
 
-static void free_type_data(expr *node)
-{
-    sum_expr *snode = (sum_expr *) node;
-    sp_free(snode->idx_map);
-}
-
 expr *new_sum(expr *child, int axis)
 {
     int d2 = 0;
@@ -183,7 +177,7 @@ expr *new_sum(expr *child, int axis)
        sum with an axis argument as a row vector */
     init_expr(node, 1, d2, child->n_vars, forward, jacobian_init_impl,
               eval_jacobian_impl, is_affine, wsum_hess_init_impl,
-              eval_wsum_hess_impl, free_type_data);
+              eval_wsum_hess_impl, NULL);
     node->left = child;
     expr_retain(child);