Bin points with a WebGPU compute shader, with CPU fallback - concept-collection/random-points-in-disk

concept-collection / random-points-in-disk

Bin points with a WebGPU compute shader, with CPU fallback

Jeremy Magland <jmagland@flatironinstitute.org> committed 27 days ago commit 35a88f5a39d6 parent 393eebb Browse files

2 changed files+212−5

README.mdmodified+14−0View file

@@ -22,6 +22,20 @@ and how many points a target noise level would require.

2222 Statistics are computed only over voxels lying *entirely* inside the disk; boundary

2323 voxels have a smaller expected count and would otherwise inflate the measured spread.

2424

25+## Sampling backend

26+

27+Binning the points is the entire cost here and it parallelizes perfectly, so when WebGPU

28+is available a compute shader draws the points and accumulates the histogram with atomics

29+— which lifts the usable range from ~30M points to a billion. Coarse grids would

30+serialize on a handful of global counters, so for 32×32 and below each workgroup

31+accumulates into a private histogram in workgroup memory and flushes once per bin at the

32+end. Random numbers come from a PCG hash seeded per thread; since this demo is *about*

33+uniformity, that generator was checked against the 1/√μ prediction before being trusted

34+(a hash with visible structure would paint that structure straight into the heatmap).

35+

36+There is a plain-JS fallback for browsers without WebGPU, and the badge next to

37+"New sample" shows which backend is live. The point count is capped lower on the CPU path.

38+

2539 ## Motivation

2640

2741 This is the discretization question behind isochromat-based MRI simulation: a voxel's

index.htmlmodified+198−5View file

@@ -48,6 +48,12 @@

4848 select { width: 100%; }

4949 button { cursor: pointer; white-space: nowrap; }

5050 button:hover { border-color: var(--accent); }

51+ .badge {

52+ font-size: .72rem; letter-spacing: .06em; text-transform: uppercase;

53+ padding: 3px 8px; border-radius: 20px; white-space: nowrap;

54+ border: 1px solid var(--line); color: var(--dim);

55+ }

56+ .badge.on { border-color: #2e7d5b; color: #6ddba4; }

5157

5258 /* ---- panels ---- */

5359 .panels { display: grid; grid-template-columns: repeat(3, 1fr); gap: 14px; margin: 14px 0; }

@@ -106,7 +112,10 @@

106112 <option value="512">512 × 512</option>

107113 </select>

108114 </div>

109- <div><button id="resample">New sample</button></div>

115+ <div style="display:flex;align-items:center;gap:10px">

116+ <button id="resample">New sample</button>

117+ <span class="badge" id="backend" title="Points are binned by a WebGPU compute shader when available, otherwise on the CPU">CPU</span>

118+ </div>

110119 </div>

111120

112121 <div class="panels">

@@ -294,6 +303,153 @@ function fmtPct(x) {

294303 return (x * 100).toFixed(2) + '%';

295304 }

296305

306+// ---------- GPU sampling ----------

307+// The histogram is the whole cost of this demo, and it is embarrassingly parallel: one

308+// compute shader draws the points and atomically bins them, so N can reach a billion.

309+// The CPU path above stays as the fallback.

310+const MAX_G = 512, MAX_BINS = MAX_G * MAX_G;

311+const WG_SIZE = 256, WORKGROUPS = 1024; // 262144 grid-stride threads

312+const LOCAL_BINS = 1024; // workgroup-private histogram, for G <= 32

313+

314+const WGSL = `

315+struct Params { g: u32, n: u32, seed: u32, use_local: u32 }

316+@group(0) @binding(0) var<uniform> P: Params;

317+@group(0) @binding(1) var<storage, read_write> counts: array<atomic<u32>>;

318+

319+// Coarse grids would serialize on a handful of global counters, so each workgroup

320+// accumulates privately and flushes once per bin at the end.

321+var<workgroup> local_hist: array<atomic<u32>, ${LOCAL_BINS}>;

322+

323+// PCG (rxs-m-xs). Validated against the Poisson prediction before being trusted here:

324+// a hash with visible structure would paint that structure straight into the heatmap.

325+fn hash(x: u32) -> u32 {

326+ var s = x * 747796405u + 2891336453u;

327+ s = ((s >> ((s >> 28u) + 4u)) ^ s) * 277803737u;

328+ return (s >> 22u) ^ s;

329+}

330+fn pcg(state: ptr<function, u32>) -> u32 {

331+ let old = *state;

332+ *state = old * 747796405u + 2891336453u;

333+ let word = ((old >> ((old >> 28u) + 4u)) ^ old) * 277803737u;

334+ return (word >> 22u) ^ word;

335+}

336+

337+const INV32: f32 = 2.3283064365386963e-10;

338+const TAU: f32 = 6.283185307179586;

339+

340+@compute @workgroup_size(${WG_SIZE})

341+fn main(@builtin(global_invocation_id) gid: vec3<u32>,

342+ @builtin(local_invocation_id) lid: vec3<u32>,

343+ @builtin(num_workgroups) nwg: vec3<u32>) {

344+ let nbins = P.g * P.g;

345+ let use_local = P.use_local == 1u;

346+

347+ if (use_local) {

348+ for (var i = lid.x; i < nbins; i = i + ${WG_SIZE}u) {

349+ atomicStore(&local_hist[i], 0u);

350+ }

351+ }

352+ workgroupBarrier();

353+

354+ var state = hash(gid.x ^ hash(P.seed));

355+ let stride = nwg.x * ${WG_SIZE}u;

356+ let half = f32(P.g) * 0.5;

357+ let top = f32(P.g) - 1.0;

358+

359+ for (var idx = gid.x; idx < P.n; idx = idx + stride) {

360+ let r = sqrt(f32(pcg(&state)) * INV32); // sqrt for uniform area density

361+ let t = TAU * f32(pcg(&state)) * INV32;

362+ let ix = u32(clamp((r * cos(t) + 1.0) * half, 0.0, top));

363+ let iy = u32(clamp((r * sin(t) + 1.0) * half, 0.0, top));

364+ let bin = iy * P.g + ix;

365+ if (use_local) { atomicAdd(&local_hist[bin], 1u); }

366+ else { atomicAdd(&counts[bin], 1u); }

367+ }

368+

369+ workgroupBarrier();

370+ if (use_local) {

371+ for (var i = lid.x; i < nbins; i = i + ${WG_SIZE}u) {

372+ let v = atomicLoad(&local_hist[i]);

373+ if (v > 0u) { atomicAdd(&counts[i], v); }

374+ }

375+ }

376+}

377+`;

378+

379+let gpu = null;

380+let seed = 1;

381+

382+async function initGPU() {

383+ if (!navigator.gpu) return null;

384+ try {

385+ const adapter = await navigator.gpu.requestAdapter();

386+ if (!adapter) return null;

387+ const device = await adapter.requestDevice();

388+ const module = device.createShaderModule({ code: WGSL });

389+ if (module.getCompilationInfo) {

390+ const info = await module.getCompilationInfo();

391+ const errs = info.messages.filter((m) => m.type === 'error');

392+ if (errs.length) { console.error('WGSL compilation failed', errs); return null; }

393+ }

394+ const pipeline = device.createComputePipeline({

395+ layout: 'auto', compute: { module, entryPoint: 'main' },

396+ });

397+ const U = GPUBufferUsage;

398+ const params = device.createBuffer({ size: 16, usage: U.UNIFORM | U.COPY_DST });

399+ const counts = device.createBuffer({ size: MAX_BINS * 4, usage: U.STORAGE | U.COPY_SRC | U.COPY_DST });

400+ const stage = device.createBuffer({ size: MAX_BINS * 4, usage: U.COPY_DST | U.MAP_READ });

401+ const bind = device.createBindGroup({

402+ layout: pipeline.getBindGroupLayout(0),

403+ entries: [

404+ { binding: 0, resource: { buffer: params } },

405+ { binding: 1, resource: { buffer: counts } },

406+ ],

407+ });

408+ device.addEventListener('uncapturederror', (e) => console.error('WebGPU error', e.error));

409+ device.lost.then(() => { gpu = null; showBackend(); });

410+ return { device, pipeline, params, counts, stage, bind, host: new Uint32Array(MAX_BINS) };

411+ } catch (e) {

412+ console.warn('WebGPU unavailable', e);

413+ return null;

414+ }

415+}

416+

417+async function gpuCounts(N, G) {

418+ const g = gpu;

419+ const bytes = G * G * 4;

420+ g.device.queue.writeBuffer(g.params, 0,

421+ new Uint32Array([G, N, seed, G * G <= LOCAL_BINS ? 1 : 0]));

422+ const enc = g.device.createCommandEncoder();

423+ enc.clearBuffer(g.counts, 0, bytes);

424+ const pass = enc.beginComputePass();

425+ pass.setPipeline(g.pipeline);

426+ pass.setBindGroup(0, g.bind);

427+ pass.dispatchWorkgroups(WORKGROUPS);

428+ pass.end();

429+ enc.copyBufferToBuffer(g.counts, 0, g.stage, 0, bytes);

430+ g.device.queue.submit([enc.finish()]);

431+ await g.stage.mapAsync(GPUMapMode.READ, 0, bytes);

432+ g.host.set(new Uint32Array(g.stage.getMappedRange(0, bytes)));

433+ g.stage.unmap();

434+ return g.host; // only [0, G*G) is meaningful

435+}

436+

437+// Falls back permanently to the CPU if the GPU path ever throws mid-session. Returns null

438+// in that case rather than sampling: N may be far above what the CPU can take, so the

439+// caller re-runs after showBackend() has clamped the slider back to the CPU range.

440+async function getCounts(N, G) {

441+ if (gpu) {

442+ try { return await gpuCounts(N, G); }

443+ catch (e) {

444+ console.warn('GPU sampling failed, falling back to CPU', e);

445+ gpu = null;

446+ showBackend();

447+ return null;

448+ }

449+ }

450+ return sampleCounts(N, G);

451+}

452+

297453 // ---------- app ----------

298454 const el = (id) => document.getElementById(id);

299455 const nSlider = el('n'), resSel = el('res'), targetSel = el('target');

@@ -314,9 +470,21 @@ function showN(N) {

314470 el('nInline').textContent = fmtCount(N);

315471 }

316472

473+// The CPU tops out around 30M points before a drag becomes unpleasant; the GPU

474+// comfortably reaches a billion, which is where fine grids get interesting.

475+const CPU_MAX = 750, GPU_MAX = 900; // slider units, = log10(N) * 100

476+

477+function showBackend() {

478+ el('backend').textContent = gpu ? 'WebGPU' : 'CPU';

479+ el('backend').className = 'badge' + (gpu ? ' on' : '');

480+ const max = gpu ? GPU_MAX : CPU_MAX;

481+ nSlider.max = String(max);

482+ if (+nSlider.value > max) nSlider.value = String(max);

483+}

484+

317485 let lastElapsed = 0;

318486

319-function update() {

487+async function draw() {

320488 const N = sliderToN(+nSlider.value);

321489 const G = +resSel.value;

322490

@@ -324,7 +492,8 @@ function update() {

324492

325493 const frac = areaFractions(G);

326494 const t0 = performance.now();

327- const counts = sampleCounts(N, G);

495+ const counts = await getCounts(N, G);

496+ if (!counts) { queued = true; return; } // backend just changed; redraw within the new limits

328497 const elapsed = lastElapsed = performance.now() - t0;

329498

330499 // Expected count in a fully interior voxel: N * (voxel area / disk area).

@@ -363,7 +532,22 @@ function update() {

363532 const need = Math.PI * G * G / (4 * c * c);

364533 el('sNeed').innerHTML = fmtCount(Math.round(need)) + ' <small>points</small>';

365534

366- el('sTime').innerHTML = (elapsed < 10 ? elapsed.toFixed(1) : elapsed.toFixed(0)) + ' <small>ms</small>';

535+ el('sTime').innerHTML = (elapsed < 10 ? elapsed.toFixed(1) : elapsed.toFixed(0))

536+ + ' <small>ms · ' + (gpu ? 'GPU' : 'CPU') + '</small>';

537+}

538+

539+// Sampling is async now, so serialize runs and collapse anything that arrives while one

540+// is in flight down to a single trailing update.

541+let running = false, queued = false;

542+async function update() {

543+ if (running) { queued = true; return; }

544+ running = true;

545+ try { await draw(); }

546+ catch (e) { console.error(e); }

547+ finally {

548+ running = false;

549+ if (queued) { queued = false; schedule(); }

550+ }

367551 }

368552

369553 // Coalesce rapid input into one update per frame; once a sample gets slow enough to

@@ -384,8 +568,17 @@ function schedule() {

384568 nSlider.addEventListener('input', () => { showN(sliderToN(+nSlider.value)); schedule(); });

385569 resSel.addEventListener('change', schedule);

386570 targetSel.addEventListener('change', schedule);

387-el('resample').addEventListener('click', update);

571+el('resample').addEventListener('click', () => { seed++; update(); });

572+

573+// Draw immediately on the CPU so the page is never blank, then upgrade to the GPU.

574+showBackend();

388575 update();

576+initGPU().then((g) => {

577+ if (!g) return;

578+ gpu = g;

579+ showBackend();

580+ update();

581+});

389582 </script>

390583 </body>

391584 </html>