1/**
2 * The non-elementwise kernels: tiled GEMM, tiled transpose, and reductions.
3 *
4 * Everything is column-major f32, matching MATLAB's layout, so `A(:)` walks
5 * the same linear buffer the kernels index. All shapes are compile-time
6 * constants baked into the WGSL — nothing here reads a dims uniform, which is
7 * what lets the planner prebuild every pipeline and bind group.
8 *
9 * The GEMM is the classic 16x16 shared-memory tile (adapted to column-major
10 * from matmul-bench's row-major sgemm). It will not beat a tuned BLAS, but it
11 * is the honest baseline for "what does A*B cost in a browser".
12 *
13 * Reductions take a *fused loader*: the per-element expression emitted by
14 * wgsl.ts, so `sum(a .* b + c)` reads its operands exactly once, in one pass.
15 */
16import type { FusedLoader } from './wgsl.ts';
18const TILE = 16;
20/** C(m x n) = A(m x k) * B(k x n), column-major. Bindings: 0=C, 1=A, 2=B. */
21export function gemmKernel(m: number, k: number, n: number): string {
22 return `
23@group(0) @binding(0) var<storage, read_write> C: array<f32>;
24@group(0) @binding(1) var<storage, read> A: array<f32>;
25@group(0) @binding(2) var<storage, read> B: array<f32>;
27var<workgroup> tileA: array<array<f32, ${TILE}>, ${TILE}>;
28var<workgroup> tileB: array<array<f32, ${TILE}>, ${TILE}>;
30@compute @workgroup_size(${TILE}, ${TILE})
31fn main(
32 @builtin(global_invocation_id) gid: vec3<u32>,
33 @builtin(local_invocation_id) lid: vec3<u32>,
34) {
35 let row = gid.x;
36 let col = gid.y;
37 let lx = lid.x;
38 let ly = lid.y;
39 var acc: f32 = 0.0;
40 let numTiles = ${Math.ceil(k / TILE)}u;
41 for (var t: u32 = 0u; t < numTiles; t = t + 1u) {
42 // Consecutive lx reads consecutive addresses in both loads (column-major).
43 let aCol = t * ${TILE}u + ly;
44 let bRow = t * ${TILE}u + lx;
45 tileA[lx][ly] = select(0.0, A[row + aCol * ${m}u], row < ${m}u && aCol < ${k}u);
46 tileB[lx][ly] = select(0.0, B[bRow + col * ${k}u], bRow < ${k}u && col < ${n}u);
47 workgroupBarrier();
48 for (var p: u32 = 0u; p < ${TILE}u; p = p + 1u) {
49 acc = acc + tileA[lx][p] * tileB[p][ly];
50 }
51 workgroupBarrier();
52 }
53 if (row < ${m}u && col < ${n}u) {
54 C[row + col * ${m}u] = acc;
55 }
56}
57`;
58}
60export const gemmDispatch = (m: number, n: number): [number, number] => [
61 Math.ceil(m / TILE),
62 Math.ceil(n / TILE),
63];
65/** out(n x m) = in(m x n)', column-major, staged through a tile so both the
66 * read and the write are coalesced. Bindings: 0=out, 1=in. */
67export function transposeKernel(m: number, n: number): string {
68 return `
69@group(0) @binding(0) var<storage, read_write> out: array<f32>;
70@group(0) @binding(1) var<storage, read> src: array<f32>;
72var<workgroup> tile: array<array<f32, ${TILE}>, ${TILE}>;
74@compute @workgroup_size(${TILE}, ${TILE})
75fn main(
76 @builtin(workgroup_id) wg: vec3<u32>,
77 @builtin(local_invocation_id) lid: vec3<u32>,
78) {
79 let lx = lid.x;
80 let ly = lid.y;
81 // Read block (wg.x, wg.y) of src: rows wg.x*T.., cols wg.y*T..
82 let sr = wg.x * ${TILE}u + lx;
83 let sc = wg.y * ${TILE}u + ly;
84 if (sr < ${m}u && sc < ${n}u) {
85 tile[ly][lx] = src[sr + sc * ${m}u];
86 }
87 workgroupBarrier();
88 // Write the transposed block: rows of out are cols of src.
89 let dr = wg.y * ${TILE}u + lx;
90 let dc = wg.x * ${TILE}u + ly;
91 if (dr < ${n}u && dc < ${m}u) {
92 out[dr + dc * ${n}u] = tile[lx][ly];
93 }
94}
95`;
96}
98export const transposeDispatch = (m: number, n: number): [number, number] => [
99 Math.ceil(m / TILE),
100 Math.ceil(n / TILE),
101];
103// ── Reductions ──────────────────────────────────────────────────────────
105export type Combine = 'add' | 'mul' | 'max' | 'min';
106/** Applied to each loaded element before combining. `sq` serves norm/dot. */
107export type MapKind = 'id' | 'sq';
108/** Applied to the final value. `scale` divides (mean); `sqrt` closes norm. */
109export type Epilogue = { kind: 'none' } | { kind: 'scale'; by: number } | { kind: 'sqrt' };
111const REDUCE_WG = 256;
113const IDENT: Record<Combine, string> = {
114 add: '0.0',
115 mul: '1.0',
116 max: '-3.4028234663852886e+38',
117 min: '3.4028234663852886e+38',
118};
120const comb = (c: Combine, a: string, b: string): string =>
121 c === 'add' ? `${a} + ${b}`
122 : c === 'mul' ? `${a} * ${b}`
123 : `${c}(${a}, ${b})`;
125const mapped = (m: MapKind, v: string): string => (m === 'sq' ? `${v} * ${v}` : v);
127const epilogued = (e: Epilogue, v: string): string =>
128 e.kind === 'scale' ? `(${v}) * ${e.by}` : e.kind === 'sqrt' ? `sqrt(${v})` : v;
130const treeReduce = (c: Combine): string => `
131 sdata[li] = acc;
132 workgroupBarrier();
133 var stride = ${REDUCE_WG / 2}u;
134 while (stride > 0u) {
135 if (li < stride) {
136 sdata[li] = ${comb(c, 'sdata[li]', 'sdata[li + stride]')};
137 }
138 workgroupBarrier();
139 stride = stride / 2u;
140 }`;
142/** Number of pass-1 partials for a full reduction over `count` elements. */
143export function reducePartials(count: number): number {
144 return Math.max(1, Math.min(1024, Math.ceil(count / (REDUCE_WG * 8))));
145}
147/**
148 * Full reduction, pass 1: `numWg` workgroups grid-stride over `count`
149 * elements of the fused loader, leaving one partial each. Loader binding
150 * declarations (`decls`) come from wgsl.ts's bindingDecls, whose binding 0
151 * ("out") is the partials buffer here.
152 */
153export function reduceFullPass1(
154 decls: string[],
155 loader: FusedLoader,
156 count: number,
157 numWg: number,
158 combine: Combine,
159 map: MapKind,
160): string {
161 return `${decls.join('\n')}
162${loader.helpers}
163var<workgroup> sdata: array<f32, ${REDUCE_WG}>;
165@compute @workgroup_size(${REDUCE_WG})
166fn main(
167 @builtin(workgroup_id) wg: vec3<u32>,
168 @builtin(local_invocation_id) lid: vec3<u32>,
169) {
170 let li = lid.x;
171 var acc: f32 = ${IDENT[combine]};
172 var i = wg.x * ${REDUCE_WG}u + li;
173 while (i < ${count}u) {
174 let v = ${loader.body};
175 acc = ${comb(combine, 'acc', mapped(map, 'v'))};
176 i = i + ${numWg * REDUCE_WG}u;
177 }
178${treeReduce(combine)}
179 if (li == 0u) { out[wg.x] = sdata[0]; }
180}
181`;
182}
184/** Full reduction, pass 2: one workgroup folds the partials and applies the
185 * epilogue. Bindings: 0=result (1 element), 1=partials. */
186export function reduceFullPass2(
187 numPartials: number,
188 combine: Combine,
189 epilogue: Epilogue,
190): string {
191 return `
192@group(0) @binding(0) var<storage, read_write> out: array<f32>;
193@group(0) @binding(1) var<storage, read> partials: array<f32>;
194var<workgroup> sdata: array<f32, ${REDUCE_WG}>;
196@compute @workgroup_size(${REDUCE_WG})
197fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
198 let li = lid.x;
199 var acc: f32 = ${IDENT[combine]};
200 var i = li;
201 while (i < ${numPartials}u) {
202 acc = ${comb(combine, 'acc', 'partials[i]')};
203 i = i + ${REDUCE_WG}u;
204 }
205${treeReduce(combine)}
206 if (li == 0u) { out[0] = ${epilogued(epilogue, 'sdata[0]')}; }
207}
208`;
209}
211/**
212 * Column-wise reduction of an (m x n) input: one workgroup per column,
213 * threads grid-stride down the column (contiguous in column-major), leaving
214 * out[column]. Dispatch n workgroups.
215 */
216export function reduceColumns(
217 decls: string[],
218 loader: FusedLoader,
219 m: number,
220 combine: Combine,
221 map: MapKind,
222 epilogue: Epilogue,
223): string {
224 return `${decls.join('\n')}
225${loader.helpers}
226var<workgroup> sdata: array<f32, ${REDUCE_WG}>;
228@compute @workgroup_size(${REDUCE_WG})
229fn main(
230 @builtin(workgroup_id) wg: vec3<u32>,
231 @builtin(local_invocation_id) lid: vec3<u32>,
232) {
233 let li = lid.x;
234 let col = wg.x;
235 var acc: f32 = ${IDENT[combine]};
236 var r = li;
237 while (r < ${m}u) {
238 let i = r + col * ${m}u;
239 let v = ${loader.body};
240 acc = ${comb(combine, 'acc', mapped(map, 'v'))};
241 r = r + ${REDUCE_WG}u;
242 }
243${treeReduce(combine)}
244 if (li == 0u) { out[col] = ${epilogued(epilogue, 'sdata[0]')}; }
245}
246`;
247}