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