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 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 throw new UnsupportedOnGpu(
199 `'${e.name}' cannot be evaluated element-wise on the GPU`,
200 e.span,
201 );
202 }
203 return `${fn}(${e.args.map((a) => emitExpr(a, ctx)).join(', ')})`;
204 }
206 default:
207 throw new UnsupportedOnGpu(`'${e.kind}' is not supported on the GPU`, e.span);
208 }
209}
211/**
212 * `x.^k`. WGSL's `pow` is undefined for a negative base, and these fields go
213 * negative routinely, so expand small non-negative integer exponents into
214 * repeated multiplication — which is also what makes `u.^2` free.
215 */
216function emitPower(base: IRExpr, exponent: IRExpr, ctx: Ctx, span: unknown): string {
217 const k =
218 exponent.kind === 'NumLit'
219 ? exponent.value
220 : isNumeric(exponent.ty) && typeof exponent.ty.exact === 'number'
221 ? exponent.ty.exact
222 : undefined;
223 const b = emitExpr(base, ctx);
224 if (k !== undefined && Number.isInteger(k) && k >= 0 && k <= 8) {
225 if (k === 0) return '1.0';
226 // bind once so a compound base expression is not re-evaluated k times
227 return `pow_i${k}(${b})`;
228 }
229 if (k !== undefined && Number.isInteger(k) && k < 0 && k >= -8) {
230 return `(1.0 / pow_i${-k}(${b}))`;
231 }
232 throw new UnsupportedOnGpu(
233 `'.^' needs a literal integer exponent in [-8, 8] (got ` +
234 `${k === undefined ? 'a runtime value' : k}); a negative base makes ` +
235 `WGSL's pow() undefined`,
236 span,
237 );
238}
240/** Fixed-exponent power helpers, emitted only when used. */
241function powHelpers(used: Set<number>): string {
242 const out: string[] = [];
243 for (const k of [...used].sort((a, b) => a - b)) {
244 const body =
245 k === 1 ? 'x' : `x${' * x'.repeat(k - 1)}`;
246 out.push(`fn pow_i${k}(x: f32) -> f32 { return ${body}; }`);
247 }
248 return out.join('\n');
249}
251/**
252 * Reject implicit expansion (broadcasting).
253 *
254 * numbl's lowering permits it — `2x4096 .* 1x4096` lowers happily with MATLAB
255 * expansion semantics — but a kernel that walks one linear index across every
256 * operand would quietly compute the wrong thing. So every multi-element
257 * operand must have exactly the target's shape. Scalars are fine: they are
258 * read from the params buffer or folded in as literals.
259 */
260function checkShapes(e: IRExpr, target: NumericType, name: string): void {
261 const want = target.shape;
262 const same = (t: NumericType): boolean => {
263 const got = t.shape;
264 return (
265 !!want && !!got && want.length === got.length &&
266 want.every((d, i) => d === got[i])
267 );
268 };
269 const walk = (x: IRExpr): void => {
270 if (isNumeric(x.ty) && isMultiElement(x.ty) && !same(x.ty)) {
271 const got = x.ty.shape?.join('x') ?? 'dynamic';
272 throw new UnsupportedOnGpu(
273 `'${name}' would need implicit expansion: an operand is ${got} but the ` +
274 `result is ${want?.join('x') ?? 'dynamic'}. Expand it explicitly ` +
275 `(the GPU kernel walks one index across every operand).`,
276 x.span,
277 );
278 }
279 switch (x.kind) {
280 case 'Binary':
281 walk(x.left);
282 walk(x.right);
283 return;
284 case 'Unary':
285 walk(x.operand);
286 return;
287 case 'Call':
288 // A shape constructor's own arguments are sizes, not data.
289 if (x.name !== 'ones' && x.name !== 'zeros') x.args.forEach(walk);
290 return;
291 default:
292 return;
293 }
294 };
295 walk(e);
296}
298export const WORKGROUP_SIZE = 64;
300export interface Kernel {
301 code: string;
302 /** Number of output elements. */
303 count: number;
304 label: string;
305}
307/**
308 * Build the kernel for one element-wise `Assign`. `io` must already map every
309 * tensor operand cName to a binding index and every runtime scalar to a
310 * params slot; the output is binding 0 and the params buffer is the binding
311 * after the last input.
312 */
313export function buildKernel(
314 stmt: Assign,
315 io: KernelInputs,
316 count: number,
317 label: string,
318): Kernel {
319 if (!isNumeric(stmt.ty)) {
320 throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric array`, stmt.span);
321 }
322 if (stmt.ty.isComplex) {
323 throw new UnsupportedOnGpu(
324 `'${stmt.name}' is complex; the GPU backend is real-only (a spectral ` +
325 `field is carried as a real 2 x nlm array)`,
326 stmt.span,
327 );
328 }
330 checkShapes(stmt.expr, stmt.ty, stmt.name);
331 const ctx: Ctx = { io, prologue: [], bound: new Map() };
332 const body = emitExpr(stmt.expr, ctx);
334 // pow_iK helpers are discovered during emission; scan the result for them.
335 const used = new Set<number>();
336 const emitted = [...ctx.prologue, body].join('\n');
337 for (const m of emitted.matchAll(/\bpow_i(\d+)\(/g)) used.add(Number(m[1]));
339 const decls = [`@group(0) @binding(0) var<storage, read_write> out: array<f32>;`];
340 for (const [, slot] of io.tensors) {
341 decls.push(
342 `@group(0) @binding(${slot + 1}) var<storage, read> in${slot}: array<f32>;`,
343 );
344 }
345 // Params live in a read-only storage buffer rather than a uniform block:
346 // uniform arrays would need 16-byte element stride.
347 const prmBinding = io.tensors.size + 1;
348 decls.push(
349 `@group(0) @binding(${prmBinding}) var<storage, read> prm: array<f32>;`,
350 );
352 const code = `${decls.join('\n')}
354${powHelpers(used)}
356@compute @workgroup_size(${WORKGROUP_SIZE})
357fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
358 let i = gid.x;
359 if (i >= ${count}u) { return; }
360${ctx.prologue.length ? `${ctx.prologue.join('\n')}\n` : ''} out[i] = ${body};
361}
362`;
363 return { code, count, label };
364}