1/**
2 * WGSL compute kernels for the Helmholtz layer potentials on quadrilateral
3 * patch surfaces, in single precision. A port of fmm3dbie_gpu/cuda_ops.py
4 * from the CUDA proof of concept, with two changes that single precision
5 * requires:
6 *
7 * 1. For a target on the surface, the displacement y - x_t from the target to
8 * a quadrature node is formed in the coordinates of the patch chart rather
9 * than as a difference of two rounded positions, so that it carries full
10 * relative precision however small it is. Without this the double-layer
11 * kernels, whose numerators (y - x_t).n are O(r^2), lose most of their
12 * digits near the target. There are two charts: the superellipsoid face
13 * chart of the rounded cube (gmode 0), and for any other surface the
14 * Legendre interpolant of the node data (gmode 1), where the displacement
15 * comes from divided differences of the Legendre polynomials.
16 * 2. The weakly singular difference D'_k - D'_kappa is evaluated from a short
17 * power series when k r < 1, where the direct formula subtracts two nearly
18 * equal O(1/r^3) terms.
19 *
20 * Six kernels are handled together, in the order used by fmm3dbie's
21 * helm_rpcomb_imp solver:
22 *
23 * 0: S'_k 1: S_kappa 2: S'_kappa 3: D'_k - D'_kappa 4: S_k 5: D_k
24 *
25 * with kappa = i|k| and k real, G_k = exp(ikr)/(4 pi r), S' the normal
26 * derivative at the target, D the normal derivative at the source.
27 */
29export const NT = 64; // threads per workgroup in near_weights
30export const NORD1_MAX = 11; // norder <= 10
31export const NPOLS_MAX = 121;
32export const DENSE_WG = 128;
34const COMMON = /* wgsl */ `
35const OVER4PI: f32 = 0.07957747154594767;
37fn cmul(a: vec2f, b: vec2f) -> vec2f {
38 return vec2f(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
39}
41// (cos x, sin x) to a few ulp. WGSL only asks sin and cos for an absolute error
42// of 2^-11 on [-pi, pi], and some implementations (SwiftShader among them) use
43// approximations that coarse, so we reduce by pi/2 in three parts (Cody-Waite)
44// and use the minimax polynomials of Cephes on [-pi/4, pi/4].
45fn sincos(x: f32) -> vec2f {
46 let q = round(x * 0.63661977236758134);
47 var r = x - q * 1.5703125;
48 r = r - q * 4.837512969970703125e-4;
49 r = r - q * 7.54978995489188216e-8;
50 let z = r * r;
51 let s = r + r * z * (-1.6666654611e-1 + z * (8.3321608736e-3 + z * (-1.9515295891e-4)));
52 let c = 1.0 - 0.5 * z + z * z * (4.166664568298827e-2 + z * (-1.388731625493765e-3 + z * 2.443315711809948e-5));
53 let iq = i32(q) & 3;
54 switch (iq) {
55 case 0: { return vec2f(c, s); }
56 case 1: { return vec2f(-s, c); }
57 case 2: { return vec2f(-c, -s); }
58 default: { return vec2f(s, -c); }
59 }
60}
62fn ipow(x: f32, n: i32) -> f32 {
63 var r = 1.0;
64 for (var i = 0; i < n; i++) { r *= x; }
65 return r;
66}
68// sum_{j=0}^{m} c^j d^{m-j}, so that c^{m+1} - d^{m+1} = (c - d) * geomsum(c, d, m)
69fn geomsum(c: f32, d: f32, m: i32) -> f32 {
70 var acc = 1.0;
71 var dm = 1.0;
72 for (var j = 1; j <= m; j++) { dm *= d; acc = c * acc + dm; }
73 return acc;
74}
76// The six kernels at one (target, source) pair, times a real weight w.
77// dx = x_t - y. K[3] is only computed when need3 is set.
78fn kernels6(dx: vec3f, nt: vec3f, ns: vec3f, w: f32, k: f32, need3: bool) -> array<vec2f, 6> {
79 let r2 = dot(dx, dx);
80 let r = sqrt(r2);
81 let rinv = 1.0 / r;
82 let r2inv = rinv * rinv;
83 let dnt = dot(dx, nt);
84 let dns = dot(dx, ns);
85 let kr = k * r;
86 let E0 = OVER4PI * rinv * w;
87 let ek = sincos(kr);
88 let ec = exp(-kr);
89 let Ek = E0 * ek;
90 let Ec = E0 * ec;
91 let fk = vec2f(1.0, -kr);
92 let fc = 1.0 + kr;
93 let fkEk = cmul(fk, Ek);
94 var K: array<vec2f, 6>;
95 K[4] = Ek;
96 K[1] = vec2f(Ec, 0.0);
97 K[0] = -dnt * r2inv * fkEk;
98 K[2] = vec2f(-dnt * fc * Ec * r2inv, 0.0);
99 K[5] = dns * r2inv * fkEk;
100 K[3] = vec2f(0.0, 0.0);
101 if (need3) {
102 let nn = dot(ns, nt);
103 var A: vec2f; // ((1 - ikr) e^{ikr} - (1 + kr) e^{-kr})
104 var B: vec2f; // ((3 - 3ikr - k^2r^2) e^{ikr} - (3 + 3kr + k^2r^2) e^{-kr})
105 if (kr < 1.0) {
106 // both are sum_{n>=2} b_n (kr)^n (i^n - (-1)^n), with b_n = (1-n)/n! and (n-1)(n-3)/n!
107 A = vec2f(0.0, 0.0);
108 B = vec2f(0.0, 0.0);
109 var t = 0.5 * kr * kr;
110 for (var n = 2; n <= 14; n++) {
111 var cn: vec2f;
112 switch (n % 4) {
113 case 0: { cn = vec2f(0.0, 0.0); }
114 case 1: { cn = vec2f(1.0, 1.0); }
115 case 2: { cn = vec2f(-2.0, 0.0); }
116 default: { cn = vec2f(1.0, -1.0); }
117 }
118 let nf = f32(n);
119 A += ((1.0 - nf) * t) * cn;
120 B += ((nf - 1.0) * (nf - 3.0) * t) * cn;
121 t *= kr / (nf + 1.0);
122 }
123 } else {
124 A = cmul(fk, ek) - vec2f(fc * ec, 0.0);
125 let gk = vec2f(3.0 - kr * kr, -3.0 * kr);
126 let gc = 3.0 + 3.0 * kr + kr * kr;
127 B = cmul(gk, ek) - vec2f(gc * ec, 0.0);
128 }
129 K[3] = E0 * (nn * r2inv * A - dns * dnt * r2inv * r2inv * B);
130 }
131 return K;
132}
134fn legendre(x: f32, norder: i32, P: ptr<function, array<f32, ${NORD1_MAX}>>) {
135 (*P)[0] = 1.0;
136 if (norder >= 1) { (*P)[1] = x; }
137 for (var n = 1; n < norder; n++) {
138 (*P)[n + 1] = ((2.0 * f32(n) + 1.0) * x * (*P)[n] - f32(n) * (*P)[n - 1]) / f32(n + 1);
139 }
140}
141`;
143/** The patch charts; they need a gpar binding. */
144const CHART = /* wgsl */ `
145struct Geo { x: vec3f, ru: vec3f, rv: vec3f, d: vec3f }
147// xyz, r_u, r_v of the superellipsoid face chart of patch ip at (u, v), and the
148// displacement d = y - x_t. With rel set, x_t is the chart point at (ut, vt) and d
149// is formed in chart coordinates; otherwise d = y - xt for a given xt.
150// gpar: 4 vec4 per patch, [n, s0], [e1, s1], [e2, t0], [0, t1].
151fn chart_rel(ip: i32, u: f32, v: f32, ut: f32, vt: f32, rel: bool, xt: vec3f, ga: f32, gpe: i32) -> Geo {
152 let g0 = gpar[4 * ip]; let g1 = gpar[4 * ip + 1]; let g2 = gpar[4 * ip + 2]; let g3 = gpar[4 * ip + 3];
153 let nf = g0.xyz; let e1 = g1.xyz; let e2 = g2.xyz;
154 let hs = 0.5 * (g1.w - g0.w); let ht = 0.5 * (g3.w - g2.w);
155 let s = g0.w + (u + 1.0) * hs; let t = g2.w + (v + 1.0) * ht;
156 let c = nf + s * e1 + t * e2;
157 let S = ipow(c.x, gpe) + ipow(c.y, gpe) + ipow(c.z, gpe);
158 let N = pow(S, 1.0 / f32(gpe));
159 let gN = vec3f(ipow(c.x / N, gpe - 1), ipow(c.y / N, gpe - 1), ipow(c.z / N, gpe - 1));
160 let dNu = dot(gN, e1); let dNv = dot(gN, e2);
161 var out: Geo;
162 out.x = ga * c / N;
163 out.ru = ga * hs * (e1 / N - c * dNu / (N * N));
164 out.rv = ga * ht * (e2 / N - c * dNv / (N * N));
165 if (rel) {
166 let st = g0.w + (ut + 1.0) * hs; let tt = g2.w + (vt + 1.0) * ht;
167 let ct = nf + st * e1 + tt * e2;
168 let dc = ((u - ut) * hs) * e1 + ((v - vt) * ht) * e2;
169 let St = ipow(ct.x, gpe) + ipow(ct.y, gpe) + ipow(ct.z, gpe);
170 let Nt = pow(St, 1.0 / f32(gpe));
171 // N^p - Nt^p, then N - Nt, both without cancellation
172 let dP = dc.x * geomsum(c.x, ct.x, gpe - 1) + dc.y * geomsum(c.y, ct.y, gpe - 1) + dc.z * geomsum(c.z, ct.z, gpe - 1);
173 let dN = dP / geomsum(N, Nt, gpe - 1);
174 out.d = ga * (dc * Nt - ct * dN) / (N * Nt);
175 } else {
176 out.d = out.x - xt;
177 }
178 return out;
179}
181// P_n(x) - P_n(xt) = (x - xt) * dd[n], from the three-term recurrence:
182// dd[n+1] = ((2n+1) (P_n(xt) + x dd[n]) - n dd[n-1]) / (n+1).
183fn legendre_dd(x: f32, norder: i32, Pt: ptr<function, array<f32, ${NORD1_MAX}>>, dd: ptr<function, array<f32, ${NORD1_MAX}>>) {
184 (*dd)[0] = 0.0;
185 if (norder >= 1) { (*dd)[1] = 1.0; }
186 for (var n = 1; n < norder; n++) {
187 (*dd)[n + 1] = ((2.0 * f32(n) + 1.0) * ((*Pt)[n] + x * (*dd)[n]) - f32(n) * (*dd)[n - 1]) / f32(n + 1);
188 }
189}
191fn gcoef(i: u32) -> f32 {
192 return gpar[i >> 2u][i & 3u];
193}
195// xyz, r_u, r_v of the Legendre interpolant of patch ip at (u, v), and d = y - x_t.
196// gpar holds the coefficients of every patch, 9 * npols floats each (xyz, r_u, r_v).
197// With rel set, d = x(u, v) - x(ut, vt) + e is formed with divided differences, where
198// e = x(ut, vt) - x_t is the (small) gap between the target and the interpolant at
199// its preimage, computed on the CPU; otherwise d = y - xt.
200fn chart_poly(ip: i32, u: f32, v: f32, ut: f32, vt: f32, rel: bool, xt: vec3f, e: vec3f, norder: i32) -> Geo {
201 let n1 = norder + 1;
202 let npols = n1 * n1;
203 let base = u32(ip) * u32(9 * npols);
204 var Pu: array<f32, ${NORD1_MAX}>;
205 var Pv: array<f32, ${NORD1_MAX}>;
206 legendre(u, norder, &Pu);
207 legendre(v, norder, &Pv);
208 var out: Geo;
209 out.x = vec3f(0.0); out.ru = vec3f(0.0); out.rv = vec3f(0.0);
210 for (var b = 0; b < n1; b++) {
211 for (var a = 0; a < n1; a++) {
212 let P = Pu[a] * Pv[b];
213 let n = u32(b * n1 + a);
214 let np = u32(npols);
215 out.x += P * vec3f(gcoef(base + n), gcoef(base + np + n), gcoef(base + 2u * np + n));
216 out.ru += P * vec3f(gcoef(base + 3u * np + n), gcoef(base + 4u * np + n), gcoef(base + 5u * np + n));
217 out.rv += P * vec3f(gcoef(base + 6u * np + n), gcoef(base + 7u * np + n), gcoef(base + 8u * np + n));
218 }
219 }
220 if (rel) {
221 var Put: array<f32, ${NORD1_MAX}>;
222 var Pvt: array<f32, ${NORD1_MAX}>;
223 var du: array<f32, ${NORD1_MAX}>;
224 var dv: array<f32, ${NORD1_MAX}>;
225 legendre(ut, norder, &Put);
226 legendre(vt, norder, &Pvt);
227 legendre_dd(u, norder, &Put, &du);
228 legendre_dd(v, norder, &Pvt, &dv);
229 // P_a(u) P_b(v) - P_a(ut) P_b(vt) = (u - ut) du_a P_b(v) + (v - vt) P_a(ut) dv_b
230 let hu = u - ut;
231 let hv = v - vt;
232 var d = e;
233 for (var b = 0; b < n1; b++) {
234 for (var a = 0; a < n1; a++) {
235 let w = hu * du[a] * Pv[b] + hv * Put[a] * dv[b];
236 let n = u32(b * n1 + a);
237 let np = u32(npols);
238 d += w * vec3f(gcoef(base + n), gcoef(base + np + n), gcoef(base + 2u * np + n));
239 }
240 }
241 out.d = d;
242 } else {
243 out.d = out.x - xt;
244 }
245 return out;
246}
248`;
250/**
251 * Quadtree leaves on the reference square of a near (non-self) patch, refined
252 * while the target is within rfac_q times the radius of a leaf or the leaf is
253 * large compared with the wavelength. Two passes: counts only (fill = 0), then
254 * fill using offsets. Self pairs are skipped; their items come from the CPU.
255 * Items are 2 vec4: [kind, cu, cv, h], [0, 0, 0, 0] with kind 0 for a leaf.
256 */
257export const BUILD_LEAVES = /* wgsl */ `
258struct Params {
259 n: u32,
260 pair_offset: u32,
261 fill: u32,
262 nlmax: i32,
263 rfac_q: f32,
264 zk: f32,
265 kR_max: f32,
266 ga: f32,
267 gpe: i32,
268 gmode: i32,
269 norder: i32,
270}
271@group(0) @binding(0) var<uniform> U: Params;
272@group(0) @binding(1) var<storage, read> pairs_i: array<vec2i>;
273@group(0) @binding(2) var<storage, read> pairs_f: array<vec4f>;
274@group(0) @binding(3) var<storage, read> targ: array<vec4f>;
275@group(0) @binding(4) var<storage, read> gpar: array<vec4f>;
276@group(0) @binding(5) var<storage, read_write> counts: array<u32>;
277@group(0) @binding(6) var<storage, read> offsets: array<u32>;
278@group(0) @binding(7) var<storage, read_write> items: array<vec4f>;
279${COMMON}
280${CHART}
282@compute @workgroup_size(64)
283fn build_leaves(@builtin(global_invocation_id) gid: vec3u) {
284 let pl = gid.x;
285 if (pl >= U.n) { return; }
286 let p = pl + U.pair_offset;
287 let ip = pairs_i[p].y;
288 let it = pairs_i[p].x;
289 let kind = i32(pairs_f[2u * p].z);
290 if (kind == 2) {
291 if (U.fill == 0u) { counts[pl] = 0u; }
292 return;
293 }
294 let tx = targ[2 * it].xyz;
295 var su: array<f32, 64>;
296 var sv: array<f32, 64>;
297 var sh: array<f32, 64>;
298 var sd: array<i32, 64>;
299 var sp = 1;
300 su[0] = 0.0; sv[0] = 0.0; sh[0] = 1.0; sd[0] = 0;
301 var n = 0u;
302 let rq2 = U.rfac_q * U.rfac_q;
303 let kmax2 = U.kR_max * U.kR_max;
304 let k2 = U.zk * U.zk;
305 while (sp > 0) {
306 sp--;
307 let cu = su[sp]; let cv = sv[sp]; let h = sh[sp]; let d = sd[sp];
308 var refine = false;
309 if (d < U.nlmax) {
310 // centroid and radius from a 3x3 sample grid
311 var xc: array<vec3f, 9>;
312 var cen = vec3f(0.0);
313 for (var q = 0; q < 9; q++) {
314 let uq = cu + h * f32(q % 3 - 1);
315 let vq = cv + h * f32(q / 3 - 1);
316 var g: Geo;
317 if (U.gmode == 1) { g = chart_poly(ip, uq, vq, 0.0, 0.0, false, vec3f(0.0), vec3f(0.0), U.norder); }
318 else { g = chart_rel(ip, uq, vq, 0.0, 0.0, false, vec3f(0.0), U.ga, U.gpe); }
319 xc[q] = g.x;
320 cen += g.x / 9.0;
321 }
322 var R2 = 0.0;
323 for (var q = 0; q < 9; q++) { let e = xc[q] - cen; R2 = max(R2, dot(e, e)); }
324 let e = tx - cen;
325 let d2 = dot(e, e);
326 refine = (d2 < rq2 * R2) || (k2 * R2 > kmax2);
327 }
328 if (refine && sp <= 60) {
329 let hh = 0.5 * h;
330 su[sp] = cu - hh; sv[sp] = cv - hh; sh[sp] = hh; sd[sp] = d + 1; sp++;
331 su[sp] = cu + hh; sv[sp] = cv - hh; sh[sp] = hh; sd[sp] = d + 1; sp++;
332 su[sp] = cu + hh; sv[sp] = cv + hh; sh[sp] = hh; sd[sp] = d + 1; sp++;
333 su[sp] = cu - hh; sv[sp] = cv + hh; sh[sp] = hh; sd[sp] = d + 1; sp++;
334 } else {
335 if (U.fill != 0u) {
336 let o = 2u * (offsets[pl] + n);
337 items[o] = vec4f(0.0, cu, cv, h);
338 items[o + 1u] = vec4f(0.0);
339 }
340 n++;
341 }
342 }
343 if (U.fill == 0u) { counts[pl] = n; }
344}
345`;
347/**
348 * Near-field weights, one workgroup per (target, patch) pair. The nodes of the
349 * pair's items (self-patch polar segments or quadtree leaves) are evaluated
350 * ${NT} at a time; each thread then accumulates the Legendre moments
351 * I_n = sum K P_n |J| w for its basis indices n, and finally the moments are
352 * turned into node weights w_j = sum_n U[n, j] I_n.
353 *
354 * Self items are 2 vec4: [1, edge, xi0, xi1], [h, xc, heff, 0]: on edge e the
355 * polar rule uses rho in [0, 1] along the ray and the sinh substitution
356 * x = xc + heff sinh(xi), xi in [xi0, xi1], along the edge.
357 * tables: umat (npols*npols) followed by the GL rules gl_x, gl_w (nq), glr_x,
358 * glr_w (nqr, on [0, 1]) and glx_x, glx_w (nqx).
359 */
360export const NEAR_WEIGHTS = /* wgsl */ `
361struct Params {
362 n: u32,
363 pair_offset: u32,
364 norder: i32,
365 npols: i32,
366 nq: i32,
367 nqr: i32,
368 nqx: i32,
369 zk: f32,
370 ga: f32,
371 gpe: i32,
372 gmode: i32,
373}
374@group(0) @binding(0) var<uniform> U: Params;
375@group(0) @binding(1) var<storage, read> pairs_i: array<vec2i>;
376@group(0) @binding(2) var<storage, read> pairs_f: array<vec4f>;
377@group(0) @binding(3) var<storage, read> targ: array<vec4f>;
378@group(0) @binding(4) var<storage, read> gpar: array<vec4f>;
379@group(0) @binding(5) var<storage, read> item_ptr: array<u32>;
380@group(0) @binding(6) var<storage, read> items: array<vec4f>;
381@group(0) @binding(7) var<storage, read> tables: array<f32>;
382@group(0) @binding(8) var<storage, read_write> W: array<vec2f>;
383${COMMON}
384${CHART}
386const NT = ${NT};
387var<workgroup> sK: array<vec2f, ${6 * NT}>;
388var<workgroup> sLu: array<f32, ${NORD1_MAX * NT}>;
389var<workgroup> sLv: array<f32, ${NORD1_MAX * NT}>;
390var<workgroup> sM: array<vec2f, ${6 * NPOLS_MAX}>;
391var<workgroup> sNtot: i32;
393@compute @workgroup_size(${NT})
394fn near_weights(@builtin(workgroup_id) wid: vec3u, @builtin(local_invocation_id) lid: vec3u) {
395 let pl = wid.x;
396 let p = pl + U.pair_offset;
397 let t = i32(lid.x);
398 let it = pairs_i[p].x;
399 let ip = pairs_i[p].y;
400 let ut = pairs_f[2u * p].x;
401 let vt = pairs_f[2u * p].y;
402 let kind = i32(pairs_f[2u * p].z);
403 let egap = pairs_f[2u * p + 1u].xyz;
404 let is_self = kind == 2;
405 let onsurf = kind >= 1;
406 let xt = targ[2 * it].xyz;
407 let nt = targ[2 * it + 1].xyz;
408 let n1 = U.norder + 1;
409 let npols = U.npols;
410 let i0 = item_ptr[pl];
411 let nodes_per_item = select(U.nq * U.nq, U.nqr * U.nqx, is_self);
412 if (t == 0) { sNtot = i32(item_ptr[pl + 1u] - i0) * nodes_per_item; }
413 let ntot = workgroupUniformLoad(&sNtot);
414 let oq = npols * npols; // gl tables start after umat
415 let or = oq + 2 * U.nq; // glr_x
416 let ox = or + 2 * U.nqr; // glx_x
418 var Ir: array<vec2f, 12>;
419 for (var a = 0; a < 12; a++) { Ir[a] = vec2f(0.0, 0.0); }
421 for (var base = 0; base < ntot; base += NT) {
422 let g = base + t;
423 var K: array<vec2f, 6>;
424 var Lu: array<f32, ${NORD1_MAX}>;
425 var Lv: array<f32, ${NORD1_MAX}>;
426 let live = g < ntot;
427 if (live) {
428 let item = g / nodes_per_item;
429 let loc = g - item * nodes_per_item;
430 let it0 = items[2u * (i0 + u32(item))];
431 let it1 = items[2u * (i0 + u32(item)) + 1u];
432 var u: f32; var v: f32; var wq: f32;
433 if (is_self) {
434 let e = i32(it0.y);
435 let ir = loc / U.nqx;
436 let ix = loc - ir * U.nqx;
437 let h = it1.x; let xc = it1.y; let heff = it1.z;
438 let rho = tables[or + ir];
439 let wr = tables[or + U.nqr + ir];
440 let xi = it0.z + 0.5 * (it0.w - it0.z) * (tables[ox + ix] + 1.0);
441 let wxi = 0.5 * (it0.w - it0.z) * tables[ox + U.nqx + ix];
442 let x = xc + heff * sinh(xi);
443 let dx = heff * cosh(xi) * wxi;
444 if (e < 2) { u = ut + rho * h; v = vt + rho * x; } else { u = ut + rho * x; v = vt + rho * h; }
445 wq = abs(h) * rho * wr * dx;
446 } else {
447 let a = loc / U.nq;
448 let b = loc - a * U.nq;
449 u = it0.y + it0.w * tables[oq + a];
450 v = it0.z + it0.w * tables[oq + b];
451 wq = it0.w * it0.w * tables[oq + U.nq + a] * tables[oq + U.nq + b];
452 }
453 var G: Geo;
454 if (U.gmode == 1) { G = chart_poly(ip, u, v, ut, vt, onsurf, xt, egap, U.norder); }
455 else { G = chart_rel(ip, u, v, ut, vt, onsurf, xt, U.ga, U.gpe); }
456 var ns = cross(G.ru, G.rv);
457 let jac = length(ns);
458 ns = ns / jac;
459 K = kernels6(-G.d, nt, ns, wq * jac, U.zk, true);
460 legendre(u, U.norder, &Lu);
461 legendre(v, U.norder, &Lv);
462 }
463 for (var m = 0; m < 6; m++) { sK[m * NT + t] = select(vec2f(0.0, 0.0), K[m], live); }
464 for (var a = 0; a < n1; a++) {
465 sLu[a * NT + t] = select(0.0, Lu[a], live);
466 sLv[a * NT + t] = select(0.0, Lv[a], live);
467 }
468 workgroupBarrier();
469 for (var acc = 0; acc < 2; acc++) {
470 let n = t + acc * NT;
471 if (n < npols) {
472 let a = n % n1;
473 let b = n / n1;
474 for (var i = 0; i < NT; i++) {
475 let P = sLu[a * NT + i] * sLv[b * NT + i];
476 for (var m = 0; m < 6; m++) { Ir[acc * 6 + m] += P * sK[m * NT + i]; }
477 }
478 }
479 }
480 workgroupBarrier();
481 }
483 // moments -> node weights, w_j = sum_n U[n, j] I_n
484 for (var acc = 0; acc < 2; acc++) {
485 let n = t + acc * NT;
486 if (n < npols) { for (var m = 0; m < 6; m++) { sM[m * ${NPOLS_MAX} + n] = Ir[acc * 6 + m]; } }
487 }
488 workgroupBarrier();
489 for (var j = t; j < npols; j += NT) {
490 for (var m = 0; m < 6; m++) {
491 var w = vec2f(0.0, 0.0);
492 for (var n = 0; n < npols; n++) { w += tables[n * npols + j] * sM[m * ${NPOLS_MAX} + n]; }
493 W[(u32(p) * 6u + u32(m)) * u32(npols) + u32(j)] = w;
494 }
495 }
496}
497`;
499/**
500 * Far-field pass over the patches j0 <= j < j1: out[m][i] += sum over patches j
501 * not near target i of sum_l K_m(x_i, y_l) w_l q_l over the oversampled nodes
502 * of patch j, for the kernels selected in selmask (4 bits per kernel: 0 skip,
503 * 1 density qa, 2 density qb). A patch is near target i when it appears in
504 * the target's near list (row_ptr, pair_p, ascending), so the split between
505 * this sum and the near weights is exact.
506 * srcgeo: 2 vec4 per node, [x, w], [n, 0]; srcq: vec4 per node (qa, qb).
507 */
508export const DENSE_APPLY = /* wgsl */ `
509struct Params {
510 ntarg: u32,
511 j0: u32,
512 j1: u32,
513 npols_o: u32,
514 zk: f32,
515 selmask: u32,
516}
517@group(0) @binding(0) var<uniform> U: Params;
518@group(0) @binding(1) var<storage, read> targ: array<vec4f>;
519@group(0) @binding(2) var<storage, read> srcgeo: array<vec4f>;
520@group(0) @binding(3) var<storage, read> srcq: array<vec4f>;
521@group(0) @binding(4) var<storage, read> row_ptr: array<u32>;
522@group(0) @binding(5) var<storage, read> pair_p: array<u32>;
523@group(0) @binding(6) var<storage, read_write> out: array<vec2f>;
524${COMMON}
526const WG = ${DENSE_WG}u;
527var<workgroup> tile: array<vec4f, ${3 * DENSE_WG}>;
529@compute @workgroup_size(${DENSE_WG})
530fn dense_apply(@builtin(global_invocation_id) gid: vec3u, @builtin(local_invocation_id) lid: vec3u) {
531 let i = gid.x;
532 let valid = i < U.ntarg;
533 var xt = vec3f(0.0);
534 var nt = vec3f(0.0);
535 var cur = 0u;
536 var end = 0u;
537 if (valid) {
538 xt = targ[2u * i].xyz;
539 nt = targ[2u * i + 1u].xyz;
540 cur = row_ptr[i];
541 end = row_ptr[i + 1u];
542 while (cur < end && pair_p[cur] < U.j0) { cur++; }
543 }
544 var sel: array<u32, 6>;
545 for (var m = 0u; m < 6u; m++) { sel[m] = (U.selmask >> (4u * m)) & 3u; }
546 let need3 = sel[3] != 0u;
547 var acc: array<vec2f, 6>;
548 for (var m = 0; m < 6; m++) { acc[m] = vec2f(0.0, 0.0); }
549 for (var j = U.j0; j < U.j1; j++) {
550 var far = valid;
551 if (valid && cur < end && pair_p[cur] == j) { far = false; cur++; }
552 for (var base = 0u; base < U.npols_o; base += WG) {
553 let l = base + lid.x;
554 if (l < U.npols_o) {
555 let s = j * U.npols_o + l;
556 tile[3u * lid.x] = srcgeo[2u * s];
557 tile[3u * lid.x + 1u] = srcgeo[2u * s + 1u];
558 tile[3u * lid.x + 2u] = srcq[s];
559 }
560 workgroupBarrier();
561 if (far) {
562 let nl = min(WG, U.npols_o - base);
563 var part: array<vec2f, 6>;
564 for (var m = 0; m < 6; m++) { part[m] = vec2f(0.0, 0.0); }
565 for (var s = 0u; s < nl; s++) {
566 let g0 = tile[3u * s];
567 let g1 = tile[3u * s + 1u];
568 let q = tile[3u * s + 2u];
569 let K = kernels6(xt - g0.xyz, nt, g1.xyz, g0.w, U.zk, need3);
570 for (var m = 0u; m < 6u; m++) {
571 if (sel[m] == 1u) { part[m] += cmul(K[m], q.xy); }
572 else if (sel[m] == 2u) { part[m] += cmul(K[m], q.zw); }
573 }
574 }
575 for (var m = 0; m < 6; m++) { acc[m] += part[m]; }
576 }
577 workgroupBarrier();
578 }
579 }
580 if (valid) {
581 for (var m = 0u; m < 6u; m++) {
582 if (sel[m] != 0u) { out[m * U.ntarg + i] += acc[m]; }
583 }
584 }
585}
586`;
588/**
589 * Near contributions: out[m][i] += sum over the near pairs p of target i of
590 * sum_l W[p][m][l] q[d][patch nodes l], for up to four requests (m, d) packed
591 * as m | d << 8 in req.
592 */
593export const SPARSE_APPLY = /* wgsl */ `
594struct Params {
595 ntarg: u32,
596 npols: u32,
597 npts: u32,
598 nreq: u32,
599 req: vec4u,
600}
601@group(0) @binding(0) var<uniform> U: Params;
602@group(0) @binding(1) var<storage, read> row_ptr: array<u32>;
603@group(0) @binding(2) var<storage, read> pair_p: array<u32>;
604@group(0) @binding(3) var<storage, read> W: array<vec2f>;
605@group(0) @binding(4) var<storage, read> q: array<vec2f>;
606@group(0) @binding(5) var<storage, read_write> out: array<vec2f>;
608fn cmul(a: vec2f, b: vec2f) -> vec2f {
609 return vec2f(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
610}
612@compute @workgroup_size(128)
613fn sparse_apply(@builtin(global_invocation_id) gid: vec3u) {
614 let i = gid.x;
615 if (i >= U.ntarg) { return; }
616 var acc: array<vec2f, 4>;
617 for (var r = 0u; r < 4u; r++) { acc[r] = vec2f(0.0, 0.0); }
618 for (var p = row_ptr[i]; p < row_ptr[i + 1u]; p++) {
619 let j = pair_p[p];
620 for (var r = 0u; r < U.nreq; r++) {
621 let m = U.req[r] & 0xffu;
622 let d = U.req[r] >> 8u;
623 let wbase = (p * 6u + m) * U.npols;
624 let qbase = d * U.npts + j * U.npols;
625 var s = vec2f(0.0, 0.0);
626 for (var l = 0u; l < U.npols; l++) { s += cmul(W[wbase + l], q[qbase + l]); }
627 acc[r] += s;
628 }
629 }
630 for (var r = 0u; r < U.nreq; r++) {
631 let m = U.req[r] & 0xffu;
632 out[m * U.ntarg + i] += acc[r];
633 }
634}
635`;