concept-collection / turing-sphere-2
turing-sphere-2 / src / mgpu / numbl.d.ts
242 lines · 7.2 KBBlameHistoryRaw
1/**
2 * The numbl compiler surface this project depends on.
3 *
4 * We reach past numbl's published entry points into its JIT internals (parser,
5 * lowerer, IR, inline pass), which its package `exports` map does not expose.
6 * Those imports resolve through the `numbl-src` alias in vite.config.ts; these
7 * declarations are what TypeScript checks against.
8 *
9 * Declaring the surface here rather than type-checking numbl's sources
10 * directly keeps this project's compiler settings independent of numbl's, and
11 * pins the exact contract we rely on. If numbl changes one of these shapes,
12 * the build breaks here with a clear diff rather than deep inside its tree.
13 *
14 * Only the nodes the WGSL backend actually walks are spelled out; every other
15 * IR kind is collapsed into a catch-all so that unhandled constructs are
16 * rejected with a message instead of being silently mis-compiled.
17 */
19declare module 'numbl-src/numbl-core/jit/lowering/types.ts' {
20 export type Sign =
21 | 'positive' | 'nonneg' | 'negative' | 'nonpositive'
22 | 'zero' | 'nonzero' | 'unknown';
24 export type DimInfo = { kind: 'exact'; value: number } | { kind: 'unknown' };
26 export type NumericExact =
27 | number
28 | Float64Array
29 | { re: number; im: number }
30 | { re: Float64Array; im: Float64Array };
32 export interface NumericType {
33 kind: 'Numeric';
34 elem: 'double' | 'logical' | 'char' | string;
35 isComplex: boolean;
36 dims: DimInfo[];
37 /** Present iff every dim is exact. */
38 shape?: number[];
39 sign: Sign;
40 exact?: NumericExact;
41 }
43 /** Everything the WGSL backend rejects. */
44 export interface NonNumericType {
45 kind: 'Void' | 'Unknown' | 'String' | 'Handle' | 'Struct' | 'Class' | 'Cell';
46 }
48 export type Type = NumericType | NonNumericType;
50 export function isMultiElement(t: NumericType): boolean;
51 export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
52 export function scalarDouble(sign?: Sign, exact?: number): NumericType;
55declare module 'numbl-src/numbl-core/jit/lowering/ir.ts' {
56 import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
58 export interface Span {
59 file: string;
60 start: number;
61 end: number;
62 }
64 export interface NumLit {
65 kind: 'NumLit';
66 value: number;
67 ty: Type;
68 span: Span;
69 }
70 export interface Var {
71 kind: 'Var';
72 name: string;
73 cName: string;
74 ty: Type;
75 span: Span;
76 }
77 export interface Binary {
78 kind: 'Binary';
79 builtin: string;
80 left: IRExpr;
81 right: IRExpr;
82 ty: Type;
83 span: Span;
84 }
85 export interface Unary {
86 kind: 'Unary';
87 builtin: string;
88 operand: IRExpr;
89 ty: Type;
90 span: Span;
91 }
92 export interface Call {
93 kind: 'Call';
94 cName: string;
95 name: string;
96 args: IRExpr[];
97 ty: Type;
98 span: Span;
99 }
100 /** Any other IR expression kind — rejected by the WGSL emitter. */
101 export interface OtherExpr {
102 kind:
103 | 'ImagLit' | 'StringLit' | 'TensorBuild' | 'TensorConcat' | 'CellLit'
104 | 'CellEmpty' | 'CellIndexLoad' | 'HandleLit' | 'HandleCaptureLoad'
105 | 'StructLit' | 'MemberLoad' | 'IndexLoad' | 'IndexSlice' | 'EndRef'
106 | 'MakeRange';
107 ty: Type;
108 span: Span;
109 }
111 export type IRExpr = NumLit | Var | Binary | Unary | Call | OtherExpr;
113 export interface Assign {
114 kind: 'Assign';
115 name: string;
116 cName: string;
117 ty: Type;
118 expr: IRExpr;
119 span: Span;
120 }
121 /** Any other IR statement kind — rejected by the planner. */
122 export interface OtherStmt {
123 kind:
124 | 'ExprStmt' | 'If' | 'While' | 'For' | 'ReturnFromFunction' | 'Break'
125 | 'Continue' | 'TypeComment' | 'MemberStore' | 'MultiAssignCall'
126 | 'IndexStore' | 'IndexSliceStore' | 'CellIndexStore';
127 span: Span;
128 }
130 export type IRStmt = Assign | OtherStmt;
132 export interface IRFunc {
133 name: string;
134 cName: string;
135 /** Parameter source names. */
136 params: string[];
137 /** Parameter cNames, parallel to `params`. */
138 cParams: string[];
139 paramTypes: Type[];
140 /** Output source names. */
141 outputs: string[];
142 /** Output cNames, parallel to `outputs`. */
143 cOutputs: string[];
144 outputTypes: Type[];
145 body: IRStmt[];
146 span: Span;
147 }
149 export interface IRProgram {
150 topLevelStmts: IRStmt[];
151 functions: Map<string, IRFunc>;
152 }
155declare module 'numbl-src/numbl-core/parser/index.ts' {
156 export interface AbstractSyntaxTree {
157 body: unknown[];
158 }
159 export function parseMFile(input: string, fileName?: string): AbstractSyntaxTree;
160 export class SyntaxError extends Error {}
163declare module 'numbl-src/numbl-core/jit/index.ts' {
164 import type { AbstractSyntaxTree } from 'numbl-src/numbl-core/parser/index.ts';
165 import type { IRProgram, IRFunc, Span } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
166 import type { Type, NumericType, Sign } from 'numbl-src/numbl-core/jit/lowering/types.ts';
168 export interface WorkspaceFile {
169 name: string;
170 source: string;
171 ast?: AbstractSyntaxTree;
172 }
174 export class Workspace {
175 constructor(mainFile: string, searchPaths?: ReadonlyArray<string>);
176 addFile(file: WorkspaceFile): void;
177 finalize(): void;
178 }
180 export interface EnvEntry {
181 cName: string;
182 ty: Type;
183 maybeUnassigned?: boolean;
184 }
186 export class Lowerer {
187 constructor(workspace: Workspace);
188 /** Pre-bindable variable scope: seed host-provided values here. */
189 env: Map<string, EnvEntry>;
190 specializations: Map<string, IRFunc>;
191 lowerProgram(ast: AbstractSyntaxTree): IRProgram;
192 }
194 /** Thrown for MATLAB the JIT pipeline cannot lower; carries a source span. */
195 export class UnsupportedConstruct extends Error {
196 span?: Span;
197 }
198 export class JitTypeError extends Error {
199 span?: Span;
200 }
202 export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
203 export function scalarDouble(sign?: Sign, exact?: number): NumericType;
204 export function isMultiElement(t: NumericType): boolean;
207declare module 'numbl-src/numbl-core/jit/lowering/specialize.ts' {
208 import type { Lowerer } from 'numbl-src/numbl-core/jit/index.ts';
209 import type { IRFunc, IRExpr, Span } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
210 import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
212 /**
213 * Lower one user function for a concrete argument-type signature. Called with
214 * a `Lowerer` as `this` (numbl's own JIT does the same), so specializations
215 * accumulate in `lowerer.specializations`.
216 */
217 export function specializeUserFunction(
218 this: Lowerer,
219 decl: unknown,
220 argTypes: Type[],
221 specSource?: string,
222 definingFile?: string,
223 preSeedOutput?: { name: string; ty: Type; initExpr: IRExpr },
224 nargout?: number,
225 callSiteSpan?: Span,
226 ): IRFunc;
229declare module 'numbl-src/numbl-core/jit/codegen/inlinePass.ts' {
230 import type { IRProgram } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
231 /** Folds single-use ANF temps into their consumer, in place. */
232 export function inlinePass(prog: IRProgram): void;
235declare module 'numbl-src/numbl-core/jit/builtins/index.ts' {
236 export interface Builtin {
237 name: string;
238 /** Safe to evaluate one output element from one input element per slot. */
239 elementwise?: boolean;
240 }
241 export function getBuiltin(name: string): Builtin | undefined;