1/**
2 * `dot(x, y)` — the one reduction the solvers need, as a GPU operation.
3 *
4 * A Krylov solver's scalars (rho, alpha, omega, ...) are inner products of the
5 * spectral state, and the whole point of the plan architecture is that no
6 * value crosses back to the CPU mid-step — so the dot product must produce a
7 * GPU-resident scalar: a 1-element buffer that later kernels read (the
8 * planner binds any single-element value as `in<slot>[0]`, see
9 * src/mgpu/wgsl.ts).
10 *
11 * One workgroup does the whole reduction: each thread accumulates a strided
12 * partial sum, then a shared-memory tree combines them and thread 0 writes
13 * the result. A single dispatch, no multi-pass bookkeeping, and — because
14 * the striding is fixed — a bit-deterministic summation order on a given
15 * device. n up to a few hundred thousand is a short loop per thread, far
16 * below anything the solver grids produce.
17 */
19const WG = 256;
21/** One dot call site: its bind group and the pipeline for its length. */
22export interface DotBinding {
23 readonly pipeline: GPUComputePipeline;
24 readonly bindGroup: GPUBindGroup;
25}
27const dotWGSL = (n: number): string => `
28@group(0) @binding(0) var<storage, read_write> out: array<f32>;
29@group(0) @binding(1) var<storage, read> a: array<f32>;
30@group(0) @binding(2) var<storage, read> b: array<f32>;
32var<workgroup> partials: array<f32, ${WG}>;
34@compute @workgroup_size(${WG})
35fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
36 var s = 0.0;
37 var i = lid.x;
38 loop {
39 if (i >= ${n}u) { break; }
40 s = s + a[i] * b[i];
41 i = i + ${WG}u;
42 }
43 partials[lid.x] = s;
44 workgroupBarrier();
45 var stride = ${WG / 2}u;
46 loop {
47 if (stride == 0u) { break; }
48 if (lid.x < stride) {
49 partials[lid.x] = partials[lid.x] + partials[lid.x + stride];
50 }
51 workgroupBarrier();
52 stride = stride / 2u;
53 }
54 if (lid.x == 0u) {
55 out[0] = partials[0];
56 }
57}
58`;
60export class ReducePlan {
61 #device: GPUDevice;
62 #layout: GPUBindGroupLayout;
63 /** Element count is baked into the shader, so pipelines cache per length. */
64 #byN = new Map<number, GPUComputePipeline>();
66 constructor(device: GPUDevice) {
67 this.#device = device;
68 const entry = (
69 binding: number,
70 type: GPUBufferBindingType,
71 ): GPUBindGroupLayoutEntry => ({
72 binding,
73 visibility: GPUShaderStage.COMPUTE,
74 buffer: { type },
75 });
76 this.#layout = device.createBindGroupLayout({
77 entries: [entry(0, 'storage'), entry(1, 'read-only-storage'), entry(2, 'read-only-storage')],
78 });
79 }
81 async #pipeline(n: number): Promise<GPUComputePipeline> {
82 const cached = this.#byN.get(n);
83 if (cached) return cached;
84 const device = this.#device;
85 device.pushErrorScope('validation');
86 const code = dotWGSL(n);
87 const module = device.createShaderModule({ code, label: `dot-${n}` });
88 const info = await module.getCompilationInfo();
89 const errors = info.messages.filter((m) => m.type === 'error');
90 if (errors.length) {
91 throw new Error(
92 `WGSL compile error in dot(${n}):\n` +
93 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
94 );
95 }
96 const pipeline = await device.createComputePipelineAsync({
97 layout: device.createPipelineLayout({ bindGroupLayouts: [this.#layout] }),
98 compute: { module, entryPoint: 'main' },
99 label: `dot-${n}`,
100 });
101 const err = await device.popErrorScope();
102 if (err) throw new Error(`pipeline dot(${n}): ${err.message}`);
103 this.#byN.set(n, pipeline);
104 return pipeline;
105 }
107 async createDotBinding(
108 a: GPUBuffer,
109 b: GPUBuffer,
110 out: GPUBuffer,
111 n: number,
112 ): Promise<DotBinding> {
113 const pipeline = await this.#pipeline(n);
114 const bindGroup = this.#device.createBindGroup({
115 layout: this.#layout,
116 entries: [
117 { binding: 0, resource: { buffer: out } },
118 { binding: 1, resource: { buffer: a } },
119 { binding: 2, resource: { buffer: b } },
120 ],
121 });
122 return { pipeline, bindGroup };
123 }
125 encodeDotInto(pass: GPUComputePassEncoder, binding: DotBinding): void {
126 pass.setPipeline(binding.pipeline);
127 pass.setBindGroup(0, binding.bindGroup);
128 pass.dispatchWorkgroups(1);
129 }
130}