2 * IR expression tree -> one WGSL compute kernel.
3 *
4 * This is the WebGPU counterpart of numbl's C-side fused emitter
5 * (`codegen/emitTensorFused.ts`): for an `Assign` whose right-hand side is
6 * purely element-wise over operands of the target's shape, emit a single
7 * kernel that computes one output element per invocation. Because numbl's
8 * inline pass has already folded the ANF temps back together, one source line
9 * of MATLAB becomes one kernel.
10 *
11 * Everything is f32, matching the existing fp32 WebGPU transform backend.
12 */
13import { getBuiltin } from 'numbl-src/numbl-core/jit/builtins/index.ts';
14import { isMultiElement } from 'numbl-src/numbl-core/jit/lowering/types.ts';
15import type { IRExpr, Assign } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
16import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
18/** Raised for a construct the WGSL backend cannot express. Mirrors numbl's
19 * own decline discipline: fail at compile time with a source span, never
20 * silently produce something that computes the wrong thing. */
21export class UnsupportedOnGpu extends Error {
22 readonly span?: unknown;
23 constructor(message: string, span?: unknown) {
24 super(message);
25 this.name = 'UnsupportedOnGpu';
26 this.span = span;
27 }
28}
30const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
31const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
33/** Element-wise binary builtins -> WGSL infix operator. */
34const BINARY_OPS: Record<string, string> = {
35 plus: '+',
36 minus: '-',
37 times: '*',
38 rdivide: '/',
39 // Degenerate to element-wise when at least one side is a scalar; the
40 // both-tensor (true matrix) case is rejected below.
41 mtimes: '*',
42 mrdivide: '/',
43};
45/** Element-wise unary builtins -> WGSL prefix operator. */
46const UNARY_OPS: Record<string, string> = { uminus: '-', uplus: '+' };
48/** Element-wise builtin calls -> WGSL builtin of the same arity. */
49const CALL_FNS: Record<string, string> = {
50 abs: 'abs',
51 acos: 'acos',
52 asin: 'asin',
53 atan: 'atan',
54 atan2: 'atan2',
55 ceil: 'ceil',
56 cos: 'cos',
57 cosh: 'cosh',
58 exp: 'exp',
59 floor: 'floor',
60 log: 'log',
61 log2: 'log2',
62 max: 'max',
63 min: 'min',
64 round: 'round',
65 sign: 'sign',
66 sin: 'sin',
67 sinh: 'sinh',
68 sqrt: 'sqrt',
69 tan: 'tan',
70 tanh: 'tanh',
71};
73/** WGSL f32 literal. Must always carry a decimal point or exponent, or WGSL
74 * infers AbstractInt and rejects the mixed-type arithmetic. */
75function f32Lit(v: number): string {
76 if (!Number.isFinite(v)) {
77 throw new UnsupportedOnGpu(`cannot emit non-finite literal ${v}`);
78 }
79 return Number.isInteger(v) && Math.abs(v) < 1e21
80 ? `${v}.0`
81 : String(v).includes('e')
82 ? `${v}f`
83 : String(v);
84}
86/** How a scalar or tensor operand is read inside the kernel. */
87export interface KernelInputs {
88 /** cName -> storage binding index, for multi-element tensor operands. */
89 tensors: Map<string, number>;
90 /** cName -> slot in the params storage buffer, for runtime scalars. */
91 params: Map<string, number>;
92 /** cName -> defining expression, for scalars the .m computes from
93 * parameters (`us = a + b`). These have no buffer and no param slot; they
94 * become `let` bindings in the prologue of every kernel that reads them. */
95 scalars: Map<string, { name: string; expr: IRExpr }>;
96}
98/** Mutable state while emitting one kernel. */
99interface Ctx {
100 io: KernelInputs;
101 /** `let` lines to emit before the body, in dependency order. */
102 prologue: string[];
103 /** cName -> WGSL identifier, for scalars already bound in the prologue. */
104 bound: Map<string, string>;
105}
107/** WGSL identifier for a derived scalar. Avoids a leading underscore, which
108 * WGSL reserves. */
109const scalarIdent = (cName: string): string =>
110 `s_${cName.replace(/[^A-Za-z0-9_]/g, '_')}`;
112/**
113 * Bind a .m-derived scalar in the prologue (once), after whatever it depends
114 * on, and return its identifier.
115 */
116function bindScalar(cName: string, ctx: Ctx): string {
117 const already = ctx.bound.get(cName);
118 if (already) return already;
119 const def = ctx.io.scalars.get(cName)!;
120 const ident = scalarIdent(cName);
121 // Claim the name before emitting the RHS so a (malformed) self-reference
122 // cannot recurse forever.
123 ctx.bound.set(cName, ident);
124 const rhs = emitExpr(def.expr, ctx);
125 ctx.prologue.push(` let ${ident} = ${rhs};`);
126 return ident;
127}
129/**
130 * Emit the per-element WGSL expression for `e`. `i` is the element index
131 * variable in scope.
132 */
133function emitExpr(e: IRExpr, ctx: Ctx): string {
134 const io = ctx.io;
135 switch (e.kind) {
136 case 'NumLit':
137 return f32Lit(e.value);
139 case 'Var': {
140 if (isTensor(e.ty)) {
141 const slot = io.tensors.get(e.cName);
142 if (slot === undefined) {
143 throw new UnsupportedOnGpu(`no buffer bound for '${e.name}'`, e.span);
144 }
145 return `in${slot}[i]`;
146 }
147 // Scalar: either an exact compile-time value or a runtime parameter.
148 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') {
149 return f32Lit(e.ty.exact);
150 }
151 const slot = io.params.get(e.cName);
152 if (slot !== undefined) return `prm[${slot}]`;
153 if (io.scalars.has(e.cName)) return bindScalar(e.cName, ctx);
154 throw new UnsupportedOnGpu(
155 `scalar '${e.name}' is not a constant, a parameter, or computed in ` +
156 `this model`,
157 e.span,
158 );
159 }
161 case 'Binary': {
162 if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
163 isTensor(e.left.ty) && isTensor(e.right.ty)) {
164 throw new UnsupportedOnGpu(
165 `matrix '${e.builtin === 'mtimes' ? '*' : '/'}' is not supported; ` +
166 `use the element-wise form ('.${e.builtin === 'mtimes' ? '*' : '/'}')`,
167 e.span,
168 );
169 }
170 if (e.builtin === 'power' || e.builtin === 'mpower') {
171 return emitPower(e.left, e.right, ctx, e.span);
172 }
173 const op = BINARY_OPS[e.builtin];
174 if (!op) {
175 throw new UnsupportedOnGpu(`operator '${e.builtin}' is not supported`, e.span);
176 }
177 return `(${emitExpr(e.left, ctx)} ${op} ${emitExpr(e.right, ctx)})`;
178 }
180 case 'Unary': {
181 const op = UNARY_OPS[e.builtin];
182 if (!op) {
183 throw new UnsupportedOnGpu(`unary '${e.builtin}' is not supported`, e.span);
184 }
185 return `(${op}${emitExpr(e.operand, ctx)})`;
186 }
188 case 'Call': {
189 // A shape constructor used inside an element-wise expression
190 // contributes the same constant at every slot, so it needs no buffer.
191 // (The shape itself is validated against the target by checkShapes.)
192 if (e.name === 'ones') return '1.0';
193 if (e.name === 'zeros') return '0.0';
195 const fn = CALL_FNS[e.name];
196 const b = getBuiltin(e.name);
197 if (!fn || !b?.elementwise) {
198 // A call numbl resolved to another function in the file gets a mangled
199 // specialization name; a builtin keeps its source-level name. Only the
200 // model's entry points are compiled, so a helper is a distinct failure
201 // from an unsupported builtin and deserves to say so.
202 const isUserFunction = e.cName !== e.name;
203 throw new UnsupportedOnGpu(
204 isUserFunction
205 ? `'${e.name}' is a function defined in this model. Only init and ` +
206 `step are compiled — inline its body into the caller.`
207 : `'${e.name}' cannot be evaluated element-wise on the GPU`,
208 e.span,
209 );
210 }
211 return `${fn}(${e.args.map((a) => emitExpr(a, ctx)).join(', ')})`;
212 }
214 default:
215 throw new UnsupportedOnGpu(`'${e.kind}' is not supported on the GPU`, e.span);
216 }
217}
219/**
220 * `x.^k`. WGSL's `pow` is undefined for a negative base, and these fields go
221 * negative routinely, so expand small non-negative integer exponents into
222 * repeated multiplication — which is also what makes `u.^2` free.
223 */
224function emitPower(base: IRExpr, exponent: IRExpr, ctx: Ctx, span: unknown): string {
225 const k =
226 exponent.kind === 'NumLit'
227 ? exponent.value
228 : isNumeric(exponent.ty) && typeof exponent.ty.exact === 'number'
229 ? exponent.ty.exact
230 : undefined;
231 const b = emitExpr(base, ctx);
232 if (k !== undefined && Number.isInteger(k) && k >= 0 && k <= 8) {
233 if (k === 0) return '1.0';
234 // bind once so a compound base expression is not re-evaluated k times
235 return `pow_i${k}(${b})`;
236 }
237 if (k !== undefined && Number.isInteger(k) && k < 0 && k >= -8) {
238 return `(1.0 / pow_i${-k}(${b}))`;
239 }
240 throw new UnsupportedOnGpu(
241 `'.^' needs a literal integer exponent in [-8, 8] (got ` +
242 `${k === undefined ? 'a runtime value' : k}); a negative base makes ` +
243 `WGSL's pow() undefined`,
244 span,
245 );
246}
248/** Fixed-exponent power helpers, emitted only when used. */
249function powHelpers(used: Set<number>): string {
250 const out: string[] = [];
251 for (const k of [...used].sort((a, b) => a - b)) {
252 const body =
253 k === 1 ? 'x' : `x${' * x'.repeat(k - 1)}`;
254 out.push(`fn pow_i${k}(x: f32) -> f32 { return ${body}; }`);
255 }
256 return out.join('\n');
257}
259/**
260 * Reject implicit expansion (broadcasting).
261 *
262 * numbl's lowering permits it — `2x4096 .* 1x4096` lowers happily with MATLAB
263 * expansion semantics — but a kernel that walks one linear index across every
264 * operand would quietly compute the wrong thing. So every multi-element
265 * operand must have exactly the target's shape. Scalars are fine: they are
266 * read from the params buffer or folded in as literals.
267 */
268function checkShapes(e: IRExpr, target: NumericType, name: string): void {
269 const want = target.shape;
270 const same = (t: NumericType): boolean => {
271 const got = t.shape;
272 return (
273 !!want && !!got && want.length === got.length &&
274 want.every((d, i) => d === got[i])
275 );
276 };
277 const walk = (x: IRExpr): void => {
278 if (isNumeric(x.ty) && isMultiElement(x.ty) && !same(x.ty)) {
279 const got = x.ty.shape?.join('x') ?? 'dynamic';
280 throw new UnsupportedOnGpu(
281 `'${name}' would need implicit expansion: an operand is ${got} but the ` +
282 `result is ${want?.join('x') ?? 'dynamic'}. Expand it explicitly ` +
283 `(the GPU kernel walks one index across every operand).`,
284 x.span,
285 );
286 }
287 switch (x.kind) {
288 case 'Binary':
289 walk(x.left);
290 walk(x.right);
291 return;
292 case 'Unary':
293 walk(x.operand);
294 return;
295 case 'Call':
296 // A shape constructor's own arguments are sizes, not data.
297 if (x.name !== 'ones' && x.name !== 'zeros') x.args.forEach(walk);
298 return;
299 default:
300 return;
301 }
302 };
303 walk(e);
304}
306export const WORKGROUP_SIZE = 64;
308export interface Kernel {
309 code: string;
310 /** Number of output elements. */
311 count: number;
312 label: string;
313}
315/**
316 * Build the kernel for one element-wise `Assign`. `io` must already map every
317 * tensor operand cName to a binding index and every runtime scalar to a
318 * params slot; the output is binding 0 and the params buffer is the binding
319 * after the last input.
320 */
321export function buildKernel(
322 stmt: Assign,
323 io: KernelInputs,
324 count: number,
325 label: string,
326): Kernel {
327 if (!isNumeric(stmt.ty)) {
328 throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric array`, stmt.span);
329 }
330 if (stmt.ty.isComplex) {
331 throw new UnsupportedOnGpu(
332 `'${stmt.name}' is complex; the GPU backend is real-only (a spectral ` +
333 `field is carried as a real 2 x nlm array)`,
334 stmt.span,
335 );
336 }
338 checkShapes(stmt.expr, stmt.ty, stmt.name);
339 const ctx: Ctx = { io, prologue: [], bound: new Map() };
340 const body = emitExpr(stmt.expr, ctx);
342 // pow_iK helpers are discovered during emission; scan the result for them.
343 const used = new Set<number>();
344 const emitted = [...ctx.prologue, body].join('\n');
345 for (const m of emitted.matchAll(/\bpow_i(\d+)\(/g)) used.add(Number(m[1]));
347 const decls = [`@group(0) @binding(0) var<storage, read_write> out: array<f32>;`];
348 for (const [, slot] of io.tensors) {
349 decls.push(
350 `@group(0) @binding(${slot + 1}) var<storage, read> in${slot}: array<f32>;`,
351 );
352 }
353 // Params live in a read-only storage buffer rather than a uniform block:
354 // uniform arrays would need 16-byte element stride.
355 const prmBinding = io.tensors.size + 1;
356 decls.push(
357 `@group(0) @binding(${prmBinding}) var<storage, read> prm: array<f32>;`,
358 );
360 const code = `${decls.join('\n')}
362${powHelpers(used)}
364@compute @workgroup_size(${WORKGROUP_SIZE})
365fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
366 let i = gid.x;
367 if (i >= ${count}u) { return; }
368${ctx.prologue.length ? `${ctx.prologue.join('\n')}\n` : ''} out[i] = ${body};
369}
370`;
371 return { code, count, label };
372}