1/* @ajax/mesh-rt-linalg: dense linear algebra at node scope in the exact mode (a GEMM tile and
2 * a GEMV slab-staged over shared memory).
3 *
4 * Every output is the sequential sum over k from +0.0 of a correctly rounded
5 * multiply, then a correctly rounded add (the build refuses a fused op), so a
6 * tile's bits do not depend on the tile shape, the node or the GPU.
7 * mesh_ref::matmul and mesh_ref::gpt2 are the CPU references. Used by
8 * programs/matmul (GEMM, GEMV) and programs/gpt2x (GPT-2's projections).
9 * Ported from alpha-rust-mesh crates/mesh-rt/src/linalg.rs. */
10#ifndef MESH_RT_LINALG_H
11#define MESH_RT_LINALG_H
1213#include <stdbool.h>
14#include <stdint.h>
1516#include "mesh/rt/block.h"
17#include "mesh/rt/ctx.h"
1819/* C is computed in TILE x TILE tiles, k in slabs of BK through shared memory. */
20#define MESH_LINALG_TILE 64u
21#define MESH_LINALG_BK 16u
22/* The tile kernel's node: 256 threads (8 warps), each 4 x 4 outputs. */
23#define MESH_LINALG_TILE_THREADS 256u
24/* Shared f32s a tile needs: A's and B's slab. */
25#define MESH_LINALG_TILE_SHARED (2u * MESH_LINALG_BK * MESH_LINALG_TILE)
26/* Columns a GEMV node computes, and the rows of B it stages per slab. */
27#define MESH_LINALG_GEMV_COLS 64u
28#define MESH_LINALG_GEMV_SLAB 64u
2930typedef void (*mesh_linalg_store2)(uint32_t r, uint32_t c, float v, void *user);
31typedef void (*mesh_linalg_store1)(uint32_t c, float v, void *user);
3233MESH_RT_INLINE uint32_t mesh_linalg_div_ceil(uint32_t a, uint32_t b) { return a / b + (a % b != 0); }
34MESH_RT_INLINE uint32_t mesh_linalg_min(uint32_t a, uint32_t b) { return a < b ? a : b; }
3536/* C tile `tile` (row-major tiles) of A (m x k, row-major) times B (k x n,
37 * row-major), on every thread of a 256-thread node; store(r, c, acc) for each
38 * of the thread's outputs inside C (rows ty + 16 i, columns tx + 16 j). sa, sb:
39 * BK x TILE f32s each. False (nothing done) for a tile outside C.
40 * a and b hold m x k and k x n f32s. */
41MESH_RT_INLINE bool mesh_linalg_gemm_tile(const mesh_ctx *ctx, const float *a, const float *b, uint32_t m, uint32_t k, uint32_t n,
42 uint32_t tile, mesh_shared sa, mesh_shared sb, mesh_linalg_store2 store, void *user) {
43 const uint32_t TILE = MESH_LINALG_TILE, BK = MESH_LINALG_BK, TT = MESH_LINALG_TILE_THREADS;
44 uint32_t tn = mesh_linalg_div_ceil(n, TILE);
45 if (tile >= mesh_linalg_div_ceil(m, TILE) * tn) return false;
46 uint32_t r0 = (tile / tn) * TILE, c0 = (tile % tn) * TILE;
47 uint32_t t = mesh_ctx_thread(ctx);
48 uint32_t ty = t / 16, tx = t % 16;
49 float acc[4][4];
50 for (int i = 0; i < 4; i++)
51 for (int j = 0; j < 4; j++) acc[i][j] = 0.0f;
52 for (uint32_t k0 = 0; k0 < k; k0 += BK) {
53 uint32_t kmax = mesh_linalg_min(k - k0, BK);
54 for (uint32_t q = 0; q < 4; q++) {
55 uint32_t e = t + TT * q;
56 /* A: 64 rows x 16 k, stored k-major (As[kk][row]); B: 16 k x 64 columns. */
57 uint32_t row = e / BK, kk = e % BK;
58 float va = (r0 + row < m && kk < kmax) ? a[(size_t)(r0 + row) * k + k0 + kk] : 0.0f;
59 mesh_shared_set_unchecked_f32(sa, kk * TILE + row, va);
60 uint32_t kk2 = e / TILE, col = e % TILE;
61 float vb = (kk2 < kmax && c0 + col < n) ? b[(size_t)(k0 + kk2) * n + c0 + col] : 0.0f;
62 mesh_shared_set_unchecked_f32(sb, kk2 * TILE + col, vb);
63 }
64 mesh_ctx_block_barrier(ctx);
65 for (uint32_t kk = 0; kk < kmax; kk++) {
66 float av[4], bv[4];
67 for (uint32_t i = 0; i < 4; i++) av[i] = mesh_shared_get_unchecked_f32(sa, kk * TILE + ty + 16 * i);
68 for (uint32_t j = 0; j < 4; j++) bv[j] = mesh_shared_get_unchecked_f32(sb, kk * TILE + tx + 16 * j);
69 for (int i = 0; i < 4; i++)
70 for (int j = 0; j < 4; j++) acc[i][j] = acc[i][j] + av[i] * bv[j];
71 }
72 mesh_ctx_block_barrier(ctx);
73 }
74 for (uint32_t i = 0; i < 4; i++)
75 for (uint32_t j = 0; j < 4; j++) {
76 uint32_t r = r0 + ty + 16 * i, c = c0 + tx + 16 * j;
77 if (r < m && c < n) store(r, c, acc[i][j], user);
78 }
79 return true;
80}
8182/* Thread t's PER values of slab k0 of B (element e = t + 256 q: row e / COLS,
83 * column e % COLS) and x[k0 + t]. */
84#define MESH_LINALG_GEMV_PER (MESH_LINALG_GEMV_SLAB * MESH_LINALG_GEMV_COLS / MESH_LINALG_TILE_THREADS)
85MESH_RT_INLINE void mesh_linalg_gemv_load(const float *b, const float *x, uint32_t k, uint32_t n, uint32_t c0, uint32_t t,
86 uint32_t k0, float v[MESH_LINALG_GEMV_PER], float *xv) {
87 uint32_t rows = mesh_linalg_min(k - k0, MESH_LINALG_GEMV_SLAB);
88 for (uint32_t q = 0; q < MESH_LINALG_GEMV_PER; q++) {
89 uint32_t e = t + MESH_LINALG_TILE_THREADS * q;
90 uint32_t r = e / MESH_LINALG_GEMV_COLS, c = e % MESH_LINALG_GEMV_COLS;
91 v[q] = (r < rows && c0 + c < n) ? b[(size_t)(k0 + r) * n + c0 + c] : 0.0f;
92 }
93 *xv = t < rows ? x[k0 + t] : 0.0f;
94}
9596/* Row `row` of A (m x k) times B (k x n), columns c0 .. c0 + GEMV_COLS: the
97 * same sequential sum as gemm_tile's, bitwise, for the single-row case
98 * (decode). The 256 threads of the node stage slabs of GEMV_SLAB rows of B
99 * (and of the row of A) through registers into shared memory -- the next
100 * slab's loads in flight while the first GEMV_COLS threads sum their column
101 * over this one in k order. sb: GEMV_SLAB x GEMV_COLS f32s, sx: GEMV_SLAB.
102 * store(c, acc). The node has TILE_THREADS threads; row < m. */
103MESH_RT_INLINE void mesh_linalg_gemv_cols(const mesh_ctx *ctx, const float *a, const float *b, uint32_t k, uint32_t n, uint32_t row,
104 uint32_t c0, mesh_shared sb, mesh_shared sx, mesh_linalg_store1 store, void *user) {
105 const uint32_t COLS = MESH_LINALG_GEMV_COLS, SLAB = MESH_LINALG_GEMV_SLAB, TT = MESH_LINALG_TILE_THREADS;
106 uint32_t t = mesh_ctx_thread(ctx);
107 const float *x = a + (size_t)row * k;
108 float acc = 0.0f;
109 float regs[MESH_LINALG_GEMV_PER], rx;
110 mesh_linalg_gemv_load(b, x, k, n, c0, t, 0, regs, &rx);
111 for (uint32_t k0 = 0; k0 < k; k0 += SLAB) {
112 uint32_t rows = mesh_linalg_min(k - k0, SLAB);
113 for (uint32_t q = 0; q < MESH_LINALG_GEMV_PER; q++) mesh_shared_set_unchecked_f32(sb, t + TT * q, regs[q]);
114 if (t < SLAB) mesh_shared_set_unchecked_f32(sx, t, rx);
115 mesh_ctx_block_barrier(ctx);
116 if (k0 + SLAB < k) mesh_linalg_gemv_load(b, x, k, n, c0, t, k0 + SLAB, regs, &rx);
117 if (t < COLS) {
118 /* 8 rows' operands read first, then their adds in order (the chain is the adds only). */
119 uint32_t r = 0;
120 while (r + 8 <= rows) {
121 float xs[8], bs[8];
122 for (uint32_t i = 0; i < 8; i++) xs[i] = mesh_shared_get_unchecked_f32(sx, r + i);
123 for (uint32_t i = 0; i < 8; i++) bs[i] = mesh_shared_get_unchecked_f32(sb, (r + i) * COLS + t);
124 for (int i = 0; i < 8; i++) acc = acc + xs[i] * bs[i];
125 r += 8;
126 }
127 while (r < rows) {
128 acc = acc + mesh_shared_get_unchecked_f32(sx, r) * mesh_shared_get_unchecked_f32(sb, r * COLS + t);
129 r += 1;
130 }
131 }
132 mesh_ctx_block_barrier(ctx);
133 }
134 if (t < COLS && c0 + t < n) store(c0 + t, acc, user);
135}
136137#endif