1/**
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 buffer-backed operands. A
89 * multi-element tensor is read at the output's index; a single-element
90 * value (a `dot` result, or a scalar computed from one) is read at [0],
91 * which broadcasts it across the output. */
92 tensors: Map<string, number>;
93 /** cName -> slot in the params storage buffer, for runtime scalars. */
94 params: Map<string, number>;
95 /** cName -> defining expression, for scalars the .m computes from
96 * parameters (`us = a + b`). These have no buffer and no param slot; they
97 * become `let` bindings in the prologue of every kernel that reads them. */
98 scalars: Map<string, { name: string; expr: IRExpr }>;
99}
101/** Mutable state while emitting one kernel. */
102interface Ctx {
103 io: KernelInputs;
104 /** `let` lines to emit before the body, in dependency order. */
105 prologue: string[];
106 /** cName -> WGSL identifier, for scalars already bound in the prologue. */
107 bound: Map<string, string>;
108}
110/** WGSL identifier for a derived scalar. Avoids a leading underscore, which
111 * WGSL reserves. */
112const scalarIdent = (cName: string): string =>
113 `s_${cName.replace(/[^A-Za-z0-9_]/g, '_')}`;
115/**
116 * Bind a .m-derived scalar in the prologue (once), after whatever it depends
117 * on, and return its identifier.
118 */
119function bindScalar(cName: string, ctx: Ctx): string {
120 const already = ctx.bound.get(cName);
121 if (already) return already;
122 const def = ctx.io.scalars.get(cName)!;
123 const ident = scalarIdent(cName);
124 // Claim the name before emitting the RHS so a (malformed) self-reference
125 // cannot recurse forever.
126 ctx.bound.set(cName, ident);
127 const rhs = emitExpr(def.expr, ctx);
128 ctx.prologue.push(` let ${ident} = ${rhs};`);
129 return ident;
130}
132/**
133 * Emit the per-element WGSL expression for `e`. `i` is the element index
134 * variable in scope.
135 */
136function emitExpr(e: IRExpr, ctx: Ctx): string {
137 const io = ctx.io;
138 switch (e.kind) {
139 case 'NumLit':
140 return f32Lit(e.value);
142 case 'Var': {
143 // Buffer-backed operands first: a single-element value in a buffer
144 // shadows any compile-time definition the same name had earlier (a
145 // scalar seeded `rho = 1` and then updated from a dot result inside
146 // the solve loop reads as a buffer from the update on).
147 const bound = io.tensors.get(e.cName);
148 if (bound !== undefined) {
149 return isTensor(e.ty) ? `in${bound}[i]` : `in${bound}[0]`;
150 }
151 if (isTensor(e.ty)) {
152 throw new UnsupportedOnGpu(`no buffer bound for '${e.name}'`, e.span);
153 }
154 // Scalar: either an exact compile-time value or a runtime parameter.
155 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') {
156 return f32Lit(e.ty.exact);
157 }
158 const slot = io.params.get(e.cName);
159 if (slot !== undefined) return `prm[${slot}]`;
160 if (io.scalars.has(e.cName)) return bindScalar(e.cName, ctx);
161 throw new UnsupportedOnGpu(
162 `scalar '${e.name}' is not a constant, a parameter, or computed in ` +
163 `this model`,
164 e.span,
165 );
166 }
168 case 'Binary': {
169 if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
170 isTensor(e.left.ty) && isTensor(e.right.ty)) {
171 throw new UnsupportedOnGpu(
172 `matrix '${e.builtin === 'mtimes' ? '*' : '/'}' is not supported; ` +
173 `use the element-wise form ('.${e.builtin === 'mtimes' ? '*' : '/'}')`,
174 e.span,
175 );
176 }
177 if (e.builtin === 'power' || e.builtin === 'mpower') {
178 return emitPower(e.left, e.right, ctx, e.span);
179 }
180 const op = BINARY_OPS[e.builtin];
181 if (!op) {
182 throw new UnsupportedOnGpu(`operator '${e.builtin}' is not supported`, e.span);
183 }
184 return `(${emitExpr(e.left, ctx)} ${op} ${emitExpr(e.right, ctx)})`;
185 }
187 case 'Unary': {
188 const op = UNARY_OPS[e.builtin];
189 if (!op) {
190 throw new UnsupportedOnGpu(`unary '${e.builtin}' is not supported`, e.span);
191 }
192 return `(${op}${emitExpr(e.operand, ctx)})`;
193 }
195 case 'Call': {
196 // A shape constructor used inside an element-wise expression
197 // contributes the same constant at every slot, so it needs no buffer.
198 // (The shape itself is validated against the target by checkShapes.)
199 if (e.name === 'ones') return '1.0';
200 if (e.name === 'zeros') return '0.0';
202 const fn = CALL_FNS[e.name];
203 const b = getBuiltin(e.name);
204 if (!fn || !b?.elementwise) {
205 // A call numbl resolved to another function in the workspace gets a
206 // mangled specialization name; a builtin keeps its source-level name.
207 // User-function calls are expanded into the caller before planning
208 // (src/mgpu/inlineCalls.ts), so one surviving to kernel emission means
209 // the expansion did not reach it — a distinct failure from an
210 // unsupported builtin, and worth saying so.
211 const isUserFunction = e.cName !== e.name;
212 throw new UnsupportedOnGpu(
213 isUserFunction
214 ? `the call to '${e.name}' was not expanded into the caller — ` +
215 `assign its result to a variable on its own line`
216 : `'${e.name}' cannot be evaluated element-wise on the GPU`,
217 e.span,
218 );
219 }
220 return `${fn}(${e.args.map((a) => emitExpr(a, ctx)).join(', ')})`;
221 }
223 default:
224 throw new UnsupportedOnGpu(`'${e.kind}' is not supported on the GPU`, e.span);
225 }
226}
228/**
229 * `x.^k`. WGSL's `pow` is undefined for a negative base, and these fields go
230 * negative routinely, so expand small non-negative integer exponents into
231 * repeated multiplication — which is also what makes `u.^2` free.
232 */
233function emitPower(base: IRExpr, exponent: IRExpr, ctx: Ctx, span: unknown): string {
234 const k =
235 exponent.kind === 'NumLit'
236 ? exponent.value
237 : isNumeric(exponent.ty) && typeof exponent.ty.exact === 'number'
238 ? exponent.ty.exact
239 : undefined;
240 const b = emitExpr(base, ctx);
241 if (k !== undefined && Number.isInteger(k) && k >= 0 && k <= 8) {
242 if (k === 0) return '1.0';
243 // bind once so a compound base expression is not re-evaluated k times
244 return `pow_i${k}(${b})`;
245 }
246 if (k !== undefined && Number.isInteger(k) && k < 0 && k >= -8) {
247 return `(1.0 / pow_i${-k}(${b}))`;
248 }
249 throw new UnsupportedOnGpu(
250 `'.^' needs a literal integer exponent in [-8, 8] (got ` +
251 `${k === undefined ? 'a runtime value' : k}); a negative base makes ` +
252 `WGSL's pow() undefined`,
253 span,
254 );
255}
257/** Fixed-exponent power helpers, emitted only when used. */
258function powHelpers(used: Set<number>): string {
259 const out: string[] = [];
260 for (const k of [...used].sort((a, b) => a - b)) {
261 const body =
262 k === 1 ? 'x' : `x${' * x'.repeat(k - 1)}`;
263 out.push(`fn pow_i${k}(x: f32) -> f32 { return ${body}; }`);
264 }
265 return out.join('\n');
266}
268/**
269 * Reject implicit expansion (broadcasting).
270 *
271 * numbl's lowering permits it — `2x4096 .* 1x4096` lowers happily with MATLAB
272 * expansion semantics — but a kernel that walks one linear index across every
273 * operand would quietly compute the wrong thing. So every multi-element
274 * operand must have exactly the target's shape. Scalars are fine: they are
275 * read from the params buffer or folded in as literals.
276 */
277function checkShapes(e: IRExpr, target: NumericType, name: string): void {
278 const want = target.shape;
279 const same = (t: NumericType): boolean => {
280 const got = t.shape;
281 return (
282 !!want && !!got && want.length === got.length &&
283 want.every((d, i) => d === got[i])
284 );
285 };
286 const walk = (x: IRExpr): void => {
287 if (isNumeric(x.ty) && isMultiElement(x.ty) && !same(x.ty)) {
288 const got = x.ty.shape?.join('x') ?? 'dynamic';
289 throw new UnsupportedOnGpu(
290 `'${name}' would need implicit expansion: an operand is ${got} but the ` +
291 `result is ${want?.join('x') ?? 'dynamic'}. Expand it explicitly ` +
292 `(the GPU kernel walks one index across every operand).`,
293 x.span,
294 );
295 }
296 switch (x.kind) {
297 case 'Binary':
298 walk(x.left);
299 walk(x.right);
300 return;
301 case 'Unary':
302 walk(x.operand);
303 return;
304 case 'Call':
305 // A shape constructor's own arguments are sizes, not data.
306 if (x.name !== 'ones' && x.name !== 'zeros') x.args.forEach(walk);
307 return;
308 default:
309 return;
310 }
311 };
312 walk(e);
313}
315export const WORKGROUP_SIZE = 64;
317export interface Kernel {
318 code: string;
319 /** Number of output elements. */
320 count: number;
321 label: string;
322}
324/**
325 * Build the kernel for one element-wise `Assign`. `io` must already map every
326 * tensor operand cName to a binding index and every runtime scalar to a
327 * params slot; the output is binding 0 and the params buffer is the binding
328 * after the last input.
329 */
330export function buildKernel(
331 stmt: Assign,
332 io: KernelInputs,
333 count: number,
334 label: string,
335): Kernel {
336 if (!isNumeric(stmt.ty)) {
337 throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric array`, stmt.span);
338 }
339 if (stmt.ty.isComplex) {
340 throw new UnsupportedOnGpu(
341 `'${stmt.name}' is complex; the GPU backend is real-only (a spectral ` +
342 `field is carried as a real 2 x nlm array)`,
343 stmt.span,
344 );
345 }
347 checkShapes(stmt.expr, stmt.ty, stmt.name);
348 const ctx: Ctx = { io, prologue: [], bound: new Map() };
349 const body = emitExpr(stmt.expr, ctx);
351 // pow_iK helpers are discovered during emission; scan the result for them.
352 const used = new Set<number>();
353 const emitted = [...ctx.prologue, body].join('\n');
354 for (const m of emitted.matchAll(/\bpow_i(\d+)\(/g)) used.add(Number(m[1]));
356 const decls = [`@group(0) @binding(0) var<storage, read_write> out: array<f32>;`];
357 for (const [, slot] of io.tensors) {
358 decls.push(
359 `@group(0) @binding(${slot + 1}) var<storage, read> in${slot}: array<f32>;`,
360 );
361 }
362 // Params live in a read-only storage buffer rather than a uniform block:
363 // uniform arrays would need 16-byte element stride.
364 const prmBinding = io.tensors.size + 1;
365 decls.push(
366 `@group(0) @binding(${prmBinding}) var<storage, read> prm: array<f32>;`,
367 );
369 const code = `${decls.join('\n')}
371${powHelpers(used)}
373@compute @workgroup_size(${WORKGROUP_SIZE})
374fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
375 let i = gid.x;
376 if (i >= ${count}u) { return; }
377${ctx.prologue.length ? `${ctx.prologue.join('\n')}\n` : ''} out[i] = ${body};
378}
379`;
380 return { code, count, label };
381}