1/**
2 * `randnfun3` — a smooth random function in 3D, evaluated at the surface.
3 *
4 * chebfun's randnfun3 is a random trig series on a box: a few thousand
5 * Fourier modes with independent normal coefficients, confined to a ball for
6 * isotropy and normalized to unit variance. Restricting it to a surface is
7 * just evaluating it at the surface's points, which is what a model's `init`
8 * wants for a seeded initial condition (surfacefun seeds exactly this way).
9 *
10 * The work splits in two, and the split is forced rather than chosen:
11 *
12 * - **Drawing the modes needs `randn`**, which the compiled WGSL dialect has
13 * no counterpart for, and `sqrt(nnz)` normalization, which is a reduction.
14 * Both are a few lines of MATLAB, so the draw lives in
15 * `tools/randnfun3.m` and runs in numbl's interpreter — a few thousand
16 * numbers, ~5 ms.
17 * - **Evaluating is npts x nmodes**, ~6e7 terms at the default lambda. That
18 * is the whole cost, and it is what this file's kernel does on the GPU.
19 *
20 * So the .m calls `f = randnfun3(lambda, gx, gy, gz)` — chebfun's signature,
21 * lambda in and values out — and the coefficient table is filled in behind it
22 * by the host, the way `synth` hides its Legendre matrices. lambda is not
23 * decorative: the plan records which parameter the .m passed, and the host
24 * draws the table from *that* parameter's value (src/mgpu/plan.ts,
25 * `randnfun3Lambda`), so changing it in the .m changes the field.
26 */
27import { executeCode } from 'numbl-src/numbl-core/executeCode.ts';
28import { isRuntimeTensor } from 'numbl-src/numbl-core/runtime/types.ts';
29import { toolFiles } from '../tools.ts';
31/**
32 * Dispatches the mode sum is split across.
33 *
34 * lambda is an absolute length and the mode count goes as its inverse cube,
35 * so halving lambda costs eight times the work — there is no natural ceiling
36 * to put on that, and nothing in the method breaks as it grows. It just gets
37 * slower, which is the caller's business. What is *not* the caller's business
38 * is a browser's GPU-process watchdog, which kills the device outright when a
39 * single dispatch runs too long; a fine wavelength would otherwise turn "this
40 * takes a while" into "device lost".
41 *
42 * So the sum is split into a fixed number of dispatches, each covering its own
43 * slice of the table and accumulating into the same output. The count is fixed
44 * at plan time (the op sequence has no runtime branching) and the slice bounds
45 * come from the table's header, so one plan serves any wavelength. Slices that
46 * fall past the end of a small table exit immediately, which is why a coarse
47 * wavelength pays nothing for the split.
48 */
49const CHUNKS = 16;
51/** Floats the table needs for `nmodes` modes. */
52export const modeTableLength = (nmodes: number): number =>
53 HEADER + STRIDE * nmodes;
55/** Modes a table holds, from its header. */
56export const modeCount = (table: Float32Array): number => table[0];
58/**
59 * Largest table this will try to build, in f32. Not a policy about how fine a
60 * wavelength is sensible — that is the caller's call, and a fine one is
61 * merely slow — but the point past which the draw would fail anyway: the
62 * host-side Float32Array alone would be 8 GB. The device's own
63 * storage-buffer limit is checked separately, when the buffer is allocated.
64 */
65const MAX_TABLE_FLOATS = 2 ** 31;
67/** What the table starts at, before any seed has been drawn. Big enough for
68 * the default wavelength on the shipped surfaces, so the common case never
69 * reallocates. */
70export const INITIAL_MODES = 4096;
72/** Wavelength of the seeded field when the app names none. Fine enough to
73 * give a Turing pattern plenty to grow from, coarse enough that the draw is
74 * ~1,400 modes rather than the ~11,500 of the slider's finest setting. */
75export const DEFAULT_LAMBDA = 0.5;
77/** Floats before the first mode: `[nmodes, 0, 0, 0]`. The count travels in
78 * the buffer rather than a second binding, so the kernel needs one storage
79 * buffer and the host one write. */
80const HEADER = 4;
81/** Floats per mode: kx, ky, kz, real, imag. */
82const STRIDE = 5;
84/** The name the coefficient buffer takes in the plan's HostBuffers. */
85export const MODE_BUFFER = 'randnfun3_modes';
87/** How many dispatches `randnfun3WGSL` must be planned as. */
88export const randnfun3Chunks = CHUNKS;
90/**
91 * One thread per surface point, summing this chunk's slice of the modes.
92 *
93 * The inner loop is a dot product, a cos, a sin and two multiply-adds, over a
94 * table small enough (~1,400 modes at the default lambda) to sit in cache for
95 * every thread. Chunk 0 initializes the output and the rest accumulate onto
96 * it; dispatches within one compute pass are ordered, so the reads see the
97 * previous chunk's writes. Nothing here is per-step work: `init` runs once a
98 * seed.
99 */
100export function randnfun3WGSL(npts: number, chunk: number): string {
101 return `
102@group(0) @binding(0) var<storage, read_write> outf: array<f32>;
103@group(0) @binding(1) var<storage, read> px: array<f32>;
104@group(0) @binding(2) var<storage, read> py: array<f32>;
105@group(0) @binding(3) var<storage, read> pz: array<f32>;
106// [nmodes, _, _, _], then kx, ky, kz, re, im per mode.
107@group(0) @binding(4) var<storage, read> modes: array<f32>;
109@compute @workgroup_size(64)
110fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
111 let i = gid.x;
112 if (i >= ${npts}u) { return; }
113 let n = u32(modes[0]);
114 // This chunk's slice. Ceiling division, so the last slices are the short
115 // ones and an empty slice costs a single comparison.
116 let per = (n + ${CHUNKS}u - 1u) / ${CHUNKS}u;
117 let lo = min(${chunk}u * per, n);
118 let hi = min(lo + per, n);
119 var acc = 0.0;
120 if (lo < hi) {
121 let x = px[i];
122 let y = py[i];
123 let z = pz[i];
124 for (var m = lo; m < hi; m = m + 1u) {
125 let b = ${HEADER}u + m * ${STRIDE}u;
126 let t = modes[b] * x + modes[b + 1u] * y + modes[b + 2u] * z;
127 acc = acc + modes[b + 3u] * cos(t) - modes[b + 4u] * sin(t);
128 }
129 }
130${chunk === 0 ? ' outf[i] = acc;' : ' outf[i] = outf[i] + acc;'}
131}
132`;
133}
135/**
136 * Modes a wavelength will draw on a box, without drawing them: chebfun's
137 * cube size, times the fraction its isotropy ball keeps (pi/6 of a cube,
138 * approached from below at small m). Used to price a wavelength up front.
139 */
140function plannedModes(lambda: number, box: BoundingBox): number {
141 const side = (w: number): number => 2 * Math.round((1.2 * w) / lambda + 2) + 1;
142 const cube =
143 side(box.x1 - box.x0) * side(box.y1 - box.y0) * side(box.z1 - box.z0);
144 return Math.ceil((Math.PI / 6) * cube);
145}
147/** The box a random field is drawn over: the surface's own bounding box. */
148export interface BoundingBox {
149 x0: number; x1: number;
150 y0: number; y1: number;
151 z0: number; z1: number;
152}
154/** The bounding box of a surface, as `Geometry` holds its coordinates. */
155export function boundingBox(
156 x: Float32Array,
157 y: Float32Array,
158 z: Float32Array,
159): BoundingBox {
160 const box = {
161 x0: Infinity, x1: -Infinity,
162 y0: Infinity, y1: -Infinity,
163 z0: Infinity, z1: -Infinity,
164 };
165 for (let i = 0; i < x.length; i++) {
166 if (x[i] < box.x0) box.x0 = x[i];
167 if (x[i] > box.x1) box.x1 = x[i];
168 if (y[i] < box.y0) box.y0 = y[i];
169 if (y[i] > box.y1) box.y1 = y[i];
170 if (z[i] < box.z0) box.z0 = z[i];
171 if (z[i] > box.z1) box.z1 = z[i];
172 }
173 return box;
174}
176/**
177 * Draw a field's modes and pack them for the GPU: `tools/randnfun3.m` run
178 * through the interpreter, seeded, then interleaved into the buffer layout
179 * above. Column-major out of MATLAB, interleaved on the way in.
180 */
181export function drawModes(
182 lambda: number,
183 box: BoundingBox,
184 seed: number,
185 /** Points the field will be summed at, for the cost budget. */
186 npts: number,
187): Float32Array {
188 if (!(lambda > 0) || !Number.isFinite(lambda)) {
189 throw new Error(`randnfun3: lambda must be a positive number, got ${lambda}`);
190 }
191 // The mode count follows from lambda and the box alone, so a table that
192 // cannot be built is refused before anything is drawn. The only ceiling is
193 // what fits: how slow a fine wavelength is, is the caller's to decide.
194 const planned = plannedModes(lambda, box);
195 if (modeTableLength(planned) > MAX_TABLE_FLOATS) {
196 throw new Error(
197 `randnfun3: lambda ${lambda} needs about ` +
198 `${planned.toLocaleString()} Fourier modes on this surface, a ` +
199 `${((4 * modeTableLength(planned)) / 1e9).toFixed(1)} GB table. ` +
200 `lambda is an absolute length, so a larger surface needs more modes ` +
201 `for the same value, and halving it costs eight times as many.`,
202 );
203 }
204 const result = executeCode(
205 'rng(seed); [k, c] = randnfun3(lambda, [x0 x1 y0 y1 z0 z1]);',
206 {
207 initialVariableValues: { lambda, seed, ...box },
208 displayResults: false,
209 implicitCwdPath: null,
210 },
211 toolFiles,
212 'randnfun3-driver.m',
213 );
214 const k = result.variableValues['k'];
215 const c = result.variableValues['c'];
216 if (!k || !c || !isRuntimeTensor(k) || !isRuntimeTensor(c)) {
217 throw new Error("randnfun3: tools/randnfun3.m did not return [k, c] arrays");
218 }
219 const nmodes = k.shape[0];
220 const out = new Float32Array(modeTableLength(nmodes));
221 out[0] = nmodes;
222 for (let i = 0; i < nmodes; i++) {
223 const b = HEADER + STRIDE * i;
224 out[b] = k.data[i]; // kx
225 out[b + 1] = k.data[nmodes + i]; // ky
226 out[b + 2] = k.data[2 * nmodes + i]; // kz
227 out[b + 3] = c.data[i]; // real
228 out[b + 4] = c.data[nmodes + i]; // imag
229 }
230 return out;
231}
233/**
234 * `drawModes` on a worker thread, so a fine wavelength does not freeze the
235 * page (src/mgpu/randnfun3.worker.ts).
236 *
237 * Falls back to drawing in place where there is no `Worker` — the node test
238 * runner and the desktop benchmark, neither of which has an event loop it
239 * would matter to. Failures surface as a rejection either way, so a caller
240 * never has to know which path ran.
241 */
242export function drawModesAsync(
243 lambda: number,
244 box: BoundingBox,
245 seed: number,
246 npts: number,
247): Promise<Float32Array> {
248 if (typeof Worker === 'undefined') {
249 try {
250 return Promise.resolve(drawModes(lambda, box, seed, npts));
251 } catch (e) {
252 return Promise.reject(e instanceof Error ? e : new Error(String(e)));
253 }
254 }
255 const w = drawWorker();
256 const id = nextDrawId++;
257 return new Promise((resolve, reject) => {
258 pendingDraws.set(id, { resolve, reject });
259 w.postMessage({ id, lambda, box, seed, npts });
260 });
261}
263let worker: Worker | null = null;
264let nextDrawId = 1;
265const pendingDraws = new Map<
266 number,
267 { resolve: (t: Float32Array) => void; reject: (e: Error) => void }
268>();
270/** The draw worker, started on first use and kept for the session — starting
271 * one re-parses numbl, which costs more than a coarse draw does. */
272function drawWorker(): Worker {
273 if (worker) return worker;
274 worker = new Worker(new URL('./randnfun3.worker.ts', import.meta.url), {
275 type: 'module',
276 });
277 worker.onmessage = (e: MessageEvent<{ id: number; table?: Float32Array; error?: string }>): void => {
278 const waiting = pendingDraws.get(e.data.id);
279 if (!waiting) return;
280 pendingDraws.delete(e.data.id);
281 if (e.data.error !== undefined) waiting.reject(new Error(e.data.error));
282 else waiting.resolve(e.data.table!);
283 };
284 worker.onerror = (e: ErrorEvent): void => {
285 // A worker that died takes every outstanding draw with it.
286 for (const [, waiting] of pendingDraws) {
287 waiting.reject(new Error(`randnfun3 draw worker failed: ${e.message}`));
288 }
289 pendingDraws.clear();
290 worker?.terminate();
291 worker = null;
292 };
293 return worker;
294}