==> tests/gpt2x.c <==1/* gpt2x on the simulator, op by op against mesh_ref::gpt2's pieces, bitwise:
2 * the window (PUT, COPY), EMBED, NORM (normalize_row), GEMM and GEMV with
3 * each epilogue (project_row, project_cols), KV and ATTN (attend_row), LOGITS
4 * (a lane's sequential sum, then the warp butterfly), and the scalar
5 * functions themselves (exp, tanh, gelu, sqrt). Synthetic data: a seeded
6 * hash, not the real weights. */
7#include "mesh/check.h"
8#include "mesh/ref/block.h"
9#include "mesh/ref/gpt2_attend.h"
10#include "mesh/ref/gpt2_linear.h"
11#include "mesh/ref/gpt2_norm.h"
12#include "mesh/ref/gpt2_scalar.h"
13#include "mesh/rt/sim_mailbox.h"
14
15#include "../src/gpt2x.c"
16
17#include <stdlib.h>
18
19#define W GX_WIDTH
20static uint8_t *mb, *view;
21static float *sc;
22
23/* Scratch layout (floats). */
24#define O_WTE 0u
25#define O_WPE (O_WTE + GX_VOCAB * W)
26#define O_X (O_WPE + GX_CTX * W)
27#define O_N (O_X + 8 * W)
28#define O_G (O_N + 8 * W)
29#define O_B (O_G + W)
30#define O_A (O_B + W)
31#define O_MW (O_A + 8 * 300)
32#define O_BIAS (O_MW + 300 * 200)
33#define O_C (O_BIAS + 200)
34#define O_QKV (O_C + 8 * 200)
35#define O_KC (O_QKV + 8 * 3 * W)
36#define O_VC (O_KC + GX_CTX * W)
37#define O_Y (O_VC + GX_CTX * W)
38#define O_LOG (O_Y + 8 * W)
39#define O_END (O_LOG + GX_VOCAB)
40
41static float val(uint32_t seed, uint32_t i) {
42 uint32_t h = mesh_block_mix(seed * 0x9e3779b9u + i * 0x85ebca6bu);
43 return (float)(h >> 8) * (1.0f / 8388608.0f) - 1.0f;
44}
45static void fill(uint32_t off, uint32_t n, uint32_t seed, float scale) { for (uint32_t i = 0; i < n; i++) sc[off + i] = val(seed, i) * scale; }
46static void params(const uint32_t *kv, size_t n) {
47 for (uint32_t i = 0; i < MESH_GPT2X_PARAM_WORDS; i++) mesh_sim_set_u32(view + MESH_GPT2X_X_PARAMS + 4 * i, 0);
48 for (size_t i = 0; i < n; i += 2) mesh_sim_set_u32(view + MESH_GPT2X_X_PARAMS + 4 * kv[i], kv[i + 1]);
49}
50static uint32_t run(uint32_t op, uint32_t arg) { return mesh_sim_call(mb, 0, op, arg); }
51static uint32_t bits(float f) { return mesh_f32_bits(f); }
52
53static void gemm_case(uint32_t m, uint32_t k, uint32_t n, uint32_t epi, bool gemv) {
54 fill(O_A, m * k, 11 + epi, 1.0f);
55 fill(O_MW, k * n, 12 + epi, 0.25f);
56 fill(O_BIAS, n, 13 + epi, 0.5f);
57 fill(O_C, m * n, 14 + epi, 1.0f);
58 float *want = malloc(sizeof(float) * m * n);
59 memcpy(want, sc + O_C, sizeof(float) * m * n);
60 for (uint32_t r = 0; r < m; r++) mesh_gpt2_project_row(sc + O_A + r * k, k, sc + O_MW, sc + O_BIAS, n, (mesh_gpt2_epi)epi, want + r * n);
61 const uint32_t kv[] = {MESH_GPT2X_P_A, O_A, MESH_GPT2X_P_B, O_MW, MESH_GPT2X_P_BIAS, O_BIAS, MESH_GPT2X_P_C, O_C,
62 MESH_GPT2X_P_M, m, MESH_GPT2X_P_K, k, MESH_GPT2X_P_N, n, MESH_GPT2X_P_EPI, epi};
63 params(kv, sizeof kv / sizeof *kv);
64 if (gemv) {
65 uint32_t nb = (n + 63) / 64;
66 for (uint32_t a = 0; a < m * nb; a++) CHECK_U(run(MESH_GPT2X_OP_GEMV, a), MESH_R_NONE);
67 CHECK_U(run(MESH_GPT2X_OP_GEMV, m * nb), MESH_R_BAD_OP);
68 } else {
69 uint32_t tiles = ((m + 63) / 64) * ((n + 63) / 64);
70 for (uint32_t t = 0; t < tiles; t++) CHECK_U(run(MESH_GPT2X_OP_GEMM, t), MESH_R_NONE);
71 CHECK_U(run(MESH_GPT2X_OP_GEMM, tiles), MESH_R_BAD_OP);
72 }
73 uint32_t bad = 0;
74 for (uint32_t i = 0; i < m * n; i++) bad += bits(sc[O_C + i]) != bits(want[i]);
75 CHECK_U(bad, 0);
76 free(want);
77}
78
79int main(void) {
80 /* The scalar functions against the reference over a sweep. */
81 uint32_t bad = 0;
82 for (int i = -4000; i <= 4000; i++) {
83 float x = (float)i * 0.0251f;
84 bad += bits(gx_exp_poly(x)) != bits(mesh_gpt2_exp_poly(x));
85 bad += bits(gx_tanh_poly(x)) != bits(mesh_gpt2_tanh_poly(x));
86 bad += bits(gx_gelu_tanh(x)) != bits(mesh_gpt2_gelu_tanh(x));
87 if (x > 0) bad += bits(gx_sqrt_newton(x)) != bits(mesh_gpt2_sqrt_newton(x));
88 }
89 CHECK_U(bad, 0);
90
91 mb = mesh_sim_mailbox();
92 sc = calloc(O_END, sizeof(float));
93 CHECK(sc != NULL);
94 mesh_sim_prog p = {.name = "gpt2x", .view_bytes = MESH_GPT2X_VIEW_BYTES, .nodes = 2, .warps = 8, .shared_bytes = 0x4200,
95 .scratch = (uint8_t *)sc, .scratch_bytes = 4u * O_END};
96 view = mesh_sim_install(mb, 0, &p);
97 mesh_sim_launch *l = mesh_sim_launch_prog(mb, 0, &p, mesh_main);
98 CHECK_U(run(MESH_OP_INIT, 0), MESH_R_NONE);
99 CHECK_U(mesh_sim_u32(view + MESH_GPT2X_X_MAGIC), MESH_GPT2X_MAGIC);
100 /* PUT then COPY through the window. */
101 for (uint32_t i = 0; i < 100; i++) mesh_sim_set_f32(view + MESH_GPT2X_X_WINDOW + 4 * i, (float)i * 1.5f);
102 const uint32_t put[] = {MESH_GPT2X_P_OFF, O_X, MESH_GPT2X_P_LEN, 100};
103 params(put, 4);
104 CHECK_U(run(MESH_GPT2X_OP_PUT, 0), MESH_R_NONE);
105 CHECK_U(bits(sc[O_X + 99]), bits(148.5f));
106 for (uint32_t i = 0; i < 100; i++) mesh_sim_set_f32(view + MESH_GPT2X_X_WINDOW + 4 * i, 0.0f);
107 CHECK_U(run(MESH_GPT2X_OP_COPY, 0), MESH_R_NONE);
108 CHECK_U(mesh_sim_u32(view + MESH_GPT2X_X_WINDOW + 4 * 7), bits(10.5f));
109 CHECK_U(run(MESH_GPT2X_OP_COPY, 1), MESH_R_BAD_OP);
110 /* EMBED: token rows plus position rows. */
111 fill(O_WTE, 4 * W, 1, 0.1f);
112 fill(O_WTE + 50000 * W, W, 2, 0.1f);
113 fill(O_WPE, 8 * W, 3, 0.01f);
114 mesh_sim_set_u32(view + MESH_GPT2X_X_TOKENS, 3);
115 mesh_sim_set_u32(view + MESH_GPT2X_X_TOKENS + 4, 50000);
116 const uint32_t emb[] = {MESH_GPT2X_P_T, 2, MESH_GPT2X_P_P0, 5, MESH_GPT2X_P_WTE, O_WTE, MESH_GPT2X_P_WPE, O_WPE, MESH_GPT2X_P_X, O_X};
117 params(emb, 10);
118 CHECK_U(run(MESH_GPT2X_OP_EMBED, 0), MESH_R_NONE);
119 CHECK_U(run(MESH_GPT2X_OP_EMBED, 1), MESH_R_NONE);
120 CHECK_U(run(MESH_GPT2X_OP_EMBED, 2), MESH_R_BAD_OP);
121 bad = 0;
122 for (uint32_t i = 0; i < W; i++) {
123 bad += bits(sc[O_X + i]) != bits(sc[O_WTE + 3 * W + i] + sc[O_WPE + 5 * W + i]);
124 bad += bits(sc[O_X + W + i]) != bits(sc[O_WTE + 50000 * W + i] + sc[O_WPE + 6 * W + i]);
125 }
126 CHECK_U(bad, 0);
127 /* NORM. */
128 fill(O_X, 3 * W, 4, 2.0f);
129 fill(O_G, W, 5, 1.0f);
130 fill(O_B, W, 6, 0.1f);
131 const uint32_t nrm[] = {MESH_GPT2X_P_A, O_X, MESH_GPT2X_P_B, O_G, MESH_GPT2X_P_BIAS, O_B, MESH_GPT2X_P_C, O_N, MESH_GPT2X_P_T, 3};
132 params(nrm, 10);
133 for (uint32_t r = 0; r < 3; r++) CHECK_U(run(MESH_GPT2X_OP_NORM, r), MESH_R_NONE);
134 float want[W];
135 bad = 0;
136 for (uint32_t r = 0; r < 3; r++) {
137 mesh_gpt2_normalize_row(sc + O_X + r * W, sc + O_G, sc + O_B, want);
138 for (uint32_t i = 0; i < W; i++) bad += bits(sc[O_N + r * W + i]) != bits(want[i]);
139 }
140 CHECK_U(bad, 0);
141 /* GEMM and GEMV with each epilogue (dimensions off the tile grid). */
142 for (uint32_t epi = 0; epi < 3; epi++) {
143 gemm_case(5, 300, 130, epi, false);
144 gemm_case(3, 300, 200, epi, true);
145 }
146 /* KV and ATTN: 4 rows at positions 9..12, causal over the cache. */
147 fill(O_QKV, 4 * 3 * W, 7, 1.0f);
148 fill(O_KC, 9 * W, 8, 1.0f);
149 fill(O_VC, 9 * W, 9, 1.0f);
150 const uint32_t kvp[] = {MESH_GPT2X_P_T, 4, MESH_GPT2X_P_P0, 9, MESH_GPT2X_P_QKV, O_QKV, MESH_GPT2X_P_KC, O_KC, MESH_GPT2X_P_VC, O_VC,
151 MESH_GPT2X_P_Y, O_Y, MESH_GPT2X_P_QB, 3};
152 params(kvp, 14);
153 for (uint32_t r = 0; r < 4; r++) CHECK_U(run(MESH_GPT2X_OP_KV, r), MESH_R_NONE);
154 uint32_t nq = 2; /* ceil(4 / 3) */
155 for (uint32_t a = 0; a < GX_HEADS * nq; a++) CHECK_U(run(MESH_GPT2X_OP_ATTN, a), MESH_R_NONE);
156 CHECK_U(run(MESH_GPT2X_OP_ATTN, GX_HEADS * nq), MESH_R_BAD_OP);
157 bad = 0;
158 for (uint32_t r = 0; r < 4; r++)
159 for (uint32_t h = 0; h < GX_HEADS; h++) {
160 float out[GX_HEAD_DIM];
161 uint32_t hd = h * GX_HEAD_DIM;
162 CHECK(mesh_gpt2_attend_row(sc + O_QKV + r * 3 * W + hd, 9 + r + 1, sc + O_KC + hd, W, sc + O_VC + hd, W, out));
163 for (uint32_t d = 0; d < GX_HEAD_DIM; d++) bad += bits(sc[O_Y + r * W + hd + d]) != bits(out[d]);
164 }
165 CHECK_U(bad, 0);
166 /* LOGITS: the first blocks and the last (rows past VOCAB clamped). */
167 fill(O_N, W, 10, 1.0f);
168 fill(O_WTE, 600 * W, 15, 0.05f);
169 const uint32_t lg[] = {MESH_GPT2X_P_A, O_N, MESH_GPT2X_P_B, O_WTE, MESH_GPT2X_P_C, O_LOG};
170 params(lg, 6);
171 uint32_t last = (GX_VOCAB - 1) / MESH_GPT2X_LOGIT_ROWS;
172 const uint32_t args[3] = {0, 1, last};
173 for (int i = 0; i < 3; i++) CHECK_U(run(MESH_GPT2X_OP_LOGITS, args[i]), MESH_R_NONE);
174 CHECK_U(run(MESH_GPT2X_OP_LOGITS, last + 1), MESH_R_BAD_OP);
175 bad = 0;
176 for (int ai = 0; ai < 3; ai++)
177 for (uint32_t r = args[ai] * MESH_GPT2X_LOGIT_ROWS; r < (args[ai] + 1) * MESH_GPT2X_LOGIT_ROWS && r < GX_VOCAB; r++) {
178 float lanes[32];
179 for (uint32_t ln = 0; ln < 32; ln++) {
180 float a = 0.0f;
181 for (uint32_t k = ln; k < W; k += 32) a = a + sc[O_WTE + r * W + k] * sc[O_N + k];
182 lanes[ln] = a;
183 }
184 mesh_block_warp_reduce_f32(lanes, mesh_block_add_f32);
185 bad += bits(sc[O_LOG + r]) != bits(lanes[0]);
186 }
187 CHECK_U(bad, 0);
188 CHECK_U(run(MESH_OP_TASK, 0), MESH_R_BAD_OP);
189 CHECK_U(mesh_sim_u32(mesh_sim_ctl(mb, 0) + MESH_P_FAULT), 0);
190 mesh_sim_halt(mb, 0, l);
191 free(sc);
192 mesh_sim_mailbox_free(mb);
193 return mesh_check_done();
194}