1/**
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;
22}
24export class ModelSession {
25 readonly device: GPUDevice;
26 readonly model: MModel;
27 readonly cfg: ShtConfig;
28 readonly sht: ShtPlan;
29 readonly gpu: GpuModel;
30 readonly npts: number;
32 /** Model time and step count since the last seeding. */
33 t = 0;
34 steps = 0;
36 #params: ModelParams;
38 private constructor(init: {
39 device: GPUDevice;
40 model: MModel;
41 cfg: ShtConfig;
42 sht: ShtPlan;
43 gpu: GpuModel;
44 params: ModelParams;
45 }) {
46 this.device = init.device;
47 this.model = init.model;
48 this.cfg = init.cfg;
49 this.sht = init.sht;
50 this.gpu = init.gpu;
51 this.npts = init.cfg.nlat * init.cfg.nphi;
52 this.#params = init.params;
53 }
55 static async create(opts: ModelSessionOptions): Promise<ModelSession> {
56 const { device, model, params, lmax } = opts;
57 const { nlat, nphi } = gridForLmax(lmax, model.pdeg);
58 const cfg = { lmax, mmax: lmax, nlat, nphi };
59 const sht = await ShtPlan.create(device, cfg);
60 try {
61 const gpu = await GpuModel.create({
62 device,
63 sht,
64 cfg,
65 source: opts.source ?? model.source,
66 paramNames: model.params.map((p) => p.key),
67 state: model.state,
68 view: model.species,
69 });
70 gpu.setParams(params);
71 return new ModelSession({ device, model, cfg, sht, gpu, params });
72 } catch (e) {
73 // The transform plan owns GPU buffers; do not leak them on a compile error.
74 sht.destroy();
75 throw e;
76 }
77 }
79 /** Run `init` from a seeded perturbation, resetting model time. */
80 seed(seed: number): void {
81 this.gpu.init(seededNoise(this.npts, this.model.seedAmp, seed));
82 this.t = 0;
83 this.steps = 0;
84 }
86 setParams(params: ModelParams): void {
87 this.#params = params;
88 this.gpu.setParams(params);
89 }
91 /** Advance `n` steps. Synchronous: records and submits, nothing read back. */
92 step(n = 1): void {
93 this.gpu.step(n);
94 this.t += n * (this.#params.dt ?? 0);
95 this.steps += n;
96 }
98 /**
99 * Wait for the submitted steps to finish, without reading anything back.
100 * This is the honest way to time the solver: a readback would add a GPU->CPU
101 * round trip, which in a browser also crosses a process boundary and can cost
102 * more than the steps themselves.
103 */
104 sync(): Promise<undefined> {
105 return this.device.queue.onSubmittedWorkDone();
106 }
108 /** Read a named value (a grid field or the spectral state). */
109 read(name: string): Promise<Float32Array> {
110 return this.gpu.read(name);
111 }
113 describe(): { init: string[]; step: string[] } {
114 return this.gpu.describe();
115 }
117 destroy(): void {
118 this.gpu.destroy();
119 this.sht.destroy();
120 }
121}