==> src/matmul.c <==1/* matmul: the exact-mode GPU GEMM and GEMV (program.toml; the tile kernel is
2 * mesh_linalg_gemm_tile, shared with programs/gpt2x). A node is 8 warps (256
3 * threads): a 64 x 64 tile of C, each thread 4 x 4 outputs (rows ty + 16 i,
4 * columns tx + 16 j), the k dimension in slabs of 16 through shared memory.
5 * Every output is summed sequentially in k from +0.0 -- a slab only splits the
6 * loop, never reorders it -- with a multiply and an add, both correctly rounded
7 * and FTZ (the build refuses a fused op).
8 * Ported from alpha-rust-mesh programs/matmul/src/lib.rs. */
9#include "mesh/rt.h"
10#include "mesh/rt/linalg.h"
11
12#define MATMUL_MAGIC 0x4d544d4du /* "MMTM" */
13#define MATMUL_V_MAGIC 0x0u
14#define MATMUL_V_DIMS 0x4u /* M, K, N, seed A, seed B */
15#define MATMUL_V_TILES 0x18u
16#define MATMUL_V_WINDOW 0x1000u
17#define MATMUL_CHUNK 65536u
18#define MATMUL_MAXD 3072u
19#define MATMUL_MAXE ((size_t)MATMUL_MAXD * MATMUL_MAXD)
20#define MATMUL_OP_SET 3u
21#define MATMUL_OP_GEN 4u
22#define MATMUL_OP_MM 5u
23#define MATMUL_OP_MV 6u
24#define MATMUL_OP_COPY 7u
25#define MATMUL_THREADS MESH_LINALG_TILE_THREADS
26/* Scratch bytes the matrices need: A, B, C of up to MAXD x MAXD, then x and y. */
27#define MATMUL_SCRATCH_MIN (4 * (3 * MATMUL_MAXE + 2 * 4096))
28
29/* mesh_ref::block::mix. */
30static uint32_t matmul_mix(uint32_t x) {
31 x ^= x >> 16;
32 x *= 0x7feb352du;
33 x ^= x >> 15;
34 x *= 0x846ca68bu;
35 return x ^ (x >> 16);
36}
37/* mesh_ref::matmul::value. */
38static float matmul_value(uint32_t seed, uint32_t i) {
39 uint32_t h = matmul_mix(((seed * 0x9e3779b9u) + i) * 0x85ebca6bu ^ seed);
40 return (float)(h >> 8) * (1.0f / 8388608.0f) - 1.0f;
41}
42
43typedef struct {
44 float *a, *b, *c, *x, *y;
45} matmul_mats;
46
47static bool matmul_get_mats(const mesh_ctx *ctx, matmul_mats *mt) {
48 uint32_t len;
49 uint8_t *p = mesh_ctx_scratch_ptr(ctx, &len);
50 if ((size_t)len < MATMUL_SCRATCH_MIN) return false;
51 float *a = (float *)(void *)p;
52 *mt = (matmul_mats){a, a + MATMUL_MAXE, a + 2 * MATMUL_MAXE, a + 3 * MATMUL_MAXE, a + 3 * MATMUL_MAXE + 4096};
53 return true;
54}
55
56static bool matmul_dims(mesh_view v, uint32_t *m, uint32_t *k, uint32_t *n) {
57 *m = mesh_view_u32(v, MATMUL_V_DIMS);
58 *k = mesh_view_u32(v, MATMUL_V_DIMS + 4);
59 *n = mesh_view_u32(v, MATMUL_V_DIMS + 8);
60 return *m >= 1 && *m <= MATMUL_MAXD && *k >= 1 && *k <= MATMUL_MAXD && *n >= 1 && *n <= MATMUL_MAXD;
61}
62
63static void matmul_init(const mesh_ctx *ctx) {
64 mesh_view v = mesh_ctx_view(ctx);
65 if (mesh_ctx_leader(ctx)) {
66 const uint32_t d[5] = {768, 768, 768, 1, 2};
67 for (uint32_t i = 0; i < 5; i++) mesh_view_set_u32(v, MATMUL_V_DIMS + 4 * i, d[i]);
68 mesh_view_set_u32(v, MATMUL_V_MAGIC, MATMUL_MAGIC);
69 }
70}
71
72typedef struct {
73 float *c;
74 uint32_t n;
75} matmul_store;
76static void matmul_store_c(uint32_t r, uint32_t col, float v, void *user) {
77 matmul_store *s = user;
78 s->c[(size_t)r * s->n + col] = v;
79}
80
81static uint32_t matmul_task(const mesh_ctx *ctx, uint32_t op, uint32_t arg) {
82 mesh_view v = mesh_ctx_view(ctx);
83 if (op == MESH_OP_INIT) {
84 matmul_init(ctx);
85 return MESH_R_NONE;
86 }
87 if (op == MATMUL_OP_SET) {
88 if (arg >> 24 > 4) return MESH_R_BAD_OP;
89 if (mesh_ctx_leader(ctx)) mesh_view_set_u32(v, MATMUL_V_DIMS + 4 * (arg >> 24), arg & 0x00ffffffu);
90 return MESH_R_NONE;
91 }
92 /* The rest needs the node's shape, the scratch and valid dimensions (the same answer on every thread). */
93 matmul_mats mt;
94 uint32_t m, k, n;
95 if (!matmul_get_mats(ctx, &mt) || !matmul_dims(v, &m, &k, &n)) return MESH_R_BAD_OP;
96 if (mesh_ctx_threads(ctx) != MATMUL_THREADS) return MESH_R_BAD_OP;
97 bool ok;
98 switch (op) {
99 case MATMUL_OP_GEN: {
100 uint32_t which = arg >> 24, chunk = arg & 0x00ffffffu, len, seed;
101 float *dst;
102 if (which == 0) { dst = mt.a; len = m * k; seed = mesh_view_u32(v, MATMUL_V_DIMS + 12); }
103 else if (which == 1) { dst = mt.b; len = k * n; seed = mesh_view_u32(v, MATMUL_V_DIMS + 16); }
104 else if (which == 2) { dst = mt.x; len = k; seed = mesh_view_u32(v, MATMUL_V_DIMS + 16) ^ 0x55555555u; }
105 else return MESH_R_BAD_OP;
106 uint32_t lo = chunk * MATMUL_CHUNK;
107 if (lo >= len) return MESH_R_BAD_OP;
108 uint32_t hi = lo + MATMUL_CHUNK < len ? lo + MATMUL_CHUNK : len;
109 MESH_CTX_FOR_STRIDED(ctx, hi - lo, i) dst[lo + i] = matmul_value(seed, lo + i);
110 if (mesh_ctx_leader(ctx)) mesh_ctx_count(ctx, 2, 1);
111 ok = true;
112 break;
113 }
114 case MATMUL_OP_MM: {
115 mesh_shared sa, sb;
116 if (!mesh_ctx_shared(ctx, 0, MESH_LINALG_BK * MESH_LINALG_TILE, 4, &sa) ||
117 !mesh_ctx_shared(ctx, 4 * MESH_LINALG_BK * MESH_LINALG_TILE, MESH_LINALG_BK * MESH_LINALG_TILE, 4, &sb))
118 return MESH_R_BAD_OP;
119 matmul_store st = {mt.c, n};
120 ok = mesh_linalg_gemm_tile(ctx, mt.a, mt.b, m, k, n, arg, sa, sb, matmul_store_c, &st);
121 if (ok && mesh_ctx_leader(ctx)) {
122 mesh_ctx_count(ctx, 0, 1);
123 mesh_view w;
124 if (mesh_view_sub(v, MATMUL_V_TILES, 4, &w)) mesh_rt_atom_add(mesh_view_ptr(w), 1);
125 }
126 break;
127 }
128 case MATMUL_OP_MV: {
129 /* Rows arg x 64 ..: warp w takes rows + 8 w .. + 8, one row at a time. */
130 uint32_t r0 = arg * MESH_LINALG_TILE;
131 if (r0 >= m) return MESH_R_BAD_OP;
132 uint32_t l = mesh_ctx_lane(ctx);
133 for (uint32_t rr = 0; rr < 8; rr++) {
134 uint32_t r = r0 + 8 * mesh_ctx_warp(ctx) + rr;
135 if (r >= m) break;
136 float s = 0.0f;
137 for (uint32_t kk = l; kk < k; kk += 32) s = s + mt.a[(size_t)r * k + kk] * mt.x[kk];
138 s = mesh_rt_warp_reduce_f32(s, mesh_rt_add_f32);
139 if (l == 0) mt.y[r] = s;
140 }
141 if (mesh_ctx_leader(ctx)) mesh_ctx_count(ctx, 1, 1);
142 ok = true;
143 break;
144 }
145 case MATMUL_OP_COPY: {
146 uint32_t which = arg >> 24, chunk = arg & 0x00ffffffu, len;
147 const float *src;
148 if (which == 0) { src = mt.c; len = m * n; }
149 else if (which == 1) { src = mt.y; len = m; }
150 else return MESH_R_BAD_OP;
151 uint32_t lo = chunk * MATMUL_CHUNK;
152 if (lo >= len) return MESH_R_BAD_OP;
153 uint32_t hi = lo + MATMUL_CHUNK < len ? lo + MATMUL_CHUNK : len;
154 MESH_CTX_FOR_STRIDED(ctx, hi - lo, i) mesh_view_set_f32(v, MATMUL_V_WINDOW + 4 * i, src[lo + i]);
155 if (mesh_ctx_leader(ctx)) mesh_ctx_count(ctx, 3, 1);
156 ok = true;
157 break;
158 }
159 default: ok = false; break;
160 }
161 return ok ? MESH_R_NONE : MESH_R_BAD_OP;
162}
163
164MESH_PROGRAM(matmul_init, matmul_task);