1/**
2 * Type-rule patches for numbl JIT builtins, applied to the in-process builtin
3 * registry before lowering a script.
4 *
5 * numbl's JIT declines a handful of constructs its C/JS backends do not
6 * implement yet (tensor `&`/`|`/`~`, two-arg `max`/`min` on tensors, matrix
7 * `rand`/`randn`), and its comparison builtins deliberately return a
8 * shape-unknown logical tensor because the lowerer has no static broadcast
9 * helper. This backend executes the IR on WebGPU — numbl's own emitters never
10 * run here — so all this sandbox needs from those builtins is a precise
11 * *type*: everything else is supplied by the WGSL kernels.
12 *
13 * `registerBuiltin` overwrites by name (by design, for HMR), which makes it
14 * the sanctioned way to swap a builtin's rules. The patches only ever widen:
15 * every case the original transfer accepts is delegated to it unchanged, so
16 * scalar behavior (including exact-value folding) is untouched.
17 *
18 * IMPORTANT: the patch is process-global to this module instance of numbl.
19 * The CPU comparison runner executes scripts through numbl's normal engine in
20 * a separate Web Worker, whose module instances are its own — it never sees
21 * these patches.
22 */
23import {
24 getBuiltin,
25 registerBuiltin,
26 type Builtin,
27} from 'numbl-src/numbl-core/jit/builtins/index.ts';
28import {
29 isNumeric,
30 isMultiElement,
31 type NumericType,
32 type Type,
33} from 'numbl-src/numbl-core/jit/lowering/types.ts';
34import { UnsupportedOnGpu } from './errors.ts';
36const isRealTensor = (t: Type): t is NumericType =>
37 isNumeric(t) && !t.isComplex && isMultiElement(t);
38const isRealScalarish = (t: Type): t is NumericType =>
39 isNumeric(t) && !t.isComplex && !isMultiElement(t);
41const logicalTensor = (shape: number[]): NumericType => ({
42 kind: 'Numeric',
43 elem: 'logical',
44 isComplex: false,
45 dims: shape.map((n) => ({ kind: 'exact', value: n })),
46 shape: [...shape],
47 sign: 'nonneg',
48});
50const doubleTensor = (shape: number[], sign: NumericType['sign']): NumericType => ({
51 kind: 'Numeric',
52 elem: 'double',
53 isComplex: false,
54 dims: shape.map((n) => ({ kind: 'exact', value: n })),
55 shape: [...shape],
56 sign,
57});
59const sameShape = (a: number[], b: number[]): boolean =>
60 a.length === b.length && a.every((d, i) => d === b[i]);
62/**
63 * The elementwise result shape for two operands, under the same rule the
64 * kernels implement: operands are either scalars or tensors of one common
65 * shape. Implicit expansion (`n x 1` with `1 x m`) is rejected here, matching
66 * the emitter's refusal to broadcast.
67 */
68function elementwiseShape(name: string, argTypes: Type[]): number[] | null {
69 let shape: number[] | null = null;
70 for (const t of argTypes) {
71 if (!isNumeric(t) || t.isComplex) {
72 throw new UnsupportedOnGpu(`'${name}': operands must be real numeric`);
73 }
74 if (!isMultiElement(t)) continue;
75 if (!t.shape) {
76 throw new UnsupportedOnGpu(
77 `'${name}': operand shape is not known at compile time`,
78 );
79 }
80 if (shape && !sameShape(shape, t.shape)) {
81 throw new UnsupportedOnGpu(
82 `'${name}': operands are ${shape.join('x')} and ${t.shape.join('x')}; ` +
83 `the GPU kernels do not broadcast — expand explicitly`,
84 );
85 }
86 shape = t.shape;
87 }
88 return shape;
89}
91/** Wrap `transfer` so tensor operands get an elementwise result type instead
92 * of the original's decline (or, for comparisons, its shape-unknown type). */
93function widenElementwise(
94 name: string,
95 makeResult: (shape: number[]) => NumericType,
96 /** Arity the elementwise form requires. `max(x)` with one tensor arg is the
97 * reduction form and must stay with the original transfer. */
98 arity?: number,
99): void {
100 const orig = getBuiltin(name);
101 if (!orig) throw new Error(`patch: no builtin named '${name}'`);
102 const origTransfer = orig.transfer.bind(orig);
103 registerBuiltin({
104 ...orig,
105 transfer(argTypes: Type[], nargout: number): Type[] {
106 const anyTensor = argTypes.some((t) => isNumeric(t) && isMultiElement(t));
107 if (!anyTensor || (arity !== undefined && argTypes.length !== arity)) {
108 return origTransfer(argTypes, nargout);
109 }
110 if (nargout > 1) {
111 throw new UnsupportedOnGpu(`'${name}' returns one value`);
112 }
113 const shape = elementwiseShape(name, argTypes);
114 // anyTensor guaranteed a tensor operand, so shape is non-null.
115 return [makeResult(shape!)];
116 },
117 } as Builtin);
118}
120/** Comparisons: `lt` et al. accept tensors already but return unknown dims;
121 * replace the tensor case with the precise shape. */
122const COMPARISONS = ['lt', 'le', 'gt', 'ge', 'eq', 'ne'];
124/** Eager elementwise `&` / `|` / `xor` / `~`. (`&&`/`||` stay scalar-only —
125 * that is MATLAB semantics, not a backend gap.) */
126const LOGICALS = ['and', 'or', 'xor', 'not'];
128/** Two-arg elementwise forms of `max`/`min`. The one-arg reduction form is
129 * typed by the original transfer. */
130const MINMAX = ['max', 'min'];
132/** `rand(n)` / `rand(m, n)` / `randn(...)`: numbl's JS-JIT supports only the
133 * scalar form; here any exact-sized matrix form is a fill kernel. */
134function widenRand(name: string, sign: NumericType['sign']): void {
135 // `randn` has no JIT builtin at all in numbl yet; register it whole.
136 const orig = getBuiltin(name);
137 const origTransfer = orig
138 ? orig.transfer.bind(orig)
139 : (): Type[] => [{
140 kind: 'Numeric', elem: 'double', isComplex: false,
141 dims: [{ kind: 'exact', value: 1 }, { kind: 'exact', value: 1 }],
142 shape: [1, 1], sign,
143 } satisfies NumericType];
144 registerBuiltin({
145 ...(orig ?? { name }),
146 name,
147 transfer(argTypes: Type[], nargout: number): Type[] {
148 if (argTypes.length === 0) return origTransfer(argTypes, nargout);
149 if (nargout > 1) throw new UnsupportedOnGpu(`'${name}' returns one value`);
150 if (argTypes.length > 2) {
151 throw new UnsupportedOnGpu(
152 `'${name}': only the 2-D forms ${name}(n) / ${name}(m, n) are supported`,
153 );
154 }
155 const dims = argTypes.map((t) => {
156 if (!isRealScalarish(t) || typeof t.exact !== 'number') {
157 throw new UnsupportedOnGpu(
158 `'${name}': array sizes must be known at compile time (assign the ` +
159 `size from a literal, e.g. n = 1024)`,
160 );
161 }
162 if (!Number.isInteger(t.exact) || t.exact < 0) {
163 throw new UnsupportedOnGpu(`'${name}': sizes must be whole numbers`);
164 }
165 return t.exact;
166 });
167 const shape = dims.length === 1 ? [dims[0], dims[0]] : dims;
168 return [doubleTensor(shape, sign)];
169 },
170 } as Builtin);
171}
173let applied = false;
175/** Apply every patch, once per module instance. */
176export function applyBuiltinPatches(): void {
177 if (applied) return;
178 applied = true;
179 for (const name of COMPARISONS) widenElementwise(name, logicalTensor);
180 for (const name of LOGICALS) widenElementwise(name, logicalTensor);
181 for (const name of MINMAX) {
182 widenElementwise(name, (shape) => doubleTensor(shape, 'unknown'), 2);
183 }
184 widenRand('rand', 'positive');
185 widenRand('randn', 'unknown');
186}