Skip to content

Commit 71fa1d8

Browse files
dance858dance858
authored
Broadcasting (#25)
* some progress on broadcasting * progress on broadcast
1 parent 3ae98b5 commit 71fa1d8

10 files changed

Lines changed: 715 additions & 3 deletions

File tree

‎include/affine.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,5 +20,6 @@ expr *new_variable(int d1, int d2, int var_id, int n_vars);
2020

2121
expr *new_index(expr *child, const int *indices, int n_idxs);
2222
expr *new_reshape(expr *child, int d1, int d2);
23+
expr *new_broadcast(expr *child, int target_d1, int target_d2);
2324

2425
#endif /* AFFINE_H */

‎include/subexpr.h‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,4 +118,20 @@ typedef struct index_expr
118118
bool has_duplicates; /* True if indices have duplicates (affects Hessian path) */
119119
} index_expr;
120120

121+
/* Broadcast types */
122+
typedef enum
123+
{
124+
BROADCAST_ROW, /* (1, n) -> (m, n) */
125+
BROADCAST_COL, /* (m, 1) -> (m, n) */
126+
BROADCAST_SCALAR /* (1, 1) -> (m, n) */
127+
} broadcast_type;
128+
129+
typedef struct broadcast_expr
130+
{
131+
expr base;
132+
broadcast_type type;
133+
int m; /* target rows */
134+
int n; /* target cols */
135+
} broadcast_expr;
136+
121137
#endif /* SUBEXPR_H */

‎include/utils/mini_numpy.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@ void repeat(double *result, const double *a, int len, int repeats);
1111
* Example: a = [1, 2], len = 2, tiles = 3
1212
* result = [1, 2, 1, 2, 1, 2]
1313
*/
14-
void tile(double *result, const double *a, int len, int tiles);
14+
void tile_double(double *result, const double *a, int len, int tiles);
15+
void tile_int(int *result, const int *a, int len, int tiles);
1516

1617
/* Fill array with 'size' copies of 'value'
1718
* Example: size = 5, value = 3.0

‎src/affine/broadcast.c‎

Lines changed: 275 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,275 @@
1+
#include "affine.h"
2+
#include "subexpr.h"
3+
#include "utils/mini_numpy.h"
4+
#include <assert.h>
5+
#include <stdio.h>
6+
#include <stdlib.h>
7+
#include <string.h>
8+
9+
/* Broadcast expands an array to a larger shape by replicating along dimensions.
10+
* Supports three types:
11+
* 1. "row": (1, n) -> (m, n) - replicate rows
12+
* 2. "col": (m, 1) -> (m, n) - replicate columns
13+
* 3. "scalar": (1, 1) -> (m, n) - replicate in both dimensions
14+
*/
15+
16+
static void forward(expr *node, const double *u)
17+
{
18+
expr *x = node->left;
19+
broadcast_expr *bcast = (broadcast_expr *) node;
20+
21+
x->forward(x, u);
22+
23+
if (bcast->type == BROADCAST_ROW)
24+
{
25+
/* (1, n) -> (m, n): replicate row m times */
26+
for (int j = 0; j < bcast->n; j++)
27+
{
28+
for (int i = 0; i < bcast->m; i++)
29+
{
30+
node->value[i + j * bcast->m] = x->value[j];
31+
}
32+
}
33+
}
34+
else if (bcast->type == BROADCAST_COL)
35+
{
36+
/* (m, 1) -> (m, n): replicate column n times */
37+
for (int j = 0; j < bcast->n; j++)
38+
{
39+
memcpy(node->value + j * bcast->m, x->value, bcast->m * sizeof(double));
40+
}
41+
}
42+
else
43+
{
44+
/* (1, 1) -> (m, n): fill with scalar value */
45+
for (int k = 0; k < node->size; k++)
46+
{
47+
node->value[k] = x->value[0];
48+
}
49+
}
50+
}
51+
52+
static void jacobian_init(expr *node)
53+
{
54+
expr *x = node->left;
55+
x->jacobian_init(x);
56+
broadcast_expr *bcast = (broadcast_expr *) node;
57+
int total_nnz;
58+
59+
// --------------------------------------------------------------------
60+
// count number of nonzeros
61+
// --------------------------------------------------------------------
62+
if (bcast->type == BROADCAST_ROW)
63+
{
64+
/* Row broadcast: (1, n) -> (m, n) */
65+
total_nnz = x->jacobian->nnz * bcast->m;
66+
}
67+
else if (bcast->type == BROADCAST_COL)
68+
{
69+
/* Column broadcast: (m, 1) -> (m, n) */
70+
total_nnz = x->jacobian->nnz * bcast->n;
71+
}
72+
else
73+
{
74+
/* Scalar broadcast: (1, 1) -> (m, n) */
75+
total_nnz = x->jacobian->nnz * bcast->m * bcast->n;
76+
}
77+
78+
node->jacobian = new_csr_matrix(node->size, node->n_vars, total_nnz);
79+
80+
// ---------------------------------------------------------------------
81+
// fill sparsity pattern
82+
// ---------------------------------------------------------------------
83+
CSR_Matrix *Jx = x->jacobian;
84+
CSR_Matrix *J = node->jacobian;
85+
J->nnz = 0;
86+
87+
if (bcast->type == BROADCAST_ROW)
88+
{
89+
for (int i = 0; i < bcast->n; i++)
90+
{
91+
int nnz_in_row = Jx->p[i + 1] - Jx->p[i];
92+
93+
/* copy columns indices */
94+
tile_int(J->i + J->nnz, Jx->i + Jx->p[i], nnz_in_row, bcast->m);
95+
96+
/* set row pointers */
97+
for (int rep = 0; rep < bcast->m; rep++)
98+
{
99+
J->p[i * bcast->m + rep] = J->nnz;
100+
J->nnz += nnz_in_row;
101+
}
102+
}
103+
}
104+
else if (bcast->type == BROADCAST_COL)
105+
{
106+
107+
/* copy column indices */
108+
tile_int(J->i, Jx->i, Jx->nnz, bcast->n);
109+
110+
/* set row pointers */
111+
int offset = 0;
112+
for (int i = 0; i < bcast->n; i++)
113+
{
114+
for (int j = 0; j < bcast->m; j++)
115+
{
116+
J->p[i * bcast->m + j] = offset;
117+
offset += Jx->p[1] - Jx->p[0];
118+
}
119+
}
120+
assert(offset == total_nnz);
121+
J->p[node->size] = total_nnz;
122+
}
123+
else
124+
{
125+
/* copy column indices */
126+
tile_int(J->i, Jx->i, Jx->nnz, bcast->m * bcast->n);
127+
128+
/* set row pointers */
129+
int offset = 0;
130+
int nnz = Jx->p[1] - Jx->p[0];
131+
for (int i = 0; i < bcast->m * bcast->n; i++)
132+
{
133+
J->p[i] = offset;
134+
offset += nnz;
135+
}
136+
assert(offset == total_nnz);
137+
J->p[node->size] = total_nnz;
138+
}
139+
}
140+
141+
static void eval_jacobian(expr *node)
142+
{
143+
node->left->eval_jacobian(node->left);
144+
145+
broadcast_expr *bcast = (broadcast_expr *) node;
146+
CSR_Matrix *Jx = node->left->jacobian;
147+
CSR_Matrix *J = node->jacobian;
148+
J->nnz = 0;
149+
150+
if (bcast->type == BROADCAST_ROW)
151+
{
152+
for (int i = 0; i < bcast->n; i++)
153+
{
154+
int nnz_in_row = Jx->p[i + 1] - Jx->p[i];
155+
tile_double(J->x + J->nnz, Jx->x + Jx->p[i], nnz_in_row, bcast->m);
156+
J->nnz += nnz_in_row * bcast->m;
157+
}
158+
}
159+
else if (bcast->type == BROADCAST_COL)
160+
{
161+
tile_double(J->x, Jx->x, Jx->nnz, bcast->n);
162+
}
163+
else
164+
{
165+
tile_double(J->x, Jx->x, Jx->nnz, bcast->m * bcast->n);
166+
}
167+
}
168+
169+
static void wsum_hess_init(expr *node)
170+
{
171+
expr *x = node->left;
172+
x->wsum_hess_init(x);
173+
174+
/* Same sparsity as child - weights get summed */
175+
node->wsum_hess = new_csr_matrix(node->n_vars, node->n_vars, x->wsum_hess->nnz);
176+
memcpy(node->wsum_hess->p, x->wsum_hess->p, (x->wsum_hess->m + 1) * sizeof(int));
177+
memcpy(node->wsum_hess->i, x->wsum_hess->i, x->wsum_hess->nnz * sizeof(int));
178+
179+
/* allocate space for weight vector */
180+
node->dwork = malloc(node->size * sizeof(double));
181+
}
182+
183+
static void eval_wsum_hess(expr *node, const double *w)
184+
{
185+
broadcast_expr *bcast = (broadcast_expr *) node;
186+
expr *x = node->left;
187+
188+
/* Zero out the work array first */
189+
memset(node->dwork, 0, x->size * sizeof(double));
190+
191+
if (bcast->type == BROADCAST_ROW)
192+
{
193+
/* (1, n) -> (m, n): each input element has m weights to sum */
194+
for (int j = 0; j < bcast->n; j++)
195+
{
196+
for (int i = 0; i < bcast->m; i++)
197+
{
198+
node->dwork[j] += w[i + j * bcast->m];
199+
}
200+
}
201+
}
202+
else if (bcast->type == BROADCAST_COL)
203+
{
204+
/* (m, 1) -> (m, n): each input element has n weights to sum */
205+
for (int j = 0; j < bcast->n; j++)
206+
{
207+
for (int i = 0; i < bcast->m; i++)
208+
{
209+
node->dwork[i] += w[i + j * bcast->m];
210+
}
211+
}
212+
}
213+
else
214+
{
215+
/* (1, 1) -> (m, n): scalar has m*n weights to sum */
216+
node->dwork[0] = 0.0;
217+
for (int k = 0; k < bcast->m * bcast->n; k++)
218+
{
219+
node->dwork[0] += w[k];
220+
}
221+
}
222+
223+
x->eval_wsum_hess(x, node->dwork);
224+
memcpy(node->wsum_hess->x, x->wsum_hess->x, x->wsum_hess->nnz * sizeof(double));
225+
}
226+
227+
static bool is_affine(const expr *node)
228+
{
229+
return node->left->is_affine(node->left);
230+
}
231+
232+
expr *new_broadcast(expr *child, int target_d1, int target_d2)
233+
{
234+
// ---------------------------------------------------------------------------
235+
// determine broadcast type
236+
// ---------------------------------------------------------------------------
237+
broadcast_type type;
238+
int m = target_d1;
239+
int n = target_d2;
240+
241+
if (child->d1 == 1 && child->d2 == n)
242+
{
243+
type = BROADCAST_ROW;
244+
}
245+
else if (child->d1 == m && child->d2 == 1)
246+
{
247+
type = BROADCAST_COL;
248+
}
249+
else if (child->d1 == 1 && child->d2 == 1)
250+
{
251+
type = BROADCAST_SCALAR;
252+
}
253+
else
254+
{
255+
assert(false);
256+
}
257+
258+
broadcast_expr *bcast = (broadcast_expr *) calloc(1, sizeof(broadcast_expr));
259+
expr *node = (expr *) bcast;
260+
261+
// --------------------------------------------------------------------------
262+
// initialize the rest of the expression
263+
// --------------------------------------------------------------------------
264+
init_expr(node, target_d1, target_d2, child->n_vars, forward, jacobian_init,
265+
eval_jacobian, is_affine, NULL);
266+
node->left = child;
267+
expr_retain(child);
268+
node->wsum_hess_init = wsum_hess_init;
269+
node->eval_wsum_hess = eval_wsum_hess;
270+
bcast->type = type;
271+
bcast->m = m;
272+
bcast->n = n;
273+
274+
return node;
275+
}

‎src/affine/sum.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ static void eval_wsum_hess(expr *node, const double *w)
135135
}
136136
else if (axis == 1)
137137
{
138-
tile(node->dwork, w, x->d1, x->d2);
138+
tile_double(node->dwork, w, x->d1, x->d2);
139139
}
140140

141141
x->eval_wsum_hess(x, node->dwork);

‎src/utils/mini_numpy.c‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,20 @@ void repeat(double *result, const double *a, int len, int repeats)
1212
}
1313
}
1414

15-
void tile(double *result, const double *a, int len, int tiles)
15+
/* TODO: we can use memcpy here */
16+
void tile_double(double *result, const double *a, int len, int tiles)
17+
{
18+
int idx = 0;
19+
for (int i = 0; i < tiles; i++)
20+
{
21+
for (int j = 0; j < len; j++)
22+
{
23+
result[idx++] = a[j];
24+
}
25+
}
26+
}
27+
28+
void tile_int(int *result, const int *a, int len, int tiles)
1629
{
1730
int idx = 0;
1831
for (int i = 0; i < tiles; i++)

‎tests/all_tests.c‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
/* Include all test headers */
66
#include "forward_pass/affine/test_add.h"
7+
#include "forward_pass/affine/test_broadcast.h"
78
#include "forward_pass/affine/test_hstack.h"
89
#include "forward_pass/affine/test_linear_op.h"
910
#include "forward_pass/affine/test_neg.h"
@@ -13,6 +14,7 @@
1314
#include "forward_pass/composite/test_composite.h"
1415
#include "forward_pass/elementwise/test_exp.h"
1516
#include "forward_pass/elementwise/test_log.h"
17+
#include "jacobian_tests/test_broadcast.h"
1618
#include "jacobian_tests/test_composite.h"
1719
#include "jacobian_tests/test_const_scalar_mult.h"
1820
#include "jacobian_tests/test_const_vector_mult.h"
@@ -41,6 +43,7 @@
4143
#include "wsum_hess/elementwise/test_power.h"
4244
#include "wsum_hess/elementwise/test_trig.h"
4345
#include "wsum_hess/elementwise/test_xexp.h"
46+
#include "wsum_hess/test_broadcast.h"
4447
#include "wsum_hess/test_const_scalar_mult.h"
4548
#include "wsum_hess/test_const_vector_mult.h"
4649
#include "wsum_hess/test_hstack.h"
@@ -76,6 +79,9 @@ int main(void)
7679
mu_run_test(test_sum_axis_1, tests_run);
7780
mu_run_test(test_hstack_forward_vectors, tests_run);
7881
mu_run_test(test_hstack_forward_matrix, tests_run);
82+
mu_run_test(test_broadcast_row, tests_run);
83+
mu_run_test(test_broadcast_col, tests_run);
84+
mu_run_test(test_broadcast_matrix, tests_run);
7985

8086
printf("\n--- Jacobian Tests ---\n");
8187
mu_run_test(test_neg_jacobian, tests_run);
@@ -121,6 +127,9 @@ int main(void)
121127
mu_run_test(test_sum_of_index, tests_run);
122128
mu_run_test(test_promote_scalar_jacobian, tests_run);
123129
mu_run_test(test_promote_scalar_to_matrix_jacobian, tests_run);
130+
mu_run_test(test_broadcast_row_jacobian, tests_run);
131+
mu_run_test(test_broadcast_col_jacobian, tests_run);
132+
mu_run_test(test_broadcast_scalar_to_matrix_jacobian, tests_run);
124133
mu_run_test(test_wsum_hess_multiply_1, tests_run);
125134
mu_run_test(test_wsum_hess_multiply_2, tests_run);
126135
mu_run_test(test_jacobian_trace_variable, tests_run);
@@ -177,6 +186,9 @@ int main(void)
177186
mu_run_test(test_wsum_hess_left_matmul_composite, tests_run);
178187
mu_run_test(test_wsum_hess_right_matmul, tests_run);
179188
mu_run_test(test_wsum_hess_right_matmul_vector, tests_run);
189+
mu_run_test(test_wsum_hess_broadcast_row, tests_run);
190+
mu_run_test(test_wsum_hess_broadcast_col, tests_run);
191+
mu_run_test(test_wsum_hess_broadcast_scalar_to_matrix, tests_run);
180192
// This test leads to seg fault
181193
// mu_run_test(test_wsum_hess_trace_variable, tests_run);
182194

0 commit comments

Comments
 (0)