acdea26Add mfs-gpu: the MFS on WebGPU, and a TypeScript form of the solver interfaceJeremy Magland 1// The WGSL for the WebGPU method of fundamental solutions.
2//
3// Everything here is f32, because that is what WebGPU has: WGSL's floating
4// point types are f32 and f16, there is no f64, and no extension in the
5// standard adds one. That single fact is the most interesting thing about
6// this solver, because the MFS is conditioning-limited: its accuracy stops
7// improving where rounding in the collocation solve overtakes the
8// approximation error, and in f32 that happens six to nine orders short of
9// where it happens in double. See the solver's manifest entry for the
10// measured numbers.
11//
12// Note that software emulation of double precision (carrying each value as
13// an unevaluated sum of two f32, the "double-single" trick) is not a way
14// out. It depends on error-free transformations such as
15// s = a + b; err = b - (s - a), which are only error-free if the compiler
16// evaluates them exactly as written; WGSL permits an implementation to
17// compute with greater precision or to reassociate, and the Mesa driver
18// this was tried on does exactly that, returning f32-accurate results for
19// values loaded from a buffer. So f32 is what a WebGPU solver gets.
21/** Workgroup sizes: the two-dimensional kernels (assembly and the trailing
22 * update), the one-dimensional ones, and the evaluation. */
23export const TILE = 16;
24export const LANES = 256;
25export const EVAL_GROUP = 64;
27/**
28 * The whole solver as one shader module.
29 *
30 * The matrix is the augmented system [A | b], n rows of ld = n + 1 entries,
31 * row-major, so that the elimination applies to the right-hand side as it
32 * goes and no separate forward substitution is needed. The multipliers are
33 * never stored: nothing reads them again.
34 *
35 * One read-only buffer holds every input, because the default WebGPU limit
36 * is eight storage buffers per stage. Its layout, in f32 elements, is
37 * derived from n rather than passed: collocation points (x, y interleaved)
38 * at 0, charges at 2n, the Dirichlet data at 4n, the targets (interleaved)
39 * at 5n.
40 */
41export function mfsShader(): string {
42 return `// generated by src/solvers/mfs-gpu/wgsl.ts
44struct Dims {
45 n: u32, // charges = collocation points
46 ld: u32, // row stride of the augmented matrix, n + 1
47 m: u32, // evaluation targets
48 pad: u32,
49};
51struct Step { k: u32, pad0: u32, pad1: u32, pad2: u32 };
53@group(0) @binding(0) var<uniform> dims: Dims;
54@group(0) @binding(1) var<uniform> step: Step;
55@group(0) @binding(2) var<storage, read> data: array<f32>;
56@group(0) @binding(3) var<storage, read_write> mat: array<f32>;
57@group(0) @binding(4) var<storage, read_write> lcol: array<f32>;
58@group(0) @binding(5) var<storage, read_write> sol: array<f32>;
59@group(0) @binding(6) var<storage, read_write> out: array<f32>;
61fn pt_x(i: u32) -> f32 { return data[2u * i]; }
62fn pt_y(i: u32) -> f32 { return data[2u * i + 1u]; }
63fn ch_x(j: u32) -> f32 { return data[2u * dims.n + 2u * j]; }
64fn ch_y(j: u32) -> f32 { return data[2u * dims.n + 2u * j + 1u]; }
65fn bdata(i: u32) -> f32 { return data[4u * dims.n + i]; }
66fn tg_x(t: u32) -> f32 { return data[5u * dims.n + 2u * t]; }
67fn tg_y(t: u32) -> f32 { return data[5u * dims.n + 2u * t + 1u]; }
69/** The MFS kernel log|p - q|, written as half the log of the squared
70 * distance so that no square root is taken. */
71fn logdist(ax: f32, ay: f32, bx: f32, by: f32) -> f32 {
72 let dx = ax - bx;
73 let dy = ay - by;
74 return 0.5 * log(dx * dx + dy * dy);
75}
77@compute @workgroup_size(${TILE}, ${TILE})
78fn assemble(@builtin(global_invocation_id) gid: vec3<u32>) {
79 let i = gid.y;
80 let j = gid.x;
81 if (i >= dims.n || j >= dims.ld) { return; }
82 if (j == dims.n) {
83 mat[i * dims.ld + j] = bdata(i);
84 return;
85 }
86 mat[i * dims.ld + j] = logdist(pt_x(i), pt_y(i), ch_x(j), ch_y(j));
87}
89var<workgroup> best_val: array<f32, ${LANES}>;
90var<workgroup> best_idx: array<u32, ${LANES}>;
92/** Partial pivoting for step k: find the largest |A(i,k)| over i >= k and
93 * swap that row with row k. One workgroup, so the search reduces in shared
94 * memory and the swap needs no second dispatch. */
95@compute @workgroup_size(${LANES})
96fn pivot(@builtin(local_invocation_id) lid: vec3<u32>) {
97 let k = step.k;
98 let n = dims.n;
99 let ld = dims.ld;
100 var bv = -1.0;
101 var bi = k;
102 for (var i = k + lid.x; i < n; i = i + ${LANES}u) {
103 let v = abs(mat[i * ld + k]);
104 if (v > bv) { bv = v; bi = i; }
105 }
106 best_val[lid.x] = bv;
107 best_idx[lid.x] = bi;
108 workgroupBarrier();
109 for (var s = ${LANES / 2}u; s > 0u; s = s >> 1u) {
110 if (lid.x < s && best_val[lid.x + s] > best_val[lid.x]) {
111 best_val[lid.x] = best_val[lid.x + s];
112 best_idx[lid.x] = best_idx[lid.x + s];
113 }
114 workgroupBarrier();
115 }
116 let p = best_idx[0];
117 if (p != k) {
118 for (var j = k + lid.x; j < ld; j = j + ${LANES}u) {
119 let a = mat[k * ld + j];
120 mat[k * ld + j] = mat[p * ld + j];
121 mat[p * ld + j] = a;
122 }
123 }
124}
126/** The multipliers of step k, computed once rather than once per element of
127 * the trailing update. */
128@compute @workgroup_size(${LANES})
129fn multipliers(@builtin(global_invocation_id) gid: vec3<u32>) {
130 let k = step.k;
131 let i = k + 1u + gid.x;
132 if (i >= dims.n) { return; }
133 lcol[i] = mat[i * dims.ld + k] / mat[k * dims.ld + k];
134}
136/** The rank-one update of the trailing submatrix, columns k+1 .. ld-1, so
137 * the augmented right-hand column is eliminated along with the rest. */
138@compute @workgroup_size(${TILE}, ${TILE})
139fn update(@builtin(global_invocation_id) gid: vec3<u32>) {
140 let k = step.k;
141 let i = k + 1u + gid.y;
142 let j = k + 1u + gid.x;
143 if (i >= dims.n || j >= dims.ld) { return; }
144 let ld = dims.ld;
145 mat[i * ld + j] = mat[i * ld + j] - lcol[i] * mat[k * ld + j];
146}
148/** Back substitution, one step per dispatch: x(k) from row k of U and the
149 * running right-hand column, then that column updated above row k. Every
150 * lane recomputes x(k) so that no barrier is needed. */
151@compute @workgroup_size(${LANES})
152fn backsub(@builtin(local_invocation_id) lid: vec3<u32>) {
153 let k = step.k;
154 let ld = dims.ld;
155 let xk = mat[k * ld + dims.n] / mat[k * ld + k];
156 if (lid.x == 0u) { sol[k] = xk; }
157 for (var i = lid.x; i < k; i = i + ${LANES}u) {
158 mat[i * ld + dims.n] = mat[i * ld + dims.n] - mat[i * ld + k] * xk;
159 }
160}
162/** The potential at the targets: one thread per target, summing the
163 * charges. */
164@compute @workgroup_size(${EVAL_GROUP})
165fn evaluate(@builtin(global_invocation_id) gid: vec3<u32>) {
166 let t = gid.x;
167 if (t >= dims.m) { return; }
168 let ax = tg_x(t);
169 let ay = tg_y(t);
170 var acc = 0.0;
171 for (var j = 0u; j < dims.n; j = j + 1u) {
172 acc = acc + sol[j] * logdist(ax, ay, ch_x(j), ch_y(j));
173 }
174 out[t] = acc;
175}
176`;
177}