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);
208 }
210 /** The plan whose grid `readSpecies` samples on — the display plan when
211 * oversampling, otherwise the solver's. Its cosTheta/nphi define the mesh. */
212 get viewSht(): ShtPlan {
213 return this.#displaySht ?? this.sht;
214 }
216 /**
217 * Change the display oversampling in place. Display-only: the simulation
218 * state, time and parameters are untouched, so the run continues seamlessly
219 * on the new render grid. The caller must not have a readSpecies in flight —
220 * its readback maps a buffer of the plan being destroyed.
221 */
222 async setOversample(oversample: number): Promise<void> {
223 const os = Math.max(1, Math.round(oversample));
224 if (os === this.#oversample) return;
225 const next =
226 os > 1
227 ? await ShtPlan.create(this.device, {
228 lmax: this.cfg.lmax,
229 mmax: this.cfg.mmax,
230 nlat: os * this.cfg.nlat,
231 nphi: os * this.cfg.nphi,
232 })
233 : null;
234 const old = this.#displaySht;
235 this.#displaySht = next;
236 this.#oversample = os;
237 old?.destroy();
238 }
240 /** Run `init` from a seeded perturbation, resetting model time. */
241 seed(seed: number): void {
242 this.gpu.init(seededNoise(this.npts, this.model.seedAmp, seed));
243 this.t = 0;
244 this.steps = 0;
245 }
247 setParams(params: ModelParams): void {
248 this.#params = params;
249 this.gpu.setParams(params);
250 }
252 /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
253 step(n = 1): void {
254 this.gpu.step(n);
255 this.t += n * (this.#params.dt ?? 0);
256 this.steps += n;
257 }
259 /**
260 * Wait for the submitted steps to finish, without reading anything back.
261 * This is the honest way to time the solver: a readback would add a GPU->CPU
262 * round trip, which in a browser also crosses a process boundary and can cost
263 * more than the steps themselves.
264 */
265 sync(): Promise<undefined> {
266 return this.device.queue.onSubmittedWorkDone();
267 }
269 /**
270 * Time a batch of `n` steps and return ms/step, leaving the simulation
271 * exactly where it was: the spectral state is snapshotted before the batch
272 * and restored after, and `t`/`steps` do not advance. One sync amortized
273 * over the batch — the same measurement the desktop benchmark makes. The
274 * grid view fields hold the batch's output until the next real step, so
275 * step before reading them.
276 */
277 async measure(n: number): Promise<number> {
278 this.gpu.snapshotState();
279 const t0 = performance.now();
280 this.gpu.step(n);
281 await this.sync();
282 const ms = (performance.now() - t0) / n;
283 this.gpu.restoreState();
284 return ms;
285 }
287 /** Read a named value (a grid field or the spectral state). */
288 read(name: string): Promise<Float32Array> {
289 return this.gpu.read(name);
290 }
292 /**
293 * Read species `k` at render resolution (`viewSht`'s grid). Without
294 * oversampling this is the grid field the .m returned. With oversampling the
295 * spectral state is synthesized on the finer grid instead — the same field,
296 * since the models define each species as synth of its state, evaluated
297 * exactly on more points.
298 */
299 readSpecies(k: number): Promise<Float32Array> {
300 if (!this.#displaySht) return this.read(this.model.species[k]);
301 const state = this.model.state[k];
302 const buf = this.gpu.valueBuffer(state);
303 if (!buf) throw new Error(`readSpecies: no buffer for state '${state}'`);
304 return this.#displaySht.synthFrom(buf);
305 }
307 describe(): { init: string[]; step: string[] } {
308 return this.gpu.describe();
309 }
311 destroy(): void {
312 this.gpu.destroy();
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 313 this.#deriv.destroy();
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 314 this.#displaySht?.destroy();
315 this.sht.destroy();
316 }
317}