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';
10import { gridForLmax, type ShtConfig } from '../sht/layout.ts';
11import { GpuModel, type ModelParams } from './model.ts';
12import { seededNoise } from './noise.ts';
13import type { MModel } from './registry.ts';
15export interface ModelSessionOptions {
16 device: GPUDevice;
17 model: MModel;
18 params: ModelParams;
19 lmax: number;
20 /** Override the model source — the editor's working copy. */
21 source?: string;
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 22 /** Linear render oversampling: read the species fields on a grid this many
23 * times finer than the solver's in each direction (default 1). The state is
24 * band-limited at lmax, so the finer evaluation is exact interpolation. */
25 oversample?: number;
28export class ModelSession {
0f3abbdSeparate solver time from frame time, and add a cross-environment checkJeremy Magland 29 readonly device: GPUDevice;
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 30 readonly model: MModel;
31 readonly cfg: ShtConfig;
32 readonly sht: ShtPlan;
33 readonly gpu: GpuModel;
34 readonly npts: number;
36 /** Model time and step count since the last seeding. */
37 t = 0;
38 steps = 0;
40 #params: ModelParams;
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 41 /** Display-only transforms on the oversampled grid; null at 1x. */
42 #displaySht: ShtPlan | null;
43 #oversample: number;
45 private constructor(init: {
0f3abbdSeparate solver time from frame time, and add a cross-environment checkJeremy Magland 46 device: GPUDevice;
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 47 model: MModel;
48 cfg: ShtConfig;
49 sht: ShtPlan;
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 50 displaySht: ShtPlan | null;
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 51 gpu: GpuModel;
52 params: ModelParams;
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 53 oversample: number;
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 54 }) {
0f3abbdSeparate solver time from frame time, and add a cross-environment checkJeremy Magland 55 this.device = init.device;
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 56 this.model = init.model;
57 this.cfg = init.cfg;
58 this.sht = init.sht;
59 this.gpu = init.gpu;
60 this.npts = init.cfg.nlat * init.cfg.nphi;
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 61 this.#oversample = init.oversample;
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 62 this.#params = init.params;
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 63 this.#displaySht = init.displaySht;
64 }
66 /** Linear render oversampling factor (1 = read on the solver grid). */
67 get oversample(): number {
68 return this.#oversample;
71 static async create(opts: ModelSessionOptions): Promise<ModelSession> {
72 const { device, model, params, lmax } = opts;
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 73 const oversample = Math.max(1, Math.round(opts.oversample ?? 1));
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 74 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
75 const cfg = { lmax, mmax: lmax, nlat, nphi };
76 const sht = await ShtPlan.create(device, cfg);
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 77 let displaySht: ShtPlan | null = null;
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 78 try {
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 79 // The display plan shares nothing with the solver's beyond the
80 // coefficients copied into it per readback; its grid is the solver's
81 // scaled by the oversampling factor, so nphi stays a power of two (the
82 // FFT path) for power-of-two factors.
83 if (oversample > 1) {
84 displaySht = await ShtPlan.create(device, {
85 lmax,
86 mmax: lmax,
87 nlat: oversample * nlat,
88 nphi: oversample * nphi,
89 });
90 }
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 91 const gpu = await GpuModel.create({
92 device,
93 sht,
94 cfg,
95 source: opts.source ?? model.source,
96 paramNames: model.params.map((p) => p.key),
97 state: model.state,
98 view: model.species,
99 });
100 gpu.setParams(params);
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 101 return new ModelSession({
102 device, model, cfg, sht, displaySht, gpu, params, oversample,
103 });
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 104 } catch (e) {
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 105 // The transform plans own GPU buffers; do not leak them on a compile error.
106 displaySht?.destroy();
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 107 sht.destroy();
108 throw e;
109 }
110 }
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 112 /** The plan whose grid `readSpecies` samples on — the display plan when
113 * oversampling, otherwise the solver's. Its cosTheta/nphi define the mesh. */
114 get viewSht(): ShtPlan {
115 return this.#displaySht ?? this.sht;
116 }
118 /**
119 * Change the display oversampling in place. Display-only: the simulation
120 * state, time and parameters are untouched, so the run continues seamlessly
121 * on the new render grid. The caller must not have a readSpecies in flight —
122 * its readback maps a buffer of the plan being destroyed.
123 */
124 async setOversample(oversample: number): Promise<void> {
125 const os = Math.max(1, Math.round(oversample));
126 if (os === this.#oversample) return;
127 const next =
128 os > 1
129 ? await ShtPlan.create(this.device, {
130 lmax: this.cfg.lmax,
131 mmax: this.cfg.mmax,
132 nlat: os * this.cfg.nlat,
133 nphi: os * this.cfg.nphi,
134 })
135 : null;
136 const old = this.#displaySht;
137 this.#displaySht = next;
138 this.#oversample = os;
139 old?.destroy();
140 }
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 142 /** Run `init` from a seeded perturbation, resetting model time. */
143 seed(seed: number): void {
144 this.gpu.init(seededNoise(this.npts, this.model.seedAmp, seed));
145 this.t = 0;
146 this.steps = 0;
147 }
149 setParams(params: ModelParams): void {
150 this.#params = params;
151 this.gpu.setParams(params);
152 }
154 /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
155 step(n = 1): void {
156 this.gpu.step(n);
157 this.t += n * (this.#params.dt ?? 0);
158 this.steps += n;
159 }
0f3abbdSeparate solver time from frame time, and add a cross-environment checkJeremy Magland 161 /**
162 * Wait for the submitted steps to finish, without reading anything back.
163 * This is the honest way to time the solver: a readback would add a GPU->CPU
164 * round trip, which in a browser also crosses a process boundary and can cost
165 * more than the steps themselves.
166 */
167 sync(): Promise<undefined> {
168 return this.device.queue.onSubmittedWorkDone();
169 }
172 * Time a batch of `n` steps and return ms/step, leaving the simulation
173 * exactly where it was: the spectral state is snapshotted before the batch
174 * and restored after, and `t`/`steps` do not advance. One sync amortized
175 * over the batch — the same measurement the desktop benchmark makes. The
176 * grid view fields hold the batch's output until the next real step, so
177 * step before reading them.
178 */
179 async measure(n: number): Promise<number> {
180 this.gpu.snapshotState();
181 const t0 = performance.now();
182 this.gpu.step(n);
183 await this.sync();
184 const ms = (performance.now() - t0) / n;
185 this.gpu.restoreState();
186 return ms;
187 }
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 189 /** Read a named value (a grid field or the spectral state). */
190 read(name: string): Promise<Float32Array> {
191 return this.gpu.read(name);
192 }
195 * Read species `k` at render resolution (`viewSht`'s grid). Without
196 * oversampling this is the grid field the .m returned. With oversampling the
197 * spectral state is synthesized on the finer grid instead — the same field,
198 * since the models define each species as synth of its state, evaluated
199 * exactly on more points.
200 */
201 readSpecies(k: number): Promise<Float32Array> {
202 if (!this.#displaySht) return this.read(this.model.species[k]);
203 const state = this.model.state[k];
204 const buf = this.gpu.valueBuffer(state);
205 if (!buf) throw new Error(`readSpecies: no buffer for state '${state}'`);
206 return this.#displaySht.synthFrom(buf);
207 }
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 209 describe(): { init: string[]; step: string[] } {
210 return this.gpu.describe();
211 }
213 destroy(): void {
214 this.gpu.destroy();
163ec45Render on demand, display oversampling, and jump-free solver timingJeremy Magland 215 this.#displaySht?.destroy();
35d91faDelete the TypeScript solver; the .m models are the only implementationJeremy Magland 216 this.sht.destroy();
217 }
218}