==> src/forward.c <==1/* GPT-2's forward pass on the CPU (mesh-ref gpt2.rs `forward_layers`, `generate`, `par_linear`,
2 * `par_rows`, `project_vocab`). The threads only split independent outputs. */
3#include "mesh/ref/gpt2.h"
4
5#include "mesh/ftz.h"
6#include "mesh/ref/gpt2_attend.h"
7#include "mesh/ref/gpt2_linear.h"
8#include "mesh/ref/gpt2_logits.h"
9#include "mesh/ref/gpt2_norm.h"
10#include "mesh/ref/matmul.h"
11
12#include <pthread.h>
13#include <stdatomic.h>
14#include <stdlib.h>
15#include <string.h>
16
17#define W MESH_GPT2_WIDTH
18#define HD MESH_GPT2_HEAD_DIM
19
20/* ---- a pool of units, each an index, taken from a shared counter by up to `threads` threads ---- */
21
22typedef void (*unit_fn)(size_t unit, void *ctx);
23
24typedef struct {
25 unit_fn fn;
26 void *ctx;
27 size_t units;
28 atomic_size_t next;
29} pool;
30
31static void *pool_worker(void *arg) {
32 pool *p = arg;
33 for (size_t u; (u = atomic_fetch_add(&p->next, 1)) < p->units;) p->fn(u, p->ctx);
34 return NULL;
35}
36
37static void run_units(size_t units, size_t threads, unit_fn fn, void *ctx) {
38 pool p = {fn, ctx, units, 0};
39 if (threads > units) threads = units;
40 if (threads <= 1) {
41 pool_worker(&p);
42 return;
43 }
44 pthread_t tid[threads];
45 size_t started = 0;
46 for (size_t i = 0; i < threads; i++)
47 if (pthread_create(&tid[started], NULL, pool_worker, &p) == 0) started++;
48 pool_worker(&p); /* this thread works too, and finishes whatever no thread took */
49 for (size_t i = 0; i < started; i++) pthread_join(tid[i], NULL);
50}
51
52/* ---- a linear layer over rows: (row, column block) units ---- */
53
54typedef struct {
55 const float *x, *w, *b;
56 size_t kd, nout, bw, blocks;
57 mesh_gpt2_epi epi;
58 float *out;
59} linear;
60
61static void linear_unit(size_t u, void *ctx) {
62 linear *c = ctx;
63 size_t r = u / c->blocks, c0 = (u % c->blocks) * c->bw;
64 size_t n = c0 + c->bw <= c->nout ? c->bw : c->nout - c0;
65 mesh_gpt2_project_cols(c->x + r * c->kd, c->kd, c->w, c->b, c->nout, c->epi, c0, n, c->out + r * c->nout + c0);
66}
67
68static void par_linear(const float *x, size_t kd, const float *w, const float *b, size_t nout, mesh_gpt2_epi epi,
69 float *out, size_t rows, size_t threads) {
70 size_t blocks = ((threads ? threads : 1) * 4 + rows - 1) / rows;
71 size_t max_blocks = (nout + 63) / 64;
72 if (blocks < 1) blocks = 1;
73 if (blocks > max_blocks) blocks = max_blocks;
74 size_t bw = (nout + blocks - 1) / blocks;
75 blocks = (nout + bw - 1) / bw;
76 linear c = {x, w, b, kd, nout, bw, blocks, epi, out};
77 run_units(rows * blocks, threads, linear_unit, &c);
78}
79
80/* ---- row-wise steps ---- */
81
82typedef struct {
83 const float *x, *g, *b;
84 float *out;
85} norm_ctx;
86
87static void norm_unit(size_t r, void *ctx) {
88 norm_ctx *c = ctx;
89 mesh_gpt2_normalize_row(c->x + r * W, c->g, c->b, c->out + r * W);
90}
91
92typedef struct {
93 const float *qkv, *kc, *vc;
94 size_t p0;
95 float *y;
96} attend_ctx;
97
98static void attend_unit(size_t r, void *ctx) {
99 attend_ctx *c = ctx;
100 for (size_t hh = 0; hh < MESH_GPT2_HEADS; hh++)
101 mesh_gpt2_attend_row(c->qkv + r * 3 * W + hh * HD, c->p0 + r + 1, c->kc + hh * HD, W, c->vc + hh * HD, W,
102 c->y + r * W + hh * HD);
103}
104
105typedef struct {
106 const float *wte, *h;
107 float *out;
108} vocab_ctx;
109
110static void vocab_unit(size_t i, void *ctx) {
111 vocab_ctx *c = ctx;
112 mesh_matvec(c->wte + i * W, c->h, 1, W, &c->out[i]);
113}
114
115void mesh_gpt2_project_vocab(const float *wte, const float *h, size_t threads, float *out) {
116 vocab_ctx c = {wte, h, out};
117 run_units(MESH_GPT2_VOCAB, threads, vocab_unit, &c);
118}
119
120bool mesh_gpt2_forward(const float *w, const mesh_gpt2_layout *lay, const uint32_t *tokens, size_t t,
121 mesh_gpt2_kv *kv, size_t threads, float *logits) {
122 return mesh_gpt2_forward_layers(w, lay, tokens, t, kv, threads, NULL, NULL, logits);
123}
124
125bool mesh_gpt2_forward_layers(const float *w, const mesh_gpt2_layout *lay, const uint32_t *tokens, size_t t,
126 mesh_gpt2_kv *kv, size_t threads, mesh_gpt2_after_layer after, void *ctx,
127 float *logits) {
128 size_t p0 = mesh_gpt2_kv_len(kv);
129 if (t == 0 || p0 + t > MESH_GPT2_CTX || !mesh_gpt2_kv_reserve(kv, p0 + t)) return false;
130 float *x = malloc(t * W * sizeof *x), *h = malloc(t * W * sizeof *h), *qkv = malloc(t * 3 * W * sizeof *qkv),
131 *y = malloc(t * W * sizeof *y), *f = malloc(t * MESH_GPT2_FF * sizeof *f), *last = malloc(W * sizeof *last);
132 bool ok = x && h && qkv && y && f && last;
133 if (ok) {
134 for (size_t r = 0; r < t; r++) {
135 size_t e = lay->wte + (size_t)tokens[r] * W, p = lay->wpe + (p0 + r) * W;
136 for (size_t i = 0; i < W; i++) x[r * W + i] = mesh_ftz_add(w[e + i], w[p + i]);
137 }
138 for (size_t l = 0; l < MESH_GPT2_LAYERS; l++) {
139 const mesh_gpt2_block_off *b = &lay->blocks[l];
140 norm_ctx n1 = {x, w + b->ln1_g, w + b->ln1_b, h};
141 run_units(t, threads, norm_unit, &n1);
142 par_linear(h, W, w + b->attn_w, w + b->attn_b, 3 * W, MESH_GPT2_EPI_BIAS, qkv, t, threads);
143 for (size_t r = 0; r < t; r++) {
144 memcpy(kv->k[l] + (p0 + r) * W, qkv + r * 3 * W + W, W * sizeof(float));
145 memcpy(kv->v[l] + (p0 + r) * W, qkv + r * 3 * W + 2 * W, W * sizeof(float));
146 }
147 kv->rows[l] = p0 + t;
148 attend_ctx at = {qkv, kv->k[l], kv->v[l], p0, y};
149 run_units(t, threads, attend_unit, &at);
150 par_linear(y, W, w + b->proj_w, w + b->proj_b, W, MESH_GPT2_EPI_BIAS_RESIDUAL, x, t, threads);
151 norm_ctx n2 = {x, w + b->ln2_g, w + b->ln2_b, h};
152 run_units(t, threads, norm_unit, &n2);
153 par_linear(h, W, w + b->fc_w, w + b->fc_b, MESH_GPT2_FF, MESH_GPT2_EPI_BIAS_GELU, f, t, threads);
154 par_linear(f, MESH_GPT2_FF, w + b->fc2_w, w + b->fc2_b, W, MESH_GPT2_EPI_BIAS_RESIDUAL, x, t, threads);
155 if (after) after(l, kv, ctx);
156 }
157 mesh_gpt2_normalize_row(x + (t - 1) * W, w + lay->lnf_g, w + lay->lnf_b, last);
158 mesh_gpt2_project_vocab(w + lay->wte, last, threads, logits);
159 }
160 free(x), free(h), free(qkv), free(y), free(f), free(last);
161 return ok;
162}
163
164bool mesh_gpt2_generate(const float *w, const mesh_gpt2_layout *lay, const uint32_t *prompt, size_t n,
165 size_t steps, size_t threads, uint32_t *toks, uint64_t *digest, mesh_gpt2_kv *kv) {
166 mesh_gpt2_kv_init(kv);
167 float *l = malloc(MESH_GPT2_VOCAB * sizeof *l);
168 if (!l || !mesh_gpt2_forward(w, lay, prompt, n, kv, threads, l)) {
169 free(l);
170 return false;
171 }
172 uint64_t h = mesh_gpt2_digest_extend(MESH_GPT2_DIGEST_SEED, l, MESH_GPT2_VOCAB);
173 toks[0] = mesh_gpt2_argmax(l, MESH_GPT2_VOCAB);
174 for (size_t s = 0; s < steps; s++) {
175 if (!mesh_gpt2_forward(w, lay, &toks[s], 1, kv, threads, l)) {
176 free(l);
177 return false;
178 }
179 h = mesh_gpt2_digest_extend(h, l, MESH_GPT2_VOCAB);
180 toks[s + 1] = mesh_gpt2_argmax(l, MESH_GPT2_VOCAB);
181 }
182 free(l);
183 *digest = h;
184 return true;
185}