Remove "atom specific" functions from matrix classes - PR 3

include/subexpr.h

@@ -216,6 +216,14 @@ typedef struct index_expr
     bool has_duplicates; /* True if indices have duplicates (affects Hessian path) */
 } index_expr;
 
+/* Broadcast shape used by the broadcast atom. */
+typedef enum
+{
+    BROADCAST_ROW,   /* (1, n) -> (m, n) */
+    BROADCAST_COL,   /* (m, 1) -> (m, n) */
+    BROADCAST_SCALAR /* (1, 1) -> (m, n) */
+} broadcast_type;
+
 typedef struct broadcast_expr
 {
     expr base;

include/utils/matrix.h

@@ -23,14 +23,6 @@
 #include <stdbool.h>
 #include <stdint.h>
 
-/* Broadcast shape used by the broadcast atom and its vtable methods. */
-typedef enum
-{
-    BROADCAST_ROW,   /* (1, n) -> (m, n) */
-    BROADCAST_COL,   /* (m, 1) -> (m, n) */
-    BROADCAST_SCALAR /* (1, 1) -> (m, n) */
-} broadcast_type;
-
 /* Polymorphic matrix base. Concrete types embed `matrix` as their first
    member and implement the vtable slots below. Currently implemented:
        1. sparse_matrix  — generic CSR_matrix-backed matrix.
@@ -89,13 +81,6 @@ typedef matrix *(*matrix_row_gather_alloc_fn)(const matrix *A, const int *map,
 /* Fill values of C = A[map, :] */
 typedef void (*matrix_row_gather_fill_values_fn)(const matrix *A, matrix *C);
 
-/* Broadcast: lift the child Jacobian of a broadcast atom into the output
-   Jacobian. `type` is the broadcast variant; (d1, d2) is the output shape. */
-typedef matrix *(*matrix_broadcast_alloc_fn)(matrix *A, broadcast_type type, int d1,
-                                             int d2);
-typedef void (*matrix_broadcast_fill_values_fn)(matrix *A, broadcast_type type,
-                                                int d1, int d2, matrix *out);
-
 /* diag_vec: A is an (n, A->n) Jacobian for a length-n vector; output is
    (n*n, A->n) where row i lands at output row i*(n+1) (column-major
    diagonal positions). Other output rows are structurally zero. */
@@ -152,8 +137,6 @@ struct matrix
     /* Atom-specific ops */
     matrix_row_gather_alloc_fn row_gather_alloc;
     matrix_row_gather_fill_values_fn row_gather_fill_values;
-    matrix_broadcast_alloc_fn broadcast_alloc;
-    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_row_partition_alloc_fn sum_row_partition_alloc;

include/utils/mini_numpy.h

@@ -25,7 +25,6 @@ void repeat(double *result, const double *a, int len, int repeats);
 
 /* Example: a = [1, 2], len = 2, tiles = 3, result = [1, 2, 1, 2, 1, 2] */
 void tile_double(double *result, const double *a, int len, int tiles);
-void tile_int(int *result, const int *a, int len, int tiles);
 
 /* Example: size = 5, value = 3.0, result = [3.0, 3.0, 3.0, 3.0, 3.0] */
 void scaled_ones(double *result, int size, double value);

include/utils/permuted_dense.h

@@ -94,14 +94,6 @@ matrix *new_permuted_dense_full(int m, int n, const double *data);
    place; contents are NOT preserved. */
 void permuted_dense_ensure_kernel_dwork(const permuted_dense *A, size_t size);
 
-/* Allocate C = broadcast(A, type, d1, d2), where A and C are permuted dense. */
-matrix *broadcast_pd_alloc(const permuted_dense *A, broadcast_type type, int d1,
-                           int d2);
-
-/* Fill values of C = broadcast(A, type, d1, d2). */
-void broadcast_pd_fill_values(const permuted_dense *A, broadcast_type type, int d1,
-                              int d2, permuted_dense *C);
-
 /* Allocate C = A[map, :], where A and C are permuted dense. C stores map
    internally, so the fill takes none. */
 matrix *row_gather_pd_alloc(const permuted_dense *A, const int *map, int m_out);

src/atoms/affine/broadcast.c

@@ -71,20 +71,34 @@ static void jacobian_init_impl(expr *node)
     expr *x = node->left;
     jacobian_init(x);
 
-    /* allocate sparsity for the broadcast output; output type matches child's. */
+    /* Every output row (i, j) at column-major index i + j*d1 is a copy of one
+       child row: j for ROW ((1, d2) child), i for COL ((d1, 1) child), 0 for
+       SCALAR. A row gather with that map; the map is bound to node->jacobian
+       and not kept. */
     broadcast_expr *bcast = (broadcast_expr *) node;
-    node->jacobian =
-        x->jacobian->broadcast_alloc(x->jacobian, bcast->type, node->d1, node->d2);
+    int d1 = node->d1;
+    int *map = (int *) sp_malloc(node->size * sizeof(int));
+    for (int j = 0; j < node->d2; j++)
+    {
+        for (int i = 0; i < d1; i++)
+        {
+            int src = 0;
+            if (bcast->type == BROADCAST_ROW) src = j;
+            if (bcast->type == BROADCAST_COL) src = i;
+            map[i + j * d1] = src;
+        }
+    }
+    node->jacobian = x->jacobian->row_gather_alloc(x->jacobian, map, node->size);
+    sp_free(map);
 }
 
 static void eval_jacobian_impl(expr *node)
 {
     eval_jacobian(node->left);
 
     /* fill values into the preallocated output. */
-    broadcast_expr *bcast = (broadcast_expr *) node;
-    node->left->jacobian->broadcast_fill_values(node->left->jacobian, bcast->type,
-                                                node->d1, node->d2, node->jacobian);
+    node->left->jacobian->row_gather_fill_values(node->left->jacobian,
+                                                 node->jacobian);
 }
 
 static void wsum_hess_init_impl(expr *node)

src/utils/mini_numpy.c

@@ -38,14 +38,6 @@ void tile_double(double *result, const double *a, int len, int tiles)
     }
 }
 
-void tile_int(int *result, const int *a, int len, int tiles)
-{
-    for (int i = 0; i < tiles; i++)
-    {
-        memcpy(result + i * len, a, len * sizeof(int));
-    }
-}
-
 void scaled_ones(double *result, int size, double value)
 {
     for (int i = 0; i < size; i++)

src/utils/permuted_dense.c

@@ -164,119 +164,6 @@ static void permuted_dense_vtable_row_gather_fill_values(const matrix *self,
     row_gather_pd_fill_values((const permuted_dense *) self, (permuted_dense *) out);
 }
 
-matrix *broadcast_pd_alloc(const permuted_dense *A, broadcast_type type, int d1,
-                           int d2)
-{
-    int out_m = d1 * d2;
-
-    int new_m0;
-    if (type == BROADCAST_SCALAR)
-    {
-        new_m0 = (A->m0 == 0) ? 0 : out_m;
-    }
-    else if (type == BROADCAST_ROW)
-    {
-        new_m0 = d1 * A->m0;
-    }
-    else /* BROADCAST_COL */
-    {
-        new_m0 = d2 * A->m0;
-    }
-
-    if (new_m0 == 0)
-    {
-        return new_permuted_dense(out_m, A->base.n, 0, A->n0, NULL, A->col_perm,
-                                  NULL);
-    }
-
-    int *new_row_perm = (int *) sp_malloc(new_m0 * sizeof(int));
-    int k = 0;
-    if (type == BROADCAST_SCALAR)
-    {
-        for (int i = 0; i < out_m; i++)
-        {
-            new_row_perm[k++] = i;
-        }
-    }
-    else if (type == BROADCAST_ROW)
-    {
-        for (int j_ii = 0; j_ii < A->m0; j_ii++)
-        {
-            int j_old = A->row_perm[j_ii];
-            for (int i = 0; i < d1; i++)
-            {
-                new_row_perm[k++] = j_old * d1 + i;
-            }
-        }
-    }
-    else /* BROADCAST_COL */
-    {
-        for (int j = 0; j < d2; j++)
-        {
-            for (int ii_old = 0; ii_old < A->m0; ii_old++)
-            {
-                new_row_perm[k++] = j * d1 + A->row_perm[ii_old];
-            }
-        }
-    }
-
-    matrix *out = new_permuted_dense(out_m, A->base.n, new_m0, A->n0, new_row_perm,
-                                     A->col_perm, NULL);
-    sp_free(new_row_perm);
-    return out;
-}
-
-void broadcast_pd_fill_values(const permuted_dense *A, broadcast_type type, int d1,
-                              int d2, permuted_dense *C)
-{
-    if (A->m0 == 0)
-    {
-        return;
-    }
-    int n0 = A->n0;
-
-    if (type == BROADCAST_SCALAR)
-    {
-        for (int k = 0; k < C->m0; k++)
-        {
-            memcpy(C->X + k * n0, A->X, n0 * sizeof(double));
-        }
-    }
-    else if (type == BROADCAST_ROW)
-    {
-        /* output row k corresponds to child dense row (k / d1). */
-        (void) d2;
-        for (int k = 0; k < C->m0; k++)
-        {
-            memcpy(C->X + k * n0, A->X + (k / d1) * n0, n0 * sizeof(double));
-        }
-    }
-    else /* BROADCAST_COL */
-    {
-        (void) d1;
-        size_t child_block = A->m0 * n0;
-        for (int j = 0; j < d2; j++)
-        {
-            memcpy(C->X + j * child_block, A->X, child_block * sizeof(double));
-        }
-    }
-}
-
-static matrix *permuted_dense_vtable_broadcast_alloc(matrix *self,
-                                                     broadcast_type type, int d1,
-                                                     int d2)
-{
-    return broadcast_pd_alloc((const permuted_dense *) self, type, d1, d2);
-}
-
-static void permuted_dense_vtable_broadcast_fill_values(matrix *self,
-                                                        broadcast_type type, int d1,
-                                                        int d2, matrix *out)
-{
-    broadcast_pd_fill_values((const permuted_dense *) self, type, d1, d2,
-                             (permuted_dense *) out);
-}
-
 matrix *diag_vec_pd_alloc(const permuted_dense *A)
 {
     int n = A->base.m;
@@ -525,8 +412,6 @@ static void wire_vtable(permuted_dense *pd)
     pd->base.transpose_fill_values = permuted_dense_vtable_transpose_fill_values;
     pd->base.row_gather_alloc = permuted_dense_vtable_row_gather_alloc;
     pd->base.row_gather_fill_values = permuted_dense_vtable_row_gather_fill_values;
-    pd->base.broadcast_alloc = permuted_dense_vtable_broadcast_alloc;
-    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_row_partition_alloc = permuted_dense_vtable_sum_row_partition_alloc;

src/utils/sparse_matrix.c

@@ -185,94 +185,6 @@ static void sparse_row_gather_fill_values(const matrix *self, matrix *out)
     }
 }
 
-static matrix *sparse_broadcast_alloc(matrix *self, broadcast_type type, int d1,
-                                      int d2)
-{
-    CSR_matrix *Jx = ((sparse_matrix *) self)->csr;
-    int out_m = d1 * d2;
-    int total_nnz;
-    if (type == BROADCAST_ROW)
-    {
-        total_nnz = Jx->nnz * d1;
-    }
-    else if (type == BROADCAST_COL)
-    {
-        total_nnz = Jx->nnz * d2;
-    }
-    else /* BROADCAST_SCALAR */
-    {
-        total_nnz = Jx->nnz * out_m;
-    }
-
-    CSR_matrix *J = new_CSR_matrix(out_m, self->n, total_nnz);
-
-    if (type == BROADCAST_ROW)
-    {
-        int acc = 0;
-        for (int i = 0; i < d2; i++)
-        {
-            int nnz_in_row = Jx->p[i + 1] - Jx->p[i];
-            tile_int(J->i + acc, Jx->i + Jx->p[i], nnz_in_row, d1);
-            for (int rep = 0; rep < d1; rep++)
-            {
-                J->p[i * d1 + rep] = acc;
-                acc += nnz_in_row;
-            }
-        }
-        J->p[out_m] = total_nnz;
-    }
-    else if (type == BROADCAST_COL)
-    {
-        tile_int(J->i, Jx->i, Jx->nnz, d2);
-        int offset = 0;
-        for (int i = 0; i < d2; i++)
-        {
-            for (int j = 0; j < d1; j++)
-            {
-                int nnz_in_row = Jx->p[j + 1] - Jx->p[j];
-                J->p[i * d1 + j] = offset;
-                offset += nnz_in_row;
-            }
-        }
-        J->p[out_m] = total_nnz;
-    }
-    else /* BROADCAST_SCALAR */
-    {
-        tile_int(J->i, Jx->i, Jx->nnz, out_m);
-        int row_nnz = Jx->nnz;
-        for (int i = 0; i < out_m; i++)
-        {
-            J->p[i] = i * row_nnz;
-        }
-        J->p[out_m] = total_nnz;
-    }
-    return new_sparse_matrix(J);
-}
-
-static void sparse_broadcast_fill_values(matrix *self, broadcast_type type, int d1,
-                                         int d2, matrix *out)
-{
-    CSR_matrix *Jx = ((sparse_matrix *) self)->csr;
-    if (type == BROADCAST_ROW)
-    {
-        int acc = 0;
-        for (int i = 0; i < d2; i++)
-        {
-            int nnz_in_row = Jx->p[i + 1] - Jx->p[i];
-            tile_double(out->x + acc, Jx->x + Jx->p[i], nnz_in_row, d1);
-            acc += nnz_in_row * d1;
-        }
-    }
-    else if (type == BROADCAST_COL)
-    {
-        tile_double(out->x, Jx->x, Jx->nnz, d2);
-    }
-    else /* BROADCAST_SCALAR */
-    {
-        tile_double(out->x, Jx->x, Jx->nnz, d1 * d2);
-    }
-}
-
 static matrix *sparse_diag_vec_alloc(matrix *self)
 {
     CSR_matrix *Jx = ((sparse_matrix *) self)->csr;
@@ -376,8 +288,6 @@ static void wire_vtable(sparse_matrix *sm)
     sm->base.transpose_fill_values = sparse_transpose_fill_values;
     sm->base.row_gather_alloc = sparse_row_gather_alloc;
     sm->base.row_gather_fill_values = sparse_row_gather_fill_values;
-    sm->base.broadcast_alloc = sparse_broadcast_alloc;
-    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_row_partition_alloc = sparse_sum_row_partition_alloc;