2 * One running model: grid, transforms, compiled .m, seeded state.
3 *
4 * Everything that is not rendering. The app, the desktop benchmark and the
5 * tests all go through this, so there is one place that decides how a model is
6 * turned into something running on the GPU — and nothing about it is
7 * browser-specific beyond needing a GPUDevice.
8 */
9import { ShtPlan } from '../sht/sht.ts';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 10import { DerivPlan } from '../sht/deriv.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 11import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
12import { GpuModel, type ModelParams } from './model.ts';
13import { seededNoise } from './noise.ts';
14import type { MModel } from './registry.ts';
15import { Geometry } from '../geom/geometry.ts';
16import { mGeometryByKey, defaultGeometryParams, SPHERE_KEY, type MGeometry } from '../geom/registry.ts';
18export interface ModelSessionOptions {
19 device: GPUDevice;
20 model: MModel;
21 params: ModelParams;
22 lmax: number;
23 /** Override the model source — the editor's working copy. */
24 source?: string;
25 /** Linear render oversampling: read the species fields on a grid this many
26 * times finer than the solver's in each direction (default 1). The state is
27 * band-limited at lmax, so the finer evaluation is exact interpolation. */
28 oversample?: number;
29 /** The surface to solve on. Defaults to the unit sphere. */
30 geometry?: MGeometry;
31 geometryParams?: ModelParams;
32 /** Override the geometry source — the editor's working copy. */
33 geometrySource?: string;
34 /**
35 * Iterations of the .m's implicit solve. Structural, not tunable: the loop
36 * is unrolled into the op sequence, so a change recompiles.
37 */
38 niter?: number;
39}
41export class ModelSession {
42 readonly device: GPUDevice;
43 readonly model: MModel;
44 readonly cfg: ShtConfig;
45 readonly sht: ShtPlan;
46 readonly gpu: GpuModel;
47 readonly npts: number;
48 /** Iterations of the implicit solve compiled into the step. */
49 readonly niter: number;
51 /** The surface being solved on, as spherical-harmonic coefficients. */
52 #geometry: Geometry;
53 #geometryModel: MGeometry;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 54 /** Computes the theta/phi derivatives a geometry's metric quantities need. */
55 #deriv: DerivPlan;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 56
57 /** Model time and step count since the last seeding. */
58 t = 0;
59 steps = 0;
61 #params: ModelParams;
62 /** Display-only transforms on the oversampled grid; null at 1x. */
63 #displaySht: ShtPlan | null;
64 #oversample: number;
66 private constructor(init: {
67 device: GPUDevice;
68 model: MModel;
69 cfg: ShtConfig;
70 sht: ShtPlan;
71 displaySht: ShtPlan | null;
72 gpu: GpuModel;
73 params: ModelParams;
74 oversample: number;
75 geometry: Geometry;
76 geometryModel: MGeometry;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 77 deriv: DerivPlan;
79 }) {
80 this.device = init.device;
81 this.model = init.model;
82 this.cfg = init.cfg;
83 this.sht = init.sht;
84 this.gpu = init.gpu;
85 this.npts = init.cfg.nlat * init.cfg.nphi;
86 this.#oversample = init.oversample;
87 this.#params = init.params;
88 this.#displaySht = init.displaySht;
89 this.#geometry = init.geometry;
90 this.#geometryModel = init.geometryModel;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 91 this.#deriv = init.deriv;
93 }
95 get geometry(): Geometry {
96 return this.#geometry;
97 }
99 get geometryModel(): MGeometry {
100 return this.#geometryModel;
101 }
103 /** Linear render oversampling factor (1 = read on the solver grid). */
104 get oversample(): number {
105 return this.#oversample;
106 }
108 static async create(opts: ModelSessionOptions): Promise<ModelSession> {
109 const { device, model, params, lmax } = opts;
110 const oversample = Math.max(1, Math.round(opts.oversample ?? 1));
111 const niter = Math.max(0, Math.round(opts.niter ?? 1));
112 const geometryModel = opts.geometry ?? mGeometryByKey(SPHERE_KEY)!;
113 const geometryParams = opts.geometryParams ?? defaultGeometryParams(geometryModel);
114 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
115 const cfg = { lmax, mmax: lmax, nlat, nphi };
116 const sht = await ShtPlan.create(device, cfg);
117 let displaySht: ShtPlan | null = null;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 118 let deriv: DerivPlan | null = null;
120 // The display plan shares nothing with the solver's beyond the
121 // coefficients copied into it per readback; its grid is the solver's
122 // scaled by the oversampling factor, so nphi stays a power of two (the
123 // FFT path) for power-of-two factors.
124 if (oversample > 1) {
125 displaySht = await ShtPlan.create(device, {
126 lmax,
127 mmax: lmax,
128 nlat: oversample * nlat,
129 nphi: oversample * nphi,
130 });
131 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 132 // Computes the theta/phi derivatives the geometry's metric quantities
133 // (and, per step, the surface Laplace-Beltrami correction) need.
134 deriv = await DerivPlan.create(device, sht);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 135 // The surface is built before the model, because the model takes it as
136 // an argument. It is a one-off: compiled, evaluated, read back, and its
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 137 // plan discarded — nothing of it survives into the timestep but sixteen
138 // buffers of numbers (the embedding, and both metric formulations built
139 // on it: the inverse metric quantities and the flux-form weights).
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 140 const geometry = await Geometry.create({
141 device,
142 sht,
143 cfg,
144 source: opts.geometrySource ?? geometryModel.source,
145 paramNames: geometryModel.params.map((p) => p.key),
146 params: geometryParams,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 147 deriv,
149 const gpu = await GpuModel.create({
150 device,
151 sht,
152 cfg,
153 source: opts.source ?? model.source,
154 paramNames: model.params.map((p) => p.key),
155 state: model.state,
156 view: model.species,
157 geometry,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 158 deriv,
160 });
161 gpu.setParams(params);
162 return new ModelSession({
163 device, model, cfg, sht, displaySht, gpu, params, oversample,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 164 geometry, geometryModel, deriv, niter,
166 } catch (e) {
167 // The transform plans own GPU buffers; do not leak them on a compile error.
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 168 deriv?.destroy();
170 sht.destroy();
171 throw e;
172 }
173 }
175 /**
176 * Vertex positions for the current render grid: the surface synthesized on
177 * `viewSht`, interleaved xyz. Exact interpolation of the same coefficients
178 * the solver sees, so the drawn surface is the one being solved on however
179 * finely it is sampled.
180 */
181 renderPositions(): Promise<Float32Array> {
182 return this.#geometry.positionsOn(this.viewSht);
183 }
185 /**
186 * Change the surface in place, without recompiling or disturbing the run.
187 * The geometry's shape in the bindings depends only on the grid, so the
188 * compiled step does not change — only the numbers it reads. The caller
189 * still has to rebuild the mesh from `renderPositions()`.
190 */
191 async setGeometry(
192 geometryModel: MGeometry,
193 params: ModelParams,
194 source?: string,
195 ): Promise<void> {
196 const next = await Geometry.create({
197 device: this.device,
198 sht: this.sht,
199 cfg: this.cfg,
200 source: source ?? geometryModel.source,
201 paramNames: geometryModel.params.map((p) => p.key),
202 params,
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 203 deriv: this.#deriv,
205 this.#geometry = next;
206 this.#geometryModel = geometryModel;
207 this.gpu.uploadGeometry(next);
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 208 // The new surface brings a new preconditioner scale (GpuModel folds its
209 // current geometry's jhat into every params upload).
210 this.gpu.setParams(this.#params);
213 /** The plan whose grid `readSpecies` samples on — the display plan when
214 * oversampling, otherwise the solver's. Its cosTheta/nphi define the mesh. */
215 get viewSht(): ShtPlan {
216 return this.#displaySht ?? this.sht;
217 }
219 /**
220 * Change the display oversampling in place. Display-only: the simulation
221 * state, time and parameters are untouched, so the run continues seamlessly
222 * on the new render grid. The caller must not have a readSpecies in flight —
223 * its readback maps a buffer of the plan being destroyed.
224 */
225 async setOversample(oversample: number): Promise<void> {
226 const os = Math.max(1, Math.round(oversample));
227 if (os === this.#oversample) return;
228 const next =
229 os > 1
230 ? await ShtPlan.create(this.device, {
231 lmax: this.cfg.lmax,
232 mmax: this.cfg.mmax,
233 nlat: os * this.cfg.nlat,
234 nphi: os * this.cfg.nphi,
235 })
236 : null;
237 const old = this.#displaySht;
238 this.#displaySht = next;
239 this.#oversample = os;
240 old?.destroy();
241 }
243 /** Run `init` from a seeded perturbation, resetting model time. */
244 seed(seed: number): void {
245 this.gpu.init(seededNoise(this.npts, this.model.seedAmp, seed));
246 this.t = 0;
247 this.steps = 0;
248 }
250 setParams(params: ModelParams): void {
251 this.#params = params;
252 this.gpu.setParams(params);
253 }
255 /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
256 step(n = 1): void {
257 this.gpu.step(n);
258 this.t += n * (this.#params.dt ?? 0);
259 this.steps += n;
260 }
262 /**
263 * Wait for the submitted steps to finish, without reading anything back.
264 * This is the honest way to time the solver: a readback would add a GPU->CPU
265 * round trip, which in a browser also crosses a process boundary and can cost
266 * more than the steps themselves.
267 */
268 sync(): Promise<undefined> {
269 return this.device.queue.onSubmittedWorkDone();
270 }
272 /**
273 * Time a batch of `n` steps and return ms/step, leaving the simulation
274 * exactly where it was: the spectral state is snapshotted before the batch
275 * and restored after, and `t`/`steps` do not advance. One sync amortized
276 * over the batch — the same measurement the desktop benchmark makes. The
277 * grid view fields hold the batch's output until the next real step, so
278 * step before reading them.
279 */
280 async measure(n: number): Promise<number> {
281 this.gpu.snapshotState();
282 const t0 = performance.now();
283 this.gpu.step(n);
284 await this.sync();
285 const ms = (performance.now() - t0) / n;
286 this.gpu.restoreState();
287 return ms;
288 }
290 /** Read a named value (a grid field or the spectral state). */
291 read(name: string): Promise<Float32Array> {
292 return this.gpu.read(name);
293 }
295 /**
296 * Read species `k` at render resolution (`viewSht`'s grid). Without
297 * oversampling this is the grid field the .m returned. With oversampling the
298 * spectral state is synthesized on the finer grid instead — the same field,
299 * since the models define each species as synth of its state, evaluated
300 * exactly on more points.
301 */
302 readSpecies(k: number): Promise<Float32Array> {
303 if (!this.#displaySht) return this.read(this.model.species[k]);
304 const state = this.model.state[k];
305 const buf = this.gpu.valueBuffer(state);
306 if (!buf) throw new Error(`readSpecies: no buffer for state '${state}'`);
307 return this.#displaySht.synthFrom(buf);
308 }
310 describe(): { init: string[]; step: string[] } {
311 return this.gpu.describe();
312 }
314 destroy(): void {
315 this.gpu.destroy();
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 316 this.#deriv.destroy();
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 317 this.#displaySht?.destroy();
318 this.sht.destroy();
319 }
320}