1/* GPT-2 124M's weight layout (mesh-ref gpt2.rs `Layout`). */
2#include "mesh/ref/gpt2_shape.h"
3
4void mesh_gpt2_layout_init(mesh_gpt2_layout *l) {
5    size_t at = 0;
6#define TAKE(n) (at += (n), at - (n))
7    l->wte = TAKE((size_t)MESH_GPT2_VOCAB * MESH_GPT2_WIDTH);
8    l->wpe = TAKE((size_t)MESH_GPT2_CTX * MESH_GPT2_WIDTH);
9    for (int i = 0; i < MESH_GPT2_LAYERS; i++) {
10        mesh_gpt2_block_off *b = &l->blocks[i];
11        b->ln1_g = TAKE(MESH_GPT2_WIDTH);
12        b->ln1_b = TAKE(MESH_GPT2_WIDTH);
13        b->attn_w = TAKE((size_t)MESH_GPT2_WIDTH * 3 * MESH_GPT2_WIDTH);
14        b->attn_b = TAKE(3 * MESH_GPT2_WIDTH);
15        b->proj_w = TAKE((size_t)MESH_GPT2_WIDTH * MESH_GPT2_WIDTH);
16        b->proj_b = TAKE(MESH_GPT2_WIDTH);
17        b->ln2_g = TAKE(MESH_GPT2_WIDTH);
18        b->ln2_b = TAKE(MESH_GPT2_WIDTH);
19        b->fc_w = TAKE((size_t)MESH_GPT2_WIDTH * MESH_GPT2_FF);
20        b->fc_b = TAKE(MESH_GPT2_FF);
21        b->fc2_w = TAKE((size_t)MESH_GPT2_FF * MESH_GPT2_WIDTH);
22        b->fc2_b = TAKE(MESH_GPT2_WIDTH);
23    }
24    l->lnf_g = TAKE(MESH_GPT2_WIDTH);
25    l->lnf_b = TAKE(MESH_GPT2_WIDTH);
26    l->total = at;
27#undef TAKE
28}
29
30static void put(char *name, size_t *at, const char *s) {
31    while (*s && *at < 31) name[(*at)++] = *s++;
32    name[*at] = 0;
33}
34
35bool mesh_gpt2_tensor(const mesh_gpt2_layout *l, size_t i, char name[32], size_t *off, size_t *count) {
36    size_t at = 0;
37    name[0] = 0;
38    if (i == 0) return put(name, &at, "wte.weight"), *off = l->wte, *count = (size_t)MESH_GPT2_VOCAB * MESH_GPT2_WIDTH, true;
39    if (i == 1) return put(name, &at, "wpe.weight"), *off = l->wpe, *count = (size_t)MESH_GPT2_CTX * MESH_GPT2_WIDTH, true;
40    if (i == MESH_GPT2_TENSORS - 2) return put(name, &at, "ln_f.weight"), *off = l->lnf_g, *count = MESH_GPT2_WIDTH, true;
41    if (i == MESH_GPT2_TENSORS - 1) return put(name, &at, "ln_f.bias"), *off = l->lnf_b, *count = MESH_GPT2_WIDTH, true;
42    if (i >= MESH_GPT2_TENSORS) return false;
43    size_t layer = (i - 2) / 12, k = (i - 2) % 12;
44    const mesh_gpt2_block_off *b = &l->blocks[layer];
45    static const char *const names[12] = {"ln_1.weight",        "ln_1.bias",          "attn.c_attn.weight",
46                                          "attn.c_attn.bias",   "attn.c_proj.weight", "attn.c_proj.bias",
47                                          "ln_2.weight",        "ln_2.bias",          "mlp.c_fc.weight",
48                                          "mlp.c_fc.bias",      "mlp.c_proj.weight",  "mlp.c_proj.bias"};
49    const size_t offs[12] = {b->ln1_g, b->ln1_b, b->attn_w, b->attn_b, b->proj_w, b->proj_b,
50                             b->ln2_g, b->ln2_b, b->fc_w,   b->fc_b,   b->fc2_w,  b->fc2_b};
51    const size_t counts[12] = {MESH_GPT2_WIDTH, MESH_GPT2_WIDTH, (size_t)MESH_GPT2_WIDTH * 3 * MESH_GPT2_WIDTH,
52                               3 * MESH_GPT2_WIDTH, (size_t)MESH_GPT2_WIDTH * MESH_GPT2_WIDTH, MESH_GPT2_WIDTH,
53                               MESH_GPT2_WIDTH, MESH_GPT2_WIDTH, (size_t)MESH_GPT2_WIDTH * MESH_GPT2_FF,
54                               MESH_GPT2_FF, (size_t)MESH_GPT2_FF * MESH_GPT2_WIDTH, MESH_GPT2_WIDTH};
55    char num[4] = {0};
56    if (layer >= 10) num[0] = (char)('0' + layer / 10), num[1] = (char)('0' + layer % 10);
57    else num[0] = (char)('0' + layer);
58    put(name, &at, "h.");
59    put(name, &at, num);
60    put(name, &at, ".");
61    put(name, &at, names[k]);
62    *off = offs[k];
63    *count = counts[k];
64    return true;
65}