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/** Zero-argument builtins that are compile-time constants. numbl lowers `pi`
74 * as a call rather than folding it, so the backend is where it becomes a
75 * number. */
76const CONST_FNS: Record<string, number> = { pi: Math.PI };
78/** WGSL f32 literal. Must always carry a decimal point or exponent, or WGSL
79 * infers AbstractInt and rejects the mixed-type arithmetic. */
80function f32Lit(v: number): string {
81 if (!Number.isFinite(v)) {
82 throw new UnsupportedOnGpu(`cannot emit non-finite literal ${v}`);
83 }
84 return Number.isInteger(v) && Math.abs(v) < 1e21
85 ? `${v}.0`
86 : String(v).includes('e')
87 ? `${v}f`
88 : String(v);
89}
91/** How a scalar or tensor operand is read inside the kernel. */
92export interface KernelInputs {
93 /** cName -> storage binding index, for multi-element tensor operands. */
94 tensors: Map<string, number>;
95 /** cName -> slot in the params storage buffer, for runtime scalars. */
96 params: Map<string, number>;
97 /** cName -> defining expression, for scalars the .m computes from
98 * parameters (`us = a + b`). These have no buffer and no param slot; they
99 * become `let` bindings in the prologue of every kernel that reads them. */
100 scalars: Map<string, { name: string; expr: IRExpr }>;
101}
103/** Mutable state while emitting one kernel. */
104interface Ctx {
105 io: KernelInputs;
106 /** `let` lines to emit before the body, in dependency order. */
107 prologue: string[];
108 /** cName -> WGSL identifier, for scalars already bound in the prologue. */
109 bound: Map<string, string>;
110}
112/** WGSL identifier for a derived scalar. Avoids a leading underscore, which
113 * WGSL reserves. */
114const scalarIdent = (cName: string): string =>
115 `s_${cName.replace(/[^A-Za-z0-9_]/g, '_')}`;
117/**
118 * Bind a .m-derived scalar in the prologue (once), after whatever it depends
119 * on, and return its identifier.
120 */
121function bindScalar(cName: string, ctx: Ctx): string {
122 const already = ctx.bound.get(cName);
123 if (already) return already;
124 const def = ctx.io.scalars.get(cName)!;
125 const ident = scalarIdent(cName);
126 // Claim the name before emitting the RHS so a (malformed) self-reference
127 // cannot recurse forever.
128 ctx.bound.set(cName, ident);
129 const rhs = emitExpr(def.expr, ctx);
130 ctx.prologue.push(` let ${ident} = ${rhs};`);
131 return ident;
132}
134/**
135 * Emit the per-element WGSL expression for `e`. `i` is the element index
136 * variable in scope.
137 */
138function emitExpr(e: IRExpr, ctx: Ctx): string {
139 const io = ctx.io;
140 switch (e.kind) {
141 case 'NumLit':
142 return f32Lit(e.value);
144 case 'Var': {
145 if (isTensor(e.ty)) {
146 const slot = io.tensors.get(e.cName);
147 if (slot === undefined) {
148 throw new UnsupportedOnGpu(`no buffer bound for '${e.name}'`, e.span);
149 }
150 return `in${slot}[i]`;
151 }
152 // Scalar: either an exact compile-time value or a runtime parameter.
153 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') {
154 return f32Lit(e.ty.exact);
155 }
156 const slot = io.params.get(e.cName);
157 if (slot !== undefined) return `prm[${slot}]`;
158 if (io.scalars.has(e.cName)) return bindScalar(e.cName, ctx);
159 throw new UnsupportedOnGpu(
160 `scalar '${e.name}' is not a constant, a parameter, or computed in ` +
161 `this model`,
162 e.span,
163 );
164 }
166 case 'Binary': {
167 if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
168 isTensor(e.left.ty) && isTensor(e.right.ty)) {
169 throw new UnsupportedOnGpu(
170 `matrix '${e.builtin === 'mtimes' ? '*' : '/'}' is not supported; ` +
171 `use the element-wise form ('.${e.builtin === 'mtimes' ? '*' : '/'}')`,
172 e.span,
173 );
174 }
175 if (e.builtin === 'power' || e.builtin === 'mpower') {
176 return emitPower(e.left, e.right, ctx, e.span);
177 }
178 const op = BINARY_OPS[e.builtin];
179 if (!op) {
180 throw new UnsupportedOnGpu(`operator '${e.builtin}' is not supported`, e.span);
181 }
182 return `(${emitExpr(e.left, ctx)} ${op} ${emitExpr(e.right, ctx)})`;
183 }
185 case 'Unary': {
186 const op = UNARY_OPS[e.builtin];
187 if (!op) {
188 throw new UnsupportedOnGpu(`unary '${e.builtin}' is not supported`, e.span);
189 }
190 return `(${op}${emitExpr(e.operand, ctx)})`;
191 }
193 case 'Call': {
194 // A shape constructor used inside an element-wise expression
195 // contributes the same constant at every slot, so it needs no buffer.
196 // (The shape itself is validated against the target by checkShapes.)
197 if (e.name === 'ones') return '1.0';
198 if (e.name === 'zeros') return '0.0';
200 const konst = CONST_FNS[e.name];
201 if (konst !== undefined && e.args.length === 0) return f32Lit(konst);
203 const fn = CALL_FNS[e.name];
204 const b = getBuiltin(e.name);
205 if (!fn || !b?.elementwise) {
206 // A call numbl resolved to another function in the file gets a mangled
207 // specialization name; a builtin keeps its source-level name. Only the
208 // model's entry points are compiled, so a helper is a distinct failure
209 // from an unsupported builtin and deserves to say so.
210 const isUserFunction = e.cName !== e.name;
211 throw new UnsupportedOnGpu(
212 isUserFunction
213 ? `'${e.name}' is a function defined in this model. Only init and ` +
214 `step are compiled — inline its body into the caller.`
215 : `'${e.name}' cannot be evaluated element-wise on the GPU`,
216 e.span,
217 );
218 }
219 return `${fn}(${e.args.map((a) => emitExpr(a, ctx)).join(', ')})`;
220 }
222 default:
223 throw new UnsupportedOnGpu(`'${e.kind}' is not supported on the GPU`, e.span);
224 }
225}
227/**
228 * `x.^k`. WGSL's `pow` is undefined for a negative base, and these fields go
229 * negative routinely, so expand small non-negative integer exponents into
230 * repeated multiplication — which is also what makes `u.^2` free.
231 */
232function emitPower(base: IRExpr, exponent: IRExpr, ctx: Ctx, span: unknown): string {
233 const k =
234 exponent.kind === 'NumLit'
235 ? exponent.value
236 : isNumeric(exponent.ty) && typeof exponent.ty.exact === 'number'
237 ? exponent.ty.exact
238 : undefined;
239 const b = emitExpr(base, ctx);
240 if (k !== undefined && Number.isInteger(k) && k >= 0 && k <= 8) {
241 if (k === 0) return '1.0';
242 // bind once so a compound base expression is not re-evaluated k times
243 return `pow_i${k}(${b})`;
244 }
245 if (k !== undefined && Number.isInteger(k) && k < 0 && k >= -8) {
246 return `(1.0 / pow_i${-k}(${b}))`;
247 }
248 throw new UnsupportedOnGpu(
249 `'.^' needs a literal integer exponent in [-8, 8] (got ` +
250 `${k === undefined ? 'a runtime value' : k}); a negative base makes ` +
251 `WGSL's pow() undefined`,
252 span,
253 );
254}
256/** Fixed-exponent power helpers, emitted only when used. */
257function powHelpers(used: Set<number>): string {
258 const out: string[] = [];
259 for (const k of [...used].sort((a, b) => a - b)) {
260 const body =
261 k === 1 ? 'x' : `x${' * x'.repeat(k - 1)}`;
262 out.push(`fn pow_i${k}(x: f32) -> f32 { return ${body}; }`);
263 }
264 return out.join('\n');
265}
267/**
268 * Can this expression be evaluated inside a fused kernel?
269 *
270 * The predicate behind src/mgpu/fuse.ts, and the single statement of what the
271 * emitter above accepts: everything here is something `emitExpr` can write out
272 * per element, and everything it declines is something that needs its own
273 * dispatch. Keep the two in step.
274 */
275export function isGpuFusableExpr(e: IRExpr): boolean {
276 switch (e.kind) {
277 case 'NumLit':
278 return true;
279 case 'Var':
280 return isNumeric(e.ty);
281 case 'Binary': {
282 if (
283 (e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
284 isTensor(e.left.ty) && isTensor(e.right.ty)
285 ) {
286 return false;
287 }
288 if (e.builtin === 'power' || e.builtin === 'mpower') {
289 return isGpuFusableExpr(e.left) && isGpuFusableExpr(e.right);
290 }
291 return (
292 e.builtin in BINARY_OPS && isGpuFusableExpr(e.left) && isGpuFusableExpr(e.right)
293 );
294 }
295 case 'Unary':
296 return e.builtin in UNARY_OPS && isGpuFusableExpr(e.operand);
297 case 'Call': {
298 if (e.name === 'zeros' || e.name === 'ones') return true;
299 if (e.name in CONST_FNS) return e.args.length === 0;
300 return e.name in CALL_FNS && e.args.every(isGpuFusableExpr);
301 }
302 default:
303 return false;
304 }
305}
307/**
308 * Reject implicit expansion (broadcasting).
309 *
310 * numbl's lowering permits it — `2x4096 .* 1x4096` lowers happily with MATLAB
311 * expansion semantics — but a kernel that walks one linear index across every
312 * operand would quietly compute the wrong thing. So every multi-element
313 * operand must have exactly the target's shape. Scalars are fine: they are
314 * read from the params buffer or folded in as literals.
315 */
316export function checkShapes(e: IRExpr, target: NumericType, name: string): void {
317 const want = target.shape;
318 const same = (t: NumericType): boolean => {
319 const got = t.shape;
320 return (
321 !!want && !!got && want.length === got.length &&
322 want.every((d, i) => d === got[i])
323 );
324 };
325 const walk = (x: IRExpr): void => {
326 if (isNumeric(x.ty) && isMultiElement(x.ty) && !same(x.ty)) {
327 const got = x.ty.shape?.join('x') ?? 'dynamic';
328 throw new UnsupportedOnGpu(
329 `'${name}' would need implicit expansion: an operand is ${got} but the ` +
330 `result is ${want?.join('x') ?? 'dynamic'}. Expand it explicitly ` +
331 `(the GPU kernel walks one index across every operand).`,
332 x.span,
333 );
334 }
335 switch (x.kind) {
336 case 'Binary':
337 walk(x.left);
338 walk(x.right);
339 return;
340 case 'Unary':
341 walk(x.operand);
342 return;
343 case 'Call':
344 // A shape constructor's own arguments are sizes, not data.
345 if (x.name !== 'ones' && x.name !== 'zeros') x.args.forEach(walk);
346 return;
347 default:
348 return;
349 }
350 };
351 walk(e);
352}
354export const WORKGROUP_SIZE = 64;
356export interface Kernel {
357 code: string;
358 /** Number of output elements. */
359 count: number;
360 label: string;
361}
363/**
364 * Build the kernel for one element-wise `Assign`. `io` must already map every
365 * tensor operand cName to a binding index and every runtime scalar to a
366 * params slot; the output is binding 0 and the params buffer is the binding
367 * after the last input.
368 */
369export function buildKernel(
370 stmt: Assign,
371 io: KernelInputs,
372 count: number,
373 label: string,
374): Kernel {
375 if (!isNumeric(stmt.ty)) {
376 throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric array`, stmt.span);
377 }
378 if (stmt.ty.isComplex) {
379 throw new UnsupportedOnGpu(
380 `'${stmt.name}' is complex; the GPU backend is real-only (a spectral ` +
381 `field is carried as a real 2 x nlm array)`,
382 stmt.span,
383 );
384 }
386 checkShapes(stmt.expr, stmt.ty, stmt.name);
387 const ctx: Ctx = { io, prologue: [], bound: new Map() };
388 const body = emitExpr(stmt.expr, ctx);
390 // pow_iK helpers are discovered during emission; scan the result for them.
391 const used = new Set<number>();
392 const emitted = [...ctx.prologue, body].join('\n');
393 for (const m of emitted.matchAll(/\bpow_i(\d+)\(/g)) used.add(Number(m[1]));
395 const decls = [`@group(0) @binding(0) var<storage, read_write> out: array<f32>;`];
396 for (const [, slot] of io.tensors) {
397 decls.push(
398 `@group(0) @binding(${slot + 1}) var<storage, read> in${slot}: array<f32>;`,
399 );
400 }
401 // Params live in a read-only storage buffer rather than a uniform block:
402 // uniform arrays would need 16-byte element stride.
403 const prmBinding = io.tensors.size + 1;
404 decls.push(
405 `@group(0) @binding(${prmBinding}) var<storage, read> prm: array<f32>;`,
406 );
408 const code = `${decls.join('\n')}
410${powHelpers(used)}
412@compute @workgroup_size(${WORKGROUP_SIZE})
413fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
414 let i = gid.x;
415 if (i >= ${count}u) { return; }
416${ctx.prologue.length ? `${ctx.prologue.join('\n')}\n` : ''} out[i] = ${body};
417}
418`;
419 return { code, count, label };
420}