1// Method of fundamental solutions on WebGPU.
2//
3// The same method as src/solvers/mfs/solver.m: n logarithmic point charges
4// on a curve a fixed distance 0.3 outside the boundary, their strengths
5// found by collocating the Dirichlet data at n boundary points, and the
6// potential evaluated at the requested targets. What differs is where the
7// work happens. The assembly, the dense solve, and the evaluation all run
8// as WebGPU compute passes; the host does only the O(n) geometry, in f64,
9// the way shtns-webgpu and the rest of the WebGPU work in this collection
10// keep precomputation on the CPU.
11//
12// The solve is a right-looking LU with partial pivoting, three dispatches
13// per column: pivot search and row swap in one workgroup, the multipliers,
14// then the rank-one update of the trailing submatrix. There is no blocking
15// and no GEMM, so it is memory-bound rather than compute-bound; it is the
16// straightforward implementation, and a blocked panel factorization would
17// be the next thing to try. The right-hand side rides along as an extra
18// matrix column, which is what makes forward substitution disappear; back
19// substitution is one dispatch per column after that. A sweep at n = 768
20// therefore encodes about 3000 dispatches, all into one command buffer and
21// one submit.
22//
23// Everything is f32: WebGPU has no double precision, which is what caps
24// this solver's accuracy well short of the same method on the CPU. See
25// ./wgsl.ts.
27import { requestGpu } from "../../harness/webgpuDevice";
28import type { Laplace2dProblem } from "../../problems/laplace2d/problem";
29import { EVAL_GROUP, LANES, mfsShader, TILE } from "./wgsl";
31/** Distance of the charge curve outside the boundary, as in mfs/solver.m. */
32const DELTA = 0.3;
34const KERNELS = [
35 "assemble",
36 "pivot",
37 "multipliers",
38 "update",
39 "backsub",
40 "evaluate",
41] as const;
42type Kernel = (typeof KERNELS)[number];
44export interface GpuMfsResult {
45 uEval: Float64Array;
46 uGrid: Float64Array | null;
47}
49function ceilDiv(a: number, b: number): number {
50 return Math.ceil(a / b);
51}
53export class MfsGpu {
54 private constructor(
55 private readonly device: GPUDevice,
56 private readonly layout: GPUBindGroupLayout,
57 private readonly pipelines: Record<Kernel, GPUComputePipeline>,
58 private readonly stepStride: number,
59 /** The adapter, for the result file's environment record. */
60 readonly adapter: string,
61 readonly via: string
62 ) {}
64 /** Compile the shader and build the pipelines. Done once per device; a
65 * solver call reuses them, which is where numbl's JIT sits too. */
66 static async create(): Promise<MfsGpu> {
67 const { device, adapter, via } = await requestGpu();
68 const module = device.createShaderModule({
69 code: mfsShader(),
70 label: "mfs-gpu",
71 });
72 const info = await module.getCompilationInfo();
73 const errors = info.messages.filter((m) => m.type === "error");
74 if (errors.length > 0) {
75 throw new Error(
76 "mfs-gpu failed to compile:\n" +
77 errors.map((m) => ` ${m.lineNum}:${m.linePos} ${m.message}`).join("\n")
78 );
79 }
80 const layout = device.createBindGroupLayout({
81 entries: [
82 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
83 {
84 binding: 1,
85 visibility: GPUShaderStage.COMPUTE,
86 buffer: { type: "uniform", hasDynamicOffset: true, minBindingSize: 16 },
87 },
88 { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
89 { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
90 { binding: 4, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
91 { binding: 5, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
92 { binding: 6, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
93 ],
94 });
95 const pipelineLayout = device.createPipelineLayout({ bindGroupLayouts: [layout] });
96 const built = await Promise.all(
97 KERNELS.map((k) =>
98 device.createComputePipelineAsync({
99 layout: pipelineLayout,
100 compute: { module, entryPoint: k },
101 label: `mfs-gpu ${k}`,
102 })
103 )
104 );
105 const pipelines = Object.fromEntries(
106 KERNELS.map((k, i) => [k, built[i]])
107 ) as Record<Kernel, GPUComputePipeline>;
108 return new MfsGpu(
109 device,
110 layout,
111 pipelines,
112 Math.max(16, device.limits.minUniformBufferOffsetAlignment),
113 adapter,
114 via
115 );
116 }
118 private checked = false;
120 /**
121 * One full solve at resolution n: the timed unit of the protocol, so it
122 * includes the host geometry, the buffer allocation, the assembly, the
123 * factorization, the evaluation, and the read-back.
124 */
125 async run(
126 prob: Laplace2dProblem,
127 n: number,
128 wantGrid = false
129 ): Promise<GpuMfsResult> {
130 const { device } = this;
131 const ld = n + 1;
132 const nGrid = wantGrid ? prob.nViz : 0;
133 const m = prob.nEval + nGrid;
134 const bytes = 4;
136 // --- host geometry, in f64, rounded once on the way to the device --
137 const data = new Float32Array(5 * n + 2 * m);
138 const put = (index: number, v: number) => {
139 data[index] = v;
140 };
141 for (let j = 0; j < n; j++) {
142 const t = (2 * Math.PI * j) / n;
143 const p = prob.curve(t);
144 const d = prob.curveD(t);
145 const sp = Math.hypot(d.x, d.y);
146 // The outward unit normal of the counterclockwise curve.
147 put(2 * j, p.x);
148 put(2 * j + 1, p.y);
149 put(2 * n + 2 * j, p.x + (DELTA * d.y) / sp);
150 put(2 * n + 2 * j + 1, p.y - (DELTA * d.x) / sp);
151 put(4 * n + j, prob.g(t));
152 }
153 for (let i = 0; i < prob.nEval; i++) {
154 put(5 * n + 2 * i, prob.evalXY[2 * i]);
155 put(5 * n + 2 * i + 1, prob.evalXY[2 * i + 1]);
156 }
157 for (let i = 0; i < nGrid; i++) {
158 const o = 5 * n + 2 * (prob.nEval + i);
159 put(o, prob.vizXY[2 * i]);
160 put(o + 1, prob.vizXY[2 * i + 1]);
161 }
163 // --- buffers -----------------------------------------------------
164 const S = GPUBufferUsage.STORAGE;
165 const buffers = {
166 dims: device.createBuffer({ size: 16, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST }),
167 step: device.createBuffer({
168 size: Math.max(this.stepStride * Math.max(n, 1), this.stepStride),
169 usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
170 }),
171 data: device.createBuffer({ size: data.byteLength, usage: S | GPUBufferUsage.COPY_DST }),
172 mat: device.createBuffer({ size: n * ld * bytes, usage: S }),
173 lcol: device.createBuffer({ size: n * bytes, usage: S }),
174 sol: device.createBuffer({ size: n * bytes, usage: S }),
175 out: device.createBuffer({ size: m * bytes, usage: S | GPUBufferUsage.COPY_SRC }),
176 read: device.createBuffer({
177 size: m * bytes,
178 usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
179 }),
180 };
181 try {
182 device.queue.writeBuffer(buffers.dims, 0, new Uint32Array([n, ld, m, 0]));
183 // One slot per elimination step, so k reaches the shader through a
184 // dynamic uniform offset and the whole sweep needs one bind group.
185 const steps = new Uint32Array((this.stepStride / 4) * Math.max(n, 1));
186 for (let k = 0; k < n; k++) steps[k * (this.stepStride / 4)] = k;
187 device.queue.writeBuffer(buffers.step, 0, steps);
188 device.queue.writeBuffer(buffers.data, 0, data);
190 const bind = device.createBindGroup({
191 layout: this.layout,
192 entries: [
193 { binding: 0, resource: { buffer: buffers.dims } },
194 { binding: 1, resource: { buffer: buffers.step, size: 16 } },
195 { binding: 2, resource: { buffer: buffers.data } },
196 { binding: 3, resource: { buffer: buffers.mat } },
197 { binding: 4, resource: { buffer: buffers.lcol } },
198 { binding: 5, resource: { buffer: buffers.sol } },
199 { binding: 6, resource: { buffer: buffers.out } },
200 ],
201 });
203 // The first run of a session is checked for validation errors, which
204 // is a host round trip and so does not belong in a timed run; the
205 // protocol's untimed warmups absorb it.
206 const check = !this.checked;
207 if (check) device.pushErrorScope("validation");
209 const enc = device.createCommandEncoder();
210 const pass = enc.beginComputePass();
211 const P = this.pipelines;
212 // WebGPU orders dispatches within a pass and makes each one's writes
213 // visible to the next, so the dependences here need no barriers.
214 pass.setPipeline(P.assemble);
215 pass.setBindGroup(0, bind, [0]);
216 pass.dispatchWorkgroups(ceilDiv(ld, TILE), ceilDiv(n, TILE));
217 for (let k = 0; k < n; k++) {
218 const off = k * this.stepStride;
219 pass.setPipeline(P.pivot);
220 pass.setBindGroup(0, bind, [off]);
221 pass.dispatchWorkgroups(1);
222 const rows = n - 1 - k;
223 if (rows > 0) {
224 pass.setPipeline(P.multipliers);
225 pass.setBindGroup(0, bind, [off]);
226 pass.dispatchWorkgroups(ceilDiv(rows, LANES));
227 pass.setPipeline(P.update);
228 pass.setBindGroup(0, bind, [off]);
229 pass.dispatchWorkgroups(ceilDiv(ld - 1 - k, TILE), ceilDiv(rows, TILE));
230 }
231 }
232 for (let k = n - 1; k >= 0; k--) {
233 pass.setPipeline(P.backsub);
234 pass.setBindGroup(0, bind, [k * this.stepStride]);
235 pass.dispatchWorkgroups(1);
236 }
237 pass.setPipeline(P.evaluate);
238 pass.setBindGroup(0, bind, [0]);
239 pass.dispatchWorkgroups(ceilDiv(m, EVAL_GROUP));
240 pass.end();
241 enc.copyBufferToBuffer(buffers.out, 0, buffers.read, 0, m * bytes);
242 device.queue.submit([enc.finish()]);
244 await buffers.read.mapAsync(GPUMapMode.READ);
245 const raw = new Float32Array(buffers.read.getMappedRange().slice(0));
246 buffers.read.unmap();
248 if (check) {
249 this.checked = true;
250 const err = await device.popErrorScope();
251 if (err) throw new Error(`mfs-gpu: WebGPU validation: ${err.message}`);
252 }
254 const all = Float64Array.from(raw.subarray(0, m));
255 return {
256 uEval: all.slice(0, prob.nEval),
257 uGrid: wantGrid ? all.slice(prob.nEval) : null,
258 };
259 } finally {
260 for (const b of Object.values(buffers)) b.destroy();
261 }
262 }
263}