1/**
2 * IR expression tree -> one WGSL compute kernel.
3 *
4 * The elementwise core follows turing-surface's emitter: for an `Assign` whose
5 * right-hand side is purely element-wise over operands of the target's shape,
6 * emit a single kernel that computes one output element per invocation.
7 * Because numbl's inline pass has already folded the ANF temps back together,
8 * one source line of MATLAB becomes one kernel.
9 *
10 * The sandbox extends it with:
11 * - comparisons and eager logicals (`<`, `&`, `~`, ...), carried as f32 0/1;
12 * - inline *generators* — `rand`, `randn`, `linspace`, ranges, `zeros`,
13 * `ones`, `eye` — evaluated per element from the linear index, so
14 * `x = 2*rand(n,1) - 1` is one kernel and touches no other buffer;
15 * - runtime scalars read from 1-element storage buffers (`in3[0]`);
16 * - loop variables read from a per-loop dynamic-offset uniform, so a `for`
17 * body compiles once and replays.
18 *
19 * Everything is f32 — WGSL has no f64. Arrays are column-major linear buffers,
20 * matching MATLAB, so `A(:)` and `reshape` are views of the same buffer.
21 */
22import { isMultiElement } from 'numbl-src/numbl-core/jit/lowering/types.ts';
23import type {
24 IRExpr,
25 Assign,
26 IndexSlice,
27 Span,
28} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
29import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
30import { UnsupportedOnGpu } from './errors.ts';
32export const WORKGROUP_SIZE = 64;
34const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
35const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
36export const numel = (t: NumericType): number =>
37 (t.shape ?? []).reduce((a, b) => a * b, 1);
39/** The compile-time value of a scalar expression, if it has one. */
40export const exactValue = (e: IRExpr): number | undefined => {
41 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
42 return e.kind === 'NumLit' ? e.value : undefined;
43};
45/** Element-wise binary builtins -> WGSL infix operator. */
46const BINARY_OPS: Record<string, string> = {
47 plus: '+',
48 minus: '-',
49 times: '*',
50 rdivide: '/',
51 // Degenerate to element-wise when at least one side is a scalar; the
52 // both-tensor (true matrix) case never reaches here — the planner routes it
53 // to the GEMM kernel (mtimes) or rejects it (mrdivide).
54 mtimes: '*',
55 mrdivide: '/',
56};
58/** Comparison builtins -> WGSL comparison; result carried as f32 0/1. */
59const COMPARE_OPS: Record<string, string> = {
60 lt: '<',
61 le: '<=',
62 gt: '>',
63 ge: '>=',
64 eq: '==',
65 ne: '!=',
66};
68/** Element-wise unary builtins -> WGSL prefix operator. */
69const UNARY_OPS: Record<string, string> = { uminus: '-', uplus: '+' };
71/** Element-wise builtin calls -> WGSL builtin of the same arity. */
72const CALL_FNS: Record<string, string> = {
73 abs: 'abs',
74 acos: 'acos',
75 asin: 'asin',
76 atan: 'atan',
77 atan2: 'atan2',
78 ceil: 'ceil',
79 cos: 'cos',
80 cosh: 'cosh',
81 exp: 'exp',
82 fix: 'trunc',
83 floor: 'floor',
84 log: 'log',
85 log2: 'log2',
86 round: 'round',
87 sign: 'sign',
88 sin: 'sin',
89 sinh: 'sinh',
90 sqrt: 'sqrt',
91 tan: 'tan',
92 tanh: 'tanh',
93};
95/** Reductions the planner materializes before a kernel is built. Their 1-arg
96 * (and for dot, 2-arg) tensor forms never reach the elementwise emitter. */
97export const REDUCTIONS = new Set([
98 'sum', 'mean', 'prod', 'max', 'min', 'norm', 'dot',
99]);
101/** WGSL f32 literal. Must always carry a decimal point or exponent, or WGSL
102 * infers AbstractInt and rejects the mixed-type arithmetic. */
103function f32Lit(v: number): string {
104 if (!Number.isFinite(v)) {
105 // WGSL has no NaN/Inf literal, and 0.0/0.0 is a const-eval error;
106 // bitcast the IEEE pattern instead.
107 if (Number.isNaN(v)) return 'bitcast<f32>(0x7fc00000u)';
108 return v > 0 ? 'bitcast<f32>(0x7f800000u)' : 'bitcast<f32>(0xff800000u)';
109 }
110 return Number.isInteger(v) && Math.abs(v) < 1e21
111 ? `${v}.0`
112 : String(v).includes('e')
113 ? `${v}f`
114 : String(v);
115}
117/** `A(:)` — the only IndexSlice this backend executes. Returns the base Var
118 * expression, which reads the same buffer at the same linear index (a
119 * column-major flatten is the identity on the linear buffer). */
120export function fullColonBase(e: IndexSlice): IRExpr | null {
121 if (e.index.length !== 1 || e.index[0].kind !== 'Colon') return null;
122 return e.base;
123}
125/** How operands are read inside a kernel. */
126export interface KernelInputs {
127 /** cName -> storage binding slot, for tensors AND runtime scalars (a
128 * runtime scalar is a 1-element buffer, read as `inN[0]`). */
129 buffers: Map<string, number>;
130 /** cVar -> dense per-kernel index of enclosing loop variables, each bound
131 * as a dynamic-offset uniform (`lvK`). */
132 loopVars: Map<string, number>;
133 /** Distinct-per-call-site seeds for `rand`/`randn`. Global to the plan, so
134 * two kernels never share a stream. */
135 nextSeed: () => number;
136}
138interface Ctx {
139 io: KernelInputs;
140 /** Emitted-helper flags, gathered during emission. */
141 usesHash: boolean;
142 usedPows: Set<number>;
143 usesMod: boolean;
144 usesRem: boolean;
145}
147/** Mix every enclosing loop's iteration counter into a hash lane, so a
148 * generator inside a replayed loop draws fresh values each iteration. */
149function seedExpr(seed: number, ctx: Ctx): string {
150 let s = `${seed >>> 0}u`;
151 for (const [, k] of ctx.io.loopVars) {
152 s += ` ^ (lv${k}.it * ${[2654435761, 2246822519, 3266489917, 668265263][k % 4]}u)`;
153 }
154 return s;
155}
157/** Is `t` a value with more than one element? (logical/double both count) */
158const multi = (t: Type): boolean => isTensor(t);
160/**
161 * Emit the per-element WGSL expression for `e`. `i` is the element index
162 * variable in scope. Comparisons/logicals produce f32 0/1 so any consumer
163 * can treat them as numbers, exactly like MATLAB's logicals.
164 */
165function emitExpr(e: IRExpr, ctx: Ctx): string {
166 // Anything numbl constant-folded (pi, 2*pi, n-1, ...) is a literal, no
167 // matter what expression kind computed it.
168 const exact = exactValue(e);
169 if (exact !== undefined) return f32Lit(exact);
170 const io = ctx.io;
171 switch (e.kind) {
172 case 'NumLit':
173 return f32Lit(e.value);
175 case 'Var': {
176 const lv = io.loopVars.get(e.cName);
177 if (lv !== undefined) return `lv${lv}.v`;
178 if (isNumeric(e.ty) && typeof e.ty.exact === 'number') {
179 return f32Lit(e.ty.exact);
180 }
181 const slot = io.buffers.get(e.cName);
182 if (slot === undefined) {
183 throw new UnsupportedOnGpu(`no buffer bound for '${e.name}'`, e.span);
184 }
185 return multi(e.ty) ? `in${slot}[i]` : `in${slot}[0]`;
186 }
188 case 'IndexSlice': {
189 const base = fullColonBase(e);
190 if (!base || base.kind !== 'Var') {
191 throw new UnsupportedOnGpu(
192 `only the full linearization 'X(:)' of a variable is supported — ` +
193 `general indexing/slicing is not implemented on the GPU yet`,
194 e.span,
195 );
196 }
197 return emitExpr(base, ctx);
198 }
200 case 'MakeRange': {
201 // start + i*step; the count is already fixed in the node's type.
202 const start = emitScalar(e.start, ctx, `range start`);
203 const stepV = exactValue(e.step);
204 if (stepV === undefined) {
205 throw new UnsupportedOnGpu(`a range's step must be a compile-time value`, e.span);
206 }
207 return `(${start} + f32(i) * ${f32Lit(stepV)})`;
208 }
210 case 'Binary': {
211 if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
212 multi(e.left.ty) && multi(e.right.ty)) {
213 // The planner materializes tensor mtimes into a GEMM before building
214 // the kernel; reaching here means it could not (mrdivide) or a
215 // planner bug (mtimes).
216 throw new UnsupportedOnGpu(
217 e.builtin === 'mrdivide'
218 ? `matrix '/' (mrdivide) is not supported; use './' or a factorization`
219 : `internal: tensor '*' was not materialized as a GEMM`,
220 e.span,
221 );
222 }
223 if (e.builtin === 'power' || e.builtin === 'mpower') {
224 return emitPower(e.left, e.right, ctx, e.span);
225 }
226 const cmp = COMPARE_OPS[e.builtin];
227 if (cmp) {
228 return `select(0.0, 1.0, ${emitExpr(e.left, ctx)} ${cmp} ${emitExpr(e.right, ctx)})`;
229 }
230 if (e.builtin === 'and' || e.builtin === 'or' || e.builtin === 'andand' || e.builtin === 'oror') {
231 const op = e.builtin === 'and' || e.builtin === 'andand' ? '&&' : '||';
232 return `select(0.0, 1.0, (${emitExpr(e.left, ctx)} != 0.0) ${op} (${emitExpr(e.right, ctx)} != 0.0))`;
233 }
234 const op = BINARY_OPS[e.builtin];
235 if (!op) {
236 throw new UnsupportedOnGpu(`operator '${e.builtin}' is not supported`, e.span);
237 }
238 return `(${emitExpr(e.left, ctx)} ${op} ${emitExpr(e.right, ctx)})`;
239 }
241 case 'Unary': {
242 if (e.builtin === 'not') {
243 return `select(1.0, 0.0, ${emitExpr(e.operand, ctx)} != 0.0)`;
244 }
245 if (e.builtin === 'transpose') {
246 // A vector transpose changes orientation only — the linear buffer is
247 // identical. A matrix transpose is materialized by the planner.
248 if (isVectorish(e.operand.ty)) return emitExpr(e.operand, ctx);
249 throw new UnsupportedOnGpu(
250 `internal: matrix transpose was not materialized`,
251 e.span,
252 );
253 }
254 const op = UNARY_OPS[e.builtin];
255 if (!op) {
256 throw new UnsupportedOnGpu(`unary '${e.builtin}' is not supported`, e.span);
257 }
258 return `(${op}${emitExpr(e.operand, ctx)})`;
259 }
261 case 'Call':
262 return emitCall(e, ctx);
264 default:
265 throw new UnsupportedOnGpu(`'${e.kind}' is not supported on the GPU`, e.span);
266 }
267}
269const isVectorish = (t: Type): boolean =>
270 isNumeric(t) && (t.shape ?? []).filter((d) => d !== 1).length <= 1;
272/** A scalar-valued subexpression (range endpoints, linspace args). */
273function emitScalar(e: IRExpr, ctx: Ctx, what: string): string {
274 if (multi(e.ty)) {
275 throw new UnsupportedOnGpu(`${what} must be a scalar`, e.span);
276 }
277 return emitExpr(e, ctx);
278}
280function emitCall(e: IRExpr & { kind: 'Call' }, ctx: Ctx): string {
281 switch (e.name) {
282 case 'zeros':
283 return '0.0';
284 case 'ones':
285 return '1.0';
286 case 'eye': {
287 const t = e.ty;
288 const m = isNumeric(t) && t.shape ? t.shape[0] : undefined;
289 if (m === undefined) {
290 throw new UnsupportedOnGpu(`'eye' needs a compile-time size`, e.span);
291 }
292 return `select(0.0, 1.0, (i % ${m}u) == (i / ${m}u))`;
293 }
294 case 'rand': {
295 ctx.usesHash = true;
296 return `rand01(i, ${seedExpr(ctx.io.nextSeed(), ctx)})`;
297 }
298 case 'randn': {
299 ctx.usesHash = true;
300 const a = seedExpr(ctx.io.nextSeed(), ctx);
301 const b = seedExpr(ctx.io.nextSeed(), ctx);
302 // Box–Muller; rand01 returns (0,1) so the log is finite.
303 return `(sqrt(-2.0 * log(rand01(i, ${a}))) * cos(6.283185307179586 * rand01(i, ${b})))`;
304 }
305 case 'linspace': {
306 if (e.args.length !== 3) {
307 throw new UnsupportedOnGpu(`'linspace' needs 3 arguments here`, e.span);
308 }
309 const n = exactValue(e.args[2]);
310 if (n === undefined) {
311 throw new UnsupportedOnGpu(`'linspace' count must be a compile-time value`, e.span);
312 }
313 const a = emitScalar(e.args[0], ctx, `'linspace' start`);
314 const b = emitScalar(e.args[1], ctx, `'linspace' end`);
315 if (n <= 1) return b; // MATLAB: linspace(a, b, 1) == b
316 return `(${a} + f32(i) * ((${b} - ${a}) * ${f32Lit(1 / (n - 1))}))`;
317 }
318 case 'mod':
319 ctx.usesMod = true;
320 return `mod_m(${emitExpr(e.args[0], ctx)}, ${emitExpr(e.args[1], ctx)})`;
321 case 'rem':
322 ctx.usesRem = true;
323 return `rem_m(${emitExpr(e.args[0], ctx)}, ${emitExpr(e.args[1], ctx)})`;
324 case 'max':
325 case 'min': {
326 // Two-arg elementwise form; the one-arg reduction never reaches here.
327 if (e.args.length !== 2) {
328 throw new UnsupportedOnGpu(`internal: '${e.name}' reduction was not materialized`, e.span);
329 }
330 return `${e.name}(${emitExpr(e.args[0], ctx)}, ${emitExpr(e.args[1], ctx)})`;
331 }
332 case 'xor':
333 return `select(0.0, 1.0, (${emitExpr(e.args[0], ctx)} != 0.0) != (${emitExpr(e.args[1], ctx)} != 0.0))`;
334 case 'double':
335 case 'logical':
336 // Representation is f32 either way.
337 return emitExpr(e.args[0], ctx);
338 default: {
339 const fn = CALL_FNS[e.name];
340 if (!fn) {
341 const isUserFunction = e.cName !== e.name;
342 throw new UnsupportedOnGpu(
343 isUserFunction
344 ? `'${e.name}' is a user-defined function — not supported in the sandbox; inline it`
345 : REDUCTIONS.has(e.name)
346 ? `internal: '${e.name}' reduction was not materialized`
347 : `'${e.name}' cannot be evaluated element-wise on the GPU`,
348 e.span,
349 );
350 }
351 return `${fn}(${e.args.map((a) => emitExpr(a, ctx)).join(', ')})`;
352 }
353 }
354}
356/**
357 * `x.^k`. WGSL's `pow` is undefined for a negative base, so expand literal
358 * integer exponents into repeated multiplication — which is also what makes
359 * `u.^2` free. Non-integer exponents fall through to `pow`, defined only for
360 * a non-negative base (as in MATLAB, where a negative base goes complex —
361 * here it is NaN, and the result cross-check will show it).
362 */
363function emitPower(base: IRExpr, exponent: IRExpr, ctx: Ctx, span: Span): string {
364 const k = exactValue(exponent);
365 const b = emitExpr(base, ctx);
366 if (k !== undefined && Number.isInteger(k) && k >= 0 && k <= 16) {
367 if (k === 0) return '1.0';
368 ctx.usedPows.add(k);
369 return `pow_i${k}(${b})`;
370 }
371 if (k !== undefined && Number.isInteger(k) && k < 0 && k >= -16) {
372 ctx.usedPows.add(-k);
373 return `(1.0 / pow_i${-k}(${b}))`;
374 }
375 return `pow(${b}, ${emitExpr(exponent, ctx)})`;
376}
378/** Fixed-exponent power helpers, emitted only when used. */
379function powHelpers(used: Set<number>): string {
380 const out: string[] = [];
381 for (const k of [...used].sort((a, b) => a - b)) {
382 const body = k === 1 ? 'x' : `x${' * x'.repeat(k - 1)}`;
383 out.push(`fn pow_i${k}(x: f32) -> f32 { return ${body}; }`);
384 }
385 return out.join('\n');
386}
388/** PCG-style hash -> (0,1). Counter-based: a call site's stream is a pure
389 * function of (element index, seed), so runs are reproducible. */
390const HASH_HELPERS = `
391fn hash_u(x0: u32) -> u32 {
392 var x = x0 * 747796405u + 2891336453u;
393 x = ((x >> ((x >> 28u) + 4u)) ^ x) * 277803737u;
394 return (x >> 22u) ^ x;
395}
396fn rand01(i: u32, seed: u32) -> f32 {
397 return (f32(hash_u(i ^ (seed * 2654435769u)) & 0x00FFFFFFu) + 0.5) * (1.0 / 16777216.0);
398}`;
400const MOD_HELPER = `
401fn mod_m(a: f32, b: f32) -> f32 { return select(a - b * floor(a / b), a, b == 0.0); }`;
402const REM_HELPER = `
403fn rem_m(a: f32, b: f32) -> f32 { return select(a - b * trunc(a / b), a, b == 0.0); }`;
405/**
406 * Reject implicit expansion (broadcasting).
407 *
408 * numbl's lowering permits it — `2x4096 .* 1x4096` lowers happily with MATLAB
409 * expansion semantics — but a kernel that walks one linear index across every
410 * operand would quietly compute the wrong thing. So every multi-element
411 * operand must have exactly the target's number of elements (a flattened or
412 * transposed vector reads the same linear buffer, so only numel must match).
413 */
414function checkShapes(e: IRExpr, target: NumericType, name: string): void {
415 const want = numel(target);
416 const walk = (x: IRExpr): void => {
417 if (isNumeric(x.ty) && isMultiElement(x.ty)) {
418 const got = x.ty.shape ? numel(x.ty) : undefined;
419 if (got !== want) {
420 throw new UnsupportedOnGpu(
421 `'${name}' would need implicit expansion: an operand is ` +
422 `${x.ty.shape?.join('x') ?? 'dynamic'} but the result has ${want} ` +
423 `elements. Expand it explicitly (the GPU kernel walks one linear ` +
424 `index across every operand).`,
425 x.span,
426 );
427 }
428 // Same-numel vectors of different orientation share a linear layout;
429 // same-numel *matrices* of different shape do not (transpose is not a
430 // relayout numbl would insert silently, so shapes agree here).
431 }
432 switch (x.kind) {
433 case 'Binary':
434 walk(x.left);
435 walk(x.right);
436 return;
437 case 'Unary':
438 walk(x.operand);
439 return;
440 case 'IndexSlice':
441 return; // the base reads through the slice's own (checked) type
442 case 'Call':
443 // A generator's arguments are sizes/endpoints, not per-element data.
444 if (!['zeros', 'ones', 'eye', 'rand', 'randn', 'linspace'].includes(x.name)) {
445 x.args.forEach(walk);
446 }
447 return;
448 default:
449 return;
450 }
451 };
452 walk(e);
453}
455export interface Kernel {
456 code: string;
457 /** Number of output elements. */
458 count: number;
459 label: string;
460 /** Buffer operand cNames in binding order (bindings 1..n). */
461 buffers: string[];
462 /** Loop-variable cVars in binding order (after the buffers). */
463 loops: string[];
464}
466/** True if the expression draws random numbers anywhere. */
467function usesRandom(e: IRExpr): boolean {
468 let found = false;
469 const walk = (x: IRExpr): void => {
470 if (found) return;
471 switch (x.kind) {
472 case 'Call':
473 if (x.name === 'rand' || x.name === 'randn') found = true;
474 else x.args.forEach(walk);
475 return;
476 case 'Binary':
477 walk(x.left);
478 walk(x.right);
479 return;
480 case 'Unary':
481 walk(x.operand);
482 return;
483 case 'IndexSlice':
484 walk(x.base);
485 return;
486 default:
487 return;
488 }
489 };
490 walk(e);
491 return found;
492}
494/** Every buffer-backed variable the expression reads (tensors and runtime
495 * scalars), and every loop variable. */
496export function collectReads(
497 e: IRExpr,
498 isLoopVar: (cName: string) => boolean,
499 isExact: (x: IRExpr) => boolean,
500 visitBuffer: (cName: string) => void,
501 visitLoop: (cName: string) => void,
502): void {
503 const walk = (x: IRExpr): void => {
504 switch (x.kind) {
505 case 'Var':
506 if (isLoopVar(x.cName)) visitLoop(x.cName);
507 else if (!isExact(x)) visitBuffer(x.cName);
508 return;
509 case 'Binary':
510 walk(x.left);
511 walk(x.right);
512 return;
513 case 'Unary':
514 walk(x.operand);
515 return;
516 case 'IndexSlice':
517 walk(x.base);
518 return;
519 case 'MakeRange':
520 walk(x.start);
521 walk(x.step);
522 return;
523 case 'Call':
524 if (!['zeros', 'ones', 'eye', 'rand', 'randn'].includes(x.name)) {
525 x.args.forEach(walk);
526 }
527 return;
528 default:
529 return;
530 }
531 };
532 walk(e);
533}
535/**
536 * Can this expression live inside one fused GPU kernel?
537 *
538 * Wider than numbl's own `isPureElementwiseExpr`: the WGSL emitter fuses
539 * transcendental calls, comparisons/logicals, generators, ranges, `X(:)` and
540 * vector transposes, all of which numbl's C-side pass declines. The sandbox's
541 * fuse pass uses this to fold the temps numbl's inline pass left behind.
542 */
543export function isGpuFusableExpr(e: IRExpr): boolean {
544 if (exactValue(e) !== undefined) return true;
545 switch (e.kind) {
546 case 'NumLit':
547 return true;
548 case 'Var':
549 return isNumeric(e.ty);
550 case 'Binary': {
551 if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
552 multi(e.left.ty) && multi(e.right.ty)) {
553 return false;
554 }
555 if (e.builtin === 'mpower' && multi(e.left.ty)) return false; // matrix power
556 const known =
557 e.builtin in BINARY_OPS || e.builtin in COMPARE_OPS ||
558 ['and', 'or', 'andand', 'oror', 'power', 'mpower'].includes(e.builtin);
559 return known && isGpuFusableExpr(e.left) && isGpuFusableExpr(e.right);
560 }
561 case 'Unary': {
562 if (e.builtin === 'transpose') {
563 return isVectorish(e.operand.ty) && isGpuFusableExpr(e.operand);
564 }
565 return (e.builtin in UNARY_OPS || e.builtin === 'not') && isGpuFusableExpr(e.operand);
566 }
567 case 'Call': {
568 if (['zeros', 'ones', 'eye', 'rand', 'randn'].includes(e.name)) return true;
569 if (e.name === 'linspace') return e.args.length === 3;
570 if (['mod', 'rem', 'xor', 'atan2'].includes(e.name)) {
571 return e.args.every(isGpuFusableExpr);
572 }
573 if ((e.name === 'max' || e.name === 'min') && e.args.length === 2) {
574 return e.args.every(isGpuFusableExpr);
575 }
576 if (e.name === 'double' || e.name === 'logical') {
577 return e.args.length === 1 && isGpuFusableExpr(e.args[0]);
578 }
579 return e.name in CALL_FNS && e.args.every(isGpuFusableExpr);
580 }
581 case 'IndexSlice': {
582 const base = fullColonBase(e);
583 return !!base && base.kind === 'Var';
584 }
585 case 'MakeRange':
586 return exactValue(e.step) !== undefined;
587 default:
588 return false;
589 }
590}
592/** A fused elementwise subexpression, for embedding inside a non-elementwise
593 * kernel (a reduction's per-element load). `body` reads element `i`. */
594export interface FusedLoader {
595 body: string;
596 helpers: string;
597}
599/**
600 * Emit `e` as a per-element load for a reduction kernel. Same contract as
601 * `buildKernel`: `io.buffers` maps operands to binding slots; loop variables
602 * are bound on demand (all of them, if the expression draws random numbers).
603 */
604export function emitLoader(
605 e: IRExpr,
606 io: KernelInputs,
607 enclosingLoops: string[],
608): FusedLoader {
609 if (usesRandom(e)) {
610 for (const cVar of enclosingLoops) {
611 if (!io.loopVars.has(cVar)) io.loopVars.set(cVar, io.loopVars.size);
612 }
613 }
614 const ctx: Ctx = {
615 io,
616 usesHash: false,
617 usedPows: new Set(),
618 usesMod: false,
619 usesRem: false,
620 };
621 const body = emitExpr(e, ctx);
622 const helpers = [
623 ctx.usesHash ? HASH_HELPERS : '',
624 ctx.usesMod ? MOD_HELPER : '',
625 ctx.usesRem ? REM_HELPER : '',
626 powHelpers(ctx.usedPows),
627 ].filter(Boolean).join('\n');
628 return { body, helpers };
629}
631/** Binding declarations shared by every kernel shape: output at 0, operand
632 * buffers after it, loop-variable uniforms after those. */
633export function bindingDecls(io: KernelInputs): {
634 decls: string[];
635 buffers: string[];
636 loops: string[];
637} {
638 const decls = [`@group(0) @binding(0) var<storage, read_write> out: array<f32>;`];
639 const buffers: string[] = [];
640 for (const [cName, slot] of io.buffers) {
641 buffers[slot] = cName;
642 decls.push(
643 `@group(0) @binding(${slot + 1}) var<storage, read> in${slot}: array<f32>;`,
644 );
645 }
646 const loops: string[] = [];
647 if (io.loopVars.size) {
648 decls.push(`struct Lv { v: f32, it: u32 }`);
649 for (const [cVar, k] of io.loopVars) {
650 loops[k] = cVar;
651 decls.push(
652 `@group(0) @binding(${io.buffers.size + 1 + k}) var<uniform> lv${k}: Lv;`,
653 );
654 }
655 }
656 return { decls, buffers, loops };
657}
659/**
660 * Build the fused elementwise kernel for one `Assign`. `io.buffers` must
661 * already map every buffer operand to a binding slot; the output is binding 0
662 * and loop-variable uniforms follow the last input.
663 *
664 * `enclosingLoops` lists the cVars of the loops this statement sits inside
665 * (outermost first). Loop variables the expression reads are bound; if the
666 * expression draws random numbers, ALL enclosing loop counters are bound and
667 * mixed into the stream so each replayed iteration draws fresh values.
668 */
669export function buildKernel(
670 stmt: Pick<Assign, 'name' | 'cName' | 'ty' | 'expr' | 'span'>,
671 io: KernelInputs,
672 enclosingLoops: string[],
673 label: string,
674): Kernel {
675 if (!isNumeric(stmt.ty)) {
676 throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric array`, stmt.span);
677 }
678 if (stmt.ty.isComplex) {
679 throw new UnsupportedOnGpu(
680 `'${stmt.name}' is complex; the GPU backend is real-only (f32)`,
681 stmt.span,
682 );
683 }
684 const count = numel(stmt.ty);
685 if (isMultiElement(stmt.ty)) checkShapes(stmt.expr, stmt.ty, stmt.name);
687 // Random draws need every enclosing loop counter; make sure they are bound
688 // before emission asks for them.
689 if (usesRandom(stmt.expr)) {
690 for (const cVar of enclosingLoops) {
691 if (!io.loopVars.has(cVar)) io.loopVars.set(cVar, io.loopVars.size);
692 }
693 }
695 const ctx: Ctx = {
696 io,
697 usesHash: false,
698 usedPows: new Set(),
699 usesMod: false,
700 usesRem: false,
701 };
702 const body = emitExpr(stmt.expr, ctx);
703 const { decls, buffers, loops } = bindingDecls(io);
705 const helpers = [
706 ctx.usesHash ? HASH_HELPERS : '',
707 ctx.usesMod ? MOD_HELPER : '',
708 ctx.usesRem ? REM_HELPER : '',
709 powHelpers(ctx.usedPows),
710 ].filter(Boolean).join('\n');
712 // Dispatch is 2-D so counts past 65535 workgroups still fit: x rows of
713 // ELEMENTS_PER_ROW elements each. dispatchFor() picks matching counts.
714 const code = `${decls.join('\n')}
715${helpers}
716@compute @workgroup_size(${WORKGROUP_SIZE})
717fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
718 let i = gid.x + gid.y * ${ELEMENTS_PER_ROW}u;
719 if (i >= ${count}u) { return; }
720 out[i] = ${body};
721}
722`;
723 return { code, count, label, buffers, loops };
724}
726/** Elements covered by one row of the 2-D elementwise dispatch. */
727export const ELEMENTS_PER_ROW = 32768 * WORKGROUP_SIZE;
729/** Workgroup counts for an elementwise dispatch over `count` elements. */
730export function dispatchFor(count: number): [number, number] {
731 const rows = Math.ceil(count / ELEMENTS_PER_ROW);
732 const x = rows === 1 ? Math.ceil(count / WORKGROUP_SIZE) : 32768;
733 return [x, rows];
734}