/ proof-of-concept / fmm3dbie-webgpu-poc
Sign in
proof-of-concept / fmm3dbie-webgpu-poc
fmm3dbie-webgpu-poc / src / gpu / operator.ts
635 lines · 25.5 KBBlameHistoryRaw
1/**
2 * The discretized operator of fmm3dbie's right-preconditioned combined field
3 * impedance equation (helm_rpcomb_imp), on WebGPU. A port of
4 * fmm3dbie_gpu/solver.py from the CUDA proof of concept.
5 *
6 * u = S_k[sigma] + i alpha D_k S_kappa[sigma], kappa = i|k|,
7 *
8 * with the boundary condition du/dn + i k lambda u = f. The discretized
9 * equation is
10 *
11 * z sigma + S'_k sigma + i alpha S'_kappa phi2 + i alpha (D'_k - D'_kappa) phi1
12 * + i k lambda [S_k sigma + i alpha D_k phi1 + i alpha w phi1] = f,
13 *
14 * phi1 = S_kappa sigma, phi2 = S'_kappa sigma,
15 * z = (-1)^ifinout / 2 - i alpha / 4, w = 1/2 (exterior), -1/2 (interior).
16 *
17 * Every layer potential is the sum of a far part (smooth rule on the
18 * oversampled nodes of the patches that are not near the target, the dense
19 * kernel) and a near part (precomputed weights for the near patches, the
20 * sparse kernel). The near weights are computed once, in chunks of pairs, and
21 * kept on the GPU. GMRES runs on the CPU in double precision and calls
22 * applyK for each Krylov vector.
23 */
24import { leggauss } from '../gauss.ts';
25import { oversampleMatrix, type Surface } from '../surface.ts';
26import type { CVec } from '../problem.ts';
27import { gmres, type GmresResult } from '../gmres.ts';
28import { Gpu, Pacer, UniformWriter } from './device.ts';
29import { BUILD_LEAVES, DENSE_APPLY, DENSE_WG, NEAR_WEIGHTS, NT, SPARSE_APPLY } from './shaders.ts';
31export interface OperatorOptions {
32 /** near radius factor: a patch is near a target within rfac * (patch radius) of its centroid */
33 rfac?: number;
34 rfacQ?: number;
35 nq?: number;
36 nqr?: number;
37 nqx?: number;
38 subLen?: number;
39 oscLen?: number;
40 kRMax?: number;
41 nlmax?: number;
42 nover?: number;
43 /** seconds of GPU work per dispatch */
44 pace?: number;
45 progress?: (stage: string, fraction: number) => void;
48export interface NearPairs {
49 pairT: Int32Array;
50 pairP: Int32Array;
51 rowPtr: Uint32Array;
52 /** per pair: (u_t, v_t, kind, 0); kind 0 off-surface target, 1 on-surface, 2 self patch */
53 pairF: Float32Array;
54 npairs: number;
57/** A set of targets with everything needed to apply the operator at them. */
58interface TargetSet {
59 n: number;
60 targ: GPUBuffer; // 2 vec4 per target: [x, 0], [n, 0]
61 pairs: NearPairs;
62 pairsI: GPUBuffer;
63 pairsF: GPUBuffer;
64 rowPtr: GPUBuffer;
65 pairP: GPUBuffer;
66 W: GPUBuffer;
67 nleaves: number;
70export interface SetupInfo {
71 npts: number;
72 noversampled: number;
73 npairs: number;
74 nleaves: number;
75 weightsBytes: number;
76 setupSeconds: number;
79const now = () => performance.now() / 1000;
81export class Operator {
82 readonly k: number;
83 readonly alpha: number;
84 /** wall impedance; the near weights do not depend on it, so it may be changed between solves */
85 lam: number;
86 readonly ifinout: number;
87 readonly norder: number;
88 readonly npols: number;
89 readonly npts: number;
90 readonly nover: number;
91 readonly npolsO: number;
92 readonly rfac: number;
93 readonly quad: { nq: number; nqr: number; nqx: number; subLen: number; oscLen: number };
94 readonly leaf: { rfacQ: number; nlmax: number; kRMax: number };
95 readonly progress: (stage: string, fraction: number) => void;
96 readonly pace: number;
97 info!: SetupInfo;
99 private gpu: Gpu;
100 private surf: Surface;
101 private r2: Float64Array;
102 private overMat: Float64Array;
103 private gpar!: GPUBuffer;
104 private tables!: GPUBuffer;
105 private srcgeo!: GPUBuffer;
106 private srcq!: GPUBuffer;
107 private qbuf!: GPUBuffer;
108 private out!: GPUBuffer;
109 private wall!: TargetSet;
110 private nself = 0;
111 private selfItems: Float32Array[] = []; // per surface node, its polar segments (8 floats each)
112 private uniformLeaves: GPUBuffer;
113 private uniformWeights: GPUBuffer;
114 private uniformDense: GPUBuffer;
115 private uniformSparse: GPUBuffer;
116 private destroyed = false;
118 constructor(gpu: Gpu, surf: Surface, k: number, alpha: number, lam: number, ifinout = 0, opts: OperatorOptions = {}) {
119 this.gpu = gpu;
120 this.surf = surf;
121 this.k = k;
122 this.alpha = alpha;
123 this.lam = lam;
124 this.ifinout = ifinout;
125 this.norder = surf.norder;
126 this.npols = surf.npols;
127 this.npts = surf.npts;
128 if (this.norder > 10) throw new Error('order at most 10');
129 this.rfac = opts.rfac ?? (this.norder <= 6 ? 2.0 : 1.25); // fmm3dbie's get_rfacs
130 this.nover = opts.nover ?? this.norder + 2;
131 this.npolsO = (this.nover + 1) ** 2;
132 let radMax = 0;
133 for (let j = 0; j < surf.npatches; j++) radMax = Math.max(radMax, surf.rads[j]);
134 // the radial rule must resolve the basis and exp(ikr) along a ray
135 const nqr = opts.nqr ?? Math.floor(Math.min(Math.max(12, this.norder + 2 * Math.abs(k) * radMax + 6), 32));
136 this.quad = { nq: opts.nq ?? 8, nqr, nqx: opts.nqx ?? 8, subLen: opts.subLen ?? 1.0, oscLen: opts.oscLen ?? 2.0 };
137 this.leaf = { rfacQ: opts.rfacQ ?? 2.0, nlmax: opts.nlmax ?? 14, kRMax: opts.kRMax ?? 3.0 };
138 this.progress = opts.progress ?? (() => {});
139 this.pace = opts.pace ?? 0.15;
140 this.r2 = new Float64Array(surf.npatches);
141 for (let j = 0; j < surf.npatches; j++) this.r2[j] = (this.rfac * surf.rads[j]) ** 2;
142 this.overMat = oversampleMatrix(this.norder, this.nover);
143 this.uniformLeaves = gpu.uniform(64);
144 this.uniformWeights = gpu.uniform(64);
145 this.uniformDense = gpu.uniform(64);
146 this.uniformSparse = gpu.uniform(64);
147 }
149 /** Upload the geometry and compute the near weights for the surface nodes. */
150 async setup(): Promise<SetupInfo> {
151 const t0 = now();
152 const { gpu, surf, npols, npts } = this;
153 // patch charts: 4 vec4 per patch, [n, s0], [e1, s1], [e2, t0], [0, t1]
154 const gp = new Float32Array(16 * surf.npatches);
155 for (let j = 0; j < surf.npatches; j++) {
156 const par = surf.exact.params.subarray(13 * j, 13 * j + 13);
157 gp.set([par[0], par[1], par[2], par[9], par[3], par[4], par[5], par[10], par[6], par[7], par[8], par[11], 0, 0, 0, par[12]], 16 * j);
158 }
159 this.gpar = gpu.storage(gp);
160 // umat and the quadrature rules in one table
161 const { nq, nqr, nqx } = this.quad;
162 const gl = leggauss(nq);
163 const glr = leggauss(nqr);
164 const glx = leggauss(nqx);
165 const tab = new Float32Array(npols * npols + 2 * nq + 2 * nqr + 2 * nqx);
166 tab.set(surf.umat, 0);
167 let o = npols * npols;
168 tab.set(gl.x, o);
169 tab.set(gl.w, o + nq);
170 o += 2 * nq;
171 for (let i = 0; i < nqr; i++) {
172 tab[o + i] = 0.5 * (glr.x[i] + 1);
173 tab[o + nqr + i] = 0.5 * glr.w[i];
174 }
175 o += 2 * nqr;
176 tab.set(glx.x, o);
177 tab.set(glx.w, o + nqx);
178 this.tables = gpu.storage(tab);
179 // oversampled sources for the far field
180 const ov = surf.oversampled(this.nover);
181 const N = ov.w.length;
182 const sg = new Float32Array(8 * N);
183 for (let i = 0; i < N; i++) {
184 sg[8 * i] = ov.x[i];
185 sg[8 * i + 1] = ov.x[N + i];
186 sg[8 * i + 2] = ov.x[2 * N + i];
187 sg[8 * i + 3] = ov.w[i];
188 sg[8 * i + 4] = ov.n[i];
189 sg[8 * i + 5] = ov.n[N + i];
190 sg[8 * i + 6] = ov.n[2 * N + i];
191 }
192 this.srcgeo = gpu.storage(sg);
193 this.srcq = gpu.storage(16 * N);
194 this.qbuf = gpu.storage(2 * 8 * npts);
195 this.out = gpu.storage(6 * 8 * npts);
197 // self-patch polar segments for every node, from the first fundamental form at the node
198 this.selfItems = new Array(npts);
199 for (let i = 0; i < npts; i++) this.selfItems[i] = this.selfSegments(i);
201 const xs = surf.srcvals.subarray(0, 3 * npts);
202 const nrm = surf.srcvals.subarray(9 * npts, 12 * npts);
203 this.wall = await this.makeTargets(xs, nrm, true, 'near quadrature');
204 await gpu.sync();
205 const setupSeconds = now() - t0;
206 this.info = {
207 npts,
208 noversampled: N,
209 npairs: this.wall.pairs.npairs,
210 nleaves: this.wall.nleaves,
211 weightsBytes: this.wall.W.size,
212 setupSeconds,
213 };
214 return this.info;
215 }
217 destroy() {
218 this.destroyed = true;
219 for (const b of [this.gpar, this.tables, this.srcgeo, this.srcq, this.qbuf, this.out, this.uniformLeaves, this.uniformWeights, this.uniformDense, this.uniformSparse]) b?.destroy();
220 if (this.wall) this.freeTargets(this.wall);
221 }
223 /** Bytes of near weights for npairs pairs. */
224 weightsBytes(npairs: number): number {
225 return npairs * 6 * this.npols * 8;
226 }
228 /** Estimated bytes of near weights for the surface nodes, before setup. */
229 estimateWeightsBytes(): number {
230 const { npts } = this;
231 const xs = this.surf.srcvals.subarray(0, 3 * npts);
232 let cnt = 0;
233 const { npatches } = this.surf;
234 const cen = this.surf.centroids;
235 // sample every 16th node
236 for (let i = 0; i < npts; i += 16) {
237 for (let j = 0; j < npatches; j++) {
238 const d2 = (xs[i] - cen[j]) ** 2 + (xs[npts + i] - cen[npatches + j]) ** 2 + (xs[2 * npts + i] - cen[2 * npatches + j]) ** 2;
239 if (d2 <= this.r2[j]) cnt++;
240 }
241 }
242 return this.weightsBytes(Math.ceil((cnt * npts) / Math.ceil(npts / 16)));
243 }
245 // ---------------------------------------------------------------- near pairs
247 /** Pairs (target, patch) with |x_t - c_j|^2 <= r2_j, grouped by target, patches ascending. */
248 nearPairs(x: Float64Array, onSurface: boolean): NearPairs {
249 const n = x.length / 3;
250 const { npatches, centroids: cen, patchId, uvNodes } = this.surf;
251 const r2 = this.r2;
252 const tList: number[] = [];
253 const pList: number[] = [];
254 const fList: number[] = [];
255 const rowPtr = new Uint32Array(n + 1);
256 const params = this.surf.exact.params;
257 for (let i = 0; i < n; i++) {
258 const xi = x[i];
259 const yi = x[n + i];
260 const zi = x[2 * n + i];
261 for (let j = 0; j < npatches; j++) {
262 const d2 = (xi - cen[j]) ** 2 + (yi - cen[npatches + j]) ** 2 + (zi - cen[2 * npatches + j]) ** 2;
263 if (d2 > r2[j]) continue;
264 tList.push(i);
265 pList.push(j);
266 if (!onSurface) {
267 fList.push(0, 0, 0, 0);
268 } else if (patchId[i] === j) {
269 fList.push(uvNodes[i], uvNodes[n + i], 2, 0);
270 } else {
271 // the target mapped into this patch's face chart: c_t = x_t / (x_t . n)
272 const par = params.subarray(13 * j, 13 * j + 13);
273 const xn = xi * par[0] + yi * par[1] + zi * par[2];
274 if (xn < 0.2) {
275 fList.push(0, 0, 0, 0);
276 } else {
277 const cx = xi / xn;
278 const cy = yi / xn;
279 const cz = zi / xn;
280 const st = cx * par[3] + cy * par[4] + cz * par[5];
281 const tt = cx * par[6] + cy * par[7] + cz * par[8];
282 fList.push((2 * (st - par[9])) / (par[10] - par[9]) - 1, (2 * (tt - par[11])) / (par[12] - par[11]) - 1, 1, 0);
283 }
284 }
285 }
286 rowPtr[i + 1] = tList.length;
287 }
288 return { pairT: Int32Array.from(tList), pairP: Int32Array.from(pList), rowPtr, pairF: Float32Array.from(fList), npairs: tList.length };
289 }
291 /** The sinh-substituted segments along the four edges for the self patch of node i: 8 floats each. */
292 private selfSegments(i: number): Float32Array {
293 const { npts } = this;
294 const sv = this.surf.srcvals;
295 const u0 = this.surf.uvNodes[i];
296 const v0 = this.surf.uvNodes[npts + i];
297 const ru = [sv[3 * npts + i], sv[4 * npts + i], sv[5 * npts + i]];
298 const rv = [sv[6 * npts + i], sv[7 * npts + i], sv[8 * npts + i]];
299 const E = ru[0] ** 2 + ru[1] ** 2 + ru[2] ** 2;
300 const F = ru[0] * rv[0] + ru[1] * rv[1] + ru[2] * rv[2];
301 const G = rv[0] ** 2 + rv[1] ** 2 + rv[2] ** 2;
302 const det = Math.sqrt(Math.max(E * G - F * F, 1e-300));
303 const k = Math.abs(this.k);
304 const { subLen, oscLen } = this.quad;
305 const items: number[] = [];
306 for (let e = 0; e < 4; e++) {
307 const h = [1 - u0, -1 - u0, 1 - v0, -1 - v0][e];
308 if (Math.abs(h) < 1e-12) continue;
309 let x1: number, x2: number, xc: number, heff: number, arc: number;
310 if (e < 2) {
311 x1 = -1 - v0;
312 x2 = 1 - v0;
313 xc = (-h * F) / G;
314 heff = (Math.abs(h) * det) / G;
315 arc = Math.sqrt(G);
316 } else {
317 x1 = -1 - u0;
318 x2 = 1 - u0;
319 xc = (-h * F) / E;
320 heff = (Math.abs(h) * det) / E;
321 arc = Math.sqrt(E);
322 }
323 const xi0 = Math.asinh((x1 - xc) / heff);
324 const xi1 = Math.asinh((x2 - xc) / heff);
325 // sub-intervals of length <= subLen in xi whose extent along the edge is <= dxmax
326 const dxmax = k * arc > 0 ? oscLen / (k * arc) : Infinity;
327 let xi = xi0;
328 let cnt = 0;
329 while (xi < xi1 && cnt < 256) {
330 const xh = xc + heff * Math.sinh(xi);
331 let xn = Math.min(xi + subLen, Math.asinh((xh + dxmax - xc) / heff));
332 if (xn > xi1 || cnt === 255) xn = xi1;
333 items.push(1, e, xi, xn, h, xc, heff, 0);
334 xi = xn;
335 cnt++;
336 }
337 }
338 return Float32Array.from(items);
339 }
341 // ------------------------------------------------------------ near weights
343 /** Build the near pairs, quadtree leaves and near weights for a set of targets. */
344 private async makeTargets(x: Float64Array, nrm: Float64Array | null, onSurface: boolean, stage: string): Promise<TargetSet> {
345 const { gpu } = this;
346 const n = x.length / 3;
347 const tg = new Float32Array(8 * n);
348 for (let i = 0; i < n; i++) {
349 tg[8 * i] = x[i];
350 tg[8 * i + 1] = x[n + i];
351 tg[8 * i + 2] = x[2 * n + i];
352 if (nrm) {
353 tg[8 * i + 4] = nrm[i];
354 tg[8 * i + 5] = nrm[n + i];
355 tg[8 * i + 6] = nrm[2 * n + i];
356 }
357 }
358 const pairs = this.nearPairs(x, onSurface);
359 const wbytes = this.weightsBytes(pairs.npairs);
360 const lim = gpu.device.limits;
361 if (wbytes > lim.maxStorageBufferBindingSize || wbytes > lim.maxBufferSize) {
362 throw new Error(
363 `the near weights need ${(wbytes / 2 ** 20).toFixed(0)} MB but this GPU allows buffers of at most ` +
364 `${(Math.min(lim.maxStorageBufferBindingSize, lim.maxBufferSize) / 2 ** 20).toFixed(0)} MB; use fewer patches or a lower order`,
365 );
366 }
367 const pi = new Int32Array(2 * Math.max(1, pairs.npairs));
368 for (let p = 0; p < pairs.npairs; p++) {
369 pi[2 * p] = pairs.pairT[p];
370 pi[2 * p + 1] = pairs.pairP[p];
371 }
372 const set: TargetSet = {
373 n,
374 targ: gpu.storage(tg),
375 pairs,
376 pairsI: gpu.storage(pi),
377 pairsF: gpu.storage(pairs.npairs ? pairs.pairF : new Float32Array(4)),
378 rowPtr: gpu.storage(pairs.rowPtr),
379 pairP: gpu.storage(pairs.npairs ? new Uint32Array(pairs.pairP.buffer, pairs.pairP.byteOffset, pairs.npairs) : new Uint32Array(1)),
380 W: gpu.storage(Math.max(16, wbytes)),
381 nleaves: 0,
382 };
383 if (pairs.npairs > 0) set.nleaves = await this.computeWeights(set, stage);
384 return set;
385 }
387 private freeTargets(s: TargetSet) {
388 for (const b of [s.targ, s.pairsI, s.pairsF, s.rowPtr, s.pairP, s.W]) b.destroy();
389 }
391 /** Quadtree leaves and near weights for all pairs of a target set, in paced chunks. Returns the leaf count. */
392 private async computeWeights(set: TargetSet, stage: string): Promise<number> {
393 const { gpu, quad, leaf } = this;
394 const { npairs, pairF } = set.pairs;
395 const leavesPipe = await gpu.pipeline('leaves', BUILD_LEAVES, 'build_leaves');
396 const weightsPipe = await gpu.pipeline('weights', NEAR_WEIGHTS, 'near_weights');
397 const pacer = new Pacer(this.pace, 512);
398 let nleaves = 0;
399 let p0 = 0;
400 let itemsPerPair = 40;
401 const maxItemsBytes = 48 * 2 ** 20;
402 while (p0 < npairs) {
403 if (this.destroyed) throw new Error('operator destroyed');
404 const cap = Math.min(65535, npairs - p0, Math.max(256, Math.floor(maxItemsBytes / (32 * itemsPerPair))));
405 const n = pacer.next(cap);
406 const t0 = now();
407 // pass 1: leaf counts
408 const counts = gpu.storage(4 * n);
409 const uw = new UniformWriter(9);
410 uw.u32(n).u32(p0).u32(0).i32(leaf.nlmax).f32(leaf.rfacQ).f32(Math.abs(this.k)).f32(leaf.kRMax).f32(this.surf.exact.a).i32(this.surf.exact.p);
411 gpu.write(this.uniformLeaves, new Uint8Array(uw.bytes));
412 const dummyItems = gpu.storage(16);
413 const dummyOffsets = gpu.storage(16);
414 gpu.dispatch(leavesPipe, gpu.bind(leavesPipe, [this.uniformLeaves, set.pairsI, set.pairsF, set.targ, this.gpar, counts, dummyOffsets, dummyItems]), Math.ceil(n / 64));
415 const cnt = new Uint32Array(await gpu.read(counts, 4 * n));
416 // self pairs: the CPU segments
417 const ptr = new Uint32Array(n + 1);
418 for (let q = 0; q < n; q++) {
419 const p = p0 + q;
420 const c = pairF[4 * p + 2] === 2 ? this.selfItems[set.pairs.pairT[p]].length / 8 : cnt[q];
421 if (pairF[4 * p + 2] !== 2) nleaves += cnt[q];
422 ptr[q + 1] = ptr[q] + c;
423 }
424 const nitems = ptr[n];
425 const items = new Float32Array(8 * Math.max(1, nitems));
426 for (let q = 0; q < n; q++) {
427 const p = p0 + q;
428 if (pairF[4 * p + 2] === 2) items.set(this.selfItems[set.pairs.pairT[p]], 8 * ptr[q]);
429 }
430 const itemsBuf = gpu.storage(items);
431 const ptrBuf = gpu.storage(ptr);
432 // pass 2: fill the leaves
433 const uw2 = new UniformWriter(9);
434 uw2.u32(n).u32(p0).u32(1).i32(leaf.nlmax).f32(leaf.rfacQ).f32(Math.abs(this.k)).f32(leaf.kRMax).f32(this.surf.exact.a).i32(this.surf.exact.p);
435 gpu.write(this.uniformLeaves, new Uint8Array(uw2.bytes));
436 gpu.dispatch(leavesPipe, gpu.bind(leavesPipe, [this.uniformLeaves, set.pairsI, set.pairsF, set.targ, this.gpar, counts, ptrBuf, itemsBuf]), Math.ceil(n / 64));
437 // the weights
438 const uw3 = new UniformWriter(10);
439 uw3.u32(n).u32(p0).i32(this.norder).i32(this.npols).i32(quad.nq).i32(quad.nqr).i32(quad.nqx).f32(this.k).f32(this.surf.exact.a).i32(this.surf.exact.p);
440 gpu.write(this.uniformWeights, new Uint8Array(uw3.bytes));
441 gpu.dispatch(weightsPipe, gpu.bind(weightsPipe, [this.uniformWeights, set.pairsI, set.pairsF, set.targ, this.gpar, ptrBuf, itemsBuf, this.tables, set.W]), n);
442 await gpu.sync();
443 const dt = now() - t0;
444 pacer.record(n, dt);
445 itemsPerPair = Math.max(1, nitems / n);
446 for (const b of [counts, dummyItems, dummyOffsets, itemsBuf, ptrBuf]) b.destroy();
447 p0 += n;
448 this.progress(stage, p0 / npairs);
449 }
450 return nleaves;
451 }
453 // ------------------------------------------------------------- application
455 /** Values of q at the oversampled nodes, per patch: q_o = M q. */
456 private oversample(q: CVec, out: Float32Array, offset: number) {
457 const { npols, npolsO, overMat } = this;
458 const np = this.surf.npatches;
459 for (let j = 0; j < np; j++) {
460 for (let kk = 0; kk < npolsO; kk++) {
461 let sr = 0;
462 let si = 0;
463 for (let m = 0; m < npols; m++) {
464 const M = overMat[kk * npols + m];
465 sr += M * q[2 * (j * npols + m)];
466 si += M * q[2 * (j * npols + m) + 1];
467 }
468 const idx = 4 * (j * npolsO + kk) + offset;
469 out[idx] = sr;
470 out[idx + 1] = si;
471 }
472 }
473 }
475 /** Far and near parts of the selected kernels at a target set. sel: 6 entries in {0, 1, 2}; reqs: [kernel, density] pairs. */
476 private async applyAt(set: TargetSet, qa: CVec, qb: CVec | null, sel: number[], reqs: [number, number][]): Promise<Float32Array> {
477 const { gpu, npts, npolsO } = this;
478 const N = npolsO * this.surf.npatches;
479 const sq = new Float32Array(4 * N);
480 this.oversample(qa, sq, 0);
481 if (qb) this.oversample(qb, sq, 2);
482 gpu.write(this.srcq, sq);
483 const qd = new Float32Array(2 * 2 * npts);
484 for (let i = 0; i < 2 * npts; i++) qd[i] = qa[i];
485 if (qb) for (let i = 0; i < 2 * npts; i++) qd[2 * npts + i] = qb[i];
486 gpu.write(this.qbuf, qd);
487 const out = set.n === this.wall.n ? this.out : gpu.storage(6 * 8 * set.n);
488 gpu.clear(out);
489 const selmask = sel.reduce((acc, s, m) => acc | (s << (4 * m)), 0);
490 const densePipe = await gpu.pipeline('dense', DENSE_APPLY, 'dense_apply');
491 const group = gpu.bind(densePipe, [this.uniformDense, set.targ, this.srcgeo, this.srcq, set.rowPtr, set.pairP, out]);
492 const pacer = new Pacer(this.pace, Math.max(1, Math.floor(4e8 / (set.n * npolsO * 60))));
493 const np = this.surf.npatches;
494 let j0 = 0;
495 while (j0 < np) {
496 const n = pacer.next(np - j0);
497 const t0 = now();
498 const uw = new UniformWriter(6);
499 uw.u32(set.n).u32(j0).u32(j0 + n).u32(npolsO).f32(this.k).u32(selmask);
500 gpu.write(this.uniformDense, new Uint8Array(uw.bytes));
501 gpu.dispatch(densePipe, group, Math.ceil(set.n / DENSE_WG));
502 await gpu.sync();
503 pacer.record(n, now() - t0);
504 j0 += n;
505 }
506 if (set.pairs.npairs > 0 && reqs.length > 0) {
507 const sparsePipe = await gpu.pipeline('sparse', SPARSE_APPLY, 'sparse_apply');
508 const uw = new UniformWriter(8);
509 uw.u32(set.n).u32(this.npols).u32(npts).u32(reqs.length);
510 for (let r = 0; r < 4; r++) uw.u32(r < reqs.length ? reqs[r][0] | (reqs[r][1] << 8) : 0);
511 gpu.write(this.uniformSparse, new Uint8Array(uw.bytes));
512 gpu.dispatch(sparsePipe, gpu.bind(sparsePipe, [this.uniformSparse, set.rowPtr, set.pairP, set.W, this.qbuf, out]), Math.ceil(set.n / 128));
513 }
514 const res = new Float32Array(await gpu.read(out, 6 * 8 * set.n));
515 if (out !== this.out) out.destroy();
516 return res;
517 }
519 zid(): [number, number] {
520 return [(-1) ** this.ifinout / 2, -this.alpha / 4];
521 }
523 /** The operator without the identity term z. */
524 async applyK(sigma: CVec): Promise<CVec> {
525 const { npts, k, alpha: a, lam } = this;
526 const F = await this.applyAt(this.wall, sigma, null, [1, 1, 1, 0, 1, 0], [[0, 0], [1, 0], [2, 0], [4, 0]]);
527 const phi1 = new Float64Array(2 * npts);
528 const phi2 = new Float64Array(2 * npts);
529 const res = new Float64Array(2 * npts);
530 for (let i = 0; i < npts; i++) {
531 phi1[2 * i] = F[2 * (npts + i)];
532 phi1[2 * i + 1] = F[2 * (npts + i) + 1];
533 phi2[2 * i] = F[2 * (2 * npts + i)];
534 phi2[2 * i + 1] = F[2 * (2 * npts + i) + 1];
535 const s4r = F[2 * (4 * npts + i)];
536 const s4i = F[2 * (4 * npts + i) + 1];
537 // res = F0 + i k lam F4
538 res[2 * i] = F[2 * i] - k * lam * s4i;
539 res[2 * i + 1] = F[2 * i + 1] + k * lam * s4r;
540 }
541 const G = await this.applyAt(this.wall, phi1, phi2, [0, 0, 2, 1, 0, 1], [[2, 1], [3, 0], [5, 0]]);
542 const sgn = (-1) ** this.ifinout;
543 for (let i = 0; i < npts; i++) {
544 const g2r = G[2 * (2 * npts + i)];
545 const g2i = G[2 * (2 * npts + i) + 1];
546 const g3r = G[2 * (3 * npts + i)];
547 const g3i = G[2 * (3 * npts + i) + 1];
548 const g5r = G[2 * (5 * npts + i)];
549 const g5i = G[2 * (5 * npts + i) + 1];
550 // res += i a (G2 + G3) - k a lam G5 + sgn k a lam phi1 / 2
551 res[2 * i] += -a * (g2i + g3i) - k * a * lam * g5r + (sgn * k * a * lam * phi1[2 * i]) / 2;
552 res[2 * i + 1] += a * (g2r + g3r) - k * a * lam * g5i + (sgn * k * a * lam * phi1[2 * i + 1]) / 2;
553 }
554 return res;
555 }
557 /** phi1 = S_kappa[sigma] on the surface. */
558 async sKappa(sigma: CVec): Promise<CVec> {
559 const F = await this.applyAt(this.wall, sigma, null, [0, 1, 0, 0, 0, 0], [[1, 0]]);
560 const { npts } = this;
561 const phi1 = new Float64Array(2 * npts);
562 for (let i = 0; i < 2 * npts; i++) phi1[i] = F[2 * npts + i];
563 return phi1;
564 }
566 /** One of the six operators applied to q on the surface (principal value), for tests. */
567 async applySingle(m: number, q: CVec): Promise<CVec> {
568 const sel = [0, 0, 0, 0, 0, 0];
569 sel[m] = 1;
570 const F = await this.applyAt(this.wall, q, null, sel, [[m, 0]]);
571 const { npts } = this;
572 const out = new Float64Array(2 * npts);
573 for (let i = 0; i < 2 * npts; i++) out[i] = F[2 * m * npts + i];
574 return out;
575 }
577 async solve(rhs: CVec, epsGmres = 1e-5, numit = 200, onIter?: (it: number, rres: number) => void): Promise<GmresResult & { phi1: CVec; seconds: number }> {
578 const t0 = now();
579 const r = await gmres((x) => this.applyK(x), rhs, this.surf.wts, this.zid(), epsGmres, numit, onIter);
580 const phi1 = await this.sKappa(r.x);
581 return { ...r, phi1, seconds: now() - t0 };
582 }
584 /** u = S_k sigma + i alpha D_k phi1 at the surface nodes, as the limit from the side selected by ifinout. */
585 async evalOnSurface(sigma: CVec, phi1: CVec): Promise<CVec> {
586 const F = await this.applyAt(this.wall, sigma, phi1, [0, 0, 0, 0, 1, 2], [[4, 0], [5, 1]]);
587 return this.combine(F, this.npts, phi1);
588 }
590 /** u = S_k sigma + i alpha D_k phi1 at off-surface targets x (3, n), in chunks. */
591 async eval(x: Float64Array, sigma: CVec, phi1: CVec, stage = 'field'): Promise<CVec> {
592 const n = x.length / 3;
593 const u = new Float64Array(2 * n);
594 const chunk = 4096;
595 for (let i0 = 0; i0 < n; i0 += chunk) {
596 const m = Math.min(chunk, n - i0);
597 const xs = new Float64Array(3 * m);
598 for (let d = 0; d < 3; d++) xs.set(x.subarray(d * n + i0, d * n + i0 + m), d * m);
599 const set = await this.makeTargets(xs, null, false, stage);
600 const F = await this.applyAt(set, sigma, phi1, [0, 0, 0, 0, 1, 2], [[4, 0], [5, 1]]);
601 this.freeTargets(set);
602 u.set(this.combine(F, m, null), 2 * i0);
603 this.progress(stage, (i0 + m) / n);
604 }
605 return u;
606 }
608 private combine(F: Float32Array, n: number, phi1OnSurface: CVec | null): CVec {
609 const a = this.alpha;
610 const u = new Float64Array(2 * n);
611 const jump = phi1OnSurface ? (a * (-1) ** (this.ifinout + 1)) / 2 : 0;
612 for (let i = 0; i < n; i++) {
613 const s4r = F[2 * (4 * n + i)];
614 const s4i = F[2 * (4 * n + i) + 1];
615 const d5r = F[2 * (5 * n + i)];
616 const d5i = F[2 * (5 * n + i) + 1];
617 // u = S + i a D (+ i a (-1)^(ifinout+1) phi1 / 2 on the surface)
618 u[2 * i] = s4r - a * d5i;
619 u[2 * i + 1] = s4i + a * d5r;
620 if (phi1OnSurface) {
621 u[2 * i] += -jump * phi1OnSurface[2 * i + 1];
622 u[2 * i + 1] += jump * phi1OnSurface[2 * i];
623 }
624 }
625 return u;
626 }
628 /** The near weights of the surface nodes, for tests: (npairs, 6, npols) complex interleaved. */
629 async readWallWeights(): Promise<{ pairs: NearPairs; W: Float32Array }> {
630 const W = new Float32Array(await this.gpu.read(this.wall.W, this.weightsBytes(this.wall.pairs.npairs)));
631 return { pairs: this.wall.pairs, W };
632 }
635export { NT };
↑↓ move↵openescclose