==> src/gemm.c <==1/* gemm: f32 matrix multiply throughput as a native worker program: a
2 * cache-blocked GEMM (4 x 16 register tiles, f32 fused multiply-add) over
3 * every CPU this process may use.
4 *
5 * View (fields in program.toml): 0x00 magic, 0x04 size (the periodic task's
6 * n), 0x08 last n, 0x0c last us, 0x10 last GFLOP/s, 0x14 best, 0x18 plain
7 * loop GFLOP/s, 0x1c fast path on one thread, 0x20 impl, 0x24 threads, 0x28
8 * runs, 0x1000.. C of the last check (n <= 64). Jobs: each item (OP_JOB) is
9 * one GEMM of the view's size (`setup 7:<n>` sets it), its result the
10 * microseconds it took and its job_out record (n, GFLOP/s).
11 * Ported from alpha-rust-mesh programs/gemm/src/main.rs (its Linux path). */
12#include "gemm.h"
13
14#include "mesh/abi/ops.h"
15#include "mesh/abi/ring.h"
16#include "mesh/ref/gemm.h"
17
18#include <math.h>
19#include <pthread.h>
20#include <stdatomic.h>
21#include <stdlib.h>
22#include <string.h>
23#include <time.h>
24#include <unistd.h>
25
26#define KB 256u
27#define JB 512u
28
29static size_t min_z(size_t a, size_t b) { return a < b ? a : b; }
30
31static void tile(size_t n, const float *a, const float *b, float *c, size_t r0, size_t i, size_t j, size_t k0, size_t k1) {
32 float acc[4][16];
33 for (size_t r = 0; r < 4; r++) memcpy(acc[r], c + (i + r) * n + j, sizeof acc[r]);
34 for (size_t k = k0; k < k1; k++) {
35 const float *bk = b + k * n + j;
36 for (size_t r = 0; r < 4; r++) {
37 float x = a[(r0 + i + r) * n + k];
38 for (size_t t = 0; t < 16; t++) acc[r][t] = fmaf(x, bk[t], acc[r][t]);
39 }
40 }
41 for (size_t r = 0; r < 4; r++) memcpy(c + (i + r) * n + j, acc[r], sizeof acc[r]);
42}
43
44static void edge(size_t n, const float *a, const float *b, float *c, size_t r0, size_t i, size_t ir, size_t j, size_t jw, size_t k0, size_t k1) {
45 for (size_t r = 0; r < ir; r++)
46 for (size_t k = k0; k < k1; k++) {
47 float x = a[(r0 + i + r) * n + k];
48 for (size_t t = 0; t < jw; t++) c[(i + r) * n + j + t] = fmaf(x, b[k * n + j + t], c[(i + r) * n + j + t]);
49 }
50}
51
52/* Rows [r0, r0 + rows) of C = A B: k blocks of KB, j blocks of JB, 4 x 16 tiles. */
53static void rows(size_t n, const float *a, const float *b, float *c, size_t rows_n, size_t r0) {
54 memset(c, 0, sizeof(float) * rows_n * n);
55 for (size_t k0 = 0; k0 < n; k0 += KB) {
56 size_t k1 = min_z(k0 + KB, n);
57 for (size_t j0 = 0; j0 < n; j0 += JB) {
58 size_t j1 = min_z(j0 + JB, n);
59 for (size_t i = 0; i < rows_n;) {
60 size_t ir = min_z(rows_n - i, 4);
61 for (size_t j = j0; j < j1;) {
62 size_t jw = min_z(j1 - j, 16);
63 if (ir == 4 && jw == 16) tile(n, a, b, c, r0, i, j, k0, k1);
64 else edge(n, a, b, c, r0, i, ir, j, jw, k0, k1);
65 j += jw;
66 }
67 i += ir;
68 }
69 }
70 }
71}
72
73static _Atomic size_t threads_set;
74
75size_t gemm_threads(void) {
76 size_t t = atomic_load(&threads_set);
77 if (t) return t;
78 long c = sysconf(_SC_NPROCESSORS_ONLN);
79 return c > 0 ? (size_t)c : 1;
80}
81
82typedef struct {
83 size_t n, r0, rows;
84 const float *a, *b;
85 float *c;
86 pthread_t th;
87} chunk;
88
89static void *run_chunk(void *arg) {
90 chunk *k = arg;
91 rows(k->n, k->a, k->b, k->c + k->r0 * k->n, k->rows, k->r0);
92 return NULL;
93}
94
95void gemm_fast(size_t n, const float *a, const float *b, float *c, bool single) {
96 size_t t = single ? 1 : gemm_threads();
97 if (!single) {
98 size_t cap = (n + 3) / 4;
99 if (t > cap) t = cap;
100 if (t < 1) t = 1;
101 }
102 size_t per = ((n + t - 1) / t + 3) / 4 * 4;
103 size_t nchunks = (n + per - 1) / per;
104 chunk *ks = calloc(nchunks, sizeof *ks);
105 if (!ks) return;
106 for (size_t i = 0; i < nchunks; i++) {
107 size_t r0 = i * per;
108 ks[i] = (chunk){n, r0, min_z(per, n - r0), a, b, c, 0};
109 if (pthread_create(&ks[i].th, NULL, run_chunk, &ks[i]) != 0) run_chunk(&ks[i]), ks[i].rows = 0;
110 }
111 for (size_t i = 0; i < nchunks; i++)
112 if (ks[i].rows) pthread_join(ks[i].th, NULL);
113 free(ks);
114}
115
116void gemm_plain(size_t n, const float *a, const float *b, float *c) {
117 for (size_t i = 0; i < n; i++)
118 for (size_t j = 0; j < n; j++) {
119 float s = 0.0f;
120 for (size_t k = 0; k < n; k++) s += a[i * n + k] * b[k * n + j];
121 c[i * n + j] = s;
122 }
123}
124
125/* The size, A, B and C of the last GEMM (per node thread), freed when the
126 * thread ends, as Rust's thread_local is dropped. */
127typedef struct {
128 size_t n;
129 float *a, *b, *c;
130} bufs;
131static pthread_key_t buf_key;
132static pthread_once_t buf_once = PTHREAD_ONCE_INIT;
133
134static void free_bufs(void *p) {
135 bufs *b = p;
136 free(b->a), free(b->b), free(b->c);
137 free(b);
138}
139static void make_key(void) { pthread_key_create(&buf_key, free_bufs); }
140
141/* This thread's buffers holding size n's inputs; NULL when memory runs out. */
142static bufs *with_buf(size_t n) {
143 pthread_once(&buf_once, make_key);
144 bufs *b = pthread_getspecific(buf_key);
145 if (!b) {
146 b = calloc(1, sizeof *b);
147 if (!b || pthread_setspecific(buf_key, b) != 0) { free(b); return NULL; }
148 }
149 if (b->n == n) return b;
150 free(b->a), free(b->b), free(b->c);
151 b->n = 0;
152 b->a = malloc(sizeof(float) * n * n);
153 b->b = malloc(sizeof(float) * n * n);
154 b->c = calloc(n * n, sizeof(float));
155 if (!b->a || !b->b || !b->c) {
156 free(b->a), free(b->b), free(b->c);
157 b->a = b->b = b->c = NULL;
158 return NULL;
159 }
160 mesh_gemm_inputs(n, b->a, b->b);
161 b->n = n;
162 return b;
163}
164
165static double now_s(void) {
166 struct timespec t;
167 clock_gettime(CLOCK_MONOTONIC, &t);
168 return (double)t.tv_sec + (double)t.tv_nsec / 1e9;
169}
170static float gflops(size_t n, double secs) { return (float)(2.0 * (double)n * (double)n * (double)n / secs / 1e9); }
171
172void gemm_init(const mesh_worker_ctx *ctx, void *user) {
173 (void)user;
174 mesh_ctx_set_u32(ctx, 0, GEMM_MAGIC);
175 mesh_ctx_set_u32(ctx, GEMM_V_SIZE, mesh_ctx_param(ctx) != 0 ? mesh_ctx_param(ctx) : 1024);
176 mesh_ctx_set_u32(ctx, GEMM_V_IMPL, GEMM_IMPL);
177 mesh_ctx_set_u32(ctx, GEMM_V_THREADS, (uint32_t)gemm_threads());
178}
179
180static bool size(const mesh_worker_ctx *ctx, uint32_t arg, size_t *n) {
181 uint32_t v = arg != 0 ? arg : mesh_ctx_view_u32(ctx, GEMM_V_SIZE);
182 *n = v;
183 return v >= 1 && v <= GEMM_N_MAX;
184}
185
186/* One fast GEMM of size n; its seconds, or a negative number when memory ran out. */
187static double timed(size_t n, bool single, bool plain, bufs **used) {
188 bufs *b = with_buf(n);
189 if (!b) return -1;
190 if (used) *used = b;
191 double t0 = now_s();
192 if (plain) gemm_plain(n, b->a, b->b, b->c);
193 else gemm_fast(n, b->a, b->b, b->c, single);
194 return now_s() - t0;
195}
196
197static void record(const mesh_worker_ctx *ctx, size_t n, double secs, float g) {
198 mesh_ctx_set_u32(ctx, GEMM_V_N, (uint32_t)n);
199 mesh_ctx_set_u32(ctx, GEMM_V_US, (uint32_t)(secs * 1e6));
200 mesh_ctx_set_f32(ctx, GEMM_V_GFLOPS, g);
201 if (g > mesh_ctx_view_f32(ctx, GEMM_V_BEST)) mesh_ctx_set_f32(ctx, GEMM_V_BEST, g);
202 mesh_ctx_set_u32(ctx, GEMM_V_RUNS, mesh_ctx_view_u32(ctx, GEMM_V_RUNS) + 1);
203 mesh_ctx_count(ctx, 0, 1);
204 mesh_ctx_count(ctx, 1, (uint32_t)llround(2.0 * (double)n * (double)n * (double)n / 1e9));
205}
206
207uint32_t gemm_task(const mesh_worker_ctx *ctx, uint32_t op, uint32_t arg, void *user) {
208 (void)user;
209 size_t n;
210 switch (op) {
211 case MESH_OP_JOB: {
212 if (!size(ctx, 0, &n)) return MESH_R_JOB_FAIL;
213 double secs = timed(n, false, false, NULL);
214 if (secs < 0) return MESH_R_JOB_FAIL;
215 float g = gflops(n, secs);
216 uint32_t rec = GEMM_V_JOB_OUT + 8 * (mesh_ctx_ticket(ctx) % MESH_PROG_RESULTS);
217 mesh_ctx_set_u32(ctx, rec, (uint32_t)n);
218 mesh_ctx_set_f32(ctx, rec + 4, g);
219 record(ctx, n, secs, g);
220 uint32_t us = (uint32_t)(secs * 1e6);
221 return us < MESH_R_JOB_FAIL - 1 ? us : MESH_R_JOB_FAIL - 1;
222 }
223 case GEMM_OP_SIZE:
224 if (!size(ctx, arg, &n)) return MESH_R_BAD_OP;
225 mesh_ctx_set_u32(ctx, GEMM_V_SIZE, (uint32_t)n);
226 return (uint32_t)n;
227 case MESH_OP_TASK: {
228 if (!size(ctx, arg, &n)) return MESH_R_BAD_OP;
229 double secs = timed(n, false, false, NULL);
230 if (secs < 0) return MESH_R_BAD_OP;
231 record(ctx, n, secs, gflops(n, secs));
232 return (uint32_t)(secs * 1e6);
233 }
234 case GEMM_OP_CHECK: {
235 n = arg;
236 bufs *b = NULL;
237 if (n < 1 || n > GEMM_CHECK_MAX || timed(n, false, false, &b) < 0) return MESH_R_BAD_OP;
238 mesh_view_write_f32s(mesh_ctx_view(ctx), GEMM_V_C, b->c, n * n);
239 mesh_ctx_count(ctx, 3, 1);
240 return (uint32_t)n;
241 }
242 case GEMM_OP_PLAIN:
243 case GEMM_OP_SINGLE: {
244 if (!size(ctx, arg, &n)) return MESH_R_BAD_OP;
245 double secs = timed(n, op == GEMM_OP_SINGLE, op == GEMM_OP_PLAIN, NULL);
246 if (secs < 0) return MESH_R_BAD_OP;
247 mesh_ctx_set_f32(ctx, op == GEMM_OP_PLAIN ? GEMM_V_PLAIN : GEMM_V_SINGLE, gflops(n, secs));
248 if (op == GEMM_OP_PLAIN) mesh_ctx_count(ctx, 2, 1);
249 return (uint32_t)(secs * 1e6);
250 }
251 case GEMM_OP_THREADS:
252 atomic_store(&threads_set, (size_t)arg);
253 mesh_ctx_set_u32(ctx, GEMM_V_THREADS, (uint32_t)gemm_threads());
254 return (uint32_t)gemm_threads();
255 default: return MESH_R_BAD_OP;
256 }
257}