From b0e78e560312baff97c06006c3da4673569abc72 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Mon, 27 Jul 2026 17:41:36 +0100 Subject: [PATCH 01/19] latest --- src/lib/decimate/decimate-source.ts | 35 +- src/lib/decimate/edge-cost-cpu.ts | 328 ++++++++++-------- src/lib/decimate/priority.ts | 129 +++---- src/lib/decimate/select-recost.ts | 457 +++++++++++++++++++++++++ src/lib/decimate/select.ts | 169 +++++----- src/lib/gpu/gpu-edge-cost.ts | 498 ++++++++++------------------ test/decimate-edge-cost.test.mjs | 64 ++++ test/decimate-merge-stream.test.mjs | 18 +- test/decimate-priority.test.mjs | 11 +- test/decimate-select.test.mjs | 91 +++-- test/decimate-source.test.mjs | 15 +- test/decimate.test.mjs | 21 +- 12 files changed, 1138 insertions(+), 698 deletions(-) create mode 100644 src/lib/decimate/select-recost.ts create mode 100644 test/decimate-edge-cost.test.mjs diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index 39d58149..2064f546 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -6,13 +6,14 @@ import { mergeStream } from './merge-stream'; import { kdPartition, coherenceRuns, type ResidentPositions } from './partition'; import { runPriorityPass, HALO_CAP, type CandidateArrays } from './priority'; import { selectMerges } from './select'; +import { selectMergesRecosted, CACHE_STRIDE } from './select-recost'; import { compact, type ChunkDataPool, type ChunkSource, type ChunkSourceMetadata } from '../chunk'; -import { APP_CHUNK } from '../gpu/gpu-edge-cost'; +import { SPLAT_STRIDE } from '../gpu/gpu-edge-cost'; import { type ReadFileSystem } from '../io/read'; import { type FileSystem } from '../io/write'; import { bakeTransform } from '../ops'; @@ -30,8 +31,14 @@ const BLOCK_SIZE = 1 << 21; /** Same no-grind stall semantics as legacy: a shortfall generation must remove at least this fraction. */ const MIN_ITERATION_PROGRESS = 0.05; -/** Default resident-memory budget steering the candidate-K policy. */ -const DEFAULT_MEMORY_BUDGET = 24 * 2 ** 30; +/** Default resident-memory budget steering the candidate-K and re-costed-selection policies. */ +const DEFAULT_MEMORY_BUDGET = 48 * 2 ** 30; + +// Per-gaussian residency of re-costed selection beyond the base state: splat +// cache (16 f32) + neighbour ids (k u32) + f64 cluster moments/colour/error + +// union-find/chains/heap. Conservative round-up; used by the per-generation +// gate that falls back to one-shot selectMerges when over budget. +const RECOST_BYTES_PER_GAUSSIAN = (k: number) => CACHE_STRIDE * 4 + k * 4 + 200; /** Coherence heuristic: gap (rows) merged into one run / runs-per-block considered scattered. */ const COHERENCE_GAP_ROWS = 64; @@ -59,7 +66,7 @@ type DecimateOptions = { createDevice?: DeviceCreator; /** Spill destination for over-budget intermediate generations. */ spill?: DecimateSpill; - /** Resident-memory budget driving the candidate-K policy (default 24 GiB). */ + /** Resident-memory budget driving the candidate-K and re-costed-selection policies (default 48 GiB). */ memoryBudgetBytes?: number; }; @@ -177,7 +184,7 @@ const decimateSource = async ( const bindingLimit = (device as unknown as { limits?: { maxStorageBufferBindingSize?: number } } | undefined) ?.limits?.maxStorageBufferBindingSize; if (typeof bindingLimit === 'number') { - const largestBinding = (bs: number) => bs * (1 + HALO_CAP) * Math.max(Math.min(APP_CHUNK, colorDim) * 4, 36); + const largestBinding = (bs: number) => bs * (1 + HALO_CAP) * SPLAT_STRIDE * 4; while (blockSize > (1 << 16) && largestBinding(blockSize) > bindingLimit) { blockSize >>= 1; } @@ -206,9 +213,19 @@ const decimateSource = async ( cost: new Float32Array(N * K).fill(Infinity) }; + // Re-costed selection (exact within-generation greedy) when its + // resident state fits the budget alongside the base state; one-shot + // selection otherwise. Gated per generation, so large scenes regain + // re-costing as soon as the cascade shrinks under the budget. + const k = Math.min(KNN_K, Math.max(1, N - 1)); + const baseBytes = N * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; + const recost = baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; + const cacheOut = recost ? new Float32Array(N * CACHE_STRIDE) : undefined; + const neighborsOut = recost ? new Uint32Array(N * k) : undefined; + const priorityBar = logger.bar('computing merge priorities', N); await runPriorityPass( - { source: src, pool, pos: positions, order, blocks, device, K, k: Math.min(KNN_K, Math.max(1, N - 1)) }, + { source: src, pool, pos: positions, order, blocks, device, K, k, cacheOut, neighborsOut }, cand, n => priorityBar.tick(n) ); @@ -216,8 +233,10 @@ const decimateSource = async ( const generationTarget = Math.max(targetCount, N - Math.floor(N / 2)); const needed = N - generationTarget; - const selectSub = logger.group('Selecting merges'); - const selection = selectMerges(cand, N, K, needed); + const selectSub = logger.group(recost ? 'Selecting merges (re-costed)' : 'Selecting merges'); + const selection = cacheOut ? + selectMergesRecosted({ cand, K, splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed }) : + selectMerges(cand, N, K, needed); selectSub.end(); if (selection.removed === 0) { diff --git a/src/lib/decimate/edge-cost-cpu.ts b/src/lib/decimate/edge-cost-cpu.ts index c0ae5835..8e5ec9f5 100644 --- a/src/lib/decimate/edge-cost-cpu.ts +++ b/src/lib/decimate/edge-cost-cpu.ts @@ -1,123 +1,172 @@ /** - * CPU edge-cost path for chunk-native decimation — a port of the legacy - * `computeEdgeCost` + `buildPerSplatCache` (CPU variant) over a - * {@link SplatView}. Formula-identical to the legacy implementation and to - * the `GpuEdgeCost` WGSL kernel; used as the no-device fallback and by the - * GPU parity tests. + * CPU edge-cost path for chunk-native decimation. + * + * PROTOTYPE — principled L2 field error. The cost of merging splats i and j + * into the moment-matched splat m is the squared L2 norm of the difference + * between the original emitted field (αc·G for each splat, G the unnormalized + * Gaussian kernel) and the merged field: + * + * E = ∫ ‖ αᵢcᵢGᵢ + αⱼcⱼGⱼ − αₘcₘGₘ ‖² dx + * + * which expands to a closed form in Gaussian–Gaussian L2 products + * ⟨G_a,G_b⟩ = (2π)^{3/2}|Σ_a|^{1/2}|Σ_b|^{1/2}|Σ_a+Σ_b|^{-1/2} + * exp(−½ dᵀ(Σ_a+Σ_b)⁻¹ d), d = μ_a − μ_b. + * + * Every term scales with √|Σ| (Gaussian volume), so merging large (distant + * sky) Gaussians costs far more than a geometrically-similar tiny merge — + * unlike the scale-invariant KL cost this replaces. The metric is ≥ 0 by + * construction (a squared norm). Merged parameters (μ, Σ, α, colour) are the + * faithful moment-matched values from {@link mergeGroup}, so the cost predicts + * the real output error. + * + * Modelling choices (prototype): amplitude is α × base colour (0.5 + C0·f_dc), + * i.e. higher-order SH is carried through the merge but not scored in the + * ranking cost. * * Engine-free. */ import { EPS_COV, - LOG2PI, - logAddExp, sigmoid, ellipsoidArea, quatToRotmat, sigmaFromRotVar, - det3, - gaussLogpdfDiagrot, - type SplatView, - type MergeScratch + type SplatView } from './moment-match'; +/** SH band-0 constant (f_dc → base colour: 0.5 + C0·f_dc). */ +const C0 = 0.28209479177387814; + +/** + * Scale-free colour dissimilarity weight. Unlike the field-L2's own colour + * sensitivity (which vanishes ∝σ³ for faint splats), this term keeps + * light-vs-dark pairing selective at any scale — without it, fine texture + * (validated on grass/snow speckle) merges colour-blind and washes to mush. + * Calibrated as λ=1e-6 over raw f_dc coefficients (crop+full-scale study); + * base colour = 0.5 + C0·f_dc and C0² = 1/(4π), so in base-colour space the + * constant is 4π·1e-6. Mirrored in the GpuEdgeCost WGSL kernel. + */ +const COLOR_WEIGHT = 4 * Math.PI * 1e-6; + +/** ⟨G,G⟩ self-product constant: (2π)^{3/2}·2^{-3/2} = π^{3/2}. */ +const PI_1_5 = Math.PI ** 1.5; + +/** Cross-product constant (2π)^{3/2}. */ +const TWO_PI_1_5 = (2 * Math.PI) ** 1.5; + /** - * Per-splat derived quantities for the cost function (legacy - * `buildPerSplatCache`, forGpu = false). `mass` uses the cost-path epsilon - * (+1e-12), matching legacy exactly. + * Per-splat derived quantities for the L2 field cost. `sig` holds each + * covariance's 6 unique components [xx, xy, xz, yy, yz, zz]; `sqrtDet` is + * √|Σ| = sx·sy·sz; `alpha` is the linear opacity; `base` is the 3-channel base + * colour; `baseN2` is |base|²; `mass` is the area·α merge weight. */ type CostCache = { - R: Float32Array; - v: Float32Array; - invdiag: Float32Array; - logdet: Float32Array; - sigma: Float32Array; + sig: Float32Array; + sqrtDet: Float32Array; + alpha: Float32Array; + base: Float32Array; + baseN2: Float32Array; mass: Float32Array; }; const buildCostCache = (view: SplatView): CostCache => { - const { geo } = view; + const { geo, color, colorDim } = view; const n = geo.length / 8; - const R = new Float32Array(n * 9); - const v = new Float32Array(n * 3); - const invdiag = new Float32Array(n * 3); - const logdet = new Float32Array(n); - const sigma = new Float32Array(n * 9); + const sig = new Float32Array(n * 6); + const sqrtDet = new Float32Array(n); + const alpha = new Float32Array(n); + const base = new Float32Array(n * 3); + const baseN2 = new Float32Array(n); const mass = new Float32Array(n); + const R = new Float32Array(9); + const S = new Float32Array(9); for (let i = 0; i < n; i++) { - const i3 = 3 * i; const i8 = 8 * i; - const i9 = 9 * i; - const linAlpha = sigmoid(geo[i8 + 7]); + const a = sigmoid(geo[i8 + 7]); const sx = Math.max(Math.exp(geo[i8 + 4]), 1e-12); const sy = Math.max(Math.exp(geo[i8 + 5]), 1e-12); const sz = Math.max(Math.exp(geo[i8 + 6]), 1e-12); - const vx = sx * sx + EPS_COV; - const vy = sy * sy + EPS_COV; - const vz = sz * sz + EPS_COV; - - v[i3] = vx; v[i3 + 1] = vy; v[i3 + 2] = vz; - invdiag[i3] = 1 / Math.max(vx, 1e-30); - invdiag[i3 + 1] = 1 / Math.max(vy, 1e-30); - invdiag[i3 + 2] = 1 / Math.max(vz, 1e-30); - logdet[i] = Math.log(Math.max(vx, 1e-30)) + Math.log(Math.max(vy, 1e-30)) + Math.log(Math.max(vz, 1e-30)); - let qw = geo[i8], qx = geo[i8 + 1], qy = geo[i8 + 2], qz = geo[i8 + 3]; - const qn = Math.hypot(qw, qx, qy, qz); - const invq = 1 / Math.max(qn, 1e-12); + const invq = 1 / Math.max(Math.hypot(qw, qx, qy, qz), 1e-12); qw *= invq; qx *= invq; qy *= invq; qz *= invq; - quatToRotmat(qw, qx, qy, qz, R, i9); - sigmaFromRotVar(R, i9, vx, vy, vz, sigma, i9); + quatToRotmat(qw, qx, qy, qz, R, 0); + sigmaFromRotVar(R, 0, sx * sx, sy * sy, sz * sz, S, 0); - mass[i] = linAlpha * ellipsoidArea(sx, sy, sz) + 1e-12; + const i6 = 6 * i; + sig[i6] = S[0]; sig[i6 + 1] = S[1]; sig[i6 + 2] = S[2]; + sig[i6 + 3] = S[4]; sig[i6 + 4] = S[5]; sig[i6 + 5] = S[8]; + + sqrtDet[i] = sx * sy * sz; + alpha[i] = a; + mass[i] = a * ellipsoidArea(sx, sy, sz) + 1e-30; + + const i3 = 3 * i; + const br = 0.5 + C0 * color[i * colorDim]; + const bg = 0.5 + C0 * color[i * colorDim + 1]; + const bb = 0.5 + C0 * color[i * colorDim + 2]; + base[i3] = br; base[i3 + 1] = bg; base[i3 + 2] = bb; + baseN2[i] = br * br + bg * bg + bb * bb; } - return { R, v, invdiag, logdet, sigma, mass }; + return { sig, sqrtDet, alpha, base, baseN2, mass }; +}; + +// dᵀ(A)⁻¹d and √|A| for a symmetric 3×3 A = [xx,xy,xz,yy,yz,zz]; returns the +// Gaussian cross-product ⟨G_a,G_b⟩ scaled by the caller's √|Σ_a|·√|Σ_b|. +const crossG = ( + sdA: number, sdB: number, + xx: number, xy: number, xz: number, yy: number, yz: number, zz: number, + dx: number, dy: number, dz: number +): number => { + // Signed cofactors (adj is symmetric); det via row-0 expansion. + const c00 = yy * zz - yz * yz; + const c01 = xz * yz - xy * zz; + const c02 = xy * yz - xz * yy; + const c11 = xx * zz - xz * xz; + const c12 = xy * xz - xx * yz; + const c22 = xx * yy - xy * xy; + const det = Math.max(xx * c00 + xy * c01 + xz * c02, 1e-30); + const quadAdj = c00 * dx * dx + c11 * dy * dy + c22 * dz * dz + + 2 * (c01 * dx * dy + c02 * dx * dz + c12 * dy * dz); + const quad = quadAdj / det; + return TWO_PI_1_5 * sdA * sdB / Math.sqrt(det) * Math.exp(-0.5 * quad); }; /** - * Edge cost between splats `i` and `j` of the view: KL-style geometric term - * (single MC sample) + L2 over the color/SH coefficients. Legacy - * `computeEdgeCost`, verbatim. + * L2 field-error cost of merging splats `i` and `j` of the view. * * @param view - Splat columns. * @param cache - Per-splat cache from {@link buildCostCache}. * @param i - First splat (view row). * @param j - Second splat (view row). - * @param Z - MC samples (legacy: one sample, seed 0). - * @param scratch - Merge scratch (uses `sigm`). - * @returns The edge cost. + * @returns The edge cost (≥ 0). */ const computeEdgeCostView = ( view: SplatView, cache: CostCache, i: number, - j: number, - Z: Float64Array[], - scratch: MergeScratch + j: number ): number => { - const { pos, color, colorDim } = view; + const { pos } = view; + const { sig, sqrtDet, alpha, base, baseN2, mass } = cache; + const i3 = 3 * i, j3 = 3 * j; - const i9 = 9 * i, j9 = 9 * j; + const i6 = 6 * i, j6 = 6 * j; const mux = pos[i3], muy = pos[i3 + 1], muz = pos[i3 + 2]; const mvx = pos[j3], mvy = pos[j3 + 1], mvz = pos[j3 + 2]; - const wi = cache.mass[i], wj = cache.mass[j]; - const W = wi + wj; - const Wsafe = W > 0 ? W : 1; - - let pi = wi / Wsafe; - pi = Math.max(1e-12, Math.min(1 - 1e-12, pi)); + const mi = mass[i], mj = mass[j]; + const W = mi + mj; + const pi = W > 0 ? mi / W : 0.5; const pj = 1 - pi; - const logPi = Math.log(pi); - const logPj = Math.log(pj); + // Merged mean. const mmx = pi * mux + pj * mvx; const mmy = pi * muy + pj * mvy; const mmz = pi * muz + pj * mvz; @@ -125,78 +174,87 @@ const computeEdgeCostView = ( const dix = mux - mmx, diy = muy - mmy, diz = muz - mmz; const djx = mvx - mmx, djy = mvy - mmy, djz = mvz - mmz; - const sigm = scratch.sigm; - for (let a = 0; a < 9; a++) { - sigm[a] = pi * cache.sigma[i9 + a] + pj * cache.sigma[j9 + a]; - } - sigm[0] += pi * dix * dix + pj * djx * djx; - sigm[1] += pi * dix * diy + pj * djx * djy; - sigm[2] += pi * dix * diz + pj * djx * djz; - sigm[3] += pi * diy * dix + pj * djy * djx; - sigm[4] += pi * diy * diy + pj * djy * djy; - sigm[5] += pi * diy * diz + pj * djy * djz; - sigm[6] += pi * diz * dix + pj * djz * djx; - sigm[7] += pi * diz * diy + pj * djz * djy; - sigm[8] += pi * diz * diz + pj * djz * djz; - - sigm[1] = sigm[3] = 0.5 * (sigm[1] + sigm[3]); - sigm[2] = sigm[6] = 0.5 * (sigm[2] + sigm[6]); - sigm[5] = sigm[7] = 0.5 * (sigm[5] + sigm[7]); - sigm[0] += EPS_COV; - sigm[4] += EPS_COV; - sigm[8] += EPS_COV; - - const detm = Math.max(det3(sigm, 0), 1e-30); - const logdetm = Math.log(detm); - - const EpNegLogQ = 0.5 * (3 * LOG2PI + logdetm + 3); - - const stdix = Math.sqrt(Math.max(cache.v[i3], 0)); - const stdiy = Math.sqrt(Math.max(cache.v[i3 + 1], 0)); - const stdiz = Math.sqrt(Math.max(cache.v[i3 + 2], 0)); - const stdjx = Math.sqrt(Math.max(cache.v[j3], 0)); - const stdjy = Math.sqrt(Math.max(cache.v[j3 + 1], 0)); - const stdjz = Math.sqrt(Math.max(cache.v[j3 + 2], 0)); - - let sumLogpOnI = 0; - let sumLogpOnJ = 0; - - for (let s = 0; s < Z.length; s++) { - const z0 = Z[s][0], z1 = Z[s][1], z2 = Z[s][2]; - - const xix = mux + z0 * stdix * cache.R[i9 + 0] + z1 * stdiy * cache.R[i9 + 1] + z2 * stdiz * cache.R[i9 + 2]; - const xiy = muy + z0 * stdix * cache.R[i9 + 3] + z1 * stdiy * cache.R[i9 + 4] + z2 * stdiz * cache.R[i9 + 5]; - const xiz = muz + z0 * stdix * cache.R[i9 + 6] + z1 * stdiy * cache.R[i9 + 7] + z2 * stdiz * cache.R[i9 + 8]; - - const xjx = mvx + z0 * stdjx * cache.R[j9 + 0] + z1 * stdjy * cache.R[j9 + 1] + z2 * stdjz * cache.R[j9 + 2]; - const xjy = mvy + z0 * stdjx * cache.R[j9 + 3] + z1 * stdjy * cache.R[j9 + 4] + z2 * stdjz * cache.R[j9 + 5]; - const xjz = mvz + z0 * stdjx * cache.R[j9 + 6] + z1 * stdjy * cache.R[j9 + 7] + z2 * stdjz * cache.R[j9 + 8]; - - const logNiOnI = gaussLogpdfDiagrot(xix, xiy, xiz, mux, muy, muz, - cache.R, i9, cache.invdiag[i3], cache.invdiag[i3 + 1], cache.invdiag[i3 + 2], cache.logdet[i]); - const logNjOnI = gaussLogpdfDiagrot(xix, xiy, xiz, mvx, mvy, mvz, - cache.R, j9, cache.invdiag[j3], cache.invdiag[j3 + 1], cache.invdiag[j3 + 2], cache.logdet[j]); - sumLogpOnI += logAddExp(logPi + logNiOnI, logPj + logNjOnI); - - const logNiOnJ = gaussLogpdfDiagrot(xjx, xjy, xjz, mux, muy, muz, - cache.R, i9, cache.invdiag[i3], cache.invdiag[i3 + 1], cache.invdiag[i3 + 2], cache.logdet[i]); - const logNjOnJ = gaussLogpdfDiagrot(xjx, xjy, xjz, mvx, mvy, mvz, - cache.R, j9, cache.invdiag[j3], cache.invdiag[j3 + 1], cache.invdiag[j3 + 2], cache.logdet[j]); - sumLogpOnJ += logAddExp(logPi + logNiOnJ, logPj + logNjOnJ); - } - - const Ei = sumLogpOnI / Z.length; - const Ej = sumLogpOnJ / Z.length; - const EpLogp = pi * Ei + pj * Ej; - const geoCost = EpLogp + EpNegLogQ; - - let cSh = 0; - for (let c = 0; c < colorDim; c++) { - const d = color[i * colorDim + c] - color[j * colorDim + c]; - cSh += d * d; + // Merged covariance: Σ pₖ(δₖδₖᵀ + Σₖ) + EPS_COV·I (law of total variance). + const sxx = pi * (dix * dix + sig[i6]) + pj * (djx * djx + sig[j6]) + EPS_COV; + const sxy = pi * (dix * diy + sig[i6 + 1]) + pj * (djx * djy + sig[j6 + 1]); + const sxz = pi * (dix * diz + sig[i6 + 2]) + pj * (djx * djz + sig[j6 + 2]); + const syy = pi * (diy * diy + sig[i6 + 3]) + pj * (djy * djy + sig[j6 + 3]) + EPS_COV; + const syz = pi * (diy * diz + sig[i6 + 4]) + pj * (djy * djz + sig[j6 + 4]); + const szz = pi * (diz * diz + sig[i6 + 5]) + pj * (djz * djz + sig[j6 + 5]) + EPS_COV; + + const detm = Math.max( + sxx * (syy * szz - syz * syz) - sxy * (sxy * szz - syz * sxz) + sxz * (sxy * syz - syy * sxz), + 1e-30 + ); + const sqrtDetM = Math.sqrt(detm); + + // Merged opacity: mass-conserving, capped at 1 (needs merged scales, i.e. + // eigenvalues of Σ_m — Smith's closed form for a symmetric 3×3). + const q = (sxx + syy + szz) / 3; + const p1 = sxy * sxy + sxz * sxz + syz * syz; + let e0: number, e1: number, e2: number; + if (p1 <= 1e-30) { + e0 = sxx; e1 = syy; e2 = szz; + } else { + const p2 = (sxx - q) * (sxx - q) + (syy - q) * (syy - q) + (szz - q) * (szz - q) + 2 * p1; + const p = Math.sqrt(p2 / 6); + const ip = 1 / p; + const b00 = (sxx - q) * ip, b11 = (syy - q) * ip, b22 = (szz - q) * ip; + const b01 = sxy * ip, b02 = sxz * ip, b12 = syz * ip; + const detB = b00 * (b11 * b22 - b12 * b12) - b01 * (b01 * b22 - b12 * b02) + b02 * (b01 * b12 - b11 * b02); + let r = detB / 2; + r = r < -1 ? -1 : (r > 1 ? 1 : r); + const phi = Math.acos(r) / 3; + e0 = q + 2 * p * Math.cos(phi); + e2 = q + 2 * p * Math.cos(phi + 2 * Math.PI / 3); + e1 = 3 * q - e0 - e2; } - - return geoCost + cSh; + const s0 = Math.sqrt(Math.max(e0, 1e-18)); + const s1 = Math.sqrt(Math.max(e1, 1e-18)); + const s2 = Math.sqrt(Math.max(e2, 1e-18)); + const alphaM = Math.min(1, W / Math.max(ellipsoidArea(s0, s1, s2), 1e-30)); + + // Base colours and their dots (merged colour = mass-weighted average). + const bir = base[i3], big = base[i3 + 1], bib = base[i3 + 2]; + const bjr = base[j3], bjg = base[j3 + 1], bjb = base[j3 + 2]; + const bij = bir * bjr + big * bjg + bib * bjb; // base_i · base_j + const bni = baseN2[i]; // base_i · base_i + const bnj = baseN2[j]; // base_j · base_j + const bim = pi * bni + pj * bij; // base_i · base_m + const bjm = pi * bij + pj * bnj; // base_j · base_m + const bnm = pi * pi * bni + 2 * pi * pj * bij + pj * pj * bnj; // base_m · base_m + + const ai = alpha[i], aj = alpha[j], am = alphaM; + const sdi = sqrtDet[i], sdj = sqrtDet[j]; + + // Self terms ⟨G,G⟩ = π^{3/2}·√|Σ|. + const selfI = ai * ai * bni * PI_1_5 * sdi; + const selfJ = aj * aj * bnj * PI_1_5 * sdj; + const selfM = am * am * bnm * PI_1_5 * sqrtDetM; + + // Cross terms (amplitude dot × Gaussian overlap). + const cIJ = crossG(sdi, sdj, + sig[i6] + sig[j6], sig[i6 + 1] + sig[j6 + 1], sig[i6 + 2] + sig[j6 + 2], + sig[i6 + 3] + sig[j6 + 3], sig[i6 + 4] + sig[j6 + 4], sig[i6 + 5] + sig[j6 + 5], + mux - mvx, muy - mvy, muz - mvz); + const cIM = crossG(sdi, sqrtDetM, + sig[i6] + sxx, sig[i6 + 1] + sxy, sig[i6 + 2] + sxz, + sig[i6 + 3] + syy, sig[i6 + 4] + syz, sig[i6 + 5] + szz, + dix, diy, diz); + const cJM = crossG(sdj, sqrtDetM, + sig[j6] + sxx, sig[j6 + 1] + sxy, sig[j6 + 2] + sxz, + sig[j6 + 3] + syy, sig[j6 + 4] + syz, sig[j6 + 5] + szz, + djx, djy, djz); + + const E = selfI + selfJ + selfM + + 2 * ai * aj * bij * cIJ - + 2 * ai * am * bim * cIM - + 2 * aj * am * bjm * cJM; + + // Clamp float-noise negatives (E is a squared norm ≥ 0) but preserve + // NaN/Inf from degenerate inputs so the caller can fail loud. + // |Δbase|² = |b_i|² + |b_j|² − 2·b_i·b_j — no extra loads needed. + return (E < 0 ? 0 : E) + COLOR_WEIGHT * (bni + bnj - 2 * bij); }; -export { buildCostCache, computeEdgeCostView, type CostCache }; +export { buildCostCache, computeEdgeCostView, COLOR_WEIGHT, type CostCache }; diff --git a/src/lib/decimate/priority.ts b/src/lib/decimate/priority.ts index 5d1b27ad..626e0b27 100644 --- a/src/lib/decimate/priority.ts +++ b/src/lib/decimate/priority.ts @@ -3,10 +3,10 @@ import { type GraphicsDevice } from 'playcanvas'; import { buildCostCache, computeEdgeCostView } from './edge-cost-cpu'; import { collectBlock, verifyAndFixKnn, toGlobalNeighbors, KNN_FIXED, type BlockLocals } from './knn-blocks'; import { KNN_SENTINEL } from './knn-core'; -import { createMergeScratch, makeGaussianSamples, sigmoid, ellipsoidArea, type SplatView } from './moment-match'; +import { type SplatView } from './moment-match'; import { type BlockRange, type ResidentPositions } from './partition'; import { type ChunkData, type ChunkDataPool, type ChunkSource } from '../chunk'; -import { APP_CHUNK, GpuEdgeCost, type EdgeCostCache } from '../gpu/gpu-edge-cost'; +import { GpuEdgeCost, SPLAT_STRIDE, type EdgeCostCache } from '../gpu/gpu-edge-cost'; import { GpuKnn } from '../gpu/gpu-knn'; import { WorkerQueue } from '../workers'; @@ -39,6 +39,14 @@ type PriorityContext = { K: number; /** Neighbours per query (16). */ k: number; + /** + * Optional resident splat-cache output for re-costed selection + * (SPLAT_STRIDE floats per gaussian, packGpuCache row layout), filled for + * every owned gaussian. + */ + cacheOut?: Float32Array; + /** Optional resident neighbour-id output (k per gaussian, KNN_SENTINEL padded). */ + neighborsOut?: Uint32Array; }; /** @@ -148,62 +156,26 @@ const indexOfSorted = (sorted: Uint32Array, g: number): number => { return -1; }; -// Pack the block view into the GpuEdgeCost cache layout (legacy packing: -// posScalars 8-wide, rotR from normalized quats, appearance in ≤APP_CHUNK -// column chunks with live-width strides). +// Pack the per-splat cache for the GPU kernel: build the CPU cache once and +// interleave it into the kernel's SPLAT_STRIDE-wide layout (mean from the view, +// the rest from the cache), so the GPU reads byte-identical per-splat inputs. const packGpuCache = (view: SplatView): EdgeCostCache => { - const { pos, geo, color, colorDim } = view; - const n = geo.length / 8; - const posScalars = new Float32Array(n * 8); - const rotR = new Float32Array(n * 9); - const rot = new Float32Array(9); - + const { pos } = view; + const c = buildCostCache(view); + const n = view.geo.length / 8; + const d = new Float32Array(n * SPLAT_STRIDE); for (let i = 0; i < n; i++) { - const i8 = i * 8; - const o = i * 8; - const linAlpha = sigmoid(geo[i8 + 7]); - const sx = Math.max(Math.exp(geo[i8 + 4]), 1e-12); - const sy = Math.max(Math.exp(geo[i8 + 5]), 1e-12); - const sz = Math.max(Math.exp(geo[i8 + 6]), 1e-12); - const vx = sx * sx + 1e-8; - const vy = sy * sy + 1e-8; - const vz = sz * sz + 1e-8; - posScalars[o] = pos[i * 3]; - posScalars[o + 1] = pos[i * 3 + 1]; - posScalars[o + 2] = pos[i * 3 + 2]; - posScalars[o + 3] = linAlpha * ellipsoidArea(sx, sy, sz) + 1e-12; - posScalars[o + 4] = Math.log(Math.max(vx, 1e-30)) + Math.log(Math.max(vy, 1e-30)) + Math.log(Math.max(vz, 1e-30)); - posScalars[o + 5] = vx; - posScalars[o + 6] = vy; - posScalars[o + 7] = vz; - - let qw = geo[i8], qx = geo[i8 + 1], qy = geo[i8 + 2], qz = geo[i8 + 3]; - const invq = 1 / Math.max(Math.hypot(qw, qx, qy, qz), 1e-12); - qw *= invq; qx *= invq; qy *= invq; qz *= invq; - const xx = qx * qx, yy = qy * qy, zz = qz * qz; - const wx = qw * qx, wy = qw * qy, wz = qw * qz; - const xy = qx * qy, xz = qx * qz, yz = qy * qz; - rot[0] = 1 - 2 * (yy + zz); rot[1] = 2 * (xy - wz); rot[2] = 2 * (xz + wy); - rot[3] = 2 * (xy + wz); rot[4] = 1 - 2 * (xx + zz); rot[5] = 2 * (yz - wx); - rot[6] = 2 * (xz - wy); rot[7] = 2 * (yz + wx); rot[8] = 1 - 2 * (xx + yy); - rotR.set(rot, i * 9); - } - - const numChunks = Math.ceil(colorDim / APP_CHUNK); - const appChunks: Float32Array[] = []; - for (let ch = 0; ch < numChunks; ch++) { - const kStart = ch * APP_CHUNK; - const width = Math.min(APP_CHUNK, colorDim - kStart); - const chunk = new Float32Array(n * width); - for (let s = 0; s < n; s++) { - const dst = s * width; - const src = s * colorDim + kStart; - for (let kk = 0; kk < width; kk++) chunk[dst + kk] = color[src + kk]; - } - appChunks.push(chunk); + const o = i * SPLAT_STRIDE, i6 = i * 6, i3 = i * 3; + d[o] = pos[i3]; d[o + 1] = pos[i3 + 1]; d[o + 2] = pos[i3 + 2]; + d[o + 3] = c.sig[i6]; d[o + 4] = c.sig[i6 + 1]; d[o + 5] = c.sig[i6 + 2]; + d[o + 6] = c.sig[i6 + 3]; d[o + 7] = c.sig[i6 + 4]; d[o + 8] = c.sig[i6 + 5]; + d[o + 9] = c.sqrtDet[i]; + d[o + 10] = c.alpha[i]; + d[o + 11] = c.mass[i]; + d[o + 12] = c.base[i3]; d[o + 13] = c.base[i3 + 1]; d[o + 14] = c.base[i3 + 2]; + d[o + 15] = c.baseN2[i]; } - - return { posScalars, rotR, appChunks, numAppCols: colorDim, numSplats: n }; + return { splatData: d, numSplats: n }; }; /** @@ -221,9 +193,6 @@ const runPriorityPass = async ( tick?: (n: number) => void ): Promise => { const { pos, order, blocks, device, K, k } = ctx; - const Z = makeGaussianSamples(1, 0); - const z = new Float32Array([Z[0][0], Z[0][1], Z[0][2]]); - const colorDim = ctx.source.meta.layouts.color!.stride >> 2; let maxOwned = 0; for (const b of blocks) maxOwned = Math.max(maxOwned, b.end - b.start); @@ -258,7 +227,7 @@ const runPriorityPass = async ( try { if (device) { gpuKnn = new GpuKnn(device, maxLocalN, k); - gpuCost = new GpuEdgeCost(device, maxLocalN, maxOwned * k, colorDim); + gpuCost = new GpuEdgeCost(device, maxLocalN, maxOwned * k); } let next: Prepared | null = blocks.length > 0 ? prepare(0) : null; @@ -297,7 +266,7 @@ const runPriorityPass = async ( if (gpuCost && viewN > gpuCostCapacity) { gpuCost.destroy(); gpuCostCapacity = Math.ceil(viewN * 1.1); - gpuCost = new GpuEdgeCost(device!, gpuCostCapacity, maxOwned * k, colorDim); + gpuCost = new GpuEdgeCost(device!, gpuCostCapacity, maxOwned * k); } const { view } = await gatherBlockView(ctx, bi, extraGlobals); @@ -330,13 +299,45 @@ const runPriorityPass = async ( edgeOf[nOwned] = e; const costs = new Float32Array(e); + const packed = device ? packGpuCache(view) : undefined; + const cpuCache = device ? undefined : buildCostCache(view); if (device) { - await gpuCost!.execute(packGpuCache(view), edgeI.subarray(0, e), edgeJ.subarray(0, e), z, costs); + await gpuCost!.execute(packed!, edgeI.subarray(0, e), edgeJ.subarray(0, e), costs); } else { - const cache = buildCostCache(view); - const scratch = createMergeScratch(); for (let i = 0; i < e; i++) { - costs[i] = computeEdgeCostView(view, cache, edgeI[i], edgeJ[i], Z, scratch); + costs[i] = computeEdgeCostView(view, cpuCache!, edgeI[i], edgeJ[i]); + } + } + + // Persist owned rows for re-costed selection: the packed splat + // cache (identical layout on both paths) and the global neighbour + // ids (sentinel-padded). + if (ctx.cacheOut) { + const CO = ctx.cacheOut; + if (packed) { + for (let qi = 0; qi < nOwned; qi++) { + CO.set(packed.splatData.subarray(qi * SPLAT_STRIDE, (qi + 1) * SPLAT_STRIDE), owned[qi] * SPLAT_STRIDE); + } + } else { + const c = cpuCache!; + for (let qi = 0; qi < nOwned; qi++) { + const o = owned[qi] * SPLAT_STRIDE; + const q6 = qi * 6, q3 = qi * 3; + CO[o] = view.pos[q3]; CO[o + 1] = view.pos[q3 + 1]; CO[o + 2] = view.pos[q3 + 2]; + CO[o + 3] = c.sig[q6]; CO[o + 4] = c.sig[q6 + 1]; CO[o + 5] = c.sig[q6 + 2]; + CO[o + 6] = c.sig[q6 + 3]; CO[o + 7] = c.sig[q6 + 4]; CO[o + 8] = c.sig[q6 + 5]; + CO[o + 9] = c.sqrtDet[qi]; + CO[o + 10] = c.alpha[qi]; + CO[o + 11] = c.mass[qi]; + CO[o + 12] = c.base[q3]; CO[o + 13] = c.base[q3 + 1]; CO[o + 14] = c.base[q3 + 2]; + CO[o + 15] = c.baseN2[qi]; + } + } + } + if (ctx.neighborsOut) { + const NO = ctx.neighborsOut; + for (let qi = 0; qi < nOwned; qi++) { + NO.set(nbGlobal.subarray(qi * k, (qi + 1) * k), owned[qi] * k); } } diff --git a/src/lib/decimate/select-recost.ts b/src/lib/decimate/select-recost.ts new file mode 100644 index 00000000..115b6a09 --- /dev/null +++ b/src/lib/decimate/select-recost.ts @@ -0,0 +1,457 @@ +/** + * Re-costed merge selection: exact greedy agglomeration within a generation. + * + * Where {@link selectMerges} consumes the priority pass's pairwise costs + * one-shot (costs go stale as groups form), this selection re-evaluates a + * cluster's candidates after every merge against its CURRENT moments, so the + * cheapest-first order is always true. Validated against the reference + * implementation (tools/decimate-exact.mjs) as the production form of the + * study winner: within-generation exact re-costing closes the remaining + * ~1–2 dB at fine levels vs one-shot selection. + * + * Cost of a cluster = exact field-L2 vs its generation-input members (their + * parameters come from the resident splat cache emitted by the priority + * pass), plus the scale-free DC colour term ({@link COLOR_WEIGHT}) — the same + * cost definition as the pairwise kernel, so the precomputed candidate costs + * seed the heap directly. + * + * Memory: ~200 B per gaussian resident (splat cache, neighbour ids, cluster + * moments, heap). decimate-source gates this path by memory budget and falls + * back to {@link selectMerges} when it does not fit. + * + * Engine-free; pure resident-array computation, no IO. + */ + +import { COLOR_WEIGHT } from './edge-cost-cpu'; +import { KNN_SENTINEL } from './knn-core'; +import { EPS_COV, ellipsoidArea } from './moment-match'; +import { type CandidateArrays } from './priority'; +import { MAX_GROUP, type SelectionResult } from './select'; + +/** Floats per splat in the resident cache (packGpuCache layout). */ +export const CACHE_STRIDE = 16; + +const NO_CANDIDATE = 0xFFFFFFFF; +const NIL = 0xFFFFFFFF; + +/** Skip Gaussian products whose exponent bound exceeds this (e^-60 ≈ 9e-27). */ +const CULL_QUAD = 120; + +const PI_1_5 = Math.PI ** 1.5; +const TWO_PI_1_5 = (2 * Math.PI) ** 1.5; + +/** Inputs for {@link selectMergesRecosted}. */ +type RecostInputs = { + /** Per-gaussian best-K candidates from the priority pass (seed costs include the colour term). */ + cand: CandidateArrays; + /** Candidates per gaussian. */ + K: number; + /** + * Resident per-splat cache, {@link CACHE_STRIDE} floats per splat + * (pos 3, Σ 6 [xx xy xz yy yz zz], √|Σ|, α, mass, base colour 3, |base|²) + * — the packGpuCache row layout, persisted by the priority pass. + */ + splatCache: Float32Array; + /** Global neighbour ids per splat (D per row, KNN_SENTINEL padded). */ + neighbors: Uint32Array; + /** Neighbours per splat. */ + D: number; + /** Gaussian count. */ + N: number; + /** Target removal count for this generation. */ + mergesNeeded: number; +}; + +// ⟨G_a,G_b⟩ scaled by √|Σa|·√|Σb| for M = Σa+Σb (6 comps) and offset d. +const crossG = ( + sdAB: number, + m0: number, m1: number, m2: number, m3: number, m4: number, m5: number, + dx: number, dy: number, dz: number +): number => { + const c00 = m3 * m5 - m4 * m4; + const c01 = m2 * m4 - m1 * m5; + const c02 = m1 * m4 - m2 * m3; + const c11 = m0 * m5 - m2 * m2; + const c12 = m1 * m2 - m0 * m4; + const c22 = m0 * m3 - m1 * m1; + const det = Math.max(m0 * c00 + m1 * c01 + m2 * c02, 1e-60); + const quad = (c00 * dx * dx + c11 * dy * dy + c22 * dz * dz + + 2 * (c01 * dx * dy + c02 * dx * dz + c12 * dy * dz)) / det; + if (!(quad < CULL_QUAD)) return 0; + return TWO_PI_1_5 * sdAB / Math.sqrt(det) * Math.exp(-0.5 * quad); +}; + +// Smith closed-form eigenvalues of a symmetric 3×3 (6 comps), descending. +const eig3 = (m0: number, m1: number, m2: number, m3: number, m4: number, m5: number, out: Float64Array): void => { + const q = (m0 + m3 + m5) / 3; + const p1 = m1 * m1 + m2 * m2 + m4 * m4; + if (p1 <= 1e-30) { + out[0] = Math.max(m0, m3, m5); + out[2] = Math.min(m0, m3, m5); + out[1] = m0 + m3 + m5 - out[0] - out[2]; + return; + } + const p2 = (m0 - q) * (m0 - q) + (m3 - q) * (m3 - q) + (m5 - q) * (m5 - q) + 2 * p1; + const p = Math.sqrt(p2 / 6); + const ip = 1 / p; + const b00 = (m0 - q) * ip, b11 = (m3 - q) * ip, b22 = (m5 - q) * ip; + const b01 = m1 * ip, b02 = m2 * ip, b12 = m4 * ip; + const detB = b00 * (b11 * b22 - b12 * b12) - b01 * (b01 * b22 - b12 * b02) + b02 * (b01 * b12 - b11 * b02); + let r = detB / 2; + r = r < -1 ? -1 : (r > 1 ? 1 : r); + const phi = Math.acos(r) / 3; + const e0 = q + 2 * p * Math.cos(phi); + const e2 = q + 2 * p * Math.cos(phi + (2 * Math.PI) / 3); + out[0] = e0; out[1] = 3 * q - e0 - e2; out[2] = e2; +}; + +/** + * Select merges with exact within-generation re-costing. + * + * @param inputs - See {@link RecostInputs}. + * @returns The selection (same contract as {@link selectMerges}). + */ +const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { + const { cand, K, splatCache: SC, neighbors, D, N, mergesNeeded } = inputs; + + // ---- Cluster state (indexed by union-find root). + const parent = new Uint32Array(N); + const size = new Uint32Array(N).fill(1); + const W = new Float64Array(N); + const mx = new Float64Array(N), my = new Float64Array(N), mz = new Float64Array(N); + const M2 = new Float64Array(N * 6); + const baseW = new Float64Array(N * 3); + const Sself = new Float64Array(N); + const Err = new Float64Array(N); + const version = new Uint32Array(N).fill(1); + const lastSeq = new Uint32Array(N); + const mHead = new Uint32Array(N), mTail = new Uint32Array(N); + const mNext = new Uint32Array(N).fill(NIL); + const stamp = new Uint32Array(N); + let stampGen = 0; + let seqCounter = 0; + let liveCount = N; + + for (let i = 0; i < N; i++) { + parent[i] = i; mHead[i] = i; mTail[i] = i; + const o = i * CACHE_STRIDE, i6 = i * 6, i3 = i * 3; + const mass = SC[o + 11]; + W[i] = mass; + mx[i] = SC[o]; my[i] = SC[o + 1]; mz[i] = SC[o + 2]; + // Cache Σ carries EPS on its diagonal; moments accumulate Σ without it + // (production mergeGroup member math). + M2[i6] = mass * (SC[o + 3] - EPS_COV); + M2[i6 + 1] = mass * SC[o + 4]; + M2[i6 + 2] = mass * SC[o + 5]; + M2[i6 + 3] = mass * (SC[o + 6] - EPS_COV); + M2[i6 + 4] = mass * SC[o + 7]; + M2[i6 + 5] = mass * (SC[o + 8] - EPS_COV); + baseW[i3] = mass * SC[o + 12]; + baseW[i3 + 1] = mass * SC[o + 13]; + baseW[i3 + 2] = mass * SC[o + 14]; + Sself[i] = SC[o + 10] * SC[o + 10] * SC[o + 15] * PI_1_5 * SC[o + 9]; + Err[i] = 0; + } + + const find = (x: number): number => { + while (parent[x] !== x) { + parent[x] = parent[parent[x]]; + x = parent[x]; + } + return x; + }; + + // ---- Exact merge cost: ΔE of A∪B vs generation-input members + colour term. + const eigOut = new Float64Array(3); + let gatherBuf = new Uint32Array(1 << 12); + const gatherMembers = (root: number): number => { + let cnt = 0; + for (let m = mHead[root]; m !== NIL; m = mNext[m]) { + if (cnt === gatherBuf.length) { + const g = new Uint32Array(gatherBuf.length * 2); + g.set(gatherBuf); gatherBuf = g; + } + gatherBuf[cnt++] = m; + } + return cnt; + }; + + const evalOut = { E: 0, Scross: 0 }; + const evalMerge = (A: number, B: number): number => { + const WA = W[A], WB = W[B], WC = WA + WB; + const iw = 1 / WC; + const mcx = (WA * mx[A] + WB * mx[B]) * iw; + const mcy = (WA * my[A] + WB * my[B]) * iw; + const mcz = (WA * mz[A] + WB * mz[B]) * iw; + const dax = mx[A] - mcx, day = my[A] - mcy, daz = mz[A] - mcz; + const dbx = mx[B] - mcx, dby = my[B] - mcy, dbz = mz[B] - mcz; + const a6 = A * 6, b6 = B * 6; + + const sm0 = (M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx) * iw + EPS_COV; + const sm1 = (M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby) * iw; + const sm2 = (M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz) * iw; + const sm3 = (M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby) * iw + EPS_COV; + const sm4 = (M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz) * iw; + const sm5 = (M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz) * iw + EPS_COV; + + const detm = Math.max( + sm0 * (sm3 * sm5 - sm4 * sm4) - sm1 * (sm1 * sm5 - sm4 * sm2) + sm2 * (sm1 * sm4 - sm3 * sm2), + 1e-60 + ); + const sdC = Math.sqrt(detm); + + eig3(sm0, sm1, sm2, sm3, sm4, sm5, eigOut); + const s0 = Math.sqrt(Math.max(eigOut[0], 1e-18)); + const s1 = Math.sqrt(Math.max(eigOut[1], 1e-18)); + const s2 = Math.sqrt(Math.max(eigOut[2], 1e-18)); + const alphaC = Math.min(1, WC / Math.max(ellipsoidArea(s0, s1, s2), 1e-30)); + + const a3 = A * 3, b3 = B * 3; + const bc0 = (baseW[a3] + baseW[b3]) * iw; + const bc1 = (baseW[a3 + 1] + baseW[b3 + 1]) * iw; + const bc2 = (baseW[a3 + 2] + baseW[b3 + 2]) * iw; + const bn2C = bc0 * bc0 + bc1 * bc1 + bc2 * bc2; + + const selfM = alphaC * alphaC * bn2C * PI_1_5 * sdC; + + // ⟨Σ member fields, f_m⟩ over both chains. + let memfm = 0; + for (let pass = 0; pass < 2; pass++) { + for (let m = pass === 0 ? mHead[A] : mHead[B]; m !== NIL; m = mNext[m]) { + const o = m * CACHE_STRIDE; + const wgt = SC[o + 10] * alphaC * + (SC[o + 12] * bc0 + SC[o + 13] * bc1 + SC[o + 14] * bc2); + if (wgt === 0) continue; + memfm += wgt * crossG(SC[o + 9] * sdC, + SC[o + 3] + sm0, SC[o + 4] + sm1, SC[o + 5] + sm2, + SC[o + 6] + sm3, SC[o + 7] + sm4, SC[o + 8] + sm5, + SC[o] - mcx, SC[o + 1] - mcy, SC[o + 2] - mcz); + } + } + + // Scross(A,B) = Σ_{a∈A,b∈B}⟨f_a,f_b⟩ with distance culling. + const na = gatherMembers(A); + let scross = 0; + for (let b = mHead[B]; b !== NIL; b = mNext[b]) { + const ob = b * CACHE_STRIDE; + const bxp = SC[ob], byp = SC[ob + 1], bzp = SC[ob + 2]; + const trb = SC[ob + 3] + SC[ob + 6] + SC[ob + 8]; + const alb = SC[ob + 10], sdb = SC[ob + 9]; + const cb0 = SC[ob + 12], cb1 = SC[ob + 13], cb2 = SC[ob + 14]; + for (let t = 0; t < na; t++) { + const a = gatherBuf[t]; + const oa = a * CACHE_STRIDE; + const dx = SC[oa] - bxp, dy = SC[oa + 1] - byp, dz = SC[oa + 2] - bzp; + const d2 = dx * dx + dy * dy + dz * dz; + if (d2 > CULL_QUAD * (SC[oa + 3] + SC[oa + 6] + SC[oa + 8] + trb)) continue; + const wgt = SC[oa + 10] * alb * + (SC[oa + 12] * cb0 + SC[oa + 13] * cb1 + SC[oa + 14] * cb2); + if (wgt === 0) continue; + scross += wgt * crossG(SC[oa + 9] * sdb, + SC[oa + 3] + SC[ob + 3], SC[oa + 4] + SC[ob + 4], SC[oa + 5] + SC[ob + 5], + SC[oa + 6] + SC[ob + 6], SC[oa + 7] + SC[ob + 7], SC[oa + 8] + SC[ob + 8], + dx, dy, dz); + } + } + + const E = Sself[A] + Sself[B] + 2 * scross - 2 * memfm + selfM; + evalOut.E = E; + evalOut.Scross = scross; + + // Scale-free colour term between the cluster mean base colours (same + // definition as the pairwise kernel, so seed costs are consistent). + const iwA = 1 / W[A], iwB = 1 / W[B]; + const d0 = baseW[a3] * iwA - baseW[b3] * iwB; + const d1 = baseW[a3 + 1] * iwA - baseW[b3 + 1] * iwB; + const d2c = baseW[a3 + 2] * iwA - baseW[b3 + 2] * iwB; + return (E - Err[A] - Err[B]) + COLOR_WEIGHT * (d0 * d0 + d1 * d1 + d2c * d2c); + }; + + // ---- Lazy min-heap of candidate edges (one live entry per cluster). + let heapCap = Math.ceil(N * 1.25) + 16; + let hCost = new Float64Array(heapCap); + let hA = new Uint32Array(heapCap), hB = new Uint32Array(heapCap); + let hSeq = new Uint32Array(heapCap), hVb = new Uint32Array(heapCap); + let heapSize = 0; + + const swap = (i: number, j: number): void => { + let t; + t = hCost[i]; hCost[i] = hCost[j]; hCost[j] = t; + t = hA[i]; hA[i] = hA[j]; hA[j] = t; + t = hB[i]; hB[i] = hB[j]; hB[j] = t; + t = hSeq[i]; hSeq[i] = hSeq[j]; hSeq[j] = t; + t = hVb[i]; hVb[i] = hVb[j]; hVb[j] = t; + }; + + const heapPush = (cost: number, a: number, b: number, seq: number, vb: number): void => { + if (heapSize === heapCap) { + const nc = heapCap * 2; + const c2 = new Float64Array(nc); c2.set(hCost); hCost = c2; + const g = (old: Uint32Array) => { + const x = new Uint32Array(nc); x.set(old); return x; + }; + hA = g(hA); hB = g(hB); hSeq = g(hSeq); hVb = g(hVb); + heapCap = nc; + } + let i = heapSize++; + hCost[i] = cost; hA[i] = a; hB[i] = b; hSeq[i] = seq; hVb[i] = vb; + while (i > 0) { + const p = (i - 1) >> 1; + if (hCost[p] <= hCost[i]) break; + swap(i, p); i = p; + } + }; + + const popOut = { cost: 0, a: 0, b: 0, seq: 0, vb: 0 }; + const heapPop = (): boolean => { + if (heapSize === 0) return false; + popOut.cost = hCost[0]; popOut.a = hA[0]; popOut.b = hB[0]; popOut.seq = hSeq[0]; popOut.vb = hVb[0]; + heapSize--; + if (heapSize > 0) { + hCost[0] = hCost[heapSize]; hA[0] = hA[heapSize]; hB[0] = hB[heapSize]; + hSeq[0] = hSeq[heapSize]; hVb[0] = hVb[heapSize]; + let i = 0; + for (;;) { + const l = 2 * i + 1, r = l + 1; + let m = i; + if (l < heapSize && hCost[l] < hCost[m]) m = l; + if (r < heapSize && hCost[r] < hCost[m]) m = r; + if (m === i) break; + swap(i, m); i = m; + } + } + return true; + }; + + // Candidate derivation: live clusters owning any neighbour of any member. + let candBuf = new Uint32Array(1 << 12); + const deriveCandidates = (root: number): number => { + stampGen++; + const gen = stampGen; + let cnt = 0; + for (let m = mHead[root]; m !== NIL; m = mNext[m]) { + const base = m * D; + for (let s = 0; s < D; s++) { + const nb = neighbors[base + s]; + if (nb === KNN_SENTINEL) continue; + const r = find(nb); + if (r === root || stamp[r] === gen) continue; + stamp[r] = gen; + if (cnt === candBuf.length) { + const g = new Uint32Array(candBuf.length * 2); + g.set(candBuf); candBuf = g; + } + candBuf[cnt++] = r; + } + } + return cnt; + }; + + const pushBestEdge = (root: number): void => { + const cnt = deriveCandidates(root); + lastSeq[root] = ++seqCounter; + if (cnt === 0) return; + const sz = size[root]; + let bc = Infinity, bp = -1, bv = 0; + for (let t = 0; t < cnt; t++) { + const c = candBuf[t]; + if (sz + size[c] > MAX_GROUP) continue; + const d = evalMerge(root, c); + if (d < bc) { + bc = d; bp = c; bv = version[c]; + } + } + if (bp >= 0) heapPush(bc, root, bp, lastSeq[root], bv); + }; + + // Seed: the priority pass's cheapest candidate per gaussian (same cost + // definition, already sorted ascending — no re-evaluation needed). + for (let i = 0; i < N; i++) { + const j = cand.idx[i * K]; + const c = cand.cost[i * K]; + lastSeq[i] = ++seqCounter; + if (j !== NO_CANDIDATE && Number.isFinite(c)) { + heapPush(c, i, j, lastSeq[i], version[j]); + } + } + + // ---- Greedy loop. + let removed = 0; + while (removed < mergesNeeded) { + if (!heapPop()) break; + const a = popOut.a; + if (parent[a] !== a || popOut.seq !== lastSeq[a]) continue; + const b = popOut.b; + if (parent[b] !== b || version[b] !== popOut.vb || size[a] + size[b] > MAX_GROUP) { + pushBestEdge(a); + continue; + } + + // Commit a+b (exact recompute for the error bookkeeping). + evalMerge(a, b); + const E = evalOut.E, scross = evalOut.Scross; + const keep = size[a] >= size[b] ? a : b; + const lose = keep === a ? b : a; + + const WA = W[a], WB = W[b], WC = WA + WB; + const iw = 1 / WC; + const mcx = (WA * mx[a] + WB * mx[b]) * iw; + const mcy = (WA * my[a] + WB * my[b]) * iw; + const mcz = (WA * mz[a] + WB * mz[b]) * iw; + const dax = mx[a] - mcx, day = my[a] - mcy, daz = mz[a] - mcz; + const dbx = mx[b] - mcx, dby = my[b] - mcy, dbz = mz[b] - mcz; + const a6 = a * 6, b6 = b * 6, k6 = keep * 6; + const n0 = M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx; + const n1 = M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby; + const n2 = M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz; + const n3 = M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby; + const n4 = M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz; + const n5 = M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz; + M2[k6] = n0; M2[k6 + 1] = n1; M2[k6 + 2] = n2; M2[k6 + 3] = n3; M2[k6 + 4] = n4; M2[k6 + 5] = n5; + W[keep] = WC; mx[keep] = mcx; my[keep] = mcy; mz[keep] = mcz; + const ka = keep * 3, aa = a * 3, bb = b * 3; + const bw0 = baseW[aa] + baseW[bb], bw1 = baseW[aa + 1] + baseW[bb + 1], bw2 = baseW[aa + 2] + baseW[bb + 2]; + baseW[ka] = bw0; baseW[ka + 1] = bw1; baseW[ka + 2] = bw2; + Sself[keep] = Sself[a] + Sself[b] + 2 * scross; + Err[keep] = E; + size[keep] += size[lose]; + mNext[mTail[keep]] = mHead[lose]; + mTail[keep] = mTail[lose]; + parent[lose] = keep; + version[keep]++; + liveCount--; + removed++; + + pushBestEdge(keep); + } + + // ---- CSR assembly (identical contract to selectMerges). + const memberGroup = new Int32Array(N).fill(-1); + const rootGroup = new Int32Array(N).fill(-1); + let G = 0; + for (let i = 0; i < N; i++) { + const r = find(i); + if (size[r] > 1) { + if (rootGroup[r] < 0) rootGroup[r] = G++; + memberGroup[i] = rootGroup[r]; + } + } + const groupOffsets = new Uint32Array(G + 1); + for (let i = 0; i < N; i++) { + const g = memberGroup[i]; + if (g >= 0) groupOffsets[g + 1]++; + } + for (let g = 0; g < G; g++) groupOffsets[g + 1] += groupOffsets[g]; + const groupMembers = new Uint32Array(groupOffsets[G]); + const fill = groupOffsets.slice(0, G); + for (let i = 0; i < N; i++) { + const g = memberGroup[i]; + if (g >= 0) groupMembers[fill[g]++] = i; + } + const groupMin = new Uint32Array(G); + for (let g = 0; g < G; g++) groupMin[g] = groupMembers[groupOffsets[g]]; + + return { groupOffsets, groupMembers, memberGroup, groupMin, mergedGroups: G, removed }; +}; + +export { selectMergesRecosted, type RecostInputs }; diff --git a/src/lib/decimate/select.ts b/src/lib/decimate/select.ts index f7cfc72e..8fcc8cf5 100644 --- a/src/lib/decimate/select.ts +++ b/src/lib/decimate/select.ts @@ -1,19 +1,28 @@ /** - * Global merge selection over the resident candidate arrays: bucketed greedy - * disjoint matching in cost order, plus chain closure that attaches - * still-unmatched gaussians to a candidate's group (≤3, relief cap 4) so a - * 50% target completes in one generation instead of a mop-up pass. + * Global merge selection over the resident candidate arrays. * - * Bucket walk ≈ exact cost-sorted greedy up to 1/SELECT_BUCKETS-of-range - * quantization — the selection-semantics match with the legacy algorithm. + * PROTOTYPE — cost-ordered agglomeration (replaces the 50%-matching + chain + * closure). Candidate edges are walked cheapest-first and unioned (union-find, + * capped at {@link MAX_GROUP} members) until the generation's target removal + * count is reached. Because a single-linkage cluster can absorb many members, + * cheap (redundant) regions collapse deeply while expensive (large / distinct) + * gaussians stay as singletons — so removal follows local error instead of a + * uniform 50% everywhere. Combined with the L2 field cost, large distant + * gaussians (sky) survive far longer than tiny detail. + * + * Cost ordering is log-scaled bucketed (the L2 cost spans many orders of + * magnitude); ties within a bucket resolve in arbitrary order. * * Engine-free; pure resident-array computation, no IO. */ import { type CandidateArrays } from './priority'; -/** Cost-histogram buckets for the ordered greedy walk. */ -const SELECT_BUCKETS = 1024; +/** Log-scaled cost buckets for the ordered agglomeration walk. */ +const COST_BUCKETS = 4096; + +/** Max gaussians merged into one group in a single generation. */ +const MAX_GROUP = 4; const NO_CANDIDATE = 0xFFFFFFFF; @@ -49,106 +58,104 @@ type SelectionResult = { const selectMerges = (cand: CandidateArrays, N: number, K: number, mergesNeeded: number): SelectionResult => { const E = N * K; - // Pass 1: finite cost range. - let lo = Infinity, hi = -Infinity; + // Finite candidate range (costs are ≥ 0; log-scale the positive ones so + // cost ordering keeps resolution across the metric's huge dynamic range). + let hi = 0, loPos = Infinity, anyFinite = false; for (let e = 0; e < E; e++) { + if (cand.idx[e] === NO_CANDIDATE) continue; const c = cand.cost[e]; - if (Number.isFinite(c)) { - if (c < lo) lo = c; - if (c > hi) hi = c; - } + if (!Number.isFinite(c)) continue; + anyFinite = true; + if (c > hi) hi = c; + if (c > 0 && c < loPos) loPos = c; } - const span = hi > lo ? hi - lo : 1; - const bucketOf = (c: number) => Math.min(SELECT_BUCKETS - 1, Math.floor(((c - lo) / span) * SELECT_BUCKETS)); + const useLog = anyFinite && Number.isFinite(loPos) && loPos < hi; + const logLo = useLog ? Math.log(loPos) : 0; + const logSpan = useLog ? Math.log(hi) - logLo : 1; + const bucketOf = (c: number): number => { + if (!(c > 0) || !useLog) return 0; + const b = 1 + Math.floor(((Math.log(c) - logLo) / logSpan) * (COST_BUCKETS - 2)); + return b < 0 ? 0 : (b >= COST_BUCKETS ? COST_BUCKETS - 1 : b); + }; - // Counting sort of finite candidate entries by bucket. - const counts = new Uint32Array(SELECT_BUCKETS + 1); + // Counting sort of finite candidate edges by bucket (cheapest first). + const counts = new Uint32Array(COST_BUCKETS + 1); for (let e = 0; e < E; e++) { - if (Number.isFinite(cand.cost[e])) counts[bucketOf(cand.cost[e]) + 1]++; + if (cand.idx[e] === NO_CANDIDATE) continue; + const c = cand.cost[e]; + if (Number.isFinite(c)) counts[bucketOf(c) + 1]++; } - for (let b = 0; b < SELECT_BUCKETS; b++) counts[b + 1] += counts[b]; - const orderE = new Uint32Array(counts[SELECT_BUCKETS]); - const cursor = counts.slice(0, SELECT_BUCKETS); + for (let b = 0; b < COST_BUCKETS; b++) counts[b + 1] += counts[b]; + const orderE = new Uint32Array(counts[COST_BUCKETS]); + const cursor = counts.slice(0, COST_BUCKETS); for (let e = 0; e < E; e++) { - if (Number.isFinite(cand.cost[e])) orderE[cursor[bucketOf(cand.cost[e])]++] = e; + if (cand.idx[e] === NO_CANDIDATE) continue; + const c = cand.cost[e]; + if (Number.isFinite(c)) orderE[cursor[bucketOf(c)]++] = e; } - const memberGroup = new Int32Array(N).fill(-1); + // Union-find (union by size + path halving). + const parent = new Uint32Array(N); + for (let i = 0; i < N; i++) parent[i] = i; + const size = new Uint32Array(N).fill(1); + const find = (x: number): number => { + while (parent[x] !== x) { + parent[x] = parent[parent[x]]; + x = parent[x]; + } + return x; + }; - // Primary greedy: pair free endpoints, cheapest bucket first. - const maxPairs = Math.max(0, mergesNeeded); - const pairA = new Uint32Array(maxPairs); - const pairB = new Uint32Array(maxPairs); - let pairs = 0; + // Cost-ordered agglomeration: union cheapest edges first, capping group + // size, until the target removal count is reached. let removed = 0; for (let t = 0; t < orderE.length && removed < mergesNeeded; t++) { const e = orderE[t]; const j = cand.idx[e]; if (j === NO_CANDIDATE) continue; const i = (e / K) | 0; - if (memberGroup[i] !== -1 || memberGroup[j] !== -1) continue; - memberGroup[i] = pairs; - memberGroup[j] = pairs; - pairA[pairs] = i; - pairB[pairs] = j; - pairs++; + let ri = find(i), rj = find(j); + if (ri === rj) continue; + if (size[ri] + size[rj] > MAX_GROUP) continue; + if (size[ri] < size[rj]) { + const tmp = ri; ri = rj; rj = tmp; + } + parent[rj] = ri; + size[ri] += size[rj]; removed++; } - // Chain closure: attach unmatched gaussians to a candidate's group, - // cheapest first, cap 3 — then a relief walk at cap 4. After the primary - // walk every free gaussian's candidates are all matched (else the pair - // would have been taken), so closure can almost always attach. - const groupSize = new Uint32Array(pairs).fill(2); - const joinMember = new Uint32Array(Math.max(0, mergesNeeded - removed)); - const joinGroup = new Uint32Array(joinMember.length); - let joins = 0; - for (const cap of [3, 4]) { - if (removed >= mergesNeeded) break; - for (let t = 0; t < orderE.length && removed < mergesNeeded; t++) { - const e = orderE[t]; - const j = cand.idx[e]; - if (j === NO_CANDIDATE) continue; - const i = (e / K) | 0; - if (memberGroup[i] !== -1) continue; - const g = memberGroup[j]; - if (g === -1 || groupSize[g] >= cap) continue; - memberGroup[i] = g; - groupSize[g]++; - joinMember[joins] = i; - joinGroup[joins] = g; - joins++; - removed++; + // Assemble CSR groups from the forest (roots with size > 1). + const memberGroup = new Int32Array(N).fill(-1); + const rootGroup = new Int32Array(N).fill(-1); + let G = 0; + for (let i = 0; i < N; i++) { + const r = find(i); + if (size[r] > 1) { + if (rootGroup[r] < 0) rootGroup[r] = G++; + memberGroup[i] = rootGroup[r]; } } - // CSR assembly. - const G = pairs; const groupOffsets = new Uint32Array(G + 1); - for (let g = 0; g < G; g++) groupOffsets[g + 1] = groupOffsets[g] + groupSize[g]; - const groupMembers = new Uint32Array(groupOffsets[G]); - const fill = new Uint32Array(G); - for (let g = 0; g < G; g++) { - const o = groupOffsets[g]; - groupMembers[o] = pairA[g]; - groupMembers[o + 1] = pairB[g]; - fill[g] = 2; + for (let i = 0; i < N; i++) { + const g = memberGroup[i]; + if (g >= 0) groupOffsets[g + 1]++; } - for (let t = 0; t < joins; t++) { - const g = joinGroup[t]; - groupMembers[groupOffsets[g] + fill[g]] = joinMember[t]; - fill[g]++; + for (let g = 0; g < G; g++) groupOffsets[g + 1] += groupOffsets[g]; + const groupMembers = new Uint32Array(groupOffsets[G]); + const fill = groupOffsets.slice(0, G); + for (let i = 0; i < N; i++) { + const g = memberGroup[i]; + if (g >= 0) groupMembers[fill[g]++] = i; } + + // Members were appended in ascending id order, so each group's first slot + // is its minimum id. const groupMin = new Uint32Array(G); - for (let g = 0; g < G; g++) { - let min = groupMembers[groupOffsets[g]]; - for (let m = groupOffsets[g] + 1; m < groupOffsets[g + 1]; m++) { - if (groupMembers[m] < min) min = groupMembers[m]; - } - groupMin[g] = min; - } + for (let g = 0; g < G; g++) groupMin[g] = groupMembers[groupOffsets[g]]; return { groupOffsets, groupMembers, memberGroup, groupMin, mergedGroups: G, removed }; }; -export { selectMerges, SELECT_BUCKETS, type SelectionResult }; +export { selectMerges, MAX_GROUP, type SelectionResult }; diff --git a/src/lib/gpu/gpu-edge-cost.ts b/src/lib/gpu/gpu-edge-cost.ts index 46928781..0b108c2f 100644 --- a/src/lib/gpu/gpu-edge-cost.ts +++ b/src/lib/gpu/gpu-edge-cost.ts @@ -3,7 +3,6 @@ import { BUFFERUSAGE_COPY_SRC, SHADERLANGUAGE_WGSL, SHADERSTAGE_COMPUTE, - UNIFORMTYPE_FLOAT, UNIFORMTYPE_UINT, BindGroupFormat, BindStorageBufferFormat, @@ -16,116 +15,75 @@ import { UniformFormat } from 'playcanvas'; -/** - * Appearance columns per storage chunk. The kernel exposes three appearance - * bindings (appA/appB/appC), so the layout holds up to 3·APP_CHUNK columns; at - * 16 the widest chunk reaches the ~2 GB per-binding limit around ~33.5M splats. - * The CPU-side packing in `decimate/priority.ts` imports this same constant, - * so the kernel strides and the host packing can't drift. - */ -export const APP_CHUNK = 16; +/** Per-splat interleaved stride in the `splat` storage buffer (see EdgeCostCache). */ +export const SPLAT_STRIDE = 16; /** - * WGSL kernel: per-edge KL-style cost (matches `computeEdgeCostView` in + * WGSL kernel: per-edge L2 field-error cost (mirrors `computeEdgeCostView` in * `decimate/edge-cost-cpu.ts`). * - * Each thread = one edge (i, j). Reads the per-splat cache for both - * endpoints, computes the merged Gaussian's covariance + determinant, - * runs a single Monte-Carlo sample through both component gaussians - * (the same `z` for both components, matching the CPU implementation), - * and adds an L2 distance over the appearance (SH) coefficients. + * Each thread = one edge (i, j). It reads the per-splat cache for both + * endpoints, moment-matches the merged Gaussian m (mean, covariance, opacity), + * and evaluates + * E = ‖αᵢcᵢGᵢ + αⱼcⱼGⱼ − αₘcₘGₘ‖² + * in closed form via Gaussian–Gaussian L2 products ⟨G_a,G_b⟩. No Monte-Carlo + * sampling and no appearance loop — the amplitude is α × base colour, so only + * the packed base colour (3) and covariance (6) per splat are needed. * - * @param strideA - Live column count of appearance chunk A (0 if unused). - * @param strideB - Live column count of appearance chunk B (0 if unused). - * @param strideC - Live column count of appearance chunk C (0 if unused). * @returns WGSL source. */ -const edgeCostWgsl = (strideA: number, strideB: number, strideC: number) => /* wgsl */` +const edgeCostWgsl = () => /* wgsl */` struct Uniforms { edgeCount: u32, - z0: f32, - z1: f32, - z2: f32, } @group(0) @binding(0) var uniforms: Uniforms; -// Edge list for the current dispatch batch only, split into two parallel -// arrays (avoids a host-side (i, j) interleave). The host uploads each batch's -// slice to offset 0, so we index edgesI/J[bid] directly — keeping these -// buffers batch-sized instead of N·k keeps them off the per-binding limit. +// Edge list for the current dispatch batch (host uploads each batch to offset 0). @group(0) @binding(1) var edgesI: array; @group(0) @binding(2) var edgesJ: array; -// Per-splat geometry, interleaved 8-wide: -// posScalars[8s + 0..2] = position xyz -// posScalars[8s + 3] = mass -// posScalars[8s + 4] = logdet -// posScalars[8s + 5..7] = variances (vx, vy, vz) -@group(0) @binding(3) var posScalars: array; -// Row-major 3x3 rotation matrix per splat (9 floats per splat). -@group(0) @binding(4) var rotR: array; -// Appearance, split into up to three chunks (≤16 columns each) so no single -// binding exceeds maxStorageBufferBindingSize (~2 GB). Each chunk's stride is -// its live column count (STRIDE_A/B/C below); appA holds columns 0.., appB the -// next span, appC the next. Unused chunks have stride 0, are bound to a dummy -// buffer, and are never read. -@group(0) @binding(5) var appA: array; -@group(0) @binding(6) var appB: array; -@group(0) @binding(7) var appC: array; +// Per-splat cache, interleaved ${SPLAT_STRIDE}-wide: +// [0..2] mean xyz +// [3..8] covariance Σ (xx, xy, xz, yy, yz, zz) +// [9] sqrtDet = √|Σ| +// [10] alpha (linear opacity) +// [11] mass (area·α merge weight) +// [12..14] base colour (r, g, b) +// [15] |base|² +@group(0) @binding(3) var splat: array; // Output: cost per edge. -@group(0) @binding(8) var costs: array; +@group(0) @binding(4) var costs: array; const EPS_COV: f32 = 1e-8; -const LOG2PI: f32 = 1.8378770664093453; -// Per-chunk appearance strides = live column count in each chunk (0 = unused, -// dummy-bound). Baked here because the column count is fixed for the lifetime -// of the kernel, so loop bounds and indexing resolve statically. -const STRIDE_A: u32 = ${strideA}u; -const STRIDE_B: u32 = ${strideB}u; -const STRIDE_C: u32 = ${strideC}u; - -// Symmetric 3x3 covariance helpers — we pass them around as 6 f32 (xx, xy, xz, yy, yz, zz). - -// Σ = R · diag(v) · R^T for row-major R (a 9-float array starting at offset r9). -// Variances v come from posScalars[s8 + 5..7]. Result is 6 floats: -// (xx, xy, xz, yy, yz, zz). -fn sigmaFromRotVar(r9: u32, s8: u32) -> array { - let r00 = rotR[r9 + 0u]; let r01 = rotR[r9 + 1u]; let r02 = rotR[r9 + 2u]; - let r10 = rotR[r9 + 3u]; let r11 = rotR[r9 + 4u]; let r12 = rotR[r9 + 5u]; - let r20 = rotR[r9 + 6u]; let r21 = rotR[r9 + 7u]; let r22 = rotR[r9 + 8u]; - let vx = posScalars[s8 + 5u]; - let vy = posScalars[s8 + 6u]; - let vz = posScalars[s8 + 7u]; - return array( - r00*r00*vx + r01*r01*vy + r02*r02*vz, // xx - r00*r10*vx + r01*r11*vy + r02*r12*vz, // xy - r00*r20*vx + r01*r21*vy + r02*r22*vz, // xz - r10*r10*vx + r11*r11*vy + r12*r12*vz, // yy - r10*r20*vx + r11*r21*vy + r12*r22*vz, // yz - r20*r20*vx + r21*r21*vy + r22*r22*vz // zz - ); +const PI_1_5: f32 = 5.5683279968317084; // π^{3/2} +const TWO_PI_1_5: f32 = 15.749609945722419; // (2π)^{3/2} +const TWO_PI_3: f32 = 2.0943951023931953; // 2π/3 +const ELLIP_P: f32 = 1.6075; +// Scale-free DC colour dissimilarity weight (4π·1e-6 in base-colour space — +// see COLOR_WEIGHT in decimate/edge-cost-cpu.ts, mirrored here). +const COLOR_WEIGHT: f32 = 1.2566370614359172e-5; + +// Knud Thomsen ellipsoid surface area (matches CPU ellipsoidArea). +fn ellipsoidArea(sx: f32, sy: f32, sz: f32) -> f32 { + let a = pow(sx * sy, ELLIP_P); + let b = pow(sx * sz, ELLIP_P); + let c = pow(sy * sz, ELLIP_P); + return 4.0 * 3.141592653589793 * pow((a + b + c) / 3.0, 1.0 / ELLIP_P); } -// log N(x | mu, R · diag(v) · R^T) for a diagonally-decomposed covariance. -// invDiag is (1/vx, 1/vy, 1/vz); ld is logdet of the full covariance. -// Evaluates y = R^T * (x - mu) using columns of R (= rows of Rt). -fn gaussLogpdfDiagrot( - x: vec3f, mu: vec3f, r9: u32, - invDiag: vec3f, ld: f32 -) -> f32 { - let dx = x.x - mu.x; - let dy = x.y - mu.y; - let dz = x.z - mu.z; - // y = R^T · d. R is row-major; column k of R is (R[k], R[k+3], R[k+6]). - let y0 = dx * rotR[r9 + 0u] + dy * rotR[r9 + 3u] + dz * rotR[r9 + 6u]; - let y1 = dx * rotR[r9 + 1u] + dy * rotR[r9 + 4u] + dz * rotR[r9 + 7u]; - let y2 = dx * rotR[r9 + 2u] + dy * rotR[r9 + 5u] + dz * rotR[r9 + 8u]; - let quad = y0*y0*invDiag.x + y1*y1*invDiag.y + y2*y2*invDiag.z; - return -0.5 * (3.0 * LOG2PI + ld + quad); -} - -fn logAddExp(a: f32, b: f32) -> f32 { - let m = max(a, b); - return m + log(exp(a - m) + exp(b - m)); +// Gaussian cross-product ⟨G_a,G_b⟩ for symmetric M = Σ_a+Σ_b (6: xx,xy,xz,yy,yz,zz) +// and mean offset d, scaled by √|Σ_a|·√|Σ_b| (passed as sdA·sdB). +fn crossG(sdA: f32, sdB: f32, m: array, d: vec3f) -> f32 { + let c00 = m[3] * m[5] - m[4] * m[4]; + let c01 = m[2] * m[4] - m[1] * m[5]; + let c02 = m[1] * m[4] - m[2] * m[3]; + let c11 = m[0] * m[5] - m[2] * m[2]; + let c12 = m[1] * m[2] - m[0] * m[4]; + let c22 = m[0] * m[3] - m[1] * m[1]; + let det = max(m[0] * c00 + m[1] * c01 + m[2] * c02, 1e-30); + let quadAdj = c00 * d.x * d.x + c11 * d.y * d.y + c22 * d.z * d.z + + 2.0 * (c01 * d.x * d.y + c02 * d.x * d.z + c12 * d.y * d.z); + let quad = quadAdj / det; + return TWO_PI_1_5 * sdA * sdB / sqrt(det) * exp(-0.5 * quad); } @compute @workgroup_size(64) @@ -133,142 +91,106 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { let bid = gid.x; if (bid >= uniforms.edgeCount) { return; } - let i = edgesI[bid]; - let j = edgesJ[bid]; - - let i8 = i * 8u; - let j8 = j * 8u; - let i9 = i * 9u; - let j9 = j * 9u; - - let mu_i = vec3f(posScalars[i8 + 0u], posScalars[i8 + 1u], posScalars[i8 + 2u]); - let mu_j = vec3f(posScalars[j8 + 0u], posScalars[j8 + 1u], posScalars[j8 + 2u]); - - let wi = posScalars[i8 + 3u]; - let wj = posScalars[j8 + 3u]; - let W = wi + wj; - let Wsafe = select(1.0, W, W > 0.0); - let pi_w_raw = wi / Wsafe; - let pi_w = clamp(pi_w_raw, 1e-12, 1.0 - 1e-12); - let pj_w = 1.0 - pi_w; - let logPi = log(pi_w); - let logPj = log(pj_w); - - // Merged mean. - let mm = pi_w * mu_i + pj_w * mu_j; - let di = mu_i - mm; - let dj = mu_j - mm; - - // Σ_i and Σ_j from rotation + variances. - let sig_i = sigmaFromRotVar(i9, i8); - let sig_j = sigmaFromRotVar(j9, j8); - - // Merged covariance: pi*(Σ_i + δi·δiᵀ) + pj*(Σ_j + δj·δjᵀ), + EPS on diag. - let s_xx = pi_w * (sig_i[0] + di.x*di.x) + pj_w * (sig_j[0] + dj.x*dj.x) + EPS_COV; - let s_xy = pi_w * (sig_i[1] + di.x*di.y) + pj_w * (sig_j[1] + dj.x*dj.y); - let s_xz = pi_w * (sig_i[2] + di.x*di.z) + pj_w * (sig_j[2] + dj.x*dj.z); - let s_yy = pi_w * (sig_i[3] + di.y*di.y) + pj_w * (sig_j[3] + dj.y*dj.y) + EPS_COV; - let s_yz = pi_w * (sig_i[4] + di.y*di.z) + pj_w * (sig_j[4] + dj.y*dj.z); - let s_zz = pi_w * (sig_i[5] + di.z*di.z) + pj_w * (sig_j[5] + dj.z*dj.z) + EPS_COV; - - // det of symmetric 3x3. - let det_m = s_xx * (s_yy*s_zz - s_yz*s_yz) - - s_xy * (s_xy*s_zz - s_yz*s_xz) - + s_xz * (s_xy*s_yz - s_yy*s_xz); - let logdet_m = log(max(det_m, 1e-30)); - - // Entropy of the merged Gaussian: H = 0.5 (k log(2π) + log|Σ_m| + k), k=3. - let EpNegLogQ = 0.5 * (3.0 * LOG2PI + logdet_m + 3.0); - - // Read per-axis std for each input (variances live at posScalars[s8+5..7]). - let vix = posScalars[i8 + 5u]; let viy = posScalars[i8 + 6u]; let viz = posScalars[i8 + 7u]; - let vjx = posScalars[j8 + 5u]; let vjy = posScalars[j8 + 6u]; let vjz = posScalars[j8 + 7u]; - let stdix = sqrt(max(vix, 0.0)); - let stdiy = sqrt(max(viy, 0.0)); - let stdiz = sqrt(max(viz, 0.0)); - let stdjx = sqrt(max(vjx, 0.0)); - let stdjy = sqrt(max(vjy, 0.0)); - let stdjz = sqrt(max(vjz, 0.0)); - - // Inverse diagonals (1 / variance) for the log-pdf quadratic term. - let invDi = vec3f(1.0 / max(vix, 1e-30), 1.0 / max(viy, 1e-30), 1.0 / max(viz, 1e-30)); - let invDj = vec3f(1.0 / max(vjx, 1e-30), 1.0 / max(vjy, 1e-30), 1.0 / max(vjz, 1e-30)); - let ldi = posScalars[i8 + 4u]; - let ldj = posScalars[j8 + 4u]; - - let z0 = uniforms.z0; - let z1 = uniforms.z1; - let z2 = uniforms.z2; - - // Sample x = mu + R · diag(std) · z where z ~ N(0, I). - // Row a of R is (rotR[r9+3a], rotR[r9+3a+1], rotR[r9+3a+2]). - // x[a] = mu[a] + R[a][0]*std[0]*z[0] + R[a][1]*std[1]*z[1] + R[a][2]*std[2]*z[2]. - let xix = mu_i.x + z0 * stdix * rotR[i9 + 0u] + z1 * stdiy * rotR[i9 + 1u] + z2 * stdiz * rotR[i9 + 2u]; - let xiy = mu_i.y + z0 * stdix * rotR[i9 + 3u] + z1 * stdiy * rotR[i9 + 4u] + z2 * stdiz * rotR[i9 + 5u]; - let xiz = mu_i.z + z0 * stdix * rotR[i9 + 6u] + z1 * stdiy * rotR[i9 + 7u] + z2 * stdiz * rotR[i9 + 8u]; - let xi = vec3f(xix, xiy, xiz); - - let xjx = mu_j.x + z0 * stdjx * rotR[j9 + 0u] + z1 * stdjy * rotR[j9 + 1u] + z2 * stdjz * rotR[j9 + 2u]; - let xjy = mu_j.y + z0 * stdjx * rotR[j9 + 3u] + z1 * stdjy * rotR[j9 + 4u] + z2 * stdjz * rotR[j9 + 5u]; - let xjz = mu_j.z + z0 * stdjx * rotR[j9 + 6u] + z1 * stdjy * rotR[j9 + 7u] + z2 * stdjz * rotR[j9 + 8u]; - let xj = vec3f(xjx, xjy, xjz); - - // log p_ij at samples from component i. - let logNiOnI = gaussLogpdfDiagrot(xi, mu_i, i9, invDi, ldi); - let logNjOnI = gaussLogpdfDiagrot(xi, mu_j, j9, invDj, ldj); - let logpOnI = logAddExp(logPi + logNiOnI, logPj + logNjOnI); - - let logNiOnJ = gaussLogpdfDiagrot(xj, mu_i, i9, invDi, ldi); - let logNjOnJ = gaussLogpdfDiagrot(xj, mu_j, j9, invDj, ldj); - let logpOnJ = logAddExp(logPi + logNiOnJ, logPj + logNjOnJ); - - let Ei = logpOnI; - let Ej = logpOnJ; - let EpLogp = pi_w * Ei + pj_w * Ej; - let geo = EpLogp + EpNegLogQ; - - // Appearance L2 cost, summed across the (up to three) chunks. Each chunk's - // stride is its live column count, so partial chunks store/read no padding. - // A 0 stride yields 0 iterations, so the dummy-bound appB / appC are never - // touched on inputs with fewer SH bands. - var cSh: f32 = 0.0; - let iA = i * STRIDE_A; let jA = j * STRIDE_A; - for (var k: u32 = 0u; k < STRIDE_A; k = k + 1u) { - let d = appA[iA + k] - appA[jA + k]; - cSh = cSh + d * d; - } - let iB = i * STRIDE_B; let jB = j * STRIDE_B; - for (var k: u32 = 0u; k < STRIDE_B; k = k + 1u) { - let d = appB[iB + k] - appB[jB + k]; - cSh = cSh + d * d; - } - let iC = i * STRIDE_C; let jC = j * STRIDE_C; - for (var k: u32 = 0u; k < STRIDE_C; k = k + 1u) { - let d = appC[iC + k] - appC[jC + k]; - cSh = cSh + d * d; - } + let io = edgesI[bid] * ${SPLAT_STRIDE}u; + let jo = edgesJ[bid] * ${SPLAT_STRIDE}u; + + let mui = vec3f(splat[io + 0u], splat[io + 1u], splat[io + 2u]); + let si = array(splat[io + 3u], splat[io + 4u], splat[io + 5u], splat[io + 6u], splat[io + 7u], splat[io + 8u]); + let sdi = splat[io + 9u]; let ai = splat[io + 10u]; let mi = splat[io + 11u]; + let bi = vec3f(splat[io + 12u], splat[io + 13u], splat[io + 14u]); let bni = splat[io + 15u]; + + let muj = vec3f(splat[jo + 0u], splat[jo + 1u], splat[jo + 2u]); + let sj = array(splat[jo + 3u], splat[jo + 4u], splat[jo + 5u], splat[jo + 6u], splat[jo + 7u], splat[jo + 8u]); + let sdj = splat[jo + 9u]; let aj = splat[jo + 10u]; let mj = splat[jo + 11u]; + let bj = vec3f(splat[jo + 12u], splat[jo + 13u], splat[jo + 14u]); let bnj = splat[jo + 15u]; + + let W = mi + mj; + let pi_ = select(0.5, mi / W, W > 0.0); + let pj_ = 1.0 - pi_; + + let mm = pi_ * mui + pj_ * muj; + let di = mui - mm; + let dj = muj - mm; + + // Merged covariance: Σ pₖ(δₖδₖᵀ + Σₖ) + EPS_COV·I. + let sm = array( + pi_ * (di.x * di.x + si[0]) + pj_ * (dj.x * dj.x + sj[0]) + EPS_COV, + pi_ * (di.x * di.y + si[1]) + pj_ * (dj.x * dj.y + sj[1]), + pi_ * (di.x * di.z + si[2]) + pj_ * (dj.x * dj.z + sj[2]), + pi_ * (di.y * di.y + si[3]) + pj_ * (dj.y * dj.y + sj[3]) + EPS_COV, + pi_ * (di.y * di.z + si[4]) + pj_ * (dj.y * dj.z + sj[4]), + pi_ * (di.z * di.z + si[5]) + pj_ * (dj.z * dj.z + sj[5]) + EPS_COV + ); - costs[bid] = geo + cSh; + let detm = max( + sm[0] * (sm[3] * sm[5] - sm[4] * sm[4]) - sm[1] * (sm[1] * sm[5] - sm[4] * sm[2]) + sm[2] * (sm[1] * sm[4] - sm[3] * sm[2]), + 1e-30 + ); + let sqrtDetM = sqrt(detm); + + // Merged opacity: mass-conserving, capped — needs merged scales (eigenvalues + // of Σ_m, Smith's closed form for a symmetric 3×3). + let q = (sm[0] + sm[3] + sm[5]) / 3.0; + let p1 = sm[1] * sm[1] + sm[2] * sm[2] + sm[4] * sm[4]; + var e0: f32; var e1: f32; var e2: f32; + if (p1 <= 1e-30) { + e0 = sm[0]; e1 = sm[3]; e2 = sm[5]; + } else { + let p2 = (sm[0] - q) * (sm[0] - q) + (sm[3] - q) * (sm[3] - q) + (sm[5] - q) * (sm[5] - q) + 2.0 * p1; + let p = sqrt(p2 / 6.0); + let ip = 1.0 / p; + let b00 = (sm[0] - q) * ip; let b11 = (sm[3] - q) * ip; let b22 = (sm[5] - q) * ip; + let b01 = sm[1] * ip; let b02 = sm[2] * ip; let b12 = sm[4] * ip; + let detB = b00 * (b11 * b22 - b12 * b12) - b01 * (b01 * b22 - b12 * b02) + b02 * (b01 * b12 - b11 * b02); + let r = clamp(detB * 0.5, -1.0, 1.0); + let phi = acos(r) / 3.0; + e0 = q + 2.0 * p * cos(phi); + e2 = q + 2.0 * p * cos(phi + TWO_PI_3); + e1 = 3.0 * q - e0 - e2; + } + let s0 = sqrt(max(e0, 1e-18)); + let s1 = sqrt(max(e1, 1e-18)); + let s2 = sqrt(max(e2, 1e-18)); + let am = min(1.0, W / max(ellipsoidArea(s0, s1, s2), 1e-30)); + + // Base-colour dots (merged colour = mass-weighted average). + let bij = dot(bi, bj); + let bim = pi_ * bni + pj_ * bij; + let bjm = pi_ * bij + pj_ * bnj; + let bnm = pi_ * pi_ * bni + 2.0 * pi_ * pj_ * bij + pj_ * pj_ * bnj; + + let selfI = ai * ai * bni * PI_1_5 * sdi; + let selfJ = aj * aj * bnj * PI_1_5 * sdj; + let selfM = am * am * bnm * PI_1_5 * sqrtDetM; + + let mij = array(si[0] + sj[0], si[1] + sj[1], si[2] + sj[2], si[3] + sj[3], si[4] + sj[4], si[5] + sj[5]); + let cIJ = crossG(sdi, sdj, mij, mui - muj); + let mim = array(si[0] + sm[0], si[1] + sm[1], si[2] + sm[2], si[3] + sm[3], si[4] + sm[4], si[5] + sm[5]); + let cIM = crossG(sdi, sqrtDetM, mim, di); + let mjm = array(sj[0] + sm[0], sj[1] + sm[1], sj[2] + sm[2], sj[3] + sm[3], sj[4] + sm[4], sj[5] + sm[5]); + let cJM = crossG(sdj, sqrtDetM, mjm, dj); + + let E = selfI + selfJ + selfM + + 2.0 * ai * aj * bij * cIJ - + 2.0 * ai * am * bim * cIM - + 2.0 * aj * am * bjm * cJM; + + // Clamp float-noise negatives but preserve NaN/Inf (degenerate input → the + // host fails loud on no finite merges), then add the scale-free colour + // dissimilarity term: |Δbase|² = bni + bnj − 2·(bi·bj). + costs[bid] = select(E, 0.0, E < 0.0) + COLOR_WEIGHT * (bni + bnj - 2.0 * bij); } `; /** - * Per-splat cache for the edge cost kernel. Packed layouts to stay within the - * WebGPU per-stage storage-buffer count limit (8) and the per-binding size - * limit (~2 GB) — appearance is split into 16-column chunks for the latter. + * Per-splat cache for the edge-cost kernel: a single interleaved buffer of + * {@link SPLAT_STRIDE} floats per splat (mean, covariance, √det, alpha, mass, + * base colour, |base|²), built by `packGpuCache` from the CPU `buildCostCache` + * so the GPU reads identical per-splat inputs. */ interface EdgeCostCache { - /** Per-splat geometry interleaved 8-wide: (x, y, z, mass, logdet, vx, vy, vz). */ - posScalars: Float32Array; - /** Row-major 3×3 rotation per splat (length 9N). */ - rotR: Float32Array; - /** - * Appearance in up to three chunks of ≤16 columns. Chunk c has stride - * width_c (= its live column count): appChunks[c][s * width_c + k]. - */ - appChunks: Float32Array[]; - /** Number of appearance columns C. */ - numAppCols: number; + /** Interleaved per-splat cache, length {@link SPLAT_STRIDE}·N. */ + splatData: Float32Array; /** Number of splats. */ numSplats: number; } @@ -276,27 +198,15 @@ interface EdgeCostCache { /** * GPU edge-cost evaluator. * - * Each compute thread evaluates the KL-style cost for one edge (i, j) by - * reading the per-splat cache for both endpoints, computing the merged - * Gaussian's covariance/determinant, running a single Monte-Carlo sample - * through both component PDFs, and adding an L2 distance over the - * appearance (SH) coefficients. Output is `costs[e] = cost for edge e`. - * - * Mirrors the CPU `computeEdgeCostView` in `decimate/edge-cost-cpu.ts`. + * Each compute thread evaluates the L2 field-error cost for one edge (i, j), + * mirroring the CPU `computeEdgeCostView` in `decimate/edge-cost-cpu.ts`. + * Output is `costs[e] = cost for edge e`. */ class GpuEdgeCost { - /** - * @param cache - Per-splat cache (uploaded once). - * @param edgeI - Edge u indices (length E). - * @param edgeJ - Edge v indices (length E). - * @param z - Single Monte-Carlo sample (3 floats from N(0,1)). - * @param outCosts - Destination for per-edge costs (length E). - */ execute: ( cache: EdgeCostCache, edgeI: Uint32Array, edgeJ: Uint32Array, - z: Float32Array, outCosts: Float32Array ) => Promise; destroy: () => void; @@ -305,96 +215,51 @@ class GpuEdgeCost { * @param device - PlayCanvas GraphicsDevice (WebGPU). * @param maxN - Maximum number of splats. * @param maxE - Maximum number of edges in a single dispatch. - * @param maxAppCols - Maximum appearance column count (over all bands). */ - constructor(device: GraphicsDevice, maxN: number, maxE: number, maxAppCols: number) { + constructor(device: GraphicsDevice, maxN: number, maxE: number) { const workgroupSize = 64; const edgesPerBatch = 1024 * workgroupSize; // 65,536 - // Appearance is split at fixed APP_CHUNK-column boundaries, but each - // chunk's *stride* is its live column count — only the last non-empty - // chunk is ever partial, so partial chunks neither allocate nor upload - // padding. The widest possible chunk reaches the ~2 GB limit at ~33.5M - // splats, past the ~11.2M wall the single 48-col buffer hit. The three - // kernel bindings (appA/appB/appC) cap the layout at three chunks. - // e.g. [16, 11, 0] for 27 cols, [3, 0, 0] for DC-only. - const appStrides = [0, 1, 2].map((ch) => { - return Math.min(APP_CHUNK, Math.max(0, maxAppCols - ch * APP_CHUNK)); - }); - // Non-empty chunk count the kernel reads. execute() validates the cache - // supplies exactly this many: a short count would leave a hoisted (reused - // across iterations) appearance buffer holding the previous iteration's - // data, which the kernel would then read as this iteration's appearance. - const numAppChunks = appStrides.filter(stride => stride > 0).length; const bindGroupFormat = new BindGroupFormat(device, [ new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE), new BindStorageBufferFormat('edgesI', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('edgesJ', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('posScalars', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('rotR', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('appA', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('appB', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('appC', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('splat', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('costs', SHADERSTAGE_COMPUTE) ]); const shader = new Shader(device, { name: 'compute-edge-cost', shaderLanguage: SHADERLANGUAGE_WGSL, - cshader: edgeCostWgsl(appStrides[0], appStrides[1], appStrides[2]), + cshader: edgeCostWgsl(), // @ts-ignore computeUniformBufferFormats: { uniforms: new UniformBufferFormat(device, [ - new UniformFormat('edgeCount', UNIFORMTYPE_UINT), - new UniformFormat('z0', UNIFORMTYPE_FLOAT), - new UniformFormat('z1', UNIFORMTYPE_FLOAT), - new UniformFormat('z2', UNIFORMTYPE_FLOAT) + new UniformFormat('edgeCount', UNIFORMTYPE_UINT) ]) }, // @ts-ignore computeBindGroupFormat: bindGroupFormat }); - // Pre-flight the largest per-N bindings against the device's storage - // limit so we fail with a clear message instead of a driver-side error. - // Edges are uploaded per batch (batch-sized buffers), so they can't hit - // the limit; the widest appearance chunk and rotR are the candidates. - // posScalars (8 floats/splat) is strictly smaller than rotR (9), so the - // rotR check already bounds it — no separate check needed. + // Pre-flight the per-N splat buffer against the device's storage limit + // so we fail with a clear message instead of a driver-side error. const maxStorage = (device as any).limits?.maxStorageBufferBindingSize; if (typeof maxStorage === 'number') { - const checkLimit = (label: string, bytes: number) => { - if (bytes > maxStorage) { - throw new Error( - `GpuEdgeCost: ${label} buffer (${bytes} bytes) exceeds device ` + - `maxStorageBufferBindingSize (${maxStorage})` - ); - } - }; - const maxChunkCols = Math.max(...appStrides); - checkLimit(`appearance chunk (${maxN} splats × ${maxChunkCols} cols)`, maxN * maxChunkCols * 4); - checkLimit(`rotR (${maxN} splats × 9)`, maxN * 9 * 4); + const bytes = maxN * SPLAT_STRIDE * 4; + if (bytes > maxStorage) { + throw new Error( + `GpuEdgeCost: splat buffer (${maxN} splats × ${SPLAT_STRIDE} floats = ${bytes} bytes) ` + + `exceeds device maxStorageBufferBindingSize (${maxStorage})` + ); + } } - const posScalarsBuf = new StorageBuffer(device, maxN * 8 * 4, BUFFERUSAGE_COPY_DST); - const rotRBuf = new StorageBuffer(device, maxN * 9 * 4, BUFFERUSAGE_COPY_DST); - - // One buffer per non-empty appearance chunk, sized to that chunk's live - // column count; empty slots (inputs with fewer SH bands) share a small - // dummy since WebGPU forbids a zero-size binding. The 3-binding layout - // stays fixed regardless of band count. - const appDummy = new StorageBuffer(device, 16, BUFFERUSAGE_COPY_DST); - const appBufs: StorageBuffer[] = appStrides.map((width) => { - return width > 0 ? - new StorageBuffer(device, maxN * width * 4, BUFFERUSAGE_COPY_DST) : - appDummy; - }); + const splatBuf = new StorageBuffer(device, maxN * SPLAT_STRIDE * 4, BUFFERUSAGE_COPY_DST); - // Two parallel u32 buffers, sized to a single dispatch batch (not the - // full N·k edge list): execute uploads each batch's slice before its - // dispatch. Batch-sizing keeps these ~256 KB instead of N·k·4 — off the - // ~2 GB per-binding limit (so edges never cap scene size) and ~1.6 GB - // less VRAM at 13M splats. Two parallel arrays avoid a host-side pack. + // Two parallel u32 edge buffers, sized to a single dispatch batch (not + // the full N·k edge list): execute uploads each batch's slice before its + // dispatch, keeping these off the ~2 GB per-binding limit. const edgesIBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); const edgesJBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); @@ -408,18 +273,13 @@ class GpuEdgeCost { const compute = new Compute(device, shader, 'compute-edge-cost'); compute.setParameter('edgesI', edgesIBuf); compute.setParameter('edgesJ', edgesJBuf); - compute.setParameter('posScalars', posScalarsBuf); - compute.setParameter('rotR', rotRBuf); - compute.setParameter('appA', appBufs[0]); - compute.setParameter('appB', appBufs[1]); - compute.setParameter('appC', appBufs[2]); + compute.setParameter('splat', splatBuf); compute.setParameter('costs', outBuf); this.execute = async ( cache: EdgeCostCache, edgeI: Uint32Array, edgeJ: Uint32Array, - z: Float32Array, outCosts: Float32Array ) => { const n = cache.numSplats; @@ -427,30 +287,11 @@ class GpuEdgeCost { if (n > maxN) throw new Error(`GpuEdgeCost: N=${n} exceeds maxN=${maxN}`); if (e > maxE) throw new Error(`GpuEdgeCost: E=${e} exceeds maxE=${maxE}`); - if (cache.numAppCols !== maxAppCols) { - throw new Error(`GpuEdgeCost: numAppCols=${cache.numAppCols} must equal maxAppCols=${maxAppCols} (baked into the kernel)`); - } - if (cache.appChunks.length !== numAppChunks) { - throw new Error(`GpuEdgeCost: cache supplies ${cache.appChunks.length} appearance chunks but the kernel layout expects ${numAppChunks}`); - } if (edgeJ.length !== e || outCosts.length !== e) { throw new Error('GpuEdgeCost: edgeI / edgeJ / outCosts must have same length'); } - if (z.length < 3) { - throw new Error('GpuEdgeCost: z must have at least 3 elements'); - } - // Upload per-splat cache. Each appearance chunk is row-major with - // stride = its live column count, so we upload n*width per chunk. - posScalarsBuf.write(0, cache.posScalars, 0, n * 8); - rotRBuf.write(0, cache.rotR, 0, n * 9); - for (let ch = 0; ch < cache.appChunks.length; ch++) { - appBufs[ch].write(0, cache.appChunks[ch], 0, n * appStrides[ch]); - } - - compute.setParameter('z0', z[0]); - compute.setParameter('z1', z[1]); - compute.setParameter('z2', z[2]); + splatBuf.write(0, cache.splatData, 0, n * SPLAT_STRIDE); const numBatches = Math.ceil(e / edgesPerBatch); for (let batch = 0; batch < numBatches; batch++) { @@ -458,8 +299,6 @@ class GpuEdgeCost { const edgeCount = Math.min(edgesPerBatch, e - edgeOffset); const groups = Math.ceil(edgeCount / workgroupSize); - // Upload just this batch's edges to offset 0; the kernel indexes - // edgesI/J[bid] within the batch. edgesIBuf.write(0, edgeI, edgeOffset, edgeCount); edgesJBuf.write(0, edgeJ, edgeOffset, edgeCount); @@ -475,12 +314,7 @@ class GpuEdgeCost { }; this.destroy = () => { - posScalarsBuf.destroy(); - rotRBuf.destroy(); - for (const buf of appBufs) { - if (buf !== appDummy) buf.destroy(); - } - appDummy.destroy(); + splatBuf.destroy(); edgesIBuf.destroy(); edgesJBuf.destroy(); outBuf.destroy(); diff --git a/test/decimate-edge-cost.test.mjs b/test/decimate-edge-cost.test.mjs new file mode 100644 index 00000000..a18be559 --- /dev/null +++ b/test/decimate-edge-cost.test.mjs @@ -0,0 +1,64 @@ +/** + * Field-L2 edge-cost formula properties (CPU): merging identical coincident + * splats is (near-)lossless, and — unlike the old scale-invariant KL cost — + * the cost scales with absolute Gaussian size, so a geometrically-similar merge + * of large Gaussians costs far more than one of tiny Gaussians. + */ + +import assert from 'node:assert'; +import { describe, it } from 'node:test'; + +import { buildCostCache, computeEdgeCostView } from '../src/lib/decimate/edge-cost-cpu.js'; + +// Build a minimal SplatView from per-splat specs (identity quaternion; only the +// DC colour is set, higher SH left zero). colorDim = 3 (DC only). +const makeView = (splats) => { + const n = splats.length; + const colorDim = 3; + const pos = new Float32Array(n * 3); + const geo = new Float32Array(n * 8); + const color = new Float32Array(n * colorDim); + for (let i = 0; i < n; i++) { + pos[i * 3] = splats[i].pos[0]; + pos[i * 3 + 1] = splats[i].pos[1]; + pos[i * 3 + 2] = splats[i].pos[2]; + geo[i * 8] = 1; // identity quat (w, x, y, z) + geo[i * 8 + 4] = splats[i].ls[0]; + geo[i * 8 + 5] = splats[i].ls[1]; + geo[i * 8 + 6] = splats[i].ls[2]; + geo[i * 8 + 7] = splats[i].op; + color[i * 3] = splats[i].dc[0]; + color[i * 3 + 1] = splats[i].dc[1]; + color[i * 3 + 2] = splats[i].dc[2]; + } + return { pos, geo, color, colorDim }; +}; + +const cost = (view, i, j) => computeEdgeCostView(view, buildCostCache(view), i, j); + +describe('field-L2 edge cost', () => { + it('merging identical coincident splats is (near-)lossless', () => { + const l = Math.log(0.1); + const s = { pos: [0.3, -0.2, 0.5], ls: [l, l, l], op: -2, dc: [0.1, 0.2, 0.3] }; + const view = makeView([s, { ...s }]); + assert.ok(cost(view, 0, 1) < 1e-3, `expected ~0, got ${cost(view, 0, 1)}`); + }); + + it('cost grows with absolute Gaussian size (not scale-invariant)', () => { + const small = makeView([ + { pos: [0, 0, 0], ls: [Math.log(0.1), Math.log(0.1), Math.log(0.1)], op: -4, dc: [1, 0, 0] }, + { pos: [0.2, 0, 0], ls: [Math.log(0.1), Math.log(0.1), Math.log(0.1)], op: -4, dc: [0, 0, 1] } + ]); + // Same geometry scaled ×10 (positions and scales). + const large = makeView([ + { pos: [0, 0, 0], ls: [Math.log(1.0), Math.log(1.0), Math.log(1.0)], op: -4, dc: [1, 0, 0] }, + { pos: [2, 0, 0], ls: [Math.log(1.0), Math.log(1.0), Math.log(1.0)], op: -4, dc: [0, 0, 1] } + ]); + const cSmall = cost(small, 0, 1); + const cLarge = cost(large, 0, 1); + assert.ok(cSmall > 0, `distinct merge should cost > 0 (got ${cSmall})`); + // Field L2 scales with Gaussian volume (~s³); a scale-invariant cost + // would give a ratio near 1. Require a large margin. + assert.ok(cLarge > cSmall * 100, `large/small ratio ${(cLarge / cSmall).toFixed(1)} (want >> 1)`); + }); +}); diff --git a/test/decimate-merge-stream.test.mjs b/test/decimate-merge-stream.test.mjs index a383259d..199a0457 100644 --- a/test/decimate-merge-stream.test.mjs +++ b/test/decimate-merge-stream.test.mjs @@ -27,12 +27,16 @@ describe('mergeStream', () => { const ctx = { source, pool, pos, order, blocks, K, k }; await runPriorityPass(ctx, cand); const sel = selectMerges(cand, n, K, n - target); - assert.strictEqual(sel.removed, n - target); + // One-shot selection may fall short of the requested removals on a + // sparse candidate graph (group cap); the stream contract is defined + // by what the selection actually returns. + assert.ok(sel.removed > 0 && sel.removed <= n - target); + const outCount = n - sel.removed; const nextPositions = { - x: new Float32Array(target), - y: new Float32Array(target), - z: new Float32Array(target) + x: new Float32Array(outCount), + y: new Float32Array(outCount), + z: new Float32Array(outCount) }; const rows = { pos: [], geo: [], color: [] }; for await (const payload of mergeStream({ ...ctx, selection: sel, nextPositions }, 256)) { @@ -41,7 +45,7 @@ describe('mergeStream', () => { rows.color.push(new Float32Array(payload.color)); } const emitted = rows.pos.reduce((a, p) => a + p.length / 3, 0); - assert.strictEqual(emitted, target); + assert.strictEqual(emitted, outCount); const flatPos = Float32Array.from(rows.pos.flatMap(a => [...a])); const flatGeo = Float32Array.from(rows.geo.flatMap(a => [...a])); @@ -83,7 +87,7 @@ describe('mergeStream', () => { row++; } } - assert.strictEqual(row, target); + assert.strictEqual(row, outCount); assert.ok(checkedMerges > 50 && checkedSurvivors > 50, `coverage: ${checkedMerges} merges, ${checkedSurvivors} survivors`); }); @@ -105,7 +109,7 @@ describe('mergeStream', () => { rowsOther.push(new Uint32Array(payload.other)); } const flatOther = Uint32Array.from(rowsOther.flatMap(a => [...a])); - assert.strictEqual(flatOther.length, target * otherDim); + assert.strictEqual(flatOther.length, (n - sel.removed) * otherDim); // survivors keep their tag verbatim let row = 0; for (const b of blocks) { diff --git a/test/decimate-priority.test.mjs b/test/decimate-priority.test.mjs index 3148e3a6..f61c8df7 100644 --- a/test/decimate-priority.test.mjs +++ b/test/decimate-priority.test.mjs @@ -1,19 +1,18 @@ /** * Priority pass tests (CPU path): the resident best-K candidates must match - * brute-force legacy edge costs computed over exact global KNN. + * brute-force edge costs computed over exact global KNN. */ import assert from 'node:assert'; import { describe, it } from 'node:test'; -import { legacyEdgeCost } from './fixtures/legacy-decimate-math.mjs'; import { makeSyntheticSource } from './helpers/synthetic-source.mjs'; -import { makeGaussianSamples } from '../src/lib/decimate/moment-match.js'; +import { buildCostCache, computeEdgeCostView } from '../src/lib/decimate/edge-cost-cpu.js'; import { kdPartition } from '../src/lib/decimate/partition.js'; import { runPriorityPass } from '../src/lib/decimate/priority.js'; describe('priority pass (CPU)', () => { - it('best-K candidates match brute-force legacy costs over exact KNN', async () => { + it('best-K candidates match brute-force costs over exact KNN', async () => { const n = 1500, k = 16, K = 4; const { source, pool, view, pos } = await makeSyntheticSource(n, 1, 5, { chunkSize: 256 }); const { order, blocks } = kdPartition(pos, 400); @@ -23,14 +22,14 @@ describe('priority pass (CPU)', () => { }; await runPriorityPass({ source, pool, pos, order, blocks, K, k }, cand); - const Z = makeGaussianSamples(1, 0); + const cache = buildCostCache(view); const d2 = (a, b) => (pos.x[a] - pos.x[b]) ** 2 + (pos.y[a] - pos.y[b]) ** 2 + (pos.z[a] - pos.z[b]) ** 2; for (let i = 0; i < n; i += 97) { const knn = Array.from({ length: n }, (_, j) => j) .filter(j => j !== i) .sort((a, b) => d2(i, a) - d2(i, b)) .slice(0, k); - const refCosts = knn.map(j => legacyEdgeCost(view, i, j, Z)).sort((a, b) => a - b); + const refCosts = knn.map(j => computeEdgeCostView(view, cache, i, j)).sort((a, b) => a - b); const got = []; for (let s = 0; s < K; s++) { if (cand.idx[i * K + s] !== 0xFFFFFFFF) got.push(cand.cost[i * K + s]); diff --git a/test/decimate-select.test.mjs b/test/decimate-select.test.mjs index b07d3bb1..955e61f2 100644 --- a/test/decimate-select.test.mjs +++ b/test/decimate-select.test.mjs @@ -1,12 +1,14 @@ /** - * Global selection tests: disjointness, exact targets, closure behaviour, - * bucketed-vs-exact-sorted greedy equivalence, non-finite exclusion. + * Global selection tests (cost-ordered agglomeration): disjoint clustering, + * exact targets, CSR consistency, cheap-region concentration (expensive + * gaussians spared), multi-member clusters when pairing alone can't reach the + * target, non-finite exclusion. */ import assert from 'node:assert'; import { describe, it } from 'node:test'; -import { selectMerges } from '../src/lib/decimate/select.js'; +import { selectMerges, MAX_GROUP } from '../src/lib/decimate/select.js'; const mulberry = (seed) => { let t = seed >>> 0; @@ -18,7 +20,7 @@ const mulberry = (seed) => { }; }; -// Ring candidates: each i lists i±1, i±2 — a perfect matching always exists. +// Ring candidates: each i lists i±1, i±2 — a dense graph that clusters freely. const ringCandidates = (N, K, r) => { const idx = new Uint32Array(N * K); const cost = new Float32Array(N * K); @@ -32,8 +34,8 @@ const ringCandidates = (N, K, r) => { return { idx, cost }; }; -describe('selectMerges', () => { - it('selection is disjoint, hits exact target, CSR consistent', () => { +describe('selectMerges (agglomeration)', () => { + it('clusters are disjoint, hit exact target, CSR consistent', () => { const N = 10000, K = 4, r = mulberry(9); const cand = ringCandidates(N, K, r); const needed = N / 2; @@ -42,7 +44,7 @@ describe('selectMerges', () => { const seen = new Int32Array(N).fill(-1); for (let g = 0; g < sel.mergedGroups; g++) { const size = sel.groupOffsets[g + 1] - sel.groupOffsets[g]; - assert.ok(size >= 2 && size <= 4, `group ${g} size ${size}`); + assert.ok(size >= 2 && size <= MAX_GROUP, `group ${g} size ${size}`); let min = Infinity; for (let m = sel.groupOffsets[g]; m < sel.groupOffsets[g + 1]; m++) { const id = sel.groupMembers[m]; @@ -53,6 +55,7 @@ describe('selectMerges', () => { } assert.strictEqual(sel.groupMin[g], min); } + // Each cluster collapses to one survivor: survivors + groups = N - removed. let survivors = 0; for (let i = 0; i < N; i++) { if (sel.memberGroup[i] === -1) survivors++; @@ -60,51 +63,42 @@ describe('selectMerges', () => { assert.strictEqual(survivors + sel.mergedGroups, N - needed); }); - it('bucketed greedy tracks exact-sorted greedy selection cost within quantization', () => { - const N = 4000, K = 4, r = mulberry(21); - const cand = ringCandidates(N, K, r); - const needed = Math.floor(N * 0.3); // selective regime - const sel = selectMerges(cand, N, K, needed); - assert.strictEqual(sel.removed, needed); - - // reference: exact-sort greedy over the same candidate edges - const entries = []; - for (let e = 0; e < N * K; e++) { - if (Number.isFinite(cand.cost[e])) entries.push(e); - } - entries.sort((a, b) => cand.cost[a] - cand.cost[b]); - const used = new Uint8Array(N); - let refCost = 0, taken = 0; - for (const e of entries) { - const i = Math.floor(e / K), j = cand.idx[e]; - if (used[i] || used[j]) continue; - used[i] = 1; - used[j] = 1; - refCost += cand.cost[e]; - if (++taken >= needed) break; - } - - let gotCost = 0; - for (let g = 0; g < sel.mergedGroups; g++) { - const a = sel.groupMembers[sel.groupOffsets[g]]; - const b = sel.groupMembers[sel.groupOffsets[g] + 1]; - let c = Infinity; + it('removal concentrates in the low-cost region, sparing expensive gaussians', () => { + // Two disjoint rings: a cheap region [0, H) and an expensive one [H, N). + // Cost-ordered agglomeration must exhaust the cheap region before ever + // touching an expensive edge, so with a target the cheap side can cover + // on its own, every expensive gaussian survives untouched. + const N = 2000, K = 4, H = 1000, r = mulberry(7); + const idx = new Uint32Array(N * K); + const cost = new Float32Array(N * K); + for (let i = 0; i < N; i++) { + const base = i < H ? 0 : H; + const local = i - base; + const nb = [ + base + (local + 1) % H, base + (local + H - 1) % H, + base + (local + 2) % H, base + (local + H - 2) % H + ]; + const cheap = i < H; for (let s = 0; s < K; s++) { - if (cand.idx[a * K + s] === b) c = Math.min(c, cand.cost[a * K + s]); - if (cand.idx[b * K + s] === a) c = Math.min(c, cand.cost[b * K + s]); + idx[i * K + s] = nb[s]; + cost[i * K + s] = (cheap ? 0.01 * r() : 10 + r()) + (s >= 2 ? (cheap ? 0.02 : 1) : 0); } - gotCost += c; } - assert.ok(gotCost <= refCost * 1.02, `bucketed ${gotCost} vs sorted ${refCost}`); + const needed = 400; // < cheap-region capacity + const sel = selectMerges({ idx, cost }, N, K, needed); + assert.strictEqual(sel.removed, needed); + for (let i = H; i < N; i++) { + assert.strictEqual(sel.memberGroup[i], -1, `expensive gaussian ${i} must be spared`); + } }); - it('closure attaches unmatched gaussians when pairs alone cannot reach the target', () => { - // Star topology: everyone's candidates point at a tiny hub set, so - // primary pairing exhausts quickly and closure must attach the rest. - const N = 100, K = 2; + it('reaches the target via multi-member clusters when pairing alone cannot', () => { + // Star topology: everyone points at a tiny hub set, so a disjoint + // matching tops out at ~8 removals. Agglomeration must grow hub-centred + // clusters (size > 2) to reach a target beyond that. + const N = 100, K = 2, r = mulberry(5); const idx = new Uint32Array(N * K); const cost = new Float32Array(N * K); - const r = mulberry(5); for (let i = 0; i < N; i++) { idx[i * K] = i < 8 ? (i + 1) % 8 : i % 8; // hub 0..7 idx[i * K + 1] = (i % 8 === 0) ? 1 : 0; @@ -113,11 +107,14 @@ describe('selectMerges', () => { } const needed = 12; const sel = selectMerges({ idx, cost }, N, K, needed); - assert.ok(sel.removed > 4, `closure extended removal (${sel.removed})`); + assert.strictEqual(sel.removed, needed); + let maxSize = 0; for (let g = 0; g < sel.mergedGroups; g++) { const size = sel.groupOffsets[g + 1] - sel.groupOffsets[g]; - assert.ok(size >= 2 && size <= 4); + assert.ok(size >= 2 && size <= MAX_GROUP, `group ${g} size ${size}`); + maxSize = Math.max(maxSize, size); } + assert.ok(maxSize >= 3, `expected a multi-member cluster (max size ${maxSize})`); }); it('non-finite costs are never selected', () => { diff --git a/test/decimate-source.test.mjs b/test/decimate-source.test.mjs index ceab7b8f..b4c1d3db 100644 --- a/test/decimate-source.test.mjs +++ b/test/decimate-source.test.mjs @@ -1,7 +1,7 @@ /** * decimateSource orchestrator tests: exact counts, value domains, deep * targets (multi-generation with RAM intermediates), spill path with temp - * cleanup, statistical parity vs the legacy reference, and input validation. + * cleanup, in-domain aggregate statistics, and input validation. */ import assert from 'node:assert'; @@ -10,7 +10,6 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { describe, it } from 'node:test'; -import { legacySimplify } from './fixtures/legacy-decimate.mjs'; import { makeSyntheticSource } from './helpers/synthetic-source.mjs'; import { decimateSource } from '../src/lib/decimate/index.js'; @@ -89,19 +88,17 @@ describe('decimateSource', () => { assert.strictEqual(geo.length, 500 * 8); }); - it('statistical parity with the legacy algorithm on the same scene', async () => { + it('decimated output has finite, in-domain aggregate statistics', async () => { const n = 3000; - const { source, pool, view } = await makeSyntheticSource(n, 1, 23, { chunkSize: 512 }); + const { source, pool } = await makeSyntheticSource(n, 1, 23, { chunkSize: 512 }); const out = await decimateSource(source, pool, { targetCount: 1500 }); const { geometric: oursGeo } = await readLayers(out, pool, ["geometric"]); await out.close(); - const legacy = legacySimplify(view, 1500); const a = stats(oursGeo, 1500); - const b = stats(legacy.geo, 1500); - assert.ok(Math.abs(a.opMean - b.opMean) / Math.abs(b.opMean) < 0.05, `opMean ${a.opMean} vs ${b.opMean}`); - assert.ok(Math.abs(a.scaleMean - b.scaleMean) / Math.abs(b.scaleMean) < 0.05, `scaleMean ${a.scaleMean} vs ${b.scaleMean}`); - assert.ok(Math.abs(a.opStd - b.opStd) / Math.abs(b.opStd) < 0.10, `opStd ${a.opStd} vs ${b.opStd}`); + assert.ok(a.opMean > 0 && a.opMean <= 1, `opacity mean in (0, 1] (${a.opMean})`); + assert.ok(a.opStd >= 0 && Number.isFinite(a.opStd), `opacity std finite (${a.opStd})`); + assert.ok(Number.isFinite(a.scaleMean), `scale mean finite (${a.scaleMean})`); }); it('spill path: intermediate generations write temp PLYs and clean them up', async () => { diff --git a/test/decimate.test.mjs b/test/decimate.test.mjs index 4252c97e..5e5f0c80 100644 --- a/test/decimate.test.mjs +++ b/test/decimate.test.mjs @@ -203,16 +203,19 @@ describe('decimate - merge quality invariants', () => { ); }); - it('should fail loud (throw) when the scene is too degenerate to decimate', async () => { - // Every splat coincident at the origin: identical queries tie-break to - // the same KNN hub set, so candidate lists collapse and the matching - // starves — same pathology and same fail-loud stall guard as legacy. + it('should decimate a fully-coincident scene (re-costed selection)', async () => { + // Every splat coincident at the origin: the legacy one-shot matching + // starved on the collapsed KNN hub set and failed loud. Re-costed + // selection derives candidates from the full neighbour graph as + // clusters form, so coincident scenes now merge cleanly (coincident + // merges are exactly lossless under the field-L2 cost). const testData = createGaussianTestData({ count: 600 }); - await assert.rejects( - () => processDataTable(testData, decimate(300)), - /too degenerate to merge further/, - 'should throw when coincident splats starve the matching' - ); + const result = await processDataTable(testData, decimate(300)); + assert.strictEqual(result.numRows, 300); + const op = result.getColumnByName('opacity').data; + for (let i = 0; i < 300; i++) { + assert.ok(Number.isFinite(op[i]), `row ${i} opacity finite`); + } }); it('should throw when gaussian columns are missing (legacy silently pruned)', async () => { From 2ee8822d9f133cc1333c1423a698547c3c103a88 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Mon, 27 Jul 2026 20:27:22 +0100 Subject: [PATCH 02/19] latest --- src/lib/decimate/decimate-source.ts | 8 +++++++- src/lib/decimate/edge-cost-cpu.ts | 10 ++++++---- src/lib/decimate/select-recost.ts | 11 ++++++----- src/lib/gpu/gpu-edge-cost.ts | 10 ++++++---- 4 files changed, 25 insertions(+), 14 deletions(-) diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index 2064f546..2f39ff53 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -165,6 +165,10 @@ const decimateSource = async ( let positions: ResidentPositions | null = null; // Cleanup for the CURRENT generation's input (previous spill / RAM source). let disposeCurrentInput: (() => Promise) | null = null; + // Resident footprint of the current generation's input (non-zero when the + // previous generation was materialized in RAM) — counted by the re-costed + // selection gate so the budget covers everything actually resident. + let residentInputBytes = 0; const totalGenerations = Math.max(1, Math.ceil(Math.log2(inputMeta.numGaussians / targetCount))); @@ -218,7 +222,7 @@ const decimateSource = async ( // selection otherwise. Gated per generation, so large scenes regain // re-costing as soon as the cascade shrinks under the budget. const k = Math.min(KNN_K, Math.max(1, N - 1)); - const baseBytes = N * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; + const baseBytes = residentInputBytes + N * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; const recost = baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; const cacheOut = recost ? new Float32Array(N * CACHE_STRIDE) : undefined; const neighborsOut = recost ? new Uint32Array(N * k) : undefined; @@ -322,7 +326,9 @@ const decimateSource = async ( if (estBytes <= budget / 4) { nextSrc = await compact(producer, pool); + residentInputBytes = estBytes; } else { + residentInputBytes = 0; if (!opts.spill) { throw new Error( `decimation intermediate generation needs ${fmtBytes(estBytes)}, over the in-memory budget — ` + diff --git a/src/lib/decimate/edge-cost-cpu.ts b/src/lib/decimate/edge-cost-cpu.ts index 8e5ec9f5..9e38b1c8 100644 --- a/src/lib/decimate/edge-cost-cpu.ts +++ b/src/lib/decimate/edge-cost-cpu.ts @@ -251,10 +251,12 @@ const computeEdgeCostView = ( 2 * ai * am * bim * cIM - 2 * aj * am * bjm * cJM; - // Clamp float-noise negatives (E is a squared norm ≥ 0) but preserve - // NaN/Inf from degenerate inputs so the caller can fail loud. - // |Δbase|² = |b_i|² + |b_j|² − 2·b_i·b_j — no extra loads needed. - return (E < 0 ? 0 : E) + COLOR_WEIGHT * (bni + bnj - 2 * bij); + // Clamp float-noise negatives (both terms are squared norms ≥ 0; the + // colour distance uses separately-rounded f32 norms so it too can round + // slightly negative) but preserve NaN/Inf from degenerate inputs so the + // caller can fail loud. |Δbase|² = |b_i|² + |b_j|² − 2·b_i·b_j. + const dc2 = bni + bnj - 2 * bij; + return (E < 0 ? 0 : E) + COLOR_WEIGHT * (dc2 < 0 ? 0 : dc2); }; export { buildCostCache, computeEdgeCostView, COLOR_WEIGHT, type CostCache }; diff --git a/src/lib/decimate/select-recost.ts b/src/lib/decimate/select-recost.ts index 115b6a09..04b137a3 100644 --- a/src/lib/decimate/select-recost.ts +++ b/src/lib/decimate/select-recost.ts @@ -138,14 +138,15 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { const mass = SC[o + 11]; W[i] = mass; mx[i] = SC[o]; my[i] = SC[o + 1]; mz[i] = SC[o + 2]; - // Cache Σ carries EPS on its diagonal; moments accumulate Σ without it - // (production mergeGroup member math). - M2[i6] = mass * (SC[o + 3] - EPS_COV); + // The cache Σ is unregularized (buildCostCache uses raw variances), so + // it feeds the moments directly — matching mergeGroup's member math. + // EPS_COV enters once, on the merged covariance in evalMerge. + M2[i6] = mass * SC[o + 3]; M2[i6 + 1] = mass * SC[o + 4]; M2[i6 + 2] = mass * SC[o + 5]; - M2[i6 + 3] = mass * (SC[o + 6] - EPS_COV); + M2[i6 + 3] = mass * SC[o + 6]; M2[i6 + 4] = mass * SC[o + 7]; - M2[i6 + 5] = mass * (SC[o + 8] - EPS_COV); + M2[i6 + 5] = mass * SC[o + 8]; baseW[i3] = mass * SC[o + 12]; baseW[i3 + 1] = mass * SC[o + 13]; baseW[i3 + 2] = mass * SC[o + 14]; diff --git a/src/lib/gpu/gpu-edge-cost.ts b/src/lib/gpu/gpu-edge-cost.ts index 0b108c2f..395b6883 100644 --- a/src/lib/gpu/gpu-edge-cost.ts +++ b/src/lib/gpu/gpu-edge-cost.ts @@ -175,10 +175,12 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { 2.0 * ai * am * bim * cIM - 2.0 * aj * am * bjm * cJM; - // Clamp float-noise negatives but preserve NaN/Inf (degenerate input → the - // host fails loud on no finite merges), then add the scale-free colour - // dissimilarity term: |Δbase|² = bni + bnj − 2·(bi·bj). - costs[bid] = select(E, 0.0, E < 0.0) + COLOR_WEIGHT * (bni + bnj - 2.0 * bij); + // Clamp float-noise negatives on both squared norms (the colour distance + // uses separately-rounded norms and can round slightly negative) but + // preserve NaN/Inf (degenerate input → the host fails loud on no finite + // merges). |Δbase|² = bni + bnj − 2·(bi·bj). + let dc2 = bni + bnj - 2.0 * bij; + costs[bid] = select(E, 0.0, E < 0.0) + COLOR_WEIGHT * select(dc2, 0.0, dc2 < 0.0); } `; From bc72b1f6d2fa969f70df55ed23aea02ac338a31e Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 11:51:49 +0100 Subject: [PATCH 03/19] latest --- src/cli/index.ts | 26 ++ src/lib/decimate/decimate-source.ts | 68 +++- src/lib/decimate/edge-cost-legacy.ts | 202 +++++++++++ src/lib/decimate/priority-legacy.ts | 383 ++++++++++++++++++++ src/lib/decimate/recost-core.ts | 281 ++++++++++++++ src/lib/decimate/select-legacy.ts | 154 ++++++++ src/lib/decimate/select-recost.ts | 522 ++++++++++++--------------- src/lib/gpu/gpu-edge-cost-legacy.ts | 493 +++++++++++++++++++++++++ src/lib/workers/tasks.ts | 52 +++ 9 files changed, 1873 insertions(+), 308 deletions(-) create mode 100644 src/lib/decimate/edge-cost-legacy.ts create mode 100644 src/lib/decimate/priority-legacy.ts create mode 100644 src/lib/decimate/recost-core.ts create mode 100644 src/lib/decimate/select-legacy.ts create mode 100644 src/lib/gpu/gpu-edge-cost-legacy.ts diff --git a/src/cli/index.ts b/src/cli/index.ts index b3b2b1bc..5c64c26c 100644 --- a/src/cli/index.ts +++ b/src/cli/index.ts @@ -1,4 +1,5 @@ import { lstat, mkdir, readFile as pathReadFile, unlink } from 'node:fs/promises'; +import { totalmem } from 'node:os'; import { basename, dirname, join, resolve } from 'node:path'; import process, { exit } from 'node:process'; import { parseArgs } from 'node:util'; @@ -66,6 +67,8 @@ interface CliOptions extends LibOptions { listGpus: boolean; deviceIdx: number; // -1 = auto, -2 = CPU, 0+ = GPU index scratchDir: string | undefined; // decimation spill location (default: output directory) + decimateMode: 'quality' | 'legacy'; + memoryBudgetBytes: number; // decimation residency policy ceiling (not an allocation) } const fileExists = async (filename: string) => { @@ -188,6 +191,8 @@ const cliOptionsConfig = { 'filter-box': { type: 'string', short: 'B', multiple: true }, 'filter-sphere': { type: 'string', short: 'S', multiple: true }, 'decimate': { type: 'string', short: 'd', multiple: true }, + 'decimate-mode': { type: 'string', default: 'quality' }, + 'memory-budget': { type: 'string' }, 'filter-cluster': { type: 'string', short: 'C', multiple: true }, 'filter-floaters': { type: 'string', short: 'F', multiple: true }, params: { type: 'string', short: 'p', multiple: true }, @@ -517,6 +522,18 @@ const parseArguments = async () => { listGpus: v['list-gpus'], deviceIdx, scratchDir: v['scratch-dir'], + decimateMode: (() => { + const m = v['decimate-mode']; + if (m !== 'quality' && m !== 'legacy') { + throw new Error(`Invalid --decimate-mode: ${m}. Must be 'quality' or 'legacy'.`); + } + return m; + })(), + // Residency policy ceiling for decimation (not an upfront allocation): + // default to half the machine's RAM, capped at 48 GiB. + memoryBudgetBytes: v['memory-budget'] !== undefined ? + Math.max(1, parseNumber(v['memory-budget'])) * 2 ** 30 : + Math.min(48 * 2 ** 30, Math.floor(totalmem() / 2)), lodSelect: v['select-lod'].split(',').filter(v => !!v).map(parseInteger), viewerSettingsJson: viewerSettingsPath && await readJsonFile(viewerSettingsPath), unbundled: v.unbundled, @@ -788,6 +805,13 @@ ACTIONS (executed in order; can be repeated) -V, --filter-value Keep Gaussians where ; cmp ∈ {lt,lte,gt,gte,eq,neq} -d, --decimate Simplify to n (or n%) Gaussians via merge-based decimation. + --decimate-mode quality (default): field-L2 cost + re-costed selection — best on + mixed-scale scenes (large gains on skies/distant structure). + legacy: pre-3.2 pipeline — faster, lower memory, and still slightly + better on scenes of uniformly-sized Gaussians (single objects). + --memory-budget Decimation residency policy ceiling (not an upfront allocation); + re-costed selection falls back to one-shot selection above it. + Default: min(48, half of system RAM). Must be the final action, and the output must be .ply --scratch-dir Directory for decimation spill files (deep targets on huge scenes). Default: the output file's directory @@ -1246,6 +1270,8 @@ const main = async () => { combined = await decimateSource(combined, pool, { targetCount: keepCount, createDevice: deviceCreator, + mode: options.decimateMode, + memoryBudgetBytes: options.memoryBudgetBytes, spill: { writeFs: new NodeFileSystem(), readFs: new NodeReadFileSystem(), diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index 2f39ff53..8c200395 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -5,7 +5,9 @@ import { createBlockProducerSource } from './block-producer'; import { mergeStream } from './merge-stream'; import { kdPartition, coherenceRuns, type ResidentPositions } from './partition'; import { runPriorityPass, HALO_CAP, type CandidateArrays } from './priority'; +import { runPriorityPassLegacy } from './priority-legacy'; import { selectMerges } from './select'; +import { selectMergesLegacy } from './select-legacy'; import { selectMergesRecosted, CACHE_STRIDE } from './select-recost'; import { compact, @@ -38,7 +40,7 @@ const DEFAULT_MEMORY_BUDGET = 48 * 2 ** 30; // cache (16 f32) + neighbour ids (k u32) + f64 cluster moments/colour/error + // union-find/chains/heap. Conservative round-up; used by the per-generation // gate that falls back to one-shot selectMerges when over budget. -const RECOST_BYTES_PER_GAUSSIAN = (k: number) => CACHE_STRIDE * 4 + k * 4 + 200; +const RECOST_BYTES_PER_GAUSSIAN = (k: number) => CACHE_STRIDE * 4 + k * 4 + 256; /** Coherence heuristic: gap (rows) merged into one run / runs-per-block considered scattered. */ const COHERENCE_GAP_ROWS = 64; @@ -68,6 +70,15 @@ type DecimateOptions = { spill?: DecimateSpill; /** Resident-memory budget driving the candidate-K and re-costed-selection policies (default 48 GiB). */ memoryBudgetBytes?: number; + /** + * Decimation algorithm. `'quality'` (default): field-L2 cost with the + * scale-free colour term and re-costed selection — the quality-study + * winner (large gains on mixed-scale scenes, e.g. skies). `'legacy'`: the + * pre-study pipeline (KL-style cost with full-SH colour term, uniform + * matching) — faster, lower memory, and still measurably better on scenes + * of uniformly-sized gaussians (single objects, uniform texture). + */ + mode?: 'quality' | 'legacy'; }; // Candidate-K policy: keep 4 when the resident estimate fits the budget, @@ -211,36 +222,55 @@ const decimateSource = async ( } } - const K = chooseK(N, budget); - const cand: CandidateArrays = { - idx: new Uint32Array(N * K).fill(0xFFFFFFFF), - cost: new Float32Array(N * K).fill(Infinity) - }; - // Re-costed selection (exact within-generation greedy) when its // resident state fits the budget alongside the base state; one-shot // selection otherwise. Gated per generation, so large scenes regain - // re-costing as soon as the cascade shrinks under the budget. + // re-costing as soon as the cascade shrinks under the budget. Legacy + // mode uses the pre-study pipeline throughout (no re-costing state). + const legacy = opts.mode === 'legacy'; + const K = chooseK(N, budget); const k = Math.min(KNN_K, Math.max(1, N - 1)); const baseBytes = residentInputBytes + N * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; - const recost = baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; - const cacheOut = recost ? new Float32Array(N * CACHE_STRIDE) : undefined; - const neighborsOut = recost ? new Uint32Array(N * k) : undefined; + const recost = !legacy && baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; + + // The splat cache and neighbour graph feed refresh rounds; shared + // memory lets those rounds run on the worker pool. + const sharedOk = recost && typeof SharedArrayBuffer !== 'undefined'; + const cand: CandidateArrays = { + idx: new Uint32Array(N * K).fill(0xFFFFFFFF), + cost: new Float32Array(N * K).fill(Infinity) + }; + const cacheOut = recost ? + new Float32Array(sharedOk ? new SharedArrayBuffer(N * CACHE_STRIDE * 4) : new ArrayBuffer(N * CACHE_STRIDE * 4)) : + undefined; + const neighborsOut = recost ? + new Uint32Array(sharedOk ? new SharedArrayBuffer(N * k * 4) : new ArrayBuffer(N * k * 4)) : + undefined; const priorityBar = logger.bar('computing merge priorities', N); - await runPriorityPass( - { source: src, pool, pos: positions, order, blocks, device, K, k, cacheOut, neighborsOut }, - cand, - n => priorityBar.tick(n) - ); + if (legacy) { + await runPriorityPassLegacy( + { source: src, pool, pos: positions, order, blocks, device, K, k }, + cand, + n => priorityBar.tick(n) + ); + } else { + await runPriorityPass( + { source: src, pool, pos: positions, order, blocks, device, K, k, cacheOut, neighborsOut }, + cand, + n => priorityBar.tick(n) + ); + } priorityBar.end(); const generationTarget = Math.max(targetCount, N - Math.floor(N / 2)); const needed = N - generationTarget; const selectSub = logger.group(recost ? 'Selecting merges (re-costed)' : 'Selecting merges'); - const selection = cacheOut ? - selectMergesRecosted({ cand, K, splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed }) : - selectMerges(cand, N, K, needed); + const selection = legacy ? + selectMergesLegacy(cand, N, K, needed) : + cacheOut ? + await selectMergesRecosted({ cand, K, splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed }) : + selectMerges(cand, N, K, needed); selectSub.end(); if (selection.removed === 0) { diff --git a/src/lib/decimate/edge-cost-legacy.ts b/src/lib/decimate/edge-cost-legacy.ts new file mode 100644 index 00000000..9ddf25f9 --- /dev/null +++ b/src/lib/decimate/edge-cost-legacy.ts @@ -0,0 +1,202 @@ +/** + * CPU edge-cost path for chunk-native decimation — a port of the legacy + * `computeEdgeCost` + `buildPerSplatCache` (CPU variant) over a + * {@link SplatView}. Formula-identical to the legacy implementation and to + * the `GpuEdgeCost` WGSL kernel; used as the no-device fallback and by the + * GPU parity tests. + * + * Engine-free. + */ + +import { + EPS_COV, + LOG2PI, + logAddExp, + sigmoid, + ellipsoidArea, + quatToRotmat, + sigmaFromRotVar, + det3, + gaussLogpdfDiagrot, + type SplatView, + type MergeScratch +} from './moment-match'; + +/** + * Per-splat derived quantities for the cost function (legacy + * `buildPerSplatCache`, forGpu = false). `mass` uses the cost-path epsilon + * (+1e-12), matching legacy exactly. + */ +type LegacyCostCache = { + R: Float32Array; + v: Float32Array; + invdiag: Float32Array; + logdet: Float32Array; + sigma: Float32Array; + mass: Float32Array; +}; + +const buildCostCacheLegacy = (view: SplatView): LegacyCostCache => { + const { geo } = view; + const n = geo.length / 8; + const R = new Float32Array(n * 9); + const v = new Float32Array(n * 3); + const invdiag = new Float32Array(n * 3); + const logdet = new Float32Array(n); + const sigma = new Float32Array(n * 9); + const mass = new Float32Array(n); + + for (let i = 0; i < n; i++) { + const i3 = 3 * i; + const i8 = 8 * i; + const i9 = 9 * i; + + const linAlpha = sigmoid(geo[i8 + 7]); + const sx = Math.max(Math.exp(geo[i8 + 4]), 1e-12); + const sy = Math.max(Math.exp(geo[i8 + 5]), 1e-12); + const sz = Math.max(Math.exp(geo[i8 + 6]), 1e-12); + + const vx = sx * sx + EPS_COV; + const vy = sy * sy + EPS_COV; + const vz = sz * sz + EPS_COV; + + v[i3] = vx; v[i3 + 1] = vy; v[i3 + 2] = vz; + invdiag[i3] = 1 / Math.max(vx, 1e-30); + invdiag[i3 + 1] = 1 / Math.max(vy, 1e-30); + invdiag[i3 + 2] = 1 / Math.max(vz, 1e-30); + logdet[i] = Math.log(Math.max(vx, 1e-30)) + Math.log(Math.max(vy, 1e-30)) + Math.log(Math.max(vz, 1e-30)); + + let qw = geo[i8], qx = geo[i8 + 1], qy = geo[i8 + 2], qz = geo[i8 + 3]; + const qn = Math.hypot(qw, qx, qy, qz); + const invq = 1 / Math.max(qn, 1e-12); + qw *= invq; qx *= invq; qy *= invq; qz *= invq; + + quatToRotmat(qw, qx, qy, qz, R, i9); + sigmaFromRotVar(R, i9, vx, vy, vz, sigma, i9); + + mass[i] = linAlpha * ellipsoidArea(sx, sy, sz) + 1e-12; + } + + return { R, v, invdiag, logdet, sigma, mass }; +}; + +/** + * Edge cost between splats `i` and `j` of the view: KL-style geometric term + * (single MC sample) + L2 over the color/SH coefficients. Legacy + * `computeEdgeCost`, verbatim. + * + * @param view - Splat columns. + * @param cache - Per-splat cache from {@link buildCostCacheLegacy}. + * @param i - First splat (view row). + * @param j - Second splat (view row). + * @param Z - MC samples (legacy: one sample, seed 0). + * @param scratch - Merge scratch (uses `sigm`). + * @returns The edge cost. + */ +const computeEdgeCostViewLegacy = ( + view: SplatView, + cache: LegacyCostCache, + i: number, + j: number, + Z: Float64Array[], + scratch: MergeScratch +): number => { + const { pos, color, colorDim } = view; + const i3 = 3 * i, j3 = 3 * j; + const i9 = 9 * i, j9 = 9 * j; + + const mux = pos[i3], muy = pos[i3 + 1], muz = pos[i3 + 2]; + const mvx = pos[j3], mvy = pos[j3 + 1], mvz = pos[j3 + 2]; + + const wi = cache.mass[i], wj = cache.mass[j]; + const W = wi + wj; + const Wsafe = W > 0 ? W : 1; + + let pi = wi / Wsafe; + pi = Math.max(1e-12, Math.min(1 - 1e-12, pi)); + const pj = 1 - pi; + const logPi = Math.log(pi); + const logPj = Math.log(pj); + + const mmx = pi * mux + pj * mvx; + const mmy = pi * muy + pj * mvy; + const mmz = pi * muz + pj * mvz; + + const dix = mux - mmx, diy = muy - mmy, diz = muz - mmz; + const djx = mvx - mmx, djy = mvy - mmy, djz = mvz - mmz; + + const sigm = scratch.sigm; + for (let a = 0; a < 9; a++) { + sigm[a] = pi * cache.sigma[i9 + a] + pj * cache.sigma[j9 + a]; + } + sigm[0] += pi * dix * dix + pj * djx * djx; + sigm[1] += pi * dix * diy + pj * djx * djy; + sigm[2] += pi * dix * diz + pj * djx * djz; + sigm[3] += pi * diy * dix + pj * djy * djx; + sigm[4] += pi * diy * diy + pj * djy * djy; + sigm[5] += pi * diy * diz + pj * djy * djz; + sigm[6] += pi * diz * dix + pj * djz * djx; + sigm[7] += pi * diz * diy + pj * djz * djy; + sigm[8] += pi * diz * diz + pj * djz * djz; + + sigm[1] = sigm[3] = 0.5 * (sigm[1] + sigm[3]); + sigm[2] = sigm[6] = 0.5 * (sigm[2] + sigm[6]); + sigm[5] = sigm[7] = 0.5 * (sigm[5] + sigm[7]); + sigm[0] += EPS_COV; + sigm[4] += EPS_COV; + sigm[8] += EPS_COV; + + const detm = Math.max(det3(sigm, 0), 1e-30); + const logdetm = Math.log(detm); + + const EpNegLogQ = 0.5 * (3 * LOG2PI + logdetm + 3); + + const stdix = Math.sqrt(Math.max(cache.v[i3], 0)); + const stdiy = Math.sqrt(Math.max(cache.v[i3 + 1], 0)); + const stdiz = Math.sqrt(Math.max(cache.v[i3 + 2], 0)); + const stdjx = Math.sqrt(Math.max(cache.v[j3], 0)); + const stdjy = Math.sqrt(Math.max(cache.v[j3 + 1], 0)); + const stdjz = Math.sqrt(Math.max(cache.v[j3 + 2], 0)); + + let sumLogpOnI = 0; + let sumLogpOnJ = 0; + + for (let s = 0; s < Z.length; s++) { + const z0 = Z[s][0], z1 = Z[s][1], z2 = Z[s][2]; + + const xix = mux + z0 * stdix * cache.R[i9 + 0] + z1 * stdiy * cache.R[i9 + 1] + z2 * stdiz * cache.R[i9 + 2]; + const xiy = muy + z0 * stdix * cache.R[i9 + 3] + z1 * stdiy * cache.R[i9 + 4] + z2 * stdiz * cache.R[i9 + 5]; + const xiz = muz + z0 * stdix * cache.R[i9 + 6] + z1 * stdiy * cache.R[i9 + 7] + z2 * stdiz * cache.R[i9 + 8]; + + const xjx = mvx + z0 * stdjx * cache.R[j9 + 0] + z1 * stdjy * cache.R[j9 + 1] + z2 * stdjz * cache.R[j9 + 2]; + const xjy = mvy + z0 * stdjx * cache.R[j9 + 3] + z1 * stdjy * cache.R[j9 + 4] + z2 * stdjz * cache.R[j9 + 5]; + const xjz = mvz + z0 * stdjx * cache.R[j9 + 6] + z1 * stdjy * cache.R[j9 + 7] + z2 * stdjz * cache.R[j9 + 8]; + + const logNiOnI = gaussLogpdfDiagrot(xix, xiy, xiz, mux, muy, muz, + cache.R, i9, cache.invdiag[i3], cache.invdiag[i3 + 1], cache.invdiag[i3 + 2], cache.logdet[i]); + const logNjOnI = gaussLogpdfDiagrot(xix, xiy, xiz, mvx, mvy, mvz, + cache.R, j9, cache.invdiag[j3], cache.invdiag[j3 + 1], cache.invdiag[j3 + 2], cache.logdet[j]); + sumLogpOnI += logAddExp(logPi + logNiOnI, logPj + logNjOnI); + + const logNiOnJ = gaussLogpdfDiagrot(xjx, xjy, xjz, mux, muy, muz, + cache.R, i9, cache.invdiag[i3], cache.invdiag[i3 + 1], cache.invdiag[i3 + 2], cache.logdet[i]); + const logNjOnJ = gaussLogpdfDiagrot(xjx, xjy, xjz, mvx, mvy, mvz, + cache.R, j9, cache.invdiag[j3], cache.invdiag[j3 + 1], cache.invdiag[j3 + 2], cache.logdet[j]); + sumLogpOnJ += logAddExp(logPi + logNiOnJ, logPj + logNjOnJ); + } + + const Ei = sumLogpOnI / Z.length; + const Ej = sumLogpOnJ / Z.length; + const EpLogp = pi * Ei + pj * Ej; + const geoCost = EpLogp + EpNegLogQ; + + let cSh = 0; + for (let c = 0; c < colorDim; c++) { + const d = color[i * colorDim + c] - color[j * colorDim + c]; + cSh += d * d; + } + + return geoCost + cSh; +}; + +export { buildCostCacheLegacy, computeEdgeCostViewLegacy, type LegacyCostCache }; diff --git a/src/lib/decimate/priority-legacy.ts b/src/lib/decimate/priority-legacy.ts new file mode 100644 index 00000000..a8700a2c --- /dev/null +++ b/src/lib/decimate/priority-legacy.ts @@ -0,0 +1,383 @@ +/** + * LEGACY priority pass — the pre-quality-study pipeline (KL-style cost with + * full-SH colour term, appearance-chunk GPU kernel), preserved verbatim from + * main for `--decimate-mode legacy`. Self-contained on purpose: shares only + * the KNN/partition machinery with the current pass. + */ +import { type GraphicsDevice } from 'playcanvas'; + +import { buildCostCacheLegacy, computeEdgeCostViewLegacy } from './edge-cost-legacy'; +import { collectBlock, verifyAndFixKnn, toGlobalNeighbors, KNN_FIXED, type BlockLocals } from './knn-blocks'; +import { KNN_SENTINEL } from './knn-core'; +import { createMergeScratch, makeGaussianSamples, sigmoid, ellipsoidArea, type SplatView } from './moment-match'; +import { type BlockRange, type ResidentPositions } from './partition'; +import { type ChunkData, type ChunkDataPool, type ChunkSource } from '../chunk'; +import { APP_CHUNK, GpuEdgeCostLegacy, type EdgeCostCacheLegacy } from '../gpu/gpu-edge-cost-legacy'; +import { GpuKnn } from '../gpu/gpu-knn'; +import { WorkerQueue } from '../workers'; + +/** Halo radius multiplier on the density-estimated k-NN radius. */ +const HALO_FACTOR = 2.5; + +/** Halo size cap as a multiple of a block's owned count (buffer-sizing bound). */ +const HALO_CAP = 1; + +/** + * Per-gaussian best-K merge candidates, the resident output of the priority + * pass. `idx[g * K + s]` is the global index of gaussian g's s-th cheapest + * candidate (0xFFFFFFFF when absent); `cost[g * K + s]` its cost (+Inf when + * absent). + */ +type CandidateArrays = { + idx: Uint32Array; + cost: Float32Array; +}; + +/** Everything the block passes need: baked single-LOD source + resident state. */ +type PriorityContext = { + source: ChunkSource; + pool: ChunkDataPool; + pos: ResidentPositions; + order: Uint32Array; + blocks: BlockRange[]; + device?: GraphicsDevice; + /** Candidates kept per gaussian (K). */ + K: number; + /** Neighbours per query (16). */ + k: number; +}; + +/** + * A block's gathered splat columns: owned rows first (block order), then the + * requested extra globals. Positions come from the resident arrays, never + * from the source. + */ +type BlockView = { + view: SplatView; + /** u32 `other` columns (extraDim per row), when requested and present. */ + other?: Uint32Array; + otherDim: number; + ownedCount: number; +}; + +/** + * Gather geometric + color (and optionally `other`) for a block's owned rows + * plus `extraGlobals`, into tight column arrays. Reads are batched at the + * pool's chunk size; owned and extra index lists must be sorted ascending + * for gather coalescing. + * + * @param ctx - The pass context. + * @param blockIdx - Which block. + * @param extraGlobals - Sorted out-of-block rows to append after the owned rows. + * @param includeOther - Also gather the `other` layer (merge pass only). + * @returns The gathered block view. + */ +const gatherBlockView = async ( + ctx: Pick, + blockIdx: number, + extraGlobals: Uint32Array, + includeOther = false +): Promise => { + const { source, pool, pos, order, blocks } = ctx; + const block = blocks[blockIdx]; + const owned = order.subarray(block.start, block.end); + const nOwned = owned.length; + const n = nOwned + extraGlobals.length; + const { layouts, availableLayers } = source.meta; + + const colorDim = layouts.color!.stride >> 2; + const wantOther = includeOther && availableLayers.has('other') && (layouts.other?.stride ?? 0) > 0; + const otherDim = wantOther ? layouts.other!.stride >> 2 : 0; + + const view: SplatView = { + pos: new Float32Array(n * 3), + geo: new Float32Array(n * 8), + color: new Float32Array(n * colorDim), + colorDim + }; + const other = wantOther ? new Uint32Array(n * otherDim) : undefined; + + const readInto = async (indices: Uint32Array, rowBase: number): Promise => { + const batch = pool.chunkSize; + for (let off = 0; off < indices.length; off += batch) { + const count = Math.min(batch, indices.length - off); + const geoCd = pool.acquire('geometric', layouts.geometric!, count); + const colCd = pool.acquire('color', layouts.color!, count); + const othCd: ChunkData | undefined = wantOther ? pool.acquire('other', layouts.other!, count) : undefined; + await source.read({ + indices, + indexOffset: off, + count, + geometric: geoCd, + color: colCd, + other: othCd + }); + view.geo.set(new Float32Array(geoCd.data, 0, count * 8), (rowBase + off) * 8); + view.color.set(new Float32Array(colCd.data, 0, count * colorDim), (rowBase + off) * colorDim); + if (othCd) other!.set(new Uint32Array(othCd.data, 0, count * otherDim), (rowBase + off) * otherDim); + geoCd.release(); + colCd.release(); + othCd?.release(); + } + }; + + await readInto(owned, 0); + await readInto(extraGlobals, nOwned); + + for (let i = 0; i < nOwned; i++) { + const g = owned[i]; + view.pos[i * 3] = pos.x[g]; + view.pos[i * 3 + 1] = pos.y[g]; + view.pos[i * 3 + 2] = pos.z[g]; + } + for (let i = 0; i < extraGlobals.length; i++) { + const g = extraGlobals[i]; + const r = nOwned + i; + view.pos[r * 3] = pos.x[g]; + view.pos[r * 3 + 1] = pos.y[g]; + view.pos[r * 3 + 2] = pos.z[g]; + } + + return { view, other, otherDim, ownedCount: nOwned }; +}; + +// Binary search `g` in the sorted array; -1 when absent. +const indexOfSorted = (sorted: Uint32Array, g: number): number => { + let lo = 0, hi = sorted.length - 1; + while (lo <= hi) { + const mid = (lo + hi) >> 1; + const v = sorted[mid]; + if (v === g) return mid; + if (v < g) lo = mid + 1; + else hi = mid - 1; + } + return -1; +}; + +// Pack the block view into the GpuEdgeCostLegacy cache layout (legacy packing: +// posScalars 8-wide, rotR from normalized quats, appearance in ≤APP_CHUNK +// column chunks with live-width strides). +const packGpuCacheLegacy = (view: SplatView): EdgeCostCacheLegacy => { + const { pos, geo, color, colorDim } = view; + const n = geo.length / 8; + const posScalars = new Float32Array(n * 8); + const rotR = new Float32Array(n * 9); + const rot = new Float32Array(9); + + for (let i = 0; i < n; i++) { + const i8 = i * 8; + const o = i * 8; + const linAlpha = sigmoid(geo[i8 + 7]); + const sx = Math.max(Math.exp(geo[i8 + 4]), 1e-12); + const sy = Math.max(Math.exp(geo[i8 + 5]), 1e-12); + const sz = Math.max(Math.exp(geo[i8 + 6]), 1e-12); + const vx = sx * sx + 1e-8; + const vy = sy * sy + 1e-8; + const vz = sz * sz + 1e-8; + posScalars[o] = pos[i * 3]; + posScalars[o + 1] = pos[i * 3 + 1]; + posScalars[o + 2] = pos[i * 3 + 2]; + posScalars[o + 3] = linAlpha * ellipsoidArea(sx, sy, sz) + 1e-12; + posScalars[o + 4] = Math.log(Math.max(vx, 1e-30)) + Math.log(Math.max(vy, 1e-30)) + Math.log(Math.max(vz, 1e-30)); + posScalars[o + 5] = vx; + posScalars[o + 6] = vy; + posScalars[o + 7] = vz; + + let qw = geo[i8], qx = geo[i8 + 1], qy = geo[i8 + 2], qz = geo[i8 + 3]; + const invq = 1 / Math.max(Math.hypot(qw, qx, qy, qz), 1e-12); + qw *= invq; qx *= invq; qy *= invq; qz *= invq; + const xx = qx * qx, yy = qy * qy, zz = qz * qz; + const wx = qw * qx, wy = qw * qy, wz = qw * qz; + const xy = qx * qy, xz = qx * qz, yz = qy * qz; + rot[0] = 1 - 2 * (yy + zz); rot[1] = 2 * (xy - wz); rot[2] = 2 * (xz + wy); + rot[3] = 2 * (xy + wz); rot[4] = 1 - 2 * (xx + zz); rot[5] = 2 * (yz - wx); + rot[6] = 2 * (xz - wy); rot[7] = 2 * (yz + wx); rot[8] = 1 - 2 * (xx + yy); + rotR.set(rot, i * 9); + } + + const numChunks = Math.ceil(colorDim / APP_CHUNK); + const appChunks: Float32Array[] = []; + for (let ch = 0; ch < numChunks; ch++) { + const kStart = ch * APP_CHUNK; + const width = Math.min(APP_CHUNK, colorDim - kStart); + const chunk = new Float32Array(n * width); + for (let s = 0; s < n; s++) { + const dst = s * width; + const src = s * colorDim + kStart; + for (let kk = 0; kk < width; kk++) chunk[dst + kk] = color[src + kk]; + } + appChunks.push(chunk); + } + + return { posScalars, rotR, appChunks, numAppCols: colorDim, numSplats: n }; +}; + +/** + * The priority pass (heavy read 1): per block — exact global KNN, edge costs + * for each owned gaussian's k neighbours, reduction to the best K candidates + * — written into the resident candidate arrays. + * + * @param ctx - The pass context. + * @param cand - Preallocated candidate arrays (`N*K`), filled per block. + * @param tick - Optional progress callback (owned gaussians completed). + */ +const runPriorityPassLegacy = async ( + ctx: PriorityContext, + cand: CandidateArrays, + tick?: (n: number) => void +): Promise => { + const { pos, order, blocks, device, K, k } = ctx; + const Z = makeGaussianSamples(1, 0); + const z = new Float32Array([Z[0][0], Z[0][1], Z[0][2]]); + const colorDim = ctx.source.meta.layouts.color!.stride >> 2; + + let maxOwned = 0; + for (const b of blocks) maxOwned = Math.max(maxOwned, b.end - b.start); + const maxLocalN = maxOwned * (1 + HALO_CAP); + + let gpuKnn: GpuKnn | undefined; + let gpuCost: GpuEdgeCostLegacy | undefined; + let gpuCostCapacity = maxLocalN; + + // 1-deep prefetch: the next block's halo collection + tree build runs + // while the current block computes. GpuKnn executions share one set of + // buffers, so they are serialized through `gpuKnnQueue` — the prefetched + // block's KNN starts only after the current block's has finished. + let gpuKnnQueue: Promise = Promise.resolve(); + type Prepared = { locals: BlockLocals; nb: Promise }; + const prepare = (bi: number): Prepared => { + const locals = collectBlock(pos, order, blocks, bi, k, HALO_FACTOR, HALO_CAP); + const copy = locals.positions.slice(); + if (device) { + const treePromise = WorkerQueue.run('flattenKdTree', { positions: copy }, [copy.buffer as ArrayBuffer]); + const out = new Uint32Array(locals.ownedCount * k); + const run = Promise.all([treePromise, gpuKnnQueue]).then(([flat]) => { + return gpuKnn!.execute(flat, locals.positions, locals.ids.length, locals.ownedCount, out); + }); + gpuKnnQueue = run.catch(() => { /* surfaced by the awaiting block */ }); + return { locals, nb: run.then(() => out) }; + } + const nb = WorkerQueue.run('knnBlock', { positions: copy, ownedCount: locals.ownedCount, k }, [copy.buffer as ArrayBuffer]); + return { locals, nb }; + }; + + try { + if (device) { + gpuKnn = new GpuKnn(device, maxLocalN, k); + gpuCost = new GpuEdgeCostLegacy(device, maxLocalN, maxOwned * k, colorDim); + } + + let next: Prepared | null = blocks.length > 0 ? prepare(0) : null; + + for (let bi = 0; bi < blocks.length; bi++) { + const { locals, nb: nbPromise } = next!; + next = bi + 1 < blocks.length ? prepare(bi + 1) : null; + + const nOwned = locals.ownedCount; + const owned = order.subarray(blocks[bi].start, blocks[bi].end); + const nbLocal = await nbPromise; + const nbGlobal = toGlobalNeighbors(locals, nbLocal); + verifyAndFixKnn(pos, order, blocks, bi, locals, k, nbGlobal, nbLocal); + + // Externals: referenced rows outside the owned range (halo members + // and verification-fixed neighbours), sorted for the gather. + const extRow = new Map(); + for (let s = 0; s < nOwned * k; s++) { + const l = nbLocal[s]; + if (l === KNN_SENTINEL || l < nOwned) continue; + const g = nbGlobal[s]; + if (l !== KNN_FIXED) { + if (!extRow.has(g)) extRow.set(g, 0); + } else if (indexOfSorted(owned, g) < 0 && !extRow.has(g)) { + extRow.set(g, 0); + } + } + const extraGlobals = Uint32Array.from(extRow.keys()).sort(); + for (let i = 0; i < extraGlobals.length; i++) extRow.set(extraGlobals[i], nOwned + i); + + // Verification-fixed externals are not bounded by the halo cap, so + // a pathological block's view can exceed the preallocated cost + // buffers — grow them to the actual view size when that happens + // (rare; costs one reallocation). + const viewN = nOwned + extraGlobals.length; + if (gpuCost && viewN > gpuCostCapacity) { + gpuCost.destroy(); + gpuCostCapacity = Math.ceil(viewN * 1.1); + gpuCost = new GpuEdgeCostLegacy(device!, gpuCostCapacity, maxOwned * k, colorDim); + } + + const { view } = await gatherBlockView(ctx, bi, extraGlobals); + + // Edge lists in owned-major order (view-local endpoints). + const edgeI = new Uint32Array(nOwned * k); + const edgeJ = new Uint32Array(nOwned * k); + const edgeNb = new Uint32Array(nOwned * k); // global neighbour per edge + const edgeOf = new Uint32Array(nOwned + 1); // CSR into the edge list per owned row + let e = 0; + for (let qi = 0; qi < nOwned; qi++) { + edgeOf[qi] = e; + for (let s = 0; s < k; s++) { + const l = nbLocal[qi * k + s]; + if (l === KNN_SENTINEL) continue; + const g = nbGlobal[qi * k + s]; + let row: number; + if (l !== KNN_FIXED) { + row = l < nOwned ? l : extRow.get(g)!; + } else { + const oi = indexOfSorted(owned, g); + row = oi >= 0 ? oi : extRow.get(g)!; + } + edgeI[e] = qi; + edgeJ[e] = row; + edgeNb[e] = g; + e++; + } + } + edgeOf[nOwned] = e; + + const costs = new Float32Array(e); + if (device) { + await gpuCost!.execute(packGpuCacheLegacy(view), edgeI.subarray(0, e), edgeJ.subarray(0, e), z, costs); + } else { + const cache = buildCostCacheLegacy(view); + const scratch = createMergeScratch(); + for (let i = 0; i < e; i++) { + costs[i] = computeEdgeCostViewLegacy(view, cache, edgeI[i], edgeJ[i], Z, scratch); + } + } + + // Reduce to best K candidates per owned gaussian (ascending by cost). + const bestIdx = new Uint32Array(K); + const bestCost = new Float64Array(K); + for (let qi = 0; qi < nOwned; qi++) { + let size = 0; + for (let s = edgeOf[qi]; s < edgeOf[qi + 1]; s++) { + const c = costs[s]; + if (!Number.isFinite(c)) continue; + if (size === K && c >= bestCost[K - 1]) continue; + let at = size < K ? size : K - 1; + while (at > 0 && bestCost[at - 1] > c) { + bestCost[at] = bestCost[at - 1]; + bestIdx[at] = bestIdx[at - 1]; + at--; + } + bestCost[at] = c; + bestIdx[at] = edgeNb[s]; + size = Math.min(size + 1, K); + } + const g = owned[qi]; + for (let s = 0; s < K; s++) { + cand.idx[g * K + s] = s < size ? bestIdx[s] : 0xFFFFFFFF; + cand.cost[g * K + s] = s < size ? bestCost[s] : Infinity; + } + } + + tick?.(nOwned); + } + } finally { + gpuKnn?.destroy(); + gpuCost?.destroy(); + } +}; + +export { runPriorityPassLegacy }; diff --git a/src/lib/decimate/recost-core.ts b/src/lib/decimate/recost-core.ts new file mode 100644 index 00000000..e3c6ab9c --- /dev/null +++ b/src/lib/decimate/recost-core.ts @@ -0,0 +1,281 @@ +/** + * Re-costed selection evaluation kernel, shared between the main thread and + * worker threads (bulk refresh rounds). Operates on a plain view-of-state + * object whose arrays may be backed by SharedArrayBuffers; workers treat the + * state as read-only (find() does no path compression here). + * + * Cost definition matches the pairwise kernel: exact field-L2 of the merged + * cluster vs its generation-input members (splat cache rows) plus the + * scale-free DC colour term — see select-recost.ts for the orchestration. + * + * Engine-free; no allocation in the hot paths beyond small local scratch. + */ + +import { COLOR_WEIGHT } from './edge-cost-cpu'; +import { EPS_COV, ellipsoidArea } from './moment-match'; + +/** Floats per splat in the resident cache (packGpuCache layout). */ +export const CACHE_STRIDE = 16; + +const NO_CANDIDATE = 0xFFFFFFFF; +const NIL = 0xFFFFFFFF; + +/** Skip Gaussian products whose exponent bound exceeds this (e^-60 ≈ 9e-27). */ +const CULL_QUAD = 120; + +const PI_1_5 = Math.PI ** 1.5; +const TWO_PI_1_5 = (2 * Math.PI) ** 1.5; + +/** + * The resident selection state (allocated by select-recost, possibly on + * SharedArrayBuffers so bulk refreshes can run on worker threads). + */ +type RecostState = { + /** Per-splat cache, {@link CACHE_STRIDE} floats per generation-input splat. */ + SC: Float32Array; + /** Candidate ids per splat (D per row, NO_CANDIDATE padded) — the union of members' rows forms a cluster's candidate pool. */ + cands: Uint32Array; + /** Candidate ids per splat. */ + D: number; + /** Gaussian count. */ + N: number; + /** Max original members per group. */ + maxGroup: number; + // Union-find + cluster state (indexed by root). + parent: Uint32Array; + size: Uint32Array; + W: Float64Array; + mx: Float64Array; my: Float64Array; mz: Float64Array; + M2: Float64Array; + baseW: Float64Array; + Sself: Float64Array; + Err: Float64Array; + version: Uint32Array; + mHead: Uint32Array; + mNext: Uint32Array; +}; + +/** + * Read-only find (no path compression — safe on shared state in workers). + * + * @param parent - Union-find parent array. + * @param x - Element to resolve. + * @returns The set root. + */ +const findRO = (parent: Uint32Array, x: number): number => { + while (parent[x] !== x) x = parent[x]; + return x; +}; + +// ⟨G_a,G_b⟩ scaled by √|Σa|·√|Σb| for M = Σa+Σb (6 comps) and offset d. +const crossG = ( + sdAB: number, + m0: number, m1: number, m2: number, m3: number, m4: number, m5: number, + dx: number, dy: number, dz: number +): number => { + const c00 = m3 * m5 - m4 * m4; + const c01 = m2 * m4 - m1 * m5; + const c02 = m1 * m4 - m2 * m3; + const c11 = m0 * m5 - m2 * m2; + const c12 = m1 * m2 - m0 * m4; + const c22 = m0 * m3 - m1 * m1; + const det = Math.max(m0 * c00 + m1 * c01 + m2 * c02, 1e-60); + const quad = (c00 * dx * dx + c11 * dy * dy + c22 * dz * dz + + 2 * (c01 * dx * dy + c02 * dx * dz + c12 * dy * dz)) / det; + if (!(quad < CULL_QUAD)) return 0; + return TWO_PI_1_5 * sdAB / Math.sqrt(det) * Math.exp(-0.5 * quad); +}; + +// Smith closed-form eigenvalues of a symmetric 3×3 (6 comps), descending. +const eigOut = new Float64Array(3); +const eig3 = (m0: number, m1: number, m2: number, m3: number, m4: number, m5: number): void => { + const q = (m0 + m3 + m5) / 3; + const p1 = m1 * m1 + m2 * m2 + m4 * m4; + if (p1 <= 1e-30) { + eigOut[0] = Math.max(m0, m3, m5); + eigOut[2] = Math.min(m0, m3, m5); + eigOut[1] = m0 + m3 + m5 - eigOut[0] - eigOut[2]; + return; + } + const p2 = (m0 - q) * (m0 - q) + (m3 - q) * (m3 - q) + (m5 - q) * (m5 - q) + 2 * p1; + const p = Math.sqrt(p2 / 6); + const ip = 1 / p; + const b00 = (m0 - q) * ip, b11 = (m3 - q) * ip, b22 = (m5 - q) * ip; + const b01 = m1 * ip, b02 = m2 * ip, b12 = m4 * ip; + const detB = b00 * (b11 * b22 - b12 * b12) - b01 * (b01 * b22 - b12 * b02) + b02 * (b01 * b12 - b11 * b02); + let r = detB / 2; + r = r < -1 ? -1 : (r > 1 ? 1 : r); + const phi = Math.acos(r) / 3; + const e0 = q + 2 * p * Math.cos(phi); + const e2 = q + 2 * p * Math.cos(phi + (2 * Math.PI) / 3); + eigOut[0] = e0; eigOut[1] = 3 * q - e0 - e2; eigOut[2] = e2; +}; + +let gatherBuf = new Uint32Array(1 << 12); +const gatherMembers = (st: RecostState, root: number): number => { + const { mHead, mNext } = st; + let cnt = 0; + for (let m = mHead[root]; m !== NIL; m = mNext[m]) { + if (cnt === gatherBuf.length) { + const g = new Uint32Array(gatherBuf.length * 2); + g.set(gatherBuf); gatherBuf = g; + } + gatherBuf[cnt++] = m; + } + return cnt; +}; + +/** Outputs of {@link evalMergeCore} beyond the cost (reused object). */ +const evalOut = { E: 0, Scross: 0 }; + +/** + * Marginal cost ΔE of merging clusters A and B (exact field-L2 vs the + * generation-input members) plus the scale-free colour term. Fills + * {@link evalOut} with E(A∪B) and Scross(A,B). + * + * @param st - Selection state. + * @param A - First cluster root. + * @param B - Second cluster root. + * @returns The merge cost. + */ +const evalMergeCore = (st: RecostState, A: number, B: number): number => { + const { SC, W, mx, my, mz, M2, baseW, Sself, Err, mHead, mNext } = st; + const WA = W[A], WB = W[B], WC = WA + WB; + const iw = 1 / WC; + const mcx = (WA * mx[A] + WB * mx[B]) * iw; + const mcy = (WA * my[A] + WB * my[B]) * iw; + const mcz = (WA * mz[A] + WB * mz[B]) * iw; + const dax = mx[A] - mcx, day = my[A] - mcy, daz = mz[A] - mcz; + const dbx = mx[B] - mcx, dby = my[B] - mcy, dbz = mz[B] - mcz; + const a6 = A * 6, b6 = B * 6; + + const sm0 = (M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx) * iw + EPS_COV; + const sm1 = (M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby) * iw; + const sm2 = (M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz) * iw; + const sm3 = (M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby) * iw + EPS_COV; + const sm4 = (M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz) * iw; + const sm5 = (M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz) * iw + EPS_COV; + + const detm = Math.max( + sm0 * (sm3 * sm5 - sm4 * sm4) - sm1 * (sm1 * sm5 - sm4 * sm2) + sm2 * (sm1 * sm4 - sm3 * sm2), + 1e-60 + ); + const sdC = Math.sqrt(detm); + + eig3(sm0, sm1, sm2, sm3, sm4, sm5); + const s0 = Math.sqrt(Math.max(eigOut[0], 1e-18)); + const s1 = Math.sqrt(Math.max(eigOut[1], 1e-18)); + const s2 = Math.sqrt(Math.max(eigOut[2], 1e-18)); + const alphaC = Math.min(1, WC / Math.max(ellipsoidArea(s0, s1, s2), 1e-30)); + + const a3 = A * 3, b3 = B * 3; + const bc0 = (baseW[a3] + baseW[b3]) * iw; + const bc1 = (baseW[a3 + 1] + baseW[b3 + 1]) * iw; + const bc2 = (baseW[a3 + 2] + baseW[b3 + 2]) * iw; + const bn2C = bc0 * bc0 + bc1 * bc1 + bc2 * bc2; + + const selfM = alphaC * alphaC * bn2C * PI_1_5 * sdC; + + // ⟨Σ member fields, f_m⟩ over both chains. + let memfm = 0; + for (let pass = 0; pass < 2; pass++) { + for (let m = pass === 0 ? mHead[A] : mHead[B]; m !== NIL; m = mNext[m]) { + const o = m * CACHE_STRIDE; + const wgt = SC[o + 10] * alphaC * + (SC[o + 12] * bc0 + SC[o + 13] * bc1 + SC[o + 14] * bc2); + if (wgt === 0) continue; + memfm += wgt * crossG(SC[o + 9] * sdC, + SC[o + 3] + sm0, SC[o + 4] + sm1, SC[o + 5] + sm2, + SC[o + 6] + sm3, SC[o + 7] + sm4, SC[o + 8] + sm5, + SC[o] - mcx, SC[o + 1] - mcy, SC[o + 2] - mcz); + } + } + + // Scross(A,B) = Σ_{a∈A,b∈B}⟨f_a,f_b⟩ with distance culling. + const na = gatherMembers(st, A); + let scross = 0; + for (let b = mHead[B]; b !== NIL; b = mNext[b]) { + const ob = b * CACHE_STRIDE; + const bxp = SC[ob], byp = SC[ob + 1], bzp = SC[ob + 2]; + const trb = SC[ob + 3] + SC[ob + 6] + SC[ob + 8]; + const alb = SC[ob + 10], sdb = SC[ob + 9]; + const cb0 = SC[ob + 12], cb1 = SC[ob + 13], cb2 = SC[ob + 14]; + for (let t = 0; t < na; t++) { + const a = gatherBuf[t]; + const oa = a * CACHE_STRIDE; + const dx = SC[oa] - bxp, dy = SC[oa + 1] - byp, dz = SC[oa + 2] - bzp; + const d2 = dx * dx + dy * dy + dz * dz; + if (d2 > CULL_QUAD * (SC[oa + 3] + SC[oa + 6] + SC[oa + 8] + trb)) continue; + const wgt = SC[oa + 10] * alb * + (SC[oa + 12] * cb0 + SC[oa + 13] * cb1 + SC[oa + 14] * cb2); + if (wgt === 0) continue; + scross += wgt * crossG(SC[oa + 9] * sdb, + SC[oa + 3] + SC[ob + 3], SC[oa + 4] + SC[ob + 4], SC[oa + 5] + SC[ob + 5], + SC[oa + 6] + SC[ob + 6], SC[oa + 7] + SC[ob + 7], SC[oa + 8] + SC[ob + 8], + dx, dy, dz); + } + } + + const E = Sself[A] + Sself[B] + 2 * scross - 2 * memfm + selfM; + evalOut.E = E; + evalOut.Scross = scross; + + const iwA = 1 / W[A], iwB = 1 / W[B]; + const d0 = baseW[a3] * iwA - baseW[b3] * iwB; + const d1 = baseW[a3 + 1] * iwA - baseW[b3 + 1] * iwB; + const d2c = baseW[a3 + 2] * iwA - baseW[b3 + 2] * iwB; + return (E - Err[A] - Err[B]) + COLOR_WEIGHT * (d0 * d0 + d1 * d1 + d2c * d2c); +}; + +/** Result of {@link bestEdgeFor} (reused object). */ +const bestOut = { partner: -1, vb: 0, cost: 0, E: 0, S: 0 }; + +/** + * Compute the cheapest legal merge for `root`: candidates are the live + * clusters owning any member's candidate ids, deduplicated linearly (pools + * are small), respecting the group cap. + * + * @param st - Selection state. + * @param root - Cluster root to refresh. + * @returns True when a legal candidate exists (result in {@link bestOut}). + */ +const candScratch = new Uint32Array(256); + +const bestEdgeFor = (st: RecostState, root: number): boolean => { + const { cands, D, parent, size, version, mHead, mNext, maxGroup } = st; + const sz = size[root]; + // Derive + dedup candidates (linear scan — pools are ≤ a few dozen). + let cnt = 0; + const cbuf = candScratch; + for (let m = mHead[root]; m !== NIL; m = mNext[m]) { + const base = m * D; + for (let s = 0; s < D; s++) { + const nb = cands[base + s]; + if (nb === NO_CANDIDATE) continue; + const r = findRO(parent, nb); + if (r === root) continue; + let seen = false; + for (let t = 0; t < cnt; t++) { + if (cbuf[t] === r) { + seen = true; break; + } + } + if (seen) continue; + if (cnt < cbuf.length) cbuf[cnt++] = r; + } + } + let bc = Infinity, bp = -1, bv = 0, bE = 0, bS = 0; + for (let t = 0; t < cnt; t++) { + const c = cbuf[t]; + if (sz + size[c] > maxGroup) continue; + const d = evalMergeCore(st, root, c); + if (d < bc) { + bc = d; bp = c; bv = version[c]; bE = evalOut.E; bS = evalOut.Scross; + } + } + if (bp < 0) return false; + bestOut.partner = bp; bestOut.vb = bv; bestOut.cost = bc; bestOut.E = bE; bestOut.S = bS; + return true; +}; + +export { evalMergeCore, bestEdgeFor, findRO, evalOut, bestOut, NO_CANDIDATE, type RecostState }; diff --git a/src/lib/decimate/select-legacy.ts b/src/lib/decimate/select-legacy.ts new file mode 100644 index 00000000..9d6c5318 --- /dev/null +++ b/src/lib/decimate/select-legacy.ts @@ -0,0 +1,154 @@ +/** + * Global merge selection over the resident candidate arrays: bucketed greedy + * disjoint matching in cost order, plus chain closure that attaches + * still-unmatched gaussians to a candidate's group (≤3, relief cap 4) so a + * 50% target completes in one generation instead of a mop-up pass. + * + * Bucket walk ≈ exact cost-sorted greedy up to 1/SELECT_BUCKETS-of-range + * quantization — the selection-semantics match with the legacy algorithm. + * + * Engine-free; pure resident-array computation, no IO. + */ + +import { type CandidateArrays } from './priority'; + +/** Cost-histogram buckets for the ordered greedy walk. */ +const SELECT_BUCKETS = 1024; + +const NO_CANDIDATE = 0xFFFFFFFF; + +/** + * The selected merge groups. + * + * Groups are CSR-packed: group g's members are + * `groupMembers[groupOffsets[g] .. groupOffsets[g+1])` (global gaussian + * ids). `memberGroup[i]` is gaussian i's group or -1 (survivor). + * `groupMin[g]` is the group's minimum member id — the merge stream emits + * each group exactly once, at that member's position in block order. + * `removed` is the gaussian-count reduction achieved (Σ size-1); callers + * compare it to `mergesNeeded` and decide stall vs next generation. + */ +type SelectionResult = { + groupOffsets: Uint32Array; + groupMembers: Uint32Array; + memberGroup: Int32Array; + groupMin: Uint32Array; + mergedGroups: number; + removed: number; +}; + +/** + * Select merges from the candidate arrays. + * + * @param cand - Per-gaussian best-K candidates from the priority pass. + * @param N - Gaussian count. + * @param K - Candidates per gaussian. + * @param mergesNeeded - Target removal count for this generation. + * @returns The selection. + */ +const selectMergesLegacy = (cand: CandidateArrays, N: number, K: number, mergesNeeded: number): SelectionResult => { + const E = N * K; + + // Pass 1: finite cost range. + let lo = Infinity, hi = -Infinity; + for (let e = 0; e < E; e++) { + const c = cand.cost[e]; + if (Number.isFinite(c)) { + if (c < lo) lo = c; + if (c > hi) hi = c; + } + } + const span = hi > lo ? hi - lo : 1; + const bucketOf = (c: number) => Math.min(SELECT_BUCKETS - 1, Math.floor(((c - lo) / span) * SELECT_BUCKETS)); + + // Counting sort of finite candidate entries by bucket. + const counts = new Uint32Array(SELECT_BUCKETS + 1); + for (let e = 0; e < E; e++) { + if (Number.isFinite(cand.cost[e])) counts[bucketOf(cand.cost[e]) + 1]++; + } + for (let b = 0; b < SELECT_BUCKETS; b++) counts[b + 1] += counts[b]; + const orderE = new Uint32Array(counts[SELECT_BUCKETS]); + const cursor = counts.slice(0, SELECT_BUCKETS); + for (let e = 0; e < E; e++) { + if (Number.isFinite(cand.cost[e])) orderE[cursor[bucketOf(cand.cost[e])]++] = e; + } + + const memberGroup = new Int32Array(N).fill(-1); + + // Primary greedy: pair free endpoints, cheapest bucket first. + const maxPairs = Math.max(0, mergesNeeded); + const pairA = new Uint32Array(maxPairs); + const pairB = new Uint32Array(maxPairs); + let pairs = 0; + let removed = 0; + for (let t = 0; t < orderE.length && removed < mergesNeeded; t++) { + const e = orderE[t]; + const j = cand.idx[e]; + if (j === NO_CANDIDATE) continue; + const i = (e / K) | 0; + if (memberGroup[i] !== -1 || memberGroup[j] !== -1) continue; + memberGroup[i] = pairs; + memberGroup[j] = pairs; + pairA[pairs] = i; + pairB[pairs] = j; + pairs++; + removed++; + } + + // Chain closure: attach unmatched gaussians to a candidate's group, + // cheapest first, cap 3 — then a relief walk at cap 4. After the primary + // walk every free gaussian's candidates are all matched (else the pair + // would have been taken), so closure can almost always attach. + const groupSize = new Uint32Array(pairs).fill(2); + const joinMember = new Uint32Array(Math.max(0, mergesNeeded - removed)); + const joinGroup = new Uint32Array(joinMember.length); + let joins = 0; + for (const cap of [3, 4]) { + if (removed >= mergesNeeded) break; + for (let t = 0; t < orderE.length && removed < mergesNeeded; t++) { + const e = orderE[t]; + const j = cand.idx[e]; + if (j === NO_CANDIDATE) continue; + const i = (e / K) | 0; + if (memberGroup[i] !== -1) continue; + const g = memberGroup[j]; + if (g === -1 || groupSize[g] >= cap) continue; + memberGroup[i] = g; + groupSize[g]++; + joinMember[joins] = i; + joinGroup[joins] = g; + joins++; + removed++; + } + } + + // CSR assembly. + const G = pairs; + const groupOffsets = new Uint32Array(G + 1); + for (let g = 0; g < G; g++) groupOffsets[g + 1] = groupOffsets[g] + groupSize[g]; + const groupMembers = new Uint32Array(groupOffsets[G]); + const fill = new Uint32Array(G); + for (let g = 0; g < G; g++) { + const o = groupOffsets[g]; + groupMembers[o] = pairA[g]; + groupMembers[o + 1] = pairB[g]; + fill[g] = 2; + } + for (let t = 0; t < joins; t++) { + const g = joinGroup[t]; + groupMembers[groupOffsets[g] + fill[g]] = joinMember[t]; + fill[g]++; + } + const groupMin = new Uint32Array(G); + for (let g = 0; g < G; g++) { + let min = groupMembers[groupOffsets[g]]; + for (let m = groupOffsets[g] + 1; m < groupOffsets[g + 1]; m++) { + if (groupMembers[m] < min) min = groupMembers[m]; + } + groupMin[g] = min; + } + + return { groupOffsets, groupMembers, memberGroup, groupMin, mergedGroups: G, removed }; +}; + +export { selectMergesLegacy, SELECT_BUCKETS, type SelectionResult }; diff --git a/src/lib/decimate/select-recost.ts b/src/lib/decimate/select-recost.ts index 04b137a3..a548b4df 100644 --- a/src/lib/decimate/select-recost.ts +++ b/src/lib/decimate/select-recost.ts @@ -1,45 +1,43 @@ /** - * Re-costed merge selection: exact greedy agglomeration within a generation. + * Re-costed merge selection: exact greedy agglomeration within a generation, + * parallelized as commit/refresh rounds. * * Where {@link selectMerges} consumes the priority pass's pairwise costs * one-shot (costs go stale as groups form), this selection re-evaluates a - * cluster's candidates after every merge against its CURRENT moments, so the - * cheapest-first order is always true. Validated against the reference - * implementation (tools/decimate-exact.mjs) as the production form of the - * study winner: within-generation exact re-costing closes the remaining - * ~1–2 dB at fine levels vs one-shot selection. + * cluster's candidates after it changes, so merges always execute at costs + * evaluated against current state. To make the (dominant) evaluation work + * parallel, refreshes are batched into rounds: * - * Cost of a cluster = exact field-L2 vs its generation-input members (their - * parameters come from the resident splat cache emitted by the priority - * pass), plus the scale-free DC colour term ({@link COLOR_WEIGHT}) — the same - * cost definition as the pairwise kernel, so the precomputed candidate costs - * seed the heap directly. + * 1. Drain the heap, committing every entry that is still valid; clusters + * that changed (merged) or whose best partner changed are queued. + * 2. Bulk-evaluate all queued clusters' best edges — on the worker pool + * when the state is SharedArrayBuffer-backed, inline otherwise — and + * push the results. + * 3. Repeat until the generation target is reached or no merges remain. * - * Memory: ~200 B per gaussian resident (splat cache, neighbour ids, cluster - * moments, heap). decimate-source gates this path by memory budget and falls - * back to {@link selectMerges} when it does not fit. + * Relative to strictly-eager greedy, refreshed clusters re-enter the heap at + * round boundaries instead of immediately; commits still only execute at + * costs validated against the exact current state (version checks), so the + * deviation is confined to near-tie ordering (re-certified on the evaluation + * harness). + * + * Cost definition and evaluation kernel live in {@link recost-core} (shared + * with the worker task). Seeds come from the priority pass's candidate + * arrays; refresh candidate pools come from the persisted neighbour graph. + * + * decimate-source gates this path by memory budget (~330 B per gaussian + * resident) and falls back to {@link selectMerges} when over. * * Engine-free; pure resident-array computation, no IO. */ -import { COLOR_WEIGHT } from './edge-cost-cpu'; -import { KNN_SENTINEL } from './knn-core'; -import { EPS_COV, ellipsoidArea } from './moment-match'; import { type CandidateArrays } from './priority'; +import { evalMergeCore, bestEdgeFor, evalOut, bestOut, NO_CANDIDATE, CACHE_STRIDE, type RecostState } from './recost-core'; import { MAX_GROUP, type SelectionResult } from './select'; +import { WorkerQueue } from '../workers'; -/** Floats per splat in the resident cache (packGpuCache layout). */ -export const CACHE_STRIDE = 16; - -const NO_CANDIDATE = 0xFFFFFFFF; const NIL = 0xFFFFFFFF; -/** Skip Gaussian products whose exponent bound exceeds this (e^-60 ≈ 9e-27). */ -const CULL_QUAD = 120; - -const PI_1_5 = Math.PI ** 1.5; -const TWO_PI_1_5 = (2 * Math.PI) ** 1.5; - /** Inputs for {@link selectMergesRecosted}. */ type RecostInputs = { /** Per-gaussian best-K candidates from the priority pass (seed costs include the colour term). */ @@ -48,11 +46,17 @@ type RecostInputs = { K: number; /** * Resident per-splat cache, {@link CACHE_STRIDE} floats per splat - * (pos 3, Σ 6 [xx xy xz yy yz zz], √|Σ|, α, mass, base colour 3, |base|²) - * — the packGpuCache row layout, persisted by the priority pass. + * (packGpuCache row layout), persisted by the priority pass. + * SharedArrayBuffer-backed for the parallel path. */ splatCache: Float32Array; - /** Global neighbour ids per splat (D per row, KNN_SENTINEL padded). */ + /** + * Global neighbour ids per splat (D per row, sentinel padded) — the union + * of a cluster's members' rows forms its refresh candidate pool (the full + * neighbour graph, not just the seed candidates: top-K lists collapse onto + * shared hubs in degenerate/coincident regions and would starve refreshes). + * SharedArrayBuffer-backed for the parallel path. + */ neighbors: Uint32Array; /** Neighbours per splat. */ D: number; @@ -62,85 +66,84 @@ type RecostInputs = { mergesNeeded: number; }; -// ⟨G_a,G_b⟩ scaled by √|Σa|·√|Σb| for M = Σa+Σb (6 comps) and offset d. -const crossG = ( - sdAB: number, - m0: number, m1: number, m2: number, m3: number, m4: number, m5: number, - dx: number, dy: number, dz: number -): number => { - const c00 = m3 * m5 - m4 * m4; - const c01 = m2 * m4 - m1 * m5; - const c02 = m1 * m4 - m2 * m3; - const c11 = m0 * m5 - m2 * m2; - const c12 = m1 * m2 - m0 * m4; - const c22 = m0 * m3 - m1 * m1; - const det = Math.max(m0 * c00 + m1 * c01 + m2 * c02, 1e-60); - const quad = (c00 * dx * dx + c11 * dy * dy + c22 * dz * dz + - 2 * (c01 * dx * dy + c02 * dx * dz + c12 * dy * dz)) / det; - if (!(quad < CULL_QUAD)) return 0; - return TWO_PI_1_5 * sdAB / Math.sqrt(det) * Math.exp(-0.5 * quad); -}; - -// Smith closed-form eigenvalues of a symmetric 3×3 (6 comps), descending. -const eig3 = (m0: number, m1: number, m2: number, m3: number, m4: number, m5: number, out: Float64Array): void => { - const q = (m0 + m3 + m5) / 3; - const p1 = m1 * m1 + m2 * m2 + m4 * m4; - if (p1 <= 1e-30) { - out[0] = Math.max(m0, m3, m5); - out[2] = Math.min(m0, m3, m5); - out[1] = m0 + m3 + m5 - out[0] - out[2]; - return; - } - const p2 = (m0 - q) * (m0 - q) + (m3 - q) * (m3 - q) + (m5 - q) * (m5 - q) + 2 * p1; - const p = Math.sqrt(p2 / 6); - const ip = 1 / p; - const b00 = (m0 - q) * ip, b11 = (m3 - q) * ip, b22 = (m5 - q) * ip; - const b01 = m1 * ip, b02 = m2 * ip, b12 = m4 * ip; - const detB = b00 * (b11 * b22 - b12 * b12) - b01 * (b01 * b22 - b12 * b02) + b02 * (b01 * b12 - b11 * b02); - let r = detB / 2; - r = r < -1 ? -1 : (r > 1 ? 1 : r); - const phi = Math.acos(r) / 3; - const e0 = q + 2 * p * Math.cos(phi); - const e2 = q + 2 * p * Math.cos(phi + (2 * Math.PI) / 3); - out[0] = e0; out[1] = 3 * q - e0 - e2; out[2] = e2; +/** + * Allocate a typed array, on shared memory when requested. + * + * @param ctor - Typed-array constructor. + * @param bytes - Buffer size in bytes. + * @param shared - Back with a SharedArrayBuffer (worker-visible state). + * @returns The array view. + */ +const alloc = ( + ctor: new (buffer: ArrayBufferLike) => T, bytes: number, shared: boolean +): T => { + return new ctor(shared ? new SharedArrayBuffer(bytes) : new ArrayBuffer(bytes)); }; /** - * Select merges with exact within-generation re-costing. + * Select merges with within-generation re-costing (round-parallel). * * @param inputs - See {@link RecostInputs}. * @returns The selection (same contract as {@link selectMerges}). */ -const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { +const selectMergesRecosted = async (inputs: RecostInputs): Promise => { const { cand, K, splatCache: SC, neighbors, D, N, mergesNeeded } = inputs; + const shared = typeof SharedArrayBuffer !== 'undefined' && + SC.buffer instanceof SharedArrayBuffer && + neighbors.buffer instanceof SharedArrayBuffer && + !WorkerQueue.isInline; + // ---- Cluster state (indexed by union-find root). - const parent = new Uint32Array(N); - const size = new Uint32Array(N).fill(1); - const W = new Float64Array(N); - const mx = new Float64Array(N), my = new Float64Array(N), mz = new Float64Array(N); - const M2 = new Float64Array(N * 6); - const baseW = new Float64Array(N * 3); - const Sself = new Float64Array(N); - const Err = new Float64Array(N); - const version = new Uint32Array(N).fill(1); + const parent = alloc(Uint32Array, N * 4, shared); + const size = alloc(Uint32Array, N * 4, shared); + const W = alloc(Float64Array, N * 8, shared); + const mx = alloc(Float64Array, N * 8, shared); + const my = alloc(Float64Array, N * 8, shared); + const mz = alloc(Float64Array, N * 8, shared); + const M2 = alloc(Float64Array, N * 6 * 8, shared); + const baseW = alloc(Float64Array, N * 3 * 8, shared); + const Sself = alloc(Float64Array, N * 8, shared); + const Err = alloc(Float64Array, N * 8, shared); + const version = alloc(Uint32Array, N * 4, shared); + const mHead = alloc(Uint32Array, N * 4, shared); + const mNext = alloc(Uint32Array, N * 4, shared); + // Main-thread-only state. + const mTail = new Uint32Array(N); const lastSeq = new Uint32Array(N); - const mHead = new Uint32Array(N), mTail = new Uint32Array(N); - const mNext = new Uint32Array(N).fill(NIL); - const stamp = new Uint32Array(N); - let stampGen = 0; - let seqCounter = 0; - let liveCount = N; + const st: RecostState = { + SC, + cands: neighbors, + D, + N, + maxGroup: MAX_GROUP, + parent, + size, + W, + mx, + my, + mz, + M2, + baseW, + Sself, + Err, + version, + mHead, + mNext + }; + + size.fill(1); + version.fill(1); + mNext.fill(NIL); for (let i = 0; i < N; i++) { parent[i] = i; mHead[i] = i; mTail[i] = i; const o = i * CACHE_STRIDE, i6 = i * 6, i3 = i * 3; const mass = SC[o + 11]; W[i] = mass; mx[i] = SC[o]; my[i] = SC[o + 1]; mz[i] = SC[o + 2]; - // The cache Σ is unregularized (buildCostCache uses raw variances), so - // it feeds the moments directly — matching mergeGroup's member math. - // EPS_COV enters once, on the merged covariance in evalMerge. + // The cache Σ is unregularized; moments accumulate it directly + // (mergeGroup member math). EPS enters once, on the merged covariance. M2[i6] = mass * SC[o + 3]; M2[i6 + 1] = mass * SC[o + 4]; M2[i6 + 2] = mass * SC[o + 5]; @@ -150,7 +153,7 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { baseW[i3] = mass * SC[o + 12]; baseW[i3 + 1] = mass * SC[o + 13]; baseW[i3 + 2] = mass * SC[o + 14]; - Sself[i] = SC[o + 10] * SC[o + 10] * SC[o + 15] * PI_1_5 * SC[o + 9]; + Sself[i] = SC[o + 10] * SC[o + 10] * SC[o + 15] * (Math.PI ** 1.5) * SC[o + 9]; Err[i] = 0; } @@ -162,118 +165,15 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { return x; }; - // ---- Exact merge cost: ΔE of A∪B vs generation-input members + colour term. - const eigOut = new Float64Array(3); - let gatherBuf = new Uint32Array(1 << 12); - const gatherMembers = (root: number): number => { - let cnt = 0; - for (let m = mHead[root]; m !== NIL; m = mNext[m]) { - if (cnt === gatherBuf.length) { - const g = new Uint32Array(gatherBuf.length * 2); - g.set(gatherBuf); gatherBuf = g; - } - gatherBuf[cnt++] = m; - } - return cnt; - }; - - const evalOut = { E: 0, Scross: 0 }; - const evalMerge = (A: number, B: number): number => { - const WA = W[A], WB = W[B], WC = WA + WB; - const iw = 1 / WC; - const mcx = (WA * mx[A] + WB * mx[B]) * iw; - const mcy = (WA * my[A] + WB * my[B]) * iw; - const mcz = (WA * mz[A] + WB * mz[B]) * iw; - const dax = mx[A] - mcx, day = my[A] - mcy, daz = mz[A] - mcz; - const dbx = mx[B] - mcx, dby = my[B] - mcy, dbz = mz[B] - mcz; - const a6 = A * 6, b6 = B * 6; - - const sm0 = (M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx) * iw + EPS_COV; - const sm1 = (M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby) * iw; - const sm2 = (M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz) * iw; - const sm3 = (M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby) * iw + EPS_COV; - const sm4 = (M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz) * iw; - const sm5 = (M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz) * iw + EPS_COV; - - const detm = Math.max( - sm0 * (sm3 * sm5 - sm4 * sm4) - sm1 * (sm1 * sm5 - sm4 * sm2) + sm2 * (sm1 * sm4 - sm3 * sm2), - 1e-60 - ); - const sdC = Math.sqrt(detm); - - eig3(sm0, sm1, sm2, sm3, sm4, sm5, eigOut); - const s0 = Math.sqrt(Math.max(eigOut[0], 1e-18)); - const s1 = Math.sqrt(Math.max(eigOut[1], 1e-18)); - const s2 = Math.sqrt(Math.max(eigOut[2], 1e-18)); - const alphaC = Math.min(1, WC / Math.max(ellipsoidArea(s0, s1, s2), 1e-30)); - - const a3 = A * 3, b3 = B * 3; - const bc0 = (baseW[a3] + baseW[b3]) * iw; - const bc1 = (baseW[a3 + 1] + baseW[b3 + 1]) * iw; - const bc2 = (baseW[a3 + 2] + baseW[b3 + 2]) * iw; - const bn2C = bc0 * bc0 + bc1 * bc1 + bc2 * bc2; - - const selfM = alphaC * alphaC * bn2C * PI_1_5 * sdC; - - // ⟨Σ member fields, f_m⟩ over both chains. - let memfm = 0; - for (let pass = 0; pass < 2; pass++) { - for (let m = pass === 0 ? mHead[A] : mHead[B]; m !== NIL; m = mNext[m]) { - const o = m * CACHE_STRIDE; - const wgt = SC[o + 10] * alphaC * - (SC[o + 12] * bc0 + SC[o + 13] * bc1 + SC[o + 14] * bc2); - if (wgt === 0) continue; - memfm += wgt * crossG(SC[o + 9] * sdC, - SC[o + 3] + sm0, SC[o + 4] + sm1, SC[o + 5] + sm2, - SC[o + 6] + sm3, SC[o + 7] + sm4, SC[o + 8] + sm5, - SC[o] - mcx, SC[o + 1] - mcy, SC[o + 2] - mcz); - } - } - - // Scross(A,B) = Σ_{a∈A,b∈B}⟨f_a,f_b⟩ with distance culling. - const na = gatherMembers(A); - let scross = 0; - for (let b = mHead[B]; b !== NIL; b = mNext[b]) { - const ob = b * CACHE_STRIDE; - const bxp = SC[ob], byp = SC[ob + 1], bzp = SC[ob + 2]; - const trb = SC[ob + 3] + SC[ob + 6] + SC[ob + 8]; - const alb = SC[ob + 10], sdb = SC[ob + 9]; - const cb0 = SC[ob + 12], cb1 = SC[ob + 13], cb2 = SC[ob + 14]; - for (let t = 0; t < na; t++) { - const a = gatherBuf[t]; - const oa = a * CACHE_STRIDE; - const dx = SC[oa] - bxp, dy = SC[oa + 1] - byp, dz = SC[oa + 2] - bzp; - const d2 = dx * dx + dy * dy + dz * dz; - if (d2 > CULL_QUAD * (SC[oa + 3] + SC[oa + 6] + SC[oa + 8] + trb)) continue; - const wgt = SC[oa + 10] * alb * - (SC[oa + 12] * cb0 + SC[oa + 13] * cb1 + SC[oa + 14] * cb2); - if (wgt === 0) continue; - scross += wgt * crossG(SC[oa + 9] * sdb, - SC[oa + 3] + SC[ob + 3], SC[oa + 4] + SC[ob + 4], SC[oa + 5] + SC[ob + 5], - SC[oa + 6] + SC[ob + 6], SC[oa + 7] + SC[ob + 7], SC[oa + 8] + SC[ob + 8], - dx, dy, dz); - } - } - - const E = Sself[A] + Sself[B] + 2 * scross - 2 * memfm + selfM; - evalOut.E = E; - evalOut.Scross = scross; - - // Scale-free colour term between the cluster mean base colours (same - // definition as the pairwise kernel, so seed costs are consistent). - const iwA = 1 / W[A], iwB = 1 / W[B]; - const d0 = baseW[a3] * iwA - baseW[b3] * iwB; - const d1 = baseW[a3 + 1] * iwA - baseW[b3 + 1] * iwB; - const d2c = baseW[a3 + 2] * iwA - baseW[b3 + 2] * iwB; - return (E - Err[A] - Err[B]) + COLOR_WEIGHT * (d0 * d0 + d1 * d1 + d2c * d2c); - }; - - // ---- Lazy min-heap of candidate edges (one live entry per cluster). + // ---- Min-heap of candidate edges. Entries carry the winning + // evaluation's E/Scross (NaN E = seed entry, computed at commit). let heapCap = Math.ceil(N * 1.25) + 16; let hCost = new Float64Array(heapCap); let hA = new Uint32Array(heapCap), hB = new Uint32Array(heapCap); let hSeq = new Uint32Array(heapCap), hVb = new Uint32Array(heapCap); + let hE = new Float64Array(heapCap), hS = new Float64Array(heapCap); let heapSize = 0; + let seqCounter = 0; const swap = (i: number, j: number): void => { let t; @@ -282,20 +182,25 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { t = hB[i]; hB[i] = hB[j]; hB[j] = t; t = hSeq[i]; hSeq[i] = hSeq[j]; hSeq[j] = t; t = hVb[i]; hVb[i] = hVb[j]; hVb[j] = t; + t = hE[i]; hE[i] = hE[j]; hE[j] = t; + t = hS[i]; hS[i] = hS[j]; hS[j] = t; }; - const heapPush = (cost: number, a: number, b: number, seq: number, vb: number): void => { + const heapPush = (cost: number, a: number, b: number, seq: number, vb: number, e: number, s: number): void => { if (heapSize === heapCap) { const nc = heapCap * 2; - const c2 = new Float64Array(nc); c2.set(hCost); hCost = c2; + const gf = (old: Float64Array) => { + const x = new Float64Array(nc); x.set(old); return x; + }; const g = (old: Uint32Array) => { const x = new Uint32Array(nc); x.set(old); return x; }; + hCost = gf(hCost); hE = gf(hE); hS = gf(hS); hA = g(hA); hB = g(hB); hSeq = g(hSeq); hVb = g(hVb); heapCap = nc; } let i = heapSize++; - hCost[i] = cost; hA[i] = a; hB[i] = b; hSeq[i] = seq; hVb[i] = vb; + hCost[i] = cost; hA[i] = a; hB[i] = b; hSeq[i] = seq; hVb[i] = vb; hE[i] = e; hS[i] = s; while (i > 0) { const p = (i - 1) >> 1; if (hCost[p] <= hCost[i]) break; @@ -303,14 +208,16 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { } }; - const popOut = { cost: 0, a: 0, b: 0, seq: 0, vb: 0 }; + const popOut = { cost: 0, a: 0, b: 0, seq: 0, vb: 0, E: 0, S: 0 }; const heapPop = (): boolean => { if (heapSize === 0) return false; popOut.cost = hCost[0]; popOut.a = hA[0]; popOut.b = hB[0]; popOut.seq = hSeq[0]; popOut.vb = hVb[0]; + popOut.E = hE[0]; popOut.S = hS[0]; heapSize--; if (heapSize > 0) { hCost[0] = hCost[heapSize]; hA[0] = hA[heapSize]; hB[0] = hB[heapSize]; hSeq[0] = hSeq[heapSize]; hVb[0] = hVb[heapSize]; + hE[0] = hE[heapSize]; hS[0] = hS[heapSize]; let i = 0; for (;;) { const l = 2 * i + 1, r = l + 1; @@ -324,106 +231,143 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { return true; }; - // Candidate derivation: live clusters owning any neighbour of any member. - let candBuf = new Uint32Array(1 << 12); - const deriveCandidates = (root: number): number => { - stampGen++; - const gen = stampGen; - let cnt = 0; - for (let m = mHead[root]; m !== NIL; m = mNext[m]) { - const base = m * D; - for (let s = 0; s < D; s++) { - const nb = neighbors[base + s]; - if (nb === KNN_SENTINEL) continue; - const r = find(nb); - if (r === root || stamp[r] === gen) continue; - stamp[r] = gen; - if (cnt === candBuf.length) { - const g = new Uint32Array(candBuf.length * 2); - g.set(candBuf); candBuf = g; - } - candBuf[cnt++] = r; - } - } - return cnt; - }; - - const pushBestEdge = (root: number): void => { - const cnt = deriveCandidates(root); + // Refresh queue: clusters whose best edge must be re-evaluated. Queuing + // bumps lastSeq so stale heap entries for the cluster discard on pop; + // queuedRound dedupes within a round. + let pending = new Uint32Array(1 << 16); + let pendingCount = 0; + const queuedRound = new Uint32Array(N); + let round = 1; + const queueRefresh = (root: number): void => { lastSeq[root] = ++seqCounter; - if (cnt === 0) return; - const sz = size[root]; - let bc = Infinity, bp = -1, bv = 0; - for (let t = 0; t < cnt; t++) { - const c = candBuf[t]; - if (sz + size[c] > MAX_GROUP) continue; - const d = evalMerge(root, c); - if (d < bc) { - bc = d; bp = c; bv = version[c]; - } + if (queuedRound[root] === round) return; + queuedRound[root] = round; + if (pendingCount === pending.length) { + const g = new Uint32Array(pending.length * 2); + g.set(pending); pending = g; } - if (bp >= 0) heapPush(bc, root, bp, lastSeq[root], bv); + pending[pendingCount++] = root; }; - // Seed: the priority pass's cheapest candidate per gaussian (same cost - // definition, already sorted ascending — no re-evaluation needed). + // Seed: the priority pass's cheapest candidate per gaussian. for (let i = 0; i < N; i++) { const j = cand.idx[i * K]; const c = cand.cost[i * K]; lastSeq[i] = ++seqCounter; if (j !== NO_CANDIDATE && Number.isFinite(c)) { - heapPush(c, i, j, lastSeq[i], version[j]); + heapPush(c, i, j, lastSeq[i], version[j], NaN, 0); } } - // ---- Greedy loop. + // ---- Round loop. Each round commits at most WAVE merges before the + // bulk refresh, so refreshed edges (notably cheap continuation merges in + // redundant regions — the concentration behaviour the quality depends on) + // re-enter the heap at most one wave late. An unbounded drain would spend + // the budget up the cost curve before any refresh returns. + const WAVE = 4096; let removed = 0; while (removed < mergesNeeded) { - if (!heapPop()) break; - const a = popOut.a; - if (parent[a] !== a || popOut.seq !== lastSeq[a]) continue; - const b = popOut.b; - if (parent[b] !== b || version[b] !== popOut.vb || size[a] + size[b] > MAX_GROUP) { - pushBestEdge(a); - continue; - } - - // Commit a+b (exact recompute for the error bookkeeping). - evalMerge(a, b); - const E = evalOut.E, scross = evalOut.Scross; - const keep = size[a] >= size[b] ? a : b; - const lose = keep === a ? b : a; - - const WA = W[a], WB = W[b], WC = WA + WB; - const iw = 1 / WC; - const mcx = (WA * mx[a] + WB * mx[b]) * iw; - const mcy = (WA * my[a] + WB * my[b]) * iw; - const mcz = (WA * mz[a] + WB * mz[b]) * iw; - const dax = mx[a] - mcx, day = my[a] - mcy, daz = mz[a] - mcz; - const dbx = mx[b] - mcx, dby = my[b] - mcy, dbz = mz[b] - mcz; - const a6 = a * 6, b6 = b * 6, k6 = keep * 6; - const n0 = M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx; - const n1 = M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby; - const n2 = M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz; - const n3 = M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby; - const n4 = M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz; - const n5 = M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz; - M2[k6] = n0; M2[k6 + 1] = n1; M2[k6 + 2] = n2; M2[k6 + 3] = n3; M2[k6 + 4] = n4; M2[k6 + 5] = n5; - W[keep] = WC; mx[keep] = mcx; my[keep] = mcy; mz[keep] = mcz; - const ka = keep * 3, aa = a * 3, bb = b * 3; - const bw0 = baseW[aa] + baseW[bb], bw1 = baseW[aa + 1] + baseW[bb + 1], bw2 = baseW[aa + 2] + baseW[bb + 2]; - baseW[ka] = bw0; baseW[ka + 1] = bw1; baseW[ka + 2] = bw2; - Sself[keep] = Sself[a] + Sself[b] + 2 * scross; - Err[keep] = E; - size[keep] += size[lose]; - mNext[mTail[keep]] = mHead[lose]; - mTail[keep] = mTail[lose]; - parent[lose] = keep; - version[keep]++; - liveCount--; - removed++; + // Drain: commit still-valid entries, up to the wave budget. + let wave = 0; + while (removed < mergesNeeded && wave < WAVE && heapPop()) { + const a = popOut.a; + if (parent[a] !== a || popOut.seq !== lastSeq[a]) continue; + const b = popOut.b; + if (parent[b] !== b || version[b] !== popOut.vb || size[a] + size[b] > MAX_GROUP) { + queueRefresh(a); + continue; + } - pushBestEdge(keep); + let E = popOut.E, scross = popOut.S; + if (Number.isNaN(E)) { + evalMergeCore(st, a, b); + E = evalOut.E; scross = evalOut.Scross; + } + const keep = size[a] >= size[b] ? a : b; + const lose = keep === a ? b : a; + + const WA = W[a], WB = W[b], WC = WA + WB; + const iw = 1 / WC; + const mcx = (WA * mx[a] + WB * mx[b]) * iw; + const mcy = (WA * my[a] + WB * my[b]) * iw; + const mcz = (WA * mz[a] + WB * mz[b]) * iw; + const dax = mx[a] - mcx, day = my[a] - mcy, daz = mz[a] - mcz; + const dbx = mx[b] - mcx, dby = my[b] - mcy, dbz = mz[b] - mcz; + const a6 = a * 6, b6 = b * 6, k6 = keep * 6; + const n0 = M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx; + const n1 = M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby; + const n2 = M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz; + const n3 = M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby; + const n4 = M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz; + const n5 = M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz; + M2[k6] = n0; M2[k6 + 1] = n1; M2[k6 + 2] = n2; M2[k6 + 3] = n3; M2[k6 + 4] = n4; M2[k6 + 5] = n5; + W[keep] = WC; mx[keep] = mcx; my[keep] = mcy; mz[keep] = mcz; + const ka = keep * 3, aa = a * 3, bb = b * 3; + const bw0 = baseW[aa] + baseW[bb], bw1 = baseW[aa + 1] + baseW[bb + 1], bw2 = baseW[aa + 2] + baseW[bb + 2]; + baseW[ka] = bw0; baseW[ka + 1] = bw1; baseW[ka + 2] = bw2; + Sself[keep] = Sself[a] + Sself[b] + 2 * scross; + Err[keep] = E; + size[keep] += size[lose]; + mNext[mTail[keep]] = mHead[lose]; + mTail[keep] = mTail[lose]; + parent[lose] = keep; + version[keep]++; + removed++; + wave++; + + queueRefresh(keep); + } + if (removed >= mergesNeeded) break; + if (pendingCount === 0 && heapSize === 0) break; + + // Bulk refresh of queued clusters (parallel when shared). + if (shared) { + const workers = Math.max(1, Math.min(8, WorkerQueue.maxWorkers ?? 4)); + const chunk = Math.ceil(pendingCount / workers); + const jobs: Promise[] = []; + for (let off = 0; off < pendingCount; off += chunk) { + const roots = pending.slice(off, Math.min(off + chunk, pendingCount)); + jobs.push(WorkerQueue.run('recostBestEdges', { + sc: SC, + cands: neighbors, + d: D, + n: N, + maxGroup: MAX_GROUP, + parent, + size, + w: W, + mx, + my, + mz, + m2: M2, + baseW, + sself: Sself, + err: Err, + version, + mHead, + mNext, + roots + }, [roots.buffer as ArrayBuffer])); + } + const results = await Promise.all(jobs); + for (const res of results) { + for (let i = 0; i < res.length; i += 6) { + const root = res[i]; + const partner = res[i + 1]; + if (partner < 0) continue; + heapPush(res[i + 3], root, partner, lastSeq[root], res[i + 2], res[i + 4], res[i + 5]); + } + } + } else { + for (let p = 0; p < pendingCount; p++) { + const root = pending[p]; + if (bestEdgeFor(st, root)) { + heapPush(bestOut.cost, root, bestOut.partner, lastSeq[root], bestOut.vb, bestOut.E, bestOut.S); + } + } + } + pendingCount = 0; + round++; } // ---- CSR assembly (identical contract to selectMerges). @@ -455,4 +399,4 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { return { groupOffsets, groupMembers, memberGroup, groupMin, mergedGroups: G, removed }; }; -export { selectMergesRecosted, type RecostInputs }; +export { selectMergesRecosted, CACHE_STRIDE, type RecostInputs }; diff --git a/src/lib/gpu/gpu-edge-cost-legacy.ts b/src/lib/gpu/gpu-edge-cost-legacy.ts new file mode 100644 index 00000000..aed46874 --- /dev/null +++ b/src/lib/gpu/gpu-edge-cost-legacy.ts @@ -0,0 +1,493 @@ +import { + BUFFERUSAGE_COPY_DST, + BUFFERUSAGE_COPY_SRC, + SHADERLANGUAGE_WGSL, + SHADERSTAGE_COMPUTE, + UNIFORMTYPE_FLOAT, + UNIFORMTYPE_UINT, + BindGroupFormat, + BindStorageBufferFormat, + BindUniformBufferFormat, + Compute, + GraphicsDevice, + Shader, + StorageBuffer, + UniformBufferFormat, + UniformFormat +} from 'playcanvas'; + +/** + * Appearance columns per storage chunk. The kernel exposes three appearance + * bindings (appA/appB/appC), so the layout holds up to 3·APP_CHUNK columns; at + * 16 the widest chunk reaches the ~2 GB per-binding limit around ~33.5M splats. + * The CPU-side packing in `decimate/priority.ts` imports this same constant, + * so the kernel strides and the host packing can't drift. + */ +export const APP_CHUNK = 16; + +/** + * WGSL kernel: per-edge KL-style cost (matches `computeEdgeCostView` in + * `decimate/edge-cost-cpu.ts`). + * + * Each thread = one edge (i, j). Reads the per-splat cache for both + * endpoints, computes the merged Gaussian's covariance + determinant, + * runs a single Monte-Carlo sample through both component gaussians + * (the same `z` for both components, matching the CPU implementation), + * and adds an L2 distance over the appearance (SH) coefficients. + * + * @param strideA - Live column count of appearance chunk A (0 if unused). + * @param strideB - Live column count of appearance chunk B (0 if unused). + * @param strideC - Live column count of appearance chunk C (0 if unused). + * @returns WGSL source. + */ +const edgeCostWgsl = (strideA: number, strideB: number, strideC: number) => /* wgsl */` +struct Uniforms { + edgeCount: u32, + z0: f32, + z1: f32, + z2: f32, +} + +@group(0) @binding(0) var uniforms: Uniforms; +// Edge list for the current dispatch batch only, split into two parallel +// arrays (avoids a host-side (i, j) interleave). The host uploads each batch's +// slice to offset 0, so we index edgesI/J[bid] directly — keeping these +// buffers batch-sized instead of N·k keeps them off the per-binding limit. +@group(0) @binding(1) var edgesI: array; +@group(0) @binding(2) var edgesJ: array; +// Per-splat geometry, interleaved 8-wide: +// posScalars[8s + 0..2] = position xyz +// posScalars[8s + 3] = mass +// posScalars[8s + 4] = logdet +// posScalars[8s + 5..7] = variances (vx, vy, vz) +@group(0) @binding(3) var posScalars: array; +// Row-major 3x3 rotation matrix per splat (9 floats per splat). +@group(0) @binding(4) var rotR: array; +// Appearance, split into up to three chunks (≤16 columns each) so no single +// binding exceeds maxStorageBufferBindingSize (~2 GB). Each chunk's stride is +// its live column count (STRIDE_A/B/C below); appA holds columns 0.., appB the +// next span, appC the next. Unused chunks have stride 0, are bound to a dummy +// buffer, and are never read. +@group(0) @binding(5) var appA: array; +@group(0) @binding(6) var appB: array; +@group(0) @binding(7) var appC: array; +// Output: cost per edge. +@group(0) @binding(8) var costs: array; + +const EPS_COV: f32 = 1e-8; +const LOG2PI: f32 = 1.8378770664093453; +// Per-chunk appearance strides = live column count in each chunk (0 = unused, +// dummy-bound). Baked here because the column count is fixed for the lifetime +// of the kernel, so loop bounds and indexing resolve statically. +const STRIDE_A: u32 = ${strideA}u; +const STRIDE_B: u32 = ${strideB}u; +const STRIDE_C: u32 = ${strideC}u; + +// Symmetric 3x3 covariance helpers — we pass them around as 6 f32 (xx, xy, xz, yy, yz, zz). + +// Σ = R · diag(v) · R^T for row-major R (a 9-float array starting at offset r9). +// Variances v come from posScalars[s8 + 5..7]. Result is 6 floats: +// (xx, xy, xz, yy, yz, zz). +fn sigmaFromRotVar(r9: u32, s8: u32) -> array { + let r00 = rotR[r9 + 0u]; let r01 = rotR[r9 + 1u]; let r02 = rotR[r9 + 2u]; + let r10 = rotR[r9 + 3u]; let r11 = rotR[r9 + 4u]; let r12 = rotR[r9 + 5u]; + let r20 = rotR[r9 + 6u]; let r21 = rotR[r9 + 7u]; let r22 = rotR[r9 + 8u]; + let vx = posScalars[s8 + 5u]; + let vy = posScalars[s8 + 6u]; + let vz = posScalars[s8 + 7u]; + return array( + r00*r00*vx + r01*r01*vy + r02*r02*vz, // xx + r00*r10*vx + r01*r11*vy + r02*r12*vz, // xy + r00*r20*vx + r01*r21*vy + r02*r22*vz, // xz + r10*r10*vx + r11*r11*vy + r12*r12*vz, // yy + r10*r20*vx + r11*r21*vy + r12*r22*vz, // yz + r20*r20*vx + r21*r21*vy + r22*r22*vz // zz + ); +} + +// log N(x | mu, R · diag(v) · R^T) for a diagonally-decomposed covariance. +// invDiag is (1/vx, 1/vy, 1/vz); ld is logdet of the full covariance. +// Evaluates y = R^T * (x - mu) using columns of R (= rows of Rt). +fn gaussLogpdfDiagrot( + x: vec3f, mu: vec3f, r9: u32, + invDiag: vec3f, ld: f32 +) -> f32 { + let dx = x.x - mu.x; + let dy = x.y - mu.y; + let dz = x.z - mu.z; + // y = R^T · d. R is row-major; column k of R is (R[k], R[k+3], R[k+6]). + let y0 = dx * rotR[r9 + 0u] + dy * rotR[r9 + 3u] + dz * rotR[r9 + 6u]; + let y1 = dx * rotR[r9 + 1u] + dy * rotR[r9 + 4u] + dz * rotR[r9 + 7u]; + let y2 = dx * rotR[r9 + 2u] + dy * rotR[r9 + 5u] + dz * rotR[r9 + 8u]; + let quad = y0*y0*invDiag.x + y1*y1*invDiag.y + y2*y2*invDiag.z; + return -0.5 * (3.0 * LOG2PI + ld + quad); +} + +fn logAddExp(a: f32, b: f32) -> f32 { + let m = max(a, b); + return m + log(exp(a - m) + exp(b - m)); +} + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3u) { + let bid = gid.x; + if (bid >= uniforms.edgeCount) { return; } + + let i = edgesI[bid]; + let j = edgesJ[bid]; + + let i8 = i * 8u; + let j8 = j * 8u; + let i9 = i * 9u; + let j9 = j * 9u; + + let mu_i = vec3f(posScalars[i8 + 0u], posScalars[i8 + 1u], posScalars[i8 + 2u]); + let mu_j = vec3f(posScalars[j8 + 0u], posScalars[j8 + 1u], posScalars[j8 + 2u]); + + let wi = posScalars[i8 + 3u]; + let wj = posScalars[j8 + 3u]; + let W = wi + wj; + let Wsafe = select(1.0, W, W > 0.0); + let pi_w_raw = wi / Wsafe; + let pi_w = clamp(pi_w_raw, 1e-12, 1.0 - 1e-12); + let pj_w = 1.0 - pi_w; + let logPi = log(pi_w); + let logPj = log(pj_w); + + // Merged mean. + let mm = pi_w * mu_i + pj_w * mu_j; + let di = mu_i - mm; + let dj = mu_j - mm; + + // Σ_i and Σ_j from rotation + variances. + let sig_i = sigmaFromRotVar(i9, i8); + let sig_j = sigmaFromRotVar(j9, j8); + + // Merged covariance: pi*(Σ_i + δi·δiᵀ) + pj*(Σ_j + δj·δjᵀ), + EPS on diag. + let s_xx = pi_w * (sig_i[0] + di.x*di.x) + pj_w * (sig_j[0] + dj.x*dj.x) + EPS_COV; + let s_xy = pi_w * (sig_i[1] + di.x*di.y) + pj_w * (sig_j[1] + dj.x*dj.y); + let s_xz = pi_w * (sig_i[2] + di.x*di.z) + pj_w * (sig_j[2] + dj.x*dj.z); + let s_yy = pi_w * (sig_i[3] + di.y*di.y) + pj_w * (sig_j[3] + dj.y*dj.y) + EPS_COV; + let s_yz = pi_w * (sig_i[4] + di.y*di.z) + pj_w * (sig_j[4] + dj.y*dj.z); + let s_zz = pi_w * (sig_i[5] + di.z*di.z) + pj_w * (sig_j[5] + dj.z*dj.z) + EPS_COV; + + // det of symmetric 3x3. + let det_m = s_xx * (s_yy*s_zz - s_yz*s_yz) + - s_xy * (s_xy*s_zz - s_yz*s_xz) + + s_xz * (s_xy*s_yz - s_yy*s_xz); + let logdet_m = log(max(det_m, 1e-30)); + + // Entropy of the merged Gaussian: H = 0.5 (k log(2π) + log|Σ_m| + k), k=3. + let EpNegLogQ = 0.5 * (3.0 * LOG2PI + logdet_m + 3.0); + + // Read per-axis std for each input (variances live at posScalars[s8+5..7]). + let vix = posScalars[i8 + 5u]; let viy = posScalars[i8 + 6u]; let viz = posScalars[i8 + 7u]; + let vjx = posScalars[j8 + 5u]; let vjy = posScalars[j8 + 6u]; let vjz = posScalars[j8 + 7u]; + let stdix = sqrt(max(vix, 0.0)); + let stdiy = sqrt(max(viy, 0.0)); + let stdiz = sqrt(max(viz, 0.0)); + let stdjx = sqrt(max(vjx, 0.0)); + let stdjy = sqrt(max(vjy, 0.0)); + let stdjz = sqrt(max(vjz, 0.0)); + + // Inverse diagonals (1 / variance) for the log-pdf quadratic term. + let invDi = vec3f(1.0 / max(vix, 1e-30), 1.0 / max(viy, 1e-30), 1.0 / max(viz, 1e-30)); + let invDj = vec3f(1.0 / max(vjx, 1e-30), 1.0 / max(vjy, 1e-30), 1.0 / max(vjz, 1e-30)); + let ldi = posScalars[i8 + 4u]; + let ldj = posScalars[j8 + 4u]; + + let z0 = uniforms.z0; + let z1 = uniforms.z1; + let z2 = uniforms.z2; + + // Sample x = mu + R · diag(std) · z where z ~ N(0, I). + // Row a of R is (rotR[r9+3a], rotR[r9+3a+1], rotR[r9+3a+2]). + // x[a] = mu[a] + R[a][0]*std[0]*z[0] + R[a][1]*std[1]*z[1] + R[a][2]*std[2]*z[2]. + let xix = mu_i.x + z0 * stdix * rotR[i9 + 0u] + z1 * stdiy * rotR[i9 + 1u] + z2 * stdiz * rotR[i9 + 2u]; + let xiy = mu_i.y + z0 * stdix * rotR[i9 + 3u] + z1 * stdiy * rotR[i9 + 4u] + z2 * stdiz * rotR[i9 + 5u]; + let xiz = mu_i.z + z0 * stdix * rotR[i9 + 6u] + z1 * stdiy * rotR[i9 + 7u] + z2 * stdiz * rotR[i9 + 8u]; + let xi = vec3f(xix, xiy, xiz); + + let xjx = mu_j.x + z0 * stdjx * rotR[j9 + 0u] + z1 * stdjy * rotR[j9 + 1u] + z2 * stdjz * rotR[j9 + 2u]; + let xjy = mu_j.y + z0 * stdjx * rotR[j9 + 3u] + z1 * stdjy * rotR[j9 + 4u] + z2 * stdjz * rotR[j9 + 5u]; + let xjz = mu_j.z + z0 * stdjx * rotR[j9 + 6u] + z1 * stdjy * rotR[j9 + 7u] + z2 * stdjz * rotR[j9 + 8u]; + let xj = vec3f(xjx, xjy, xjz); + + // log p_ij at samples from component i. + let logNiOnI = gaussLogpdfDiagrot(xi, mu_i, i9, invDi, ldi); + let logNjOnI = gaussLogpdfDiagrot(xi, mu_j, j9, invDj, ldj); + let logpOnI = logAddExp(logPi + logNiOnI, logPj + logNjOnI); + + let logNiOnJ = gaussLogpdfDiagrot(xj, mu_i, i9, invDi, ldi); + let logNjOnJ = gaussLogpdfDiagrot(xj, mu_j, j9, invDj, ldj); + let logpOnJ = logAddExp(logPi + logNiOnJ, logPj + logNjOnJ); + + let Ei = logpOnI; + let Ej = logpOnJ; + let EpLogp = pi_w * Ei + pj_w * Ej; + let geo = EpLogp + EpNegLogQ; + + // Appearance L2 cost, summed across the (up to three) chunks. Each chunk's + // stride is its live column count, so partial chunks store/read no padding. + // A 0 stride yields 0 iterations, so the dummy-bound appB / appC are never + // touched on inputs with fewer SH bands. + var cSh: f32 = 0.0; + let iA = i * STRIDE_A; let jA = j * STRIDE_A; + for (var k: u32 = 0u; k < STRIDE_A; k = k + 1u) { + let d = appA[iA + k] - appA[jA + k]; + cSh = cSh + d * d; + } + let iB = i * STRIDE_B; let jB = j * STRIDE_B; + for (var k: u32 = 0u; k < STRIDE_B; k = k + 1u) { + let d = appB[iB + k] - appB[jB + k]; + cSh = cSh + d * d; + } + let iC = i * STRIDE_C; let jC = j * STRIDE_C; + for (var k: u32 = 0u; k < STRIDE_C; k = k + 1u) { + let d = appC[iC + k] - appC[jC + k]; + cSh = cSh + d * d; + } + + costs[bid] = geo + cSh; +} +`; + +/** + * Per-splat cache for the edge cost kernel. Packed layouts to stay within the + * WebGPU per-stage storage-buffer count limit (8) and the per-binding size + * limit (~2 GB) — appearance is split into 16-column chunks for the latter. + */ +interface EdgeCostCacheLegacy { + /** Per-splat geometry interleaved 8-wide: (x, y, z, mass, logdet, vx, vy, vz). */ + posScalars: Float32Array; + /** Row-major 3×3 rotation per splat (length 9N). */ + rotR: Float32Array; + /** + * Appearance in up to three chunks of ≤16 columns. Chunk c has stride + * width_c (= its live column count): appChunks[c][s * width_c + k]. + */ + appChunks: Float32Array[]; + /** Number of appearance columns C. */ + numAppCols: number; + /** Number of splats. */ + numSplats: number; +} + +/** + * GPU edge-cost evaluator. + * + * Each compute thread evaluates the KL-style cost for one edge (i, j) by + * reading the per-splat cache for both endpoints, computing the merged + * Gaussian's covariance/determinant, running a single Monte-Carlo sample + * through both component PDFs, and adding an L2 distance over the + * appearance (SH) coefficients. Output is `costs[e] = cost for edge e`. + * + * Mirrors the CPU `computeEdgeCostView` in `decimate/edge-cost-cpu.ts`. + */ +class GpuEdgeCostLegacy { + /** + * @param cache - Per-splat cache (uploaded once). + * @param edgeI - Edge u indices (length E). + * @param edgeJ - Edge v indices (length E). + * @param z - Single Monte-Carlo sample (3 floats from N(0,1)). + * @param outCosts - Destination for per-edge costs (length E). + */ + execute: ( + cache: EdgeCostCacheLegacy, + edgeI: Uint32Array, + edgeJ: Uint32Array, + z: Float32Array, + outCosts: Float32Array + ) => Promise; + destroy: () => void; + + /** + * @param device - PlayCanvas GraphicsDevice (WebGPU). + * @param maxN - Maximum number of splats. + * @param maxE - Maximum number of edges in a single dispatch. + * @param maxAppCols - Maximum appearance column count (over all bands). + */ + constructor(device: GraphicsDevice, maxN: number, maxE: number, maxAppCols: number) { + const workgroupSize = 64; + const edgesPerBatch = 1024 * workgroupSize; // 65,536 + // Appearance is split at fixed APP_CHUNK-column boundaries, but each + // chunk's *stride* is its live column count — only the last non-empty + // chunk is ever partial, so partial chunks neither allocate nor upload + // padding. The widest possible chunk reaches the ~2 GB limit at ~33.5M + // splats, past the ~11.2M wall the single 48-col buffer hit. The three + // kernel bindings (appA/appB/appC) cap the layout at three chunks. + // e.g. [16, 11, 0] for 27 cols, [3, 0, 0] for DC-only. + const appStrides = [0, 1, 2].map((ch) => { + return Math.min(APP_CHUNK, Math.max(0, maxAppCols - ch * APP_CHUNK)); + }); + // Non-empty chunk count the kernel reads. execute() validates the cache + // supplies exactly this many: a short count would leave a hoisted (reused + // across iterations) appearance buffer holding the previous iteration's + // data, which the kernel would then read as this iteration's appearance. + const numAppChunks = appStrides.filter(stride => stride > 0).length; + + const bindGroupFormat = new BindGroupFormat(device, [ + new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE), + new BindStorageBufferFormat('edgesI', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('edgesJ', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('posScalars', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('rotR', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('appA', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('appB', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('appC', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('costs', SHADERSTAGE_COMPUTE) + ]); + + const shader = new Shader(device, { + name: 'compute-edge-cost-legacy', + shaderLanguage: SHADERLANGUAGE_WGSL, + cshader: edgeCostWgsl(appStrides[0], appStrides[1], appStrides[2]), + // @ts-ignore + computeUniformBufferFormats: { + uniforms: new UniformBufferFormat(device, [ + new UniformFormat('edgeCount', UNIFORMTYPE_UINT), + new UniformFormat('z0', UNIFORMTYPE_FLOAT), + new UniformFormat('z1', UNIFORMTYPE_FLOAT), + new UniformFormat('z2', UNIFORMTYPE_FLOAT) + ]) + }, + // @ts-ignore + computeBindGroupFormat: bindGroupFormat + }); + + // Pre-flight the largest per-N bindings against the device's storage + // limit so we fail with a clear message instead of a driver-side error. + // Edges are uploaded per batch (batch-sized buffers), so they can't hit + // the limit; the widest appearance chunk and rotR are the candidates. + // posScalars (8 floats/splat) is strictly smaller than rotR (9), so the + // rotR check already bounds it — no separate check needed. + const maxStorage = (device as any).limits?.maxStorageBufferBindingSize; + if (typeof maxStorage === 'number') { + const checkLimit = (label: string, bytes: number) => { + if (bytes > maxStorage) { + throw new Error( + `GpuEdgeCostLegacy: ${label} buffer (${bytes} bytes) exceeds device ` + + `maxStorageBufferBindingSize (${maxStorage})` + ); + } + }; + const maxChunkCols = Math.max(...appStrides); + checkLimit(`appearance chunk (${maxN} splats × ${maxChunkCols} cols)`, maxN * maxChunkCols * 4); + checkLimit(`rotR (${maxN} splats × 9)`, maxN * 9 * 4); + } + + const posScalarsBuf = new StorageBuffer(device, maxN * 8 * 4, BUFFERUSAGE_COPY_DST); + const rotRBuf = new StorageBuffer(device, maxN * 9 * 4, BUFFERUSAGE_COPY_DST); + + // One buffer per non-empty appearance chunk, sized to that chunk's live + // column count; empty slots (inputs with fewer SH bands) share a small + // dummy since WebGPU forbids a zero-size binding. The 3-binding layout + // stays fixed regardless of band count. + const appDummy = new StorageBuffer(device, 16, BUFFERUSAGE_COPY_DST); + const appBufs: StorageBuffer[] = appStrides.map((width) => { + return width > 0 ? + new StorageBuffer(device, maxN * width * 4, BUFFERUSAGE_COPY_DST) : + appDummy; + }); + + // Two parallel u32 buffers, sized to a single dispatch batch (not the + // full N·k edge list): execute uploads each batch's slice before its + // dispatch. Batch-sizing keeps these ~256 KB instead of N·k·4 — off the + // ~2 GB per-binding limit (so edges never cap scene size) and ~1.6 GB + // less VRAM at 13M splats. Two parallel arrays avoid a host-side pack. + const edgesIBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); + const edgesJBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); + + const outBuf = new StorageBuffer( + device, + edgesPerBatch * 4, + BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST + ); + const outScratch = new Float32Array(edgesPerBatch); + + const compute = new Compute(device, shader, 'compute-edge-cost-legacy'); + compute.setParameter('edgesI', edgesIBuf); + compute.setParameter('edgesJ', edgesJBuf); + compute.setParameter('posScalars', posScalarsBuf); + compute.setParameter('rotR', rotRBuf); + compute.setParameter('appA', appBufs[0]); + compute.setParameter('appB', appBufs[1]); + compute.setParameter('appC', appBufs[2]); + compute.setParameter('costs', outBuf); + + this.execute = async ( + cache: EdgeCostCacheLegacy, + edgeI: Uint32Array, + edgeJ: Uint32Array, + z: Float32Array, + outCosts: Float32Array + ) => { + const n = cache.numSplats; + const e = edgeI.length; + + if (n > maxN) throw new Error(`GpuEdgeCostLegacy: N=${n} exceeds maxN=${maxN}`); + if (e > maxE) throw new Error(`GpuEdgeCostLegacy: E=${e} exceeds maxE=${maxE}`); + if (cache.numAppCols !== maxAppCols) { + throw new Error(`GpuEdgeCostLegacy: numAppCols=${cache.numAppCols} must equal maxAppCols=${maxAppCols} (baked into the kernel)`); + } + if (cache.appChunks.length !== numAppChunks) { + throw new Error(`GpuEdgeCostLegacy: cache supplies ${cache.appChunks.length} appearance chunks but the kernel layout expects ${numAppChunks}`); + } + if (edgeJ.length !== e || outCosts.length !== e) { + throw new Error('GpuEdgeCostLegacy: edgeI / edgeJ / outCosts must have same length'); + } + if (z.length < 3) { + throw new Error('GpuEdgeCostLegacy: z must have at least 3 elements'); + } + + // Upload per-splat cache. Each appearance chunk is row-major with + // stride = its live column count, so we upload n*width per chunk. + posScalarsBuf.write(0, cache.posScalars, 0, n * 8); + rotRBuf.write(0, cache.rotR, 0, n * 9); + for (let ch = 0; ch < cache.appChunks.length; ch++) { + appBufs[ch].write(0, cache.appChunks[ch], 0, n * appStrides[ch]); + } + + compute.setParameter('z0', z[0]); + compute.setParameter('z1', z[1]); + compute.setParameter('z2', z[2]); + + const numBatches = Math.ceil(e / edgesPerBatch); + for (let batch = 0; batch < numBatches; batch++) { + const edgeOffset = batch * edgesPerBatch; + const edgeCount = Math.min(edgesPerBatch, e - edgeOffset); + const groups = Math.ceil(edgeCount / workgroupSize); + + // Upload just this batch's edges to offset 0; the kernel indexes + // edgesI/J[bid] within the batch. + edgesIBuf.write(0, edgeI, edgeOffset, edgeCount); + edgesJBuf.write(0, edgeJ, edgeOffset, edgeCount); + + compute.setParameter('edgeCount', edgeCount); + + compute.setupDispatch(groups); + device.computeDispatch([compute], `edge-cost-dispatch-${batch}`); + + const readBytes = edgeCount * 4; + await outBuf.read(0, readBytes, outScratch, true); + outCosts.set(outScratch.subarray(0, edgeCount), edgeOffset); + } + }; + + this.destroy = () => { + posScalarsBuf.destroy(); + rotRBuf.destroy(); + for (const buf of appBufs) { + if (buf !== appDummy) buf.destroy(); + } + appDummy.destroy(); + edgesIBuf.destroy(); + edgesJBuf.destroy(); + outBuf.destroy(); + shader.destroy(); + bindGroupFormat.destroy(); + }; + } +} + +export { GpuEdgeCostLegacy, type EdgeCostCacheLegacy }; diff --git a/src/lib/workers/tasks.ts b/src/lib/workers/tasks.ts index 215fe19e..5a9bdca9 100644 --- a/src/lib/workers/tasks.ts +++ b/src/lib/workers/tasks.ts @@ -1,6 +1,7 @@ import type { TypedArray } from '../data-table/data-table'; import { knnQueryBlock } from '../decimate/knn-core'; import { mergeGroup, createMergeScratch, splatMass } from '../decimate/moment-match'; +import { bestEdgeFor, bestOut, type RecostState } from '../decimate/recost-core'; import { KdTree, type FlatKdTree } from '../spatial/kd-tree'; import { quantize1dColumns, type QuantizedColumns } from '../spatial/quantize-1d-core'; import { WebPCodec } from '../utils/webp-codec'; @@ -68,6 +69,57 @@ const taskHandlers = { return { result, transfer: [result.buffer as ArrayBuffer] }; }, + // Re-costed selection bulk refresh: best merge edge for each queued + // cluster root, evaluated against SharedArrayBuffer-backed state that the + // host freezes for the duration of the round (read-only here). Packed + // result: 6 f64 per root — [root, partner|-1, partnerVersion, cost, E, Scross]. + recostBestEdges: (args: { + sc: Float32Array, cands: Uint32Array, d: number, n: number, maxGroup: number, + parent: Uint32Array, size: Uint32Array, + w: Float64Array, mx: Float64Array, my: Float64Array, mz: Float64Array, + m2: Float64Array, baseW: Float64Array, sself: Float64Array, err: Float64Array, + version: Uint32Array, mHead: Uint32Array, mNext: Uint32Array, + roots: Uint32Array + }): TaskOutput => { + const st: RecostState = { + SC: args.sc, + cands: args.cands, + D: args.d, + N: args.n, + maxGroup: args.maxGroup, + parent: args.parent, + size: args.size, + W: args.w, + mx: args.mx, + my: args.my, + mz: args.mz, + M2: args.m2, + baseW: args.baseW, + Sself: args.sself, + Err: args.err, + version: args.version, + mHead: args.mHead, + mNext: args.mNext + }; + const roots = args.roots; + const out = new Float64Array(roots.length * 6); + for (let i = 0; i < roots.length; i++) { + const root = roots[i]; + const o = i * 6; + out[o] = root; + if (bestEdgeFor(st, root)) { + out[o + 1] = bestOut.partner; + out[o + 2] = bestOut.vb; + out[o + 3] = bestOut.cost; + out[o + 4] = bestOut.E; + out[o + 5] = bestOut.S; + } else { + out[o + 1] = -1; + } + } + return { result: out, transfer: [out.buffer as ArrayBuffer] }; + }, + // Decimation merge stream: n-ary moment match of packed member-major // groups. Inputs are member-major (pos 3 / geo 8 / color colorDim floats // per member, groups back to back per `sizes`); outputs are group-major. From 1dbbf2082dd5fbc470e3101e13d3c7f60b1f40e2 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 14:05:11 +0100 Subject: [PATCH 04/19] latest --- src/lib/decimate/priority-legacy.ts | 2 +- src/lib/decimate/priority.ts | 2 +- src/lib/gpu/gpu-knn.ts | 33 +----- src/lib/spatial/kd-tree.ts | 176 +++++++++++++--------------- src/lib/workers/tasks.ts | 13 +- test/kd-tree.test.mjs | 46 ++++++-- 6 files changed, 134 insertions(+), 138 deletions(-) diff --git a/src/lib/decimate/priority-legacy.ts b/src/lib/decimate/priority-legacy.ts index a8700a2c..7bcd6afd 100644 --- a/src/lib/decimate/priority-legacy.ts +++ b/src/lib/decimate/priority-legacy.ts @@ -249,7 +249,7 @@ const runPriorityPassLegacy = async ( const locals = collectBlock(pos, order, blocks, bi, k, HALO_FACTOR, HALO_CAP); const copy = locals.positions.slice(); if (device) { - const treePromise = WorkerQueue.run('flattenKdTree', { positions: copy }, [copy.buffer as ArrayBuffer]); + const treePromise = WorkerQueue.run('buildFlatKdTree', { positions: copy }, [copy.buffer as ArrayBuffer]); const out = new Uint32Array(locals.ownedCount * k); const run = Promise.all([treePromise, gpuKnnQueue]).then(([flat]) => { return gpuKnn!.execute(flat, locals.positions, locals.ids.length, locals.ownedCount, out); diff --git a/src/lib/decimate/priority.ts b/src/lib/decimate/priority.ts index 626e0b27..c941f6bc 100644 --- a/src/lib/decimate/priority.ts +++ b/src/lib/decimate/priority.ts @@ -212,7 +212,7 @@ const runPriorityPass = async ( const locals = collectBlock(pos, order, blocks, bi, k, HALO_FACTOR, HALO_CAP); const copy = locals.positions.slice(); if (device) { - const treePromise = WorkerQueue.run('flattenKdTree', { positions: copy }, [copy.buffer as ArrayBuffer]); + const treePromise = WorkerQueue.run('buildFlatKdTree', { positions: copy }, [copy.buffer as ArrayBuffer]); const out = new Uint32Array(locals.ownedCount * k); const run = Promise.all([treePromise, gpuKnnQueue]).then(([flat]) => { return gpuKnn!.execute(flat, locals.positions, locals.ids.length, locals.ownedCount, out); diff --git a/src/lib/gpu/gpu-knn.ts b/src/lib/gpu/gpu-knn.ts index 2718f683..6ef6803f 100644 --- a/src/lib/gpu/gpu-knn.ts +++ b/src/lib/gpu/gpu-knn.ts @@ -176,8 +176,8 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { * candidate-rejection path is a single compare. Same O(N log N) total work * as the CPU KD-tree the kernel mirrors, just parallelised across queries. * - * The flattened tree is built by the caller (`KdTree.flatten`, typically - * off-thread via the `flattenKdTree` worker task) — this class only uploads + * The flat tree is built by the caller (`buildFlatKdTree`, typically + * off-thread via the worker task of the same name) — this class only uploads * and traverses it. * * Memory footprint: ~24 N bytes for the flattened tree (3 floats + 3 @@ -185,8 +185,8 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { */ class GpuKnn { /** - * @param tree - Prebuilt flattened KD-tree over the `n` local points - * (see `KdTree.flatten`; node splat ids are LOCAL indices). + * @param tree - Prebuilt flat KD-tree over the `n` local points + * (see `buildFlatKdTree`; node splat ids are LOCAL indices). * @param positions - Interleaved xyz for all `n` local points; queries * are the first `queryCount` of them (owned-first ordering). * @param n - Total local point count (tree size). @@ -268,11 +268,6 @@ class GpuKnn { compute.setParameter('nodeChildren', nChildrenBuf); compute.setParameter('outIndices', outBuf); - // Pack scratches reused across execute() calls — avoids allocating - // O(N) every call when the decimator runs KNN per block. - let nodePosPacked = new Float32Array(0); - let nodeChildrenPacked = new Uint32Array(0); - this.execute = async ( tree: FlatKdTree, positions: Float32Array, @@ -293,26 +288,10 @@ class GpuKnn { throw new Error(`GpuKnn: outNeighbours length ${outNeighbours.length} must be queryCount*k = ${queryCount * k}`); } - // Pack node positions xyz-interleaved. - if (nodePosPacked.length < n * 3) nodePosPacked = new Float32Array(n * 3); - const nodeX = tree.nodeX, nodeY = tree.nodeY, nodeZ = tree.nodeZ; - for (let i = 0; i < n; i++) { - nodePosPacked[i * 3 + 0] = nodeX[i]; - nodePosPacked[i * 3 + 1] = nodeY[i]; - nodePosPacked[i * 3 + 2] = nodeZ[i]; - } - // Pack node children (left, right) pairs. - if (nodeChildrenPacked.length < n * 2) nodeChildrenPacked = new Uint32Array(n * 2); - const nodeLeft = tree.nodeLeft, nodeRight = tree.nodeRight; - for (let i = 0; i < n; i++) { - nodeChildrenPacked[i * 2 + 0] = nodeLeft[i]; - nodeChildrenPacked[i * 2 + 1] = nodeRight[i]; - } - positionsBuf.write(0, positions, 0, n * 3); nSplatIdxBuf.write(0, tree.nodeSplatIdx, 0, n); - nPositionsBuf.write(0, nodePosPacked, 0, n * 3); - nChildrenBuf.write(0, nodeChildrenPacked, 0, n * 2); + nPositionsBuf.write(0, tree.nodePositions, 0, n * 3); + nChildrenBuf.write(0, tree.nodeChildren, 0, n * 2); compute.setParameter('rootIdx', tree.rootIdx); const numBatches = Math.ceil(queryCount / queriesPerBatch); diff --git a/src/lib/spatial/kd-tree.ts b/src/lib/spatial/kd-tree.ts index 654d7814..d494ac19 100644 --- a/src/lib/spatial/kd-tree.ts +++ b/src/lib/spatial/kd-tree.ts @@ -1,14 +1,14 @@ /** - * KD-tree over raw column arrays (build, nearest / k-nearest queries, GPU - * flatten). + * KD-tree over raw column arrays: a pointer-node tree for CPU queries + * (`KdTree`), and a direct-to-flat builder (`buildFlatKdTree`) emitting the + * GPU traversal layout without intermediate node objects. * * Engine-free by contract: worker tasks build trees off-thread, and the * worker bundle inlines its whole import graph — an engine import here would * embed playcanvas into dist/worker.mjs (see the note atop workers/tasks.ts). * - * Dimensionality is the number of columns: 3 for spatial consumers, arbitrary - * for k-means centroid assignment. `flatten()` assumes the first three - * columns are x, y, z. + * `KdTree` dimensionality is the number of columns: 3 for spatial consumers, + * arbitrary for k-means centroid assignment. `buildFlatKdTree` is 3D only. */ interface KdTreeNode { @@ -19,19 +19,18 @@ interface KdTreeNode { } /** - * The kd-tree flattened into GPU-friendly parallel typed arrays. For tree - * index `t`, the node holds splat `nodeSplatIdx[t]` whose position is - * `(nodeX[t], nodeY[t], nodeZ[t])`; children live at `nodeLeft[t]` / - * `nodeRight[t]` with the sentinel `0xFFFFFFFF` for missing children. The - * root is at index 0. + * A kd-tree in GPU-friendly flat typed arrays (the exact layout GpuKnn + * uploads). For tree index `t`, the node holds splat `nodeSplatIdx[t]` at + * position `nodePositions[t*3 + 0/1/2]`; children are + * `nodeChildren[t*2 + 0/1]` (left, right) with the sentinel `0xFFFFFFFF` + * for missing children. The root is at index 0. Node positions are + * denormalised (rather than indirected through `nodeSplatIdx` + the source + * position arrays) so a tree-walk does one read per visit instead of two. */ type FlatKdTree = { nodeSplatIdx: Uint32Array; - nodeX: Float32Array; - nodeY: Float32Array; - nodeZ: Float32Array; - nodeLeft: Uint32Array; - nodeRight: Uint32Array; + nodePositions: Float32Array; + nodeChildren: Uint32Array; rootIdx: number; }; @@ -73,6 +72,73 @@ const nthElement = (arr: Uint32Array, lo: number, hi: number, k: number, values: } }; +/** + * Build a 3D kd-tree directly into the flat GPU layout — no intermediate + * pointer-node graph (which costs ~50-80 B/node of JS objects and GC + * pressure, prohibitive at part-scale point counts). + * + * Same structure as `KdTree`'s build: median split via `nthElement`, axis + * cycling x→y→z by depth, two-point ranges ordered so the smaller value is + * the node and the larger its right child. Nodes are emitted in pre-order + * DFS (left subtree before right), root at index 0. Recursion depth is + * bounded by the median split (≈ log2 n). + * + * @param x - Point x column. + * @param y - Point y column. + * @param z - Point z column. + * @returns The flat tree over all points. + */ +const buildFlatKdTree = (x: Float32Array, y: Float32Array, z: Float32Array): FlatKdTree => { + const n = x.length; + const cols = [x, y, z]; + const indices = new Uint32Array(n); + for (let i = 0; i < n; i++) indices[i] = i; + + const nodeSplatIdx = new Uint32Array(n); + const nodePositions = new Float32Array(n * 3); + const nodeChildren = new Uint32Array(n * 2).fill(0xFFFFFFFF); + + let cursor = 0; + const emit = (splat: number): number => { + const t = cursor++; + nodeSplatIdx[t] = splat; + const t3 = t * 3; + nodePositions[t3] = x[splat]; + nodePositions[t3 + 1] = y[splat]; + nodePositions[t3 + 2] = z[splat]; + return t; + }; + + const build = (lo: number, hi: number, depth: number): number => { + const count = hi - lo + 1; + + if (count === 1) return emit(indices[lo]); + + const values = cols[depth % 3]; + + if (count === 2) { + if (values[indices[lo]] > values[indices[hi]]) { + const tmp = indices[lo]; indices[lo] = indices[hi]; indices[hi] = tmp; + } + const t = emit(indices[lo]); + nodeChildren[t * 2 + 1] = emit(indices[hi]); + return t; + } + + const mid = lo + (count >> 1); + nthElement(indices, lo, hi, mid, values); + + const t = emit(indices[mid]); + nodeChildren[t * 2] = build(lo, mid - 1, depth + 1); + nodeChildren[t * 2 + 1] = build(mid + 1, hi, depth + 1); + return t; + }; + + if (n > 0) build(0, n - 1, 0); + + return { nodeSplatIdx, nodePositions, nodeChildren, rootIdx: 0 }; +}; + class KdTree { root: KdTreeNode; readonly colData: ArrayLike[]; @@ -271,84 +337,6 @@ class KdTree { return { indices: resultIndices, distances: resultDist }; } - - /** - * Flatten the tree into GPU-friendly typed arrays (see {@link FlatKdTree}). - * Each tree node is assigned a tree-index in pre-order DFS. - * - * Positions are denormalised at each tree node (rather than indirected - * through `nodeSplatIdx` + the source position arrays) so a tree-walk - * does one read per visit instead of two. Costs 12 bytes/node extra. - * - * Layout assumes the first three columns are `x`, `y`, `z`. Callers with - * other dimensionalities must not call this. - * - * @returns Parallel arrays of length N where N = number of points. - */ - flatten(): FlatKdTree { - const n = this.numRows; - const nodeSplatIdx = new Uint32Array(n); - const nodeX = new Float32Array(n); - const nodeY = new Float32Array(n); - const nodeZ = new Float32Array(n); - const nodeLeft = new Uint32Array(n); - const nodeRight = new Uint32Array(n); - nodeLeft.fill(0xFFFFFFFF); - nodeRight.fill(0xFFFFFFFF); - - const x = this.colData[0], y = this.colData[1], z = this.colData[2]; - - // Iterative pre-order DFS: assign tree indices, then patch the parent's - // left/right slot when each child is visited. JS recursion blows the - // stack on heavily unbalanced trees, so we maintain the work stack - // ourselves. Encoded entries: nodeRef + (parentTreeIdx, side) where - // side ∈ {0 = left of parent, 1 = right of parent, 2 = root}. - // - // Max DFS depth is the tree's height. `build` is recursive and splits - // at the nthElement median, so the tree is near-balanced and its - // height is bounded by JS's recursion limit (~10K). A fixed 64 - // entries is enough for any tree this codebase can actually build - // (2^64 ≫ 10K) and avoids an `n+1`-sized scratch (~85 MB at N=17.9M). - const stackCap = 64; - const stackNode: KdTreeNode[] = [this.root]; - const stackParent = new Int32Array(stackCap); - const stackSide = new Uint8Array(stackCap); - stackParent[0] = -1; - stackSide[0] = 2; - let sp = 1; - - let cursor = 0; - const rootIdx = cursor; - while (sp > 0) { - sp--; - const node = stackNode[sp]; - const parent = stackParent[sp]; - const side = stackSide[sp]; - const treeIdx = cursor++; - const splat = node.index; - nodeSplatIdx[treeIdx] = splat; - nodeX[treeIdx] = x[splat]; - nodeY[treeIdx] = y[splat]; - nodeZ[treeIdx] = z[splat]; - if (side === 0) nodeLeft[parent] = treeIdx; - else if (side === 1) nodeRight[parent] = treeIdx; - // Push right then left so left is popped first (pre-order). - if (node.right) { - stackNode[sp] = node.right; - stackParent[sp] = treeIdx; - stackSide[sp] = 1; - sp++; - } - if (node.left) { - stackNode[sp] = node.left; - stackParent[sp] = treeIdx; - stackSide[sp] = 0; - sp++; - } - } - - return { nodeSplatIdx, nodeX, nodeY, nodeZ, nodeLeft, nodeRight, rootIdx }; - } } -export { KdTree, type KdTreeNode, type FlatKdTree }; +export { KdTree, buildFlatKdTree, type KdTreeNode, type FlatKdTree }; diff --git a/src/lib/workers/tasks.ts b/src/lib/workers/tasks.ts index 5a9bdca9..903d53ec 100644 --- a/src/lib/workers/tasks.ts +++ b/src/lib/workers/tasks.ts @@ -2,7 +2,7 @@ import type { TypedArray } from '../data-table/data-table'; import { knnQueryBlock } from '../decimate/knn-core'; import { mergeGroup, createMergeScratch, splatMass } from '../decimate/moment-match'; import { bestEdgeFor, bestOut, type RecostState } from '../decimate/recost-core'; -import { KdTree, type FlatKdTree } from '../spatial/kd-tree'; +import { buildFlatKdTree, type FlatKdTree } from '../spatial/kd-tree'; import { quantize1dColumns, type QuantizedColumns } from '../spatial/quantize-1d-core'; import { WebPCodec } from '../utils/webp-codec'; @@ -40,9 +40,9 @@ const taskHandlers = { return { result: webp, transfer: [webp.buffer as ArrayBuffer] }; }, - // Build + flatten a KD-tree over interleaved local positions (decimation - // GPU path: the flattened arrays upload straight into GpuKnn). - flattenKdTree: (args: { positions: Float32Array }): TaskOutput => { + // Build a flat KD-tree over interleaved local positions (decimation GPU + // path: the flat arrays upload straight into GpuKnn). + buildFlatKdTree: (args: { positions: Float32Array }): TaskOutput => { const n = args.positions.length / 3; const x = new Float32Array(n); const y = new Float32Array(n); @@ -52,12 +52,11 @@ const taskHandlers = { y[i] = args.positions[i * 3 + 1]; z[i] = args.positions[i * 3 + 2]; } - const flat = new KdTree([x, y, z]).flatten(); + const flat = buildFlatKdTree(x, y, z); return { result: flat, transfer: [ - flat.nodeSplatIdx.buffer, flat.nodeX.buffer, flat.nodeY.buffer, - flat.nodeZ.buffer, flat.nodeLeft.buffer, flat.nodeRight.buffer + flat.nodeSplatIdx.buffer, flat.nodePositions.buffer, flat.nodeChildren.buffer ] as ArrayBuffer[] }; }, diff --git a/test/kd-tree.test.mjs b/test/kd-tree.test.mjs index a4bb20a8..0c6ef01f 100644 --- a/test/kd-tree.test.mjs +++ b/test/kd-tree.test.mjs @@ -13,7 +13,7 @@ import assert from 'node:assert'; import { performance } from 'node:perf_hooks'; import { describe, it } from 'node:test'; -import { KdTree } from '../src/lib/spatial/kd-tree.js'; +import { KdTree, buildFlatKdTree } from '../src/lib/spatial/kd-tree.js'; const mulberry = (seed) => { let t = seed >>> 0; @@ -61,10 +61,10 @@ describe('KdTree', () => { } }); - it('flatten produces a traversable tree covering all points exactly once', () => { + it('buildFlatKdTree produces a traversable tree covering all points exactly once', () => { const n = 257; const [x, y, z] = randomCols(n, 3, 99); - const flat = new KdTree([x, y, z]).flatten(); + const flat = buildFlatKdTree(x, y, z); assert.strictEqual(flat.rootIdx, 0); const seen = new Set(); const stack = [flat.rootIdx]; @@ -73,15 +73,45 @@ describe('KdTree', () => { const splat = flat.nodeSplatIdx[t]; assert.ok(!seen.has(splat), 'splat appears once'); seen.add(splat); - assert.strictEqual(flat.nodeX[t], x[splat]); - assert.strictEqual(flat.nodeY[t], y[splat]); - assert.strictEqual(flat.nodeZ[t], z[splat]); - if (flat.nodeLeft[t] !== 0xFFFFFFFF) stack.push(flat.nodeLeft[t]); - if (flat.nodeRight[t] !== 0xFFFFFFFF) stack.push(flat.nodeRight[t]); + assert.strictEqual(flat.nodePositions[t * 3], x[splat]); + assert.strictEqual(flat.nodePositions[t * 3 + 1], y[splat]); + assert.strictEqual(flat.nodePositions[t * 3 + 2], z[splat]); + if (flat.nodeChildren[t * 2] !== 0xFFFFFFFF) stack.push(flat.nodeChildren[t * 2]); + if (flat.nodeChildren[t * 2 + 1] !== 0xFFFFFFFF) stack.push(flat.nodeChildren[t * 2 + 1]); } assert.strictEqual(seen.size, n); }); + it('buildFlatKdTree splitting planes are consistent with the cycling axis rule', () => { + // The GPU kernel derives each node's axis from traversal depth + // (x→y→z cycling), so the flat tree must respect it: every splat in + // the left subtree sits at-or-below the node on that axis, right + // subtree at-or-above. Ties may land on either side of an equal + // block, so compare with <= / >=. + const n = 1000; + const [x, y, z] = randomCols(n, 3, 4242); + const cols = [x, y, z]; + const flat = buildFlatKdTree(x, y, z); + const check = (t, depth) => { + if (t === 0xFFFFFFFF) return; + const axis = depth % 3; + const v = cols[axis][flat.nodeSplatIdx[t]]; + const left = flat.nodeChildren[t * 2]; + const right = flat.nodeChildren[t * 2 + 1]; + const walk = (s, cmp) => { + if (s === 0xFFFFFFFF) return; + assert.ok(cmp(cols[axis][flat.nodeSplatIdx[s]], v), `axis ${axis} split violated at node ${t}`); + walk(flat.nodeChildren[s * 2], cmp); + walk(flat.nodeChildren[s * 2 + 1], cmp); + }; + walk(left, (a, b) => a <= b); + walk(right, (a, b) => a >= b); + check(left, depth + 1); + check(right, depth + 1); + }; + check(flat.rootIdx, 0); + }); + it('builds in sub-quadratic time over a large all-identical point set', { timeout: 5000 }, () => { // Pre-fix the 2-way Lomuto partition degenerated to O(N^2) in the KD-tree // *build* (~5e9 ops at N=100k → tens of seconds), which is exactly how the From 27cb4b3ccc4156bbdc9bf814132749c7d1e48877 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 14:47:10 +0100 Subject: [PATCH 05/19] latest --- src/lib/decimate/edge-cost-cpu.ts | 134 +++++++++++----------- src/lib/decimate/priority.ts | 181 ++++++++++++++---------------- src/lib/decimate/recost-core.ts | 5 +- src/lib/decimate/select-recost.ts | 2 +- src/lib/gpu/gpu-edge-cost.ts | 152 +++++++++++++------------ src/lib/gpu/index.ts | 1 - test/decimate-edge-cost.test.mjs | 8 +- test/decimate-priority.test.mjs | 7 +- 8 files changed, 241 insertions(+), 249 deletions(-) diff --git a/src/lib/decimate/edge-cost-cpu.ts b/src/lib/decimate/edge-cost-cpu.ts index 9e38b1c8..cfb43af4 100644 --- a/src/lib/decimate/edge-cost-cpu.ts +++ b/src/lib/decimate/edge-cost-cpu.ts @@ -38,6 +38,13 @@ import { /** SH band-0 constant (f_dc → base colour: 0.5 + C0·f_dc). */ const C0 = 0.28209479177387814; +/** + * Floats per splat in the interleaved cost cache (see {@link buildSplatCache} + * for the layout). Single source of truth — the GPU kernels and the re-costed + * selection state share this exact layout. + */ +const CACHE_STRIDE = 16; + /** * Scale-free colour dissimilarity weight. Unlike the field-L2's own colour * sensitivity (which vanishes ∝σ³ for faint splats), this term keeps @@ -56,29 +63,28 @@ const PI_1_5 = Math.PI ** 1.5; const TWO_PI_1_5 = (2 * Math.PI) ** 1.5; /** - * Per-splat derived quantities for the L2 field cost. `sig` holds each - * covariance's 6 unique components [xx, xy, xz, yy, yz, zz]; `sqrtDet` is - * √|Σ| = sx·sy·sz; `alpha` is the linear opacity; `base` is the 3-channel base - * colour; `baseN2` is |base|²; `mass` is the area·α merge weight. + * Build the per-splat cost cache: {@link CACHE_STRIDE} interleaved floats per + * splat, written into `out` (reusable scratch; only the leading n·16 floats + * are touched). Layout per row: + * + * [0..2] mean xyz + * [3..8] covariance Σ (xx, xy, xz, yy, yz, zz) — unregularized + * [9] sqrtDet = √|Σ| = sx·sy·sz + * [10] alpha (linear opacity) + * [11] mass (area·α merge weight) + * [12..14] base colour (0.5 + C0·f_dc per channel) + * [15] |base|² + * + * This exact buffer uploads to the GPU kernels and persists per-gaussian for + * the re-costed selection, so CPU and GPU read identical per-splat inputs. + * + * @param view - Splat columns (only the first 3 colour components are read). + * @param out - Destination, at least n·{@link CACHE_STRIDE} floats. + * @returns The splat count n. */ -type CostCache = { - sig: Float32Array; - sqrtDet: Float32Array; - alpha: Float32Array; - base: Float32Array; - baseN2: Float32Array; - mass: Float32Array; -}; - -const buildCostCache = (view: SplatView): CostCache => { - const { geo, color, colorDim } = view; +const buildSplatCache = (view: SplatView, out: Float32Array): number => { + const { pos, geo, color, colorDim } = view; const n = geo.length / 8; - const sig = new Float32Array(n * 6); - const sqrtDet = new Float32Array(n); - const alpha = new Float32Array(n); - const base = new Float32Array(n * 3); - const baseN2 = new Float32Array(n); - const mass = new Float32Array(n); const R = new Float32Array(9); const S = new Float32Array(9); @@ -97,23 +103,23 @@ const buildCostCache = (view: SplatView): CostCache => { quatToRotmat(qw, qx, qy, qz, R, 0); sigmaFromRotVar(R, 0, sx * sx, sy * sy, sz * sz, S, 0); - const i6 = 6 * i; - sig[i6] = S[0]; sig[i6 + 1] = S[1]; sig[i6 + 2] = S[2]; - sig[i6 + 3] = S[4]; sig[i6 + 4] = S[5]; sig[i6 + 5] = S[8]; - - sqrtDet[i] = sx * sy * sz; - alpha[i] = a; - mass[i] = a * ellipsoidArea(sx, sy, sz) + 1e-30; - + const o = i * CACHE_STRIDE; const i3 = 3 * i; + out[o] = pos[i3]; out[o + 1] = pos[i3 + 1]; out[o + 2] = pos[i3 + 2]; + out[o + 3] = S[0]; out[o + 4] = S[1]; out[o + 5] = S[2]; + out[o + 6] = S[4]; out[o + 7] = S[5]; out[o + 8] = S[8]; + out[o + 9] = sx * sy * sz; + out[o + 10] = a; + out[o + 11] = a * ellipsoidArea(sx, sy, sz) + 1e-30; + const br = 0.5 + C0 * color[i * colorDim]; const bg = 0.5 + C0 * color[i * colorDim + 1]; const bb = 0.5 + C0 * color[i * colorDim + 2]; - base[i3] = br; base[i3 + 1] = bg; base[i3 + 2] = bb; - baseN2[i] = br * br + bg * bg + bb * bb; + out[o + 12] = br; out[o + 13] = bg; out[o + 14] = bb; + out[o + 15] = br * br + bg * bg + bb * bb; } - return { sig, sqrtDet, alpha, base, baseN2, mass }; + return n; }; // dᵀ(A)⁻¹d and √|A| for a symmetric 3×3 A = [xx,xy,xz,yy,yz,zz]; returns the @@ -138,30 +144,24 @@ const crossG = ( }; /** - * L2 field-error cost of merging splats `i` and `j` of the view. + * L2 field-error cost of merging splats `i` and `j`. * - * @param view - Splat columns. - * @param cache - Per-splat cache from {@link buildCostCache}. - * @param i - First splat (view row). - * @param j - Second splat (view row). + * @param cache - Per-splat cache from {@link buildSplatCache}. + * @param i - First splat (cache row). + * @param j - Second splat (cache row). * @returns The edge cost (≥ 0). */ -const computeEdgeCostView = ( - view: SplatView, - cache: CostCache, +const computeEdgeCost = ( + cache: Float32Array, i: number, j: number ): number => { - const { pos } = view; - const { sig, sqrtDet, alpha, base, baseN2, mass } = cache; - - const i3 = 3 * i, j3 = 3 * j; - const i6 = 6 * i, j6 = 6 * j; + const io = i * CACHE_STRIDE, jo = j * CACHE_STRIDE; - const mux = pos[i3], muy = pos[i3 + 1], muz = pos[i3 + 2]; - const mvx = pos[j3], mvy = pos[j3 + 1], mvz = pos[j3 + 2]; + const mux = cache[io], muy = cache[io + 1], muz = cache[io + 2]; + const mvx = cache[jo], mvy = cache[jo + 1], mvz = cache[jo + 2]; - const mi = mass[i], mj = mass[j]; + const mi = cache[io + 11], mj = cache[jo + 11]; const W = mi + mj; const pi = W > 0 ? mi / W : 0.5; const pj = 1 - pi; @@ -175,12 +175,12 @@ const computeEdgeCostView = ( const djx = mvx - mmx, djy = mvy - mmy, djz = mvz - mmz; // Merged covariance: Σ pₖ(δₖδₖᵀ + Σₖ) + EPS_COV·I (law of total variance). - const sxx = pi * (dix * dix + sig[i6]) + pj * (djx * djx + sig[j6]) + EPS_COV; - const sxy = pi * (dix * diy + sig[i6 + 1]) + pj * (djx * djy + sig[j6 + 1]); - const sxz = pi * (dix * diz + sig[i6 + 2]) + pj * (djx * djz + sig[j6 + 2]); - const syy = pi * (diy * diy + sig[i6 + 3]) + pj * (djy * djy + sig[j6 + 3]) + EPS_COV; - const syz = pi * (diy * diz + sig[i6 + 4]) + pj * (djy * djz + sig[j6 + 4]); - const szz = pi * (diz * diz + sig[i6 + 5]) + pj * (djz * djz + sig[j6 + 5]) + EPS_COV; + const sxx = pi * (dix * dix + cache[io + 3]) + pj * (djx * djx + cache[jo + 3]) + EPS_COV; + const sxy = pi * (dix * diy + cache[io + 4]) + pj * (djx * djy + cache[jo + 4]); + const sxz = pi * (dix * diz + cache[io + 5]) + pj * (djx * djz + cache[jo + 5]); + const syy = pi * (diy * diy + cache[io + 6]) + pj * (djy * djy + cache[jo + 6]) + EPS_COV; + const syz = pi * (diy * diz + cache[io + 7]) + pj * (djy * djz + cache[jo + 7]); + const szz = pi * (diz * diz + cache[io + 8]) + pj * (djz * djz + cache[jo + 8]) + EPS_COV; const detm = Math.max( sxx * (syy * szz - syz * syz) - sxy * (sxy * szz - syz * sxz) + sxz * (sxy * syz - syy * sxz), @@ -215,17 +215,17 @@ const computeEdgeCostView = ( const alphaM = Math.min(1, W / Math.max(ellipsoidArea(s0, s1, s2), 1e-30)); // Base colours and their dots (merged colour = mass-weighted average). - const bir = base[i3], big = base[i3 + 1], bib = base[i3 + 2]; - const bjr = base[j3], bjg = base[j3 + 1], bjb = base[j3 + 2]; + const bir = cache[io + 12], big = cache[io + 13], bib = cache[io + 14]; + const bjr = cache[jo + 12], bjg = cache[jo + 13], bjb = cache[jo + 14]; const bij = bir * bjr + big * bjg + bib * bjb; // base_i · base_j - const bni = baseN2[i]; // base_i · base_i - const bnj = baseN2[j]; // base_j · base_j + const bni = cache[io + 15]; // base_i · base_i + const bnj = cache[jo + 15]; // base_j · base_j const bim = pi * bni + pj * bij; // base_i · base_m const bjm = pi * bij + pj * bnj; // base_j · base_m const bnm = pi * pi * bni + 2 * pi * pj * bij + pj * pj * bnj; // base_m · base_m - const ai = alpha[i], aj = alpha[j], am = alphaM; - const sdi = sqrtDet[i], sdj = sqrtDet[j]; + const ai = cache[io + 10], aj = cache[jo + 10], am = alphaM; + const sdi = cache[io + 9], sdj = cache[jo + 9]; // Self terms ⟨G,G⟩ = π^{3/2}·√|Σ|. const selfI = ai * ai * bni * PI_1_5 * sdi; @@ -234,16 +234,16 @@ const computeEdgeCostView = ( // Cross terms (amplitude dot × Gaussian overlap). const cIJ = crossG(sdi, sdj, - sig[i6] + sig[j6], sig[i6 + 1] + sig[j6 + 1], sig[i6 + 2] + sig[j6 + 2], - sig[i6 + 3] + sig[j6 + 3], sig[i6 + 4] + sig[j6 + 4], sig[i6 + 5] + sig[j6 + 5], + cache[io + 3] + cache[jo + 3], cache[io + 4] + cache[jo + 4], cache[io + 5] + cache[jo + 5], + cache[io + 6] + cache[jo + 6], cache[io + 7] + cache[jo + 7], cache[io + 8] + cache[jo + 8], mux - mvx, muy - mvy, muz - mvz); const cIM = crossG(sdi, sqrtDetM, - sig[i6] + sxx, sig[i6 + 1] + sxy, sig[i6 + 2] + sxz, - sig[i6 + 3] + syy, sig[i6 + 4] + syz, sig[i6 + 5] + szz, + cache[io + 3] + sxx, cache[io + 4] + sxy, cache[io + 5] + sxz, + cache[io + 6] + syy, cache[io + 7] + syz, cache[io + 8] + szz, dix, diy, diz); const cJM = crossG(sdj, sqrtDetM, - sig[j6] + sxx, sig[j6 + 1] + sxy, sig[j6 + 2] + sxz, - sig[j6 + 3] + syy, sig[j6 + 4] + syz, sig[j6 + 5] + szz, + cache[jo + 3] + sxx, cache[jo + 4] + sxy, cache[jo + 5] + sxz, + cache[jo + 6] + syy, cache[jo + 7] + syz, cache[jo + 8] + szz, djx, djy, djz); const E = selfI + selfJ + selfM + @@ -259,4 +259,4 @@ const computeEdgeCostView = ( return (E < 0 ? 0 : E) + COLOR_WEIGHT * (dc2 < 0 ? 0 : dc2); }; -export { buildCostCache, computeEdgeCostView, COLOR_WEIGHT, type CostCache }; +export { buildSplatCache, computeEdgeCost, CACHE_STRIDE, COLOR_WEIGHT }; diff --git a/src/lib/decimate/priority.ts b/src/lib/decimate/priority.ts index c941f6bc..fbf39268 100644 --- a/src/lib/decimate/priority.ts +++ b/src/lib/decimate/priority.ts @@ -1,12 +1,12 @@ import { type GraphicsDevice } from 'playcanvas'; -import { buildCostCache, computeEdgeCostView } from './edge-cost-cpu'; +import { buildSplatCache, computeEdgeCost, CACHE_STRIDE } from './edge-cost-cpu'; import { collectBlock, verifyAndFixKnn, toGlobalNeighbors, KNN_FIXED, type BlockLocals } from './knn-blocks'; import { KNN_SENTINEL } from './knn-core'; import { type SplatView } from './moment-match'; import { type BlockRange, type ResidentPositions } from './partition'; import { type ChunkData, type ChunkDataPool, type ChunkSource } from '../chunk'; -import { GpuEdgeCost, SPLAT_STRIDE, type EdgeCostCache } from '../gpu/gpu-edge-cost'; +import { GpuEdgeCost } from '../gpu/gpu-edge-cost'; import { GpuKnn } from '../gpu/gpu-knn'; import { WorkerQueue } from '../workers'; @@ -41,8 +41,8 @@ type PriorityContext = { k: number; /** * Optional resident splat-cache output for re-costed selection - * (SPLAT_STRIDE floats per gaussian, packGpuCache row layout), filled for - * every owned gaussian. + * (CACHE_STRIDE floats per gaussian, buildSplatCache row layout), filled + * for every owned gaussian. */ cacheOut?: Float32Array; /** Optional resident neighbour-id output (k per gaussian, KNN_SENTINEL padded). */ @@ -72,13 +72,18 @@ type BlockView = { * @param blockIdx - Which block. * @param extraGlobals - Sorted out-of-block rows to append after the owned rows. * @param includeOther - Also gather the `other` layer (merge pass only). + * @param colorComponents - Colour components to copy per row (default: all). + * The quality cost reads only DC, so its pass gathers 3 — at SH band 3 that + * is 16× less colour RAM and copy traffic per block. The read itself still + * decodes whole rows (row-interleaved sources); only the view copy narrows. * @returns The gathered block view. */ const gatherBlockView = async ( ctx: Pick, blockIdx: number, extraGlobals: Uint32Array, - includeOther = false + includeOther = false, + colorComponents?: number ): Promise => { const { source, pool, pos, order, blocks } = ctx; const block = blocks[blockIdx]; @@ -88,14 +93,15 @@ const gatherBlockView = async ( const { layouts, availableLayers } = source.meta; const colorDim = layouts.color!.stride >> 2; + const viewColorDim = Math.min(colorComponents ?? colorDim, colorDim); const wantOther = includeOther && availableLayers.has('other') && (layouts.other?.stride ?? 0) > 0; const otherDim = wantOther ? layouts.other!.stride >> 2 : 0; const view: SplatView = { pos: new Float32Array(n * 3), geo: new Float32Array(n * 8), - color: new Float32Array(n * colorDim), - colorDim + color: new Float32Array(n * viewColorDim), + colorDim: viewColorDim }; const other = wantOther ? new Uint32Array(n * otherDim) : undefined; @@ -115,7 +121,15 @@ const gatherBlockView = async ( other: othCd }); view.geo.set(new Float32Array(geoCd.data, 0, count * 8), (rowBase + off) * 8); - view.color.set(new Float32Array(colCd.data, 0, count * colorDim), (rowBase + off) * colorDim); + const colSrc = new Float32Array(colCd.data, 0, count * colorDim); + if (viewColorDim === colorDim) { + view.color.set(colSrc, (rowBase + off) * colorDim); + } else { + for (let r = 0; r < count; r++) { + const src = r * colorDim, dst = (rowBase + off + r) * viewColorDim; + for (let c = 0; c < viewColorDim; c++) view.color[dst + c] = colSrc[src + c]; + } + } if (othCd) other!.set(new Uint32Array(othCd.data, 0, count * otherDim), (rowBase + off) * otherDim); geoCd.release(); colCd.release(); @@ -156,28 +170,6 @@ const indexOfSorted = (sorted: Uint32Array, g: number): number => { return -1; }; -// Pack the per-splat cache for the GPU kernel: build the CPU cache once and -// interleave it into the kernel's SPLAT_STRIDE-wide layout (mean from the view, -// the rest from the cache), so the GPU reads byte-identical per-splat inputs. -const packGpuCache = (view: SplatView): EdgeCostCache => { - const { pos } = view; - const c = buildCostCache(view); - const n = view.geo.length / 8; - const d = new Float32Array(n * SPLAT_STRIDE); - for (let i = 0; i < n; i++) { - const o = i * SPLAT_STRIDE, i6 = i * 6, i3 = i * 3; - d[o] = pos[i3]; d[o + 1] = pos[i3 + 1]; d[o + 2] = pos[i3 + 2]; - d[o + 3] = c.sig[i6]; d[o + 4] = c.sig[i6 + 1]; d[o + 5] = c.sig[i6 + 2]; - d[o + 6] = c.sig[i6 + 3]; d[o + 7] = c.sig[i6 + 4]; d[o + 8] = c.sig[i6 + 5]; - d[o + 9] = c.sqrtDet[i]; - d[o + 10] = c.alpha[i]; - d[o + 11] = c.mass[i]; - d[o + 12] = c.base[i3]; d[o + 13] = c.base[i3 + 1]; d[o + 14] = c.base[i3 + 2]; - d[o + 15] = c.baseN2[i]; - } - return { splatData: d, numSplats: n }; -}; - /** * The priority pass (heavy read 1): per block — exact global KNN, edge costs * for each owned gaussian's k neighbours, reduction to the best K candidates @@ -227,7 +219,7 @@ const runPriorityPass = async ( try { if (device) { gpuKnn = new GpuKnn(device, maxLocalN, k); - gpuCost = new GpuEdgeCost(device, maxLocalN, maxOwned * k); + gpuCost = new GpuEdgeCost(device, maxLocalN, k); } let next: Prepared | null = blocks.length > 0 ? prepare(0) : null; @@ -242,21 +234,32 @@ const runPriorityPass = async ( const nbGlobal = toGlobalNeighbors(locals, nbLocal); verifyAndFixKnn(pos, order, blocks, bi, locals, k, nbGlobal, nbLocal); + const slots = nOwned * k; + // Externals: referenced rows outside the owned range (halo members - // and verification-fixed neighbours), sorted for the gather. - const extRow = new Map(); - for (let s = 0; s < nOwned * k; s++) { + // and verification-fixed neighbours) — count, collect, sort, dedup. + let extCount = 0; + for (let s = 0; s < slots; s++) { + const l = nbLocal[s]; + if (l === KNN_SENTINEL || l < nOwned) continue; + if (l === KNN_FIXED && indexOfSorted(owned, nbGlobal[s]) >= 0) continue; + extCount++; + } + const extSorted = new Uint32Array(extCount); + extCount = 0; + for (let s = 0; s < slots; s++) { const l = nbLocal[s]; if (l === KNN_SENTINEL || l < nOwned) continue; const g = nbGlobal[s]; - if (l !== KNN_FIXED) { - if (!extRow.has(g)) extRow.set(g, 0); - } else if (indexOfSorted(owned, g) < 0 && !extRow.has(g)) { - extRow.set(g, 0); - } + if (l === KNN_FIXED && indexOfSorted(owned, g) >= 0) continue; + extSorted[extCount++] = g; + } + extSorted.sort(); + let uniq = 0; + for (let i = 0; i < extCount; i++) { + if (i === 0 || extSorted[i] !== extSorted[i - 1]) extSorted[uniq++] = extSorted[i]; } - const extraGlobals = Uint32Array.from(extRow.keys()).sort(); - for (let i = 0; i < extraGlobals.length; i++) extRow.set(extraGlobals[i], nOwned + i); + const extraGlobals = extSorted.subarray(0, uniq); // Verification-fixed externals are not bounded by the halo cap, so // a pathological block's view can exceed the preallocated cost @@ -266,72 +269,49 @@ const runPriorityPass = async ( if (gpuCost && viewN > gpuCostCapacity) { gpuCost.destroy(); gpuCostCapacity = Math.ceil(viewN * 1.1); - gpuCost = new GpuEdgeCost(device!, gpuCostCapacity, maxOwned * k); + gpuCost = new GpuEdgeCost(device!, gpuCostCapacity, k); } - const { view } = await gatherBlockView(ctx, bi, extraGlobals); + const { view } = await gatherBlockView(ctx, bi, extraGlobals, false, 3); - // Edge lists in owned-major order (view-local endpoints). - const edgeI = new Uint32Array(nOwned * k); - const edgeJ = new Uint32Array(nOwned * k); - const edgeNb = new Uint32Array(nOwned * k); // global neighbour per edge - const edgeOf = new Uint32Array(nOwned + 1); // CSR into the edge list per owned row - let e = 0; - for (let qi = 0; qi < nOwned; qi++) { - edgeOf[qi] = e; - for (let s = 0; s < k; s++) { - const l = nbLocal[qi * k + s]; - if (l === KNN_SENTINEL) continue; - const g = nbGlobal[qi * k + s]; - let row: number; - if (l !== KNN_FIXED) { - row = l < nOwned ? l : extRow.get(g)!; - } else { - const oi = indexOfSorted(owned, g); - row = oi >= 0 ? oi : extRow.get(g)!; - } - edgeI[e] = qi; - edgeJ[e] = row; - edgeNb[e] = g; - e++; + // Per-splat cost cache, built once for the GPU upload, the CPU + // path, and the re-costed selection's resident copy alike. + const cache = new Float32Array(viewN * CACHE_STRIDE); + buildSplatCache(view, cache); + + // Translate neighbour slots in place: local/global → view rows + // (dense-slot edge model; sentinel slots stay sentinel). + for (let s = 0; s < slots; s++) { + const l = nbLocal[s]; + if (l === KNN_SENTINEL || l < nOwned) continue; + const g = nbGlobal[s]; + if (l === KNN_FIXED) { + const oi = indexOfSorted(owned, g); + nbLocal[s] = oi >= 0 ? oi : nOwned + indexOfSorted(extraGlobals, g); + } else { + nbLocal[s] = nOwned + indexOfSorted(extraGlobals, g); } } - edgeOf[nOwned] = e; - const costs = new Float32Array(e); - const packed = device ? packGpuCache(view) : undefined; - const cpuCache = device ? undefined : buildCostCache(view); + const blockCosts = new Float32Array(slots); if (device) { - await gpuCost!.execute(packed!, edgeI.subarray(0, e), edgeJ.subarray(0, e), costs); + await gpuCost!.execute(cache, viewN, nbLocal, blockCosts); } else { - for (let i = 0; i < e; i++) { - costs[i] = computeEdgeCostView(view, cpuCache!, edgeI[i], edgeJ[i]); + for (let s = 0; s < slots; s++) { + const row = nbLocal[s]; + blockCosts[s] = row === KNN_SENTINEL ? + 0 : + computeEdgeCost(cache, (s / k) | 0, row); } } - // Persist owned rows for re-costed selection: the packed splat - // cache (identical layout on both paths) and the global neighbour - // ids (sentinel-padded). + // Persist owned rows for re-costed selection: the splat cache + // (identical layout on both paths) and the global neighbour ids + // (sentinel-padded). if (ctx.cacheOut) { const CO = ctx.cacheOut; - if (packed) { - for (let qi = 0; qi < nOwned; qi++) { - CO.set(packed.splatData.subarray(qi * SPLAT_STRIDE, (qi + 1) * SPLAT_STRIDE), owned[qi] * SPLAT_STRIDE); - } - } else { - const c = cpuCache!; - for (let qi = 0; qi < nOwned; qi++) { - const o = owned[qi] * SPLAT_STRIDE; - const q6 = qi * 6, q3 = qi * 3; - CO[o] = view.pos[q3]; CO[o + 1] = view.pos[q3 + 1]; CO[o + 2] = view.pos[q3 + 2]; - CO[o + 3] = c.sig[q6]; CO[o + 4] = c.sig[q6 + 1]; CO[o + 5] = c.sig[q6 + 2]; - CO[o + 6] = c.sig[q6 + 3]; CO[o + 7] = c.sig[q6 + 4]; CO[o + 8] = c.sig[q6 + 5]; - CO[o + 9] = c.sqrtDet[qi]; - CO[o + 10] = c.alpha[qi]; - CO[o + 11] = c.mass[qi]; - CO[o + 12] = c.base[q3]; CO[o + 13] = c.base[q3 + 1]; CO[o + 14] = c.base[q3 + 2]; - CO[o + 15] = c.baseN2[qi]; - } + for (let qi = 0; qi < nOwned; qi++) { + CO.set(cache.subarray(qi * CACHE_STRIDE, (qi + 1) * CACHE_STRIDE), owned[qi] * CACHE_STRIDE); } } if (ctx.neighborsOut) { @@ -341,13 +321,17 @@ const runPriorityPass = async ( } } - // Reduce to best K candidates per owned gaussian (ascending by cost). + // Reduce to best K candidates per owned gaussian (ascending by + // cost; sentinel slots are skipped by id — their cost values are + // never read). const bestIdx = new Uint32Array(K); const bestCost = new Float64Array(K); for (let qi = 0; qi < nOwned; qi++) { let size = 0; - for (let s = edgeOf[qi]; s < edgeOf[qi + 1]; s++) { - const c = costs[s]; + const base = qi * k; + for (let s = 0; s < k; s++) { + if (nbLocal[base + s] === KNN_SENTINEL) continue; + const c = blockCosts[base + s]; if (!Number.isFinite(c)) continue; if (size === K && c >= bestCost[K - 1]) continue; let at = size < K ? size : K - 1; @@ -357,7 +341,7 @@ const runPriorityPass = async ( at--; } bestCost[at] = c; - bestIdx[at] = edgeNb[s]; + bestIdx[at] = nbGlobal[base + s]; size = Math.min(size + 1, K); } const g = owned[qi]; @@ -378,7 +362,6 @@ const runPriorityPass = async ( export { runPriorityPass, gatherBlockView, - packGpuCache, indexOfSorted, HALO_FACTOR, HALO_CAP, diff --git a/src/lib/decimate/recost-core.ts b/src/lib/decimate/recost-core.ts index e3c6ab9c..0b69f6cb 100644 --- a/src/lib/decimate/recost-core.ts +++ b/src/lib/decimate/recost-core.ts @@ -11,11 +11,10 @@ * Engine-free; no allocation in the hot paths beyond small local scratch. */ -import { COLOR_WEIGHT } from './edge-cost-cpu'; +import { CACHE_STRIDE, COLOR_WEIGHT } from './edge-cost-cpu'; import { EPS_COV, ellipsoidArea } from './moment-match'; -/** Floats per splat in the resident cache (packGpuCache layout). */ -export const CACHE_STRIDE = 16; +export { CACHE_STRIDE }; const NO_CANDIDATE = 0xFFFFFFFF; const NIL = 0xFFFFFFFF; diff --git a/src/lib/decimate/select-recost.ts b/src/lib/decimate/select-recost.ts index a548b4df..50c3cc01 100644 --- a/src/lib/decimate/select-recost.ts +++ b/src/lib/decimate/select-recost.ts @@ -46,7 +46,7 @@ type RecostInputs = { K: number; /** * Resident per-splat cache, {@link CACHE_STRIDE} floats per splat - * (packGpuCache row layout), persisted by the priority pass. + * (buildSplatCache row layout), persisted by the priority pass. * SharedArrayBuffer-backed for the parallel path. */ splatCache: Float32Array; diff --git a/src/lib/gpu/gpu-edge-cost.ts b/src/lib/gpu/gpu-edge-cost.ts index 395b6883..2fa63ad2 100644 --- a/src/lib/gpu/gpu-edge-cost.ts +++ b/src/lib/gpu/gpu-edge-cost.ts @@ -15,32 +15,39 @@ import { UniformFormat } from 'playcanvas'; -/** Per-splat interleaved stride in the `splat` storage buffer (see EdgeCostCache). */ -export const SPLAT_STRIDE = 16; +import { CACHE_STRIDE } from '../decimate/edge-cost-cpu'; + +/** Per-splat interleaved stride in the `splat` storage buffer (= the CPU cost-cache layout). */ +export const SPLAT_STRIDE = CACHE_STRIDE; /** - * WGSL kernel: per-edge L2 field-error cost (mirrors `computeEdgeCostView` in + * WGSL kernel: per-slot L2 field-error cost (mirrors `computeEdgeCost` in * `decimate/edge-cost-cpu.ts`). * - * Each thread = one edge (i, j). It reads the per-splat cache for both - * endpoints, moment-matches the merged Gaussian m (mean, covariance, opacity), - * and evaluates + * Dense-slot edge model: slot s belongs to owned row qi = s / K and holds a + * view-local neighbour row (or the 0xFFFFFFFF sentinel for an empty slot, + * whose cost output is never read — the reduction skips sentinel slots by + * id). Each thread = one slot. + * It reads the per-splat cache for both endpoints, moment-matches the merged + * Gaussian m (mean, covariance, opacity), and evaluates * E = ‖αᵢcᵢGᵢ + αⱼcⱼGⱼ − αₘcₘGₘ‖² * in closed form via Gaussian–Gaussian L2 products ⟨G_a,G_b⟩. No Monte-Carlo * sampling and no appearance loop — the amplitude is α × base colour, so only * the packed base colour (3) and covariance (6) per splat are needed. * + * @param k - Compile-time K, neighbour slots per owned row. * @returns WGSL source. */ -const edgeCostWgsl = () => /* wgsl */` +const edgeCostWgsl = (k: number) => /* wgsl */` struct Uniforms { - edgeCount: u32, + slotBase: u32, + slotCount: u32, } @group(0) @binding(0) var uniforms: Uniforms; -// Edge list for the current dispatch batch (host uploads each batch to offset 0). -@group(0) @binding(1) var edgesI: array; -@group(0) @binding(2) var edgesJ: array; +// Neighbour rows for the current dispatch batch (host uploads each batch to +// offset 0): view-local row per slot, 0xFFFFFFFF for empty slots. +@group(0) @binding(1) var nbRow: array; // Per-splat cache, interleaved ${SPLAT_STRIDE}-wide: // [0..2] mean xyz // [3..8] covariance Σ (xx, xy, xz, yy, yz, zz) @@ -49,9 +56,12 @@ struct Uniforms { // [11] mass (area·α merge weight) // [12..14] base colour (r, g, b) // [15] |base|² -@group(0) @binding(3) var splat: array; -// Output: cost per edge. -@group(0) @binding(4) var costs: array; +@group(0) @binding(2) var splat: array; +// Output: cost per slot. +@group(0) @binding(3) var costs: array; + +const K: u32 = ${k}u; +const SENTINEL: u32 = 0xFFFFFFFFu; const EPS_COV: f32 = 1e-8; const PI_1_5: f32 = 5.5683279968317084; // π^{3/2} @@ -89,10 +99,18 @@ fn crossG(sdA: f32, sdB: f32, m: array, d: vec3f) -> f32 { @compute @workgroup_size(64) fn main(@builtin(global_invocation_id) gid: vec3u) { let bid = gid.x; - if (bid >= uniforms.edgeCount) { return; } + if (bid >= uniforms.slotCount) { return; } + + // Empty slot: the reduction skips sentinel slots by id, so the cost value + // is never read (WGSL cannot materialize an f32 infinity to park here). + let row = nbRow[bid]; + if (row == SENTINEL) { + costs[bid] = 0.0; + return; + } - let io = edgesI[bid] * ${SPLAT_STRIDE}u; - let jo = edgesJ[bid] * ${SPLAT_STRIDE}u; + let io = ((uniforms.slotBase + bid) / K) * ${SPLAT_STRIDE}u; + let jo = row * ${SPLAT_STRIDE}u; let mui = vec3f(splat[io + 0u], splat[io + 1u], splat[io + 2u]); let si = array(splat[io + 3u], splat[io + 4u], splat[io + 5u], splat[io + 6u], splat[io + 7u], splat[io + 8u]); @@ -184,31 +202,19 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { } `; -/** - * Per-splat cache for the edge-cost kernel: a single interleaved buffer of - * {@link SPLAT_STRIDE} floats per splat (mean, covariance, √det, alpha, mass, - * base colour, |base|²), built by `packGpuCache` from the CPU `buildCostCache` - * so the GPU reads identical per-splat inputs. - */ -interface EdgeCostCache { - /** Interleaved per-splat cache, length {@link SPLAT_STRIDE}·N. */ - splatData: Float32Array; - /** Number of splats. */ - numSplats: number; -} - /** * GPU edge-cost evaluator. * - * Each compute thread evaluates the L2 field-error cost for one edge (i, j), - * mirroring the CPU `computeEdgeCostView` in `decimate/edge-cost-cpu.ts`. - * Output is `costs[e] = cost for edge e`. + * Each compute thread evaluates the L2 field-error cost for one dense slot + * (owned row qi = slot / K vs its slot's neighbour row), mirroring the CPU + * `computeEdgeCost` in `decimate/edge-cost-cpu.ts`. Empty (sentinel) slots + * cost +Inf. Output is `costs[s] = cost for slot s`. */ class GpuEdgeCost { execute: ( - cache: EdgeCostCache, - edgeI: Uint32Array, - edgeJ: Uint32Array, + splatData: Float32Array, + numSplats: number, + nbRows: Uint32Array, outCosts: Float32Array ) => Promise; destroy: () => void; @@ -216,16 +222,18 @@ class GpuEdgeCost { /** * @param device - PlayCanvas GraphicsDevice (WebGPU). * @param maxN - Maximum number of splats. - * @param maxE - Maximum number of edges in a single dispatch. + * @param k - Neighbour slots per owned row. */ - constructor(device: GraphicsDevice, maxN: number, maxE: number) { + constructor(device: GraphicsDevice, maxN: number, k: number) { const workgroupSize = 64; - const edgesPerBatch = 1024 * workgroupSize; // 65,536 + // Slots per dispatch: bounded by the 65,535 workgroups-per-dimension + // limit. One blocking readback per ~4.2M slots (a 2M-row block costs + // 8 round trips, not 512 as with 65,536-edge batches). + const slotsPerBatch = 65535 * workgroupSize; // 4,194,240 const bindGroupFormat = new BindGroupFormat(device, [ new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE), - new BindStorageBufferFormat('edgesI', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('edgesJ', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('nbRow', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('splat', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('costs', SHADERSTAGE_COMPUTE) ]); @@ -233,11 +241,12 @@ class GpuEdgeCost { const shader = new Shader(device, { name: 'compute-edge-cost', shaderLanguage: SHADERLANGUAGE_WGSL, - cshader: edgeCostWgsl(), + cshader: edgeCostWgsl(k), // @ts-ignore computeUniformBufferFormats: { uniforms: new UniformBufferFormat(device, [ - new UniformFormat('edgeCount', UNIFORMTYPE_UINT) + new UniformFormat('slotBase', UNIFORMTYPE_UINT), + new UniformFormat('slotCount', UNIFORMTYPE_UINT) ]) }, // @ts-ignore @@ -259,66 +268,63 @@ class GpuEdgeCost { const splatBuf = new StorageBuffer(device, maxN * SPLAT_STRIDE * 4, BUFFERUSAGE_COPY_DST); - // Two parallel u32 edge buffers, sized to a single dispatch batch (not - // the full N·k edge list): execute uploads each batch's slice before its - // dispatch, keeping these off the ~2 GB per-binding limit. - const edgesIBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); - const edgesJBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); + // Neighbour-row buffer sized to a single dispatch batch (not the full + // N·k slot list): execute uploads each batch's slice before its + // dispatch, keeping it off the ~2 GB per-binding limit. + const nbRowBuf = new StorageBuffer(device, slotsPerBatch * 4, BUFFERUSAGE_COPY_DST); const outBuf = new StorageBuffer( device, - edgesPerBatch * 4, + slotsPerBatch * 4, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST ); - const outScratch = new Float32Array(edgesPerBatch); + const outScratch = new Float32Array(slotsPerBatch); const compute = new Compute(device, shader, 'compute-edge-cost'); - compute.setParameter('edgesI', edgesIBuf); - compute.setParameter('edgesJ', edgesJBuf); + compute.setParameter('nbRow', nbRowBuf); compute.setParameter('splat', splatBuf); compute.setParameter('costs', outBuf); this.execute = async ( - cache: EdgeCostCache, - edgeI: Uint32Array, - edgeJ: Uint32Array, + splatData: Float32Array, + numSplats: number, + nbRows: Uint32Array, outCosts: Float32Array ) => { - const n = cache.numSplats; - const e = edgeI.length; + const n = numSplats; + const s = nbRows.length; if (n > maxN) throw new Error(`GpuEdgeCost: N=${n} exceeds maxN=${maxN}`); - if (e > maxE) throw new Error(`GpuEdgeCost: E=${e} exceeds maxE=${maxE}`); - if (edgeJ.length !== e || outCosts.length !== e) { - throw new Error('GpuEdgeCost: edgeI / edgeJ / outCosts must have same length'); + if (outCosts.length !== s) { + throw new Error('GpuEdgeCost: nbRows / outCosts must have same length'); } + if (s % k !== 0) throw new Error(`GpuEdgeCost: slot count ${s} must be a multiple of k=${k}`); - splatBuf.write(0, cache.splatData, 0, n * SPLAT_STRIDE); + splatBuf.write(0, splatData, 0, n * SPLAT_STRIDE); - const numBatches = Math.ceil(e / edgesPerBatch); + const numBatches = Math.ceil(s / slotsPerBatch); for (let batch = 0; batch < numBatches; batch++) { - const edgeOffset = batch * edgesPerBatch; - const edgeCount = Math.min(edgesPerBatch, e - edgeOffset); - const groups = Math.ceil(edgeCount / workgroupSize); + const slotBase = batch * slotsPerBatch; + const slotCount = Math.min(slotsPerBatch, s - slotBase); + const groups = Math.ceil(slotCount / workgroupSize); - edgesIBuf.write(0, edgeI, edgeOffset, edgeCount); - edgesJBuf.write(0, edgeJ, edgeOffset, edgeCount); + nbRowBuf.write(0, nbRows, slotBase, slotCount); - compute.setParameter('edgeCount', edgeCount); + compute.setParameter('slotBase', slotBase); + compute.setParameter('slotCount', slotCount); compute.setupDispatch(groups); device.computeDispatch([compute], `edge-cost-dispatch-${batch}`); - const readBytes = edgeCount * 4; + const readBytes = slotCount * 4; await outBuf.read(0, readBytes, outScratch, true); - outCosts.set(outScratch.subarray(0, edgeCount), edgeOffset); + outCosts.set(outScratch.subarray(0, slotCount), slotBase); } }; this.destroy = () => { splatBuf.destroy(); - edgesIBuf.destroy(); - edgesJBuf.destroy(); + nbRowBuf.destroy(); outBuf.destroy(); shader.destroy(); bindGroupFormat.destroy(); @@ -326,4 +332,4 @@ class GpuEdgeCost { } } -export { GpuEdgeCost, type EdgeCostCache }; +export { GpuEdgeCost }; diff --git a/src/lib/gpu/index.ts b/src/lib/gpu/index.ts index a5241878..3877ff7a 100644 --- a/src/lib/gpu/index.ts +++ b/src/lib/gpu/index.ts @@ -1,6 +1,5 @@ export { GpuDilation } from './gpu-dilation'; export { GpuEdgeCost } from './gpu-edge-cost'; -export type { EdgeCostCache } from './gpu-edge-cost'; export { GpuKmeans } from './gpu-kmeans'; export { GpuKnn } from './gpu-knn'; export { GpuSplatRasterizer } from './gpu-splat-rasterizer'; diff --git a/test/decimate-edge-cost.test.mjs b/test/decimate-edge-cost.test.mjs index a18be559..f123b45a 100644 --- a/test/decimate-edge-cost.test.mjs +++ b/test/decimate-edge-cost.test.mjs @@ -8,7 +8,7 @@ import assert from 'node:assert'; import { describe, it } from 'node:test'; -import { buildCostCache, computeEdgeCostView } from '../src/lib/decimate/edge-cost-cpu.js'; +import { buildSplatCache, computeEdgeCost, CACHE_STRIDE } from '../src/lib/decimate/edge-cost-cpu.js'; // Build a minimal SplatView from per-splat specs (identity quaternion; only the // DC colour is set, higher SH left zero). colorDim = 3 (DC only). @@ -34,7 +34,11 @@ const makeView = (splats) => { return { pos, geo, color, colorDim }; }; -const cost = (view, i, j) => computeEdgeCostView(view, buildCostCache(view), i, j); +const cost = (view, i, j) => { + const cache = new Float32Array((view.geo.length / 8) * CACHE_STRIDE); + buildSplatCache(view, cache); + return computeEdgeCost(cache, i, j); +}; describe('field-L2 edge cost', () => { it('merging identical coincident splats is (near-)lossless', () => { diff --git a/test/decimate-priority.test.mjs b/test/decimate-priority.test.mjs index f61c8df7..40bd6edd 100644 --- a/test/decimate-priority.test.mjs +++ b/test/decimate-priority.test.mjs @@ -7,7 +7,7 @@ import assert from 'node:assert'; import { describe, it } from 'node:test'; import { makeSyntheticSource } from './helpers/synthetic-source.mjs'; -import { buildCostCache, computeEdgeCostView } from '../src/lib/decimate/edge-cost-cpu.js'; +import { buildSplatCache, computeEdgeCost, CACHE_STRIDE } from '../src/lib/decimate/edge-cost-cpu.js'; import { kdPartition } from '../src/lib/decimate/partition.js'; import { runPriorityPass } from '../src/lib/decimate/priority.js'; @@ -22,14 +22,15 @@ describe('priority pass (CPU)', () => { }; await runPriorityPass({ source, pool, pos, order, blocks, K, k }, cand); - const cache = buildCostCache(view); + const cache = new Float32Array(n * CACHE_STRIDE); + buildSplatCache(view, cache); const d2 = (a, b) => (pos.x[a] - pos.x[b]) ** 2 + (pos.y[a] - pos.y[b]) ** 2 + (pos.z[a] - pos.z[b]) ** 2; for (let i = 0; i < n; i += 97) { const knn = Array.from({ length: n }, (_, j) => j) .filter(j => j !== i) .sort((a, b) => d2(i, a) - d2(i, b)) .slice(0, k); - const refCosts = knn.map(j => computeEdgeCostView(view, cache, i, j)).sort((a, b) => a - b); + const refCosts = knn.map(j => computeEdgeCost(cache, i, j)).sort((a, b) => a - b); const got = []; for (let s = 0; s < K; s++) { if (cand.idx[i * K + s] !== 0xFFFFFFFF) got.push(cand.cost[i * K + s]); From fada6268ad5e97e7b03c0efef648dcb5cd7a04e8 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 15:47:00 +0100 Subject: [PATCH 06/19] latest --- src/lib/decimate/decimate-source.ts | 13 +- src/lib/decimate/recost-core.ts | 381 +++++++++++++++++++--------- src/lib/decimate/select-recost.ts | 205 +++------------ src/lib/workers/tasks.ts | 52 ---- test/decimate-recost.test.mjs | 269 ++++++++++++++++++++ 5 files changed, 580 insertions(+), 340 deletions(-) create mode 100644 test/decimate-recost.test.mjs diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index 8c200395..b7dcb759 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -233,19 +233,12 @@ const decimateSource = async ( const baseBytes = residentInputBytes + N * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; const recost = !legacy && baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; - // The splat cache and neighbour graph feed refresh rounds; shared - // memory lets those rounds run on the worker pool. - const sharedOk = recost && typeof SharedArrayBuffer !== 'undefined'; const cand: CandidateArrays = { idx: new Uint32Array(N * K).fill(0xFFFFFFFF), cost: new Float32Array(N * K).fill(Infinity) }; - const cacheOut = recost ? - new Float32Array(sharedOk ? new SharedArrayBuffer(N * CACHE_STRIDE * 4) : new ArrayBuffer(N * CACHE_STRIDE * 4)) : - undefined; - const neighborsOut = recost ? - new Uint32Array(sharedOk ? new SharedArrayBuffer(N * k * 4) : new ArrayBuffer(N * k * 4)) : - undefined; + const cacheOut = recost ? new Float32Array(N * CACHE_STRIDE) : undefined; + const neighborsOut = recost ? new Uint32Array(N * k) : undefined; const priorityBar = logger.bar('computing merge priorities', N); if (legacy) { @@ -269,7 +262,7 @@ const decimateSource = async ( const selection = legacy ? selectMergesLegacy(cand, N, K, needed) : cacheOut ? - await selectMergesRecosted({ cand, K, splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed }) : + selectMergesRecosted({ cand, K, splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed }) : selectMerges(cand, N, K, needed); selectSub.end(); diff --git a/src/lib/decimate/recost-core.ts b/src/lib/decimate/recost-core.ts index 0b69f6cb..8b49842c 100644 --- a/src/lib/decimate/recost-core.ts +++ b/src/lib/decimate/recost-core.ts @@ -1,14 +1,28 @@ /** - * Re-costed selection evaluation kernel, shared between the main thread and - * worker threads (bulk refresh rounds). Operates on a plain view-of-state - * object whose arrays may be backed by SharedArrayBuffers; workers treat the - * state as read-only (find() does no path compression here). + * Re-costed selection evaluation kernel — stateless over members. * * Cost definition matches the pairwise kernel: exact field-L2 of the merged * cluster vs its generation-input members (splat cache rows) plus the - * scale-free DC colour term — see select-recost.ts for the orchestration. + * scale-free DC colour term. Because groups are capped at MAX_GROUP original + * members, a cluster's moments (W, mean, M2, colour sums) are recomputed on + * the fly from its ≤ maxGroup immutable cache rows — no per-root float state + * is stored anywhere, and commits are pure integer structure updates. * - * Engine-free; no allocation in the hot paths beyond small local scratch. + * The marginal cost ΔE = E(A∪B) − E(A) − E(B) is evaluated in the cancelled + * form (the within-cluster pairwise sums Sself drop out identically): + * + * ΔE = 2·scross(A,B) − 2·memfm(A∪B) + selfM(A∪B) + * + (|A|≥2 ? 2·memfm(A) − selfM(A) : self_A) + * + (|B|≥2 ? 2·memfm(B) − selfM(B) : self_B) + * + * where memfm(C) = Σ_{k∈C} ⟨f_k, f_C⟩, selfM(C) = ⟨f_C, f_C⟩, scross(A,B) = + * Σ_{a∈A,b∈B} ⟨f_a, f_b⟩, and self_x is a singleton's self product (its + * E is exactly 0 by convention — the merged Gaussian of one splat is itself). + * Beyond deleting all per-root aggregates, the cancelled form is also + * better-conditioned: the mutually-cancelling Sself magnitudes never enter. + * + * Engine-free; single-threaded (module-scope scratch, no allocation in the + * hot paths). */ import { CACHE_STRIDE, COLOR_WEIGHT } from './edge-cost-cpu'; @@ -26,8 +40,9 @@ const PI_1_5 = Math.PI ** 1.5; const TWO_PI_1_5 = (2 * Math.PI) ** 1.5; /** - * The resident selection state (allocated by select-recost, possibly on - * SharedArrayBuffers so bulk refreshes can run on worker threads). + * The resident selection state (allocated by select-recost). Only integer + * structure — the union-find, the member chains, and the immutable per-splat + * cache/neighbour inputs from the priority pass. */ type RecostState = { /** Per-splat cache, {@link CACHE_STRIDE} floats per generation-input splat. */ @@ -40,22 +55,16 @@ type RecostState = { N: number; /** Max original members per group. */ maxGroup: number; - // Union-find + cluster state (indexed by root). + // Union-find + cluster structure (indexed by root). parent: Uint32Array; size: Uint32Array; - W: Float64Array; - mx: Float64Array; my: Float64Array; mz: Float64Array; - M2: Float64Array; - baseW: Float64Array; - Sself: Float64Array; - Err: Float64Array; version: Uint32Array; mHead: Uint32Array; mNext: Uint32Array; }; /** - * Read-only find (no path compression — safe on shared state in workers). + * Read-only find (no path compression). * * @param parent - Union-find parent array. * @param x - Element to resolve. @@ -110,171 +119,315 @@ const eig3 = (m0: number, m1: number, m2: number, m3: number, m4: number, m5: nu eigOut[0] = e0; eigOut[1] = 3 * q - e0 - e2; eigOut[2] = e2; }; -let gatherBuf = new Uint32Array(1 << 12); -const gatherMembers = (st: RecostState, root: number): number => { - const { mHead, mNext } = st; - let cnt = 0; - for (let m = mHead[root]; m !== NIL; m = mNext[m]) { - if (cnt === gatherBuf.length) { - const g = new Uint32Array(gatherBuf.length * 2); - g.set(gatherBuf); gatherBuf = g; - } - gatherBuf[cnt++] = m; - } - return cnt; +/** + * A cluster's moments composed from its member cache rows: raw mass-scaled + * aggregates (the parallel-axis form the old per-root state stored) plus the + * finished merged-Gaussian quantities the field terms need. + */ +type Comp = { + // Raw aggregates. + W: number; + mx: number; my: number; mz: number; // mass-weighted mean + M2: Float64Array; // Σ mass·(Σₖ + δδᵀ), 6 comps + bw0: number; bw1: number; bw2: number; // Σ mass·base + // Finished merged Gaussian (valid after finishComp). + sm: Float64Array; // M2/W + EPS_COV·I, 6 comps + sd: number; // √|Σ_C| + alpha: number; // min(1, W / area) + bc0: number; bc1: number; bc2: number; // mean base colour (bw/W) + bn2: number; // |mean base|² + selfM: number; // ⟨f_C, f_C⟩ }; -/** Outputs of {@link evalMergeCore} beyond the cost (reused object). */ -const evalOut = { E: 0, Scross: 0 }; +const makeComp = (): Comp => ({ + W: 0, + mx: 0, + my: 0, + mz: 0, + M2: new Float64Array(6), + bw0: 0, + bw1: 0, + bw2: 0, + sm: new Float64Array(6), + sd: 0, + alpha: 0, + bc0: 0, + bc1: 0, + bc2: 0, + bn2: 0, + selfM: 0 +}); -/** - * Marginal cost ΔE of merging clusters A and B (exact field-L2 vs the - * generation-input members) plus the scale-free colour term. Fills - * {@link evalOut} with E(A∪B) and Scross(A,B). - * - * @param st - Selection state. - * @param A - First cluster root. - * @param B - Second cluster root. - * @returns The merge cost. - */ -const evalMergeCore = (st: RecostState, A: number, B: number): number => { - const { SC, W, mx, my, mz, M2, baseW, Sself, Err, mHead, mNext } = st; - const WA = W[A], WB = W[B], WC = WA + WB; +// Module-scope scratch — single-threaded by design. +const compA = makeComp(); +const compB = makeComp(); +const compAB = makeComp(); + +// Compose raw aggregates from member cache rows (two passes: mean, then +// central second moments — algebraically identical to the parallel-axis +// accumulation the formation tree would produce). +const composeRaw = (SC: Float32Array, members: Uint32Array, count: number, out: Comp): void => { + let W = 0, sx = 0, sy = 0, sz = 0, b0 = 0, b1 = 0, b2 = 0; + for (let t = 0; t < count; t++) { + const o = members[t] * CACHE_STRIDE; + const m = SC[o + 11]; + W += m; + sx += m * SC[o]; sy += m * SC[o + 1]; sz += m * SC[o + 2]; + b0 += m * SC[o + 12]; b1 += m * SC[o + 13]; b2 += m * SC[o + 14]; + } + const iw = 1 / W; + const mx = sx * iw, my = sy * iw, mz = sz * iw; + const M2 = out.M2; + M2.fill(0); + for (let t = 0; t < count; t++) { + const o = members[t] * CACHE_STRIDE; + const m = SC[o + 11]; + const dx = SC[o] - mx, dy = SC[o + 1] - my, dz = SC[o + 2] - mz; + M2[0] += m * (SC[o + 3] + dx * dx); + M2[1] += m * (SC[o + 4] + dx * dy); + M2[2] += m * (SC[o + 5] + dx * dz); + M2[3] += m * (SC[o + 6] + dy * dy); + M2[4] += m * (SC[o + 7] + dy * dz); + M2[5] += m * (SC[o + 8] + dz * dz); + } + out.W = W; out.mx = mx; out.my = my; out.mz = mz; + out.bw0 = b0; out.bw1 = b1; out.bw2 = b2; +}; + +// Compose A∪B's raw aggregates from A's and B's (parallel-axis identity — +// the exact arithmetic the merged-covariance step has always used). +const composeUnion = (a: Comp, b: Comp, out: Comp): void => { + const WA = a.W, WB = b.W, WC = WA + WB; const iw = 1 / WC; - const mcx = (WA * mx[A] + WB * mx[B]) * iw; - const mcy = (WA * my[A] + WB * my[B]) * iw; - const mcz = (WA * mz[A] + WB * mz[B]) * iw; - const dax = mx[A] - mcx, day = my[A] - mcy, daz = mz[A] - mcz; - const dbx = mx[B] - mcx, dby = my[B] - mcy, dbz = mz[B] - mcz; - const a6 = A * 6, b6 = B * 6; - - const sm0 = (M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx) * iw + EPS_COV; - const sm1 = (M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby) * iw; - const sm2 = (M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz) * iw; - const sm3 = (M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby) * iw + EPS_COV; - const sm4 = (M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz) * iw; - const sm5 = (M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz) * iw + EPS_COV; + const mcx = (WA * a.mx + WB * b.mx) * iw; + const mcy = (WA * a.my + WB * b.my) * iw; + const mcz = (WA * a.mz + WB * b.mz) * iw; + const dax = a.mx - mcx, day = a.my - mcy, daz = a.mz - mcz; + const dbx = b.mx - mcx, dby = b.my - mcy, dbz = b.mz - mcz; + const MA = a.M2, MB = b.M2, MC = out.M2; + MC[0] = MA[0] + MB[0] + WA * dax * dax + WB * dbx * dbx; + MC[1] = MA[1] + MB[1] + WA * dax * day + WB * dbx * dby; + MC[2] = MA[2] + MB[2] + WA * dax * daz + WB * dbx * dbz; + MC[3] = MA[3] + MB[3] + WA * day * day + WB * dby * dby; + MC[4] = MA[4] + MB[4] + WA * day * daz + WB * dby * dbz; + MC[5] = MA[5] + MB[5] + WA * daz * daz + WB * dbz * dbz; + out.W = WC; out.mx = mcx; out.my = mcy; out.mz = mcz; + out.bw0 = a.bw0 + b.bw0; out.bw1 = a.bw1 + b.bw1; out.bw2 = a.bw2 + b.bw2; +}; + +// Derive the merged Gaussian's field quantities from the raw aggregates. +const finishComp = (c: Comp): void => { + const iw = 1 / c.W; + const sm = c.sm; + sm[0] = c.M2[0] * iw + EPS_COV; + sm[1] = c.M2[1] * iw; + sm[2] = c.M2[2] * iw; + sm[3] = c.M2[3] * iw + EPS_COV; + sm[4] = c.M2[4] * iw; + sm[5] = c.M2[5] * iw + EPS_COV; const detm = Math.max( - sm0 * (sm3 * sm5 - sm4 * sm4) - sm1 * (sm1 * sm5 - sm4 * sm2) + sm2 * (sm1 * sm4 - sm3 * sm2), + sm[0] * (sm[3] * sm[5] - sm[4] * sm[4]) - sm[1] * (sm[1] * sm[5] - sm[4] * sm[2]) + sm[2] * (sm[1] * sm[4] - sm[3] * sm[2]), 1e-60 ); - const sdC = Math.sqrt(detm); + c.sd = Math.sqrt(detm); - eig3(sm0, sm1, sm2, sm3, sm4, sm5); + eig3(sm[0], sm[1], sm[2], sm[3], sm[4], sm[5]); const s0 = Math.sqrt(Math.max(eigOut[0], 1e-18)); const s1 = Math.sqrt(Math.max(eigOut[1], 1e-18)); const s2 = Math.sqrt(Math.max(eigOut[2], 1e-18)); - const alphaC = Math.min(1, WC / Math.max(ellipsoidArea(s0, s1, s2), 1e-30)); - - const a3 = A * 3, b3 = B * 3; - const bc0 = (baseW[a3] + baseW[b3]) * iw; - const bc1 = (baseW[a3 + 1] + baseW[b3 + 1]) * iw; - const bc2 = (baseW[a3 + 2] + baseW[b3 + 2]) * iw; - const bn2C = bc0 * bc0 + bc1 * bc1 + bc2 * bc2; - - const selfM = alphaC * alphaC * bn2C * PI_1_5 * sdC; - - // ⟨Σ member fields, f_m⟩ over both chains. - let memfm = 0; - for (let pass = 0; pass < 2; pass++) { - for (let m = pass === 0 ? mHead[A] : mHead[B]; m !== NIL; m = mNext[m]) { - const o = m * CACHE_STRIDE; - const wgt = SC[o + 10] * alphaC * - (SC[o + 12] * bc0 + SC[o + 13] * bc1 + SC[o + 14] * bc2); - if (wgt === 0) continue; - memfm += wgt * crossG(SC[o + 9] * sdC, - SC[o + 3] + sm0, SC[o + 4] + sm1, SC[o + 5] + sm2, - SC[o + 6] + sm3, SC[o + 7] + sm4, SC[o + 8] + sm5, - SC[o] - mcx, SC[o + 1] - mcy, SC[o + 2] - mcz); + c.alpha = Math.min(1, c.W / Math.max(ellipsoidArea(s0, s1, s2), 1e-30)); + + c.bc0 = c.bw0 * iw; c.bc1 = c.bw1 * iw; c.bc2 = c.bw2 * iw; + c.bn2 = c.bc0 * c.bc0 + c.bc1 * c.bc1 + c.bc2 * c.bc2; + c.selfM = c.alpha * c.alpha * c.bn2 * PI_1_5 * c.sd; +}; + +// memfm(C over `members`) = Σ ⟨f_k, f_C⟩ for the finished cluster `c`. +const memfm = (SC: Float32Array, members: Uint32Array, count: number, c: Comp): number => { + const sm = c.sm; + let acc = 0; + for (let t = 0; t < count; t++) { + const o = members[t] * CACHE_STRIDE; + const wgt = SC[o + 10] * c.alpha * + (SC[o + 12] * c.bc0 + SC[o + 13] * c.bc1 + SC[o + 14] * c.bc2); + if (wgt === 0) continue; + acc += wgt * crossG(SC[o + 9] * c.sd, + SC[o + 3] + sm[0], SC[o + 4] + sm[1], SC[o + 5] + sm[2], + SC[o + 6] + sm[3], SC[o + 7] + sm[4], SC[o + 8] + sm[5], + SC[o] - c.mx, SC[o + 1] - c.my, SC[o + 2] - c.mz); + } + return acc; +}; + +// A singleton's self product ⟨f, f⟩ (its cluster error is exactly 0). +const selfRow = (SC: Float32Array, row: number): number => { + const o = row * CACHE_STRIDE; + return SC[o + 10] * SC[o + 10] * SC[o + 15] * PI_1_5 * SC[o + 9]; +}; + +// Member buffers (grow-only; production maxGroup is 4). +let aBuf: Uint32Array = new Uint32Array(1 << 12); +let bBuf: Uint32Array = new Uint32Array(1 << 12); +const gatherChain = (st: RecostState, root: number, into: Uint32Array): { buf: Uint32Array, count: number } => { + const { mHead, mNext } = st; + let buf = into; + let cnt = 0; + for (let m = mHead[root]; m !== NIL; m = mNext[m]) { + if (cnt === buf.length) { + const g = new Uint32Array(buf.length * 2); + g.set(buf); buf = g; } + buf[cnt++] = m; } + return { buf, count: cnt }; +}; - // Scross(A,B) = Σ_{a∈A,b∈B}⟨f_a,f_b⟩ with distance culling. - const na = gatherMembers(st, A); - let scross = 0; - for (let b = mHead[B]; b !== NIL; b = mNext[b]) { - const ob = b * CACHE_STRIDE; +// One side's term of the cancelled cost form, given the side's raw +// aggregates (finishes `c` when the side is a real cluster). +const sideTerm = (SC: Float32Array, members: Uint32Array, count: number, c: Comp): number => { + if (count < 2) return selfRow(SC, members[0]); + finishComp(c); + return 2 * memfm(SC, members, count, c) - c.selfM; +}; + +// scross(A,B) = Σ_{a∈A,b∈B} ⟨f_a, f_b⟩ with distance culling. +const scrossPairs = ( + SC: Float32Array, + aMembers: Uint32Array, na: number, + bMembers: Uint32Array, nb: number +): number => { + let acc = 0; + for (let u = 0; u < nb; u++) { + const ob = bMembers[u] * CACHE_STRIDE; const bxp = SC[ob], byp = SC[ob + 1], bzp = SC[ob + 2]; const trb = SC[ob + 3] + SC[ob + 6] + SC[ob + 8]; const alb = SC[ob + 10], sdb = SC[ob + 9]; const cb0 = SC[ob + 12], cb1 = SC[ob + 13], cb2 = SC[ob + 14]; for (let t = 0; t < na; t++) { - const a = gatherBuf[t]; - const oa = a * CACHE_STRIDE; + const oa = aMembers[t] * CACHE_STRIDE; const dx = SC[oa] - bxp, dy = SC[oa + 1] - byp, dz = SC[oa + 2] - bzp; const d2 = dx * dx + dy * dy + dz * dz; if (d2 > CULL_QUAD * (SC[oa + 3] + SC[oa + 6] + SC[oa + 8] + trb)) continue; const wgt = SC[oa + 10] * alb * (SC[oa + 12] * cb0 + SC[oa + 13] * cb1 + SC[oa + 14] * cb2); if (wgt === 0) continue; - scross += wgt * crossG(SC[oa + 9] * sdb, + acc += wgt * crossG(SC[oa + 9] * sdb, SC[oa + 3] + SC[ob + 3], SC[oa + 4] + SC[ob + 4], SC[oa + 5] + SC[ob + 5], SC[oa + 6] + SC[ob + 6], SC[oa + 7] + SC[ob + 7], SC[oa + 8] + SC[ob + 8], dx, dy, dz); } } + return acc; +}; - const E = Sself[A] + Sself[B] + 2 * scross - 2 * memfm + selfM; - evalOut.E = E; - evalOut.Scross = scross; +// Evaluate ΔE for A (context already composed in compA/aBuf, side term +// `aTerm`) against cluster B, in the cancelled form. +const evalAgainst = (st: RecostState, na: number, aTerm: number, B: number): number => { + const { SC } = st; + const b = gatherChain(st, B, bBuf); + bBuf = b.buf; + const nb = b.count; + composeRaw(SC, bBuf, nb, compB); + const bTerm = sideTerm(SC, bBuf, nb, compB); - const iwA = 1 / W[A], iwB = 1 / W[B]; - const d0 = baseW[a3] * iwA - baseW[b3] * iwB; - const d1 = baseW[a3 + 1] * iwA - baseW[b3 + 1] * iwB; - const d2c = baseW[a3 + 2] * iwA - baseW[b3 + 2] * iwB; - return (E - Err[A] - Err[B]) + COLOR_WEIGHT * (d0 * d0 + d1 * d1 + d2c * d2c); + composeUnion(compA, compB, compAB); + finishComp(compAB); + + const memfmAB = memfm(SC, aBuf, na, compAB) + memfm(SC, bBuf, nb, compAB); + const scross = scrossPairs(SC, aBuf, na, bBuf, nb); + + // Scale-free DC colour term between the clusters' mean base colours. + const iwA = 1 / compA.W, iwB = 1 / compB.W; + const d0 = compA.bw0 * iwA - compB.bw0 * iwB; + const d1 = compA.bw1 * iwA - compB.bw1 * iwB; + const d2c = compA.bw2 * iwA - compB.bw2 * iwB; + + return (2 * scross - 2 * memfmAB + compAB.selfM + aTerm + bTerm) + + COLOR_WEIGHT * (d0 * d0 + d1 * d1 + d2c * d2c); +}; + +/** + * Marginal cost ΔE of merging clusters A and B (exact field-L2 vs the + * generation-input members) plus the scale-free colour term — recomputed + * statelessly from the members' cache rows. + * + * @param st - Selection state. + * @param A - First cluster root. + * @param B - Second cluster root. + * @returns The merge cost. + */ +const evalMergeCore = (st: RecostState, A: number, B: number): number => { + const a = gatherChain(st, A, aBuf); + aBuf = a.buf; + const na = a.count; + composeRaw(st.SC, aBuf, na, compA); + const aTerm = sideTerm(st.SC, aBuf, na, compA); + return evalAgainst(st, na, aTerm, B); }; /** Result of {@link bestEdgeFor} (reused object). */ -const bestOut = { partner: -1, vb: 0, cost: 0, E: 0, S: 0 }; +const bestOut = { partner: -1, vb: 0, cost: 0 }; + +const candScratch = new Uint32Array(256); /** * Compute the cheapest legal merge for `root`: candidates are the live * clusters owning any member's candidate ids, deduplicated linearly (pools - * are small), respecting the group cap. + * are small), respecting the group cap. The root's composition and side + * term are hoisted — computed once, reused for every candidate. * * @param st - Selection state. * @param root - Cluster root to refresh. * @returns True when a legal candidate exists (result in {@link bestOut}). */ -const candScratch = new Uint32Array(256); - const bestEdgeFor = (st: RecostState, root: number): boolean => { - const { cands, D, parent, size, version, mHead, mNext, maxGroup } = st; + const { SC, cands, D, parent, size, version, maxGroup } = st; + if (maxGroup * D > candScratch.length) { + throw new Error(`recost: candidate pool bound ${maxGroup * D} exceeds scratch (${candScratch.length})`); + } const sz = size[root]; - // Derive + dedup candidates (linear scan — pools are ≤ a few dozen). + + const a = gatherChain(st, root, aBuf); + aBuf = a.buf; + const na = a.count; + + // Derive + dedup candidates (linear scan — pools are ≤ maxGroup·D). let cnt = 0; const cbuf = candScratch; - for (let m = mHead[root]; m !== NIL; m = mNext[m]) { - const base = m * D; + for (let t = 0; t < na; t++) { + const base = aBuf[t] * D; for (let s = 0; s < D; s++) { const nb = cands[base + s]; if (nb === NO_CANDIDATE) continue; const r = findRO(parent, nb); if (r === root) continue; let seen = false; - for (let t = 0; t < cnt; t++) { - if (cbuf[t] === r) { + for (let u = 0; u < cnt; u++) { + if (cbuf[u] === r) { seen = true; break; } } - if (seen) continue; - if (cnt < cbuf.length) cbuf[cnt++] = r; + if (!seen) cbuf[cnt++] = r; } } - let bc = Infinity, bp = -1, bv = 0, bE = 0, bS = 0; + if (cnt === 0) return false; + + composeRaw(SC, aBuf, na, compA); + const aTerm = sideTerm(SC, aBuf, na, compA); + + let bc = Infinity, bp = -1, bv = 0; for (let t = 0; t < cnt; t++) { const c = cbuf[t]; if (sz + size[c] > maxGroup) continue; - const d = evalMergeCore(st, root, c); + const d = evalAgainst(st, na, aTerm, c); if (d < bc) { - bc = d; bp = c; bv = version[c]; bE = evalOut.E; bS = evalOut.Scross; + bc = d; bp = c; bv = version[c]; } } if (bp < 0) return false; - bestOut.partner = bp; bestOut.vb = bv; bestOut.cost = bc; bestOut.E = bE; bestOut.S = bS; + bestOut.partner = bp; bestOut.vb = bv; bestOut.cost = bc; return true; }; -export { evalMergeCore, bestEdgeFor, findRO, evalOut, bestOut, NO_CANDIDATE, type RecostState }; +export { evalMergeCore, bestEdgeFor, bestOut, NO_CANDIDATE, type RecostState }; diff --git a/src/lib/decimate/select-recost.ts b/src/lib/decimate/select-recost.ts index 50c3cc01..ef1e98d8 100644 --- a/src/lib/decimate/select-recost.ts +++ b/src/lib/decimate/select-recost.ts @@ -1,18 +1,15 @@ /** * Re-costed merge selection: exact greedy agglomeration within a generation, - * parallelized as commit/refresh rounds. + * batched into commit/refresh rounds. * * Where {@link selectMerges} consumes the priority pass's pairwise costs * one-shot (costs go stale as groups form), this selection re-evaluates a * cluster's candidates after it changes, so merges always execute at costs - * evaluated against current state. To make the (dominant) evaluation work - * parallel, refreshes are batched into rounds: + * evaluated against current state: * * 1. Drain the heap, committing every entry that is still valid; clusters * that changed (merged) or whose best partner changed are queued. - * 2. Bulk-evaluate all queued clusters' best edges — on the worker pool - * when the state is SharedArrayBuffer-backed, inline otherwise — and - * push the results. + * 2. Re-evaluate all queued clusters' best edges and push the results. * 3. Repeat until the generation target is reached or no merges remain. * * Relative to strictly-eager greedy, refreshed clusters re-enter the heap at @@ -21,20 +18,23 @@ * deviation is confined to near-tie ordering (re-certified on the evaluation * harness). * - * Cost definition and evaluation kernel live in {@link recost-core} (shared - * with the worker task). Seeds come from the priority pass's candidate - * arrays; refresh candidate pools come from the persisted neighbour graph. + * The evaluation kernel ({@link recost-core}) is stateless over members: + * cluster moments are recomputed from the ≤ MAX_GROUP member cache rows on + * every evaluation, so this module keeps no per-root float state at all — + * a commit is a chain splice plus union-find/size/version updates, and the + * heap carries only (cost, a, b, seq, versionB). Seeds come from the + * priority pass's candidate arrays; refresh candidate pools come from the + * persisted neighbour graph. * - * decimate-source gates this path by memory budget (~330 B per gaussian - * resident) and falls back to {@link selectMerges} when over. + * decimate-source gates this path by memory budget and falls back to + * {@link selectMerges} when over. * - * Engine-free; pure resident-array computation, no IO. + * Engine-free; single-threaded, pure resident-array computation, no IO. */ import { type CandidateArrays } from './priority'; -import { evalMergeCore, bestEdgeFor, evalOut, bestOut, NO_CANDIDATE, CACHE_STRIDE, type RecostState } from './recost-core'; +import { bestEdgeFor, bestOut, NO_CANDIDATE, CACHE_STRIDE, type RecostState } from './recost-core'; import { MAX_GROUP, type SelectionResult } from './select'; -import { WorkerQueue } from '../workers'; const NIL = 0xFFFFFFFF; @@ -47,7 +47,6 @@ type RecostInputs = { /** * Resident per-splat cache, {@link CACHE_STRIDE} floats per splat * (buildSplatCache row layout), persisted by the priority pass. - * SharedArrayBuffer-backed for the parallel path. */ splatCache: Float32Array; /** @@ -55,7 +54,6 @@ type RecostInputs = { * of a cluster's members' rows forms its refresh candidate pool (the full * neighbour graph, not just the seed candidates: top-K lists collapse onto * shared hubs in degenerate/coincident regions and would starve refreshes). - * SharedArrayBuffer-backed for the parallel path. */ neighbors: Uint32Array; /** Neighbours per splat. */ @@ -67,48 +65,20 @@ type RecostInputs = { }; /** - * Allocate a typed array, on shared memory when requested. - * - * @param ctor - Typed-array constructor. - * @param bytes - Buffer size in bytes. - * @param shared - Back with a SharedArrayBuffer (worker-visible state). - * @returns The array view. - */ -const alloc = ( - ctor: new (buffer: ArrayBufferLike) => T, bytes: number, shared: boolean -): T => { - return new ctor(shared ? new SharedArrayBuffer(bytes) : new ArrayBuffer(bytes)); -}; - -/** - * Select merges with within-generation re-costing (round-parallel). + * Select merges with within-generation re-costing (round-batched). * * @param inputs - See {@link RecostInputs}. * @returns The selection (same contract as {@link selectMerges}). */ -const selectMergesRecosted = async (inputs: RecostInputs): Promise => { +const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { const { cand, K, splatCache: SC, neighbors, D, N, mergesNeeded } = inputs; - const shared = typeof SharedArrayBuffer !== 'undefined' && - SC.buffer instanceof SharedArrayBuffer && - neighbors.buffer instanceof SharedArrayBuffer && - !WorkerQueue.isInline; - - // ---- Cluster state (indexed by union-find root). - const parent = alloc(Uint32Array, N * 4, shared); - const size = alloc(Uint32Array, N * 4, shared); - const W = alloc(Float64Array, N * 8, shared); - const mx = alloc(Float64Array, N * 8, shared); - const my = alloc(Float64Array, N * 8, shared); - const mz = alloc(Float64Array, N * 8, shared); - const M2 = alloc(Float64Array, N * 6 * 8, shared); - const baseW = alloc(Float64Array, N * 3 * 8, shared); - const Sself = alloc(Float64Array, N * 8, shared); - const Err = alloc(Float64Array, N * 8, shared); - const version = alloc(Uint32Array, N * 4, shared); - const mHead = alloc(Uint32Array, N * 4, shared); - const mNext = alloc(Uint32Array, N * 4, shared); - // Main-thread-only state. + // ---- Cluster structure (indexed by union-find root) — integers only. + const parent = new Uint32Array(N); + const size = new Uint32Array(N); + const version = new Uint32Array(N); + const mHead = new Uint32Array(N); + const mNext = new Uint32Array(N); const mTail = new Uint32Array(N); const lastSeq = new Uint32Array(N); @@ -120,14 +90,6 @@ const selectMergesRecosted = async (inputs: RecostInputs): Promise { @@ -165,13 +110,13 @@ const selectMergesRecosted = async (inputs: RecostInputs): Promise { + const heapPush = (cost: number, a: number, b: number, seq: number, vb: number): void => { if (heapSize === heapCap) { const nc = heapCap * 2; - const gf = (old: Float64Array) => { - const x = new Float64Array(nc); x.set(old); return x; + const gf = (old: Float32Array) => { + const x = new Float32Array(nc); x.set(old); return x; }; const g = (old: Uint32Array) => { const x = new Uint32Array(nc); x.set(old); return x; }; - hCost = gf(hCost); hE = gf(hE); hS = gf(hS); + hCost = gf(hCost); hA = g(hA); hB = g(hB); hSeq = g(hSeq); hVb = g(hVb); heapCap = nc; } let i = heapSize++; - hCost[i] = cost; hA[i] = a; hB[i] = b; hSeq[i] = seq; hVb[i] = vb; hE[i] = e; hS[i] = s; + hCost[i] = cost; hA[i] = a; hB[i] = b; hSeq[i] = seq; hVb[i] = vb; while (i > 0) { const p = (i - 1) >> 1; if (hCost[p] <= hCost[i]) break; @@ -208,16 +151,14 @@ const selectMergesRecosted = async (inputs: RecostInputs): Promise { if (heapSize === 0) return false; popOut.cost = hCost[0]; popOut.a = hA[0]; popOut.b = hB[0]; popOut.seq = hSeq[0]; popOut.vb = hVb[0]; - popOut.E = hE[0]; popOut.S = hS[0]; heapSize--; if (heapSize > 0) { hCost[0] = hCost[heapSize]; hA[0] = hA[heapSize]; hB[0] = hB[heapSize]; hSeq[0] = hSeq[heapSize]; hVb[0] = hVb[heapSize]; - hE[0] = hE[heapSize]; hS[0] = hS[heapSize]; let i = 0; for (;;) { const l = 2 * i + 1, r = l + 1; @@ -255,7 +196,7 @@ const selectMergesRecosted = async (inputs: RecostInputs): Promise= size[b] ? a : b; const lose = keep === a ? b : a; - - const WA = W[a], WB = W[b], WC = WA + WB; - const iw = 1 / WC; - const mcx = (WA * mx[a] + WB * mx[b]) * iw; - const mcy = (WA * my[a] + WB * my[b]) * iw; - const mcz = (WA * mz[a] + WB * mz[b]) * iw; - const dax = mx[a] - mcx, day = my[a] - mcy, daz = mz[a] - mcz; - const dbx = mx[b] - mcx, dby = my[b] - mcy, dbz = mz[b] - mcz; - const a6 = a * 6, b6 = b * 6, k6 = keep * 6; - const n0 = M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx; - const n1 = M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby; - const n2 = M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz; - const n3 = M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby; - const n4 = M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz; - const n5 = M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz; - M2[k6] = n0; M2[k6 + 1] = n1; M2[k6 + 2] = n2; M2[k6 + 3] = n3; M2[k6 + 4] = n4; M2[k6 + 5] = n5; - W[keep] = WC; mx[keep] = mcx; my[keep] = mcy; mz[keep] = mcz; - const ka = keep * 3, aa = a * 3, bb = b * 3; - const bw0 = baseW[aa] + baseW[bb], bw1 = baseW[aa + 1] + baseW[bb + 1], bw2 = baseW[aa + 2] + baseW[bb + 2]; - baseW[ka] = bw0; baseW[ka + 1] = bw1; baseW[ka + 2] = bw2; - Sself[keep] = Sself[a] + Sself[b] + 2 * scross; - Err[keep] = E; - size[keep] += size[lose]; mNext[mTail[keep]] = mHead[lose]; mTail[keep] = mTail[lose]; parent[lose] = keep; + size[keep] += size[lose]; version[keep]++; removed++; wave++; @@ -320,50 +236,11 @@ const selectMergesRecosted = async (inputs: RecostInputs): Promise= mergesNeeded) break; if (pendingCount === 0 && heapSize === 0) break; - // Bulk refresh of queued clusters (parallel when shared). - if (shared) { - const workers = Math.max(1, Math.min(8, WorkerQueue.maxWorkers ?? 4)); - const chunk = Math.ceil(pendingCount / workers); - const jobs: Promise[] = []; - for (let off = 0; off < pendingCount; off += chunk) { - const roots = pending.slice(off, Math.min(off + chunk, pendingCount)); - jobs.push(WorkerQueue.run('recostBestEdges', { - sc: SC, - cands: neighbors, - d: D, - n: N, - maxGroup: MAX_GROUP, - parent, - size, - w: W, - mx, - my, - mz, - m2: M2, - baseW, - sself: Sself, - err: Err, - version, - mHead, - mNext, - roots - }, [roots.buffer as ArrayBuffer])); - } - const results = await Promise.all(jobs); - for (const res of results) { - for (let i = 0; i < res.length; i += 6) { - const root = res[i]; - const partner = res[i + 1]; - if (partner < 0) continue; - heapPush(res[i + 3], root, partner, lastSeq[root], res[i + 2], res[i + 4], res[i + 5]); - } - } - } else { - for (let p = 0; p < pendingCount; p++) { - const root = pending[p]; - if (bestEdgeFor(st, root)) { - heapPush(bestOut.cost, root, bestOut.partner, lastSeq[root], bestOut.vb, bestOut.E, bestOut.S); - } + // Refresh queued clusters' best edges. + for (let p = 0; p < pendingCount; p++) { + const root = pending[p]; + if (bestEdgeFor(st, root)) { + heapPush(bestOut.cost, root, bestOut.partner, lastSeq[root], bestOut.vb); } } pendingCount = 0; diff --git a/src/lib/workers/tasks.ts b/src/lib/workers/tasks.ts index 903d53ec..f3b35f9e 100644 --- a/src/lib/workers/tasks.ts +++ b/src/lib/workers/tasks.ts @@ -1,7 +1,6 @@ import type { TypedArray } from '../data-table/data-table'; import { knnQueryBlock } from '../decimate/knn-core'; import { mergeGroup, createMergeScratch, splatMass } from '../decimate/moment-match'; -import { bestEdgeFor, bestOut, type RecostState } from '../decimate/recost-core'; import { buildFlatKdTree, type FlatKdTree } from '../spatial/kd-tree'; import { quantize1dColumns, type QuantizedColumns } from '../spatial/quantize-1d-core'; import { WebPCodec } from '../utils/webp-codec'; @@ -68,57 +67,6 @@ const taskHandlers = { return { result, transfer: [result.buffer as ArrayBuffer] }; }, - // Re-costed selection bulk refresh: best merge edge for each queued - // cluster root, evaluated against SharedArrayBuffer-backed state that the - // host freezes for the duration of the round (read-only here). Packed - // result: 6 f64 per root — [root, partner|-1, partnerVersion, cost, E, Scross]. - recostBestEdges: (args: { - sc: Float32Array, cands: Uint32Array, d: number, n: number, maxGroup: number, - parent: Uint32Array, size: Uint32Array, - w: Float64Array, mx: Float64Array, my: Float64Array, mz: Float64Array, - m2: Float64Array, baseW: Float64Array, sself: Float64Array, err: Float64Array, - version: Uint32Array, mHead: Uint32Array, mNext: Uint32Array, - roots: Uint32Array - }): TaskOutput => { - const st: RecostState = { - SC: args.sc, - cands: args.cands, - D: args.d, - N: args.n, - maxGroup: args.maxGroup, - parent: args.parent, - size: args.size, - W: args.w, - mx: args.mx, - my: args.my, - mz: args.mz, - M2: args.m2, - baseW: args.baseW, - Sself: args.sself, - Err: args.err, - version: args.version, - mHead: args.mHead, - mNext: args.mNext - }; - const roots = args.roots; - const out = new Float64Array(roots.length * 6); - for (let i = 0; i < roots.length; i++) { - const root = roots[i]; - const o = i * 6; - out[o] = root; - if (bestEdgeFor(st, root)) { - out[o + 1] = bestOut.partner; - out[o + 2] = bestOut.vb; - out[o + 3] = bestOut.cost; - out[o + 4] = bestOut.E; - out[o + 5] = bestOut.S; - } else { - out[o + 1] = -1; - } - } - return { result: out, transfer: [out.buffer as ArrayBuffer] }; - }, - // Decimation merge stream: n-ary moment match of packed member-major // groups. Inputs are member-major (pos 3 / geo 8 / color colorDim floats // per member, groups back to back per `sizes`); outputs are group-major. diff --git a/test/decimate-recost.test.mjs b/test/decimate-recost.test.mjs new file mode 100644 index 00000000..01d6580e --- /dev/null +++ b/test/decimate-recost.test.mjs @@ -0,0 +1,269 @@ +/** + * Stateless re-costed selection kernel: + * + * 1. Singleton-pair parity — the stateless cancelled-form eval must agree + * with the pairwise kernel (they are the same formula, differently + * associated). + * 2. Marginal-cost decomposition — accumulated ΔE (minus the ordering-only + * colour terms) telescopes to E(final cluster), so it must be independent + * of the merge path taken. + * 3. selectMergesRecosted behaviour — target respected, CSR well-formed, + * group cap enforced, and coincident inputs concentrate into full groups + * (continuation merges via refresh — the concentration behaviour the + * quality study validated). + */ + +import assert from 'node:assert'; +import { describe, it } from 'node:test'; + +import { buildSplatCache, computeEdgeCost, CACHE_STRIDE, COLOR_WEIGHT } from '../src/lib/decimate/edge-cost-cpu.js'; +import { evalMergeCore } from '../src/lib/decimate/recost-core.js'; +import { MAX_GROUP } from '../src/lib/decimate/select.js'; +import { selectMergesRecosted } from '../src/lib/decimate/select-recost.js'; + +const NIL = 0xFFFFFFFF; + +const mulberry = (seed) => { + let t = seed >>> 0; + return () => { + t += 0x6d2b79f5; + let r = Math.imul(t ^ (t >>> 15), t | 1); + r ^= r + Math.imul(r ^ (r >>> 7), r | 61); + return ((r ^ (r >>> 14)) >>> 0) / 4294967296; + }; +}; + +// Random splats in a unit-ish box; identity quats, DC-only colour. +const makeCache = (n, seed) => { + const rand = mulberry(seed); + const pos = new Float32Array(n * 3); + const geo = new Float32Array(n * 8); + const color = new Float32Array(n * 3); + for (let i = 0; i < n; i++) { + pos[i * 3] = rand() * 2; + pos[i * 3 + 1] = rand() * 2; + pos[i * 3 + 2] = rand() * 2; + geo[i * 8] = 1; + geo[i * 8 + 4] = Math.log(0.05 + rand() * 0.1); + geo[i * 8 + 5] = Math.log(0.05 + rand() * 0.1); + geo[i * 8 + 6] = Math.log(0.05 + rand() * 0.1); + geo[i * 8 + 7] = rand() * 4 - 2; + color[i * 3] = rand() - 0.5; + color[i * 3 + 1] = rand() - 0.5; + color[i * 3 + 2] = rand() - 0.5; + } + const cache = new Float32Array(n * CACHE_STRIDE); + buildSplatCache({ pos, geo, color, colorDim: 3 }, cache); + return cache; +}; + +// Brute-force k nearest neighbour ids per splat (sentinel padded). +const bruteNeighbors = (cache, n, k) => { + const nb = new Uint32Array(n * k).fill(NIL); + const d2 = (a, b) => { + const oa = a * CACHE_STRIDE, ob = b * CACHE_STRIDE; + return (cache[oa] - cache[ob]) ** 2 + (cache[oa + 1] - cache[ob + 1]) ** 2 + (cache[oa + 2] - cache[ob + 2]) ** 2; + }; + for (let i = 0; i < n; i++) { + const ids = Array.from({ length: n }, (_, j) => j) + .filter(j => j !== i) + .sort((a, b) => d2(i, a) - d2(i, b)) + .slice(0, k); + for (let s = 0; s < ids.length; s++) nb[i * k + s] = ids[s]; + } + return nb; +}; + +// A hand-driven selection state (mirrors select-recost's structure). +const makeState = (cache, neighbors, D, n) => { + const parent = new Uint32Array(n); + const size = new Uint32Array(n).fill(1); + const version = new Uint32Array(n).fill(1); + const mHead = new Uint32Array(n); + const mNext = new Uint32Array(n).fill(NIL); + const mTail = new Uint32Array(n); + for (let i = 0; i < n; i++) { + parent[i] = i; mHead[i] = i; mTail[i] = i; + } + return { + st: { SC: cache, cands: neighbors, D, N: n, maxGroup: MAX_GROUP, parent, size, version, mHead, mNext }, + mTail + }; +}; + +// Commit a merge exactly as select-recost does (keep = larger side). +const commit = ({ st, mTail }, a, b) => { + const keep = st.size[a] >= st.size[b] ? a : b; + const lose = keep === a ? b : a; + st.mNext[mTail[keep]] = st.mHead[lose]; + mTail[keep] = mTail[lose]; + st.parent[lose] = keep; + st.size[keep] += st.size[lose]; + st.version[keep]++; + return keep; +}; + +// The ordering-only colour term between two clusters (mass-weighted mean +// base colours) — subtracted so accumulated costs telescope to pure E. +const colorTerm = (cache, membersA, membersB) => { + const mean = (members) => { + let w = 0, b0 = 0, b1 = 0, b2 = 0; + for (const m of members) { + const o = m * CACHE_STRIDE; + const mass = cache[o + 11]; + w += mass; + b0 += mass * cache[o + 12]; + b1 += mass * cache[o + 13]; + b2 += mass * cache[o + 14]; + } + return [b0 / w, b1 / w, b2 / w]; + }; + const a = mean(membersA), b = mean(membersB); + const d0 = a[0] - b[0], d1 = a[1] - b[1], d2 = a[2] - b[2]; + return COLOR_WEIGHT * (d0 * d0 + d1 * d1 + d2 * d2); +}; + +describe('stateless re-costed eval', () => { + it('singleton pair cost matches the pairwise kernel', () => { + const n = 200, k = 8; + const cache = makeCache(n, 42); + const neighbors = bruteNeighbors(cache, n, k); + const { st } = makeState(cache, neighbors, k, n); + // Same gate as the reference tool's selftest (rel 1e-3): near-zero + // costs are differences of large self terms, so association noise + // amplifies relatively — the accepted near-tie class. + let worst = 0; + for (let i = 0; i < n; i += 3) { + const j = neighbors[i * k]; + const mine = evalMergeCore(st, i, j); + const lib = computeEdgeCost(cache, i, j); + const rel = Math.abs(mine - lib) / Math.max(1e-12, Math.abs(lib)); + worst = Math.max(worst, rel); + } + assert.ok(worst < 1e-3, `worst singleton parity rel diff ${worst.toExponential(2)}`); + }); + + it('accumulated marginal cost is merge-path independent (telescopes to E)', () => { + const n = 60, k = 8; + const cache = makeCache(n, 7); + const neighbors = bruteNeighbors(cache, n, k); + + for (let a = 0; a < n; a += 7) { + const b = neighbors[a * k]; + const c = neighbors[a * k + 1]; + + // Path 1: (a+b) then (ab+c). + const s1 = makeState(cache, neighbors, k, n); + const e1ab = evalMergeCore(s1.st, a, b) - colorTerm(cache, [a], [b]); + const ab = commit(s1, a, b); + const e1c = evalMergeCore(s1.st, ab, c) - colorTerm(cache, [a, b], [c]); + + // Path 2: (a+c) then (ac+b). + const s2 = makeState(cache, neighbors, k, n); + const e2ac = evalMergeCore(s2.st, a, c) - colorTerm(cache, [a], [c]); + const ac = commit(s2, a, c); + const e2b = evalMergeCore(s2.st, ac, b) - colorTerm(cache, [a, c], [b]); + + const sum1 = e1ab + e1c; + const sum2 = e2ac + e2b; + assert.ok( + Math.abs(sum1 - sum2) <= Math.max(1e-12, 1e-8 * Math.abs(sum1)), + `triple (${a},${b},${c}): path sums ${sum1} vs ${sum2}` + ); + } + }); +}); + +describe('selectMergesRecosted', () => { + // Build seeds the way decimate-source does: best candidate per splat. + const makeSeeds = (cache, neighbors, n, k, K) => { + const cand = { + idx: new Uint32Array(n * K).fill(NIL), + cost: new Float32Array(n * K).fill(Infinity) + }; + for (let i = 0; i < n; i++) { + let bc = Infinity, bj = NIL; + for (let s = 0; s < k; s++) { + const j = neighbors[i * k + s]; + if (j === NIL) continue; + const cst = computeEdgeCost(cache, i, j); + if (cst < bc) { + bc = cst; bj = j; + } + } + cand.idx[i * K] = bj; + cand.cost[i * K] = bc; + } + return cand; + }; + + it('hits the target with a well-formed capped CSR', () => { + const n = 256, k = 8, K = 4; + const cache = makeCache(n, 99); + const neighbors = bruteNeighbors(cache, n, k); + const cand = makeSeeds(cache, neighbors, n, k, K); + const needed = n >> 1; + const sel = selectMergesRecosted({ cand, K, splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); + + assert.strictEqual(sel.removed, needed); + assert.strictEqual(sel.groupOffsets[sel.mergedGroups], sel.groupMembers.length); + const seen = new Set(); + let members = 0; + for (let g = 0; g < sel.mergedGroups; g++) { + const sz = sel.groupOffsets[g + 1] - sel.groupOffsets[g]; + assert.ok(sz >= 2 && sz <= MAX_GROUP, `group ${g} size ${sz}`); + members += sz; + for (let s = sel.groupOffsets[g]; s < sel.groupOffsets[g + 1]; s++) { + const m = sel.groupMembers[s]; + assert.ok(!seen.has(m), 'member in one group only'); + seen.add(m); + assert.strictEqual(sel.memberGroup[m], g); + } + assert.strictEqual(sel.groupMin[g], sel.groupMembers[sel.groupOffsets[g]]); + } + assert.strictEqual(members - sel.mergedGroups, needed); + }); + + it('concentrates coincident splats into full groups (refresh continuation)', () => { + const n = 64, k = 16, K = 4; + // All splats identical and coincident. + const pos = new Float32Array(n * 3).fill(0.5); + const geo = new Float32Array(n * 8); + const color = new Float32Array(n * 3).fill(0.25); + for (let i = 0; i < n; i++) { + geo[i * 8] = 1; + geo[i * 8 + 4] = geo[i * 8 + 5] = geo[i * 8 + 6] = Math.log(0.1); + geo[i * 8 + 7] = 0; + } + const cache = new Float32Array(n * CACHE_STRIDE); + buildSplatCache({ pos, geo, color, colorDim: 3 }, cache); + // Ring neighbour graph (i ± 1..8 mod n): coincident points have + // arbitrary-but-diverse KNN in the real pipeline; a hub-shaped graph + // (everyone pointing at the same few ids) would starve pools by + // construction, which is a fixture artifact, not pipeline behaviour. + const neighbors = new Uint32Array(n * k); + for (let i = 0; i < n; i++) { + for (let s = 0; s < k; s++) { + const off = (s >> 1) + 1; + neighbors[i * k + s] = (i + (s & 1 ? n - off : off)) % n; + } + } + const cand = makeSeeds(cache, neighbors, n, k, K); + const needed = n - (n / MAX_GROUP); // 48: reachable only by filling groups to the cap + const sel = selectMergesRecosted({ cand, K, splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); + + // Disjoint pairing alone tops out at n/2 = 32 removals; anything well + // beyond that requires continuation merges via refresh re-entry (the + // concentration behaviour). The ring graph's bounded reach (±8) can + // strand a final under-full pair, so demand near-full rather than + // perfect packing. + assert.ok(sel.removed >= needed - MAX_GROUP, `removed ${sel.removed} of ${needed}`); + let full = 0; + for (let g = 0; g < sel.mergedGroups; g++) { + const sz = sel.groupOffsets[g + 1] - sel.groupOffsets[g]; + assert.ok(sz <= MAX_GROUP, `group ${g} size ${sz} over cap`); + if (sz === MAX_GROUP) full++; + } + assert.ok(full >= (n / MAX_GROUP) - 2, `only ${full} full groups`); + }); +}); From 5bfbb592d77a23ed8802ad6a1d44c99294fac4b6 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 16:11:47 +0100 Subject: [PATCH 07/19] latest --- src/lib/decimate/decimate-source.ts | 29 ++--- src/lib/decimate/priority.ts | 161 +++++++++++++++------------- src/lib/decimate/select-recost.ts | 34 +++--- test/decimate-priority.test.mjs | 39 +++++++ test/decimate-recost.test.mjs | 32 +----- 5 files changed, 160 insertions(+), 135 deletions(-) diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index b7dcb759..b7d8eb8b 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -37,10 +37,11 @@ const MIN_ITERATION_PROGRESS = 0.05; const DEFAULT_MEMORY_BUDGET = 48 * 2 ** 30; // Per-gaussian residency of re-costed selection beyond the base state: splat -// cache (16 f32) + neighbour ids (k u32) + f64 cluster moments/colour/error + -// union-find/chains/heap. Conservative round-up; used by the per-generation -// gate that falls back to one-shot selectMerges when over budget. -const RECOST_BYTES_PER_GAUSSIAN = (k: number) => CACHE_STRIDE * 4 + k * 4 + 256; +// cache (16 f32) + neighbour ids (k u32) + integer structure (union-find, +// chains, seq/round bookkeeping: 8×u32) + heap (5 arrays × 1.25N ≈ 25 B). +// Conservative round-up; used by the per-generation gate that falls back to +// one-shot selectMerges when over budget. +const RECOST_BYTES_PER_GAUSSIAN = (k: number) => CACHE_STRIDE * 4 + k * 4 + 80; /** Coherence heuristic: gap (rows) merged into one run / runs-per-block considered scattered. */ const COHERENCE_GAP_ROWS = 64; @@ -233,10 +234,14 @@ const decimateSource = async ( const baseBytes = residentInputBytes + N * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; const recost = !legacy && baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; - const cand: CandidateArrays = { - idx: new Uint32Array(N * K).fill(0xFFFFFFFF), - cost: new Float32Array(N * K).fill(Infinity) - }; + // Re-costed selection seeds itself from the neighbour graph (wave 0), + // so candidate arrays exist only for the one-shot/legacy selections. + const cand: CandidateArrays | undefined = recost ? + undefined : + { + idx: new Uint32Array(N * K).fill(0xFFFFFFFF), + cost: new Float32Array(N * K).fill(Infinity) + }; const cacheOut = recost ? new Float32Array(N * CACHE_STRIDE) : undefined; const neighborsOut = recost ? new Uint32Array(N * k) : undefined; @@ -244,7 +249,7 @@ const decimateSource = async ( if (legacy) { await runPriorityPassLegacy( { source: src, pool, pos: positions, order, blocks, device, K, k }, - cand, + cand!, n => priorityBar.tick(n) ); } else { @@ -260,10 +265,10 @@ const decimateSource = async ( const needed = N - generationTarget; const selectSub = logger.group(recost ? 'Selecting merges (re-costed)' : 'Selecting merges'); const selection = legacy ? - selectMergesLegacy(cand, N, K, needed) : + selectMergesLegacy(cand!, N, K, needed) : cacheOut ? - selectMergesRecosted({ cand, K, splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed }) : - selectMerges(cand, N, K, needed); + selectMergesRecosted({ splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed }) : + selectMerges(cand!, N, K, needed); selectSub.end(); if (selection.removed === 0) { diff --git a/src/lib/decimate/priority.ts b/src/lib/decimate/priority.ts index fbf39268..e5eb286b 100644 --- a/src/lib/decimate/priority.ts +++ b/src/lib/decimate/priority.ts @@ -175,13 +175,19 @@ const indexOfSorted = (sorted: Uint32Array, g: number): number => { * for each owned gaussian's k neighbours, reduction to the best K candidates * — written into the resident candidate arrays. * + * Without `cand` (re-costed selection seeds itself from the neighbour graph) + * the pass only persists the splat cache + neighbour ids: no externals, no + * edge costs, no top-K — the gather narrows to owned rows and the GPU cost + * kernel is never constructed. + * * @param ctx - The pass context. - * @param cand - Preallocated candidate arrays (`N*K`), filled per block. + * @param cand - Candidate arrays (`N*K`) to fill, or undefined to skip all + * cost work (cacheOut/neighborsOut persistence only). * @param tick - Optional progress callback (owned gaussians completed). */ const runPriorityPass = async ( ctx: PriorityContext, - cand: CandidateArrays, + cand: CandidateArrays | undefined, tick?: (n: number) => void ): Promise => { const { pos, order, blocks, device, K, k } = ctx; @@ -219,7 +225,7 @@ const runPriorityPass = async ( try { if (device) { gpuKnn = new GpuKnn(device, maxLocalN, k); - gpuCost = new GpuEdgeCost(device, maxLocalN, k); + if (cand) gpuCost = new GpuEdgeCost(device, maxLocalN, k); } let next: Prepared | null = blocks.length > 0 ? prepare(0) : null; @@ -238,28 +244,33 @@ const runPriorityPass = async ( // Externals: referenced rows outside the owned range (halo members // and verification-fixed neighbours) — count, collect, sort, dedup. - let extCount = 0; - for (let s = 0; s < slots; s++) { - const l = nbLocal[s]; - if (l === KNN_SENTINEL || l < nOwned) continue; - if (l === KNN_FIXED && indexOfSorted(owned, nbGlobal[s]) >= 0) continue; - extCount++; - } - const extSorted = new Uint32Array(extCount); - extCount = 0; - for (let s = 0; s < slots; s++) { - const l = nbLocal[s]; - if (l === KNN_SENTINEL || l < nOwned) continue; - const g = nbGlobal[s]; - if (l === KNN_FIXED && indexOfSorted(owned, g) >= 0) continue; - extSorted[extCount++] = g; - } - extSorted.sort(); - let uniq = 0; - for (let i = 0; i < extCount; i++) { - if (i === 0 || extSorted[i] !== extSorted[i - 1]) extSorted[uniq++] = extSorted[i]; + // Only cost evaluation needs them; the persisted cache covers owned + // rows only (every gaussian is owned by exactly one block). + let extraGlobals = new Uint32Array(0); + if (cand) { + let extCount = 0; + for (let s = 0; s < slots; s++) { + const l = nbLocal[s]; + if (l === KNN_SENTINEL || l < nOwned) continue; + if (l === KNN_FIXED && indexOfSorted(owned, nbGlobal[s]) >= 0) continue; + extCount++; + } + const extSorted = new Uint32Array(extCount); + extCount = 0; + for (let s = 0; s < slots; s++) { + const l = nbLocal[s]; + if (l === KNN_SENTINEL || l < nOwned) continue; + const g = nbGlobal[s]; + if (l === KNN_FIXED && indexOfSorted(owned, g) >= 0) continue; + extSorted[extCount++] = g; + } + extSorted.sort(); + let uniq = 0; + for (let i = 0; i < extCount; i++) { + if (i === 0 || extSorted[i] !== extSorted[i - 1]) extSorted[uniq++] = extSorted[i]; + } + extraGlobals = extSorted.subarray(0, uniq); } - const extraGlobals = extSorted.subarray(0, uniq); // Verification-fixed externals are not bounded by the halo cap, so // a pathological block's view can exceed the preallocated cost @@ -279,29 +290,61 @@ const runPriorityPass = async ( const cache = new Float32Array(viewN * CACHE_STRIDE); buildSplatCache(view, cache); - // Translate neighbour slots in place: local/global → view rows - // (dense-slot edge model; sentinel slots stay sentinel). - for (let s = 0; s < slots; s++) { - const l = nbLocal[s]; - if (l === KNN_SENTINEL || l < nOwned) continue; - const g = nbGlobal[s]; - if (l === KNN_FIXED) { - const oi = indexOfSorted(owned, g); - nbLocal[s] = oi >= 0 ? oi : nOwned + indexOfSorted(extraGlobals, g); + if (cand) { + // Translate neighbour slots in place: local/global → view rows + // (dense-slot edge model; sentinel slots stay sentinel). + for (let s = 0; s < slots; s++) { + const l = nbLocal[s]; + if (l === KNN_SENTINEL || l < nOwned) continue; + const g = nbGlobal[s]; + if (l === KNN_FIXED) { + const oi = indexOfSorted(owned, g); + nbLocal[s] = oi >= 0 ? oi : nOwned + indexOfSorted(extraGlobals, g); + } else { + nbLocal[s] = nOwned + indexOfSorted(extraGlobals, g); + } + } + + const blockCosts = new Float32Array(slots); + if (device) { + await gpuCost!.execute(cache, viewN, nbLocal, blockCosts); } else { - nbLocal[s] = nOwned + indexOfSorted(extraGlobals, g); + for (let s = 0; s < slots; s++) { + const row = nbLocal[s]; + blockCosts[s] = row === KNN_SENTINEL ? + 0 : + computeEdgeCost(cache, (s / k) | 0, row); + } } - } - const blockCosts = new Float32Array(slots); - if (device) { - await gpuCost!.execute(cache, viewN, nbLocal, blockCosts); - } else { - for (let s = 0; s < slots; s++) { - const row = nbLocal[s]; - blockCosts[s] = row === KNN_SENTINEL ? - 0 : - computeEdgeCost(cache, (s / k) | 0, row); + // Reduce to best K candidates per owned gaussian (ascending by + // cost; sentinel slots are skipped by id — their cost values + // are never read). + const bestIdx = new Uint32Array(K); + const bestCost = new Float64Array(K); + for (let qi = 0; qi < nOwned; qi++) { + let size = 0; + const base = qi * k; + for (let s = 0; s < k; s++) { + if (nbLocal[base + s] === KNN_SENTINEL) continue; + const c = blockCosts[base + s]; + if (!Number.isFinite(c)) continue; + if (size === K && c >= bestCost[K - 1]) continue; + let at = size < K ? size : K - 1; + while (at > 0 && bestCost[at - 1] > c) { + bestCost[at] = bestCost[at - 1]; + bestIdx[at] = bestIdx[at - 1]; + at--; + } + bestCost[at] = c; + bestIdx[at] = nbGlobal[base + s]; + size = Math.min(size + 1, K); + } + const g = owned[qi]; + for (let s = 0; s < K; s++) { + cand.idx[g * K + s] = s < size ? bestIdx[s] : 0xFFFFFFFF; + cand.cost[g * K + s] = s < size ? bestCost[s] : Infinity; + } } } @@ -321,36 +364,6 @@ const runPriorityPass = async ( } } - // Reduce to best K candidates per owned gaussian (ascending by - // cost; sentinel slots are skipped by id — their cost values are - // never read). - const bestIdx = new Uint32Array(K); - const bestCost = new Float64Array(K); - for (let qi = 0; qi < nOwned; qi++) { - let size = 0; - const base = qi * k; - for (let s = 0; s < k; s++) { - if (nbLocal[base + s] === KNN_SENTINEL) continue; - const c = blockCosts[base + s]; - if (!Number.isFinite(c)) continue; - if (size === K && c >= bestCost[K - 1]) continue; - let at = size < K ? size : K - 1; - while (at > 0 && bestCost[at - 1] > c) { - bestCost[at] = bestCost[at - 1]; - bestIdx[at] = bestIdx[at - 1]; - at--; - } - bestCost[at] = c; - bestIdx[at] = nbGlobal[base + s]; - size = Math.min(size + 1, K); - } - const g = owned[qi]; - for (let s = 0; s < K; s++) { - cand.idx[g * K + s] = s < size ? bestIdx[s] : 0xFFFFFFFF; - cand.cost[g * K + s] = s < size ? bestCost[s] : Infinity; - } - } - tick?.(nOwned); } } finally { diff --git a/src/lib/decimate/select-recost.ts b/src/lib/decimate/select-recost.ts index ef1e98d8..71dd6db0 100644 --- a/src/lib/decimate/select-recost.ts +++ b/src/lib/decimate/select-recost.ts @@ -22,9 +22,9 @@ * cluster moments are recomputed from the ≤ MAX_GROUP member cache rows on * every evaluation, so this module keeps no per-root float state at all — * a commit is a chain splice plus union-find/size/version updates, and the - * heap carries only (cost, a, b, seq, versionB). Seeds come from the - * priority pass's candidate arrays; refresh candidate pools come from the - * persisted neighbour graph. + * heap carries only (cost, a, b, seq, versionB). Seeding is wave 0: every + * gaussian starts queued, so the first refresh evaluates each singleton's + * best edge over its neighbour rows — no candidate arrays are consumed. * * decimate-source gates this path by memory budget and falls back to * {@link selectMerges} when over. @@ -32,18 +32,13 @@ * Engine-free; single-threaded, pure resident-array computation, no IO. */ -import { type CandidateArrays } from './priority'; -import { bestEdgeFor, bestOut, NO_CANDIDATE, CACHE_STRIDE, type RecostState } from './recost-core'; +import { bestEdgeFor, bestOut, CACHE_STRIDE, type RecostState } from './recost-core'; import { MAX_GROUP, type SelectionResult } from './select'; const NIL = 0xFFFFFFFF; /** Inputs for {@link selectMergesRecosted}. */ type RecostInputs = { - /** Per-gaussian best-K candidates from the priority pass (seed costs include the colour term). */ - cand: CandidateArrays; - /** Candidates per gaussian. */ - K: number; /** * Resident per-splat cache, {@link CACHE_STRIDE} floats per splat * (buildSplatCache row layout), persisted by the priority pass. @@ -71,7 +66,7 @@ type RecostInputs = { * @returns The selection (same contract as {@link selectMerges}). */ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { - const { cand, K, splatCache: SC, neighbors, D, N, mergesNeeded } = inputs; + const { splatCache: SC, neighbors, D, N, mergesNeeded } = inputs; // ---- Cluster structure (indexed by union-find root) — integers only. const parent = new Uint32Array(N); @@ -175,7 +170,7 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { // Refresh queue: clusters whose best edge must be re-evaluated. Queuing // bumps lastSeq so stale heap entries for the cluster discard on pop; // queuedRound dedupes within a round. - let pending = new Uint32Array(1 << 16); + const pending = new Uint32Array(N); let pendingCount = 0; const queuedRound = new Uint32Array(N); let round = 1; @@ -183,22 +178,19 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { lastSeq[root] = ++seqCounter; if (queuedRound[root] === round) return; queuedRound[root] = round; - if (pendingCount === pending.length) { - const g = new Uint32Array(pending.length * 2); - g.set(pending); pending = g; - } pending[pendingCount++] = root; }; - // Seed: the priority pass's cheapest candidate per gaussian. + // Seed = wave 0: every gaussian starts queued, so the first refresh + // evaluates each singleton's best edge over its neighbour rows (identical + // work to a per-edge cost pass + argmin, and the same eval as every later + // refresh). The first drain is a no-op on the empty heap. for (let i = 0; i < N; i++) { - const j = cand.idx[i * K]; - const c = cand.cost[i * K]; lastSeq[i] = ++seqCounter; - if (j !== NO_CANDIDATE && Number.isFinite(c)) { - heapPush(c, i, j, lastSeq[i], version[j]); - } + queuedRound[i] = round; + pending[i] = i; } + pendingCount = N; // ---- Round loop. Each round commits at most WAVE merges before the // bulk refresh, so refreshed edges (notably cheap continuation merges in diff --git a/test/decimate-priority.test.mjs b/test/decimate-priority.test.mjs index 40bd6edd..d3b7a098 100644 --- a/test/decimate-priority.test.mjs +++ b/test/decimate-priority.test.mjs @@ -46,6 +46,45 @@ describe('priority pass (CPU)', () => { } }); + it('persist-only mode (no cand) fills cache + neighbours without cost work', async () => { + const n = 900, k = 16, K = 4; + const { source, pool, view, pos } = await makeSyntheticSource(n, 1, 21, { chunkSize: 256 }); + const { order, blocks } = kdPartition(pos, 300); + const cacheOut = new Float32Array(n * CACHE_STRIDE); + const neighborsOut = new Uint32Array(n * k); + await runPriorityPass( + { source, pool, pos, order, blocks, K, k, cacheOut, neighborsOut }, + undefined + ); + + // Cache rows must equal buildSplatCache over the same splats. + const ref = new Float32Array(n * CACHE_STRIDE); + buildSplatCache(view, ref); + for (let i = 0; i < n; i += 53) { + for (let c = 0; c < CACHE_STRIDE; c++) { + const got = cacheOut[i * CACHE_STRIDE + c]; + const want = ref[i * CACHE_STRIDE + c]; + assert.ok( + Math.abs(got - want) <= Math.max(1e-6, Math.abs(want) * 1e-6), + `cache row ${i} field ${c}: ${got} vs ${want}` + ); + } + } + + // Neighbour rows must be the exact global KNN as a set. + const d2 = (a, b) => (pos.x[a] - pos.x[b]) ** 2 + (pos.y[a] - pos.y[b]) ** 2 + (pos.z[a] - pos.z[b]) ** 2; + for (let i = 0; i < n; i += 97) { + const brute = new Set(Array.from({ length: n }, (_, j) => j) + .filter(j => j !== i) + .sort((a, b) => d2(i, a) - d2(i, b)) + .slice(0, k)); + for (let s = 0; s < k; s++) { + const g = neighborsOut[i * k + s]; + assert.ok(brute.has(g), `query ${i} slot ${s}: ${g} not in brute-force KNN`); + } + } + }); + it('candidate ids are real neighbours (no self, no sentinels leaking as ids)', async () => { const n = 600, k = 16, K = 2; const { source, pool, pos } = await makeSyntheticSource(n, 0, 9, { chunkSize: 128 }); diff --git a/test/decimate-recost.test.mjs b/test/decimate-recost.test.mjs index 01d6580e..e7eba702 100644 --- a/test/decimate-recost.test.mjs +++ b/test/decimate-recost.test.mjs @@ -175,35 +175,12 @@ describe('stateless re-costed eval', () => { }); describe('selectMergesRecosted', () => { - // Build seeds the way decimate-source does: best candidate per splat. - const makeSeeds = (cache, neighbors, n, k, K) => { - const cand = { - idx: new Uint32Array(n * K).fill(NIL), - cost: new Float32Array(n * K).fill(Infinity) - }; - for (let i = 0; i < n; i++) { - let bc = Infinity, bj = NIL; - for (let s = 0; s < k; s++) { - const j = neighbors[i * k + s]; - if (j === NIL) continue; - const cst = computeEdgeCost(cache, i, j); - if (cst < bc) { - bc = cst; bj = j; - } - } - cand.idx[i * K] = bj; - cand.cost[i * K] = bc; - } - return cand; - }; - it('hits the target with a well-formed capped CSR', () => { - const n = 256, k = 8, K = 4; + const n = 256, k = 8; const cache = makeCache(n, 99); const neighbors = bruteNeighbors(cache, n, k); - const cand = makeSeeds(cache, neighbors, n, k, K); const needed = n >> 1; - const sel = selectMergesRecosted({ cand, K, splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); + const sel = selectMergesRecosted({ splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); assert.strictEqual(sel.removed, needed); assert.strictEqual(sel.groupOffsets[sel.mergedGroups], sel.groupMembers.length); @@ -225,7 +202,7 @@ describe('selectMergesRecosted', () => { }); it('concentrates coincident splats into full groups (refresh continuation)', () => { - const n = 64, k = 16, K = 4; + const n = 64, k = 16; // All splats identical and coincident. const pos = new Float32Array(n * 3).fill(0.5); const geo = new Float32Array(n * 8); @@ -248,9 +225,8 @@ describe('selectMergesRecosted', () => { neighbors[i * k + s] = (i + (s & 1 ? n - off : off)) % n; } } - const cand = makeSeeds(cache, neighbors, n, k, K); const needed = n - (n / MAX_GROUP); // 48: reachable only by filling groups to the cap - const sel = selectMergesRecosted({ cand, K, splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); + const sel = selectMergesRecosted({ splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); // Disjoint pairing alone tops out at n/2 = 32 removals; anything well // beyond that requires continuation merges via refresh re-entry (the From 7c11b45f348e00d219cce3812997c53cce53226e Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 16:29:34 +0100 Subject: [PATCH 08/19] latest --- src/lib/gpu/compute-kernel.ts | 66 +++++++++++++++++++++++ src/lib/gpu/gpu-edge-cost.ts | 44 +++------------ src/lib/gpu/gpu-kmeans.ts | 51 +----------------- src/lib/gpu/shaders/chunks/gaussian-l2.ts | 51 ++++++++++++++++++ 4 files changed, 125 insertions(+), 87 deletions(-) create mode 100644 src/lib/gpu/compute-kernel.ts create mode 100644 src/lib/gpu/shaders/chunks/gaussian-l2.ts diff --git a/src/lib/gpu/compute-kernel.ts b/src/lib/gpu/compute-kernel.ts new file mode 100644 index 00000000..7026298b --- /dev/null +++ b/src/lib/gpu/compute-kernel.ts @@ -0,0 +1,66 @@ +import { + SHADERLANGUAGE_WGSL, + SHADERSTAGE_COMPUTE, + UNIFORMTYPE_UINT, + BindGroupFormat, + BindStorageBufferFormat, + BindUniformBufferFormat, + Compute, + GraphicsDevice, + Shader, + UniformBufferFormat, + UniformFormat +} from 'playcanvas'; + +type Kernel = { + compute: Compute; + destroy: () => void; +}; + +/** + * Shader + bind group format + compute boilerplate shared by compute kernels. + * Uniforms are u32 (the struct in the WGSL must match `uniformNames` order); + * storage bindings are `[name, readOnly]` pairs in binding order after the + * uniform buffer. + * + * @param device - PlayCanvas GraphicsDevice (WebGPU). + * @param name - Kernel name (shader + compute label). + * @param source - WGSL source. + * @param uniformNames - u32 uniform names, struct order. + * @param storageBindings - Storage buffer bindings, `[name, readOnly]`. + * @returns The compute wrapper and its destroy. + */ +const makeKernel = ( + device: GraphicsDevice, + name: string, + source: string, + uniformNames: string[], + storageBindings: [string, boolean][] +): Kernel => { + const bindGroupFormat = new BindGroupFormat(device, [ + new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE), + ...storageBindings.map(([bname, readOnly]) => new BindStorageBufferFormat(bname, SHADERSTAGE_COMPUTE, readOnly)) + ]); + + const shader = new Shader(device, { + name, + shaderLanguage: SHADERLANGUAGE_WGSL, + cshader: source, + // @ts-ignore + computeUniformBufferFormats: { + uniforms: new UniformBufferFormat(device, uniformNames.map(u => new UniformFormat(u, UNIFORMTYPE_UINT))) + }, + // @ts-ignore + computeBindGroupFormat: bindGroupFormat + }); + + return { + compute: new Compute(device, shader, name), + destroy: () => { + shader.destroy(); + bindGroupFormat.destroy(); + } + }; +}; + +export { makeKernel, type Kernel }; diff --git a/src/lib/gpu/gpu-edge-cost.ts b/src/lib/gpu/gpu-edge-cost.ts index 2fa63ad2..0bc8d3e5 100644 --- a/src/lib/gpu/gpu-edge-cost.ts +++ b/src/lib/gpu/gpu-edge-cost.ts @@ -16,6 +16,7 @@ import { } from 'playcanvas'; import { CACHE_STRIDE } from '../decimate/edge-cost-cpu'; +import { gaussianL2Wgsl } from './shaders/chunks/gaussian-l2'; /** Per-splat interleaved stride in the `splat` storage buffer (= the CPU cost-cache layout). */ export const SPLAT_STRIDE = CACHE_STRIDE; @@ -62,23 +63,7 @@ struct Uniforms { const K: u32 = ${k}u; const SENTINEL: u32 = 0xFFFFFFFFu; - -const EPS_COV: f32 = 1e-8; -const PI_1_5: f32 = 5.5683279968317084; // π^{3/2} -const TWO_PI_1_5: f32 = 15.749609945722419; // (2π)^{3/2} -const TWO_PI_3: f32 = 2.0943951023931953; // 2π/3 -const ELLIP_P: f32 = 1.6075; -// Scale-free DC colour dissimilarity weight (4π·1e-6 in base-colour space — -// see COLOR_WEIGHT in decimate/edge-cost-cpu.ts, mirrored here). -const COLOR_WEIGHT: f32 = 1.2566370614359172e-5; - -// Knud Thomsen ellipsoid surface area (matches CPU ellipsoidArea). -fn ellipsoidArea(sx: f32, sy: f32, sz: f32) -> f32 { - let a = pow(sx * sy, ELLIP_P); - let b = pow(sx * sz, ELLIP_P); - let c = pow(sy * sz, ELLIP_P); - return 4.0 * 3.141592653589793 * pow((a + b + c) / 3.0, 1.0 / ELLIP_P); -} +${gaussianL2Wgsl} // Gaussian cross-product ⟨G_a,G_b⟩ for symmetric M = Σ_a+Σ_b (6: xx,xy,xz,yy,yz,zz) // and mean offset d, scaled by √|Σ_a|·√|Σ_b| (passed as sdA·sdB). @@ -148,27 +133,10 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { // Merged opacity: mass-conserving, capped — needs merged scales (eigenvalues // of Σ_m, Smith's closed form for a symmetric 3×3). - let q = (sm[0] + sm[3] + sm[5]) / 3.0; - let p1 = sm[1] * sm[1] + sm[2] * sm[2] + sm[4] * sm[4]; - var e0: f32; var e1: f32; var e2: f32; - if (p1 <= 1e-30) { - e0 = sm[0]; e1 = sm[3]; e2 = sm[5]; - } else { - let p2 = (sm[0] - q) * (sm[0] - q) + (sm[3] - q) * (sm[3] - q) + (sm[5] - q) * (sm[5] - q) + 2.0 * p1; - let p = sqrt(p2 / 6.0); - let ip = 1.0 / p; - let b00 = (sm[0] - q) * ip; let b11 = (sm[3] - q) * ip; let b22 = (sm[5] - q) * ip; - let b01 = sm[1] * ip; let b02 = sm[2] * ip; let b12 = sm[4] * ip; - let detB = b00 * (b11 * b22 - b12 * b12) - b01 * (b01 * b22 - b12 * b02) + b02 * (b01 * b12 - b11 * b02); - let r = clamp(detB * 0.5, -1.0, 1.0); - let phi = acos(r) / 3.0; - e0 = q + 2.0 * p * cos(phi); - e2 = q + 2.0 * p * cos(phi + TWO_PI_3); - e1 = 3.0 * q - e0 - e2; - } - let s0 = sqrt(max(e0, 1e-18)); - let s1 = sqrt(max(e1, 1e-18)); - let s2 = sqrt(max(e2, 1e-18)); + let e = eig3(sm); + let s0 = sqrt(max(e.x, 1e-18)); + let s1 = sqrt(max(e.y, 1e-18)); + let s2 = sqrt(max(e.z, 1e-18)); let am = min(1.0, W / max(ellipsoidArea(s0, s1, s2), 1e-30)); // Base-colour dots (merged colour = mass-weighted average). diff --git a/src/lib/gpu/gpu-kmeans.ts b/src/lib/gpu/gpu-kmeans.ts index 006a89de..df3e8536 100644 --- a/src/lib/gpu/gpu-kmeans.ts +++ b/src/lib/gpu/gpu-kmeans.ts @@ -1,21 +1,13 @@ import { BUFFERUSAGE_COPY_DST, BUFFERUSAGE_COPY_SRC, - SHADERLANGUAGE_WGSL, - SHADERSTAGE_COMPUTE, - UNIFORMTYPE_UINT, - BindGroupFormat, - BindStorageBufferFormat, - BindUniformBufferFormat, - Compute, FloatPacking, - Shader, StorageBuffer, - UniformBufferFormat, - UniformFormat, GraphicsDevice } from 'playcanvas'; +import { makeKernel, type Kernel } from './compute-kernel'; + /** * Flash-kmeans (arXiv 2603.09229) adapted to WebGPU: a fully GPU-resident * Lloyd loop. Per iteration, all stages are recorded with zero CPU @@ -391,45 +383,6 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { } `; -type Kernel = { - compute: Compute; - destroy: () => void; -}; - -// shader + bind group format + compute boilerplate shared by all kernels -const makeKernel = ( - device: GraphicsDevice, - name: string, - source: string, - uniformNames: string[], - storageBindings: [string, boolean][] // [name, readOnly] -): Kernel => { - const bindGroupFormat = new BindGroupFormat(device, [ - new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE), - ...storageBindings.map(([bname, readOnly]) => new BindStorageBufferFormat(bname, SHADERSTAGE_COMPUTE, readOnly)) - ]); - - const shader = new Shader(device, { - name, - shaderLanguage: SHADERLANGUAGE_WGSL, - cshader: source, - // @ts-ignore - computeUniformBufferFormats: { - uniforms: new UniformBufferFormat(device, uniformNames.map(u => new UniformFormat(u, UNIFORMTYPE_UINT))) - }, - // @ts-ignore - computeBindGroupFormat: bindGroupFormat - }); - - return { - compute: new Compute(device, shader, name), - destroy: () => { - shader.destroy(); - bindGroupFormat.destroy(); - } - }; -}; - class GpuKmeans { /** * Run the full Lloyd loop on the GPU. `points` is row-major interleaved diff --git a/src/lib/gpu/shaders/chunks/gaussian-l2.ts b/src/lib/gpu/shaders/chunks/gaussian-l2.ts new file mode 100644 index 00000000..73a8dac2 --- /dev/null +++ b/src/lib/gpu/shaders/chunks/gaussian-l2.ts @@ -0,0 +1,51 @@ +/** + * Shared WGSL for the field-L2 decimation cost: constants, the Knud Thomsen + * ellipsoid area, and Smith's closed-form symmetric-3×3 eigenvalues — + * interpolated into the kernels that evaluate merged-Gaussian costs + * (GpuEdgeCost, GpuRecost). The Gaussian cross-product stays per-kernel. + * + * Mirrors the CPU implementations in decimate/edge-cost-cpu.ts and + * decimate/recost-core.ts — keep them in lockstep. + */ +const gaussianL2Wgsl = /* wgsl */` +const EPS_COV: f32 = 1e-8; +const PI_1_5: f32 = 5.5683279968317084; // π^{3/2} +const TWO_PI_1_5: f32 = 15.749609945722419; // (2π)^{3/2} +const TWO_PI_3: f32 = 2.0943951023931953; // 2π/3 +const ELLIP_P: f32 = 1.6075; +// Scale-free DC colour dissimilarity weight (4π·1e-6 in base-colour space — +// see COLOR_WEIGHT in decimate/edge-cost-cpu.ts, mirrored here). +const COLOR_WEIGHT: f32 = 1.2566370614359172e-5; + +// Knud Thomsen ellipsoid surface area (matches CPU ellipsoidArea). +fn ellipsoidArea(sx: f32, sy: f32, sz: f32) -> f32 { + let a = pow(sx * sy, ELLIP_P); + let b = pow(sx * sz, ELLIP_P); + let c = pow(sy * sz, ELLIP_P); + return 4.0 * 3.141592653589793 * pow((a + b + c) / 3.0, 1.0 / ELLIP_P); +} + +// Smith closed-form eigenvalues of a symmetric 3×3 (6 comps: xx,xy,xz,yy,yz,zz). +// Consumers feed these through ellipsoidArea (symmetric), so no ordering is +// promised beyond the Smith branch's (largest, middle, smallest). +fn eig3(m: array) -> vec3f { + let q = (m[0] + m[3] + m[5]) / 3.0; + let p1 = m[1] * m[1] + m[2] * m[2] + m[4] * m[4]; + if (p1 <= 1e-30) { + return vec3f(m[0], m[3], m[5]); + } + let p2 = (m[0] - q) * (m[0] - q) + (m[3] - q) * (m[3] - q) + (m[5] - q) * (m[5] - q) + 2.0 * p1; + let p = sqrt(p2 / 6.0); + let ip = 1.0 / p; + let b00 = (m[0] - q) * ip; let b11 = (m[3] - q) * ip; let b22 = (m[5] - q) * ip; + let b01 = m[1] * ip; let b02 = m[2] * ip; let b12 = m[4] * ip; + let detB = b00 * (b11 * b22 - b12 * b12) - b01 * (b01 * b22 - b12 * b02) + b02 * (b01 * b12 - b11 * b02); + let r = clamp(detB * 0.5, -1.0, 1.0); + let phi = acos(r) / 3.0; + let e0 = q + 2.0 * p * cos(phi); + let e2 = q + 2.0 * p * cos(phi + TWO_PI_3); + return vec3f(e0, 3.0 * q - e0 - e2, e2); +} +`; + +export { gaussianL2Wgsl }; From 62ea2d953b1f0416e423daa12e68605143bb276a Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 17:18:05 +0100 Subject: [PATCH 09/19] latest --- src/lib/decimate/decimate-source.ts | 2 +- src/lib/decimate/select-recost.ts | 129 ++++-- src/lib/gpu/gpu-recost.ts | 597 ++++++++++++++++++++++++++++ test/decimate-recost.test.mjs | 8 +- test/gpu-recost.test.mjs | 213 ++++++++++ 5 files changed, 906 insertions(+), 43 deletions(-) create mode 100644 src/lib/gpu/gpu-recost.ts create mode 100644 test/gpu-recost.test.mjs diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index b7d8eb8b..0a37faf2 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -267,7 +267,7 @@ const decimateSource = async ( const selection = legacy ? selectMergesLegacy(cand!, N, K, needed) : cacheOut ? - selectMergesRecosted({ splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed }) : + await selectMergesRecosted({ splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed, device }) : selectMerges(cand!, N, K, needed); selectSub.end(); diff --git a/src/lib/decimate/select-recost.ts b/src/lib/decimate/select-recost.ts index 71dd6db0..c2c0daf3 100644 --- a/src/lib/decimate/select-recost.ts +++ b/src/lib/decimate/select-recost.ts @@ -26,14 +26,24 @@ * gaussian starts queued, so the first refresh evaluates each singleton's * best edge over its neighbour rows — no candidate arrays are consumed. * + * With a device, refreshes run on the GPU ({@link GpuRecost}): the CPU + * drains commits into a log, the GPU replays it against its own structure + * copy and bulk-evaluates the queued roots, and the CPU pushes the returned + * (partner, cost) pairs with validation tokens from its own mirrors. Waves + * are serial by construction (each drain consumes the previous refresh's + * entries), so there is nothing to double-buffer. Without a device — or + * when the buffers don't fit the adapter's binding limits — the refresh + * loop runs inline with the same eval in f64. + * * decimate-source gates this path by memory budget and falls back to * {@link selectMerges} when over. - * - * Engine-free; single-threaded, pure resident-array computation, no IO. */ +import { type GraphicsDevice } from 'playcanvas'; + import { bestEdgeFor, bestOut, CACHE_STRIDE, type RecostState } from './recost-core'; import { MAX_GROUP, type SelectionResult } from './select'; +import { GpuRecost, COMMIT_LOG_STRIDE } from '../gpu/gpu-recost'; const NIL = 0xFFFFFFFF; @@ -57,16 +67,23 @@ type RecostInputs = { N: number; /** Target removal count for this generation. */ mergesNeeded: number; + /** Optional GPU device: bulk refreshes run on the wave engine when the buffers fit. */ + device?: GraphicsDevice; }; +// Max merges committed between refreshes. LOAD-BEARING: an unbounded drain +// runs up the cost curve before cheap continuation merges re-enter (measured +// −9.5 dB); refreshing more often is quality-safe, less often is not. +const WAVE = 4096; + /** * Select merges with within-generation re-costing (round-batched). * * @param inputs - See {@link RecostInputs}. * @returns The selection (same contract as {@link selectMerges}). */ -const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { - const { splatCache: SC, neighbors, D, N, mergesNeeded } = inputs; +const selectMergesRecosted = async (inputs: RecostInputs): Promise => { + const { splatCache: SC, neighbors, D, N, mergesNeeded, device } = inputs; // ---- Cluster structure (indexed by union-find root) — integers only. const parent = new Uint32Array(N); @@ -192,51 +209,87 @@ const selectMergesRecosted = (inputs: RecostInputs): SelectionResult => { } pendingCount = N; + // ---- GPU wave engine when the buffers fit (its workgroup shape needs + // the production k·MAX_GROUP); inline refresh otherwise. + const gpu = device && D * MAX_GROUP === 64 && GpuRecost.fits(device, N, D, WAVE) ? + new GpuRecost(device, N, D, MAX_GROUP, WAVE) : + undefined; + const commitLog = gpu ? new Uint32Array(WAVE * COMMIT_LOG_STRIDE) : undefined; + const outBest = gpu ? new Uint32Array(N * 2) : undefined; + const outCost = gpu ? new Float32Array(outBest!.buffer) : undefined; + // ---- Round loop. Each round commits at most WAVE merges before the // bulk refresh, so refreshed edges (notably cheap continuation merges in // redundant regions — the concentration behaviour the quality depends on) // re-enter the heap at most one wave late. An unbounded drain would spend // the budget up the cost curve before any refresh returns. - const WAVE = 4096; let removed = 0; - while (removed < mergesNeeded) { - // Drain: commit still-valid entries, up to the wave budget. - let wave = 0; - while (removed < mergesNeeded && wave < WAVE && heapPop()) { - const a = popOut.a; - if (parent[a] !== a || popOut.seq !== lastSeq[a]) continue; - const b = popOut.b; - if (parent[b] !== b || version[b] !== popOut.vb || size[a] + size[b] > MAX_GROUP) { - queueRefresh(a); - continue; - } + try { + gpu?.init(SC, neighbors); - // Commit: pure structure — splice the member chains, re-parent, - // bump the version so stale partner entries invalidate. - const keep = size[a] >= size[b] ? a : b; - const lose = keep === a ? b : a; - mNext[mTail[keep]] = mHead[lose]; - mTail[keep] = mTail[lose]; - parent[lose] = keep; - size[keep] += size[lose]; - version[keep]++; - removed++; - wave++; + while (removed < mergesNeeded) { + // Drain: commit still-valid entries, up to the wave budget. + let wave = 0; + while (removed < mergesNeeded && wave < WAVE && heapPop()) { + const a = popOut.a; + if (parent[a] !== a || popOut.seq !== lastSeq[a]) continue; + const b = popOut.b; + if (parent[b] !== b || version[b] !== popOut.vb || size[a] + size[b] > MAX_GROUP) { + queueRefresh(a); + continue; + } - queueRefresh(keep); - } - if (removed >= mergesNeeded) break; - if (pendingCount === 0 && heapSize === 0) break; + // Commit: pure structure — splice the member chains, re-parent, + // bump the version so stale partner entries invalidate. The GPU + // replays the same splice from the log (values pre-resolved so + // the replay does no reads). + const keep = size[a] >= size[b] ? a : b; + const lose = keep === a ? b : a; + if (commitLog) { + const o = wave * COMMIT_LOG_STRIDE; + commitLog[o] = lose; + commitLog[o + 1] = keep; + commitLog[o + 2] = mTail[keep]; + commitLog[o + 3] = mHead[lose]; + commitLog[o + 4] = size[keep] + size[lose]; + } + mNext[mTail[keep]] = mHead[lose]; + mTail[keep] = mTail[lose]; + parent[lose] = keep; + size[keep] += size[lose]; + version[keep]++; + removed++; + wave++; + + queueRefresh(keep); + } + if (removed >= mergesNeeded) break; + if (pendingCount === 0 && heapSize === 0) break; - // Refresh queued clusters' best edges. - for (let p = 0; p < pendingCount; p++) { - const root = pending[p]; - if (bestEdgeFor(st, root)) { - heapPush(bestOut.cost, root, bestOut.partner, lastSeq[root], bestOut.vb); + // Refresh queued clusters' best edges. + if (gpu && pendingCount > 0) { + await gpu.wave(commitLog!, wave, pending, pendingCount, outBest!); + for (let p = 0; p < pendingCount; p++) { + const partner = outBest![p * 2]; + if (partner === NIL) continue; + const cost = outCost![p * 2 + 1]; + if (Number.isNaN(cost)) continue; + const root = pending[p]; + heapPush(cost, root, partner, lastSeq[root], version[partner]); + } + } else { + for (let p = 0; p < pendingCount; p++) { + const root = pending[p]; + if (bestEdgeFor(st, root)) { + heapPush(bestOut.cost, root, bestOut.partner, lastSeq[root], bestOut.vb); + } + } } + pendingCount = 0; + round++; } - pendingCount = 0; - round++; + } finally { + gpu?.destroy(); } // ---- CSR assembly (identical contract to selectMerges). diff --git a/src/lib/gpu/gpu-recost.ts b/src/lib/gpu/gpu-recost.ts new file mode 100644 index 00000000..e17ebf1f --- /dev/null +++ b/src/lib/gpu/gpu-recost.ts @@ -0,0 +1,597 @@ +import { + BUFFERUSAGE_COPY_DST, + BUFFERUSAGE_COPY_SRC, + GraphicsDevice, + StorageBuffer +} from 'playcanvas'; + +import { makeKernel, type Kernel } from './compute-kernel'; +import { gaussianL2Wgsl } from './shaders/chunks/gaussian-l2'; + +/** + * GPU engine for the re-costed selection's wave rounds. + * + * The CPU keeps the heap, the commit decisions, and integer mirrors of the + * union-find/chains; the GPU holds its own copy of that structure plus the + * immutable splat cache and neighbour graph, and does the bulk work: after + * each drained wave the CPU uploads the wave's commit log, a `replay` kernel + * applies it (three scattered writes per entry — commits within a wave touch + * disjoint roots, so no atomics), and a `refresh` kernel re-evaluates each + * queued root's best edge with one 64-lane workgroup per root (lane = one + * (member, neighbour-slot) candidate; duplicate candidates cost nothing in + * lockstep). Results come back as 8 B per queued root: (partner, cost f32). + * + * The evaluation mirrors decimate/recost-core.ts's stateless cancelled form + * in f32 — keep them in lockstep. Costs order the heap only; nothing + * accumulates, so f32 error stays per-eval (near-tie ordering class). + * + * The splat cache and neighbour rows are split across two buffers by row + * range (a single binding would exceed the ~2 GiB per-binding ceiling near + * 34M splats); `fits` pre-flights all binding sizes by arithmetic because a + * device-side OOM escalates to a hard failure (node-device policy), not a + * catchable fallback. + */ + +const WG = 64; +const NONE = 0xFFFFFFFF; +const MAX_DIM = 65535; + +// Shared structure access for the replay/refresh kernels: parentMeta[i] = +// (parent, size-at-root); chain[i] = (head-at-root, next). +const refreshWgsl = (k: number, maxGroup: number, splitN: number) => /* wgsl */` +struct Uniforms { + pendingCount: u32, +} + +@group(0) @binding(0) var uniforms: Uniforms; +// Per-splat cache (CACHE_STRIDE 16 f32/row), split by row range at SPLIT. +@group(0) @binding(1) var cacheA: array; +@group(0) @binding(2) var cacheB: array; +// Neighbour ids (K u32/row, sentinel padded), same row split. +@group(0) @binding(3) var nbA: array; +@group(0) @binding(4) var nbB: array; +@group(0) @binding(5) var parentMeta: array; +@group(0) @binding(6) var chain: array; +@group(0) @binding(7) var pending: array; +// Per queued root: (best partner or 0xFFFFFFFF, bitcast(cost)). +@group(0) @binding(8) var outBest: array; + +const K: u32 = ${k}u; +const MAXG: u32 = ${maxGroup}u; +const SPLIT: u32 = ${splitN}u; +const NONE: u32 = 0xFFFFFFFFu; +const NIL: u32 = 0xFFFFFFFFu; +const F32_MAX: f32 = 3.4028234663852886e+38; +const CULL_QUAD: f32 = 120.0; +${gaussianL2Wgsl} + +fn cacheAt(row: u32, c: u32) -> f32 { + if (row < SPLIT) { return cacheA[row * 16u + c]; } + return cacheB[(row - SPLIT) * 16u + c]; +} + +fn nbAt(row: u32, s: u32) -> u32 { + if (row < SPLIT) { return nbA[row * K + s]; } + return nbB[(row - SPLIT) * K + s]; +} + +fn findRO(x0: u32) -> u32 { + var x = x0; + loop { + let p = parentMeta[x].x; + if (p == x) { return x; } + x = p; + } +} + +// ⟨G_a,G_b⟩ scaled by √|Σa|·√|Σb| for M = Σa+Σb and offset d, with the +// recost kernel's exponent cull (mirrors recost-core crossG). +fn crossG(sdAB: f32, m: array, d: vec3f) -> f32 { + let c00 = m[3] * m[5] - m[4] * m[4]; + let c01 = m[2] * m[4] - m[1] * m[5]; + let c02 = m[1] * m[4] - m[2] * m[3]; + let c11 = m[0] * m[5] - m[2] * m[2]; + let c12 = m[1] * m[2] - m[0] * m[4]; + let c22 = m[0] * m[3] - m[1] * m[1]; + let det = max(m[0] * c00 + m[1] * c01 + m[2] * c02, 1e-30); + let quad = (c00 * d.x * d.x + c11 * d.y * d.y + c22 * d.z * d.z + + 2.0 * (c01 * d.x * d.y + c02 * d.x * d.z + c12 * d.y * d.z)) / det; + if (!(quad < CULL_QUAD)) { return 0.0; } + return TWO_PI_1_5 * sdAB / sqrt(det) * exp(-0.5 * quad); +} + +// Raw mass-scaled aggregates of a member set (mirrors composeRaw). +struct Raw { + w: f32, + mean: vec3f, + m2: array, + bw: vec3f, +} + +fn composeRaw(members: array, count: u32) -> Raw { + var w = 0.0; + var s = vec3f(0.0); + var b = vec3f(0.0); + for (var t = 0u; t < count; t++) { + let o = members[t]; + let m = cacheAt(o, 11u); + w += m; + s += m * vec3f(cacheAt(o, 0u), cacheAt(o, 1u), cacheAt(o, 2u)); + b += m * vec3f(cacheAt(o, 12u), cacheAt(o, 13u), cacheAt(o, 14u)); + } + var r: Raw; + r.w = w; + r.mean = s / w; + r.bw = b; + for (var c = 0u; c < 6u; c++) { r.m2[c] = 0.0; } + for (var t = 0u; t < count; t++) { + let o = members[t]; + let m = cacheAt(o, 11u); + let d = vec3f(cacheAt(o, 0u), cacheAt(o, 1u), cacheAt(o, 2u)) - r.mean; + r.m2[0] += m * (cacheAt(o, 3u) + d.x * d.x); + r.m2[1] += m * (cacheAt(o, 4u) + d.x * d.y); + r.m2[2] += m * (cacheAt(o, 5u) + d.x * d.z); + r.m2[3] += m * (cacheAt(o, 6u) + d.y * d.y); + r.m2[4] += m * (cacheAt(o, 7u) + d.y * d.z); + r.m2[5] += m * (cacheAt(o, 8u) + d.z * d.z); + } + return r; +} + +// A∪B raw aggregates via the parallel-axis identity (mirrors composeUnion). +fn composeUnion(a: Raw, b: Raw) -> Raw { + var o: Raw; + o.w = a.w + b.w; + let iw = 1.0 / o.w; + o.mean = (a.w * a.mean + b.w * b.mean) * iw; + let da = a.mean - o.mean; + let db = b.mean - o.mean; + o.m2[0] = a.m2[0] + b.m2[0] + a.w * da.x * da.x + b.w * db.x * db.x; + o.m2[1] = a.m2[1] + b.m2[1] + a.w * da.x * da.y + b.w * db.x * db.y; + o.m2[2] = a.m2[2] + b.m2[2] + a.w * da.x * da.z + b.w * db.x * db.z; + o.m2[3] = a.m2[3] + b.m2[3] + a.w * da.y * da.y + b.w * db.y * db.y; + o.m2[4] = a.m2[4] + b.m2[4] + a.w * da.y * da.z + b.w * db.y * db.z; + o.m2[5] = a.m2[5] + b.m2[5] + a.w * da.z * da.z + b.w * db.z * db.z; + o.bw = a.bw + b.bw; + return o; +} + +// Finished merged-Gaussian quantities (mirrors finishComp). +struct Fin { + sm: array, + mean: vec3f, + sd: f32, + alpha: f32, + bc: vec3f, + selfM: f32, +} + +fn finish(r: Raw) -> Fin { + var f: Fin; + let iw = 1.0 / r.w; + f.sm[0] = r.m2[0] * iw + EPS_COV; + f.sm[1] = r.m2[1] * iw; + f.sm[2] = r.m2[2] * iw; + f.sm[3] = r.m2[3] * iw + EPS_COV; + f.sm[4] = r.m2[4] * iw; + f.sm[5] = r.m2[5] * iw + EPS_COV; + f.mean = r.mean; + let detm = max( + f.sm[0] * (f.sm[3] * f.sm[5] - f.sm[4] * f.sm[4]) - f.sm[1] * (f.sm[1] * f.sm[5] - f.sm[4] * f.sm[2]) + f.sm[2] * (f.sm[1] * f.sm[4] - f.sm[3] * f.sm[2]), + 1e-30 + ); + f.sd = sqrt(detm); + let e = eig3(f.sm); + let s0 = sqrt(max(e.x, 1e-18)); + let s1 = sqrt(max(e.y, 1e-18)); + let s2 = sqrt(max(e.z, 1e-18)); + f.alpha = min(1.0, r.w / max(ellipsoidArea(s0, s1, s2), 1e-30)); + f.bc = r.bw * iw; + f.selfM = f.alpha * f.alpha * dot(f.bc, f.bc) * PI_1_5 * f.sd; + return f; +} + +// memfm(C over members) = Σ ⟨f_k, f_C⟩ (mirrors memfm). +fn memfmOf(members: array, count: u32, f: Fin) -> f32 { + var acc = 0.0; + for (var t = 0u; t < count; t++) { + let o = members[t]; + let wgt = cacheAt(o, 10u) * f.alpha * + dot(vec3f(cacheAt(o, 12u), cacheAt(o, 13u), cacheAt(o, 14u)), f.bc); + if (wgt == 0.0) { continue; } + let m = array( + cacheAt(o, 3u) + f.sm[0], cacheAt(o, 4u) + f.sm[1], cacheAt(o, 5u) + f.sm[2], + cacheAt(o, 6u) + f.sm[3], cacheAt(o, 7u) + f.sm[4], cacheAt(o, 8u) + f.sm[5] + ); + let d = vec3f(cacheAt(o, 0u), cacheAt(o, 1u), cacheAt(o, 2u)) - f.mean; + acc += wgt * crossG(cacheAt(o, 9u) * f.sd, m, d); + } + return acc; +} + +// A singleton's self product ⟨f, f⟩ (mirrors selfRow). +fn selfRow(o: u32) -> f32 { + let a = cacheAt(o, 10u); + return a * a * cacheAt(o, 15u) * PI_1_5 * cacheAt(o, 9u); +} + +// One side's term of the cancelled cost form (mirrors sideTerm). +fn sideTerm(members: array, count: u32, r: Raw) -> f32 { + if (count < 2u) { return selfRow(members[0]); } + let f = finish(r); + return 2.0 * memfmOf(members, count, f) - f.selfM; +} + +// scross(A,B) = Σ_{a∈A,b∈B} ⟨f_a, f_b⟩ with distance culling (mirrors scrossPairs). +fn scrossPairs(am: array, na: u32, bm: array, nb: u32) -> f32 { + var acc = 0.0; + for (var u = 0u; u < nb; u++) { + let ob = bm[u]; + let bp = vec3f(cacheAt(ob, 0u), cacheAt(ob, 1u), cacheAt(ob, 2u)); + let trb = cacheAt(ob, 3u) + cacheAt(ob, 6u) + cacheAt(ob, 8u); + let alb = cacheAt(ob, 10u); + let sdb = cacheAt(ob, 9u); + let cb = vec3f(cacheAt(ob, 12u), cacheAt(ob, 13u), cacheAt(ob, 14u)); + for (var t = 0u; t < na; t++) { + let oa = am[t]; + let d = vec3f(cacheAt(oa, 0u), cacheAt(oa, 1u), cacheAt(oa, 2u)) - bp; + let d2 = dot(d, d); + if (d2 > CULL_QUAD * (cacheAt(oa, 3u) + cacheAt(oa, 6u) + cacheAt(oa, 8u) + trb)) { continue; } + let wgt = cacheAt(oa, 10u) * alb * + dot(vec3f(cacheAt(oa, 12u), cacheAt(oa, 13u), cacheAt(oa, 14u)), cb); + if (wgt == 0.0) { continue; } + let m = array( + cacheAt(oa, 3u) + cacheAt(ob, 3u), cacheAt(oa, 4u) + cacheAt(ob, 4u), cacheAt(oa, 5u) + cacheAt(ob, 5u), + cacheAt(oa, 6u) + cacheAt(ob, 6u), cacheAt(oa, 7u) + cacheAt(ob, 7u), cacheAt(oa, 8u) + cacheAt(ob, 8u) + ); + acc += wgt * crossG(cacheAt(oa, 9u) * sdb, m, d); + } + } + return acc; +} + +var wgAbort: u32; +var wgRoot: u32; +var wgSize: u32; +var wgCount: u32; +var wgMembers: array; +var wgRawW: f32; +var wgRawMean: vec3f; +var wgRawM2: array; +var wgRawBw: vec3f; +var wgATerm: f32; +var redCost: array; +var redPartner: array; + +@compute @workgroup_size(${WG}) +fn main(@builtin(workgroup_id) wgid: vec3u, @builtin(local_invocation_id) lid3: vec3u) { + let pIdx = wgid.y * ${MAX_DIM}u + wgid.x; + let lid = lid3.x; + + // Lane 0 resolves the root and hoists the candidate-independent A side. + if (lid == 0u) { + if (pIdx >= uniforms.pendingCount) { + wgAbort = 1u; + } else { + let root = pending[pIdx]; + if (parentMeta[root].x != root) { + // Stale queued root (absorbed since queuing) — no result. + outBest[pIdx] = vec2u(NONE, 0u); + wgAbort = 1u; + } else { + wgAbort = 0u; + wgRoot = root; + wgSize = parentMeta[root].y; + var cnt = 0u; + var m = chain[root].x; + while (m != NIL && cnt < MAXG) { + wgMembers[cnt] = m; + cnt++; + m = chain[m].y; + } + wgCount = cnt; + let raw = composeRaw(wgMembers, cnt); + wgRawW = raw.w; + wgRawMean = raw.mean; + for (var c = 0u; c < 6u; c++) { wgRawM2[c] = raw.m2[c]; } + wgRawBw = raw.bw; + wgATerm = sideTerm(wgMembers, cnt, raw); + } + } + } + workgroupBarrier(); + // No early return — the reduction barriers below must stay in uniform + // control flow, so aborted workgroups just run with every lane inactive. + let aborted = wgAbort == 1u; + + // One lane per (member, neighbour-slot) candidate. Duplicate candidate + // roots are evaluated redundantly (lockstep makes them free); dedup would + // only shift tie order. + var cost = F32_MAX; + var partner = NONE; + let mIdx = lid / K; + let slot = lid % K; + if (!aborted && mIdx < wgCount) { + let cand = nbAt(wgMembers[mIdx], slot); + if (cand != NONE) { + let r = findRO(cand); + if (r != wgRoot && wgSize + parentMeta[r].y <= MAXG) { + // Gather B's members and rebuild A's raw aggregates locally. + var bm: array; + var bCount = 0u; + var bmm = chain[r].x; + while (bmm != NIL && bCount < MAXG) { + bm[bCount] = bmm; + bCount++; + bmm = chain[bmm].y; + } + var aRaw: Raw; + aRaw.w = wgRawW; + aRaw.mean = wgRawMean; + for (var c = 0u; c < 6u; c++) { aRaw.m2[c] = wgRawM2[c]; } + aRaw.bw = wgRawBw; + + let bRaw = composeRaw(bm, bCount); + let bTerm = sideTerm(bm, bCount, bRaw); + let fAB = finish(composeUnion(aRaw, bRaw)); + let memfmAB = memfmOf(wgMembers, wgCount, fAB) + memfmOf(bm, bCount, fAB); + let scross = scrossPairs(wgMembers, wgCount, bm, bCount); + let dbc = aRaw.bw / aRaw.w - bRaw.bw / bRaw.w; + let c = (2.0 * scross - 2.0 * memfmAB + fAB.selfM + wgATerm + bTerm) + + COLOR_WEIGHT * dot(dbc, dbc); + // NaN loses every comparison → stays unselected (fail-loud: + // an all-NaN scene produces no pushes and the caller throws). + if (c < F32_MAX) { + cost = c; + partner = r; + } + } + } + } + + // Lexicographic (cost, partner) min-reduction — deterministic under any + // lane scheduling. + redCost[lid] = cost; + redPartner[lid] = partner; + workgroupBarrier(); + for (var s = ${WG >> 1}u; s > 0u; s >>= 1u) { + if (lid < s) { + let c2 = redCost[lid + s]; + let p2 = redPartner[lid + s]; + if (c2 < redCost[lid] || (c2 == redCost[lid] && p2 < redPartner[lid])) { + redCost[lid] = c2; + redPartner[lid] = p2; + } + } + workgroupBarrier(); + } + if (lid == 0u && !aborted) { + var p = redPartner[0]; + if (redCost[0] == F32_MAX) { p = NONE; } + outBest[pIdx] = vec2u(p, bitcast(redCost[0])); + } +} +`; + +// Apply a wave's commit log: per entry (lose, keep, tailKeep, headLose, +// newSize) — three scattered writes, no reads. Race-free because validation +// guarantees each root commits at most once per wave (entries touching a +// committed root fail their version/seq/parent checks), so writes across +// entries target disjoint words. +const replayWgsl = () => /* wgsl */` +struct Uniforms { + commitCount: u32, +} + +@group(0) @binding(0) var uniforms: Uniforms; +@group(0) @binding(1) var commitLog: array; +@group(0) @binding(2) var parentMeta: array; +@group(0) @binding(3) var chain: array; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3u) { + let e = gid.x; + if (e >= uniforms.commitCount) { return; } + let o = e * 5u; + let lose = commitLog[o]; + let keep = commitLog[o + 1u]; + let tailKeep = commitLog[o + 2u]; + let headLose = commitLog[o + 3u]; + let newSize = commitLog[o + 4u]; + chain[tailKeep].y = headLose; + parentMeta[lose].x = keep; + parentMeta[keep].y = newSize; +} +`; + +// Initialize the structure buffers: every splat a singleton root. +const initWgsl = () => /* wgsl */` +struct Uniforms { + count: u32, +} + +@group(0) @binding(0) var uniforms: Uniforms; +@group(0) @binding(1) var parentMeta: array; +@group(0) @binding(2) var chain: array; + +@compute @workgroup_size(256) +fn main(@builtin(global_invocation_id) gid: vec3u) { + let i = gid.y * ${MAX_DIM * 256}u + gid.x; + if (i >= uniforms.count) { return; } + parentMeta[i] = vec2u(i, 1u); + chain[i] = vec2u(i, 0xFFFFFFFFu); +} +`; + +/** A wave's commit log entry width in u32s. */ +const COMMIT_LOG_STRIDE = 5; + +class GpuRecost { + /** + * Upload the immutable inputs and initialize the structure buffers. + * Call once before the first wave. + */ + init: (cache: Float32Array, neighbors: Uint32Array) => void; + /** + * Run one wave: replay `commitCount` log entries, refresh + * `pendingCount` queued roots, and read back (partner, cost) pairs + * into `outBest` (2 u32 per root; cost is a bitcast f32). + */ + wave: ( + commitLog: Uint32Array, + commitCount: number, + pending: Uint32Array, + pendingCount: number, + outBest: Uint32Array + ) => Promise; + destroy: () => void; + + /** + * Whether the wave engine's buffers fit the device's binding limits for + * `n` splats (pre-flight by arithmetic — a device-side OOM is a hard + * failure, not a catchable fallback). + * + * @param device - PlayCanvas GraphicsDevice (WebGPU). + * @param n - Generation splat count. + * @param k - Neighbours per splat. + * @param wave - Max commits per wave. + * @returns True when all bindings fit. + */ + static fits(device: GraphicsDevice, n: number, k: number, wave: number): boolean { + if (n < 1024) return false; // inline is instant below this + const limits = (device as any).limits; + const maxBinding = Math.min( + typeof limits?.maxStorageBufferBindingSize === 'number' ? limits.maxStorageBufferBindingSize : 128 * 2 ** 20, + typeof limits?.maxBufferSize === 'number' ? limits.maxBufferSize : 256 * 2 ** 20 + ); + const splitN = Math.ceil(n / 2); + return splitN * 16 * 4 <= maxBinding && // cacheA/B + splitN * k * 4 <= maxBinding && // nbA/B + n * 8 <= maxBinding && // parentMeta/chain/outBest + wave * COMMIT_LOG_STRIDE * 4 <= maxBinding; + } + + /** + * @param device - PlayCanvas GraphicsDevice (WebGPU). + * @param n - Generation splat count. + * @param k - Neighbours per splat (the refresh workgroup is k·maxGroup lanes). + * @param maxGroup - Group size cap. + * @param wave - Max commits per wave (commit log capacity). + */ + constructor(device: GraphicsDevice, n: number, k: number, maxGroup: number, wave: number) { + if (k * maxGroup !== WG) { + throw new Error(`GpuRecost: k·maxGroup must be ${WG} (got ${k}·${maxGroup})`); + } + const splitN = Math.ceil(n / 2); + + const cacheABuf = new StorageBuffer(device, splitN * 16 * 4, BUFFERUSAGE_COPY_DST); + const cacheBBuf = new StorageBuffer(device, Math.max(n - splitN, 1) * 16 * 4, BUFFERUSAGE_COPY_DST); + const nbABuf = new StorageBuffer(device, splitN * k * 4, BUFFERUSAGE_COPY_DST); + const nbBBuf = new StorageBuffer(device, Math.max(n - splitN, 1) * k * 4, BUFFERUSAGE_COPY_DST); + const parentMetaBuf = new StorageBuffer(device, n * 8, BUFFERUSAGE_COPY_DST); + const chainBuf = new StorageBuffer(device, n * 8, BUFFERUSAGE_COPY_DST); + const pendingBuf = new StorageBuffer(device, n * 4, BUFFERUSAGE_COPY_DST); + const outBestBuf = new StorageBuffer(device, n * 8, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST); + const commitLogBuf = new StorageBuffer(device, wave * COMMIT_LOG_STRIDE * 4, BUFFERUSAGE_COPY_DST); + + const initKernel = makeKernel(device, 'recost-init', initWgsl(), ['count'], [ + ['parentMeta', false], + ['chain', false] + ]); + initKernel.compute.setParameter('parentMeta', parentMetaBuf); + initKernel.compute.setParameter('chain', chainBuf); + + const replayKernel = makeKernel(device, 'recost-replay', replayWgsl(), ['commitCount'], [ + ['commitLog', true], + ['parentMeta', false], + ['chain', false] + ]); + replayKernel.compute.setParameter('commitLog', commitLogBuf); + replayKernel.compute.setParameter('parentMeta', parentMetaBuf); + replayKernel.compute.setParameter('chain', chainBuf); + + const refreshKernel = makeKernel(device, 'recost-refresh', refreshWgsl(k, maxGroup, splitN), ['pendingCount'], [ + ['cacheA', true], + ['cacheB', true], + ['nbA', true], + ['nbB', true], + ['parentMeta', true], + ['chain', true], + ['pending', true], + ['outBest', false] + ]); + refreshKernel.compute.setParameter('cacheA', cacheABuf); + refreshKernel.compute.setParameter('cacheB', cacheBBuf); + refreshKernel.compute.setParameter('nbA', nbABuf); + refreshKernel.compute.setParameter('nbB', nbBBuf); + refreshKernel.compute.setParameter('parentMeta', parentMetaBuf); + refreshKernel.compute.setParameter('chain', chainBuf); + refreshKernel.compute.setParameter('pending', pendingBuf); + refreshKernel.compute.setParameter('outBest', outBestBuf); + + // Chunked uploads keep Dawn's staging allocations bounded. + const CHUNK = 1 << 24; + const writeChunked = (buf: StorageBuffer, data: Float32Array | Uint32Array, srcBase: number, count: number) => { + for (let off = 0; off < count; off += CHUNK) { + const c = Math.min(CHUNK, count - off); + buf.write(off * 4, data, srcBase + off, c); + } + }; + + this.init = (cache: Float32Array, neighbors: Uint32Array) => { + writeChunked(cacheABuf, cache, 0, splitN * 16); + writeChunked(cacheBBuf, cache, splitN * 16, (n - splitN) * 16); + writeChunked(nbABuf, neighbors, 0, splitN * k); + writeChunked(nbBBuf, neighbors, splitN * k, (n - splitN) * k); + + const groups = Math.ceil(n / 256); + initKernel.compute.setParameter('count', n); + initKernel.compute.setupDispatch(Math.min(groups, MAX_DIM), Math.ceil(groups / MAX_DIM)); + device.computeDispatch([initKernel.compute], 'recost-init'); + }; + + this.wave = async ( + commitLog: Uint32Array, + commitCount: number, + pending: Uint32Array, + pendingCount: number, + outBest: Uint32Array + ) => { + const computes = []; + if (commitCount > 0) { + commitLogBuf.write(0, commitLog, 0, commitCount * COMMIT_LOG_STRIDE); + replayKernel.compute.setParameter('commitCount', commitCount); + replayKernel.compute.setupDispatch(Math.ceil(commitCount / 64)); + computes.push(replayKernel.compute); + } + pendingBuf.write(0, pending, 0, pendingCount); + refreshKernel.compute.setParameter('pendingCount', pendingCount); + refreshKernel.compute.setupDispatch( + Math.min(pendingCount, MAX_DIM), + Math.ceil(pendingCount / MAX_DIM) + ); + computes.push(refreshKernel.compute); + device.computeDispatch(computes, 'recost-wave'); + + // Blocking readback — also the wave's submit boundary. + await outBestBuf.read(0, pendingCount * 8, outBest, true); + }; + + this.destroy = () => { + cacheABuf.destroy(); + cacheBBuf.destroy(); + nbABuf.destroy(); + nbBBuf.destroy(); + parentMetaBuf.destroy(); + chainBuf.destroy(); + pendingBuf.destroy(); + outBestBuf.destroy(); + commitLogBuf.destroy(); + initKernel.destroy(); + replayKernel.destroy(); + refreshKernel.destroy(); + }; + } +} + +export { GpuRecost, COMMIT_LOG_STRIDE }; diff --git a/test/decimate-recost.test.mjs b/test/decimate-recost.test.mjs index e7eba702..b761928f 100644 --- a/test/decimate-recost.test.mjs +++ b/test/decimate-recost.test.mjs @@ -175,12 +175,12 @@ describe('stateless re-costed eval', () => { }); describe('selectMergesRecosted', () => { - it('hits the target with a well-formed capped CSR', () => { + it('hits the target with a well-formed capped CSR', async () => { const n = 256, k = 8; const cache = makeCache(n, 99); const neighbors = bruteNeighbors(cache, n, k); const needed = n >> 1; - const sel = selectMergesRecosted({ splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); + const sel = await selectMergesRecosted({ splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); assert.strictEqual(sel.removed, needed); assert.strictEqual(sel.groupOffsets[sel.mergedGroups], sel.groupMembers.length); @@ -201,7 +201,7 @@ describe('selectMergesRecosted', () => { assert.strictEqual(members - sel.mergedGroups, needed); }); - it('concentrates coincident splats into full groups (refresh continuation)', () => { + it('concentrates coincident splats into full groups (refresh continuation)', async () => { const n = 64, k = 16; // All splats identical and coincident. const pos = new Float32Array(n * 3).fill(0.5); @@ -226,7 +226,7 @@ describe('selectMergesRecosted', () => { } } const needed = n - (n / MAX_GROUP); // 48: reachable only by filling groups to the cap - const sel = selectMergesRecosted({ splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); + const sel = await selectMergesRecosted({ splatCache: cache, neighbors, D: k, N: n, mergesNeeded: needed }); // Disjoint pairing alone tops out at n/2 = 32 removals; anything well // beyond that requires continuation merges via refresh re-entry (the diff --git a/test/gpu-recost.test.mjs b/test/gpu-recost.test.mjs new file mode 100644 index 00000000..7634be4f --- /dev/null +++ b/test/gpu-recost.test.mjs @@ -0,0 +1,213 @@ +/** + * GPU wave-engine acceptance (GPU required; suites skip without a WebGPU + * adapter): + * + * 1. Refresh-kernel parity — wave-0 (singleton) and post-commit + * (multi-member, replayed through the commit log) best-edge costs must + * match the CPU f64 stateless eval within the study's 1e-3 relative gate + * for ≥99% of roots. + * 2. Merge-set equality — on a fixture with well-separated costs (no + * near-ties), the GPU and inline paths must produce identical selections. + */ + +import assert from 'node:assert'; +import { after, before, describe, it } from 'node:test'; + +import { buildSplatCache, CACHE_STRIDE } from '../src/lib/decimate/edge-cost-cpu.js'; +import { bestEdgeFor, bestOut } from '../src/lib/decimate/recost-core.js'; +import { MAX_GROUP } from '../src/lib/decimate/select.js'; +import { selectMergesRecosted } from '../src/lib/decimate/select-recost.js'; +import { GpuRecost, COMMIT_LOG_STRIDE } from '../src/lib/gpu/gpu-recost.js'; + +const NIL = 0xFFFFFFFF; +const K = 16; + +let device = null; + +before(async () => { + try { + const { createDevice } = await import('../src/cli/node-device.js'); + device = await createDevice(); + } catch { + device = null; + } +}); + +after(() => { + device?.destroy?.(); +}); + +const mulberry = (seed) => { + let t = seed >>> 0; + return () => { + t += 0x6d2b79f5; + let r = Math.imul(t ^ (t >>> 15), t | 1); + r ^= r + Math.imul(r ^ (r >>> 7), r | 61); + return ((r ^ (r >>> 14)) >>> 0) / 4294967296; + }; +}; + +// Random splats; size spread makes costs well-separated (few near-ties). +const makeCache = (n, seed) => { + const rand = mulberry(seed); + const pos = new Float32Array(n * 3); + const geo = new Float32Array(n * 8); + const color = new Float32Array(n * 3); + for (let i = 0; i < n; i++) { + pos[i * 3] = rand() * 4; + pos[i * 3 + 1] = rand() * 4; + pos[i * 3 + 2] = rand() * 4; + geo[i * 8] = 1; + geo[i * 8 + 4] = Math.log(0.02 + rand() * 0.3); + geo[i * 8 + 5] = Math.log(0.02 + rand() * 0.3); + geo[i * 8 + 6] = Math.log(0.02 + rand() * 0.3); + geo[i * 8 + 7] = rand() * 4 - 2; + color[i * 3] = rand() - 0.5; + color[i * 3 + 1] = rand() - 0.5; + color[i * 3 + 2] = rand() - 0.5; + } + const cache = new Float32Array(n * CACHE_STRIDE); + buildSplatCache({ pos, geo, color, colorDim: 3 }, cache); + return cache; +}; + +const bruteNeighbors = (cache, n, k) => { + const nb = new Uint32Array(n * k).fill(NIL); + const d2 = (a, b) => { + const oa = a * CACHE_STRIDE, ob = b * CACHE_STRIDE; + return (cache[oa] - cache[ob]) ** 2 + (cache[oa + 1] - cache[ob + 1]) ** 2 + (cache[oa + 2] - cache[ob + 2]) ** 2; + }; + for (let i = 0; i < n; i++) { + const ids = Array.from({ length: n }, (_, j) => j) + .filter(j => j !== i) + .sort((a, b) => d2(i, a) - d2(i, b)) + .slice(0, k); + for (let s = 0; s < ids.length; s++) nb[i * k + s] = ids[s]; + } + return nb; +}; + +const makeState = (cache, neighbors, n) => { + const parent = new Uint32Array(n); + const size = new Uint32Array(n).fill(1); + const version = new Uint32Array(n).fill(1); + const mHead = new Uint32Array(n); + const mNext = new Uint32Array(n).fill(NIL); + const mTail = new Uint32Array(n); + for (let i = 0; i < n; i++) { + parent[i] = i; mHead[i] = i; mTail[i] = i; + } + return { + st: { SC: cache, cands: neighbors, D: K, N: n, maxGroup: MAX_GROUP, parent, size, version, mHead, mNext }, + mTail + }; +}; + +// Compare GPU (partner, cost f32) rows against CPU bestEdgeFor over `roots`. +const compareRefresh = (state, roots, out) => { + const outCost = new Float32Array(out.buffer); + let compared = 0, ok = 0, disagreePartner = 0; + for (let p = 0; p < roots.length; p++) { + const root = roots[p]; + const gPartner = out[p * 2]; + const has = bestEdgeFor(state.st, root); + if (!has || gPartner === NIL) { + // Both must agree there is no legal candidate. + assert.strictEqual(has, gPartner !== NIL, `root ${root}: candidate existence disagrees`); + continue; + } + compared++; + const gCost = outCost[p * 2 + 1]; + const rel = Math.abs(gCost - bestOut.cost) / Math.max(1e-12, Math.abs(bestOut.cost)); + if (rel < 1e-3) ok++; + if (gPartner !== bestOut.partner) disagreePartner++; + } + return { compared, ok, disagreePartner }; +}; + +describe('GpuRecost refresh parity', () => { + it('wave-0 singleton and post-commit cluster costs match CPU within 1e-3', async (t) => { + if (!device) return t.skip('no WebGPU adapter available'); + + const n = 2048; + const cache = makeCache(n, 1234); + const neighbors = bruteNeighbors(cache, n, K); + const state = makeState(cache, neighbors, n); + + const gpu = new GpuRecost(device, n, K, MAX_GROUP, 4096); + try { + gpu.init(cache, neighbors); + + // Wave 0: all singletons. + const pending = new Uint32Array(n); + for (let i = 0; i < n; i++) pending[i] = i; + const out = new Uint32Array(n * 2); + await gpu.wave(new Uint32Array(0), 0, pending, n, out); + const w0 = compareRefresh(state, pending, out); + assert.ok(w0.compared > n * 0.9, `wave-0 compared ${w0.compared}`); + assert.ok(w0.ok / w0.compared >= 0.99, `wave-0 parity ${w0.ok}/${w0.compared}`); + + // Commit ~300 merges on the CPU mirrors, replay the same log on + // the GPU, then compare multi-member refreshes. Each root may be + // touched at most once per wave — the production drain guarantees + // this via its version checks, and the parallel replay's + // race-freedom depends on it. + const { st, mTail } = state; + const log = new Uint32Array(4096 * COMMIT_LOG_STRIDE); + let commits = 0; + const touched = []; + const touchedSet = new Set(); + for (let a = 0; a < n && commits < 300; a += 3) { + if (st.parent[a] !== a || touchedSet.has(a)) continue; + if (!bestEdgeFor(st, a)) continue; + const b = bestOut.partner; + if (st.parent[b] !== b || touchedSet.has(b) || st.size[a] + st.size[b] > MAX_GROUP) continue; + touchedSet.add(a); + touchedSet.add(b); + const keep = st.size[a] >= st.size[b] ? a : b; + const lose = keep === a ? b : a; + const o = commits * COMMIT_LOG_STRIDE; + log[o] = lose; + log[o + 1] = keep; + log[o + 2] = mTail[keep]; + log[o + 3] = st.mHead[lose]; + log[o + 4] = st.size[keep] + st.size[lose]; + st.mNext[mTail[keep]] = st.mHead[lose]; + mTail[keep] = mTail[lose]; + st.parent[lose] = keep; + st.size[keep] += st.size[lose]; + st.version[keep]++; + touched.push(keep); + commits++; + } + assert.ok(commits >= 200, `committed ${commits}`); + + const roots = Uint32Array.from(touched); + const out2 = new Uint32Array(roots.length * 2); + await gpu.wave(log, commits, roots, roots.length, out2); + const w1 = compareRefresh(state, roots, out2); + assert.ok(w1.compared > roots.length * 0.5, `post-commit compared ${w1.compared}`); + assert.ok(w1.ok / w1.compared >= 0.99, `post-commit parity ${w1.ok}/${w1.compared}`); + } finally { + gpu.destroy(); + } + }); +}); + +describe('GpuRecost selection equality', () => { + it('GPU and inline paths produce identical selections on distinct costs', async (t) => { + if (!device) return t.skip('no WebGPU adapter available'); + + const n = 2048; + const cache = makeCache(n, 777); + const neighbors = bruteNeighbors(cache, n, K); + const needed = n >> 1; + + const inline = await selectMergesRecosted({ splatCache: cache, neighbors, D: K, N: n, mergesNeeded: needed }); + const gpu = await selectMergesRecosted({ splatCache: cache, neighbors, D: K, N: n, mergesNeeded: needed, device }); + + assert.strictEqual(gpu.removed, inline.removed); + assert.strictEqual(gpu.mergedGroups, inline.mergedGroups); + assert.deepStrictEqual(Array.from(gpu.memberGroup), Array.from(inline.memberGroup)); + }); +}); From 75be79ab574c94eab58b735235e6ec50bc513c89 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 18:06:47 +0100 Subject: [PATCH 10/19] latest --- src/lib/decimate/decimate-source.ts | 4 +- src/lib/decimate/knn-blocks.ts | 332 ----------------------- src/lib/decimate/knn-core.ts | 148 +++++++--- src/lib/decimate/priority-legacy.ts | 121 +++++---- src/lib/decimate/priority.ts | 255 ++++++++++++----- src/lib/gpu/gpu-knn.ts | 407 +++++++++++++++------------- src/lib/workers/tasks.ts | 76 ++++-- test/decimate-knn.test.mjs | 227 ++++++++++------ 8 files changed, 791 insertions(+), 779 deletions(-) delete mode 100644 src/lib/decimate/knn-blocks.ts diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index 0a37faf2..cb3e7d6c 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -4,7 +4,7 @@ import { type GraphicsDevice } from 'playcanvas'; import { createBlockProducerSource } from './block-producer'; import { mergeStream } from './merge-stream'; import { kdPartition, coherenceRuns, type ResidentPositions } from './partition'; -import { runPriorityPass, HALO_CAP, type CandidateArrays } from './priority'; +import { runPriorityPass, VIEW_GROW, type CandidateArrays } from './priority'; import { runPriorityPassLegacy } from './priority-legacy'; import { selectMerges } from './select'; import { selectMergesLegacy } from './select-legacy'; @@ -200,7 +200,7 @@ const decimateSource = async ( const bindingLimit = (device as unknown as { limits?: { maxStorageBufferBindingSize?: number } } | undefined) ?.limits?.maxStorageBufferBindingSize; if (typeof bindingLimit === 'number') { - const largestBinding = (bs: number) => bs * (1 + HALO_CAP) * SPLAT_STRIDE * 4; + const largestBinding = (bs: number) => Math.ceil(bs * VIEW_GROW) * SPLAT_STRIDE * 4; while (blockSize > (1 << 16) && largestBinding(blockSize) > bindingLimit) { blockSize >>= 1; } diff --git a/src/lib/decimate/knn-blocks.ts b/src/lib/decimate/knn-blocks.ts deleted file mode 100644 index c8874333..00000000 --- a/src/lib/decimate/knn-blocks.ts +++ /dev/null @@ -1,332 +0,0 @@ -import { knnQueryBlock, KNN_SENTINEL } from './knn-core'; -import { type BlockRange, type ResidentPositions } from './partition'; - -/** - * A block's local point set: owned gaussians first (in the block's sorted - * owned order), then halo members from neighbouring blocks. `ids` maps local - * index → global gaussian index; `positions` is interleaved xyz; `h` is the - * halo radius the set was collected with — `-Infinity` when no covering halo - * fits the cap (halo empty, every owned query takes the exact requery). - */ -type BlockLocals = { - ids: Uint32Array; - ownedCount: number; - positions: Float32Array; - h: number; -}; - -/** Local-slot marker: neighbour was fixed by verification; resolve via its global id. */ -const KNN_FIXED = 0xFFFFFFFE; - -// Density-based halo radius: haloFactor × the Poisson estimate of the k-NN -// radius from the block's AABB volume and count. Degenerate blocks (planar / -// coincident points) drive the estimate toward 0 — that's fine, the -// verification pass is the correctness backstop; h is only an efficiency hint. -const haloRadius = (block: BlockRange, k: number, haloFactor: number): number => { - const nOwned = block.end - block.start; - if (nOwned === 0) return 0; - const ex = Math.max(block.aabb[3] - block.aabb[0], 1e-12); - const ey = Math.max(block.aabb[4] - block.aabb[1], 1e-12); - const ez = Math.max(block.aabb[5] - block.aabb[2], 1e-12); - const lambda = nOwned / (ex * ey * ez); - const rk = Math.cbrt((k * 3) / (4 * Math.PI * lambda)); - return haloFactor * rk; -}; - -// Squared distance from a point to an AABB ([minx..z, maxx..z]). -const pointAabbDist2 = (px: number, py: number, pz: number, aabb: Float32Array): number => { - const dx = Math.max(0, aabb[0] - px, px - aabb[3]); - const dy = Math.max(0, aabb[1] - py, py - aabb[4]); - const dz = Math.max(0, aabb[2] - pz, pz - aabb[5]); - return dx * dx + dy * dy + dz * dz; -}; - -// Squared distance between two AABBs (0 when overlapping). -const aabbAabbDist2 = (a: Float32Array, b: Float32Array): number => { - let d2 = 0; - for (let c = 0; c < 3; c++) { - const gap = Math.max(0, b[c] - a[3 + c], a[c] - b[3 + c]); - d2 += gap * gap; - } - return d2; -}; - -/** - * Collect a block's local point set: its owned gaussians plus a halo of - * points from neighbouring blocks within `h` of the block AABB. - * - * @param pos - Resident positions. - * @param order - Partition index array. - * @param blocks - All block ranges. - * @param blockIdx - Which block to collect. - * @param k - Neighbours per query (drives the halo radius estimate). - * @param haloFactor - Multiplier on the k-NN radius estimate. - * @param haloCap - Maximum halo size as a multiple of the owned count (buffer-sizing bound; `h` shrinks until the full halo fits the cap — members are never dropped). - * @returns The block's locals. - */ -const collectBlock = ( - pos: ResidentPositions, - order: Uint32Array, - blocks: BlockRange[], - blockIdx: number, - k: number, - haloFactor: number, - haloCap = 1 -): BlockLocals => { - const block = blocks[blockIdx]; - const nOwned = block.end - block.start; - const maxHalo = Math.ceil(nOwned * haloCap); - - // The shrink loop only asks whether the halo exceeds the cap — bail as - // soon as that is known rather than counting the whole scene. - const countHalo = (h2: number): number => { - let count = 0; - for (let b2 = 0; b2 < blocks.length; b2++) { - if (b2 === blockIdx) continue; - const other = blocks[b2]; - if (aabbAabbDist2(block.aabb, other.aabb) > h2) continue; - for (let i = other.start; i < other.end; i++) { - const g = order[i]; - if (pointAabbDist2(pos.x[g], pos.y[g], pos.z[g], block.aabb) <= h2) { - if (++count > maxHalo) return count; - } - } - } - return count; - }; - - // The verification rule (`d_k ≤ depth + h`) is only sound when the halo - // FULLY covers AABB ⊕ h — so the size cap must never truncate members. - // Instead, shrink h until the full halo fits the cap; points beyond the - // reduced h are then legitimately outside the covered region and boundary - // queries fall through to the brute-force requery. - let h = haloRadius(block, k, haloFactor); - let fits = false; - for (let iter = 0; iter < 40 && h > 0; iter++) { - if (countHalo(h * h) <= maxHalo) { - fits = true; - break; - } - h *= 0.7; - if (h < 1e-12) h = 0; - } - if (!fits) { - if (countHalo(0) > maxHalo) { - // No h can fit the cap: other blocks' points sit INSIDE this - // block's AABB (e.g. an outlier-residual block enveloping the - // core), so the covering guarantee is unobtainable at any radius. - // Collect no halo and disable the guarantee (h = -Infinity): - // every owned query then takes the best-first requery, which is - // exact by construction. - const ids = new Uint32Array(nOwned); - const positions = new Float32Array(nOwned * 3); - for (let i = 0; i < nOwned; i++) { - const g = order[block.start + i]; - ids[i] = g; - positions[i * 3] = pos.x[g]; - positions[i * 3 + 1] = pos.y[g]; - positions[i * 3 + 2] = pos.z[g]; - } - return { ids, ownedCount: nOwned, positions, h: -Infinity }; - } - // The shrink iterations ran out, but a zero-radius halo fits. - h = 0; - } - const h2 = h * h; - - // Halo membership: points of other blocks within h of this block's AABB. - const haloIds: number[] = []; - for (let b2 = 0; b2 < blocks.length; b2++) { - if (b2 === blockIdx) continue; - const other = blocks[b2]; - if (aabbAabbDist2(block.aabb, other.aabb) > h2) continue; - for (let i = other.start; i < other.end; i++) { - const g = order[i]; - if (pointAabbDist2(pos.x[g], pos.y[g], pos.z[g], block.aabb) <= h2) { - haloIds.push(g); - } - } - } - - const n = nOwned + haloIds.length; - const ids = new Uint32Array(n); - const positions = new Float32Array(n * 3); - for (let i = 0; i < nOwned; i++) { - const g = order[block.start + i]; - ids[i] = g; - positions[i * 3] = pos.x[g]; - positions[i * 3 + 1] = pos.y[g]; - positions[i * 3 + 2] = pos.z[g]; - } - for (let i = 0; i < haloIds.length; i++) { - const g = haloIds[i]; - const l = nOwned + i; - ids[l] = g; - positions[l * 3] = pos.x[g]; - positions[l * 3 + 1] = pos.y[g]; - positions[l * 3 + 2] = pos.z[g]; - } - return { ids, ownedCount: nOwned, positions, h }; -}; - -/** - * CPU block KNN: exact k-NN of the owned points within block ∪ halo, as - * LOCAL indices (see {@link knnQueryBlock}). - * @param locals - The block's local point set. - * @param k - Neighbours per query. - * @returns Local neighbour indices, `ownedCount * k` long. - */ -const knnBlockCpu = (locals: BlockLocals, k: number): Uint32Array => { - return knnQueryBlock(locals.positions, locals.ownedCount, k); -}; - -/** - * Map local neighbour indices to global gaussian indices (sentinels pass - * through). - * @param locals - The block's local point set. - * @param nbLocal - Local neighbour indices. - * @returns A new array of global neighbour indices. - */ -const toGlobalNeighbors = (locals: BlockLocals, nbLocal: Uint32Array): Uint32Array => { - const out = new Uint32Array(nbLocal.length); - for (let s = 0; s < nbLocal.length; s++) { - out[s] = nbLocal[s] === KNN_SENTINEL ? KNN_SENTINEL : locals.ids[nbLocal[s]]; - } - return out; -}; - -// Insert (g, d2) into the k-best arrays (ascending by distance). -const kBestInsert = (bestIds: Uint32Array, bestD2: Float64Array, size: number, k: number, g: number, d2: number): number => { - if (size === k && d2 >= bestD2[k - 1]) return size; - let at = size < k ? size : k - 1; - while (at > 0 && bestD2[at - 1] > d2) { - bestD2[at] = bestD2[at - 1]; - bestIds[at] = bestIds[at - 1]; - at--; - } - bestD2[at] = d2; - bestIds[at] = g; - return Math.min(size + 1, k); -}; - -/** - * Exactness backstop for block KNN. A query's result is guaranteed correct - * when its k-th neighbour distance fits inside the halo-covered region - * (`d_k ≤ depth(q) + h`). Queries that fail — or that carry sentinel slots - * despite the scene having ≥ k other points — are re-queried best-first: - * blocks are visited in ascending point-to-AABB distance and the scan stops - * at the first block that can no longer improve on the k-th best so far. - * Same candidate set as a full scan, making the block KNN globally exact. - * - * Fixed entries are written into `nbGlobal`; the matching `nbLocal` slots - * (when provided) are marked {@link KNN_FIXED} so callers resolve those - * neighbours by global id. - * - * @param pos - Resident positions. - * @param order - Partition index array. - * @param blocks - All block ranges. - * @param blockIdx - Which block was queried. - * @param locals - The block's local point set. - * @param k - Neighbours per query. - * @param nbGlobal - Global neighbour indices (fixed in place). - * @param nbLocal - Optional parallel local indices to mark. - * @returns The number of re-queried gaussians. - */ -const verifyAndFixKnn = ( - pos: ResidentPositions, - order: Uint32Array, - blocks: BlockRange[], - blockIdx: number, - locals: BlockLocals, - k: number, - nbGlobal: Uint32Array, - nbLocal?: Uint32Array -): number => { - const block = blocks[blockIdx]; - const N = pos.x.length; - const { h } = locals; - let fixed = 0; - - const bestIds = new Uint32Array(k); - const bestD2 = new Float64Array(k); - const blockD2 = new Float64Array(blocks.length); - const blockOrd = new Uint32Array(blocks.length); - - for (let qi = 0; qi < locals.ownedCount; qi++) { - const g = locals.ids[qi]; - const qx = pos.x[g], qy = pos.y[g], qz = pos.z[g]; - - let dkSq = 0; - let sentinels = false; - for (let s = 0; s < k; s++) { - const nb = nbGlobal[qi * k + s]; - if (nb === KNN_SENTINEL) { - sentinels = true; - continue; - } - const dx = pos.x[nb] - qx, dy = pos.y[nb] - qy, dz = pos.z[nb] - qz; - const d2 = dx * dx + dy * dy + dz * dz; - if (d2 > dkSq) dkSq = d2; - } - - // Depth of q inside the block AABB (0 at/outside the boundary). - const depth = Math.max(0, Math.min( - qx - block.aabb[0], block.aabb[3] - qx, - qy - block.aabb[1], block.aabb[4] - qy, - qz - block.aabb[2], block.aabb[5] - qz - )); - - const needFix = (sentinels && N - 1 >= k) || Math.sqrt(dkSq) > depth + h; - if (!needFix) continue; - - // Best-first re-query: blocks in ascending point-to-AABB distance - // (insertion sort — block counts are small), stopping at the first - // block beyond the unverified k-th distance (a valid upper bound on - // the true one; unknown when sentinels) or beyond the k-th best so - // far, which tightens as candidates land. Blocks the loop never - // reaches provably contain no improving candidate. - const r2 = sentinels ? Infinity : dkSq; - for (let b2 = 0; b2 < blocks.length; b2++) { - const d2b = pointAabbDist2(qx, qy, qz, blocks[b2].aabb); - blockD2[b2] = d2b; - let at = b2; - while (at > 0 && blockD2[blockOrd[at - 1]] > d2b) { - blockOrd[at] = blockOrd[at - 1]; - at--; - } - blockOrd[at] = b2; - } - let size = 0; - for (let t = 0; t < blocks.length; t++) { - const b2 = blockOrd[t]; - const d2b = blockD2[b2]; - if (d2b > r2 || (size === k && d2b >= bestD2[k - 1])) break; - const other = blocks[b2]; - for (let i = other.start; i < other.end; i++) { - const cand = order[i]; - if (cand === g) continue; - const dx = pos.x[cand] - qx, dy = pos.y[cand] - qy, dz = pos.z[cand] - qz; - const d2 = dx * dx + dy * dy + dz * dz; - if (size < k || d2 < bestD2[size - 1]) { - size = kBestInsert(bestIds, bestD2, size, k, cand, d2); - } - } - } - for (let s = 0; s < k; s++) { - nbGlobal[qi * k + s] = s < size ? bestIds[s] : KNN_SENTINEL; - if (nbLocal) nbLocal[qi * k + s] = s < size ? KNN_FIXED : KNN_SENTINEL; - } - fixed++; - } - return fixed; -}; - -export { - collectBlock, - knnBlockCpu, - toGlobalNeighbors, - verifyAndFixKnn, - haloRadius, - KNN_FIXED, - type BlockLocals -}; diff --git a/src/lib/decimate/knn-core.ts b/src/lib/decimate/knn-core.ts index 7afb94bf..ddc08d80 100644 --- a/src/lib/decimate/knn-core.ts +++ b/src/lib/decimate/knn-core.ts @@ -1,51 +1,123 @@ -import { KdTree } from '../spatial/kd-tree'; +import { type FlatKdTree } from '../spatial/kd-tree'; /** Marks an unfilled neighbour slot (fewer than k non-self points available). */ const KNN_SENTINEL = 0xFFFFFFFF; +/** A forest part: flat KD subtree (GLOBAL splat ids) plus its point AABB. */ +type ForestPart = FlatKdTree & { + /** [minx, miny, minz, maxx, maxy, maxz] over the part's points. */ + aabb: Float32Array; +}; + /** - * Exact k-nearest-neighbours for the owned prefix of a local point set. + * Exact k-nearest-neighbours of query points against a forest of flat KD + * subtrees (each part's `nodeSplatIdx` holds GLOBAL splat ids; together the + * parts cover the whole scene). * - * Engine-free (imported by worker tasks). Builds a {@link KdTree} over - * all `n` local points (owned first, then halo) and queries the first - * `ownedCount`. Output `out[q * k + s]` is a LOCAL index into `positions`, - * sorted ascending by distance, excluding the query itself, with - * {@link KNN_SENTINEL} filling surplus slots — the same contract as the - * legacy CPU KNN loop. + * The top-K state carries across parts: a part whose AABB is farther than + * the carried worst is skipped outright (the near-path descent is otherwise + * unconditional, so without this cull every part costs a full descent per + * query); within a traversed part the standard "skip the far subtree" DFS + * runs with the accumulated bound. The final top-K equals a single tree's + * over the union — exact by construction, no halos, no verification. * - * @param positions - Interleaved xyz for all local points (owned + halo). - * @param ownedCount - Number of owned points at the front; only these are queried. + * Engine-free (imported by worker tasks); mirrors the GpuKnn WGSL traversal. + * + * @param parts - The forest (global splat ids in `nodeSplatIdx`). + * @param queryPos - Interleaved xyz query positions. + * @param queryIds - Global splat id per query (self-exclusion). + * @param count - Query count. * @param k - Neighbours per query. - * @returns Local neighbour indices, `ownedCount * k` long. + * @param out - Destination, `count * k` global ids (UNSORTED within a row), + * {@link KNN_SENTINEL} filling surplus slots. */ -const knnQueryBlock = (positions: Float32Array, ownedCount: number, k: number): Uint32Array => { - const n = positions.length / 3; - const x = new Float32Array(n); - const y = new Float32Array(n); - const z = new Float32Array(n); - for (let i = 0; i < n; i++) { - x[i] = positions[i * 3]; - y[i] = positions[i * 3 + 1]; - z[i] = positions[i * 3 + 2]; - } - const tree = new KdTree([x, y, z]); - const out = new Uint32Array(ownedCount * k).fill(KNN_SENTINEL); - const q = new Float32Array(3); - for (let i = 0; i < ownedCount; i++) { - q[0] = x[i]; - q[1] = y[i]; - q[2] = z[i]; - // Request k+1 because the tree returns the query itself (distance 0). - const res = tree.findKNearest(q, k + 1); - let outPos = 0; - for (let m = 0; m < res.indices.length && outPos < k; m++) { - const j = res.indices[m]; - if (j === i) continue; - out[i * k + outPos] = j; - outPos++; +const knnForestQuery = ( + parts: ForestPart[], + queryPos: Float32Array, + queryIds: Uint32Array, + count: number, + k: number, + out: Uint32Array +): void => { + const topDist = new Float32Array(k); + const topIdx = new Uint32Array(k); + const stack = new Uint32Array(48); + + for (let q = 0; q < count; q++) { + const qx = queryPos[q * 3]; + const qy = queryPos[q * 3 + 1]; + const qz = queryPos[q * 3 + 2]; + const qid = queryIds[q]; + + topDist.fill(Infinity); + topIdx.fill(KNN_SENTINEL); + let worst = Infinity; + let worstIdx = 0; + + for (const part of parts) { + const { nodeSplatIdx, nodePositions, nodeChildren, aabb } = part; + + // Part-entry cull: skip when the AABB cannot improve the top-K. + const ddx = Math.max(aabb[0] - qx, qx - aabb[3], 0); + const ddy = Math.max(aabb[1] - qy, qy - aabb[4], 0); + const ddz = Math.max(aabb[2] - qz, qz - aabb[5], 0); + if (ddx * ddx + ddy * ddy + ddz * ddz >= worst) continue; + // Stack packs axis (top 2 bits) with the node index, mirroring + // the GPU kernel: axis cycles x→y→z with depth. + let sp = 0; + stack[sp++] = part.rootIdx; + while (sp > 0) { + const packed = stack[--sp]; + const nodeIdx = packed & 0x3FFFFFFF; + const axis = packed >>> 30; + + const np = nodeIdx * 3; + const nx = nodePositions[np]; + const ny = nodePositions[np + 1]; + const nz = nodePositions[np + 2]; + const splatId = nodeSplatIdx[nodeIdx]; + + if (splatId !== qid) { + const dx = nx - qx, dy = ny - qy, dz = nz - qz; + const d2 = dx * dx + dy * dy + dz * dz; + if (d2 < worst) { + topDist[worstIdx] = d2; + topIdx[worstIdx] = splatId; + let w = topDist[0]; + let wi = 0; + for (let i = 1; i < k; i++) { + if (topDist[i] > w) { + w = topDist[i]; wi = i; + } + } + worst = w; + worstIdx = wi; + } + } + + const qAxisVal = axis === 0 ? qx : (axis === 1 ? qy : qz); + const nAxisVal = axis === 0 ? nx : (axis === 1 ? ny : nz); + const delta = qAxisVal - nAxisVal; + const nextAxis = axis + 1 >= 3 ? 0 : axis + 1; + const nextAxisPacked = nextAxis << 30; + + const nc = nodeIdx * 2; + const left = nodeChildren[nc]; + const right = nodeChildren[nc + 1]; + const near = delta < 0 ? left : right; + const far = delta < 0 ? right : left; + + if (far !== KNN_SENTINEL && delta * delta < worst) { + stack[sp++] = far | nextAxisPacked; + } + if (near !== KNN_SENTINEL) { + stack[sp++] = near | nextAxisPacked; + } + } } + + out.set(topIdx, q * k); } - return out; }; -export { knnQueryBlock, KNN_SENTINEL }; +export { knnForestQuery, KNN_SENTINEL, type ForestPart }; diff --git a/src/lib/decimate/priority-legacy.ts b/src/lib/decimate/priority-legacy.ts index 7bcd6afd..ee2eccd2 100644 --- a/src/lib/decimate/priority-legacy.ts +++ b/src/lib/decimate/priority-legacy.ts @@ -7,21 +7,15 @@ import { type GraphicsDevice } from 'playcanvas'; import { buildCostCacheLegacy, computeEdgeCostViewLegacy } from './edge-cost-legacy'; -import { collectBlock, verifyAndFixKnn, toGlobalNeighbors, KNN_FIXED, type BlockLocals } from './knn-blocks'; import { KNN_SENTINEL } from './knn-core'; import { createMergeScratch, makeGaussianSamples, sigmoid, ellipsoidArea, type SplatView } from './moment-match'; import { type BlockRange, type ResidentPositions } from './partition'; +import { buildForest, sortNeighborRows, VIEW_GROW } from './priority'; import { type ChunkData, type ChunkDataPool, type ChunkSource } from '../chunk'; import { APP_CHUNK, GpuEdgeCostLegacy, type EdgeCostCacheLegacy } from '../gpu/gpu-edge-cost-legacy'; import { GpuKnn } from '../gpu/gpu-knn'; import { WorkerQueue } from '../workers'; -/** Halo radius multiplier on the density-estimated k-NN radius. */ -const HALO_FACTOR = 2.5; - -/** Halo size cap as a multiple of a block's owned count (buffer-sizing bound). */ -const HALO_CAP = 1; - /** * Per-gaussian best-K merge candidates, the resident output of the priority * pass. `idx[g * K + s]` is the global index of gaussian g's s-th cheapest @@ -233,72 +227,90 @@ const runPriorityPassLegacy = async ( let maxOwned = 0; for (const b of blocks) maxOwned = Math.max(maxOwned, b.end - b.start); - const maxLocalN = maxOwned * (1 + HALO_CAP); let gpuKnn: GpuKnn | undefined; let gpuCost: GpuEdgeCostLegacy | undefined; - let gpuCostCapacity = maxLocalN; + let gpuCostCapacity = Math.ceil(maxOwned * VIEW_GROW); + + // The forest is built once per generation; blocks stay the query/IO + // batching unit (see the quality pass for the shared machinery). + const { parts: forest, blockPart } = await buildForest(pos, order, blocks, !device && !WorkerQueue.isInline); - // 1-deep prefetch: the next block's halo collection + tree build runs - // while the current block computes. GpuKnn executions share one set of - // buffers, so they are serialized through `gpuKnnQueue` — the prefetched - // block's KNN starts only after the current block's has finished. + // 1-deep prefetch: the next block's KNN runs while the current block + // gathers and costs. GpuKnn executions share one set of buffers, so they + // are serialized through `gpuKnnQueue`. let gpuKnnQueue: Promise = Promise.resolve(); - type Prepared = { locals: BlockLocals; nb: Promise }; - const prepare = (bi: number): Prepared => { - const locals = collectBlock(pos, order, blocks, bi, k, HALO_FACTOR, HALO_CAP); - const copy = locals.positions.slice(); + const prepare = (bi: number): Promise => { + const owned = order.subarray(blocks[bi].start, blocks[bi].end); + const nOwned = owned.length; + const home = blockPart[bi]; + const queryPos = new Float32Array(nOwned * 3); + for (let i = 0; i < nOwned; i++) { + const g = owned[i]; + queryPos[i * 3] = pos.x[g]; + queryPos[i * 3 + 1] = pos.y[g]; + queryPos[i * 3 + 2] = pos.z[g]; + } if (device) { - const treePromise = WorkerQueue.run('buildFlatKdTree', { positions: copy }, [copy.buffer as ArrayBuffer]); - const out = new Uint32Array(locals.ownedCount * k); - const run = Promise.all([treePromise, gpuKnnQueue]).then(([flat]) => { - return gpuKnn!.execute(flat, locals.positions, locals.ids.length, locals.ownedCount, out); - }); + const out = new Uint32Array(nOwned * k); + const run = gpuKnnQueue.then(() => gpuKnn!.execute(queryPos, owned, nOwned, out, home)); gpuKnnQueue = run.catch(() => { /* surfaced by the awaiting block */ }); - return { locals, nb: run.then(() => out) }; + return run.then(() => out); } - const nb = WorkerQueue.run('knnBlock', { positions: copy, ownedCount: locals.ownedCount, k }, [copy.buffer as ArrayBuffer]); - return { locals, nb }; + const ordered = [forest[home], ...forest.filter((_, i) => i !== home)]; + const per = Math.ceil(nOwned / 4); + const jobs: Promise[] = []; + for (let off = 0; off < nOwned; off += per) { + const cnt = Math.min(per, nOwned - off); + const qp = queryPos.slice(off * 3, (off + cnt) * 3); + const qi = owned.slice(off, off + cnt); + jobs.push(WorkerQueue.run('knnForest', { parts: ordered, queryPos: qp, queryIds: qi, k }, [ + qp.buffer as ArrayBuffer, qi.buffer as ArrayBuffer + ])); + } + return Promise.all(jobs).then((outs) => { + const out = new Uint32Array(nOwned * k); + let at = 0; + for (const o of outs) { + out.set(o, at); + at += o.length; + } + return out; + }); }; try { if (device) { - gpuKnn = new GpuKnn(device, maxLocalN, k); - gpuCost = new GpuEdgeCostLegacy(device, maxLocalN, maxOwned * k, colorDim); + gpuKnn = new GpuKnn(device, forest, k); + gpuCost = new GpuEdgeCostLegacy(device, gpuCostCapacity, maxOwned * k, colorDim); } - let next: Prepared | null = blocks.length > 0 ? prepare(0) : null; + let next: Promise | null = blocks.length > 0 ? prepare(0) : null; for (let bi = 0; bi < blocks.length; bi++) { - const { locals, nb: nbPromise } = next!; + const nbPromise = next!; next = bi + 1 < blocks.length ? prepare(bi + 1) : null; - const nOwned = locals.ownedCount; const owned = order.subarray(blocks[bi].start, blocks[bi].end); - const nbLocal = await nbPromise; - const nbGlobal = toGlobalNeighbors(locals, nbLocal); - verifyAndFixKnn(pos, order, blocks, bi, locals, k, nbGlobal, nbLocal); + const nOwned = owned.length; + const nb = await nbPromise; + sortNeighborRows(pos, owned, nb, k); - // Externals: referenced rows outside the owned range (halo members - // and verification-fixed neighbours), sorted for the gather. + // Externals: referenced ids outside the owned range, sorted for + // the gather. const extRow = new Map(); for (let s = 0; s < nOwned * k; s++) { - const l = nbLocal[s]; - if (l === KNN_SENTINEL || l < nOwned) continue; - const g = nbGlobal[s]; - if (l !== KNN_FIXED) { - if (!extRow.has(g)) extRow.set(g, 0); - } else if (indexOfSorted(owned, g) < 0 && !extRow.has(g)) { - extRow.set(g, 0); - } + const g = nb[s]; + if (g === KNN_SENTINEL) continue; + if (indexOfSorted(owned, g) < 0 && !extRow.has(g)) extRow.set(g, 0); } const extraGlobals = Uint32Array.from(extRow.keys()).sort(); for (let i = 0; i < extraGlobals.length; i++) extRow.set(extraGlobals[i], nOwned + i); - // Verification-fixed externals are not bounded by the halo cap, so - // a pathological block's view can exceed the preallocated cost - // buffers — grow them to the actual view size when that happens - // (rare; costs one reallocation). + // Cross-block neighbours aren't bounded by the initial capacity + // estimate, so a pathological block's view can exceed the + // preallocated cost buffers — grow them to the actual view size + // when that happens (rare; costs one reallocation). const viewN = nOwned + extraGlobals.length; if (gpuCost && viewN > gpuCostCapacity) { gpuCost.destroy(); @@ -317,18 +329,11 @@ const runPriorityPassLegacy = async ( for (let qi = 0; qi < nOwned; qi++) { edgeOf[qi] = e; for (let s = 0; s < k; s++) { - const l = nbLocal[qi * k + s]; - if (l === KNN_SENTINEL) continue; - const g = nbGlobal[qi * k + s]; - let row: number; - if (l !== KNN_FIXED) { - row = l < nOwned ? l : extRow.get(g)!; - } else { - const oi = indexOfSorted(owned, g); - row = oi >= 0 ? oi : extRow.get(g)!; - } + const g = nb[qi * k + s]; + if (g === KNN_SENTINEL) continue; + const oi = indexOfSorted(owned, g); edgeI[e] = qi; - edgeJ[e] = row; + edgeJ[e] = oi >= 0 ? oi : extRow.get(g)!; edgeNb[e] = g; e++; } diff --git a/src/lib/decimate/priority.ts b/src/lib/decimate/priority.ts index e5eb286b..80e07a66 100644 --- a/src/lib/decimate/priority.ts +++ b/src/lib/decimate/priority.ts @@ -1,8 +1,7 @@ import { type GraphicsDevice } from 'playcanvas'; import { buildSplatCache, computeEdgeCost, CACHE_STRIDE } from './edge-cost-cpu'; -import { collectBlock, verifyAndFixKnn, toGlobalNeighbors, KNN_FIXED, type BlockLocals } from './knn-blocks'; -import { KNN_SENTINEL } from './knn-core'; +import { KNN_SENTINEL, type ForestPart } from './knn-core'; import { type SplatView } from './moment-match'; import { type BlockRange, type ResidentPositions } from './partition'; import { type ChunkData, type ChunkDataPool, type ChunkSource } from '../chunk'; @@ -10,11 +9,17 @@ import { GpuEdgeCost } from '../gpu/gpu-edge-cost'; import { GpuKnn } from '../gpu/gpu-knn'; import { WorkerQueue } from '../workers'; -/** Halo radius multiplier on the density-estimated k-NN radius. */ -const HALO_FACTOR = 2.5; +/** + * Max points per forest part. Parts build in parallel on the worker pool + * (the build is each generation's serial prefix — smaller parts spread it + * across the pool) and their trees must fit the GPU binding limits; query + * cost is ~part-count-independent (the carried bound prunes distant parts + * at their root). + */ +const PART_SIZE_MAX = 1 << 22; -/** Halo size cap as a multiple of a block's owned count (buffer-sizing bound). */ -const HALO_CAP = 1; +/** Initial per-block view capacity multiplier (externals beyond it grow the cost buffers on demand). */ +const VIEW_GROW = 1.25; /** * Per-gaussian best-K merge candidates, the resident output of the priority @@ -170,6 +175,100 @@ const indexOfSorted = (sorted: Uint32Array, g: number): number => { return -1; }; +/** + * Build the generation's KNN forest: consecutive kdPartition blocks grouped + * to {@link PART_SIZE_MAX} points per part, each part's flat tree built on + * the worker pool with node splat ids remapped to global. With `shared` the + * arrays land on SharedArrayBuffers so the CPU query tasks read them without + * copies. + * + * @param pos - Resident position columns. + * @param order - Partition order (block-contiguous global ids). + * @param blocks - Partition blocks. + * @param shared - Allocate the flat arrays on shared memory. + * @returns The forest parts and each block's home part index (queries + * traverse their home part first so every other part culls on its AABB). + */ +const buildForest = async ( + pos: ResidentPositions, + order: Uint32Array, + blocks: BlockRange[], + shared: boolean +): Promise<{ parts: ForestPart[], blockPart: Uint32Array }> => { + const jobs: Promise[] = []; + const blockPart = new Uint32Array(blocks.length); + let bi = 0; + while (bi < blocks.length) { + const startRow = blocks[bi].start; + let endRow = blocks[bi].end; + blockPart[bi] = jobs.length; + bi++; + while (bi < blocks.length && blocks[bi].end - startRow <= PART_SIZE_MAX) { + endRow = blocks[bi].end; + blockPart[bi] = jobs.length; + bi++; + } + const cnt = endRow - startRow; + const x = new Float32Array(cnt); + const y = new Float32Array(cnt); + const z = new Float32Array(cnt); + const ids = new Uint32Array(cnt); + for (let i = 0; i < cnt; i++) { + const g = order[startRow + i]; + ids[i] = g; + x[i] = pos.x[g]; + y[i] = pos.y[g]; + z[i] = pos.z[g]; + } + jobs.push(WorkerQueue.run('buildKdForestPart', { x, y, z, ids, shared }, [ + x.buffer as ArrayBuffer, y.buffer as ArrayBuffer, z.buffer as ArrayBuffer, ids.buffer as ArrayBuffer + ])); + } + return { parts: await Promise.all(jobs), blockPart }; +}; + +/** + * Canonically order each neighbour row by (distance², id) ascending with + * sentinels last, so all downstream tie-breaking is deterministic and + * identical across the GPU and CPU KNN paths. + * + * @param pos - Resident position columns. + * @param owned - The block's owned global ids (row order). + * @param nb - Neighbour rows (`owned.length * k` global ids), sorted in place. + * @param k - Neighbours per row. + */ +const sortNeighborRows = ( + pos: ResidentPositions, + owned: Uint32Array, + nb: Uint32Array, + k: number +): void => { + const d = new Float64Array(k); + const id = new Uint32Array(k); + for (let q = 0; q < owned.length; q++) { + const g = owned[q]; + const qx = pos.x[g], qy = pos.y[g], qz = pos.z[g]; + const base = q * k; + let m = 0; + for (let s = 0; s < k; s++) { + const j = nb[base + s]; + if (j === KNN_SENTINEL) continue; + const dx = pos.x[j] - qx, dy = pos.y[j] - qy, dz = pos.z[j] - qz; + const dist = dx * dx + dy * dy + dz * dz; + let at = m; + while (at > 0 && (d[at - 1] > dist || (d[at - 1] === dist && id[at - 1] > j))) { + d[at] = d[at - 1]; + id[at] = id[at - 1]; + at--; + } + d[at] = dist; + id[at] = j; + m++; + } + for (let s = 0; s < k; s++) nb[base + s] = s < m ? id[s] : KNN_SENTINEL; + } +}; + /** * The priority pass (heavy read 1): per block — exact global KNN, edge costs * for each owned gaussian's k neighbours, reduction to the best K candidates @@ -194,74 +293,101 @@ const runPriorityPass = async ( let maxOwned = 0; for (const b of blocks) maxOwned = Math.max(maxOwned, b.end - b.start); - const maxLocalN = maxOwned * (1 + HALO_CAP); let gpuKnn: GpuKnn | undefined; let gpuCost: GpuEdgeCost | undefined; - let gpuCostCapacity = maxLocalN; + let gpuCostCapacity = Math.ceil(maxOwned * VIEW_GROW); - // 1-deep prefetch: the next block's halo collection + tree build runs - // while the current block computes. GpuKnn executions share one set of - // buffers, so they are serialized through `gpuKnnQueue` — the prefetched - // block's KNN starts only after the current block's has finished. + // The forest is built once per generation (its trees are exact and + // global, so block boundaries are invisible to the results); blocks stay + // the query/IO batching unit. Shared arrays feed the CPU query tasks. + const { parts: forest, blockPart } = await buildForest(pos, order, blocks, !device && !WorkerQueue.isInline); + + // 1-deep prefetch: the next block's KNN runs while the current block + // gathers and costs. GpuKnn executions share one set of buffers, so they + // are serialized through `gpuKnnQueue` — the prefetched block's KNN + // starts only after the current block's has finished. let gpuKnnQueue: Promise = Promise.resolve(); - type Prepared = { locals: BlockLocals; nb: Promise }; - const prepare = (bi: number): Prepared => { - const locals = collectBlock(pos, order, blocks, bi, k, HALO_FACTOR, HALO_CAP); - const copy = locals.positions.slice(); + const prepare = (bi: number): Promise => { + const owned = order.subarray(blocks[bi].start, blocks[bi].end); + const nOwned = owned.length; + const home = blockPart[bi]; + const queryPos = new Float32Array(nOwned * 3); + for (let i = 0; i < nOwned; i++) { + const g = owned[i]; + queryPos[i * 3] = pos.x[g]; + queryPos[i * 3 + 1] = pos.y[g]; + queryPos[i * 3 + 2] = pos.z[g]; + } if (device) { - const treePromise = WorkerQueue.run('buildFlatKdTree', { positions: copy }, [copy.buffer as ArrayBuffer]); - const out = new Uint32Array(locals.ownedCount * k); - const run = Promise.all([treePromise, gpuKnnQueue]).then(([flat]) => { - return gpuKnn!.execute(flat, locals.positions, locals.ids.length, locals.ownedCount, out); - }); + const out = new Uint32Array(nOwned * k); + const run = gpuKnnQueue.then(() => gpuKnn!.execute(queryPos, owned, nOwned, out, home)); gpuKnnQueue = run.catch(() => { /* surfaced by the awaiting block */ }); - return { locals, nb: run.then(() => out) }; + return run.then(() => out); } - const nb = WorkerQueue.run('knnBlock', { positions: copy, ownedCount: locals.ownedCount, k }, [copy.buffer as ArrayBuffer]); - return { locals, nb }; + // CPU path: split the block's queries across the worker pool (extra + // slices just queue when the pool is smaller). Home part first, so + // the other parts cull on their AABBs. + const ordered = [forest[home], ...forest.filter((_, i) => i !== home)]; + const per = Math.ceil(nOwned / 4); + const jobs: Promise[] = []; + for (let off = 0; off < nOwned; off += per) { + const cnt = Math.min(per, nOwned - off); + const qp = queryPos.slice(off * 3, (off + cnt) * 3); + const qi = owned.slice(off, off + cnt); + jobs.push(WorkerQueue.run('knnForest', { parts: ordered, queryPos: qp, queryIds: qi, k }, [ + qp.buffer as ArrayBuffer, qi.buffer as ArrayBuffer + ])); + } + return Promise.all(jobs).then((outs) => { + const out = new Uint32Array(nOwned * k); + let at = 0; + for (const o of outs) { + out.set(o, at); + at += o.length; + } + return out; + }); }; try { if (device) { - gpuKnn = new GpuKnn(device, maxLocalN, k); - if (cand) gpuCost = new GpuEdgeCost(device, maxLocalN, k); + gpuKnn = new GpuKnn(device, forest, k); + if (cand) gpuCost = new GpuEdgeCost(device, gpuCostCapacity, k); } - let next: Prepared | null = blocks.length > 0 ? prepare(0) : null; + let next: Promise | null = blocks.length > 0 ? prepare(0) : null; for (let bi = 0; bi < blocks.length; bi++) { - const { locals, nb: nbPromise } = next!; + const nbPromise = next!; next = bi + 1 < blocks.length ? prepare(bi + 1) : null; - const nOwned = locals.ownedCount; const owned = order.subarray(blocks[bi].start, blocks[bi].end); - const nbLocal = await nbPromise; - const nbGlobal = toGlobalNeighbors(locals, nbLocal); - verifyAndFixKnn(pos, order, blocks, bi, locals, k, nbGlobal, nbLocal); + const nOwned = owned.length; + // Global neighbour ids, canonically (d², id)-ordered per row so + // downstream tie-breaking is deterministic across KNN paths. + const nb = await nbPromise; + sortNeighborRows(pos, owned, nb, k); const slots = nOwned * k; - // Externals: referenced rows outside the owned range (halo members - // and verification-fixed neighbours) — count, collect, sort, dedup. - // Only cost evaluation needs them; the persisted cache covers owned - // rows only (every gaussian is owned by exactly one block). + // Externals: referenced ids outside the owned range — count, + // collect, sort, dedup. Only cost evaluation needs them; the + // persisted cache covers owned rows only (every gaussian is owned + // by exactly one block). let extraGlobals = new Uint32Array(0); if (cand) { let extCount = 0; for (let s = 0; s < slots; s++) { - const l = nbLocal[s]; - if (l === KNN_SENTINEL || l < nOwned) continue; - if (l === KNN_FIXED && indexOfSorted(owned, nbGlobal[s]) >= 0) continue; + const g = nb[s]; + if (g === KNN_SENTINEL || indexOfSorted(owned, g) >= 0) continue; extCount++; } const extSorted = new Uint32Array(extCount); extCount = 0; for (let s = 0; s < slots; s++) { - const l = nbLocal[s]; - if (l === KNN_SENTINEL || l < nOwned) continue; - const g = nbGlobal[s]; - if (l === KNN_FIXED && indexOfSorted(owned, g) >= 0) continue; + const g = nb[s]; + if (g === KNN_SENTINEL || indexOfSorted(owned, g) >= 0) continue; extSorted[extCount++] = g; } extSorted.sort(); @@ -272,10 +398,10 @@ const runPriorityPass = async ( extraGlobals = extSorted.subarray(0, uniq); } - // Verification-fixed externals are not bounded by the halo cap, so - // a pathological block's view can exceed the preallocated cost - // buffers — grow them to the actual view size when that happens - // (rare; costs one reallocation). + // Cross-block neighbours aren't bounded by the initial capacity + // estimate, so a pathological block's view can exceed the + // preallocated cost buffers — grow them to the actual view size + // when that happens (rare; costs one reallocation). const viewN = nOwned + extraGlobals.length; if (gpuCost && viewN > gpuCostCapacity) { gpuCost.destroy(); @@ -291,26 +417,26 @@ const runPriorityPass = async ( buildSplatCache(view, cache); if (cand) { - // Translate neighbour slots in place: local/global → view rows - // (dense-slot edge model; sentinel slots stay sentinel). + // Translate neighbour slots: global ids → view rows + // (dense-slot edge model; sentinel slots stay sentinel). `nb` + // keeps the global ids for the candidate output. + const nbRow = new Uint32Array(slots); for (let s = 0; s < slots; s++) { - const l = nbLocal[s]; - if (l === KNN_SENTINEL || l < nOwned) continue; - const g = nbGlobal[s]; - if (l === KNN_FIXED) { - const oi = indexOfSorted(owned, g); - nbLocal[s] = oi >= 0 ? oi : nOwned + indexOfSorted(extraGlobals, g); - } else { - nbLocal[s] = nOwned + indexOfSorted(extraGlobals, g); + const g = nb[s]; + if (g === KNN_SENTINEL) { + nbRow[s] = KNN_SENTINEL; + continue; } + const oi = indexOfSorted(owned, g); + nbRow[s] = oi >= 0 ? oi : nOwned + indexOfSorted(extraGlobals, g); } const blockCosts = new Float32Array(slots); if (device) { - await gpuCost!.execute(cache, viewN, nbLocal, blockCosts); + await gpuCost!.execute(cache, viewN, nbRow, blockCosts); } else { for (let s = 0; s < slots; s++) { - const row = nbLocal[s]; + const row = nbRow[s]; blockCosts[s] = row === KNN_SENTINEL ? 0 : computeEdgeCost(cache, (s / k) | 0, row); @@ -326,7 +452,7 @@ const runPriorityPass = async ( let size = 0; const base = qi * k; for (let s = 0; s < k; s++) { - if (nbLocal[base + s] === KNN_SENTINEL) continue; + if (nb[base + s] === KNN_SENTINEL) continue; const c = blockCosts[base + s]; if (!Number.isFinite(c)) continue; if (size === K && c >= bestCost[K - 1]) continue; @@ -337,7 +463,7 @@ const runPriorityPass = async ( at--; } bestCost[at] = c; - bestIdx[at] = nbGlobal[base + s]; + bestIdx[at] = nb[base + s]; size = Math.min(size + 1, K); } const g = owned[qi]; @@ -360,7 +486,7 @@ const runPriorityPass = async ( if (ctx.neighborsOut) { const NO = ctx.neighborsOut; for (let qi = 0; qi < nOwned; qi++) { - NO.set(nbGlobal.subarray(qi * k, (qi + 1) * k), owned[qi] * k); + NO.set(nb.subarray(qi * k, (qi + 1) * k), owned[qi] * k); } } @@ -374,10 +500,11 @@ const runPriorityPass = async ( export { runPriorityPass, + buildForest, gatherBlockView, + sortNeighborRows, indexOfSorted, - HALO_FACTOR, - HALO_CAP, + VIEW_GROW, type CandidateArrays, type PriorityContext, type BlockView diff --git a/src/lib/gpu/gpu-knn.ts b/src/lib/gpu/gpu-knn.ts index 6ef6803f..ede26a0a 100644 --- a/src/lib/gpu/gpu-knn.ts +++ b/src/lib/gpu/gpu-knn.ts @@ -1,90 +1,101 @@ import { BUFFERUSAGE_COPY_DST, BUFFERUSAGE_COPY_SRC, - SHADERLANGUAGE_WGSL, - SHADERSTAGE_COMPUTE, - UNIFORMTYPE_UINT, - BindGroupFormat, - BindStorageBufferFormat, - BindUniformBufferFormat, - Compute, GraphicsDevice, - Shader, - StorageBuffer, - UniformBufferFormat, - UniformFormat + StorageBuffer } from 'playcanvas'; -import { type FlatKdTree } from '../spatial/kd-tree'; +import { makeKernel, type Kernel } from './compute-kernel'; +import { type ForestPart } from '../decimate/knn-core'; /** - * WGSL kernel: iterative KD-tree K-nearest-neighbours. + * WGSL kernel: iterative KD-tree K-nearest-neighbours over one forest part. * - * Each thread runs a depth-first traversal of the flattened KD-tree with a - * fixed-size per-thread stack. Visits at most `O(K · log N)` nodes per - * query thanks to the standard "skip the far subtree if its splitting plane - * is farther than the current K-th best" pruning. Top-K is maintained - * unsorted in per-thread storage with explicit worst-index tracking, so the - * common-case "candidate is rejected against worst" path is a single - * compare-and-branch (no dynamic-indexed shift). + * Each thread runs a depth-first traversal of a flattened KD subtree with a + * fixed-size per-thread stack, carrying its query's top-K in the topDist / + * topIdx buffers (reset per batch by a separate kernel): each part loads + * and tightens them. A part whose AABB cannot improve the carried worst is + * skipped before any traversal (the near-path descent is otherwise + * unconditional — without the cull every part costs a full descent per + * query). The caller dispatches each batch's HOME part first, so the top-K + * is tight before any other part is tested and interior queries cull + * everything else. Top-K is maintained unsorted with explicit worst-index + * tracking, so the common "candidate rejected against worst" path is a + * single compare. + * + * The part's root index and AABB are baked as compile-time constants (each + * part is its own kernel), so only `queryCount` varies per dispatch. * * @param k - Compile-time K, the number of nearest neighbours per query. * @param stackSize - Compile-time per-thread DFS stack depth. + * @param rootIdx - The part's group-local root node index. + * @param aabb - The part's point AABB [minx, miny, minz, maxx, maxy, maxz]. * @returns WGSL source. */ -const knnWgsl = (k: number, stackSize: number) => /* wgsl */` +const knnWgsl = (k: number, stackSize: number, rootIdx: number, aabb: Float32Array) => /* wgsl */` struct Uniforms { - queryOffset: u32, queryCount: u32, - rootIdx: u32, } @group(0) @binding(0) var uniforms: Uniforms; -// Query positions interleaved xyz: positions[q*3 + 0/1/2]. -@group(0) @binding(1) var positions: array; -// Flattened KD-tree. Positions and children are interleaved so the kernel -// stays comfortably under the WebGPU per-stage storage-buffer minimum (8): -// nodePositions[t*3 + 0/1/2] for tree node t, nodeChildren[t*2 + 0/1] for -// (left, right). Kept separate from nodeSplatIdx to avoid mixing f32/u32. -@group(0) @binding(2) var nodeSplatIdx: array; -@group(0) @binding(3) var nodePositions: array; -@group(0) @binding(4) var nodeChildren: array; -// Output: per query, k neighbour splat indices (unsorted). -@group(0) @binding(5) var outIndices: array; +// This batch's query positions (interleaved xyz) and global splat ids. +@group(0) @binding(1) var queryPos: array; +@group(0) @binding(2) var queryIds: array; +// The binding group's concatenated forest parts: node splat ids are GLOBAL, +// node child indices are group-local (offset applied at pack time), +// positions denormalized per node. +@group(0) @binding(3) var nodeSplatIdx: array; +@group(0) @binding(4) var nodePositions: array; +@group(0) @binding(5) var nodeChildren: array; +// Carried per-query top-K state (this batch's rows). +@group(0) @binding(6) var topDist: array; +@group(0) @binding(7) var topIdx: array; const K: u32 = ${k}u; const NULL_NODE: u32 = 0xFFFFFFFFu; const F32_MAX: f32 = 3.4028234663852886e+38; -// log2(N) + slack — safe to ~2^40 nodes which is way past our limits. const STACK_SIZE: u32 = ${stackSize}u; +const ROOT_IDX: u32 = ${rootIdx}u; +const AABB_MIN: vec3f = vec3f(${aabb[0]}, ${aabb[1]}, ${aabb[2]}); +const AABB_MAX: vec3f = vec3f(${aabb[3]}, ${aabb[4]}, ${aabb[5]}); @compute @workgroup_size(64) fn main(@builtin(global_invocation_id) gid: vec3u) { let bid = gid.x; if (bid >= uniforms.queryCount) { return; } - let q = bid + uniforms.queryOffset; - - let q3 = q * 3u; - let qx = positions[q3 + 0u]; - let qy = positions[q3 + 1u]; - let qz = positions[q3 + 2u]; - - // Top-K state, unsorted. worstIdx points to the current K-th worst slot - // so accepts replace it in O(1) and we recompute worst via a fixed loop. - var topIdx: array; - var topDist: array; - var worst: f32 = F32_MAX; + + let q3 = bid * 3u; + let qx = queryPos[q3 + 0u]; + let qy = queryPos[q3 + 1u]; + let qz = queryPos[q3 + 2u]; + let qid = queryIds[bid]; + + // Load the carried top-K (unsorted, worst tracked): distances first and + // cull on the part AABB before touching anything else — a part that + // cannot improve the top-K costs one distance-row read, not a tree + // descent. (worst == F32_MAX means unfilled slots — never cull then.) + var tIdx: array; + var tDist: array; + let base = bid * K; var worstIdx: u32 = 0u; for (var i: u32 = 0u; i < K; i++) { - topDist[i] = F32_MAX; - topIdx[i] = 0u; + tDist[i] = topDist[base + i]; + } + var worst: f32 = tDist[0]; + for (var i: u32 = 1u; i < K; i++) { + if (tDist[i] > worst) { worst = tDist[i]; worstIdx = i; } + } + let dd = max(max(AABB_MIN - vec3f(qx, qy, qz), vec3f(qx, qy, qz) - AABB_MAX), vec3f(0.0)); + if (worst < F32_MAX && dot(dd, dd) >= worst) { return; } + for (var i: u32 = 0u; i < K; i++) { + tIdx[i] = topIdx[base + i]; } // Stack: (nodeIdx, axis) packed as u32. axis ∈ {0,1,2} in top 2 bits, - // nodeIdx in low 30 — supports up to ~1B nodes. + // group-local nodeIdx in low 30 — supports up to ~1B nodes per group. var stack: array; var sp: u32 = 0u; - stack[0] = uniforms.rootIdx; // axis=0 → no axis bits set + stack[0] = ROOT_IDX; // axis=0 → no axis bits set sp = 1u; while (sp > 0u) { @@ -93,37 +104,32 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { let nodeIdx = packed & 0x3FFFFFFFu; let axis = packed >> 30u; - // Read the node's position + splat id. let np = nodeIdx * 3u; let nx = nodePositions[np + 0u]; let ny = nodePositions[np + 1u]; let nz = nodePositions[np + 2u]; let splatId = nodeSplatIdx[nodeIdx]; - // Update top-K, skipping the query itself. - if (splatId != q) { + // Update top-K, skipping the query itself (by global id, so + // coincident points stay valid mutual neighbours). + if (splatId != qid) { let dx = nx - qx; let dy = ny - qy; let dz = nz - qz; let d2 = dx * dx + dy * dy + dz * dz; if (d2 < worst) { - topDist[worstIdx] = d2; - topIdx[worstIdx] = splatId; - // Recompute worst with a constant-bound loop (compiler can - // unroll → all accesses to topDist resolve statically). - var w: f32 = topDist[0]; + tDist[worstIdx] = d2; + tIdx[worstIdx] = splatId; + var w: f32 = tDist[0]; var wi: u32 = 0u; for (var i: u32 = 1u; i < K; i++) { - if (topDist[i] > w) { w = topDist[i]; wi = i; } + if (tDist[i] > w) { w = tDist[i]; wi = i; } } worst = w; worstIdx = wi; } } - // Choose near/far children based on which side of the splitting - // plane the query lies on. Walk near first (push far first so LIFO - // pops near first), with pruning on far. var qAxisVal: f32; var nAxisVal: f32; if (axis == 0u) { qAxisVal = qx; nAxisVal = nx; } @@ -140,8 +146,6 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { let near = select(rightChild, leftChild, delta < 0.0); let far = select(leftChild, rightChild, delta < 0.0); - // Push far first iff its subtree could still hold a closer point - // than the current K-th best. if (far != NULL_NODE && delta * delta < worst) { stack[sp] = far | nextAxisPacked; sp = sp + 1u; @@ -152,174 +156,215 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { } } - // Emit unsorted top-K (the decimator does not require sorted neighbours). - // Slots that never received a real candidate (n-1 < K) keep F32_MAX in - // topDist; emit the sentinel 0xFFFFFFFF for those so downstream - // edge-extraction can skip them, matching the CPU path. - let outBase = bid * K; + // Store the carried state back. Unfilled slots (n-1 < K over the whole + // forest) keep the NULL sentinel, matching the CPU path. for (var i: u32 = 0u; i < K; i++) { - if (topDist[i] == F32_MAX) { - outIndices[outBase + i] = 0xFFFFFFFFu; - } else { - outIndices[outBase + i] = topIdx[i]; - } + topDist[base + i] = tDist[i]; + topIdx[base + i] = tIdx[i]; } } `; +// Reset the batch's carried top-K rows (runs before the batch's parts). +const knnResetWgsl = () => /* wgsl */` +struct Uniforms { + slotCount: u32, +} + +@group(0) @binding(0) var uniforms: Uniforms; +@group(0) @binding(1) var topDist: array; +@group(0) @binding(2) var topIdx: array; + +@compute @workgroup_size(256) +fn main(@builtin(global_invocation_id) gid: vec3u) { + let i = gid.x; + if (i >= uniforms.slotCount) { return; } + topDist[i] = 3.4028234663852886e+38; + topIdx[i] = 0xFFFFFFFFu; +} +`; + /** - * GPU K-nearest-neighbours over a fixed point set using a flattened KD-tree. - * - * Algorithm: classic KD-tree DFS with bounded heap pruning, except the - * recursion is unrolled into an explicit per-thread stack and the top-K is - * maintained unsorted (with worst-index tracking) so the dominant - * candidate-rejection path is a single compare. Same O(N log N) total work - * as the CPU KD-tree the kernel mirrors, just parallelised across queries. + * GPU K-nearest-neighbours over a forest of flat KD subtrees. * - * The flat tree is built by the caller (`buildFlatKdTree`, typically - * off-thread via the worker task of the same name) — this class only uploads - * and traverses it. - * - * Memory footprint: ~24 N bytes for the flattened tree (3 floats + 3 - * u32 per node), plus query positions and the per-query output indices. + * The forest (built per generation by the `buildKdForestPart` worker task, + * global splat ids in `nodeSplatIdx`) is uploaded once at construction: + * parts are concatenated into binding groups sized under the device's + * storage-binding limit, child indices offset to group-local at pack time. + * Queries batch at 65,536; each batch dispatches every part in order (one + * Compute per part — uniforms never vary within a submit) with the top-K + * carried in buffers, then reads back ids only. Exactness: the carried + * worst-distance bound makes the multi-part traversal equivalent to one + * tree over the union. */ class GpuKnn { /** - * @param tree - Prebuilt flat KD-tree over the `n` local points - * (see `buildFlatKdTree`; node splat ids are LOCAL indices). - * @param positions - Interleaved xyz for all `n` local points; queries - * are the first `queryCount` of them (owned-first ordering). - * @param n - Total local point count (tree size). - * @param queryCount - How many leading points to query. - * @param outNeighbours - destination for per-query K neighbour indices, - * length `queryCount * k`. `outNeighbours[i * k + j]` is one of the k - * nearest LOCAL neighbours of point i (UNSORTED). Excludes i itself; - * sentinel 0xFFFFFFFF fills surplus slots. + * @param queryPos - Interleaved xyz for this call's queries. + * @param queryIds - Global splat id per query (self-exclusion). + * @param queryCount - Query count. + * @param outNeighbours - `queryCount * k` GLOBAL neighbour ids + * (UNSORTED within a row; 0xFFFFFFFF sentinel for surplus slots). + * @param homePart - Part index containing this call's queries; it is + * traversed first so every other part sees a tight bound and culls. */ execute: ( - tree: FlatKdTree, - positions: Float32Array, - n: number, + queryPos: Float32Array, + queryIds: Uint32Array, queryCount: number, - outNeighbours: Uint32Array + outNeighbours: Uint32Array, + homePart: number ) => Promise; destroy: () => void; /** * @param device - PlayCanvas GraphicsDevice (WebGPU). - * @param maxN - Maximum number of points the index will handle. + * @param parts - The forest (node splat ids GLOBAL, children part-local). * @param k - Number of nearest neighbours per query. */ - constructor(device: GraphicsDevice, maxN: number, k: number) { + constructor(device: GraphicsDevice, parts: ForestPart[], k: number) { const workgroupSize = 64; const queriesPerBatch = 1024 * workgroupSize; // 65,536 - // Per-thread DFS stack depth: tree depth = log2(maxN) + slack. 48 is - // safe for any N within the 30-bit nodeIdx packing limit checked below. const stackSize = 48; - if (maxN > 0x3FFFFFFF) { - throw new Error(`GpuKnn: maxN=${maxN} exceeds 30-bit nodeIdx packing limit (~1B nodes)`); - } - // 5 storage buffers + 1 uniform — comfortably under the WebGPU - // per-stage minimum (8 storage buffers). Positions and KD-tree - // arrays are interleaved (see WGSL above) to keep the count down. - const bindGroupFormat = new BindGroupFormat(device, [ - new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE), - new BindStorageBufferFormat('positions', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('nodeSplatIdx', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('nodePositions', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('nodeChildren', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('outIndices', SHADERSTAGE_COMPUTE) - ]); + // Group parts under the per-binding ceiling (positions, 12 B/node, + // is the largest array). + const limits = (device as unknown as { limits?: { maxStorageBufferBindingSize?: number, maxBufferSize?: number } }).limits; + const maxBinding = Math.min( + typeof limits?.maxStorageBufferBindingSize === 'number' ? limits.maxStorageBufferBindingSize : 128 * 2 ** 20, + typeof limits?.maxBufferSize === 'number' ? limits.maxBufferSize : 256 * 2 ** 20 + ); + const maxNodesPerGroup = Math.floor(maxBinding / 12); - const shader = new Shader(device, { - name: 'compute-knn-kdtree', - shaderLanguage: SHADERLANGUAGE_WGSL, - cshader: knnWgsl(k, stackSize), - // @ts-ignore - computeUniformBufferFormats: { - uniforms: new UniformBufferFormat(device, [ - new UniformFormat('queryOffset', UNIFORMTYPE_UINT), - new UniformFormat('queryCount', UNIFORMTYPE_UINT), - new UniformFormat('rootIdx', UNIFORMTYPE_UINT) - ]) - }, - // @ts-ignore - computeBindGroupFormat: bindGroupFormat - }); - - const positionsBuf = new StorageBuffer(device, maxN * 3 * 4, BUFFERUSAGE_COPY_DST); - const nSplatIdxBuf = new StorageBuffer(device, maxN * 4, BUFFERUSAGE_COPY_DST); - const nPositionsBuf = new StorageBuffer(device, maxN * 3 * 4, BUFFERUSAGE_COPY_DST); - const nChildrenBuf = new StorageBuffer(device, maxN * 2 * 4, BUFFERUSAGE_COPY_DST); - - const outBatchBytes = queriesPerBatch * k * 4; - const outBuf = new StorageBuffer( + type Group = { parts: { part: ForestPart, base: number }[], nodes: number }; + const groups: Group[] = []; + for (const part of parts) { + const n = part.nodeSplatIdx.length; + if (n > maxNodesPerGroup) { + throw new Error(`GpuKnn: a forest part (${n} nodes) exceeds the device binding limit — reduce the part size`); + } + if (n > 0x3FFFFFFF) { + throw new Error(`GpuKnn: part exceeds the 30-bit node packing limit (${n} nodes)`); + } + let g = groups[groups.length - 1]; + if (!g || g.nodes + n > maxNodesPerGroup) { + g = { parts: [], nodes: 0 }; + groups.push(g); + } + g.parts.push({ part, base: g.nodes }); + g.nodes += n; + } + + // Shared per-batch buffers. + const queryPosBuf = new StorageBuffer(device, queriesPerBatch * 3 * 4, BUFFERUSAGE_COPY_DST); + const queryIdBuf = new StorageBuffer(device, queriesPerBatch * 4, BUFFERUSAGE_COPY_DST); + const topDistBuf = new StorageBuffer(device, queriesPerBatch * k * 4, BUFFERUSAGE_COPY_DST); + const topIdxBuf = new StorageBuffer( device, - outBatchBytes, + queriesPerBatch * k * 4, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST ); const outScratch = new Uint32Array(queriesPerBatch * k); - const compute = new Compute(device, shader, 'compute-knn-kdtree'); - compute.setParameter('positions', positionsBuf); - compute.setParameter('nodeSplatIdx', nSplatIdxBuf); - compute.setParameter('nodePositions', nPositionsBuf); - compute.setParameter('nodeChildren', nChildrenBuf); - compute.setParameter('outIndices', outBuf); + // Per-batch top-K reset (frees the part dispatch order per batch). + const resetKernel = makeKernel(device, 'compute-knn-reset', knnResetWgsl(), ['slotCount'], [ + ['topDist', false], + ['topIdx', false] + ]); + resetKernel.compute.setParameter('topDist', topDistBuf); + resetKernel.compute.setParameter('topIdx', topIdxBuf); + + // Per group: concatenated node buffers (children offset to + // group-local); per part: a Compute with its root/AABB baked in + // (only queryCount varies per dispatch). + const kernels: Kernel[] = []; + const groupBufs: StorageBuffer[] = []; + for (const g of groups) { + const splatIdxBuf = new StorageBuffer(device, g.nodes * 4, BUFFERUSAGE_COPY_DST); + const positionsBuf = new StorageBuffer(device, g.nodes * 3 * 4, BUFFERUSAGE_COPY_DST); + const childrenBuf = new StorageBuffer(device, g.nodes * 2 * 4, BUFFERUSAGE_COPY_DST); + groupBufs.push(splatIdxBuf, positionsBuf, childrenBuf); + + const childScratch = new Uint32Array(1 << 16); + for (const { part, base } of g.parts) { + const n = part.nodeSplatIdx.length; + splatIdxBuf.write(base * 4, part.nodeSplatIdx, 0, n); + positionsBuf.write(base * 3 * 4, part.nodePositions, 0, n * 3); + // Children shift to group-local indices at upload. + let remapped = childScratch; + if (remapped.length < n * 2) remapped = new Uint32Array(n * 2); + for (let i = 0; i < n * 2; i++) { + const c = part.nodeChildren[i]; + remapped[i] = c === 0xFFFFFFFF ? c : c + base; + } + childrenBuf.write(base * 2 * 4, remapped, 0, n * 2); + + const source = knnWgsl(k, stackSize, base + part.rootIdx, part.aabb); + const kernel = makeKernel(device, 'compute-knn-forest', source, ['queryCount'], [ + ['queryPos', true], + ['queryIds', true], + ['nodeSplatIdx', true], + ['nodePositions', true], + ['nodeChildren', true], + ['topDist', false], + ['topIdx', false] + ]); + kernel.compute.setParameter('queryPos', queryPosBuf); + kernel.compute.setParameter('queryIds', queryIdBuf); + kernel.compute.setParameter('nodeSplatIdx', splatIdxBuf); + kernel.compute.setParameter('nodePositions', positionsBuf); + kernel.compute.setParameter('nodeChildren', childrenBuf); + kernel.compute.setParameter('topDist', topDistBuf); + kernel.compute.setParameter('topIdx', topIdxBuf); + kernels.push(kernel); + } + } this.execute = async ( - tree: FlatKdTree, - positions: Float32Array, - n: number, + queryPos: Float32Array, + queryIds: Uint32Array, queryCount: number, - outNeighbours: Uint32Array + outNeighbours: Uint32Array, + homePart: number ) => { - if (n > maxN) { - throw new Error(`GpuKnn: N=${n} exceeds maxN=${maxN}`); - } - if (positions.length < n * 3) { - throw new Error(`GpuKnn: positions length ${positions.length} must be at least N*3 = ${n * 3}`); - } - if (queryCount > n) { - throw new Error(`GpuKnn: queryCount=${queryCount} exceeds N=${n}`); - } if (outNeighbours.length !== queryCount * k) { throw new Error(`GpuKnn: outNeighbours length ${outNeighbours.length} must be queryCount*k = ${queryCount * k}`); } - positionsBuf.write(0, positions, 0, n * 3); - nSplatIdxBuf.write(0, tree.nodeSplatIdx, 0, n); - nPositionsBuf.write(0, tree.nodePositions, 0, n * 3); - nChildrenBuf.write(0, tree.nodeChildren, 0, n * 2); - compute.setParameter('rootIdx', tree.rootIdx); + // Home part first: it fills the top-K with true near neighbours, + // so every other part's AABB test culls for interior queries. + const order = [kernels[homePart], ...kernels.filter((_, i) => i !== homePart)]; const numBatches = Math.ceil(queryCount / queriesPerBatch); for (let batch = 0; batch < numBatches; batch++) { const queryOffset = batch * queriesPerBatch; const batchCount = Math.min(queriesPerBatch, queryCount - queryOffset); - const groups = Math.ceil(batchCount / workgroupSize); + const dispatchGroups = Math.ceil(batchCount / workgroupSize); - compute.setParameter('queryOffset', queryOffset); - compute.setParameter('queryCount', batchCount); + queryPosBuf.write(0, queryPos, queryOffset * 3, batchCount * 3); + queryIdBuf.write(0, queryIds, queryOffset, batchCount); - compute.setupDispatch(groups); - device.computeDispatch([compute], `knn-dispatch-${batch}`); + resetKernel.compute.setParameter('slotCount', batchCount * k); + resetKernel.compute.setupDispatch(Math.ceil((batchCount * k) / 256)); + for (const kernel of order) { + kernel.compute.setParameter('queryCount', batchCount); + kernel.compute.setupDispatch(dispatchGroups); + } + device.computeDispatch([resetKernel.compute, ...order.map(kn => kn.compute)], `knn-batch-${batch}`); const readBytes = batchCount * k * 4; - await outBuf.read(0, readBytes, outScratch, true); + await topIdxBuf.read(0, readBytes, outScratch, true); outNeighbours.set(outScratch.subarray(0, batchCount * k), queryOffset * k); } }; this.destroy = () => { - positionsBuf.destroy(); - nSplatIdxBuf.destroy(); - nPositionsBuf.destroy(); - nChildrenBuf.destroy(); - outBuf.destroy(); - shader.destroy(); - bindGroupFormat.destroy(); + queryPosBuf.destroy(); + queryIdBuf.destroy(); + topDistBuf.destroy(); + topIdxBuf.destroy(); + for (const b of groupBufs) b.destroy(); + resetKernel.destroy(); + for (const kernel of kernels) kernel.destroy(); }; } } diff --git a/src/lib/workers/tasks.ts b/src/lib/workers/tasks.ts index f3b35f9e..9aa818a7 100644 --- a/src/lib/workers/tasks.ts +++ b/src/lib/workers/tasks.ts @@ -1,7 +1,7 @@ import type { TypedArray } from '../data-table/data-table'; -import { knnQueryBlock } from '../decimate/knn-core'; +import { knnForestQuery, type ForestPart } from '../decimate/knn-core'; import { mergeGroup, createMergeScratch, splatMass } from '../decimate/moment-match'; -import { buildFlatKdTree, type FlatKdTree } from '../spatial/kd-tree'; +import { buildFlatKdTree } from '../spatial/kd-tree'; import { quantize1dColumns, type QuantizedColumns } from '../spatial/quantize-1d-core'; import { WebPCodec } from '../utils/webp-codec'; @@ -39,32 +39,66 @@ const taskHandlers = { return { result: webp, transfer: [webp.buffer as ArrayBuffer] }; }, - // Build a flat KD-tree over interleaved local positions (decimation GPU - // path: the flat arrays upload straight into GpuKnn). - buildFlatKdTree: (args: { positions: Float32Array }): TaskOutput => { - const n = args.positions.length / 3; - const x = new Float32Array(n); - const y = new Float32Array(n); - const z = new Float32Array(n); - for (let i = 0; i < n; i++) { - x[i] = args.positions[i * 3]; - y[i] = args.positions[i * 3 + 1]; - z[i] = args.positions[i * 3 + 2]; - } + // Build one forest part: a flat KD-tree over the part's position columns + // whose node splat ids are remapped to GLOBAL ids, plus the part's point + // AABB (the query-side part-entry cull). With `shared`, the flat arrays + // move to SharedArrayBuffers so knnForest query tasks read them without + // copies (structured clone shares SABs). + buildKdForestPart: (args: { + x: Float32Array, y: Float32Array, z: Float32Array, ids: Uint32Array, shared?: boolean + }): TaskOutput => { + const { x, y, z, ids } = args; const flat = buildFlatKdTree(x, y, z); + const splatIdx = flat.nodeSplatIdx; + for (let t = 0; t < splatIdx.length; t++) splatIdx[t] = ids[splatIdx[t]]; + const aabb = new Float32Array([Infinity, Infinity, Infinity, -Infinity, -Infinity, -Infinity]); + for (let i = 0; i < x.length; i++) { + if (x[i] < aabb[0]) aabb[0] = x[i]; + if (y[i] < aabb[1]) aabb[1] = y[i]; + if (z[i] < aabb[2]) aabb[2] = z[i]; + if (x[i] > aabb[3]) aabb[3] = x[i]; + if (y[i] > aabb[4]) aabb[4] = y[i]; + if (z[i] > aabb[5]) aabb[5] = z[i]; + } + if (args.shared && typeof SharedArrayBuffer !== 'undefined') { + const shareU32 = (a: Uint32Array): Uint32Array => { + const s = new Uint32Array(new SharedArrayBuffer(a.byteLength)); + s.set(a); + return s; + }; + const shareF32 = (a: Float32Array): Float32Array => { + const s = new Float32Array(new SharedArrayBuffer(a.byteLength)); + s.set(a); + return s; + }; + return { + result: { + nodeSplatIdx: shareU32(flat.nodeSplatIdx), + nodePositions: shareF32(flat.nodePositions), + nodeChildren: shareU32(flat.nodeChildren), + rootIdx: flat.rootIdx, + aabb + }, + transfer: [] + }; + } return { - result: flat, + result: { ...flat, aabb }, transfer: [ - flat.nodeSplatIdx.buffer, flat.nodePositions.buffer, flat.nodeChildren.buffer + flat.nodeSplatIdx.buffer, flat.nodePositions.buffer, flat.nodeChildren.buffer, aabb.buffer ] as ArrayBuffer[] }; }, - // Decimation CPU-fallback block KNN: exact k-NN of the owned prefix - // within the local point set, as local indices. - knnBlock: (args: { positions: Float32Array, ownedCount: number, k: number }): TaskOutput => { - const result = knnQueryBlock(args.positions, args.ownedCount, args.k); - return { result, transfer: [result.buffer as ArrayBuffer] }; + // Decimation CPU-fallback KNN: exact forest k-NN for a slice of queries, + // as global ids. Part arrays are SAB-backed (shared, no copies). + knnForest: (args: { + parts: ForestPart[], queryPos: Float32Array, queryIds: Uint32Array, k: number + }): TaskOutput => { + const count = args.queryIds.length; + const out = new Uint32Array(count * args.k); + knnForestQuery(args.parts, args.queryPos, args.queryIds, count, args.k, out); + return { result: out, transfer: [out.buffer as ArrayBuffer] }; }, // Decimation merge stream: n-ary moment match of packed member-major diff --git a/test/decimate-knn.test.mjs b/test/decimate-knn.test.mjs index ea7ae129..54703c51 100644 --- a/test/decimate-knn.test.mjs +++ b/test/decimate-knn.test.mjs @@ -1,13 +1,17 @@ /** - * Block KNN tests: exactness of block∪halo KNN + verification against global - * brute force, on both benign (uniform) and adversarial (clustered) scenes. + * Forest KNN tests: exactness of the multi-part carried-bound query against + * global brute force on benign and adversarial scenes — no halos, no + * verification pass, exact by construction — plus the canonical neighbour + * ordering all downstream tie-breaking depends on. */ import assert from 'node:assert'; import { describe, it } from 'node:test'; -import { collectBlock, knnBlockCpu, toGlobalNeighbors, verifyAndFixKnn } from '../src/lib/decimate/knn-blocks.js'; +import { knnForestQuery, KNN_SENTINEL } from '../src/lib/decimate/knn-core.js'; import { kdPartition } from '../src/lib/decimate/partition.js'; +import { sortNeighborRows } from '../src/lib/decimate/priority.js'; +import { buildFlatKdTree } from '../src/lib/spatial/kd-tree.js'; const mulberry = (seed) => { let t = seed >>> 0; @@ -22,69 +26,104 @@ const mulberry = (seed) => { const scenes = { uniform: (n, r) => Float32Array.from({ length: n }, () => r() * 10), clustered: (n, r) => Float32Array.from({ length: n }, (_, i) => (i % 7) + r() * 0.01), - // Bulk cluster + rare extreme flyaways: stretched block AABBs defeat the - // density-based halo estimate — the requery-heavy regime. + // Bulk cluster + rare extreme flyaways: the stretched-AABB regime that + // defeated the old halo estimate. flyaway: (n, r) => Float32Array.from({ length: n }, () => (r() < 0.01 ? (r() - 0.5) * 5000 : r() * 10)), // Integer-grid coordinates: many coincident points and exact distance ties. coincident: (n, r) => Float32Array.from({ length: n }, () => Math.floor(r() * 12)) }; -// Adversarial scenes are ALLOWED to requery heavily; the benign ones must not. -const requeryCapped = new Set(['uniform', 'clustered']); +// Build a forest exactly as the buildKdForestPart task does: part-local +// column slices from partition ranges, node splat ids remapped to global, +// point AABB alongside (the part-entry cull). +const buildForest = (pos, order, ranges) => { + return ranges.map(([start, end]) => { + const cnt = end - start; + const x = new Float32Array(cnt); + const y = new Float32Array(cnt); + const z = new Float32Array(cnt); + const ids = new Uint32Array(cnt); + const aabb = new Float32Array([Infinity, Infinity, Infinity, -Infinity, -Infinity, -Infinity]); + for (let i = 0; i < cnt; i++) { + const g = order[start + i]; + ids[i] = g; + x[i] = pos.x[g]; + y[i] = pos.y[g]; + z[i] = pos.z[g]; + aabb[0] = Math.min(aabb[0], x[i]); + aabb[1] = Math.min(aabb[1], y[i]); + aabb[2] = Math.min(aabb[2], z[i]); + aabb[3] = Math.max(aabb[3], x[i]); + aabb[4] = Math.max(aabb[4], y[i]); + aabb[5] = Math.max(aabb[5], z[i]); + } + const flat = buildFlatKdTree(x, y, z); + for (let t = 0; t < flat.nodeSplatIdx.length; t++) flat.nodeSplatIdx[t] = ids[flat.nodeSplatIdx[t]]; + return { ...flat, aabb }; + }); +}; + +// Group partition blocks into `partCount`-ish contiguous ranges. +const partRanges = (blocks, partCount) => { + const per = Math.ceil(blocks.length / partCount); + const ranges = []; + for (let b = 0; b < blocks.length; b += per) { + const last = Math.min(b + per, blocks.length) - 1; + ranges.push([blocks[b].start, blocks[last].end]); + } + return ranges; +}; -describe('block KNN with halo + verification', () => { +const queryAll = (pos, n, forest, k) => { + const queryPos = new Float32Array(n * 3); + const ids = new Uint32Array(n); + for (let i = 0; i < n; i++) { + ids[i] = i; + queryPos[i * 3] = pos.x[i]; + queryPos[i * 3 + 1] = pos.y[i]; + queryPos[i * 3 + 2] = pos.z[i]; + } + const out = new Uint32Array(n * k); + knnForestQuery(forest, queryPos, ids, n, k, out); + return out; +}; + +describe('forest KNN', () => { for (const [name, gen] of Object.entries(scenes)) { - it(`matches global brute force (${name})`, () => { + it(`matches global brute force across parts (${name})`, () => { const n = 3000, k = 8, r = mulberry(3); const pos = { x: gen(n, r), y: gen(n, r), z: gen(n, r) }; const { order, blocks } = kdPartition(pos, 500); + const forest = buildForest(pos, order, partRanges(blocks, 3)); + assert.ok(forest.length >= 2, 'multi-part coverage'); + const out = queryAll(pos, n, forest, k); + const d2 = (a, b) => (pos.x[a] - pos.x[b]) ** 2 + (pos.y[a] - pos.y[b]) ** 2 + (pos.z[a] - pos.z[b]) ** 2; - let totalFixed = 0; - for (let bi = 0; bi < blocks.length; bi++) { - const locals = collectBlock(pos, order, blocks, bi, k, 2.5); - // Buffer-sizing contract: locals never exceed owned × (1 + haloCap). - assert.ok( - locals.ids.length <= locals.ownedCount + Math.ceil(locals.ownedCount * 1), - `${name} block ${bi}: halo exceeds the cap (${locals.ids.length} locals for ${locals.ownedCount} owned)` - ); - const nbLocal = knnBlockCpu(locals, k); - const nbGlobal = toGlobalNeighbors(locals, nbLocal); - totalFixed += verifyAndFixKnn(pos, order, blocks, bi, locals, k, nbGlobal, nbLocal); - for (let q = 0; q < locals.ownedCount; q++) { - const g = locals.ids[q]; - const dists = []; - for (let i = 0; i < n; i++) { - if (i !== g) dists.push(d2(g, i)); - } - dists.sort((a, b) => a - b); - const got = []; - const seen = new Set(); - for (let s = 0; s < k; s++) { - const nb = nbGlobal[q * k + s]; - assert.notStrictEqual(nb, 0xFFFFFFFF, `${name} block ${bi} q ${q} slot ${s} sentinel`); - assert.notStrictEqual(nb, g, 'self excluded'); - assert.ok(!seen.has(nb), `${name} block ${bi} q ${q}: duplicate neighbour`); - seen.add(nb); - got.push(d2(g, nb)); - } - // Exact k-NN: the neighbour distance multiset must equal the - // true k smallest (valid under ties, where ids may differ). - got.sort((a, b) => a - b); - assert.deepStrictEqual(got, dists.slice(0, k), `${name} block ${bi} q ${q}: not the exact k nearest`); + for (let q = 0; q < n; q += 7) { + const dists = []; + for (let i = 0; i < n; i++) { + if (i !== q) dists.push(d2(q, i)); } - } - // sanity: verification exists but is not doing all the work - if (requeryCapped.has(name)) { - assert.ok(totalFixed < n * 0.5, `${name}: too many requeries (${totalFixed})`); + dists.sort((a, b) => a - b); + const got = []; + const seen = new Set(); + for (let s = 0; s < k; s++) { + const nb = out[q * k + s]; + assert.notStrictEqual(nb, KNN_SENTINEL, `${name} q ${q} slot ${s} sentinel`); + assert.notStrictEqual(nb, q, 'self excluded'); + assert.ok(!seen.has(nb), `${name} q ${q}: duplicate neighbour`); + seen.add(nb); + got.push(d2(q, nb)); + } + // Exact k-NN: the neighbour distance multiset must equal the + // true k smallest (valid under ties, where ids may differ). + got.sort((a, b) => a - b); + assert.deepStrictEqual(got, dists.slice(0, k), `${name} q ${q}: not the exact k nearest`); } }); } - it('enveloping residual block (no covering halo) stays exact via forced requery', () => { - // Bulk grid + rare extreme flyaways: the residual block's AABB - // envelops the core, so no halo radius can satisfy the size cap — - // collectBlock must fall back to an empty halo (h = -Infinity) and - // verification must recover exactness for every residual query. + it('enveloping flyaways stay exact (the old no-covering-halo regime)', () => { const k = 8; const side = 22, nBulk = side * side * side; // 10648 const fly = [ @@ -105,36 +144,25 @@ describe('block KNN with halo + verification', () => { pos.x[nBulk + j] = fx; pos.y[nBulk + j] = fy; pos.z[nBulk + j] = fz; }); const { order, blocks } = kdPartition(pos, 1024); + const forest = buildForest(pos, order, partRanges(blocks, 4)); + const out = queryAll(pos, n, forest, k); + const d2 = (a, b) => (pos.x[a] - pos.x[b]) ** 2 + (pos.y[a] - pos.y[b]) ** 2 + (pos.z[a] - pos.z[b]) ** 2; - let sawFallback = false; - for (let bi = 0; bi < blocks.length; bi++) { - const locals = collectBlock(pos, order, blocks, bi, k, 2.5); - assert.ok( - locals.ids.length <= locals.ownedCount + Math.ceil(locals.ownedCount * 1), - `block ${bi}: halo exceeds the cap` - ); - if (locals.h === -Infinity) sawFallback = true; - const nbLocal = knnBlockCpu(locals, k); - const nbGlobal = toGlobalNeighbors(locals, nbLocal); - verifyAndFixKnn(pos, order, blocks, bi, locals, k, nbGlobal, nbLocal); - for (let q = 0; q < locals.ownedCount; q++) { - const g = locals.ids[q]; - const dists = []; - for (let j = 0; j < n; j++) { - if (j !== g) dists.push(d2(g, j)); - } - dists.sort((a, b) => a - b); - const got = []; - for (let s = 0; s < k; s++) { - const nb = nbGlobal[q * k + s]; - assert.notStrictEqual(nb, 0xFFFFFFFF, `block ${bi} q ${q} sentinel`); - got.push(d2(g, nb)); - } - got.sort((a, b) => a - b); - assert.deepStrictEqual(got, dists.slice(0, k), `block ${bi} q ${q}: not the exact k nearest`); + for (let q = 0; q < n; q += 97) { + const dists = []; + for (let j = 0; j < n; j++) { + if (j !== q) dists.push(d2(q, j)); + } + dists.sort((a, b) => a - b); + const got = []; + for (let s = 0; s < k; s++) { + const nb = out[q * k + s]; + assert.notStrictEqual(nb, KNN_SENTINEL, `q ${q} sentinel`); + got.push(d2(q, nb)); } + got.sort((a, b) => a - b); + assert.deepStrictEqual(got, dists.slice(0, k), `q ${q}: not the exact k nearest`); } - assert.ok(sawFallback, 'expected at least one no-covering-halo fallback block'); }); it('tiny scene (n <= k) keeps sentinels', () => { @@ -145,17 +173,50 @@ describe('block KNN with halo + verification', () => { z: new Float32Array(4) }; const { order, blocks } = kdPartition(pos, 100); - const locals = collectBlock(pos, order, blocks, 0, k, 2.5); - const nbLocal = knnBlockCpu(locals, k); - const nbGlobal = toGlobalNeighbors(locals, nbLocal); - verifyAndFixKnn(pos, order, blocks, 0, locals, k, nbGlobal, nbLocal); + const forest = buildForest(pos, order, partRanges(blocks, 1)); + const out = queryAll(pos, n, forest, k); for (let q = 0; q < 4; q++) { const filled = []; for (let s = 0; s < k; s++) { - if (nbGlobal[q * k + s] !== 0xFFFFFFFF) filled.push(nbGlobal[q * k + s]); + if (out[q * k + s] !== KNN_SENTINEL) filled.push(out[q * k + s]); } assert.strictEqual(filled.length, 3, 'exactly n-1 real neighbours'); - assert.ok(!filled.includes(locals.ids[q]), 'self excluded'); + assert.ok(!filled.includes(q), 'self excluded'); + } + }); + + it('canonical row order: (d², id) ascending, sentinels last', () => { + const n = 500, k = 8, r = mulberry(11); + const pos = { + x: scenes.coincident(n, r), + y: scenes.coincident(n, r), + z: scenes.coincident(n, r) + }; + const { order, blocks } = kdPartition(pos, 128); + const forest = buildForest(pos, order, partRanges(blocks, 2)); + const nb = queryAll(pos, n, forest, k); + const owned = new Uint32Array(n); + for (let i = 0; i < n; i++) owned[i] = i; + sortNeighborRows(pos, owned, nb, k); + + const d2 = (a, b) => (pos.x[a] - pos.x[b]) ** 2 + (pos.y[a] - pos.y[b]) ** 2 + (pos.z[a] - pos.z[b]) ** 2; + for (let q = 0; q < n; q++) { + let prevD = -1, prevId = -1, sawSentinel = false; + for (let s = 0; s < k; s++) { + const j = nb[q * k + s]; + if (j === KNN_SENTINEL) { + sawSentinel = true; + continue; + } + assert.ok(!sawSentinel, `q ${q}: id after sentinel`); + const dist = d2(q, j); + assert.ok( + dist > prevD || (dist === prevD && j > prevId), + `q ${q} slot ${s}: (${dist}, ${j}) not after (${prevD}, ${prevId})` + ); + prevD = dist; + prevId = j; + } } }); }); From 01262a18a8145634f8003d371fdbdc8d196543a8 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 18:12:09 +0100 Subject: [PATCH 11/19] latest --- src/lib/decimate/block-producer.ts | 79 ++++++++++++------------ src/lib/decimate/merge-stream.ts | 95 +++++++++++++++-------------- test/decimate-merge-stream.test.mjs | 61 ++++++++++-------- 3 files changed, 123 insertions(+), 112 deletions(-) diff --git a/src/lib/decimate/block-producer.ts b/src/lib/decimate/block-producer.ts index b45ba9df..d310072e 100644 --- a/src/lib/decimate/block-producer.ts +++ b/src/lib/decimate/block-producer.ts @@ -1,39 +1,37 @@ -import { type ChunkData, type ChunkLayer, type ChunkSource, type ChunkSourceMetadata, type ReadRequest } from '../chunk'; +import { type ChunkData, type ChunkSource, type ChunkSourceMetadata, type ReadRequest } from '../chunk'; /** - * One output chunk of the merge stream. The views hold exactly `count` - * records at the layer strides and alias the generator's rolling scratch — - * valid only until the generator's next `next()`; the consumer's read copies - * them out (yielding views instead of slices avoids a full extra copy of - * every output byte). + * One output chunk's destination views for the merge stream — typed windows + * over the consumer's {@link ChunkData} buffers, sized to the chunk's exact + * row count. Layers the consumer didn't request are absent and skipped. */ -type ChunkPayload = { - count: number; - position: Float32Array; - geometric: Float32Array; - color: Float32Array; +type DestBuffers = { + position?: Float32Array; + geometric?: Float32Array; + color?: Float32Array; other?: Uint32Array; }; /** - * A single-sequential-pass {@link ChunkSource} over an async generator of - * chunk payloads — how the decimation merge stream feeds the PLY writer - * (or `compact` / `writePlyStreaming` for intermediate generations) without - * ever materializing the output. + * A single-sequential-pass {@link ChunkSource} over the merge-stream + * generator — how decimation feeds the PLY writer (or `compact` / + * `writePlyStreaming` for intermediate generations) without ever + * materializing the output. The generator fills the consumer's chunk + * buffers directly (no intermediate copy of the output bytes). * * Contract: chunk reads must arrive in order (0, 1, 2, …), each at most * once; gather reads are not supported. Anything else throws — decimate * output supports exactly one sequential pass. * * @param meta - Exact output metadata (counts are known before streaming). - * @param produce - Factory for the payload generator (invoked lazily on first read). + * @param produce - Factory for the merge-stream generator (invoked lazily on first read). * @returns The stream-once source. */ const createBlockProducerSource = ( meta: ChunkSourceMetadata, - produce: () => AsyncGenerator + produce: () => AsyncGenerator ): ChunkSource => { - let generator: AsyncGenerator | null = null; + let generator: AsyncGenerator | null = null; let nextChunk = 0; let done = false; @@ -52,32 +50,33 @@ const createBlockProducerSource = ( if (done) { throw new Error('decimate output exhausted'); } - generator ??= produce(); - const { value, done: exhausted } = await generator.next(); - if (exhausted || !value) { + if (!generator) { + generator = produce(); + await generator.next(); // prime: run to the first dest handshake + } + + const expected = Math.min(meta.chunkSize, meta.numGaussians - request.chunkIndex * meta.chunkSize); + const view = ( + cd: ChunkData | undefined, + ctor: new (b: ArrayBuffer, o: number, n: number) => T, + perRow: number + ): T | undefined => (cd ? new ctor(cd.data, 0, expected * perRow) : undefined); + const dest: DestBuffers = { + position: view(request.position, Float32Array, 3), + geometric: view(request.geometric, Float32Array, 8), + color: view(request.color, Float32Array, (meta.layouts.color?.stride ?? 0) >> 2), + other: view(request.other, Uint32Array, (meta.layouts.other?.stride ?? 0) >> 2) + }; + + const { value, done: exhausted } = await generator.next(dest); + if (exhausted || value === undefined) { done = true; throw new Error(`decimate output ended early at chunk ${request.chunkIndex}`); } - const payload = value; - const expected = Math.min(meta.chunkSize, meta.numGaussians - request.chunkIndex * meta.chunkSize); - if (payload.count !== expected) { - throw new Error(`decimate output chunk ${request.chunkIndex}: expected ${expected} rows, produced ${payload.count}`); + if (value !== expected) { + throw new Error(`decimate output chunk ${request.chunkIndex}: expected ${expected} rows, produced ${value}`); } - const fill = (cd: ChunkData | undefined, layer: ChunkLayer): void => { - if (!cd) return; - const src = payload[layer as 'position' | 'geometric' | 'color' | 'other']; - if (!src) { - throw new Error(`decimate output has no '${layer}' layer`); - } - const bytes = payload.count * cd.stride; - new Uint8Array(cd.data, 0, bytes).set(new Uint8Array(src.buffer, src.byteOffset, bytes)); - }; - fill(request.position, 'position'); - fill(request.geometric, 'geometric'); - fill(request.color, 'color'); - fill(request.other, 'other'); - nextChunk++; if (nextChunk >= (meta.numChunks[0] ?? 0)) done = true; }; @@ -91,4 +90,4 @@ const createBlockProducerSource = ( return { meta, read, close }; }; -export { createBlockProducerSource, type ChunkPayload }; +export { createBlockProducerSource, type DestBuffers }; diff --git a/src/lib/decimate/merge-stream.ts b/src/lib/decimate/merge-stream.ts index 12543c40..15c163db 100644 --- a/src/lib/decimate/merge-stream.ts +++ b/src/lib/decimate/merge-stream.ts @@ -1,4 +1,4 @@ -import { type ChunkPayload } from './block-producer'; +import { type DestBuffers } from './block-producer'; import { type ResidentPositions } from './partition'; import { gatherBlockView, indexOfSorted, type PriorityContext } from './priority'; import { type SelectionResult } from './select'; @@ -15,22 +15,28 @@ type MergeStreamContext = Pick void -): AsyncGenerator { +): AsyncGenerator { const { source, pos, order, blocks, selection, nextPositions } = ctx; const { memberGroup, groupMin, groupOffsets, groupMembers } = selection; const { layouts, availableLayers } = source.meta; @@ -39,25 +45,9 @@ async function *mergeStream( const hasOther = availableLayers.has('other') && (layouts.other?.stride ?? 0) > 0; const otherDim = hasOther ? layouts.other!.stride >> 2 : 0; - // Rolling output buffers (reused across payloads: the consumer copies - // before pulling the next chunk). - const outPos = new Float32Array(chunkSize * 3); - const outGeo = new Float32Array(chunkSize * 8); - const outColor = new Float32Array(chunkSize * colorDim); - const outOther = hasOther ? new Uint32Array(chunkSize * otherDim) : undefined; let rows = 0; let emitted = 0; - - const payload = (): ChunkPayload => { - const p: ChunkPayload = { - count: rows, - position: outPos.subarray(0, rows * 3), - geometric: outGeo.subarray(0, rows * 8), - color: outColor.subarray(0, rows * colorDim) - }; - if (outOther) p.other = outOther.subarray(0, rows * otherDim); - return p; - }; + let dest = yield 0; // priming handshake: the first real chunk's views for (let bi = 0; bi < blocks.length; bi++) { const block = blocks[bi]; @@ -65,21 +55,28 @@ async function *mergeStream( const nOwned = owned.length; // This block's emitted groups (min member owned here), in owned order, - // and the out-of-block members they pull in. + // and the out-of-block members they pull in — count, collect, sort, + // dedup (members are unique across groups by construction). const blockGroups: number[] = []; - const extSet = new Map(); + let extCount = 0; for (let i = 0; i < nOwned; i++) { const g = owned[i]; const mg = memberGroup[g]; if (mg === -1 || groupMin[mg] !== g) continue; blockGroups.push(mg); + for (let m = groupOffsets[mg]; m < groupOffsets[mg + 1]; m++) { + if (indexOfSorted(owned, groupMembers[m]) < 0) extCount++; + } + } + const extraGlobals = new Uint32Array(extCount); + extCount = 0; + for (const mg of blockGroups) { for (let m = groupOffsets[mg]; m < groupOffsets[mg + 1]; m++) { const member = groupMembers[m]; - if (indexOfSorted(owned, member) < 0 && !extSet.has(member)) extSet.set(member, 0); + if (indexOfSorted(owned, member) < 0) extraGlobals[extCount++] = member; } } - const extraGlobals = Uint32Array.from(extSet.keys()).sort(); - for (let i = 0; i < extraGlobals.length; i++) extSet.set(extraGlobals[i], nOwned + i); + extraGlobals.sort(); const { view, other } = await gatherBlockView(ctx, bi, extraGlobals, hasOther); @@ -103,7 +100,7 @@ async function *mergeStream( for (let m = groupOffsets[mg]; m < groupOffsets[mg + 1]; m++) { const member = groupMembers[m]; const oi = indexOfSorted(owned, member); - const row = oi >= 0 ? oi : extSet.get(member)!; + const row = oi >= 0 ? oi : nOwned + indexOfSorted(extraGlobals, member); mPos[mi * 3] = view.pos[row * 3]; mPos[mi * 3 + 1] = view.pos[row * 3 + 1]; mPos[mi * 3 + 2] = view.pos[row * 3 + 2]; @@ -130,39 +127,46 @@ async function *mergeStream( mergedOther = merged.other; } - // Emit rows in owned order. + // Emit rows in owned order, straight into the destination views. let nextMerged = 0; for (let i = 0; i < nOwned; i++) { const g = owned[i]; const mg = memberGroup[g]; if (mg !== -1 && groupMin[mg] !== g) continue; // consumed member + let px: number, py: number, pz: number; if (mg === -1) { // Survivor pass-through: position from resident arrays, // geometric/color/other block-copied from the view. - outPos[rows * 3] = pos.x[g]; - outPos[rows * 3 + 1] = pos.y[g]; - outPos[rows * 3 + 2] = pos.z[g]; - outGeo.set(view.geo.subarray(i * 8, i * 8 + 8), rows * 8); - outColor.set(view.color.subarray(i * colorDim, (i + 1) * colorDim), rows * colorDim); - if (outOther) outOther.set(other!.subarray(i * otherDim, (i + 1) * otherDim), rows * otherDim); + px = pos.x[g]; + py = pos.y[g]; + pz = pos.z[g]; + if (dest.geometric) dest.geometric.set(view.geo.subarray(i * 8, i * 8 + 8), rows * 8); + if (dest.color) dest.color.set(view.color.subarray(i * colorDim, (i + 1) * colorDim), rows * colorDim); + if (dest.other) dest.other.set(other!.subarray(i * otherDim, (i + 1) * otherDim), rows * otherDim); } else { const mi = nextMerged++; - outPos.set(mergedPos!.subarray(mi * 3, mi * 3 + 3), rows * 3); - outGeo.set(mergedGeo!.subarray(mi * 8, mi * 8 + 8), rows * 8); - outColor.set(mergedColor!.subarray(mi * colorDim, (mi + 1) * colorDim), rows * colorDim); - if (outOther) outOther.set(mergedOther!.subarray(mi * otherDim, (mi + 1) * otherDim), rows * otherDim); + px = mergedPos![mi * 3]; + py = mergedPos![mi * 3 + 1]; + pz = mergedPos![mi * 3 + 2]; + if (dest.geometric) dest.geometric.set(mergedGeo!.subarray(mi * 8, mi * 8 + 8), rows * 8); + if (dest.color) dest.color.set(mergedColor!.subarray(mi * colorDim, (mi + 1) * colorDim), rows * colorDim); + if (dest.other) dest.other.set(mergedOther!.subarray(mi * otherDim, (mi + 1) * otherDim), rows * otherDim); + } + if (dest.position) { + dest.position[rows * 3] = px; + dest.position[rows * 3 + 1] = py; + dest.position[rows * 3 + 2] = pz; } - if (nextPositions) { - nextPositions.x[emitted] = outPos[rows * 3]; - nextPositions.y[emitted] = outPos[rows * 3 + 1]; - nextPositions.z[emitted] = outPos[rows * 3 + 2]; + nextPositions.x[emitted] = px; + nextPositions.y[emitted] = py; + nextPositions.z[emitted] = pz; } rows++; emitted++; if (rows === chunkSize) { - yield payload(); + dest = yield rows; rows = 0; } } @@ -170,8 +174,7 @@ async function *mergeStream( } if (rows > 0) { - yield payload(); - rows = 0; + yield rows; } } diff --git a/test/decimate-merge-stream.test.mjs b/test/decimate-merge-stream.test.mjs index 199a0457..4eff61f1 100644 --- a/test/decimate-merge-stream.test.mjs +++ b/test/decimate-merge-stream.test.mjs @@ -15,6 +15,28 @@ import { kdPartition } from '../src/lib/decimate/partition.js'; import { runPriorityPass } from '../src/lib/decimate/priority.js'; import { selectMerges } from '../src/lib/decimate/select.js'; +// Drive the dest-filling generator the way the block producer does: prime, +// then hand each chunk fresh destination views and collect the filled rows. +const drain = async (gen, chunkSize, colorDim, otherDim = 0) => { + const rows = { pos: [], geo: [], color: [], other: [] }; + await gen.next(); // prime + for (;;) { + const dest = { + position: new Float32Array(chunkSize * 3), + geometric: new Float32Array(chunkSize * 8), + color: new Float32Array(chunkSize * colorDim) + }; + if (otherDim) dest.other = new Uint32Array(chunkSize * otherDim); + const { value, done } = await gen.next(dest); + if (done || value === undefined) break; + rows.pos.push(dest.position.slice(0, value * 3)); + rows.geo.push(dest.geometric.slice(0, value * 8)); + rows.color.push(dest.color.slice(0, value * colorDim)); + if (otherDim) rows.other.push(dest.other.slice(0, value * otherDim)); + } + return rows; +}; + describe('mergeStream', () => { it('emits exact count; merged rows equal direct mergeGroup; survivors pass through', async () => { const n = 1200, k = 16, K = 4, target = 700; @@ -38,12 +60,7 @@ describe('mergeStream', () => { y: new Float32Array(outCount), z: new Float32Array(outCount) }; - const rows = { pos: [], geo: [], color: [] }; - for await (const payload of mergeStream({ ...ctx, selection: sel, nextPositions }, 256)) { - rows.pos.push(new Float32Array(payload.position)); - rows.geo.push(new Float32Array(payload.geometric)); - rows.color.push(new Float32Array(payload.color)); - } + const rows = await drain(mergeStream({ ...ctx, selection: sel, nextPositions }, 256), 256, view.colorDim); const emitted = rows.pos.reduce((a, p) => a + p.length / 3, 0); assert.strictEqual(emitted, outCount); @@ -103,12 +120,8 @@ describe('mergeStream', () => { await runPriorityPass(ctx, cand); const sel = selectMerges(cand, n, K, n - target); - const rowsOther = []; - for await (const payload of mergeStream({ ...ctx, selection: sel }, 128)) { - assert.ok(payload.other, 'other layer present'); - rowsOther.push(new Uint32Array(payload.other)); - } - const flatOther = Uint32Array.from(rowsOther.flatMap(a => [...a])); + const rows = await drain(mergeStream({ ...ctx, selection: sel }, 128), 128, 3, otherDim); + const flatOther = Uint32Array.from(rows.other.flatMap(a => [...a])); assert.strictEqual(flatOther.length, (n - sel.removed) * otherDim); // survivors keep their tag verbatim let row = 0; @@ -152,18 +165,13 @@ describe('createBlockProducerSource', () => { }; }; - const onePayload = () => ({ - count: 1, - position: new Float32Array(3), - geometric: new Float32Array(8), - color: new Float32Array(3) - }); - it('rejects gather reads and out-of-order chunk reads', async () => { const meta = await makeMeta(); async function* gen() { - yield onePayload(); - yield onePayload(); + let dest = yield 0; + dest = yield 1; + void dest; + yield 1; } const src = createBlockProducerSource(meta, gen); await assert.rejects( @@ -174,15 +182,16 @@ describe('createBlockProducerSource', () => { await src.close(); }); - it('serves sequential chunk reads and copies payload bytes', async () => { + it('serves sequential chunk reads by filling the destination buffers', async () => { const meta = await makeMeta(); const { createChunkDataPool } = await import('../src/lib/chunk/index.js'); const pool = createChunkDataPool({ chunkSize: 1 }); async function* gen() { - const p = onePayload(); - p.position[0] = 42; - yield p; - yield onePayload(); + let dest = yield 0; + dest.position[0] = 42; + dest = yield 1; + void dest; + yield 1; } const src = createBlockProducerSource(meta, gen); const cd = pool.acquire('position', meta.layouts.position, 1); From 338a2462ad688068768f411d8c2d988c0ff1b043 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 28 Jul 2026 18:39:27 +0100 Subject: [PATCH 12/19] latest --- src/lib/decimate/decimate-source.ts | 31 ++- src/lib/decimate/edge-cost-legacy.ts | 70 ++++- src/lib/decimate/priority-legacy.ts | 388 --------------------------- src/lib/decimate/priority.ts | 77 +++++- src/lib/gpu/gpu-edge-cost-legacy.ts | 135 +++++----- 5 files changed, 224 insertions(+), 477 deletions(-) delete mode 100644 src/lib/decimate/priority-legacy.ts diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index cb3e7d6c..6f078c34 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -2,10 +2,11 @@ import { join } from 'pathe'; import { type GraphicsDevice } from 'playcanvas'; import { createBlockProducerSource } from './block-producer'; +import { buildCostCacheLegacy, computeEdgeCostViewLegacy, packGpuCacheLegacy } from './edge-cost-legacy'; import { mergeStream } from './merge-stream'; +import { createMergeScratch, makeGaussianSamples } from './moment-match'; import { kdPartition, coherenceRuns, type ResidentPositions } from './partition'; -import { runPriorityPass, VIEW_GROW, type CandidateArrays } from './priority'; -import { runPriorityPassLegacy } from './priority-legacy'; +import { runPriorityPass, VIEW_GROW, type CandidateArrays, type CostStrategy } from './priority'; import { selectMerges } from './select'; import { selectMergesLegacy } from './select-legacy'; import { selectMergesRecosted, CACHE_STRIDE } from './select-recost'; @@ -16,6 +17,7 @@ import { type ChunkSourceMetadata } from '../chunk'; import { SPLAT_STRIDE } from '../gpu/gpu-edge-cost'; +import { GpuEdgeCostLegacy } from '../gpu/gpu-edge-cost-legacy'; import { type ReadFileSystem } from '../io/read'; import { type FileSystem } from '../io/write'; import { bakeTransform } from '../ops'; @@ -90,6 +92,26 @@ const chooseK = (n: number, budget: number): number => { return estimate(4) <= budget ? 4 : 2; }; +// The pre-study KL-style cost kernel (--decimate-mode legacy): full-SH colour +// L2, single-Monte-Carlo geometric term, its own GPU cache/kernel layouts. +const createLegacyStrategy = (colorDim: number): CostStrategy => { + const Z = makeGaussianSamples(1, 0); + const z = new Float32Array([Z[0][0], Z[0][1], Z[0][2]]); + const scratch = createMergeScratch(); + return { + cacheForGpu: false, + buildCache: view => buildCostCacheLegacy(view), + createGpu(device, capacity, k) { + const gpu = new GpuEdgeCostLegacy(device, capacity, k, colorDim); + return { + execute: (view, _cache, nbRows, outCosts) => gpu.execute(packGpuCacheLegacy(view), nbRows, z, outCosts), + destroy: () => gpu.destroy() + }; + }, + cpuEdge: (view, cache, i, j) => computeEdgeCostViewLegacy(view, cache as ReturnType, i, j, Z, scratch) + }; +}; + // Read the position layer sequentially into resident columns (generation 1 // only; later generations carry positions forward from the merge stream). const extractPositions = async (source: ChunkSource, pool: ChunkDataPool): Promise => { @@ -247,10 +269,11 @@ const decimateSource = async ( const priorityBar = logger.bar('computing merge priorities', N); if (legacy) { - await runPriorityPassLegacy( + await runPriorityPass( { source: src, pool, pos: positions, order, blocks, device, K, k }, cand!, - n => priorityBar.tick(n) + n => priorityBar.tick(n), + createLegacyStrategy(src.meta.layouts.color!.stride >> 2) ); } else { await runPriorityPass( diff --git a/src/lib/decimate/edge-cost-legacy.ts b/src/lib/decimate/edge-cost-legacy.ts index 9ddf25f9..47d8633a 100644 --- a/src/lib/decimate/edge-cost-legacy.ts +++ b/src/lib/decimate/edge-cost-legacy.ts @@ -21,6 +21,16 @@ import { type SplatView, type MergeScratch } from './moment-match'; +import { type EdgeCostCacheLegacy } from '../gpu/gpu-edge-cost-legacy'; + +/** + * Appearance columns per storage chunk of the legacy GPU kernel. The kernel + * exposes three appearance bindings (appA/appB/appC), so the layout holds up + * to 3·APP_CHUNK columns; at 16 the widest chunk reaches the ~2 GB + * per-binding limit around ~33.5M splats. The kernel imports this same + * constant, so its strides and the host packing can't drift. + */ +export const APP_CHUNK = 16; /** * Per-splat derived quantities for the cost function (legacy @@ -199,4 +209,62 @@ const computeEdgeCostViewLegacy = ( return geoCost + cSh; }; -export { buildCostCacheLegacy, computeEdgeCostViewLegacy, type LegacyCostCache }; +// Pack the block view into the GpuEdgeCostLegacy cache layout (legacy packing: +// posScalars 8-wide, rotR from normalized quats, appearance in ≤APP_CHUNK +// column chunks with live-width strides). +const packGpuCacheLegacy = (view: SplatView): EdgeCostCacheLegacy => { + const { pos, geo, color, colorDim } = view; + const n = geo.length / 8; + const posScalars = new Float32Array(n * 8); + const rotR = new Float32Array(n * 9); + const rot = new Float32Array(9); + + for (let i = 0; i < n; i++) { + const i8 = i * 8; + const o = i * 8; + const linAlpha = sigmoid(geo[i8 + 7]); + const sx = Math.max(Math.exp(geo[i8 + 4]), 1e-12); + const sy = Math.max(Math.exp(geo[i8 + 5]), 1e-12); + const sz = Math.max(Math.exp(geo[i8 + 6]), 1e-12); + const vx = sx * sx + 1e-8; + const vy = sy * sy + 1e-8; + const vz = sz * sz + 1e-8; + posScalars[o] = pos[i * 3]; + posScalars[o + 1] = pos[i * 3 + 1]; + posScalars[o + 2] = pos[i * 3 + 2]; + posScalars[o + 3] = linAlpha * ellipsoidArea(sx, sy, sz) + 1e-12; + posScalars[o + 4] = Math.log(Math.max(vx, 1e-30)) + Math.log(Math.max(vy, 1e-30)) + Math.log(Math.max(vz, 1e-30)); + posScalars[o + 5] = vx; + posScalars[o + 6] = vy; + posScalars[o + 7] = vz; + + let qw = geo[i8], qx = geo[i8 + 1], qy = geo[i8 + 2], qz = geo[i8 + 3]; + const invq = 1 / Math.max(Math.hypot(qw, qx, qy, qz), 1e-12); + qw *= invq; qx *= invq; qy *= invq; qz *= invq; + const xx = qx * qx, yy = qy * qy, zz = qz * qz; + const wx = qw * qx, wy = qw * qy, wz = qw * qz; + const xy = qx * qy, xz = qx * qz, yz = qy * qz; + rot[0] = 1 - 2 * (yy + zz); rot[1] = 2 * (xy - wz); rot[2] = 2 * (xz + wy); + rot[3] = 2 * (xy + wz); rot[4] = 1 - 2 * (xx + zz); rot[5] = 2 * (yz - wx); + rot[6] = 2 * (xz - wy); rot[7] = 2 * (yz + wx); rot[8] = 1 - 2 * (xx + yy); + rotR.set(rot, i * 9); + } + + const numChunks = Math.ceil(colorDim / APP_CHUNK); + const appChunks: Float32Array[] = []; + for (let ch = 0; ch < numChunks; ch++) { + const kStart = ch * APP_CHUNK; + const width = Math.min(APP_CHUNK, colorDim - kStart); + const chunk = new Float32Array(n * width); + for (let s = 0; s < n; s++) { + const dst = s * width; + const src = s * colorDim + kStart; + for (let kk = 0; kk < width; kk++) chunk[dst + kk] = color[src + kk]; + } + appChunks.push(chunk); + } + + return { posScalars, rotR, appChunks, numAppCols: colorDim, numSplats: n }; +}; + +export { buildCostCacheLegacy, computeEdgeCostViewLegacy, packGpuCacheLegacy, type LegacyCostCache }; diff --git a/src/lib/decimate/priority-legacy.ts b/src/lib/decimate/priority-legacy.ts deleted file mode 100644 index ee2eccd2..00000000 --- a/src/lib/decimate/priority-legacy.ts +++ /dev/null @@ -1,388 +0,0 @@ -/** - * LEGACY priority pass — the pre-quality-study pipeline (KL-style cost with - * full-SH colour term, appearance-chunk GPU kernel), preserved verbatim from - * main for `--decimate-mode legacy`. Self-contained on purpose: shares only - * the KNN/partition machinery with the current pass. - */ -import { type GraphicsDevice } from 'playcanvas'; - -import { buildCostCacheLegacy, computeEdgeCostViewLegacy } from './edge-cost-legacy'; -import { KNN_SENTINEL } from './knn-core'; -import { createMergeScratch, makeGaussianSamples, sigmoid, ellipsoidArea, type SplatView } from './moment-match'; -import { type BlockRange, type ResidentPositions } from './partition'; -import { buildForest, sortNeighborRows, VIEW_GROW } from './priority'; -import { type ChunkData, type ChunkDataPool, type ChunkSource } from '../chunk'; -import { APP_CHUNK, GpuEdgeCostLegacy, type EdgeCostCacheLegacy } from '../gpu/gpu-edge-cost-legacy'; -import { GpuKnn } from '../gpu/gpu-knn'; -import { WorkerQueue } from '../workers'; - -/** - * Per-gaussian best-K merge candidates, the resident output of the priority - * pass. `idx[g * K + s]` is the global index of gaussian g's s-th cheapest - * candidate (0xFFFFFFFF when absent); `cost[g * K + s]` its cost (+Inf when - * absent). - */ -type CandidateArrays = { - idx: Uint32Array; - cost: Float32Array; -}; - -/** Everything the block passes need: baked single-LOD source + resident state. */ -type PriorityContext = { - source: ChunkSource; - pool: ChunkDataPool; - pos: ResidentPositions; - order: Uint32Array; - blocks: BlockRange[]; - device?: GraphicsDevice; - /** Candidates kept per gaussian (K). */ - K: number; - /** Neighbours per query (16). */ - k: number; -}; - -/** - * A block's gathered splat columns: owned rows first (block order), then the - * requested extra globals. Positions come from the resident arrays, never - * from the source. - */ -type BlockView = { - view: SplatView; - /** u32 `other` columns (extraDim per row), when requested and present. */ - other?: Uint32Array; - otherDim: number; - ownedCount: number; -}; - -/** - * Gather geometric + color (and optionally `other`) for a block's owned rows - * plus `extraGlobals`, into tight column arrays. Reads are batched at the - * pool's chunk size; owned and extra index lists must be sorted ascending - * for gather coalescing. - * - * @param ctx - The pass context. - * @param blockIdx - Which block. - * @param extraGlobals - Sorted out-of-block rows to append after the owned rows. - * @param includeOther - Also gather the `other` layer (merge pass only). - * @returns The gathered block view. - */ -const gatherBlockView = async ( - ctx: Pick, - blockIdx: number, - extraGlobals: Uint32Array, - includeOther = false -): Promise => { - const { source, pool, pos, order, blocks } = ctx; - const block = blocks[blockIdx]; - const owned = order.subarray(block.start, block.end); - const nOwned = owned.length; - const n = nOwned + extraGlobals.length; - const { layouts, availableLayers } = source.meta; - - const colorDim = layouts.color!.stride >> 2; - const wantOther = includeOther && availableLayers.has('other') && (layouts.other?.stride ?? 0) > 0; - const otherDim = wantOther ? layouts.other!.stride >> 2 : 0; - - const view: SplatView = { - pos: new Float32Array(n * 3), - geo: new Float32Array(n * 8), - color: new Float32Array(n * colorDim), - colorDim - }; - const other = wantOther ? new Uint32Array(n * otherDim) : undefined; - - const readInto = async (indices: Uint32Array, rowBase: number): Promise => { - const batch = pool.chunkSize; - for (let off = 0; off < indices.length; off += batch) { - const count = Math.min(batch, indices.length - off); - const geoCd = pool.acquire('geometric', layouts.geometric!, count); - const colCd = pool.acquire('color', layouts.color!, count); - const othCd: ChunkData | undefined = wantOther ? pool.acquire('other', layouts.other!, count) : undefined; - await source.read({ - indices, - indexOffset: off, - count, - geometric: geoCd, - color: colCd, - other: othCd - }); - view.geo.set(new Float32Array(geoCd.data, 0, count * 8), (rowBase + off) * 8); - view.color.set(new Float32Array(colCd.data, 0, count * colorDim), (rowBase + off) * colorDim); - if (othCd) other!.set(new Uint32Array(othCd.data, 0, count * otherDim), (rowBase + off) * otherDim); - geoCd.release(); - colCd.release(); - othCd?.release(); - } - }; - - await readInto(owned, 0); - await readInto(extraGlobals, nOwned); - - for (let i = 0; i < nOwned; i++) { - const g = owned[i]; - view.pos[i * 3] = pos.x[g]; - view.pos[i * 3 + 1] = pos.y[g]; - view.pos[i * 3 + 2] = pos.z[g]; - } - for (let i = 0; i < extraGlobals.length; i++) { - const g = extraGlobals[i]; - const r = nOwned + i; - view.pos[r * 3] = pos.x[g]; - view.pos[r * 3 + 1] = pos.y[g]; - view.pos[r * 3 + 2] = pos.z[g]; - } - - return { view, other, otherDim, ownedCount: nOwned }; -}; - -// Binary search `g` in the sorted array; -1 when absent. -const indexOfSorted = (sorted: Uint32Array, g: number): number => { - let lo = 0, hi = sorted.length - 1; - while (lo <= hi) { - const mid = (lo + hi) >> 1; - const v = sorted[mid]; - if (v === g) return mid; - if (v < g) lo = mid + 1; - else hi = mid - 1; - } - return -1; -}; - -// Pack the block view into the GpuEdgeCostLegacy cache layout (legacy packing: -// posScalars 8-wide, rotR from normalized quats, appearance in ≤APP_CHUNK -// column chunks with live-width strides). -const packGpuCacheLegacy = (view: SplatView): EdgeCostCacheLegacy => { - const { pos, geo, color, colorDim } = view; - const n = geo.length / 8; - const posScalars = new Float32Array(n * 8); - const rotR = new Float32Array(n * 9); - const rot = new Float32Array(9); - - for (let i = 0; i < n; i++) { - const i8 = i * 8; - const o = i * 8; - const linAlpha = sigmoid(geo[i8 + 7]); - const sx = Math.max(Math.exp(geo[i8 + 4]), 1e-12); - const sy = Math.max(Math.exp(geo[i8 + 5]), 1e-12); - const sz = Math.max(Math.exp(geo[i8 + 6]), 1e-12); - const vx = sx * sx + 1e-8; - const vy = sy * sy + 1e-8; - const vz = sz * sz + 1e-8; - posScalars[o] = pos[i * 3]; - posScalars[o + 1] = pos[i * 3 + 1]; - posScalars[o + 2] = pos[i * 3 + 2]; - posScalars[o + 3] = linAlpha * ellipsoidArea(sx, sy, sz) + 1e-12; - posScalars[o + 4] = Math.log(Math.max(vx, 1e-30)) + Math.log(Math.max(vy, 1e-30)) + Math.log(Math.max(vz, 1e-30)); - posScalars[o + 5] = vx; - posScalars[o + 6] = vy; - posScalars[o + 7] = vz; - - let qw = geo[i8], qx = geo[i8 + 1], qy = geo[i8 + 2], qz = geo[i8 + 3]; - const invq = 1 / Math.max(Math.hypot(qw, qx, qy, qz), 1e-12); - qw *= invq; qx *= invq; qy *= invq; qz *= invq; - const xx = qx * qx, yy = qy * qy, zz = qz * qz; - const wx = qw * qx, wy = qw * qy, wz = qw * qz; - const xy = qx * qy, xz = qx * qz, yz = qy * qz; - rot[0] = 1 - 2 * (yy + zz); rot[1] = 2 * (xy - wz); rot[2] = 2 * (xz + wy); - rot[3] = 2 * (xy + wz); rot[4] = 1 - 2 * (xx + zz); rot[5] = 2 * (yz - wx); - rot[6] = 2 * (xz - wy); rot[7] = 2 * (yz + wx); rot[8] = 1 - 2 * (xx + yy); - rotR.set(rot, i * 9); - } - - const numChunks = Math.ceil(colorDim / APP_CHUNK); - const appChunks: Float32Array[] = []; - for (let ch = 0; ch < numChunks; ch++) { - const kStart = ch * APP_CHUNK; - const width = Math.min(APP_CHUNK, colorDim - kStart); - const chunk = new Float32Array(n * width); - for (let s = 0; s < n; s++) { - const dst = s * width; - const src = s * colorDim + kStart; - for (let kk = 0; kk < width; kk++) chunk[dst + kk] = color[src + kk]; - } - appChunks.push(chunk); - } - - return { posScalars, rotR, appChunks, numAppCols: colorDim, numSplats: n }; -}; - -/** - * The priority pass (heavy read 1): per block — exact global KNN, edge costs - * for each owned gaussian's k neighbours, reduction to the best K candidates - * — written into the resident candidate arrays. - * - * @param ctx - The pass context. - * @param cand - Preallocated candidate arrays (`N*K`), filled per block. - * @param tick - Optional progress callback (owned gaussians completed). - */ -const runPriorityPassLegacy = async ( - ctx: PriorityContext, - cand: CandidateArrays, - tick?: (n: number) => void -): Promise => { - const { pos, order, blocks, device, K, k } = ctx; - const Z = makeGaussianSamples(1, 0); - const z = new Float32Array([Z[0][0], Z[0][1], Z[0][2]]); - const colorDim = ctx.source.meta.layouts.color!.stride >> 2; - - let maxOwned = 0; - for (const b of blocks) maxOwned = Math.max(maxOwned, b.end - b.start); - - let gpuKnn: GpuKnn | undefined; - let gpuCost: GpuEdgeCostLegacy | undefined; - let gpuCostCapacity = Math.ceil(maxOwned * VIEW_GROW); - - // The forest is built once per generation; blocks stay the query/IO - // batching unit (see the quality pass for the shared machinery). - const { parts: forest, blockPart } = await buildForest(pos, order, blocks, !device && !WorkerQueue.isInline); - - // 1-deep prefetch: the next block's KNN runs while the current block - // gathers and costs. GpuKnn executions share one set of buffers, so they - // are serialized through `gpuKnnQueue`. - let gpuKnnQueue: Promise = Promise.resolve(); - const prepare = (bi: number): Promise => { - const owned = order.subarray(blocks[bi].start, blocks[bi].end); - const nOwned = owned.length; - const home = blockPart[bi]; - const queryPos = new Float32Array(nOwned * 3); - for (let i = 0; i < nOwned; i++) { - const g = owned[i]; - queryPos[i * 3] = pos.x[g]; - queryPos[i * 3 + 1] = pos.y[g]; - queryPos[i * 3 + 2] = pos.z[g]; - } - if (device) { - const out = new Uint32Array(nOwned * k); - const run = gpuKnnQueue.then(() => gpuKnn!.execute(queryPos, owned, nOwned, out, home)); - gpuKnnQueue = run.catch(() => { /* surfaced by the awaiting block */ }); - return run.then(() => out); - } - const ordered = [forest[home], ...forest.filter((_, i) => i !== home)]; - const per = Math.ceil(nOwned / 4); - const jobs: Promise[] = []; - for (let off = 0; off < nOwned; off += per) { - const cnt = Math.min(per, nOwned - off); - const qp = queryPos.slice(off * 3, (off + cnt) * 3); - const qi = owned.slice(off, off + cnt); - jobs.push(WorkerQueue.run('knnForest', { parts: ordered, queryPos: qp, queryIds: qi, k }, [ - qp.buffer as ArrayBuffer, qi.buffer as ArrayBuffer - ])); - } - return Promise.all(jobs).then((outs) => { - const out = new Uint32Array(nOwned * k); - let at = 0; - for (const o of outs) { - out.set(o, at); - at += o.length; - } - return out; - }); - }; - - try { - if (device) { - gpuKnn = new GpuKnn(device, forest, k); - gpuCost = new GpuEdgeCostLegacy(device, gpuCostCapacity, maxOwned * k, colorDim); - } - - let next: Promise | null = blocks.length > 0 ? prepare(0) : null; - - for (let bi = 0; bi < blocks.length; bi++) { - const nbPromise = next!; - next = bi + 1 < blocks.length ? prepare(bi + 1) : null; - - const owned = order.subarray(blocks[bi].start, blocks[bi].end); - const nOwned = owned.length; - const nb = await nbPromise; - sortNeighborRows(pos, owned, nb, k); - - // Externals: referenced ids outside the owned range, sorted for - // the gather. - const extRow = new Map(); - for (let s = 0; s < nOwned * k; s++) { - const g = nb[s]; - if (g === KNN_SENTINEL) continue; - if (indexOfSorted(owned, g) < 0 && !extRow.has(g)) extRow.set(g, 0); - } - const extraGlobals = Uint32Array.from(extRow.keys()).sort(); - for (let i = 0; i < extraGlobals.length; i++) extRow.set(extraGlobals[i], nOwned + i); - - // Cross-block neighbours aren't bounded by the initial capacity - // estimate, so a pathological block's view can exceed the - // preallocated cost buffers — grow them to the actual view size - // when that happens (rare; costs one reallocation). - const viewN = nOwned + extraGlobals.length; - if (gpuCost && viewN > gpuCostCapacity) { - gpuCost.destroy(); - gpuCostCapacity = Math.ceil(viewN * 1.1); - gpuCost = new GpuEdgeCostLegacy(device!, gpuCostCapacity, maxOwned * k, colorDim); - } - - const { view } = await gatherBlockView(ctx, bi, extraGlobals); - - // Edge lists in owned-major order (view-local endpoints). - const edgeI = new Uint32Array(nOwned * k); - const edgeJ = new Uint32Array(nOwned * k); - const edgeNb = new Uint32Array(nOwned * k); // global neighbour per edge - const edgeOf = new Uint32Array(nOwned + 1); // CSR into the edge list per owned row - let e = 0; - for (let qi = 0; qi < nOwned; qi++) { - edgeOf[qi] = e; - for (let s = 0; s < k; s++) { - const g = nb[qi * k + s]; - if (g === KNN_SENTINEL) continue; - const oi = indexOfSorted(owned, g); - edgeI[e] = qi; - edgeJ[e] = oi >= 0 ? oi : extRow.get(g)!; - edgeNb[e] = g; - e++; - } - } - edgeOf[nOwned] = e; - - const costs = new Float32Array(e); - if (device) { - await gpuCost!.execute(packGpuCacheLegacy(view), edgeI.subarray(0, e), edgeJ.subarray(0, e), z, costs); - } else { - const cache = buildCostCacheLegacy(view); - const scratch = createMergeScratch(); - for (let i = 0; i < e; i++) { - costs[i] = computeEdgeCostViewLegacy(view, cache, edgeI[i], edgeJ[i], Z, scratch); - } - } - - // Reduce to best K candidates per owned gaussian (ascending by cost). - const bestIdx = new Uint32Array(K); - const bestCost = new Float64Array(K); - for (let qi = 0; qi < nOwned; qi++) { - let size = 0; - for (let s = edgeOf[qi]; s < edgeOf[qi + 1]; s++) { - const c = costs[s]; - if (!Number.isFinite(c)) continue; - if (size === K && c >= bestCost[K - 1]) continue; - let at = size < K ? size : K - 1; - while (at > 0 && bestCost[at - 1] > c) { - bestCost[at] = bestCost[at - 1]; - bestIdx[at] = bestIdx[at - 1]; - at--; - } - bestCost[at] = c; - bestIdx[at] = edgeNb[s]; - size = Math.min(size + 1, K); - } - const g = owned[qi]; - for (let s = 0; s < K; s++) { - cand.idx[g * K + s] = s < size ? bestIdx[s] : 0xFFFFFFFF; - cand.cost[g * K + s] = s < size ? bestCost[s] : Infinity; - } - } - - tick?.(nOwned); - } - } finally { - gpuKnn?.destroy(); - gpuCost?.destroy(); - } -}; - -export { runPriorityPassLegacy }; diff --git a/src/lib/decimate/priority.ts b/src/lib/decimate/priority.ts index 80e07a66..46885ad0 100644 --- a/src/lib/decimate/priority.ts +++ b/src/lib/decimate/priority.ts @@ -32,6 +32,49 @@ type CandidateArrays = { cost: Float32Array; }; +/** A per-slot GPU cost engine (one thread per dense neighbour slot). */ +type GpuSlotCost = { + execute(view: SplatView, cache: unknown, nbRows: Uint32Array, outCosts: Float32Array): Promise; + destroy(): void; +}; + +/** + * The cost kernel the priority pass runs: quality (field-L2, DC-only + * colour) or legacy (KL-style, full-SH colour). Strategies own their cache + * shape (each implementation casts what it created). + */ +type CostStrategy = { + /** Colour components the block view needs (undefined = all). */ + colorComponents?: number; + /** Whether the CPU cache is also the GPU upload payload (quality's 16-f32 rows). */ + cacheForGpu: boolean; + /** Build the per-view CPU cost cache. */ + buildCache(view: SplatView): unknown; + /** Per-slot GPU engine factory (capacity = max view rows). */ + createGpu(device: GraphicsDevice, capacity: number, k: number): GpuSlotCost; + /** Evaluate one edge on the CPU. */ + cpuEdge(view: SplatView, cache: unknown, i: number, j: number): number; +}; + +/** The production field-L2 kernel (see edge-cost-cpu.ts / gpu-edge-cost.ts). */ +const qualityStrategy: CostStrategy = { + colorComponents: 3, + cacheForGpu: true, + buildCache(view) { + const cache = new Float32Array((view.geo.length / 8) * CACHE_STRIDE); + buildSplatCache(view, cache); + return cache; + }, + createGpu(device, capacity, k) { + const gpu = new GpuEdgeCost(device, capacity, k); + return { + execute: (view, cache, nbRows, outCosts) => gpu.execute(cache as Float32Array, view.geo.length / 8, nbRows, outCosts), + destroy: () => gpu.destroy() + }; + }, + cpuEdge: (view, cache, i, j) => computeEdgeCost(cache as Float32Array, i, j) +}; + /** Everything the block passes need: baked single-LOD source + resident state. */ type PriorityContext = { source: ChunkSource; @@ -283,11 +326,13 @@ const sortNeighborRows = ( * @param cand - Candidate arrays (`N*K`) to fill, or undefined to skip all * cost work (cacheOut/neighborsOut persistence only). * @param tick - Optional progress callback (owned gaussians completed). + * @param strategy - The cost kernel (default: the production field-L2). */ const runPriorityPass = async ( ctx: PriorityContext, cand: CandidateArrays | undefined, - tick?: (n: number) => void + tick?: (n: number) => void, + strategy: CostStrategy = qualityStrategy ): Promise => { const { pos, order, blocks, device, K, k } = ctx; @@ -295,7 +340,7 @@ const runPriorityPass = async ( for (const b of blocks) maxOwned = Math.max(maxOwned, b.end - b.start); let gpuKnn: GpuKnn | undefined; - let gpuCost: GpuEdgeCost | undefined; + let gpuCost: GpuSlotCost | undefined; let gpuCostCapacity = Math.ceil(maxOwned * VIEW_GROW); // The forest is built once per generation (its trees are exact and @@ -353,7 +398,7 @@ const runPriorityPass = async ( try { if (device) { gpuKnn = new GpuKnn(device, forest, k); - if (cand) gpuCost = new GpuEdgeCost(device, gpuCostCapacity, k); + if (cand) gpuCost = strategy.createGpu(device, gpuCostCapacity, k); } let next: Promise | null = blocks.length > 0 ? prepare(0) : null; @@ -406,15 +451,15 @@ const runPriorityPass = async ( if (gpuCost && viewN > gpuCostCapacity) { gpuCost.destroy(); gpuCostCapacity = Math.ceil(viewN * 1.1); - gpuCost = new GpuEdgeCost(device!, gpuCostCapacity, k); + gpuCost = strategy.createGpu(device!, gpuCostCapacity, k); } - const { view } = await gatherBlockView(ctx, bi, extraGlobals, false, 3); + const { view } = await gatherBlockView(ctx, bi, extraGlobals, false, strategy.colorComponents); - // Per-splat cost cache, built once for the GPU upload, the CPU - // path, and the re-costed selection's resident copy alike. - const cache = new Float32Array(viewN * CACHE_STRIDE); - buildSplatCache(view, cache); + // Per-view cost cache: quality's 16-f32 rows double as the GPU + // upload and the re-costed selection's resident copy; legacy's is + // CPU-only (its GPU engine packs from the view internally). + const cache = (strategy.cacheForGpu || !device) ? strategy.buildCache(view) : undefined; if (cand) { // Translate neighbour slots: global ids → view rows @@ -432,14 +477,14 @@ const runPriorityPass = async ( } const blockCosts = new Float32Array(slots); - if (device) { - await gpuCost!.execute(cache, viewN, nbRow, blockCosts); + if (gpuCost) { + await gpuCost.execute(view, cache, nbRow, blockCosts); } else { for (let s = 0; s < slots; s++) { const row = nbRow[s]; blockCosts[s] = row === KNN_SENTINEL ? 0 : - computeEdgeCost(cache, (s / k) | 0, row); + strategy.cpuEdge(view, cache, (s / k) | 0, row); } } @@ -476,11 +521,13 @@ const runPriorityPass = async ( // Persist owned rows for re-costed selection: the splat cache // (identical layout on both paths) and the global neighbour ids - // (sentinel-padded). + // (sentinel-padded). Quality-only (cacheOut implies the quality + // strategy, whose cache is the 16-f32 rows). if (ctx.cacheOut) { const CO = ctx.cacheOut; + const c16 = cache as Float32Array; for (let qi = 0; qi < nOwned; qi++) { - CO.set(cache.subarray(qi * CACHE_STRIDE, (qi + 1) * CACHE_STRIDE), owned[qi] * CACHE_STRIDE); + CO.set(c16.subarray(qi * CACHE_STRIDE, (qi + 1) * CACHE_STRIDE), owned[qi] * CACHE_STRIDE); } } if (ctx.neighborsOut) { @@ -504,8 +551,10 @@ export { gatherBlockView, sortNeighborRows, indexOfSorted, + qualityStrategy, VIEW_GROW, type CandidateArrays, + type CostStrategy, type PriorityContext, type BlockView }; diff --git a/src/lib/gpu/gpu-edge-cost-legacy.ts b/src/lib/gpu/gpu-edge-cost-legacy.ts index aed46874..a115b7b5 100644 --- a/src/lib/gpu/gpu-edge-cost-legacy.ts +++ b/src/lib/gpu/gpu-edge-cost-legacy.ts @@ -16,14 +16,9 @@ import { UniformFormat } from 'playcanvas'; -/** - * Appearance columns per storage chunk. The kernel exposes three appearance - * bindings (appA/appB/appC), so the layout holds up to 3·APP_CHUNK columns; at - * 16 the widest chunk reaches the ~2 GB per-binding limit around ~33.5M splats. - * The CPU-side packing in `decimate/priority.ts` imports this same constant, - * so the kernel strides and the host packing can't drift. - */ -export const APP_CHUNK = 16; +import { APP_CHUNK } from '../decimate/edge-cost-legacy'; + +export { APP_CHUNK }; /** * WGSL kernel: per-edge KL-style cost (matches `computeEdgeCostView` in @@ -35,45 +30,47 @@ export const APP_CHUNK = 16; * (the same `z` for both components, matching the CPU implementation), * and adds an L2 distance over the appearance (SH) coefficients. * + * @param k - Compile-time K, neighbour slots per owned row. * @param strideA - Live column count of appearance chunk A (0 if unused). * @param strideB - Live column count of appearance chunk B (0 if unused). * @param strideC - Live column count of appearance chunk C (0 if unused). * @returns WGSL source. */ -const edgeCostWgsl = (strideA: number, strideB: number, strideC: number) => /* wgsl */` +const edgeCostWgsl = (k: number, strideA: number, strideB: number, strideC: number) => /* wgsl */` struct Uniforms { - edgeCount: u32, + slotBase: u32, + slotCount: u32, z0: f32, z1: f32, z2: f32, } @group(0) @binding(0) var uniforms: Uniforms; -// Edge list for the current dispatch batch only, split into two parallel -// arrays (avoids a host-side (i, j) interleave). The host uploads each batch's -// slice to offset 0, so we index edgesI/J[bid] directly — keeping these -// buffers batch-sized instead of N·k keeps them off the per-binding limit. -@group(0) @binding(1) var edgesI: array; -@group(0) @binding(2) var edgesJ: array; +// Neighbour rows for the current dispatch batch (host uploads each batch's +// slice to offset 0): view-local row per slot, 0xFFFFFFFF for empty slots. +// Slot s belongs to owned row (slotBase + s) / K (dense-slot edge model). +@group(0) @binding(1) var nbRow: array; // Per-splat geometry, interleaved 8-wide: // posScalars[8s + 0..2] = position xyz // posScalars[8s + 3] = mass // posScalars[8s + 4] = logdet // posScalars[8s + 5..7] = variances (vx, vy, vz) -@group(0) @binding(3) var posScalars: array; +@group(0) @binding(2) var posScalars: array; // Row-major 3x3 rotation matrix per splat (9 floats per splat). -@group(0) @binding(4) var rotR: array; +@group(0) @binding(3) var rotR: array; // Appearance, split into up to three chunks (≤16 columns each) so no single // binding exceeds maxStorageBufferBindingSize (~2 GB). Each chunk's stride is // its live column count (STRIDE_A/B/C below); appA holds columns 0.., appB the // next span, appC the next. Unused chunks have stride 0, are bound to a dummy // buffer, and are never read. -@group(0) @binding(5) var appA: array; -@group(0) @binding(6) var appB: array; -@group(0) @binding(7) var appC: array; -// Output: cost per edge. -@group(0) @binding(8) var costs: array; - +@group(0) @binding(4) var appA: array; +@group(0) @binding(5) var appB: array; +@group(0) @binding(6) var appC: array; +// Output: cost per slot. +@group(0) @binding(7) var costs: array; + +const K: u32 = ${k}u; +const SENTINEL: u32 = 0xFFFFFFFFu; const EPS_COV: f32 = 1e-8; const LOG2PI: f32 = 1.8378770664093453; // Per-chunk appearance strides = live column count in each chunk (0 = unused, @@ -131,10 +128,16 @@ fn logAddExp(a: f32, b: f32) -> f32 { @compute @workgroup_size(64) fn main(@builtin(global_invocation_id) gid: vec3u) { let bid = gid.x; - if (bid >= uniforms.edgeCount) { return; } - - let i = edgesI[bid]; - let j = edgesJ[bid]; + if (bid >= uniforms.slotCount) { return; } + + // Empty slot: the reduction skips sentinel slots by id, so the cost value + // is never read. + let j = nbRow[bid]; + if (j == SENTINEL) { + costs[bid] = 0.0; + return; + } + let i = (uniforms.slotBase + bid) / K; let i8 = i * 8u; let j8 = j * 8u; @@ -287,15 +290,14 @@ interface EdgeCostCacheLegacy { class GpuEdgeCostLegacy { /** * @param cache - Per-splat cache (uploaded once). - * @param edgeI - Edge u indices (length E). - * @param edgeJ - Edge v indices (length E). + * @param nbRows - Dense neighbour slots (view-local row per slot, + * 0xFFFFFFFF sentinel for empty; slot s belongs to owned row s / k). * @param z - Single Monte-Carlo sample (3 floats from N(0,1)). - * @param outCosts - Destination for per-edge costs (length E). + * @param outCosts - Destination for per-slot costs (length = slots). */ execute: ( cache: EdgeCostCacheLegacy, - edgeI: Uint32Array, - edgeJ: Uint32Array, + nbRows: Uint32Array, z: Float32Array, outCosts: Float32Array ) => Promise; @@ -304,12 +306,13 @@ class GpuEdgeCostLegacy { /** * @param device - PlayCanvas GraphicsDevice (WebGPU). * @param maxN - Maximum number of splats. - * @param maxE - Maximum number of edges in a single dispatch. + * @param k - Neighbour slots per owned row. * @param maxAppCols - Maximum appearance column count (over all bands). */ - constructor(device: GraphicsDevice, maxN: number, maxE: number, maxAppCols: number) { + constructor(device: GraphicsDevice, maxN: number, k: number, maxAppCols: number) { const workgroupSize = 64; - const edgesPerBatch = 1024 * workgroupSize; // 65,536 + // Slots per dispatch: bounded by the 65,535 workgroups-per-dimension limit. + const slotsPerBatch = 65535 * workgroupSize; // 4,194,240 // Appearance is split at fixed APP_CHUNK-column boundaries, but each // chunk's *stride* is its live column count — only the last non-empty // chunk is ever partial, so partial chunks neither allocate nor upload @@ -328,8 +331,7 @@ class GpuEdgeCostLegacy { const bindGroupFormat = new BindGroupFormat(device, [ new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE), - new BindStorageBufferFormat('edgesI', SHADERSTAGE_COMPUTE, true), - new BindStorageBufferFormat('edgesJ', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('nbRow', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('posScalars', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('rotR', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('appA', SHADERSTAGE_COMPUTE, true), @@ -341,11 +343,12 @@ class GpuEdgeCostLegacy { const shader = new Shader(device, { name: 'compute-edge-cost-legacy', shaderLanguage: SHADERLANGUAGE_WGSL, - cshader: edgeCostWgsl(appStrides[0], appStrides[1], appStrides[2]), + cshader: edgeCostWgsl(k, appStrides[0], appStrides[1], appStrides[2]), // @ts-ignore computeUniformBufferFormats: { uniforms: new UniformBufferFormat(device, [ - new UniformFormat('edgeCount', UNIFORMTYPE_UINT), + new UniformFormat('slotBase', UNIFORMTYPE_UINT), + new UniformFormat('slotCount', UNIFORMTYPE_UINT), new UniformFormat('z0', UNIFORMTYPE_FLOAT), new UniformFormat('z1', UNIFORMTYPE_FLOAT), new UniformFormat('z2', UNIFORMTYPE_FLOAT) @@ -390,24 +393,20 @@ class GpuEdgeCostLegacy { appDummy; }); - // Two parallel u32 buffers, sized to a single dispatch batch (not the - // full N·k edge list): execute uploads each batch's slice before its - // dispatch. Batch-sizing keeps these ~256 KB instead of N·k·4 — off the - // ~2 GB per-binding limit (so edges never cap scene size) and ~1.6 GB - // less VRAM at 13M splats. Two parallel arrays avoid a host-side pack. - const edgesIBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); - const edgesJBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); + // Neighbour-row buffer sized to a single dispatch batch (not the full + // N·k slot list): execute uploads each batch's slice before its + // dispatch, keeping it off the ~2 GB per-binding limit. + const nbRowBuf = new StorageBuffer(device, slotsPerBatch * 4, BUFFERUSAGE_COPY_DST); const outBuf = new StorageBuffer( device, - edgesPerBatch * 4, + slotsPerBatch * 4, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST ); - const outScratch = new Float32Array(edgesPerBatch); + const outScratch = new Float32Array(slotsPerBatch); const compute = new Compute(device, shader, 'compute-edge-cost-legacy'); - compute.setParameter('edgesI', edgesIBuf); - compute.setParameter('edgesJ', edgesJBuf); + compute.setParameter('nbRow', nbRowBuf); compute.setParameter('posScalars', posScalarsBuf); compute.setParameter('rotR', rotRBuf); compute.setParameter('appA', appBufs[0]); @@ -417,25 +416,24 @@ class GpuEdgeCostLegacy { this.execute = async ( cache: EdgeCostCacheLegacy, - edgeI: Uint32Array, - edgeJ: Uint32Array, + nbRows: Uint32Array, z: Float32Array, outCosts: Float32Array ) => { const n = cache.numSplats; - const e = edgeI.length; + const s = nbRows.length; if (n > maxN) throw new Error(`GpuEdgeCostLegacy: N=${n} exceeds maxN=${maxN}`); - if (e > maxE) throw new Error(`GpuEdgeCostLegacy: E=${e} exceeds maxE=${maxE}`); if (cache.numAppCols !== maxAppCols) { throw new Error(`GpuEdgeCostLegacy: numAppCols=${cache.numAppCols} must equal maxAppCols=${maxAppCols} (baked into the kernel)`); } if (cache.appChunks.length !== numAppChunks) { throw new Error(`GpuEdgeCostLegacy: cache supplies ${cache.appChunks.length} appearance chunks but the kernel layout expects ${numAppChunks}`); } - if (edgeJ.length !== e || outCosts.length !== e) { - throw new Error('GpuEdgeCostLegacy: edgeI / edgeJ / outCosts must have same length'); + if (outCosts.length !== s) { + throw new Error('GpuEdgeCostLegacy: nbRows / outCosts must have same length'); } + if (s % k !== 0) throw new Error(`GpuEdgeCostLegacy: slot count ${s} must be a multiple of k=${k}`); if (z.length < 3) { throw new Error('GpuEdgeCostLegacy: z must have at least 3 elements'); } @@ -452,25 +450,23 @@ class GpuEdgeCostLegacy { compute.setParameter('z1', z[1]); compute.setParameter('z2', z[2]); - const numBatches = Math.ceil(e / edgesPerBatch); + const numBatches = Math.ceil(s / slotsPerBatch); for (let batch = 0; batch < numBatches; batch++) { - const edgeOffset = batch * edgesPerBatch; - const edgeCount = Math.min(edgesPerBatch, e - edgeOffset); - const groups = Math.ceil(edgeCount / workgroupSize); + const slotBase = batch * slotsPerBatch; + const slotCount = Math.min(slotsPerBatch, s - slotBase); + const groups = Math.ceil(slotCount / workgroupSize); - // Upload just this batch's edges to offset 0; the kernel indexes - // edgesI/J[bid] within the batch. - edgesIBuf.write(0, edgeI, edgeOffset, edgeCount); - edgesJBuf.write(0, edgeJ, edgeOffset, edgeCount); + nbRowBuf.write(0, nbRows, slotBase, slotCount); - compute.setParameter('edgeCount', edgeCount); + compute.setParameter('slotBase', slotBase); + compute.setParameter('slotCount', slotCount); compute.setupDispatch(groups); device.computeDispatch([compute], `edge-cost-dispatch-${batch}`); - const readBytes = edgeCount * 4; + const readBytes = slotCount * 4; await outBuf.read(0, readBytes, outScratch, true); - outCosts.set(outScratch.subarray(0, edgeCount), edgeOffset); + outCosts.set(outScratch.subarray(0, slotCount), slotBase); } }; @@ -481,8 +477,7 @@ class GpuEdgeCostLegacy { if (buf !== appDummy) buf.destroy(); } appDummy.destroy(); - edgesIBuf.destroy(); - edgesJBuf.destroy(); + nbRowBuf.destroy(); outBuf.destroy(); shader.destroy(); bindGroupFormat.destroy(); From a5b75a0b70e1a87cc4c6aa7a19d6fb2815d44713 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Wed, 29 Jul 2026 18:17:51 +0100 Subject: [PATCH 13/19] latest --- src/lib/decimate/block-allocation.ts | 248 +++++++++ src/lib/decimate/block-merge-stream.ts | 134 +++++ src/lib/decimate/block-plan.ts | 368 ++++++++++++++ src/lib/decimate/block-prepare.ts | 101 ++++ src/lib/decimate/decimate-source.ts | 675 ++++++++++++++++++------- src/lib/decimate/partition.ts | 166 +++++- src/lib/decimate/recost-core.ts | 49 +- src/lib/gpu/gpu-recost.ts | 96 +++- test/decimate-block-plan.test.mjs | 284 +++++++++++ test/decimate-multiblock.test.mjs | 80 +++ test/decimate-partition.test.mjs | 59 ++- test/gpu-recost.test.mjs | 63 ++- 12 files changed, 2094 insertions(+), 229 deletions(-) create mode 100644 src/lib/decimate/block-allocation.ts create mode 100644 src/lib/decimate/block-merge-stream.ts create mode 100644 src/lib/decimate/block-plan.ts create mode 100644 src/lib/decimate/block-prepare.ts create mode 100644 test/decimate-block-plan.test.mjs create mode 100644 test/decimate-multiblock.test.mjs diff --git a/src/lib/decimate/block-allocation.ts b/src/lib/decimate/block-allocation.ts new file mode 100644 index 00000000..7c8009b8 --- /dev/null +++ b/src/lib/decimate/block-allocation.ts @@ -0,0 +1,248 @@ +import { join } from 'pathe'; + +import { type BlockPlan } from './block-plan'; +import { type ReadSource, type ReadStream } from '../io/read'; +import { type FileSystem } from '../io/write'; + +const RECORD_BYTES = 12; +const CURSOR_RECORDS = 4096; + +type PlanScratch = { + writeFs: FileSystem; + readFs: { + createSource(filename: string): Promise; + }; + scratchDir: string; + remove?: (path: string) => Promise; +}; + +type StoredBlockPlan = { + path: string; + count: number; + blockIndex: number; +}; + +type PlanRecord = { + a: number; + b: number; + cost: number; +}; + +const encodePlan = (plan: BlockPlan): Uint8Array => { + const count = plan.costs.length; + const bytes = new Uint8Array(count * RECORD_BYTES); + const view = new DataView(bytes.buffer); + for (let i = 0; i < count; i++) { + view.setUint32(i * RECORD_BYTES, plan.pairs[i * 2], true); + view.setUint32(i * RECORD_BYTES + 4, plan.pairs[i * 2 + 1], true); + view.setFloat32(i * RECORD_BYTES + 8, plan.costs[i], true); + } + return bytes; +}; + +/** + * Persist one block-local plan, aborting without publishing on failure. + * + * @param scratch - Scratch filesystem pair. + * @param generation - Generation number. + * @param blockIndex - Block number. + * @param plan - Local plan. + * @returns Stored-plan descriptor. + */ +const storeBlockPlan = async ( + scratch: PlanScratch, + generation: number, + blockIndex: number, + plan: BlockPlan +): Promise => { + const path = join( + scratch.scratchDir, + `.decimate-plan-g${generation}-b${blockIndex}-${Date.now().toString(36)}.tmp` + ); + const writer = await scratch.writeFs.createWriter(path); + try { + await writer.write(encodePlan(plan)); + await writer.close(); + } catch (err) { + await writer.abort(); + throw err; + } + return { path, count: plan.costs.length, blockIndex }; +}; + +const pullExact = async (stream: ReadStream, target: Uint8Array): Promise => { + let read = 0; + while (read < target.length) { + const n = await stream.pull(target.subarray(read)); + if (n === 0) break; + read += n; + } + return read; +}; + +class PlanCursor { + private readonly source: ReadSource; + private readonly stream: ReadStream; + private readonly buffer = new Uint8Array(CURSOR_RECORDS * RECORD_BYTES); + private buffered = 0; + private offset = 0; + private consumed = 0; + + constructor(source: ReadSource, readonly plan: StoredBlockPlan) { + this.source = source; + this.stream = source.read(); + } + + async next(): Promise { + if (this.consumed === this.plan.count) return null; + if (this.offset === this.buffered) { + const remaining = (this.plan.count - this.consumed) * RECORD_BYTES; + this.buffered = await pullExact(this.stream, this.buffer.subarray(0, Math.min(this.buffer.length, remaining))); + this.offset = 0; + if (this.buffered < RECORD_BYTES) throw new Error(`truncated decimation plan '${this.plan.path}'`); + } + const view = new DataView(this.buffer.buffer, this.buffer.byteOffset + this.offset, RECORD_BYTES); + const record = { + a: view.getUint32(0, true), + b: view.getUint32(4, true), + cost: view.getFloat32(8, true) + }; + this.offset += RECORD_BYTES; + this.consumed++; + return record; + } + + close(): void { + this.stream.close(); + this.source.close(); + } +} + +type HeapEntry = PlanRecord & { + cursor: PlanCursor; +}; + +/** + * Exact k-way allocation over block-plan prefixes. Only each block's next + * commit is exposed, so non-monotonic continuation costs and prefix + * dependencies are preserved. + * + * @param plans - Stored block plans. + * @param scratch - Scratch read filesystem. + * @param needed - Exact generation removal quota. + * @param onSelect - Optional observation hook for proof/reference tests. + * @returns Selected prefix per block and achieved removals. + */ +const allocatePlanPrefixes = async ( + plans: StoredBlockPlan[], + scratch: PlanScratch, + needed: number, + onSelect?: (blockIndex: number, localIndex: number, record: PlanRecord) => void +): Promise<{ prefixes: Uint32Array; removed: number }> => { + const prefixes = new Uint32Array(plans.length); + const heap: HeapEntry[] = []; + const cursors: PlanCursor[] = []; + const less = (a: HeapEntry, b: HeapEntry): boolean => a.cost < b.cost || + (a.cost === b.cost && (a.cursor.plan.blockIndex < b.cursor.plan.blockIndex || + (a.cursor.plan.blockIndex === b.cursor.plan.blockIndex && (a.a < b.a || (a.a === b.a && a.b < b.b))))); + const push = (entry: HeapEntry): void => { + let i = heap.length; + heap.push(entry); + while (i > 0) { + const p = (i - 1) >> 1; + if (!less(heap[i], heap[p])) break; + [heap[i], heap[p]] = [heap[p], heap[i]]; + i = p; + } + }; + const pop = (): HeapEntry => { + const out = heap[0]; + const tail = heap.pop()!; + if (heap.length > 0) { + heap[0] = tail; + let i = 0; + for (;;) { + const l = i * 2 + 1; + const r = l + 1; + let m = i; + if (l < heap.length && less(heap[l], heap[m])) m = l; + if (r < heap.length && less(heap[r], heap[m])) m = r; + if (m === i) break; + [heap[i], heap[m]] = [heap[m], heap[i]]; + i = m; + } + } + return out; + }; + + try { + for (const plan of plans) { + if (plan.count === 0) continue; + const source = await scratch.readFs.createSource(plan.path); + const cursor = new PlanCursor(source, plan); + cursors.push(cursor); + const first = await cursor.next(); + if (first) push({ ...first, cursor }); + } + + let removed = 0; + while (removed < needed && heap.length > 0) { + const entry = pop(); + onSelect?.( + entry.cursor.plan.blockIndex, + prefixes[entry.cursor.plan.blockIndex], + entry + ); + prefixes[entry.cursor.plan.blockIndex]++; + removed++; + const next = await entry.cursor.next(); + if (next) push({ ...next, cursor: entry.cursor }); + } + return { prefixes, removed }; + } finally { + for (const cursor of cursors) cursor.close(); + } +}; + +/** + * Read one selected block prefix for replay. + * + * @param stored - Stored plan. + * @param scratch - Scratch read filesystem. + * @param count - Selected prefix length. + * @returns Compact plan prefix. + */ +const readBlockPlanPrefix = async ( + stored: StoredBlockPlan, + scratch: PlanScratch, + count: number +): Promise => { + if (count < 0 || count > stored.count) throw new Error('invalid block plan prefix'); + const pairs = new Uint32Array(count * 2); + const costs = new Float32Array(count); + if (count === 0) return { pairs, costs, frozen: 0, unfrozen: 0 }; + const source = await scratch.readFs.createSource(stored.path); + const stream = source.read(0, count * RECORD_BYTES); + try { + const bytes = new Uint8Array(count * RECORD_BYTES); + if (await pullExact(stream, bytes) !== bytes.length) throw new Error(`truncated decimation plan '${stored.path}'`); + const view = new DataView(bytes.buffer); + for (let i = 0; i < count; i++) { + pairs[i * 2] = view.getUint32(i * RECORD_BYTES, true); + pairs[i * 2 + 1] = view.getUint32(i * RECORD_BYTES + 4, true); + costs[i] = view.getFloat32(i * RECORD_BYTES + 8, true); + } + return { pairs, costs, frozen: 0, unfrozen: 0 }; + } finally { + stream.close(); + source.close(); + } +}; + +export { + allocatePlanPrefixes, + readBlockPlanPrefix, + storeBlockPlan, + type PlanScratch, + type StoredBlockPlan +}; diff --git a/src/lib/decimate/block-merge-stream.ts b/src/lib/decimate/block-merge-stream.ts new file mode 100644 index 00000000..ac45ee55 --- /dev/null +++ b/src/lib/decimate/block-merge-stream.ts @@ -0,0 +1,134 @@ +import { readBlockPlanPrefix, type PlanScratch, type StoredBlockPlan } from './block-allocation'; +import { replayBlockPlan } from './block-plan'; +import { type DestBuffers } from './block-producer'; +import { type ResidentPositions } from './partition'; +import { gatherBlockView, type PriorityContext } from './priority'; +import { WorkerQueue } from '../workers'; + +type BlockMergeStreamContext = Pick & { + plans: StoredBlockPlan[]; + prefixes: Uint32Array; + scratch: PlanScratch; + nextPositions?: ResidentPositions; +}; + +/** + * Replay selected plans and emit one core block at a time. Selection arrays, + * gathered fields, and moment-matching inputs are all block-local. + * + * @param ctx - Generation inputs and selected plan prefixes. + * @param chunkSize - Destination chunk size. + * @param tick - Progress callback. + * @yields Filled destination row counts. + */ +async function *blockPlanMergeStream( + ctx: BlockMergeStreamContext, + chunkSize: number, + tick?: (n: number) => void +): AsyncGenerator { + const { source, pos, order, blocks, plans, prefixes, scratch, nextPositions } = ctx; + const { layouts, availableLayers } = source.meta; + const colorDim = layouts.color!.stride >> 2; + const hasOther = availableLayers.has('other') && (layouts.other?.stride ?? 0) > 0; + const otherDim = hasOther ? layouts.other!.stride >> 2 : 0; + + let rows = 0; + let emitted = 0; + let dest = yield 0; + + for (let bi = 0; bi < blocks.length; bi++) { + const owned = order.subarray(blocks[bi].start, blocks[bi].end); + const prefix = await readBlockPlanPrefix(plans[bi], scratch, prefixes[bi]); + const selection = replayBlockPlan(owned.length, prefix); + const { memberGroup, groupMin, groupOffsets, groupMembers } = selection; + const { view, other } = await gatherBlockView(ctx, bi, new Uint32Array(0), hasOther); + + let mergedPos: Float32Array | null = null; + let mergedGeo: Float32Array | null = null; + let mergedColor: Float32Array | null = null; + let mergedOther: Uint32Array | undefined; + if (selection.mergedGroups > 0) { + const totalMembers = groupMembers.length; + const mPos = new Float32Array(totalMembers * 3); + const mGeo = new Float32Array(totalMembers * 8); + const mColor = new Float32Array(totalMembers * colorDim); + const mOther = hasOther ? new Uint32Array(totalMembers * otherDim) : undefined; + const sizes = new Uint32Array(selection.mergedGroups); + let at = 0; + for (let g = 0; g < selection.mergedGroups; g++) { + sizes[g] = groupOffsets[g + 1] - groupOffsets[g]; + for (let m = groupOffsets[g]; m < groupOffsets[g + 1]; m++) { + const row = groupMembers[m]; + mPos.set(view.pos.subarray(row * 3, row * 3 + 3), at * 3); + mGeo.set(view.geo.subarray(row * 8, row * 8 + 8), at * 8); + mColor.set(view.color.subarray(row * colorDim, (row + 1) * colorDim), at * colorDim); + if (mOther) mOther.set(other!.subarray(row * otherDim, (row + 1) * otherDim), at * otherDim); + at++; + } + } + const transfer: ArrayBuffer[] = [ + mPos.buffer as ArrayBuffer, + mGeo.buffer as ArrayBuffer, + mColor.buffer as ArrayBuffer + ]; + if (mOther) transfer.push(mOther.buffer as ArrayBuffer); + const merged = await WorkerQueue.run('mergeGroups', { + pos: mPos, + geo: mGeo, + color: mColor, + sizes, + colorDim, + other: mOther, + otherDim + }, transfer); + mergedPos = merged.pos; + mergedGeo = merged.geo; + mergedColor = merged.color; + mergedOther = merged.other; + } + + let nextMerged = 0; + for (let i = 0; i < owned.length; i++) { + const group = memberGroup[i]; + if (group !== -1 && groupMin[group] !== i) continue; + + let px: number, py: number, pz: number; + if (group === -1) { + px = pos.x[owned[i]]; + py = pos.y[owned[i]]; + pz = pos.z[owned[i]]; + if (dest.geometric) dest.geometric.set(view.geo.subarray(i * 8, i * 8 + 8), rows * 8); + if (dest.color) dest.color.set(view.color.subarray(i * colorDim, (i + 1) * colorDim), rows * colorDim); + if (dest.other) dest.other.set(other!.subarray(i * otherDim, (i + 1) * otherDim), rows * otherDim); + } else { + const mi = nextMerged++; + px = mergedPos![mi * 3]; + py = mergedPos![mi * 3 + 1]; + pz = mergedPos![mi * 3 + 2]; + if (dest.geometric) dest.geometric.set(mergedGeo!.subarray(mi * 8, mi * 8 + 8), rows * 8); + if (dest.color) dest.color.set(mergedColor!.subarray(mi * colorDim, (mi + 1) * colorDim), rows * colorDim); + if (dest.other) dest.other.set(mergedOther!.subarray(mi * otherDim, (mi + 1) * otherDim), rows * otherDim); + } + if (dest.position) { + dest.position[rows * 3] = px; + dest.position[rows * 3 + 1] = py; + dest.position[rows * 3 + 2] = pz; + } + if (nextPositions) { + nextPositions.x[emitted] = px; + nextPositions.y[emitted] = py; + nextPositions.z[emitted] = pz; + } + rows++; + emitted++; + if (rows === chunkSize) { + dest = yield rows; + rows = 0; + } + } + tick?.(owned.length); + } + if (rows > 0) yield rows; +} + +export { blockPlanMergeStream, type BlockMergeStreamContext }; diff --git a/src/lib/decimate/block-plan.ts b/src/lib/decimate/block-plan.ts new file mode 100644 index 00000000..bfa20e81 --- /dev/null +++ b/src/lib/decimate/block-plan.ts @@ -0,0 +1,368 @@ +import { type GraphicsDevice } from 'playcanvas'; + +import { + bestEdgesForPartition, + partitionBestOut, + type RecostState +} from './recost-core'; +import { MAX_GROUP, type SelectionResult } from './select'; +import { GpuRecost, COMMIT_LOG_STRIDE } from '../gpu/gpu-recost'; + +const NIL = 0xFFFFFFFF; + +/** Quality-critical maximum validated commits between refreshes. */ +const WAVE = 4096; + +type BlockPlan = { + /** Interleaved local core roots `(a, b)` in commit order. */ + pairs: Uint32Array; + /** Marginal costs in commit order. */ + costs: Float32Array; + frozen: number; + unfrozen: number; +}; + +type BlockPlanInputs = { + splatCache: Float32Array; + neighbors: Uint32Array; + D: number; + coreCount: number; + totalCount: number; + device?: GraphicsDevice; + /** Production multi-block planning sets this: CPU is reference/test only. */ + requireGpu?: boolean; +}; + +/** + * Plan every legal core-core commit for one immutable core-plus-halo view. + * Halo rows participate in eligibility only: if their best cost is no worse + * than the best core continuation the root is frozen, and no halo row is + * ever unioned. Reverse adjacency makes that decision dynamic when any + * referenced core member changes ownership. + * + * @param inputs - Local cache/KNN and core boundary. + * @returns Compact commit sequence and freeze statistics. + */ +const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { + const { splatCache: SC, neighbors, D, coreCount, totalCount: N, device, requireGpu = false } = inputs; + if (neighbors.length !== N * D) throw new Error('block plan: malformed neighbour rows'); + if (coreCount < 1 || coreCount > N) throw new Error('block plan: invalid core count'); + + const parent = new Uint32Array(N); + const size = new Uint32Array(N); + const version = new Uint32Array(N); + const mHead = new Uint32Array(N); + const mNext = new Uint32Array(N); + const mTail = new Uint32Array(N); + const lastSeq = new Uint32Array(coreCount); + const frozenState = new Uint8Array(coreCount); + size.fill(1); + version.fill(1); + mNext.fill(NIL); + for (let i = 0; i < N; i++) { + parent[i] = i; + mHead[i] = i; + mTail[i] = i; + } + + const st: RecostState = { + SC, + cands: neighbors, + D, + N, + maxGroup: MAX_GROUP, + parent, + size, + version, + mHead, + mNext + }; + + const find = (x0: number): number => { + let x = x0; + while (parent[x] !== x) { + parent[x] = parent[parent[x]]; + x = parent[x]; + } + return x; + }; + + // Fixed reverse adjacency of the generation-input KNN rows. Only core + // query rows exist; halo rows are immutable and never need refresh roots. + const reverseOffsets = new Uint32Array(N + 1); + for (let q = 0; q < coreCount; q++) { + for (let s = 0; s < D; s++) { + const c = neighbors[q * D + s]; + if (c !== NIL) reverseOffsets[c + 1]++; + } + } + for (let i = 0; i < N; i++) reverseOffsets[i + 1] += reverseOffsets[i]; + const reverseRows = new Uint32Array(reverseOffsets[N]); + const reverseFill = reverseOffsets.slice(0, N); + for (let q = 0; q < coreCount; q++) { + for (let s = 0; s < D; s++) { + const c = neighbors[q * D + s]; + if (c !== NIL) reverseRows[reverseFill[c]++] = q; + } + } + + let heapCap = Math.ceil(coreCount * 1.25) + 16; + let hCost = new Float32Array(heapCap); + let hA = new Uint32Array(heapCap); + let hB = new Uint32Array(heapCap); + let hSeq = new Uint32Array(heapCap); + let hVb = new Uint32Array(heapCap); + let heapSize = 0; + let seqCounter = 0; + + const swap = (i: number, j: number): void => { + let t = hCost[i]; hCost[i] = hCost[j]; hCost[j] = t; + t = hA[i]; hA[i] = hA[j]; hA[j] = t; + t = hB[i]; hB[i] = hB[j]; hB[j] = t; + t = hSeq[i]; hSeq[i] = hSeq[j]; hSeq[j] = t; + t = hVb[i]; hVb[i] = hVb[j]; hVb[j] = t; + }; + const less = (i: number, j: number): boolean => hCost[i] < hCost[j] || + (hCost[i] === hCost[j] && (hA[i] < hA[j] || (hA[i] === hA[j] && hB[i] < hB[j]))); + const heapPush = (cost: number, a: number, b: number, seq: number, vb: number): void => { + if (heapSize === heapCap) { + const next = heapCap * 2; + const nextCost = new Float32Array(next); nextCost.set(hCost); hCost = nextCost; + const nextA = new Uint32Array(next); nextA.set(hA); hA = nextA; + const nextB = new Uint32Array(next); nextB.set(hB); hB = nextB; + const nextSeq = new Uint32Array(next); nextSeq.set(hSeq); hSeq = nextSeq; + const nextVb = new Uint32Array(next); nextVb.set(hVb); hVb = nextVb; + heapCap = next; + } + let i = heapSize++; + hCost[i] = cost; hA[i] = a; hB[i] = b; hSeq[i] = seq; hVb[i] = vb; + while (i > 0) { + const p = (i - 1) >> 1; + if (!less(i, p)) break; + swap(i, p); + i = p; + } + }; + const popped = { cost: 0, a: 0, b: 0, seq: 0, vb: 0 }; + const heapPop = (): boolean => { + if (heapSize === 0) return false; + popped.cost = hCost[0]; popped.a = hA[0]; popped.b = hB[0]; popped.seq = hSeq[0]; popped.vb = hVb[0]; + heapSize--; + if (heapSize > 0) { + hCost[0] = hCost[heapSize]; hA[0] = hA[heapSize]; hB[0] = hB[heapSize]; + hSeq[0] = hSeq[heapSize]; hVb[0] = hVb[heapSize]; + let i = 0; + for (;;) { + const l = i * 2 + 1; + const r = l + 1; + let m = i; + if (l < heapSize && less(l, m)) m = l; + if (r < heapSize && less(r, m)) m = r; + if (m === i) break; + swap(i, m); + i = m; + } + } + return true; + }; + + const pending = new Uint32Array(coreCount); + const queuedRound = new Uint32Array(coreCount); + let pendingCount = coreCount; + let round = 1; + const queueRefresh = (root0: number): void => { + const root = find(root0); + if (root >= coreCount) return; + lastSeq[root] = ++seqCounter; + if (queuedRound[root] === round) return; + queuedRound[root] = round; + pending[pendingCount++] = root; + }; + for (let i = 0; i < coreCount; i++) { + pending[i] = i; + queuedRound[i] = round; + lastSeq[i] = ++seqCounter; + } + + const gpuFits = !!device && D * MAX_GROUP === 64 && GpuRecost.fits(device, N, D, WAVE, coreCount, true); + if (requireGpu && !gpuFits) { + throw new Error( + 'multi-block quality decimation requires a WebGPU block working set that fits the adapter limits ' + + `(core ${coreCount}, halo ${N - coreCount})` + ); + } + const gpu = gpuFits ? new GpuRecost(device!, N, D, MAX_GROUP, WAVE, coreCount) : undefined; + const commitLog = gpu ? new Uint32Array(WAVE * COMMIT_LOG_STRIDE) : undefined; + const outBest = gpu ? new Uint32Array(coreCount * 4) : undefined; + const outCost = gpu ? new Float32Array(outBest!.buffer) : undefined; + + const planPairs: number[] = []; + const planCosts: number[] = []; + let frozen = 0; + let unfrozen = 0; + + const refreshResult = (root: number, corePartner: number, coreCost: number, haloPartner: number, haloCost: number): void => { + const isFrozen = haloPartner !== NIL && haloCost <= coreCost; + if (isFrozen) { + if (frozenState[root] === 0) frozen++; + frozenState[root] = 1; + return; + } + if (frozenState[root] !== 0) unfrozen++; + frozenState[root] = 0; + if (corePartner !== NIL && Number.isFinite(coreCost)) { + heapPush(coreCost, root, corePartner, lastSeq[root], version[corePartner]); + } + }; + + try { + gpu?.init(SC, neighbors); + for (;;) { + let wave = 0; + while (wave < WAVE && heapPop()) { + const a = popped.a; + if (parent[a] !== a || popped.seq !== lastSeq[a]) continue; + const b = popped.b; + if (b === a || b >= coreCount || parent[b] !== b || + version[b] !== popped.vb || size[a] + size[b] > MAX_GROUP) { + queueRefresh(a); + continue; + } + + const keep = size[a] >= size[b] ? a : b; + const lose = keep === a ? b : a; + if (commitLog) { + const o = wave * COMMIT_LOG_STRIDE; + commitLog[o] = lose; + commitLog[o + 1] = keep; + commitLog[o + 2] = mTail[keep]; + commitLog[o + 3] = mHead[lose]; + commitLog[o + 4] = size[keep] + size[lose]; + } + mNext[mTail[keep]] = mHead[lose]; + mTail[keep] = mTail[lose]; + parent[lose] = keep; + size[keep] += size[lose]; + version[keep]++; + planPairs.push(a, b); + planCosts.push(popped.cost); + wave++; + + queueRefresh(keep); + // Any root whose fixed pool references a changed member may + // have a different core/halo minimum, including frozen roots. + for (let m = mHead[keep]; m !== NIL; m = mNext[m]) { + for (let r = reverseOffsets[m]; r < reverseOffsets[m + 1]; r++) { + queueRefresh(reverseRows[r]); + } + } + } + + if (pendingCount === 0 && heapSize === 0) break; + if (gpu && pendingCount > 0) { + await gpu.wave(commitLog!, wave, pending, pendingCount, outBest!); + for (let p = 0; p < pendingCount; p++) { + const root = pending[p]; + if (parent[root] !== root) continue; + refreshResult( + root, + outBest![p * 4], + outCost![p * 4 + 1], + outBest![p * 4 + 2], + outCost![p * 4 + 3] + ); + } + } else { + for (let p = 0; p < pendingCount; p++) { + const root = pending[p]; + if (parent[root] !== root) continue; + bestEdgesForPartition(st, root, coreCount); + refreshResult( + root, + partitionBestOut.corePartner < 0 ? NIL : partitionBestOut.corePartner, + partitionBestOut.coreCost, + partitionBestOut.haloPartner < 0 ? NIL : partitionBestOut.haloPartner, + partitionBestOut.haloCost + ); + } + } + pendingCount = 0; + round++; + } + } finally { + gpu?.destroy(); + } + + return { + pairs: Uint32Array.from(planPairs), + costs: Float32Array.from(planCosts), + frozen, + unfrozen + }; +}; + +/** + * Replay a selected prefix into the standard CSR selection shape, local to + * one core. Only this block-sized structure exists during output. + * + * @param coreCount - Core row count. + * @param plan - Selected plan prefix. + * @returns Local selection. + */ +const replayBlockPlan = (coreCount: number, plan: BlockPlan): SelectionResult => { + const parent = new Uint32Array(coreCount); + const size = new Uint32Array(coreCount).fill(1); + for (let i = 0; i < coreCount; i++) parent[i] = i; + const find = (x0: number): number => { + let x = x0; + while (parent[x] !== x) { + parent[x] = parent[parent[x]]; + x = parent[x]; + } + return x; + }; + for (let i = 0; i < plan.costs.length; i++) { + const a = find(plan.pairs[i * 2]); + const b = find(plan.pairs[i * 2 + 1]); + if (a === b || size[a] + size[b] > MAX_GROUP) throw new Error('invalid decimation block plan replay'); + const keep = size[a] >= size[b] ? a : b; + const lose = keep === a ? b : a; + parent[lose] = keep; + size[keep] += size[lose]; + } + + const memberGroup = new Int32Array(coreCount).fill(-1); + const rootGroup = new Int32Array(coreCount).fill(-1); + let groups = 0; + for (let i = 0; i < coreCount; i++) { + const root = find(i); + if (size[root] > 1) { + if (rootGroup[root] < 0) rootGroup[root] = groups++; + memberGroup[i] = rootGroup[root]; + } + } + const groupOffsets = new Uint32Array(groups + 1); + for (let i = 0; i < coreCount; i++) { + const group = memberGroup[i]; + if (group >= 0) groupOffsets[group + 1]++; + } + for (let g = 0; g < groups; g++) groupOffsets[g + 1] += groupOffsets[g]; + const groupMembers = new Uint32Array(groupOffsets[groups]); + const fill = groupOffsets.slice(0, groups); + for (let i = 0; i < coreCount; i++) { + const group = memberGroup[i]; + if (group >= 0) groupMembers[fill[group]++] = i; + } + const groupMin = new Uint32Array(groups); + for (let g = 0; g < groups; g++) groupMin[g] = groupMembers[groupOffsets[g]]; + return { + groupOffsets, + groupMembers, + memberGroup, + groupMin, + mergedGroups: groups, + removed: plan.costs.length + }; +}; + +export { planBlockMerges, replayBlockPlan, WAVE, type BlockPlan, type BlockPlanInputs }; diff --git a/src/lib/decimate/block-prepare.ts b/src/lib/decimate/block-prepare.ts new file mode 100644 index 00000000..9fe31c58 --- /dev/null +++ b/src/lib/decimate/block-prepare.ts @@ -0,0 +1,101 @@ +import { type GraphicsDevice } from 'playcanvas'; + +import { buildSplatCache, CACHE_STRIDE } from './edge-cost-cpu'; +import { KNN_SENTINEL } from './knn-core'; +import { type BlockView, gatherBlockView, type PriorityContext } from './priority'; +import { GpuKnn } from '../gpu/gpu-knn'; +import { WorkerQueue } from '../workers'; + +type PreparedBlock = BlockView & { + cache: Float32Array; + /** Local row ids; rows `[0, ownedCount)` are core queries. */ + neighbors: Uint32Array; +}; + +/** + * Gather one core-plus-halo view and run exact KNN over that local view on + * WebGPU. The returned candidate rows are fixed for the lifetime of the + * block plan. + * + * @param ctx - Source, partition, and resident positions. + * @param blockIndex - Core block index. + * @param haloGlobals - Sorted immutable halo rows. + * @param device - Required WebGPU device. + * @param k - Neighbours per core query. + * @returns Local view, field-L2 cache, and sentinel-padded KNN rows. + */ +const prepareGpuBlock = async ( + ctx: Pick, + blockIndex: number, + haloGlobals: Uint32Array, + device: GraphicsDevice, + k: number +): Promise => { + const block = await gatherBlockView(ctx, blockIndex, haloGlobals, false, 3); + const { view, ownedCount } = block; + const totalCount = view.pos.length / 3; + const x = new Float32Array(totalCount); + const y = new Float32Array(totalCount); + const z = new Float32Array(totalCount); + const ids = new Uint32Array(totalCount); + for (let i = 0; i < totalCount; i++) { + ids[i] = i; + x[i] = view.pos[i * 3]; + y[i] = view.pos[i * 3 + 1]; + z[i] = view.pos[i * 3 + 2]; + } + const forest = await WorkerQueue.run('buildKdForestPart', { x, y, z, ids, shared: false }, [ + x.buffer as ArrayBuffer, + y.buffer as ArrayBuffer, + z.buffer as ArrayBuffer, + ids.buffer as ArrayBuffer + ]); + + const queryPos = view.pos.slice(0, ownedCount * 3); + const queryIds = new Uint32Array(ownedCount); + for (let i = 0; i < ownedCount; i++) queryIds[i] = i; + const coreNeighbors = new Uint32Array(ownedCount * k); + const gpuKnn = new GpuKnn(device, [forest], k); + try { + await gpuKnn.execute(queryPos, queryIds, ownedCount, coreNeighbors, 0); + } finally { + gpuKnn.destroy(); + } + + // Canonical (distance², local id) order keeps CPU/GPU test references and + // repeated executions deterministic. + const dist = new Float64Array(k); + const cand = new Uint32Array(k); + for (let q = 0; q < ownedCount; q++) { + let count = 0; + const qx = view.pos[q * 3], qy = view.pos[q * 3 + 1], qz = view.pos[q * 3 + 2]; + for (let s = 0; s < k; s++) { + const id = coreNeighbors[q * k + s]; + if (id === KNN_SENTINEL) continue; + const dx = view.pos[id * 3] - qx; + const dy = view.pos[id * 3 + 1] - qy; + const dz = view.pos[id * 3 + 2] - qz; + const d2 = dx * dx + dy * dy + dz * dz; + let at = count; + while (at > 0 && (dist[at - 1] > d2 || (dist[at - 1] === d2 && cand[at - 1] > id))) { + dist[at] = dist[at - 1]; + cand[at] = cand[at - 1]; + at--; + } + dist[at] = d2; + cand[at] = id; + count++; + } + for (let s = 0; s < k; s++) coreNeighbors[q * k + s] = s < count ? cand[s] : KNN_SENTINEL; + } + + // GpuRecost addresses all local rows. Halo rows carry no candidate pools + // because they are immutable and never refreshed. + const neighbors = new Uint32Array(totalCount * k).fill(KNN_SENTINEL); + neighbors.set(coreNeighbors); + const cache = new Float32Array(totalCount * CACHE_STRIDE); + buildSplatCache(view, cache); + return { ...block, cache, neighbors }; +}; + +export { prepareGpuBlock, type PreparedBlock }; diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index 6f078c34..00f80ba4 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -1,13 +1,21 @@ import { join } from 'pathe'; import { type GraphicsDevice } from 'playcanvas'; -import { createBlockProducerSource } from './block-producer'; +import { + allocatePlanPrefixes, + storeBlockPlan, + type StoredBlockPlan +} from './block-allocation'; +import { blockPlanMergeStream } from './block-merge-stream'; +import { planBlockMerges } from './block-plan'; +import { prepareGpuBlock, type PreparedBlock } from './block-prepare'; +import { createBlockProducerSource, type DestBuffers } from './block-producer'; import { buildCostCacheLegacy, computeEdgeCostViewLegacy, packGpuCacheLegacy } from './edge-cost-legacy'; import { mergeStream } from './merge-stream'; import { createMergeScratch, makeGaussianSamples } from './moment-match'; -import { kdPartition, coherenceRuns, type ResidentPositions } from './partition'; +import { buildBlockHalo, kdPartition, coherenceRuns, type ResidentPositions } from './partition'; import { runPriorityPass, VIEW_GROW, type CandidateArrays, type CostStrategy } from './priority'; -import { selectMerges } from './select'; +import { selectMerges, type SelectionResult } from './select'; import { selectMergesLegacy } from './select-legacy'; import { selectMergesRecosted, CACHE_STRIDE } from './select-recost'; import { @@ -20,7 +28,7 @@ import { SPLAT_STRIDE } from '../gpu/gpu-edge-cost'; import { GpuEdgeCostLegacy } from '../gpu/gpu-edge-cost-legacy'; import { type ReadFileSystem } from '../io/read'; import { type FileSystem } from '../io/write'; -import { bakeTransform } from '../ops'; +import { bakeTransform, permuteSource } from '../ops'; import { readPly } from '../readers/read-ply'; import { type DeviceCreator } from '../types'; import { fmtBytes, fmtCount, logger, Transform } from '../utils'; @@ -38,6 +46,9 @@ const MIN_ITERATION_PROGRESS = 0.05; /** Default resident-memory budget steering the candidate-K and re-costed-selection policies. */ const DEFAULT_MEMORY_BUDGET = 48 * 2 ** 30; +/** Conservative host bytes per core row for two overlapping core+halo views. */ +const MULTI_BLOCK_BYTES_PER_CORE = 1024; + // Per-gaussian residency of re-costed selection beyond the base state: splat // cache (16 f32) + neighbour ids (k u32) + integer structure (union-find, // chains, seq/round bookkeeping: 8×u32) + heap (5 arrays × 1.25N ≈ 25 B). @@ -92,6 +103,34 @@ const chooseK = (n: number, budget: number): number => { return estimate(4) <= budget ? 4 : 2; }; +const chooseBlockSize = ( + n: number, + budget: number, + residentInputBytes: number, + outputPositionBytes: number, + device?: GraphicsDevice +): number => { + const residentIndex = n * 16; + const available = Math.max(0, budget - residentInputBytes - residentIndex - outputPositionBytes); + let blockSize = Math.min(BLOCK_SIZE, Math.max(1 << 16, Math.floor(available / MULTI_BLOCK_BYTES_PER_CORE))); + const limits = (device as unknown as { + limits?: { + maxStorageBufferBindingSize?: number; + maxBufferSize?: number; + }; + } | undefined)?.limits; + const bindingLimit = Math.min( + limits?.maxStorageBufferBindingSize ?? Infinity, + limits?.maxBufferSize ?? Infinity + ); + if (Number.isFinite(bindingLimit)) { + // The largest local binding is one half of the core+halo cache: + // approximately coreCount × CACHE_STRIDE × sizeof(f32). + blockSize = Math.min(blockSize, Math.floor(bindingLimit / (CACHE_STRIDE * 4))); + } + return Math.max(1, Math.min(blockSize, n)); +}; + // The pre-study KL-style cost kernel (--decimate-mode legacy): full-SH colour // L2, single-Monte-Carlo geometric term, its own GPU cache/kernel layouts. const createLegacyStrategy = (colorDim: number): CostStrategy => { @@ -203,213 +242,479 @@ const decimateSource = async ( // previous generation was materialized in RAM) — counted by the re-costed // selection gate so the budget covers everything actually resident. let residentInputBytes = 0; + let boundaryRetries = 0; const totalGenerations = Math.max(1, Math.ceil(Math.log2(inputMeta.numGaussians / targetCount))); - for (let generation = 1; ; generation++) { - const N = src.meta.numGaussians; - const gen = logger.group('Decimate generation', { - index: Math.min(generation, totalGenerations), - total: totalGenerations - }); - - positions ??= await extractPositions(src, pool); - - // Device binding-limit clamp: the largest per-binding buffer scales - // with block size, never scene size; halve the block size until it - // fits the adapter's storage-binding limit. - let blockSize = BLOCK_SIZE; - const bindingLimit = (device as unknown as { limits?: { maxStorageBufferBindingSize?: number } } | undefined) - ?.limits?.maxStorageBufferBindingSize; - if (typeof bindingLimit === 'number') { - const largestBinding = (bs: number) => Math.ceil(bs * VIEW_GROW) * SPLAT_STRIDE * 4; - while (blockSize > (1 << 16) && largestBinding(blockSize) > bindingLimit) { - blockSize >>= 1; + try { + for (let generation = 1; ; generation++) { + const N = src.meta.numGaussians; + const gen = logger.group('Decimate generation', { + index: Math.min(generation, totalGenerations), + total: totalGenerations + }); + + positions ??= await extractPositions(src, pool); + + // Quality multi-block working sets are sized from the actual host + // budget and adapter binding limits. Legacy retains its previous + // fixed-size batching policy. + const legacy = opts.mode === 'legacy'; + const nextCount = Math.max(targetCount, N - Math.floor(N / 2)); + let blockSize = legacy ? + BLOCK_SIZE : + chooseBlockSize(N, budget, residentInputBytes, nextCount * 12, device); + if (legacy) { + const bindingLimit = (device as unknown as { limits?: { maxStorageBufferBindingSize?: number } } | undefined) + ?.limits?.maxStorageBufferBindingSize; + if (typeof bindingLimit === 'number') { + const largestBinding = (bs: number) => Math.ceil(bs * VIEW_GROW) * SPLAT_STRIDE * 4; + while (blockSize > (1 << 16) && largestBinding(blockSize) > bindingLimit) blockSize >>= 1; + } } - if (blockSize !== BLOCK_SIZE) { - logger.warn(`reducing decimate block size to ${fmtCount(blockSize)} to fit GPU binding limit ${fmtBytes(bindingLimit)}`); + if (blockSize !== BLOCK_SIZE && N > blockSize) { + logger.info(`decimate core size ${fmtCount(blockSize)} (memory/device working-set limit)`); + } + + const partSub = logger.group('Partitioning'); + let partition = kdPartition(positions, blockSize, legacy ? null : generation); + let { order, blocks } = partition; + partSub.end(); + + if (generation === 1 && N >= COHERENCE_MIN_N) { + const runs = blocks.map(b => coherenceRuns(order, b.start, b.end, COHERENCE_GAP_ROWS)).sort((a, b) => a - b); + const median = runs[runs.length >> 1] ?? 0; + if (median > INCOHERENT_RUNS_PER_BLOCK) { + if (!legacy && blocks.length > 1) { + if (!opts.spill) { + throw new Error( + 'multi-block quality decimation needs scratch storage to stage spatially incoherent input; ' + + 'provide opts.spill / --scratch-dir' + ); + } + const rowBytes = 12 + 32 + colorDim * 4 + otherStride; + logger.info( + 'spatially incoherent input: staging one KD-ordered PLY ' + + `(estimated ${fmtBytes(N * rowBytes)})` + ); + const spill = opts.spill; + const filename = join( + spill.scratchDir, + `.decimate-stage-${Date.now().toString(36)}.tmp.ply` + ); + const stagedView = permuteSource(src, order); + let plySrc: ChunkSource; + try { + await writePlyStreaming(stagedView, pool, { filename }, spill.writeFs); + const readSource = await spill.readFs.createSource(filename); + plySrc = await readPly(readSource, pool); + } catch (err) { + try { + await spill.remove?.(filename); + } catch { + // Preserve the staging failure. + } + throw err; + } + + const reordered: ResidentPositions = { + x: new Float32Array(N), + y: new Float32Array(N), + z: new Float32Array(N) + }; + for (let i = 0; i < N; i++) { + const g = order[i]; + reordered.x[i] = positions.x[g]; + reordered.y[i] = positions.y[g]; + reordered.z[i] = positions.z[g]; + } + + await src.close(); + await disposeCurrentInput?.(); + src = plySrc; + positions = reordered; + disposeCurrentInput = async () => { + await spill.remove?.(filename); + }; + partition = kdPartition(positions, blockSize, generation); + ({ order, blocks } = partition); + } else { + logger.warn( + 'input is spatially incoherent (scattered gathers expected); run a one-time --morton-order prepass for much faster IO' + ); + } + } } - } - const partSub = logger.group('Partitioning'); - const { order, blocks } = kdPartition(positions, blockSize); - partSub.end(); + // Re-costed selection (exact within-generation greedy) when its + // resident state fits the budget alongside the base state; one-shot + // selection otherwise. Gated per generation, so large scenes regain + // re-costing as soon as the cascade shrinks under the budget. Legacy + // mode uses the pre-study pipeline throughout (no re-costing state). + const K = chooseK(N, budget); + const k = Math.min(KNN_K, Math.max(1, N - 1)); + const generationTarget = Math.max(targetCount, N - Math.floor(N / 2)); + const needed = N - generationTarget; + const multiBlock = !legacy && blocks.length > 1; + let selection: SelectionResult | undefined; + let storedPlans: StoredBlockPlan[] | undefined; + let planPrefixes: Uint32Array | undefined; + let removed: number; + + if (multiBlock) { + if (!device) { + throw new Error( + `multi-block quality decimation requires WebGPU (${fmtCount(N)} splats, ` + + `${fmtCount(blockSize)}-splat cores); increase --memory-budget for the one-block path or provide a device` + ); + } + if (!opts.spill) { + throw new Error( + 'multi-block quality decimation needs scratch storage for merge plans ' + + `(approximately ${fmtBytes(N * 12)} this generation); provide opts.spill / --scratch-dir` + ); + } + if (generation === 1) { + const rowBytes = 12 + 32 + colorDim * 4 + otherStride; + logger.info( + `decimate scratch estimate: staging up to ${fmtBytes(N * rowBytes)}; ` + + `merge plans up to ${fmtBytes(N * 12)} per generation` + ); + } - if (generation === 1 && N >= COHERENCE_MIN_N) { - const runs = blocks.map(b => coherenceRuns(order, b.start, b.end, COHERENCE_GAP_ROWS)).sort((a, b) => a - b); - const median = runs[runs.length >> 1] ?? 0; - if (median > INCOHERENT_RUNS_PER_BLOCK) { - logger.warn( - 'input is spatially incoherent (scattered gathers expected); run a one-time --morton-order prepass for much faster IO' + const planBar = logger.bar('planning local merges', N); + storedPlans = new Array(blocks.length); + let cappedHalos = 0; + let frozen = 0; + let unfrozen = 0; + let knnMs = 0; + let refreshMs = 0; + let allocationMs = 0; + let preparedNext: Promise | null = null; + + // eslint-disable-next-line no-loop-func + const prepare = (bi: number): Promise => { + const coreCount = blocks[bi].end - blocks[bi].start; + const halo = buildBlockHalo(positions!, partition, bi, coreCount); + if (halo.capped) cappedHalos++; + const started = Date.now(); + return prepareGpuBlock( + { source: src, pool, pos: positions!, order, blocks }, + bi, + halo.rows, + device, + KNN_K + ).finally(() => { + knnMs += Date.now() - started; + }); + }; + + try { + preparedNext = prepare(0); + for (let bi = 0; bi < blocks.length; bi++) { + const prepared = await preparedNext; + preparedNext = bi + 1 < blocks.length ? prepare(bi + 1) : null; + const refreshStarted = Date.now(); + const plan = await planBlockMerges({ + splatCache: prepared.cache, + neighbors: prepared.neighbors, + D: KNN_K, + coreCount: prepared.ownedCount, + totalCount: prepared.view.pos.length / 3, + device, + requireGpu: true + }); + refreshMs += Date.now() - refreshStarted; + frozen += plan.frozen; + unfrozen += plan.unfrozen; + storedPlans[bi] = await storeBlockPlan(opts.spill, generation, bi, plan); + planBar.tick(prepared.ownedCount); + } + const allocationStarted = Date.now(); + const allocation = await allocatePlanPrefixes(storedPlans, opts.spill, needed); + allocationMs = Date.now() - allocationStarted; + planPrefixes = allocation.prefixes; + removed = allocation.removed; + } catch (err) { + if (preparedNext) { + try { + await preparedNext; + } catch { + // Preserve the active planning failure. + } + } + try { + await Promise.all(storedPlans.filter(Boolean).map(plan => opts.spill!.remove?.(plan.path))); + } catch { + // Preserve the planning failure. + } + throw err; + } finally { + planBar.end(); + } + logger.info( + `local merge stats: ${fmtCount(cappedHalos)} capped halo${cappedHalos === 1 ? '' : 's'}, ` + + `${fmtCount(frozen)} freezes, ${fmtCount(unfrozen)} unfreezes, ${fmtCount(removed!)} removals` + ); + logger.info( + `local timings: KNN/gather ${(knnMs / 1000).toFixed(2)}s, ` + + `refresh/plan ${(refreshMs / 1000).toFixed(2)}s, allocation ${(allocationMs / 1000).toFixed(2)}s` ); + } else { + const baseBytes = residentInputBytes + N * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; + const recost = !legacy && baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; + + // One-block quality deliberately follows the pre-existing path + // with no staging, halos, plan files, or k-way coordination. + const cand: CandidateArrays | undefined = recost ? + undefined : + { + idx: new Uint32Array(N * K).fill(0xFFFFFFFF), + cost: new Float32Array(N * K).fill(Infinity) + }; + const cacheOut = recost ? new Float32Array(N * CACHE_STRIDE) : undefined; + const neighborsOut = recost ? new Uint32Array(N * k) : undefined; + const priorityBar = logger.bar('computing merge priorities', N); + if (legacy) { + await runPriorityPass( + { source: src, pool, pos: positions, order, blocks, device, K, k }, + cand!, + n => priorityBar.tick(n), + createLegacyStrategy(src.meta.layouts.color!.stride >> 2) + ); + } else { + await runPriorityPass( + { source: src, pool, pos: positions, order, blocks, device, K, k, cacheOut, neighborsOut }, + cand, + n => priorityBar.tick(n) + ); + } + priorityBar.end(); + + const selectSub = logger.group(recost ? 'Selecting merges (re-costed)' : 'Selecting merges'); + selection = legacy ? + selectMergesLegacy(cand!, N, K, needed) : + cacheOut ? + await selectMergesRecosted({ splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed, device }) : + selectMerges(cand!, N, K, needed); + selectSub.end(); + removed = selection.removed; } - } - // Re-costed selection (exact within-generation greedy) when its - // resident state fits the budget alongside the base state; one-shot - // selection otherwise. Gated per generation, so large scenes regain - // re-costing as soon as the cascade shrinks under the budget. Legacy - // mode uses the pre-study pipeline throughout (no re-costing state). - const legacy = opts.mode === 'legacy'; - const K = chooseK(N, budget); - const k = Math.min(KNN_K, Math.max(1, N - 1)); - const baseBytes = residentInputBytes + N * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; - const recost = !legacy && baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; - - // Re-costed selection seeds itself from the neighbour graph (wave 0), - // so candidate arrays exist only for the one-shot/legacy selections. - const cand: CandidateArrays | undefined = recost ? - undefined : - { - idx: new Uint32Array(N * K).fill(0xFFFFFFFF), - cost: new Float32Array(N * K).fill(Infinity) + let plansDisposed = false; + const disposePlans = async (): Promise => { + if (plansDisposed || !storedPlans) return; + plansDisposed = true; + await Promise.all(storedPlans.map(plan => opts.spill?.remove?.(plan.path))); }; - const cacheOut = recost ? new Float32Array(N * CACHE_STRIDE) : undefined; - const neighborsOut = recost ? new Uint32Array(N * k) : undefined; - - const priorityBar = logger.bar('computing merge priorities', N); - if (legacy) { - await runPriorityPass( - { source: src, pool, pos: positions, order, blocks, device, K, k }, - cand!, - n => priorityBar.tick(n), - createLegacyStrategy(src.meta.layouts.color!.stride >> 2) - ); - } else { - await runPriorityPass( - { source: src, pool, pos: positions, order, blocks, device, K, k, cacheOut, neighborsOut }, - cand, - n => priorityBar.tick(n) - ); - } - priorityBar.end(); - - const generationTarget = Math.max(targetCount, N - Math.floor(N / 2)); - const needed = N - generationTarget; - const selectSub = logger.group(recost ? 'Selecting merges (re-costed)' : 'Selecting merges'); - const selection = legacy ? - selectMergesLegacy(cand!, N, K, needed) : - cacheOut ? - await selectMergesRecosted({ splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed, device }) : - selectMerges(cand!, N, K, needed); - selectSub.end(); - - if (selection.removed === 0) { - gen.end(); - const cause = device ? - 'the GPU step likely failed (e.g. out-of-memory) or produced non-finite costs' : - 'cost computation produced no finite merge candidates (e.g. non-finite inputs)'; - throw new Error( - `decimation found no valid merges at ${N} splats (target ${targetCount}) — ${cause}. ` + + + if (removed === 0) { + await disposePlans(); + if (multiBlock && boundaryRetries < 8) { + boundaryRetries++; + logger.warn( + `no productive local cores at jitter ${generation}; repartitioning unchanged input ` + + `(${boundaryRetries}/8 boundary retries)` + ); + gen.end(); + continue; + } + gen.end(); + const cause = device ? + 'the GPU step likely failed (e.g. out-of-memory) or produced non-finite costs' : + 'cost computation produced no finite merge candidates (e.g. non-finite inputs)'; + throw new Error( + `decimation found no valid merges at ${N} splats (target ${targetCount}) — ${cause}. ` + 'Refusing to return an incompletely-decimated scene.' - ); - } - const removedFraction = selection.removed / N; - if (selection.removed < needed && removedFraction < MIN_ITERATION_PROGRESS) { - gen.end(); - throw new Error( - `decimation stalled at ${N} splats (target ${targetCount}): a generation removed only ` + - `${selection.removed} splat${selection.removed === 1 ? '' : 's'} (${(removedFraction * 100).toFixed(3)}% of ${N}) — ` + + ); + } + boundaryRetries = 0; + const removedFraction = removed / N; + if (removed < needed && removedFraction < MIN_ITERATION_PROGRESS) { + await disposePlans(); + gen.end(); + throw new Error( + `decimation stalled at ${N} splats (target ${targetCount}): a generation removed only ` + + `${removed} splat${removed === 1 ? '' : 's'} (${(removedFraction * 100).toFixed(3)}% of ${N}) — ` + 'the nearest-neighbour graph is too degenerate to merge further (e.g. many coincident splats). ' + 'Refusing to grind toward the target.' - ); - } + ); + } - const outCount = N - selection.removed; - const outMeta: ChunkSourceMetadata = { - numGaussians: outCount, - numLods: 1, - lodCounts: [outCount], - chunkSize: src.meta.chunkSize, - numChunks: [Math.ceil(outCount / src.meta.chunkSize)], - shBands: src.meta.shBands, - extraColumns: src.meta.extraColumns, - transform: src.meta.transform, - availableLayers: src.meta.availableLayers, - layouts: src.meta.layouts - }; + const outCount = N - removed; + const outMeta: ChunkSourceMetadata = { + numGaussians: outCount, + numLods: 1, + lodCounts: [outCount], + chunkSize: src.meta.chunkSize, + numChunks: [Math.ceil(outCount / src.meta.chunkSize)], + shBands: src.meta.shBands, + extraColumns: src.meta.extraColumns, + transform: src.meta.transform, + availableLayers: src.meta.availableLayers, + layouts: src.meta.layouts + }; - const isFinal = outCount <= targetCount; - const nextPositions: ResidentPositions | undefined = isFinal ? undefined : { - x: new Float32Array(outCount), - y: new Float32Array(outCount), - z: new Float32Array(outCount) - }; + const isFinal = outCount <= targetCount; + const nextPositions: ResidentPositions | undefined = isFinal ? undefined : { + x: new Float32Array(outCount), + y: new Float32Array(outCount), + z: new Float32Array(outCount) + }; - // `src` is reassigned each generation; capture this generation's - // values for the deferred producer closures. - const genSrc = src; - const genChunkSize = genSrc.meta.chunkSize; - const streamCtx = { source: genSrc, pool, pos: positions, order, blocks, selection, nextPositions }; + // `src` is reassigned each generation; capture this generation's + // values for the deferred producer closures. + const genSrc = src; + const genChunkSize = genSrc.meta.chunkSize; + const genPositions = positions; + const createStream = ( + tick: (n: number) => void + ): AsyncGenerator => { + if (storedPlans) { + return blockPlanMergeStream({ + source: genSrc, + pool, + pos: genPositions, + order, + blocks, + plans: storedPlans, + prefixes: planPrefixes!, + scratch: opts.spill!, + nextPositions + }, genChunkSize, tick); + } + return mergeStream({ + source: genSrc, + pool, + pos: genPositions, + order, + blocks, + selection: selection!, + nextPositions + }, genChunkSize, tick); + }; - if (isFinal) { + if (isFinal) { // The producer reads the input lazily while the consumer pulls // chunks: the input chain (and any pending spill) is released on // close. The merge bar lives outside the generation group since // streaming happens after this function returns. - gen.end(); + gen.end(); + const mergeBar = logger.bar('merging', N); + const producer = createBlockProducerSource(outMeta, () => createStream(n => mergeBar.tick(n))); + const disposeSpill = disposeCurrentInput; + let closed = false; + return { + meta: producer.meta, + read: request => producer.read(request), + close: async () => { + if (closed) return; + closed = true; + mergeBar.end(); + try { + await producer.close(); + } finally { + try { + await genSrc.close(); + } finally { + try { + await disposePlans(); + } finally { + await disposeSpill?.(); + } + } + } + } + }; + } + const mergeBar = logger.bar('merging', N); - const producer = createBlockProducerSource(outMeta, () => mergeStream(streamCtx, genChunkSize, n => mergeBar.tick(n))); - const disposeSpill = disposeCurrentInput; - let closed = false; - return { - meta: producer.meta, - read: request => producer.read(request), - close: async () => { - if (closed) return; - closed = true; - mergeBar.end(); - await producer.close(); - await genSrc.close(); - await disposeSpill?.(); + const producer = createBlockProducerSource(outMeta, () => createStream(n => mergeBar.tick(n))); + + // Intermediate generation: materialize (RAM when comfortably within + // budget, else temp PLY spill), then advance the loop. + const estBytes = outCount * (12 + 32 + colorDim * 4 + otherStride); + let nextSrc: ChunkSource; + let disposeNext: (() => Promise) | null = null; + + try { + if (estBytes <= budget / 4) { + nextSrc = await compact(producer, pool); + residentInputBytes = estBytes; + } else { + residentInputBytes = 0; + if (!opts.spill) { + throw new Error( + `decimation intermediate generation needs ${fmtBytes(estBytes)}, over the in-memory budget — ` + + 'a spill location is required (opts.spill / --scratch-dir)' + ); + } + const spill = opts.spill; + const filename = join(spill.scratchDir, `.decimate-gen${generation}.${Date.now().toString(36)}.tmp.ply`); + let plySrc: ChunkSource; + try { + await writePlyStreaming(producer, pool, { filename }, spill.writeFs); + const readSource = await spill.readFs.createSource(filename); + plySrc = await readPly(readSource, pool); + } catch (err) { + try { + await spill.remove?.(filename); + } catch { + // Preserve the generation failure. + } + throw err; + } + nextSrc = plySrc; + disposeNext = async () => { + await spill.remove?.(filename); + }; } - }; - } - - const mergeBar = logger.bar('merging', N); - const producer = createBlockProducerSource(outMeta, () => mergeStream(streamCtx, genChunkSize, n => mergeBar.tick(n))); - - // Intermediate generation: materialize (RAM when comfortably within - // budget, else temp PLY spill), then advance the loop. - const estBytes = outCount * (12 + 32 + colorDim * 4 + otherStride); - let nextSrc: ChunkSource; - let disposeNext: (() => Promise) | null = null; - - if (estBytes <= budget / 4) { - nextSrc = await compact(producer, pool); - residentInputBytes = estBytes; - } else { - residentInputBytes = 0; - if (!opts.spill) { - throw new Error( - `decimation intermediate generation needs ${fmtBytes(estBytes)}, over the in-memory budget — ` + - 'a spill location is required (opts.spill / --scratch-dir)' - ); + } finally { + await disposePlans(); } - const spill = opts.spill; - const filename = join(spill.scratchDir, `.decimate-gen${generation}.${Date.now().toString(36)}.tmp.ply`); - await writePlyStreaming(producer, pool, { filename }, spill.writeFs); - const readSource = await spill.readFs.createSource(filename); - const plySrc = await readPly(readSource, pool); - nextSrc = plySrc; - disposeNext = async () => { - await plySrc.close(); - await spill.remove?.(filename); - }; + try { + mergeBar.end(); + await producer.close(); + + // The consumed input of THIS generation can now be released: + // for generation 1 that is the caller's source (we own it), + // for later generations the previous spill / RAM intermediate. + await src.close(); + await disposeCurrentInput?.(); + } catch (err) { + try { + await nextSrc.close(); + } catch { + // Preserve the transition failure. + } + try { + await disposeNext?.(); + } catch { + // Preserve the transition failure. + } + throw err; + } + disposeCurrentInput = disposeNext; + + positions = nextPositions!; + src = nextSrc; + gen.end(); + } + } catch (err) { + // Construction failures own the current input just like a returned + // decimation source does. Preserve the active error while making a + // best effort to remove any staged/intermediate generation. + try { + await src.close(); + } catch { + // Preserve the construction failure. + } + try { + await disposeCurrentInput?.(); + } catch { + // Preserve the construction failure. } - mergeBar.end(); - await producer.close(); - - // The consumed input of THIS generation can now be released: for - // generation 1 that is the caller's source (we own it), for later - // generations the previous spill / RAM intermediate. - await src.close(); - await disposeCurrentInput?.(); - disposeCurrentInput = disposeNext; - - positions = nextPositions!; - src = nextSrc; - gen.end(); + throw err; } }; diff --git a/src/lib/decimate/partition.ts b/src/lib/decimate/partition.ts index 07dfe5f3..5c8fc4c1 100644 --- a/src/lib/decimate/partition.ts +++ b/src/lib/decimate/partition.ts @@ -19,6 +19,19 @@ type BlockRange = { start: number; end: number; aabb: Float32Array; + /** Rare out-of-fence residual block. */ + residual: boolean; +}; + +type OutlierFence = { + lo: number[]; + hi: number[]; +}; + +type PartitionResult = { + order: Uint32Array; + blocks: BlockRange[]; + fence: OutlierFence; }; /** Outlier fence: expand the sampled per-axis quantile interval this much. */ @@ -32,7 +45,7 @@ const OUTLIER_SAMPLE_CAP = 1 << 20; // Per-axis fence [lo, hi] from strided-sample quantiles: mid ± factor × the // 0.1–99.9% half-spread. An axis with no spread stays unfenced (±Infinity). -const outlierFence = (pos: ResidentPositions): { lo: number[]; hi: number[] } => { +const outlierFence = (pos: ResidentPositions): OutlierFence => { const n = pos.x.length; const stride = Math.max(1, Math.ceil(n / OUTLIER_SAMPLE_CAP)); const cols = [pos.x, pos.y, pos.z]; @@ -68,9 +81,10 @@ const outlierFence = (pos: ResidentPositions): { lo: number[]; hi: number[] } => * * @param pos - Resident positions. * @param blockSize - Maximum gaussians per block. + * @param generation - Generation number used for deterministic split jitter. * @returns The permuted index array and the block ranges over it. */ -const kdPartition = (pos: ResidentPositions, blockSize: number): { order: Uint32Array; blocks: BlockRange[] } => { +const kdPartition = (pos: ResidentPositions, blockSize: number, generation: number | null = null): PartitionResult => { const n = pos.x.length; const order = new Uint32Array(n); for (let i = 0; i < n; i++) order[i] = i; @@ -90,11 +104,11 @@ const kdPartition = (pos: ResidentPositions, blockSize: number): { order: Uint32 return a; }; - const recurse = (start: number, end: number): void => { + const recurse = (start: number, end: number, depth: number, branch: number, residual: boolean): void => { const aabb = aabbOf(start, end); if (end - start <= blockSize) { order.subarray(start, end).sort(); - blocks.push({ start, end, aabb }); + blocks.push({ start, end, aabb, residual }); return; } let axis = 0, ext = -Infinity; @@ -105,18 +119,27 @@ const kdPartition = (pos: ResidentPositions, blockSize: number): { order: Uint32 axis = c; } } - const mid = start + ((end - start) >> 1); + // Alternating 3/8 and 5/8 quantiles move core boundaries between + // generations without weakening the maximum leaf-size guarantee. + // The hash uses only stable structural inputs, so repeated runs are + // deterministic. + const count = end - start; + const upper = generation !== null && ((generation + depth + branch) & 1) !== 0; + const fraction = generation === null ? 1 / 2 : (upper ? 5 / 8 : 3 / 8); + const offset = Math.max(1, Math.min(count - 1, Math.floor(count * fraction))); + const mid = start + offset; quickselect(cols[axis], order.subarray(start, end), mid - start); - recurse(start, mid); - recurse(mid, end); + recurse(start, mid, depth + 1, branch << 1, residual); + recurse(mid, end, depth + 1, (branch << 1) | 1, residual); }; // Residual split: fence classification must stay rare — a scene that is // mostly "outliers" is just sparse, and splitting it would recreate the // stretched-AABB problem inside the residual. let coreEnd = n; - if (n > 0) { - const { lo, hi } = outlierFence(pos); + const fence = outlierFence(pos); + if (n > blockSize) { + const { lo, hi } = fence; let out = 0; for (let i = 0; i < n; i++) { if (cols[0][i] < lo[0] || cols[0][i] > hi[0] || @@ -134,9 +157,118 @@ const kdPartition = (pos: ResidentPositions, blockSize: number): { order: Uint32 } } } - if (coreEnd > 0) recurse(0, coreEnd); - if (coreEnd < n) recurse(coreEnd, n); - return { order, blocks }; + if (coreEnd > 0) recurse(0, coreEnd, 0, 0, false); + if (coreEnd < n) recurse(coreEnd, n, 0, 1, true); + return { order, blocks, fence }; +}; + +const aabbDistanceSquared = (a: Float32Array, b: Float32Array): number => { + let d2 = 0; + for (let c = 0; c < 3; c++) { + const d = a[3 + c] < b[c] ? b[c] - a[3 + c] : (b[3 + c] < a[c] ? a[c] - b[3 + c] : 0); + d2 += d * d; + } + return d2; +}; + +const pointAabbDistanceSquared = (x: number, y: number, z: number, a: Float32Array): number => { + const dx = x < a[0] ? a[0] - x : (x > a[3] ? x - a[3] : 0); + const dy = y < a[1] ? a[1] - y : (y > a[4] ? y - a[4] : 0); + const dz = z < a[2] ? a[2] - z : (z > a[5] ? z - a[5] : 0); + return dx * dx + dy * dy + dz * dz; +}; + +/** + * Build a block-local immutable halo without scanning the scene. Candidate + * blocks are ordered by AABB distance and only their rows are inspected. + * Normal cores use an in-fence density radius; residual flyaway cores gather + * nearest block populations directly so a scene-scale radius is never + * inferred from their extents. + * + * @param pos - Resident positions. + * @param partition - Partition result sharing the sampled outlier fence. + * @param blockIndex - Core block. + * @param maxHaloRows - Device working-set halo cap. + * @returns Sorted global halo rows and cap diagnostics. + */ +const buildBlockHalo = ( + pos: ResidentPositions, + partition: PartitionResult, + blockIndex: number, + maxHaloRows: number +): { rows: Uint32Array; capped: boolean; radius: number } => { + const { order, blocks, fence } = partition; + const core = blocks[blockIndex]; + const coreCount = core.end - core.start; + const cap = Math.max(0, Math.min(coreCount, maxHaloRows)); + if (cap === 0 || blocks.length === 1) return { rows: new Uint32Array(0), capped: false, radius: 0 }; + + const neighbours = blocks + .map((block, index) => ({ block, index, d2: index === blockIndex ? Infinity : aabbDistanceSquared(core.aabb, block.aabb) })) + .filter(v => v.index !== blockIndex) + .sort((a, b) => a.d2 - b.d2 || a.index - b.index); + + const selected: number[] = []; + let capped = false; + let radius = 0; + + if (core.residual) { + for (const { block } of neighbours) { + for (let i = block.start; i < block.end; i++) { + if (selected.length === cap) { + capped = true; + break; + } + selected.push(order[i]); + } + if (selected.length === cap) break; + } + } else { + let inFence = 0; + const a = new Float64Array([Infinity, Infinity, Infinity, -Infinity, -Infinity, -Infinity]); + for (let i = core.start; i < core.end; i++) { + const g = order[i]; + const x = pos.x[g], y = pos.y[g], z = pos.z[g]; + if (x < fence.lo[0] || x > fence.hi[0] || + y < fence.lo[1] || y > fence.hi[1] || + z < fence.lo[2] || z > fence.hi[2]) continue; + inFence++; + a[0] = Math.min(a[0], x); a[1] = Math.min(a[1], y); a[2] = Math.min(a[2], z); + a[3] = Math.max(a[3], x); a[4] = Math.max(a[4], y); a[5] = Math.max(a[5], z); + } + if (inFence > 0) { + const ex = Math.max(0, a[3] - a[0]); + const ey = Math.max(0, a[4] - a[1]); + const ez = Math.max(0, a[5] - a[2]); + const extents = [ex, ey, ez].sort((u, v) => v - u); + const scale = extents[0] > 0 ? + (extents[2] > extents[0] * 1e-6 ? + Math.cbrt(extents[0] * extents[1] * extents[2] / inFence) : + (extents[1] > extents[0] * 1e-6 ? + Math.sqrt(extents[0] * extents[1] / inFence) : + extents[0] / inFence)) : + 0; + radius = 2.5 * scale; + } + const r2 = radius * radius; + for (const { block, d2 } of neighbours) { + if (d2 > r2) break; + for (let i = block.start; i < block.end; i++) { + const g = order[i]; + if (pointAabbDistanceSquared(pos.x[g], pos.y[g], pos.z[g], core.aabb) > r2) continue; + if (selected.length === cap) { + capped = true; + break; + } + selected.push(g); + } + if (selected.length === cap) break; + } + } + + const rows = Uint32Array.from(selected); + rows.sort(); + return { rows, capped, radius }; }; /** @@ -160,4 +292,12 @@ const coherenceRuns = (sortedIndices: Uint32Array, start: number, end: number, m return runs; }; -export { kdPartition, coherenceRuns, type BlockRange, type ResidentPositions }; +export { + kdPartition, + buildBlockHalo, + coherenceRuns, + type BlockRange, + type OutlierFence, + type PartitionResult, + type ResidentPositions +}; diff --git a/src/lib/decimate/recost-core.ts b/src/lib/decimate/recost-core.ts index 8b49842c..f7b041fe 100644 --- a/src/lib/decimate/recost-core.ts +++ b/src/lib/decimate/recost-core.ts @@ -368,6 +368,13 @@ const evalMergeCore = (st: RecostState, A: number, B: number): number => { /** Result of {@link bestEdgeFor} (reused object). */ const bestOut = { partner: -1, vb: 0, cost: 0 }; +const partitionBestOut = { + corePartner: -1, + coreVersion: 0, + coreCost: Infinity, + haloPartner: -1, + haloCost: Infinity +}; const candScratch = new Uint32Array(256); @@ -379,10 +386,16 @@ const candScratch = new Uint32Array(256); * * @param st - Selection state. * @param root - Cluster root to refresh. + * @param coreCount - Rows below this boundary are mutable core candidates. * @returns True when a legal candidate exists (result in {@link bestOut}). */ -const bestEdgeFor = (st: RecostState, root: number): boolean => { +const bestEdgesForPartition = (st: RecostState, root: number, coreCount: number): boolean => { const { SC, cands, D, parent, size, version, maxGroup } = st; + partitionBestOut.corePartner = -1; + partitionBestOut.coreVersion = 0; + partitionBestOut.coreCost = Infinity; + partitionBestOut.haloPartner = -1; + partitionBestOut.haloCost = Infinity; if (maxGroup * D > candScratch.length) { throw new Error(`recost: candidate pool bound ${maxGroup * D} exceeds scratch (${candScratch.length})`); } @@ -417,17 +430,41 @@ const bestEdgeFor = (st: RecostState, root: number): boolean => { const aTerm = sideTerm(SC, aBuf, na, compA); let bc = Infinity, bp = -1, bv = 0; + let hc = Infinity, hp = -1; for (let t = 0; t < cnt; t++) { const c = cbuf[t]; if (sz + size[c] > maxGroup) continue; const d = evalAgainst(st, na, aTerm, c); - if (d < bc) { - bc = d; bp = c; bv = version[c]; + if (c < coreCount) { + if (d < bc || (d === bc && c < bp)) { + bc = d; bp = c; bv = version[c]; + } + } else if (d < hc || (d === hc && c < hp)) { + hc = d; hp = c; } } - if (bp < 0) return false; - bestOut.partner = bp; bestOut.vb = bv; bestOut.cost = bc; + partitionBestOut.corePartner = bp; + partitionBestOut.coreVersion = bv; + partitionBestOut.coreCost = bc; + partitionBestOut.haloPartner = hp; + partitionBestOut.haloCost = hc; + return bp >= 0 || hp >= 0; +}; + +const bestEdgeFor = (st: RecostState, root: number): boolean => { + if (!bestEdgesForPartition(st, root, st.N) || partitionBestOut.corePartner < 0) return false; + bestOut.partner = partitionBestOut.corePartner; + bestOut.vb = partitionBestOut.coreVersion; + bestOut.cost = partitionBestOut.coreCost; return true; }; -export { evalMergeCore, bestEdgeFor, bestOut, NO_CANDIDATE, type RecostState }; +export { + evalMergeCore, + bestEdgeFor, + bestEdgesForPartition, + bestOut, + partitionBestOut, + NO_CANDIDATE, + type RecostState +}; diff --git a/src/lib/gpu/gpu-recost.ts b/src/lib/gpu/gpu-recost.ts index e17ebf1f..0c9154f8 100644 --- a/src/lib/gpu/gpu-recost.ts +++ b/src/lib/gpu/gpu-recost.ts @@ -38,7 +38,13 @@ const MAX_DIM = 65535; // Shared structure access for the replay/refresh kernels: parentMeta[i] = // (parent, size-at-root); chain[i] = (head-at-root, next). -const refreshWgsl = (k: number, maxGroup: number, splitN: number) => /* wgsl */` +const refreshWgsl = ( + k: number, + maxGroup: number, + splitN: number, + coreCount: number, + partitioned: boolean +) => /* wgsl */` struct Uniforms { pendingCount: u32, } @@ -53,12 +59,14 @@ struct Uniforms { @group(0) @binding(5) var parentMeta: array; @group(0) @binding(6) var chain: array; @group(0) @binding(7) var pending: array; -// Per queued root: (best partner or 0xFFFFFFFF, bitcast(cost)). -@group(0) @binding(8) var outBest: array; +// Per queued root: the global path writes (partner, cost); block-local mode +// additionally writes its best immutable-halo (partner, cost). +@group(0) @binding(8) var outBest: array<${partitioned ? 'vec4u' : 'vec2u'}>; const K: u32 = ${k}u; const MAXG: u32 = ${maxGroup}u; const SPLIT: u32 = ${splitN}u; +const CORE_COUNT: u32 = ${coreCount}u; const NONE: u32 = 0xFFFFFFFFu; const NIL: u32 = 0xFFFFFFFFu; const F32_MAX: f32 = 3.4028234663852886e+38; @@ -262,6 +270,8 @@ var wgRawBw: vec3f; var wgATerm: f32; var redCost: array; var redPartner: array; +var redHaloCost: array; +var redHaloPartner: array; @compute @workgroup_size(${WG}) fn main(@builtin(workgroup_id) wgid: vec3u, @builtin(local_invocation_id) lid3: vec3u) { @@ -276,7 +286,9 @@ fn main(@builtin(workgroup_id) wgid: vec3u, @builtin(local_invocation_id) lid3: let root = pending[pIdx]; if (parentMeta[root].x != root) { // Stale queued root (absorbed since queuing) — no result. - outBest[pIdx] = vec2u(NONE, 0u); + outBest[pIdx] = ${partitioned ? + 'vec4u(NONE, bitcast(F32_MAX), NONE, bitcast(F32_MAX))' : + 'vec2u(NONE, bitcast(F32_MAX))'}; wgAbort = 1u; } else { wgAbort = 0u; @@ -309,6 +321,8 @@ fn main(@builtin(workgroup_id) wgid: vec3u, @builtin(local_invocation_id) lid3: // only shift tie order. var cost = F32_MAX; var partner = NONE; + var haloCost = F32_MAX; + var haloPartner = NONE; let mIdx = lid / K; let slot = lid % K; if (!aborted && mIdx < wgCount) { @@ -342,8 +356,13 @@ fn main(@builtin(workgroup_id) wgid: vec3u, @builtin(local_invocation_id) lid3: // NaN loses every comparison → stays unselected (fail-loud: // an all-NaN scene produces no pushes and the caller throws). if (c < F32_MAX) { - cost = c; - partner = r; + if (r < CORE_COUNT) { + cost = c; + partner = r; + } else { + haloCost = c; + haloPartner = r; + } } } } @@ -353,6 +372,8 @@ fn main(@builtin(workgroup_id) wgid: vec3u, @builtin(local_invocation_id) lid3: // lane scheduling. redCost[lid] = cost; redPartner[lid] = partner; + redHaloCost[lid] = haloCost; + redHaloPartner[lid] = haloPartner; workgroupBarrier(); for (var s = ${WG >> 1}u; s > 0u; s >>= 1u) { if (lid < s) { @@ -362,13 +383,23 @@ fn main(@builtin(workgroup_id) wgid: vec3u, @builtin(local_invocation_id) lid3: redCost[lid] = c2; redPartner[lid] = p2; } + let hc2 = redHaloCost[lid + s]; + let hp2 = redHaloPartner[lid + s]; + if (hc2 < redHaloCost[lid] || (hc2 == redHaloCost[lid] && hp2 < redHaloPartner[lid])) { + redHaloCost[lid] = hc2; + redHaloPartner[lid] = hp2; + } } workgroupBarrier(); } if (lid == 0u && !aborted) { var p = redPartner[0]; if (redCost[0] == F32_MAX) { p = NONE; } - outBest[pIdx] = vec2u(p, bitcast(redCost[0])); + var hp = redHaloPartner[0]; + if (redHaloCost[0] == F32_MAX) { hp = NONE; } + outBest[pIdx] = ${partitioned ? + 'vec4u(p, bitcast(redCost[0]), hp, bitcast(redHaloCost[0]))' : + 'vec2u(p, bitcast(redCost[0]))'}; } } `; @@ -455,10 +486,19 @@ class GpuRecost { * @param n - Generation splat count. * @param k - Neighbours per splat. * @param wave - Max commits per wave. + * @param coreCount - Mutable leading rows when sizing block-local output. + * @param allowSmall - Permit small GPU-local blocks that the global path prefers to evaluate inline. * @returns True when all bindings fit. */ - static fits(device: GraphicsDevice, n: number, k: number, wave: number): boolean { - if (n < 1024) return false; // inline is instant below this + static fits( + device: GraphicsDevice, + n: number, + k: number, + wave: number, + coreCount = n, + allowSmall = false + ): boolean { + if (n < 1024 && !allowSmall) return false; // inline is instant below this on the global path const limits = (device as any).limits; const maxBinding = Math.min( typeof limits?.maxStorageBufferBindingSize === 'number' ? limits.maxStorageBufferBindingSize : 128 * 2 ** 20, @@ -467,7 +507,8 @@ class GpuRecost { const splitN = Math.ceil(n / 2); return splitN * 16 * 4 <= maxBinding && // cacheA/B splitN * k * 4 <= maxBinding && // nbA/B - n * 8 <= maxBinding && // parentMeta/chain/outBest + n * (coreCount < n ? 16 : 8) <= maxBinding && // outBest + n * 8 <= maxBinding && // parentMeta/chain wave * COMMIT_LOG_STRIDE * 4 <= maxBinding; } @@ -477,12 +518,15 @@ class GpuRecost { * @param k - Neighbours per splat (the refresh workgroup is k·maxGroup lanes). * @param maxGroup - Group size cap. * @param wave - Max commits per wave (commit log capacity). + * @param coreCount - Mutable leading rows; remaining rows are immutable halo. */ - constructor(device: GraphicsDevice, n: number, k: number, maxGroup: number, wave: number) { + constructor(device: GraphicsDevice, n: number, k: number, maxGroup: number, wave: number, coreCount = n) { if (k * maxGroup !== WG) { throw new Error(`GpuRecost: k·maxGroup must be ${WG} (got ${k}·${maxGroup})`); } const splitN = Math.ceil(n / 2); + const partitioned = coreCount < n; + const outputStride = partitioned ? 16 : 8; const cacheABuf = new StorageBuffer(device, splitN * 16 * 4, BUFFERUSAGE_COPY_DST); const cacheBBuf = new StorageBuffer(device, Math.max(n - splitN, 1) * 16 * 4, BUFFERUSAGE_COPY_DST); @@ -491,7 +535,7 @@ class GpuRecost { const parentMetaBuf = new StorageBuffer(device, n * 8, BUFFERUSAGE_COPY_DST); const chainBuf = new StorageBuffer(device, n * 8, BUFFERUSAGE_COPY_DST); const pendingBuf = new StorageBuffer(device, n * 4, BUFFERUSAGE_COPY_DST); - const outBestBuf = new StorageBuffer(device, n * 8, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST); + const outBestBuf = new StorageBuffer(device, n * outputStride, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST); const commitLogBuf = new StorageBuffer(device, wave * COMMIT_LOG_STRIDE * 4, BUFFERUSAGE_COPY_DST); const initKernel = makeKernel(device, 'recost-init', initWgsl(), ['count'], [ @@ -510,16 +554,22 @@ class GpuRecost { replayKernel.compute.setParameter('parentMeta', parentMetaBuf); replayKernel.compute.setParameter('chain', chainBuf); - const refreshKernel = makeKernel(device, 'recost-refresh', refreshWgsl(k, maxGroup, splitN), ['pendingCount'], [ - ['cacheA', true], - ['cacheB', true], - ['nbA', true], - ['nbB', true], - ['parentMeta', true], - ['chain', true], - ['pending', true], - ['outBest', false] - ]); + const refreshKernel = makeKernel( + device, + 'recost-refresh', + refreshWgsl(k, maxGroup, splitN, coreCount, partitioned), + ['pendingCount'], + [ + ['cacheA', true], + ['cacheB', true], + ['nbA', true], + ['nbB', true], + ['parentMeta', true], + ['chain', true], + ['pending', true], + ['outBest', false] + ] + ); refreshKernel.compute.setParameter('cacheA', cacheABuf); refreshKernel.compute.setParameter('cacheB', cacheBBuf); refreshKernel.compute.setParameter('nbA', nbABuf); @@ -574,7 +624,7 @@ class GpuRecost { device.computeDispatch(computes, 'recost-wave'); // Blocking readback — also the wave's submit boundary. - await outBestBuf.read(0, pendingCount * 8, outBest, true); + await outBestBuf.read(0, pendingCount * outputStride, outBest, true); }; this.destroy = () => { diff --git a/test/decimate-block-plan.test.mjs b/test/decimate-block-plan.test.mjs new file mode 100644 index 00000000..a7fffdef --- /dev/null +++ b/test/decimate-block-plan.test.mjs @@ -0,0 +1,284 @@ +import assert from 'node:assert'; +import { describe, it } from 'node:test'; + +import { makeSyntheticSource } from './helpers/synthetic-source.mjs'; + +import { + allocatePlanPrefixes, + readBlockPlanPrefix, + storeBlockPlan +} from '../src/lib/decimate/block-allocation.js'; +import { blockPlanMergeStream } from '../src/lib/decimate/block-merge-stream.js'; +import { planBlockMerges, replayBlockPlan } from '../src/lib/decimate/block-plan.js'; +import { createBlockProducerSource } from '../src/lib/decimate/block-producer.js'; +import { buildSplatCache, CACHE_STRIDE } from '../src/lib/decimate/edge-cost-cpu.js'; +import { kdPartition } from '../src/lib/decimate/partition.js'; +import { MemoryReadSource } from '../src/lib/io/read/memory-file-system.js'; +import { MemoryFileSystem } from '../src/lib/io/write/memory-file-system.js'; + +const NIL = 0xFFFFFFFF; + +const makeScratch = () => { + const writeFs = new MemoryFileSystem(); + const removed = []; + return { + writeFs, + readFs: { + async createSource(path) { + const bytes = writeFs.results.get(path); + if (!bytes) throw new Error(`missing ${path}`); + return new MemoryReadSource(bytes); + } + }, + scratchDir: 'scratch', + async remove(path) { + removed.push(path); + writeFs.results.delete(path); + }, + removed + }; +}; + +const manualPlan = (costs, pairs) => ({ + costs: Float32Array.from(costs), + pairs: Uint32Array.from(pairs), + frozen: 0, + unfrozen: 0 +}); + +const randomBlock = (seed) => { + let t = seed >>> 0; + const rand = () => { + t += 0x6d2b79f5; + let r = Math.imul(t ^ (t >>> 15), t | 1); + r ^= r + Math.imul(r ^ (r >>> 7), r | 61); + return ((r ^ (r >>> 14)) >>> 0) / 4294967296; + }; + const coreCount = 6; + const n = 7; + const D = 6; + const pos = new Float32Array(n * 3); + const geo = new Float32Array(n * 8); + const color = new Float32Array(n * 3); + for (let i = 0; i < n; i++) { + pos[i * 3] = rand() * 2; + pos[i * 3 + 1] = rand() * 2; + pos[i * 3 + 2] = rand() * 2; + geo[i * 8] = 1; + geo[i * 8 + 4] = geo[i * 8 + 5] = geo[i * 8 + 6] = Math.log(0.03 + rand() * 0.2); + geo[i * 8 + 7] = rand() * 4 - 2; + color[i * 3] = rand() - 0.5; + color[i * 3 + 1] = rand() - 0.5; + color[i * 3 + 2] = rand() - 0.5; + } + const cache = new Float32Array(n * CACHE_STRIDE); + buildSplatCache({ pos, geo, color, colorDim: 3 }, cache); + const neighbors = new Uint32Array(n * D).fill(NIL); + for (let i = 0; i < coreCount; i++) { + let s = 0; + for (let j = 0; j < n; j++) { + if (j !== i) neighbors[i * D + s++] = j; + } + } + return { cache, neighbors, D, coreCount, n }; +}; + +describe('block-local merge planning', () => { + it('freezes an all-halo coincident pool, then merges it when both rows become core', async () => { + const pos = new Float32Array(6); + const geo = new Float32Array(16); + const color = new Float32Array(6).fill(0.25); + for (let i = 0; i < 2; i++) { + geo[i * 8] = 1; + geo[i * 8 + 4] = geo[i * 8 + 5] = geo[i * 8 + 6] = Math.log(0.1); + } + const cache = new Float32Array(2 * CACHE_STRIDE); + buildSplatCache({ pos, geo, color, colorDim: 3 }, cache); + + const boundary = await planBlockMerges({ + splatCache: cache, + neighbors: Uint32Array.from([1, NIL]), + D: 1, + coreCount: 1, + totalCount: 2 + }); + assert.strictEqual(boundary.costs.length, 0); + assert.strictEqual(boundary.frozen, 1); + + const interior = await planBlockMerges({ + splatCache: cache, + neighbors: Uint32Array.from([1, 0]), + D: 1, + coreCount: 2, + totalCount: 2 + }); + assert.strictEqual(interior.costs.length, 1); + assert.deepStrictEqual(Array.from(interior.pairs), [0, 1]); + }); + + it('dynamically freezes and unfreezes through reverse-candidate invalidation', async () => { + const input = randomBlock(1); + const plan = await planBlockMerges({ + splatCache: input.cache, + neighbors: input.neighbors, + D: input.D, + coreCount: input.coreCount, + totalCount: input.n + }); + assert.ok(plan.frozen > 0, 'halo minimum freezes at least one root'); + assert.ok(plan.unfrozen > 0, 'changed referenced core makes a frozen root eligible again'); + }); + + it('lets halo rows affect eligibility but never merges, removes, or duplicates them', async () => { + const input = randomBlock(1); + const plan = await planBlockMerges({ + splatCache: input.cache, + neighbors: input.neighbors, + D: input.D, + coreCount: input.coreCount, + totalCount: input.n + }); + assert.ok(Array.from(plan.pairs).every(row => row < input.coreCount), 'every committed endpoint is core'); + const replay = replayBlockPlan(input.coreCount, plan); + assert.strictEqual(replay.removed, plan.costs.length); + assert.strictEqual( + new Set(replay.groupMembers).size, + replay.groupMembers.length, + 'core group members are unique' + ); + assert.ok(!Array.from(replay.groupMembers).includes(input.coreCount), 'halo row is absent from output groups'); + }); +}); + +describe('block-plan prefix allocation', () => { + it('matches restricted global greedy for non-monotonic independent block sequences and groups', async () => { + const scratch = makeScratch(); + const local = [ + manualPlan([1, 100, 2], [0, 1, 0, 2, 0, 3]), + manualPlan([3, 4], [0, 1, 2, 3]) + ]; + const stored = []; + for (let i = 0; i < local.length; i++) stored.push(await storeBlockPlan(scratch, 1, i, local[i])); + + // Single restricted-global reference: only each independent block's + // next local commit is exposed, including its prefix dependency. + const cursor = [0, 0]; + const expected = []; + while (expected.length < 4) { + let block = -1; + for (let b = 0; b < local.length; b++) { + if (cursor[b] === local[b].costs.length) continue; + if (block < 0 || local[b].costs[cursor[b]] < local[block].costs[cursor[block]]) block = b; + } + expected.push([block, cursor[block]++]); + } + + const actual = []; + const result = await allocatePlanPrefixes(stored, scratch, 4, (block, index) => actual.push([block, index])); + assert.deepStrictEqual(actual, expected, 'merge sequence matches restricted global greedy'); + assert.deepStrictEqual(Array.from(result.prefixes), cursor); + + for (let b = 0; b < local.length; b++) { + const prefix = await readBlockPlanPrefix(stored[b], scratch, result.prefixes[b]); + const replay = replayBlockPlan(4, prefix); + assert.strictEqual(replay.removed, result.prefixes[b]); + assert.ok(replay.groupMembers.every(member => member < 4)); + } + }); + + it('selects every productive prefix on capacity shortfall while retaining the exact available count', async () => { + const scratch = makeScratch(); + const plans = [ + await storeBlockPlan(scratch, 1, 0, manualPlan([5, 1], [0, 1, 0, 2])), + await storeBlockPlan(scratch, 1, 1, manualPlan([2], [0, 1])) + ]; + const result = await allocatePlanPrefixes(plans, scratch, 99); + assert.strictEqual(result.removed, 3); + assert.deepStrictEqual(Array.from(result.prefixes), [2, 1]); + }); + + it('aborts a failed plan write without publishing a partial plan', async () => { + let aborted = false; + const scratch = { + writeFs: { + createWriter() { + return { + bytesWritten: 0, + write() { + throw new Error('write failed'); + }, + close() {}, + abort() { + aborted = true; + } + }; + }, + async mkdir() {} + }, + readFs: { async createSource() { throw new Error('not published'); } }, + scratchDir: 'scratch' + }; + await assert.rejects( + storeBlockPlan(scratch, 1, 0, manualPlan([1], [0, 1])), + /write failed/ + ); + assert.ok(aborted); + }); +}); + +describe('block-plan output replay', () => { + it('moment-matches selected prefixes one core at a time and cleans plan scratch', async () => { + const n = 24; + const { source, pool, pos } = await makeSyntheticSource(n, 1, 29, { + chunkSize: 5, + extraColumns: [{ name: 'tag', type: 'uint32' }] + }); + const partition = kdPartition(pos, 8, 1); + const scratch = makeScratch(); + const plans = []; + const prefixes = new Uint32Array(partition.blocks.length).fill(1); + for (let bi = 0; bi < partition.blocks.length; bi++) { + plans.push(await storeBlockPlan(scratch, 1, bi, manualPlan([bi + 1], [0, 1]))); + } + const outCount = n - partition.blocks.length; + const meta = { + ...source.meta, + numGaussians: outCount, + lodCounts: [outCount], + numChunks: [Math.ceil(outCount / source.meta.chunkSize)] + }; + const producer = createBlockProducerSource(meta, () => blockPlanMergeStream({ + source, + pool, + pos, + order: partition.order, + blocks: partition.blocks, + plans, + prefixes, + scratch + }, source.meta.chunkSize)); + + let rows = 0; + for (let c = 0; c < meta.numChunks[0]; c++) { + const count = Math.min(meta.chunkSize, outCount - rows); + const position = pool.acquire('position', meta.layouts.position, count); + const geometric = pool.acquire('geometric', meta.layouts.geometric, count); + const color = pool.acquire('color', meta.layouts.color, count); + const other = pool.acquire('other', meta.layouts.other, count); + await producer.read({ chunkIndex: c, position, geometric, color, other }); + for (const value of new Float32Array(position.data, 0, count * 3)) assert.ok(Number.isFinite(value)); + for (const value of new Float32Array(geometric.data, 0, count * 8)) assert.ok(Number.isFinite(value)); + position.release(); + geometric.release(); + color.release(); + other.release(); + rows += count; + } + assert.strictEqual(rows, outCount); + await producer.close(); + await source.close(); + for (const plan of plans) await scratch.remove(plan.path); + assert.strictEqual(scratch.writeFs.results.size, 0); + assert.strictEqual(scratch.removed.length, plans.length); + }); +}); diff --git a/test/decimate-multiblock.test.mjs b/test/decimate-multiblock.test.mjs new file mode 100644 index 00000000..1694c297 --- /dev/null +++ b/test/decimate-multiblock.test.mjs @@ -0,0 +1,80 @@ +import assert from 'node:assert'; +import { after, before, describe, it } from 'node:test'; + +import { makeSyntheticSource } from './helpers/synthetic-source.mjs'; + +import { decimateSource } from '../src/lib/decimate/index.js'; +import { MemoryReadSource } from '../src/lib/io/read/memory-file-system.js'; +import { MemoryFileSystem } from '../src/lib/io/write/memory-file-system.js'; + +let device = null; + +before(async () => { + try { + const { createDevice } = await import('../src/cli/node-device.js'); + device = await createDevice(); + } catch { + device = null; + } +}); + +after(() => { + device?.destroy?.(); +}); + +describe('decimateSource multi-block quality path', () => { + it('fails clearly when the memory budget requires multiple blocks without WebGPU', async () => { + const { source, pool } = await makeSyntheticSource(65540, 0, 123, { chunkSize: 1024 }); + await assert.rejects( + decimateSource(source, pool, { targetCount: 65000, memoryBudgetBytes: 1 }), + /multi-block quality decimation requires WebGPU/ + ); + }); + + it('hits the exact quota through scratch plans and removes every plan on close', { timeout: 120000 }, async (t) => { + if (!device) return t.skip('no WebGPU adapter available'); + + // Dynamic sizing bottoms out at 65,536 rows. Four extra rows force + // two jittered cores without making this acceptance fixture huge. + const n = 65540; + const targetCount = 65000; + const { source, pool } = await makeSyntheticSource(n, 0, 9876, { chunkSize: 1024 }); + const writeFs = new MemoryFileSystem(); + const spill = { + writeFs, + readFs: { + async createSource(path) { + const bytes = writeFs.results.get(path); + if (!bytes) throw new Error(`missing scratch file ${path}`); + return new MemoryReadSource(bytes); + } + }, + scratchDir: 'scratch', + async remove(path) { + writeFs.results.delete(path); + } + }; + + const out = await decimateSource(source, pool, { + targetCount, + createDevice: async () => device, + memoryBudgetBytes: 1, + spill + }); + assert.strictEqual(out.meta.numGaussians, targetCount); + let rows = 0; + for (let c = 0; c < out.meta.numChunks[0]; c++) { + const count = Math.min(out.meta.chunkSize, targetCount - rows); + const position = pool.acquire('position', out.meta.layouts.position, count); + await out.read({ chunkIndex: c, position }); + for (const value of new Float32Array(position.data, 0, count * 3)) { + assert.ok(Number.isFinite(value)); + } + position.release(); + rows += count; + } + assert.strictEqual(rows, targetCount); + await out.close(); + assert.strictEqual(writeFs.results.size, 0, 'all block plans cleaned'); + }); +}); diff --git a/test/decimate-partition.test.mjs b/test/decimate-partition.test.mjs index ae19dc74..1fbb8939 100644 --- a/test/decimate-partition.test.mjs +++ b/test/decimate-partition.test.mjs @@ -5,7 +5,7 @@ import assert from 'node:assert'; import { describe, it } from 'node:test'; -import { kdPartition, coherenceRuns } from '../src/lib/decimate/partition.js'; +import { buildBlockHalo, kdPartition, coherenceRuns } from '../src/lib/decimate/partition.js'; const grid = (n) => { const x = new Float32Array(n * n * n), y = new Float32Array(n * n * n), z = new Float32Array(n * n * n); @@ -56,6 +56,62 @@ describe('kdPartition', () => { assert.strictEqual(blocks.reduce((a, b) => a + (b.end - b.start), 0), n); }); + it('jitters split planes deterministically between generations while preserving the leaf cap', () => { + const pos = grid(16); + const a = kdPartition(pos, 300, 1); + const again = kdPartition(pos, 300, 1); + const b = kdPartition(pos, 300, 2); + assert.deepStrictEqual(Array.from(a.order), Array.from(again.order), 'same generation is deterministic'); + assert.notDeepStrictEqual( + a.blocks.map(block => block.end - block.start), + b.blocks.map(block => block.end - block.start), + 'generation changes split quantiles' + ); + assert.ok(a.blocks.every(block => block.end - block.start <= 300)); + assert.ok(b.blocks.every(block => block.end - block.start <= 300)); + }); + + it('builds sorted bounded halos without duplicating core ownership', () => { + const pos = grid(16); + const partition = kdPartition(pos, 300, 1); + let foundCapped = false; + for (let bi = 0; bi < partition.blocks.length; bi++) { + const block = partition.blocks[bi]; + const core = new Set(partition.order.subarray(block.start, block.end)); + const halo = buildBlockHalo(pos, partition, bi, 10); + assert.ok(halo.rows.length <= Math.min(10, core.size)); + for (let i = 0; i < halo.rows.length; i++) { + assert.ok(!core.has(halo.rows[i]), 'halo row is not core-owned'); + if (i > 0) assert.ok(halo.rows[i] > halo.rows[i - 1], 'halo rows sorted and unique'); + } + foundCapped ||= halo.capped; + } + assert.ok(foundCapped, 'dense boundary halos exercise the cap'); + }); + + it('moves a coincident boundary pair into one interior core on the next jitter pattern', () => { + const n = 256; + const pos = { x: new Float32Array(n), y: new Float32Array(n), z: new Float32Array(n) }; + const assignments = [1, 2].map((generation) => { + const partition = kdPartition(pos, 64, generation); + const owner = new Int32Array(n); + partition.blocks.forEach((block, bi) => { + for (let i = block.start; i < block.end; i++) owner[partition.order[i]] = bi; + }); + return owner; + }); + let pair = null; + for (let a = 0; a < n && !pair; a++) { + for (let b = a + 1; b < n; b++) { + if (assignments[0][a] !== assignments[0][b] && assignments[1][a] === assignments[1][b]) { + pair = [a, b]; + break; + } + } + } + assert.ok(pair, 'a coincident pair crosses the first boundary and becomes interior after jitter'); + }); + it('rare flyaways land in residual blocks; core blocks stay tight', () => { const g = grid(22); // 10648 bulk points in [0,21]^3 const nBulk = g.x.length; @@ -79,6 +135,7 @@ describe('kdPartition', () => { if (i > b.start) assert.ok(order[i] > order[i - 1], 'owned range sorted ascending'); } assert.ok(!(hasBulk && hasFly), 'flyaways segregated from the bulk'); + assert.strictEqual(b.residual, hasFly); if (hasBulk) { const ext = Math.max(b.aabb[3] - b.aabb[0], b.aabb[4] - b.aabb[1], b.aabb[5] - b.aabb[2]); assert.ok(ext <= 21 + 1e-6, `core block stretched by flyaways (extent ${ext})`); diff --git a/test/gpu-recost.test.mjs b/test/gpu-recost.test.mjs index 7634be4f..a8fc8eac 100644 --- a/test/gpu-recost.test.mjs +++ b/test/gpu-recost.test.mjs @@ -14,7 +14,13 @@ import assert from 'node:assert'; import { after, before, describe, it } from 'node:test'; import { buildSplatCache, CACHE_STRIDE } from '../src/lib/decimate/edge-cost-cpu.js'; -import { bestEdgeFor, bestOut } from '../src/lib/decimate/recost-core.js'; +import { + bestEdgeFor, + bestEdgesForPartition, + bestOut, + partitionBestOut +} from '../src/lib/decimate/recost-core.js'; +import { planBlockMerges } from '../src/lib/decimate/block-plan.js'; import { MAX_GROUP } from '../src/lib/decimate/select.js'; import { selectMergesRecosted } from '../src/lib/decimate/select-recost.js'; import { GpuRecost, COMMIT_LOG_STRIDE } from '../src/lib/gpu/gpu-recost.js'; @@ -126,6 +132,25 @@ const compareRefresh = (state, roots, out) => { }; describe('GpuRecost refresh parity', () => { + it('accepts a tiny residual core-plus-halo block in required-GPU mode', async (t) => { + if (!device) return t.skip('no WebGPU adapter available'); + + const n = 42; + const coreCount = 21; + const cache = makeCache(n, 2468); + const neighbors = bruteNeighbors(cache, n, K); + const plan = await planBlockMerges({ + splatCache: cache, + neighbors, + D: K, + coreCount, + totalCount: n, + device, + requireGpu: true + }); + assert.ok(Array.from(plan.pairs).every(row => row < coreCount)); + }); + it('wave-0 singleton and post-commit cluster costs match CPU within 1e-3', async (t) => { if (!device) return t.skip('no WebGPU adapter available'); @@ -192,6 +217,42 @@ describe('GpuRecost refresh parity', () => { gpu.destroy(); } }); + + it('returns independent best core and immutable-halo candidates', async (t) => { + if (!device) return t.skip('no WebGPU adapter available'); + + const n = 2048; + const coreCount = 2000; + const cache = makeCache(n, 4321); + const neighbors = bruteNeighbors(cache, n, K); + for (let i = 0; i < coreCount; i++) neighbors[i * K + K - 1] = coreCount + (i % (n - coreCount)); + neighbors.fill(NIL, coreCount * K); + const state = makeState(cache, neighbors, n); + const pending = Uint32Array.from({ length: 128 }, (_, i) => i * 7); + const out = new Uint32Array(pending.length * 4); + const costs = new Float32Array(out.buffer); + + const gpu = new GpuRecost(device, n, K, MAX_GROUP, 4096, coreCount); + try { + gpu.init(cache, neighbors); + await gpu.wave(new Uint32Array(0), 0, pending, pending.length, out); + for (let p = 0; p < pending.length; p++) { + const root = pending[p]; + assert.ok(bestEdgesForPartition(state.st, root, coreCount)); + assert.strictEqual(out[p * 4], partitionBestOut.corePartner); + assert.strictEqual(out[p * 4 + 2], partitionBestOut.haloPartner); + for (const [gpuCost, cpuCost] of [ + [costs[p * 4 + 1], partitionBestOut.coreCost], + [costs[p * 4 + 3], partitionBestOut.haloCost] + ]) { + const rel = Math.abs(gpuCost - cpuCost) / Math.max(1e-12, Math.abs(cpuCost)); + assert.ok(rel < 1e-3, `root ${root} partitioned cost parity ${rel}`); + } + } + } finally { + gpu.destroy(); + } + }); }); describe('GpuRecost selection equality', () => { From 337bc7147fa00d8b3d218515585035253a3a5549 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Wed, 29 Jul 2026 21:52:17 +0100 Subject: [PATCH 14/19] latest --- src/lib/decimate/block-plan.ts | 101 ++++++++++++++++++++++------ src/lib/decimate/decimate-source.ts | 17 +++++ test/decimate-block-plan.test.mjs | 5 ++ 3 files changed, 103 insertions(+), 20 deletions(-) diff --git a/src/lib/decimate/block-plan.ts b/src/lib/decimate/block-plan.ts index bfa20e81..6a8dcd3f 100644 --- a/src/lib/decimate/block-plan.ts +++ b/src/lib/decimate/block-plan.ts @@ -20,6 +20,15 @@ type BlockPlan = { costs: Float32Array; frozen: number; unfrozen: number; + diagnostics?: BlockPlanDiagnostics; +}; + +type BlockPlanDiagnostics = { + waves: number; + refreshes: number; + reverseInvalidations: number; + heapPops: number; + staleHeapPops: number; }; type BlockPlanInputs = { @@ -112,6 +121,7 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { let hB = new Uint32Array(heapCap); let hSeq = new Uint32Array(heapCap); let hVb = new Uint32Array(heapCap); + const heapIndex = new Int32Array(coreCount).fill(-1); let heapSize = 0; let seqCounter = 0; @@ -121,10 +131,42 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { t = hB[i]; hB[i] = hB[j]; hB[j] = t; t = hSeq[i]; hSeq[i] = hSeq[j]; hSeq[j] = t; t = hVb[i]; hVb[i] = hVb[j]; hVb[j] = t; + heapIndex[hA[i]] = i; + heapIndex[hA[j]] = j; }; const less = (i: number, j: number): boolean => hCost[i] < hCost[j] || (hCost[i] === hCost[j] && (hA[i] < hA[j] || (hA[i] === hA[j] && hB[i] < hB[j]))); + const heapRemoveAt = (at: number): void => { + heapIndex[hA[at]] = -1; + heapSize--; + if (at === heapSize) return; + hCost[at] = hCost[heapSize]; hA[at] = hA[heapSize]; hB[at] = hB[heapSize]; + hSeq[at] = hSeq[heapSize]; hVb[at] = hVb[heapSize]; + heapIndex[hA[at]] = at; + let i = at; + if (i > 0 && less(i, (i - 1) >> 1)) { + while (i > 0) { + const p = (i - 1) >> 1; + if (!less(i, p)) break; + swap(i, p); + i = p; + } + } else { + for (;;) { + const l = i * 2 + 1; + const r = l + 1; + let m = i; + if (l < heapSize && less(l, m)) m = l; + if (r < heapSize && less(r, m)) m = r; + if (m === i) break; + swap(i, m); + i = m; + } + } + }; const heapPush = (cost: number, a: number, b: number, seq: number, vb: number): void => { + const existing = heapIndex[a]; + if (existing >= 0) heapRemoveAt(existing); if (heapSize === heapCap) { const next = heapCap * 2; const nextCost = new Float32Array(next); nextCost.set(hCost); hCost = nextCost; @@ -136,6 +178,7 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { } let i = heapSize++; hCost[i] = cost; hA[i] = a; hB[i] = b; hSeq[i] = seq; hVb[i] = vb; + heapIndex[a] = i; while (i > 0) { const p = (i - 1) >> 1; if (!less(i, p)) break; @@ -147,22 +190,7 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { const heapPop = (): boolean => { if (heapSize === 0) return false; popped.cost = hCost[0]; popped.a = hA[0]; popped.b = hB[0]; popped.seq = hSeq[0]; popped.vb = hVb[0]; - heapSize--; - if (heapSize > 0) { - hCost[0] = hCost[heapSize]; hA[0] = hA[heapSize]; hB[0] = hB[heapSize]; - hSeq[0] = hSeq[heapSize]; hVb[0] = hVb[heapSize]; - let i = 0; - for (;;) { - const l = i * 2 + 1; - const r = l + 1; - let m = i; - if (l < heapSize && less(l, m)) m = l; - if (r < heapSize && less(r, m)) m = r; - if (m === i) break; - swap(i, m); - i = m; - } - } + heapRemoveAt(0); return true; }; @@ -173,8 +201,10 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { const queueRefresh = (root0: number): void => { const root = find(root0); if (root >= coreCount) return; - lastSeq[root] = ++seqCounter; if (queuedRound[root] === round) return; + lastSeq[root] = ++seqCounter; + const existing = heapIndex[root]; + if (existing >= 0) heapRemoveAt(existing); queuedRound[root] = round; pending[pendingCount++] = root; }; @@ -200,6 +230,11 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { const planCosts: number[] = []; let frozen = 0; let unfrozen = 0; + let waves = 0; + let refreshes = 0; + let reverseInvalidations = 0; + let heapPops = 0; + let staleHeapPops = 0; const refreshResult = (root: number, corePartner: number, coreCost: number, haloPartner: number, haloCost: number): void => { const isFrozen = haloPartner !== NIL && haloCost <= coreCost; @@ -220,17 +255,24 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { for (;;) { let wave = 0; while (wave < WAVE && heapPop()) { + heapPops++; const a = popped.a; - if (parent[a] !== a || popped.seq !== lastSeq[a]) continue; + if (parent[a] !== a || popped.seq !== lastSeq[a]) { + staleHeapPops++; + continue; + } const b = popped.b; if (b === a || b >= coreCount || parent[b] !== b || version[b] !== popped.vb || size[a] + size[b] > MAX_GROUP) { + staleHeapPops++; queueRefresh(a); continue; } const keep = size[a] >= size[b] ? a : b; const lose = keep === a ? b : a; + const loseHeap = heapIndex[lose]; + if (loseHeap >= 0) heapRemoveAt(loseHeap); if (commitLog) { const o = wave * COMMIT_LOG_STRIDE; commitLog[o] = lose; @@ -253,12 +295,17 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { // have a different core/halo minimum, including frozen roots. for (let m = mHead[keep]; m !== NIL; m = mNext[m]) { for (let r = reverseOffsets[m]; r < reverseOffsets[m + 1]; r++) { + reverseInvalidations++; queueRefresh(reverseRows[r]); } } } if (pendingCount === 0 && heapSize === 0) break; + if (pendingCount > 0) { + waves++; + refreshes += pendingCount; + } if (gpu && pendingCount > 0) { await gpu.wave(commitLog!, wave, pending, pendingCount, outBest!); for (let p = 0; p < pendingCount; p++) { @@ -297,7 +344,14 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { pairs: Uint32Array.from(planPairs), costs: Float32Array.from(planCosts), frozen, - unfrozen + unfrozen, + diagnostics: { + waves, + refreshes, + reverseInvalidations, + heapPops, + staleHeapPops + } }; }; @@ -365,4 +419,11 @@ const replayBlockPlan = (coreCount: number, plan: BlockPlan): SelectionResult => }; }; -export { planBlockMerges, replayBlockPlan, WAVE, type BlockPlan, type BlockPlanInputs }; +export { + planBlockMerges, + replayBlockPlan, + WAVE, + type BlockPlan, + type BlockPlanDiagnostics, + type BlockPlanInputs +}; diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index 00f80ba4..b8efb441 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -387,6 +387,12 @@ const decimateSource = async ( let cappedHalos = 0; let frozen = 0; let unfrozen = 0; + let planned = 0; + let waves = 0; + let refreshes = 0; + let reverseInvalidations = 0; + let heapPops = 0; + let staleHeapPops = 0; let knnMs = 0; let refreshMs = 0; let allocationMs = 0; @@ -427,6 +433,12 @@ const decimateSource = async ( refreshMs += Date.now() - refreshStarted; frozen += plan.frozen; unfrozen += plan.unfrozen; + planned += plan.costs.length; + waves += plan.diagnostics!.waves; + refreshes += plan.diagnostics!.refreshes; + reverseInvalidations += plan.diagnostics!.reverseInvalidations; + heapPops += plan.diagnostics!.heapPops; + staleHeapPops += plan.diagnostics!.staleHeapPops; storedPlans[bi] = await storeBlockPlan(opts.spill, generation, bi, plan); planBar.tick(prepared.ownedCount); } @@ -456,6 +468,11 @@ const decimateSource = async ( `local merge stats: ${fmtCount(cappedHalos)} capped halo${cappedHalos === 1 ? '' : 's'}, ` + `${fmtCount(frozen)} freezes, ${fmtCount(unfrozen)} unfreezes, ${fmtCount(removed!)} removals` ); + logger.info( + `local planner work: ${fmtCount(planned)} planned, ${fmtCount(waves)} waves, ` + + `${fmtCount(refreshes)} refreshes, ${fmtCount(reverseInvalidations)} reverse invalidations, ` + + `${fmtCount(heapPops)} heap pops (${fmtCount(staleHeapPops)} stale)` + ); logger.info( `local timings: KNN/gather ${(knnMs / 1000).toFixed(2)}s, ` + `refresh/plan ${(refreshMs / 1000).toFixed(2)}s, allocation ${(allocationMs / 1000).toFixed(2)}s` diff --git a/test/decimate-block-plan.test.mjs b/test/decimate-block-plan.test.mjs index a7fffdef..6ea2694f 100644 --- a/test/decimate-block-plan.test.mjs +++ b/test/decimate-block-plan.test.mjs @@ -114,6 +114,9 @@ describe('block-local merge planning', () => { }); assert.strictEqual(interior.costs.length, 1); assert.deepStrictEqual(Array.from(interior.pairs), [0, 1]); + assert.ok(interior.diagnostics.waves > 0); + assert.ok(interior.diagnostics.refreshes >= interior.costs.length); + assert.ok(interior.diagnostics.heapPops >= interior.costs.length); }); it('dynamically freezes and unfreezes through reverse-candidate invalidation', async () => { @@ -127,6 +130,8 @@ describe('block-local merge planning', () => { }); assert.ok(plan.frozen > 0, 'halo minimum freezes at least one root'); assert.ok(plan.unfrozen > 0, 'changed referenced core makes a frozen root eligible again'); + assert.ok(plan.diagnostics.reverseInvalidations > 0); + assert.ok(plan.diagnostics.staleHeapPops <= plan.diagnostics.heapPops); }); it('lets halo rows affect eligibility but never merges, removes, or duplicates them', async () => { From 1c54704d78b5d1fdd1cda7236f637198f1df9880 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Thu, 30 Jul 2026 09:55:24 +0100 Subject: [PATCH 15/19] latest --- README.md | 4 ++-- src/cli/index.ts | 22 ++++++++-------------- src/lib/decimate/block-plan.ts | 11 ++++++----- src/lib/decimate/decimate-source.ts | 2 +- src/lib/gpu/gpu-recost.ts | 12 ++++++++---- test/cli.test.mjs | 21 +++++++++++++++++++++ test/gpu-recost.test.mjs | 25 ++++++++++++++++++++++++- 7 files changed, 70 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index c9735a6c..5c987704 100644 --- a/README.md +++ b/README.md @@ -110,8 +110,8 @@ Actions execute in the order specified and can be repeated. Any action may appea opacity, scale_*, f_dc_* use transformed values (linear opacity 0-1, linear scale, linear color 0-1). Append _raw for raw PLY values (e.g. opacity_raw). --d, --decimate Simplify to n Gaussians via merge-based decimation - Use n% to keep a percentage of Gaussians. +-d, --decimate Simplify with the quality decimator + Use n% for a percentage; --decimate-balanced selects the pre-3.2 algorithm. Memory-bounded and streaming: scales to scenes of 100M+ Gaussians. Must be the final action, and the output must be .ply (write a decimated PLY first, then convert in a diff --git a/src/cli/index.ts b/src/cli/index.ts index 5c64c26c..ec4fc661 100644 --- a/src/cli/index.ts +++ b/src/cli/index.ts @@ -191,7 +191,7 @@ const cliOptionsConfig = { 'filter-box': { type: 'string', short: 'B', multiple: true }, 'filter-sphere': { type: 'string', short: 'S', multiple: true }, 'decimate': { type: 'string', short: 'd', multiple: true }, - 'decimate-mode': { type: 'string', default: 'quality' }, + 'decimate-balanced': { type: 'string', multiple: true }, 'memory-budget': { type: 'string' }, 'filter-cluster': { type: 'string', short: 'C', multiple: true }, 'filter-floaters': { type: 'string', short: 'F', multiple: true }, @@ -522,13 +522,7 @@ const parseArguments = async () => { listGpus: v['list-gpus'], deviceIdx, scratchDir: v['scratch-dir'], - decimateMode: (() => { - const m = v['decimate-mode']; - if (m !== 'quality' && m !== 'legacy') { - throw new Error(`Invalid --decimate-mode: ${m}. Must be 'quality' or 'legacy'.`); - } - return m; - })(), + decimateMode: 'quality', // Residency policy ceiling for decimation (not an upfront allocation): // default to half the machine's RAM, capped at 48 GiB. memoryBudgetBytes: v['memory-budget'] !== undefined ? @@ -704,7 +698,9 @@ const parseArguments = async () => { kind: 'mortonOrder' }); break; - case 'decimate': { + case 'decimate': + case 'decimate-balanced': { + if (t.name === 'decimate-balanced') options.decimateMode = 'legacy'; const value = t.value.trim(); let count: number | null = null; let percent: number | null = null; @@ -804,11 +800,9 @@ ACTIONS (executed in order; can be repeated) -S, --filter-sphere Remove Gaussians outside sphere -V, --filter-value Keep Gaussians where ; cmp ∈ {lt,lte,gt,gte,eq,neq} - -d, --decimate Simplify to n (or n%) Gaussians via merge-based decimation. - --decimate-mode quality (default): field-L2 cost + re-costed selection — best on - mixed-scale scenes (large gains on skies/distant structure). - legacy: pre-3.2 pipeline — faster, lower memory, and still slightly - better on scenes of uniformly-sized Gaussians (single objects). + -d, --decimate Simplify with field-L2 cost and re-costed selection (quality default). + --decimate-balanced Simplify with the pre-3.2 balanced algorithm (lower memory and often + better on uniformly-sized Gaussians). --memory-budget Decimation residency policy ceiling (not an upfront allocation); re-costed selection falls back to one-shot selection above it. Default: min(48, half of system RAM). diff --git a/src/lib/decimate/block-plan.ts b/src/lib/decimate/block-plan.ts index 6a8dcd3f..21b7d32d 100644 --- a/src/lib/decimate/block-plan.ts +++ b/src/lib/decimate/block-plan.ts @@ -223,7 +223,7 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { } const gpu = gpuFits ? new GpuRecost(device!, N, D, MAX_GROUP, WAVE, coreCount) : undefined; const commitLog = gpu ? new Uint32Array(WAVE * COMMIT_LOG_STRIDE) : undefined; - const outBest = gpu ? new Uint32Array(coreCount * 4) : undefined; + const outBest = gpu ? new Uint32Array(coreCount * gpu.outputStride) : undefined; const outCost = gpu ? new Float32Array(outBest!.buffer) : undefined; const planPairs: number[] = []; @@ -311,12 +311,13 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { for (let p = 0; p < pendingCount; p++) { const root = pending[p]; if (parent[root] !== root) continue; + const o = p * gpu.outputStride; refreshResult( root, - outBest![p * 4], - outCost![p * 4 + 1], - outBest![p * 4 + 2], - outCost![p * 4 + 3] + outBest![o], + outCost![o + 1], + gpu.outputStride === 4 ? outBest![o + 2] : NIL, + gpu.outputStride === 4 ? outCost![o + 3] : Infinity ); } } else { diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index b8efb441..1dcf5855 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -131,7 +131,7 @@ const chooseBlockSize = ( return Math.max(1, Math.min(blockSize, n)); }; -// The pre-study KL-style cost kernel (--decimate-mode legacy): full-SH colour +// The pre-study KL-style cost kernel (--decimate-balanced): full-SH colour // L2, single-Monte-Carlo geometric term, its own GPU cache/kernel layouts. const createLegacyStrategy = (colorDim: number): CostStrategy => { const Z = makeGaussianSamples(1, 0); diff --git a/src/lib/gpu/gpu-recost.ts b/src/lib/gpu/gpu-recost.ts index 0c9154f8..51ec4b79 100644 --- a/src/lib/gpu/gpu-recost.ts +++ b/src/lib/gpu/gpu-recost.ts @@ -458,6 +458,9 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { const COMMIT_LOG_STRIDE = 5; class GpuRecost { + /** Number of u32 values written per refreshed root. */ + readonly outputStride: number; + /** * Upload the immutable inputs and initialize the structure buffers. * Call once before the first wave. @@ -466,7 +469,7 @@ class GpuRecost { /** * Run one wave: replay `commitCount` log entries, refresh * `pendingCount` queued roots, and read back (partner, cost) pairs - * into `outBest` (2 u32 per root; cost is a bitcast f32). + * into `outBest` (`outputStride` u32 per root; costs are bitcast f32). */ wave: ( commitLog: Uint32Array, @@ -526,7 +529,8 @@ class GpuRecost { } const splitN = Math.ceil(n / 2); const partitioned = coreCount < n; - const outputStride = partitioned ? 16 : 8; + this.outputStride = partitioned ? 4 : 2; + const outputStrideBytes = this.outputStride * 4; const cacheABuf = new StorageBuffer(device, splitN * 16 * 4, BUFFERUSAGE_COPY_DST); const cacheBBuf = new StorageBuffer(device, Math.max(n - splitN, 1) * 16 * 4, BUFFERUSAGE_COPY_DST); @@ -535,7 +539,7 @@ class GpuRecost { const parentMetaBuf = new StorageBuffer(device, n * 8, BUFFERUSAGE_COPY_DST); const chainBuf = new StorageBuffer(device, n * 8, BUFFERUSAGE_COPY_DST); const pendingBuf = new StorageBuffer(device, n * 4, BUFFERUSAGE_COPY_DST); - const outBestBuf = new StorageBuffer(device, n * outputStride, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST); + const outBestBuf = new StorageBuffer(device, n * outputStrideBytes, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST); const commitLogBuf = new StorageBuffer(device, wave * COMMIT_LOG_STRIDE * 4, BUFFERUSAGE_COPY_DST); const initKernel = makeKernel(device, 'recost-init', initWgsl(), ['count'], [ @@ -624,7 +628,7 @@ class GpuRecost { device.computeDispatch(computes, 'recost-wave'); // Blocking readback — also the wave's submit boundary. - await outBestBuf.read(0, pendingCount * outputStride, outBest, true); + await outBestBuf.read(0, pendingCount * outputStrideBytes, outBest, true); }; this.destroy = () => { diff --git a/test/cli.test.mjs b/test/cli.test.mjs index 154a7350..7b9e45fa 100644 --- a/test/cli.test.mjs +++ b/test/cli.test.mjs @@ -184,6 +184,27 @@ describe('CLI decimate (terminal PLY restriction)', () => { const inputCount = parsed.count ?? parsed.numGaussians ?? parsed.lods?.[0]?.count; assert.strictEqual(written, Math.round(inputCount / 2), `50% of ${inputCount}`); }); + + it('--decimate-balanced runs the pre-3.2 algorithm', async () => { + const { mkdtemp, readFile: readFileFs, rm } = await import('node:fs/promises'); + const { tmpdir } = await import('node:os'); + const { join } = await import('node:path'); + const dir = await mkdtemp(join(tmpdir(), 'st-decimate-balanced-cli-')); + const balancedPath = join(dir, 'balanced.ply'); + + const balanced = await runCli([ + '--gpu', 'cpu', + 'test/fixtures/splat/minimal.splat', + '--decimate-balanced', '50%', + balancedPath + ]); + assert.strictEqual(balanced.code, 0, `balanced CLI failed:\n${balanced.stderr}\n${balanced.stdout}`); + const header = (await readFileFs(balancedPath)).subarray(0, 1024).toString('ascii'); + const match = header.match(/element vertex (\d+)/); + assert.ok(match, 'output has a vertex element'); + assert.strictEqual(parseInt(match[1], 10), 2); + await rm(dir, { recursive: true, force: true }); + }); }); describe('CLI filter-nan (zero-norm rotation)', () => { diff --git a/test/gpu-recost.test.mjs b/test/gpu-recost.test.mjs index a8fc8eac..b6ea6f3d 100644 --- a/test/gpu-recost.test.mjs +++ b/test/gpu-recost.test.mjs @@ -20,7 +20,7 @@ import { bestOut, partitionBestOut } from '../src/lib/decimate/recost-core.js'; -import { planBlockMerges } from '../src/lib/decimate/block-plan.js'; +import { planBlockMerges, replayBlockPlan } from '../src/lib/decimate/block-plan.js'; import { MAX_GROUP } from '../src/lib/decimate/select.js'; import { selectMergesRecosted } from '../src/lib/decimate/select-recost.js'; import { GpuRecost, COMMIT_LOG_STRIDE } from '../src/lib/gpu/gpu-recost.js'; @@ -151,6 +151,29 @@ describe('GpuRecost refresh parity', () => { assert.ok(Array.from(plan.pairs).every(row => row < coreCount)); }); + it('plans a core-only GPU block with the compact output stride', async (t) => { + if (!device) return t.skip('no WebGPU adapter available'); + + const n = 512; + const cache = makeCache(n, 777); + const neighbors = bruteNeighbors(cache, n, K); + const inputs = { + splatCache: cache, + neighbors, + D: K, + coreCount: n, + totalCount: n + }; + const inline = await planBlockMerges(inputs); + const gpu = await planBlockMerges({ ...inputs, device, requireGpu: true }); + + assert.strictEqual(gpu.costs.length, inline.costs.length); + assert.deepStrictEqual( + Array.from(replayBlockPlan(n, gpu).memberGroup), + Array.from(replayBlockPlan(n, inline).memberGroup) + ); + }); + it('wave-0 singleton and post-commit cluster costs match CPU within 1e-3', async (t) => { if (!device) return t.skip('no WebGPU adapter available'); From f1b468b8c7db8dac03fd0eaf5f3bd3a14b106713 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Thu, 30 Jul 2026 10:32:49 +0100 Subject: [PATCH 16/19] latest --- src/cli/index.ts | 15 +++++---------- src/lib/decimate/decimate-source.ts | 2 +- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/src/cli/index.ts b/src/cli/index.ts index ec4fc661..aae14ceb 100644 --- a/src/cli/index.ts +++ b/src/cli/index.ts @@ -68,7 +68,7 @@ interface CliOptions extends LibOptions { deviceIdx: number; // -1 = auto, -2 = CPU, 0+ = GPU index scratchDir: string | undefined; // decimation spill location (default: output directory) decimateMode: 'quality' | 'legacy'; - memoryBudgetBytes: number; // decimation residency policy ceiling (not an allocation) + memoryBudgetBytes: number; // decimation residency policy ceiling (not an allocation, not user-facing) } const fileExists = async (filename: string) => { @@ -192,7 +192,6 @@ const cliOptionsConfig = { 'filter-sphere': { type: 'string', short: 'S', multiple: true }, 'decimate': { type: 'string', short: 'd', multiple: true }, 'decimate-balanced': { type: 'string', multiple: true }, - 'memory-budget': { type: 'string' }, 'filter-cluster': { type: 'string', short: 'C', multiple: true }, 'filter-floaters': { type: 'string', short: 'F', multiple: true }, params: { type: 'string', short: 'p', multiple: true }, @@ -523,11 +522,10 @@ const parseArguments = async () => { deviceIdx, scratchDir: v['scratch-dir'], decimateMode: 'quality', - // Residency policy ceiling for decimation (not an upfront allocation): - // default to half the machine's RAM, capped at 48 GiB. - memoryBudgetBytes: v['memory-budget'] !== undefined ? - Math.max(1, parseNumber(v['memory-budget'])) * 2 ** 30 : - Math.min(48 * 2 ** 30, Math.floor(totalmem() / 2)), + // Residency policy ceiling for decimation (not an upfront allocation). + // Half the machine's RAM, capped at 48 GiB — derived here because the + // library is node-free and cannot read os.totalmem() itself. + memoryBudgetBytes: Math.min(48 * 2 ** 30, Math.floor(totalmem() / 2)), lodSelect: v['select-lod'].split(',').filter(v => !!v).map(parseInteger), viewerSettingsJson: viewerSettingsPath && await readJsonFile(viewerSettingsPath), unbundled: v.unbundled, @@ -803,9 +801,6 @@ ACTIONS (executed in order; can be repeated) -d, --decimate Simplify with field-L2 cost and re-costed selection (quality default). --decimate-balanced Simplify with the pre-3.2 balanced algorithm (lower memory and often better on uniformly-sized Gaussians). - --memory-budget Decimation residency policy ceiling (not an upfront allocation); - re-costed selection falls back to one-shot selection above it. - Default: min(48, half of system RAM). Must be the final action, and the output must be .ply --scratch-dir Directory for decimation spill files (deep targets on huge scenes). Default: the output file's directory diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index 1dcf5855..316212eb 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -365,7 +365,7 @@ const decimateSource = async ( if (!device) { throw new Error( `multi-block quality decimation requires WebGPU (${fmtCount(N)} splats, ` + - `${fmtCount(blockSize)}-splat cores); increase --memory-budget for the one-block path or provide a device` + `${fmtCount(blockSize)}-splat cores); provide a device, or use --decimate-balanced` ); } if (!opts.spill) { From c2dbaf39e66c8cde422d679a91a4db402cd34d01 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Thu, 30 Jul 2026 14:27:23 +0100 Subject: [PATCH 17/19] latest --- README.md | 6 +- src/cli/index.ts | 50 ++- src/lib/decimate-uniform/README.md | 81 ++++ src/lib/decimate-uniform/block-producer.ts | 94 +++++ src/lib/decimate-uniform/decimate-source.ts | 340 +++++++++++++++ .../edge-cost-cpu.ts} | 82 +--- .../gpu-edge-cost.ts} | 159 +++---- src/lib/decimate-uniform/gpu-knn.ts | 345 ++++++++++++++++ src/lib/decimate-uniform/index.ts | 8 + src/lib/decimate-uniform/knn-blocks.ts | 332 +++++++++++++++ src/lib/decimate-uniform/knn-core.ts | 51 +++ src/lib/decimate-uniform/merge-stream.ts | 178 ++++++++ src/lib/decimate-uniform/partition.ts | 163 ++++++++ src/lib/decimate-uniform/priority.ts | 387 ++++++++++++++++++ .../select.ts} | 4 +- src/lib/decimate/block-plan.ts | 2 +- src/lib/decimate/decimate-source.ts | 104 ++--- src/lib/decimate/priority.ts | 82 +--- src/lib/index.ts | 5 + src/lib/workers/tasks.ts | 33 +- test/cli.test.mjs | 6 +- test/decimate-multiblock.test.mjs | 4 +- test/decimate-uniform-parity.test.mjs | 143 +++++++ 23 files changed, 2333 insertions(+), 326 deletions(-) create mode 100644 src/lib/decimate-uniform/README.md create mode 100644 src/lib/decimate-uniform/block-producer.ts create mode 100644 src/lib/decimate-uniform/decimate-source.ts rename src/lib/{decimate/edge-cost-legacy.ts => decimate-uniform/edge-cost-cpu.ts} (67%) rename src/lib/{gpu/gpu-edge-cost-legacy.ts => decimate-uniform/gpu-edge-cost.ts} (78%) create mode 100644 src/lib/decimate-uniform/gpu-knn.ts create mode 100644 src/lib/decimate-uniform/index.ts create mode 100644 src/lib/decimate-uniform/knn-blocks.ts create mode 100644 src/lib/decimate-uniform/knn-core.ts create mode 100644 src/lib/decimate-uniform/merge-stream.ts create mode 100644 src/lib/decimate-uniform/partition.ts create mode 100644 src/lib/decimate-uniform/priority.ts rename src/lib/{decimate/select-legacy.ts => decimate-uniform/select.ts} (96%) create mode 100644 test/decimate-uniform-parity.test.mjs diff --git a/README.md b/README.md index 5c987704..16b0a051 100644 --- a/README.md +++ b/README.md @@ -110,8 +110,10 @@ Actions execute in the order specified and can be repeated. Any action may appea opacity, scale_*, f_dc_* use transformed values (linear opacity 0-1, linear scale, linear color 0-1). Append _raw for raw PLY values (e.g. opacity_raw). --d, --decimate Simplify with the quality decimator - Use n% for a percentage; --decimate-balanced selects the pre-3.2 algorithm. +-d, --decimate Simplify, allocating removal by local error (adaptive) + Use n% for a percentage. --decimate-uniform removes at a uniform + rate instead (the pre-3.2 algorithm): lower memory, and better on + uniformly-sized Gaussians. Memory-bounded and streaming: scales to scenes of 100M+ Gaussians. Must be the final action, and the output must be .ply (write a decimated PLY first, then convert in a diff --git a/src/cli/index.ts b/src/cli/index.ts index aae14ceb..af6d21f1 100644 --- a/src/cli/index.ts +++ b/src/cli/index.ts @@ -16,6 +16,7 @@ import { DataTable, dataTableToChunkSource, decimateSource, + decimateSourceUniform, fmtBytes, fmtCount, fmtTime, @@ -67,7 +68,7 @@ interface CliOptions extends LibOptions { listGpus: boolean; deviceIdx: number; // -1 = auto, -2 = CPU, 0+ = GPU index scratchDir: string | undefined; // decimation spill location (default: output directory) - decimateMode: 'quality' | 'legacy'; + decimateUniform: boolean; // --decimate-uniform: the frozen pre-3.2 decimator memoryBudgetBytes: number; // decimation residency policy ceiling (not an allocation, not user-facing) } @@ -191,7 +192,7 @@ const cliOptionsConfig = { 'filter-box': { type: 'string', short: 'B', multiple: true }, 'filter-sphere': { type: 'string', short: 'S', multiple: true }, 'decimate': { type: 'string', short: 'd', multiple: true }, - 'decimate-balanced': { type: 'string', multiple: true }, + 'decimate-uniform': { type: 'string', multiple: true }, 'filter-cluster': { type: 'string', short: 'C', multiple: true }, 'filter-floaters': { type: 'string', short: 'F', multiple: true }, params: { type: 'string', short: 'p', multiple: true }, @@ -521,7 +522,7 @@ const parseArguments = async () => { listGpus: v['list-gpus'], deviceIdx, scratchDir: v['scratch-dir'], - decimateMode: 'quality', + decimateUniform: false, // Residency policy ceiling for decimation (not an upfront allocation). // Half the machine's RAM, capped at 48 GiB — derived here because the // library is node-free and cannot read os.totalmem() itself. @@ -697,8 +698,8 @@ const parseArguments = async () => { }); break; case 'decimate': - case 'decimate-balanced': { - if (t.name === 'decimate-balanced') options.decimateMode = 'legacy'; + case 'decimate-uniform': { + if (t.name === 'decimate-uniform') options.decimateUniform = true; const value = t.value.trim(); let count: number | null = null; let percent: number | null = null; @@ -798,9 +799,11 @@ ACTIONS (executed in order; can be repeated) -S, --filter-sphere Remove Gaussians outside sphere -V, --filter-value Keep Gaussians where ; cmp ∈ {lt,lte,gt,gte,eq,neq} - -d, --decimate Simplify with field-L2 cost and re-costed selection (quality default). - --decimate-balanced Simplify with the pre-3.2 balanced algorithm (lower memory and often - better on uniformly-sized Gaussians). + -d, --decimate Simplify, allocating removal by local error (adaptive; default). + Much better on mixed-scale content such as skies. + --decimate-uniform Simplify at a uniform rate everywhere (the pre-3.2 algorithm). + Lower memory, and better at depth on uniformly-sized + Gaussians: uniform texture, single objects, snow. Must be the final action, and the output must be .ply --scratch-dir Directory for decimation spill files (deep targets on huge scenes). Default: the output file's directory @@ -1256,18 +1259,25 @@ const main = async () => { if (keepCount < 1) { failExit(`--decimate target resolves to ${keepCount} gaussians; must keep at least 1`); } - combined = await decimateSource(combined, pool, { - targetCount: keepCount, - createDevice: deviceCreator, - mode: options.decimateMode, - memoryBudgetBytes: options.memoryBudgetBytes, - spill: { - writeFs: new NodeFileSystem(), - readFs: new NodeReadFileSystem(), - scratchDir: options.scratchDir ?? dirname(outputFilename), - remove: path => unlink(path) - } - }); + const spill = { + writeFs: new NodeFileSystem(), + readFs: new NodeReadFileSystem(), + scratchDir: options.scratchDir ?? dirname(outputFilename), + remove: (path: string) => unlink(path) + }; + combined = options.decimateUniform ? + await decimateSourceUniform(combined, pool, { + targetCount: keepCount, + createDevice: deviceCreator, + memoryBudgetBytes: options.memoryBudgetBytes, + spill + }) : + await decimateSource(combined, pool, { + targetCount: keepCount, + createDevice: deviceCreator, + memoryBudgetBytes: options.memoryBudgetBytes, + spill + }); } logger.info(`${fmtCount(combined.meta.numGaussians)} gaussians · ${combined.meta.shBands} SH bands`); diff --git a/src/lib/decimate-uniform/README.md b/src/lib/decimate-uniform/README.md new file mode 100644 index 00000000..e76c84fe --- /dev/null +++ b/src/lib/decimate-uniform/README.md @@ -0,0 +1,81 @@ +# decimate-uniform — the pre-3.2 decimator + +The decimator that shipped up to 3.1.x, reached through `--decimate-uniform` / +`decimateSourceUniform()`. `src/lib/decimate/` holds the adaptive one +(`--decimate` / `decimateSource()`). + +The names describe how each allocates removal, not a ranking. Both are +supported and neither is a fallback for the other: they win on different +content, and the choice is the user's. + +- **uniform** — KL-style pairwise cost with a full-SH colour term, uniform 50% + matching per level, so every region loses the same fraction. Lower memory, + and measurably better at depth on scenes of uniformly-sized Gaussians: + uniform texture, single objects, snow. See the `old` column in + `scenes/DECIMATION-RESULTS.md` (leads at L3–L6 on `crop-snow` and `fr-snow`). +- **adaptive** — field-L2 cost with the scale-free colour term and re-costed + selection, so removal follows local error and redundant regions collapse + deeper than distinct ones. Large wins on mixed-scale content, skies + especially (+9 to +11 dB on `fr-sky`), at higher memory cost. + +## The contract + +This directory is **bit-for-bit output-compatible with the 3.1.6 binary**. That +is its value: a decimation you can reproduce exactly against a known-good +reference, and the baseline every quality comparison in +`scenes/DECIMATION-RESULTS.md` is measured against. + +Every file is a copy of its `src/lib/decimate/` counterpart at the last 3.1.x +commit. You can prove it, per file: + +```bash +git diff main:src/lib/decimate/select.ts src/lib/decimate-uniform/select.ts +``` + +Empty output means the file is untouched. Deviations are limited to these, all +mechanical: + +- **Import paths.** `../gpu/gpu-edge-cost` → `./gpu-edge-cost`, + `../gpu/gpu-knn` → `./gpu-knn`, `./moment-match` → + `../decimate/moment-match`. +- **`gpu-knn.ts`** consumes the current `FlatKdTree` (interleaved + `nodePositions` / `nodeChildren`) and so drops the packing loops that built + that same layout internally. `buildFlatKdTree` is verified structurally + identical to the 3.1.x `KdTree.flatten()` at every size, so the uploaded + bytes are unchanged. + +## Shared dependencies + +Only two, both deliberate: + +- `../decimate/moment-match.ts` — has no diff against 3.1.x, and its + `mergeGroups` worker handler is shared. Duplicating it would mean a + duplicate worker task for no benefit. +- `../spatial/kd-tree.ts` — `KdTree`'s build and query paths are unchanged + from 3.1.x, and it is shared with k-means. + +Otherwise this directory imports nothing from `../decimate/`, so work on the +adaptive path cannot change uniform output. + +## Changing things here + +Changes are fine — bug fixes, performance work, new capability — but they are +output changes to a path whose selling point is reproducibility, so they need +to be deliberate rather than incidental. Before landing one: + +- Re-run the whole-scene comparison against the 3.1.6 binary if you expect + output to be unchanged. Equivalence was last verified on both study scenes, + `fr-sky` (5.81M, 3 SH bands, multi-block) and `fr-snow` (26.1M, DC only, 13 + blocks), six chained halvings each: every level byte-identical, PSNR matching + the published `old` columns exactly. +- If output *should* change, re-baseline `scenes/DECIMATION-RESULTS.md` — the + `old` column is this path, and the study's conclusions are stated relative + to it. +- Repin the digest in `test/decimate-legacy-frozen.test.mjs`, which is the + in-suite tripwire for accidental drift. + +## If it is ever retired + +Nothing in `src/lib/decimate/` refers to this directory, so: `rm -rf` it, drop +`decimateSourceUniform` from `src/lib/index.ts`, drop `--decimate-uniform` +from the CLI, and delete `test/decimate-legacy-frozen.test.mjs`. diff --git a/src/lib/decimate-uniform/block-producer.ts b/src/lib/decimate-uniform/block-producer.ts new file mode 100644 index 00000000..b45ba9df --- /dev/null +++ b/src/lib/decimate-uniform/block-producer.ts @@ -0,0 +1,94 @@ +import { type ChunkData, type ChunkLayer, type ChunkSource, type ChunkSourceMetadata, type ReadRequest } from '../chunk'; + +/** + * One output chunk of the merge stream. The views hold exactly `count` + * records at the layer strides and alias the generator's rolling scratch — + * valid only until the generator's next `next()`; the consumer's read copies + * them out (yielding views instead of slices avoids a full extra copy of + * every output byte). + */ +type ChunkPayload = { + count: number; + position: Float32Array; + geometric: Float32Array; + color: Float32Array; + other?: Uint32Array; +}; + +/** + * A single-sequential-pass {@link ChunkSource} over an async generator of + * chunk payloads — how the decimation merge stream feeds the PLY writer + * (or `compact` / `writePlyStreaming` for intermediate generations) without + * ever materializing the output. + * + * Contract: chunk reads must arrive in order (0, 1, 2, …), each at most + * once; gather reads are not supported. Anything else throws — decimate + * output supports exactly one sequential pass. + * + * @param meta - Exact output metadata (counts are known before streaming). + * @param produce - Factory for the payload generator (invoked lazily on first read). + * @returns The stream-once source. + */ +const createBlockProducerSource = ( + meta: ChunkSourceMetadata, + produce: () => AsyncGenerator +): ChunkSource => { + let generator: AsyncGenerator | null = null; + let nextChunk = 0; + let done = false; + + const read = async (request: ReadRequest): Promise => { + if ('indices' in request) { + throw new Error('decimate output supports a single sequential pass (gather reads are not available)'); + } + if ((request.lod ?? 0) !== 0) { + throw new Error(`decimate output has a single LOD (requested lod ${request.lod})`); + } + if (request.chunkIndex !== nextChunk) { + throw new Error( + `decimate output supports a single sequential pass (expected chunk ${nextChunk}, got ${request.chunkIndex})` + ); + } + if (done) { + throw new Error('decimate output exhausted'); + } + generator ??= produce(); + const { value, done: exhausted } = await generator.next(); + if (exhausted || !value) { + done = true; + throw new Error(`decimate output ended early at chunk ${request.chunkIndex}`); + } + const payload = value; + const expected = Math.min(meta.chunkSize, meta.numGaussians - request.chunkIndex * meta.chunkSize); + if (payload.count !== expected) { + throw new Error(`decimate output chunk ${request.chunkIndex}: expected ${expected} rows, produced ${payload.count}`); + } + + const fill = (cd: ChunkData | undefined, layer: ChunkLayer): void => { + if (!cd) return; + const src = payload[layer as 'position' | 'geometric' | 'color' | 'other']; + if (!src) { + throw new Error(`decimate output has no '${layer}' layer`); + } + const bytes = payload.count * cd.stride; + new Uint8Array(cd.data, 0, bytes).set(new Uint8Array(src.buffer, src.byteOffset, bytes)); + }; + fill(request.position, 'position'); + fill(request.geometric, 'geometric'); + fill(request.color, 'color'); + fill(request.other, 'other'); + + nextChunk++; + if (nextChunk >= (meta.numChunks[0] ?? 0)) done = true; + }; + + const close = async (): Promise => { + done = true; + await generator?.return?.(undefined as never); + generator = null; + }; + + return { meta, read, close }; +}; + +export { createBlockProducerSource, type ChunkPayload }; diff --git a/src/lib/decimate-uniform/decimate-source.ts b/src/lib/decimate-uniform/decimate-source.ts new file mode 100644 index 00000000..2ec9ac67 --- /dev/null +++ b/src/lib/decimate-uniform/decimate-source.ts @@ -0,0 +1,340 @@ +import { join } from 'pathe'; +import { type GraphicsDevice } from 'playcanvas'; + +import { createBlockProducerSource } from './block-producer'; +import { mergeStream } from './merge-stream'; +import { kdPartition, coherenceRuns, type ResidentPositions } from './partition'; +import { runPriorityPass, HALO_CAP, type CandidateArrays } from './priority'; +import { selectMerges } from './select'; +import { + compact, + type ChunkDataPool, + type ChunkSource, + type ChunkSourceMetadata +} from '../chunk'; +import { APP_CHUNK } from './gpu-edge-cost'; +import { type ReadFileSystem } from '../io/read'; +import { type FileSystem } from '../io/write'; +import { bakeTransform } from '../ops'; +import { readPly } from '../readers/read-ply'; +import { type DeviceCreator } from '../types'; +import { fmtBytes, fmtCount, logger, Transform } from '../utils'; +import { writePlyStreaming } from '../writers/write-ply-streaming'; + +/** Neighbours per query — unchanged from legacy. */ +const KNN_K = 16; + +/** Default owned gaussians per KD block. */ +const BLOCK_SIZE = 1 << 21; + +/** Same no-grind stall semantics as legacy: a shortfall generation must remove at least this fraction. */ +const MIN_ITERATION_PROGRESS = 0.05; + +/** Default resident-memory budget steering the candidate-K policy. */ +const DEFAULT_MEMORY_BUDGET = 24 * 2 ** 30; + +/** Coherence heuristic: gap (rows) merged into one run / runs-per-block considered scattered. */ +const COHERENCE_GAP_ROWS = 64; +const INCOHERENT_RUNS_PER_BLOCK = 64; + +/** Warn about incoherent input only for scenes big enough for it to matter. */ +const COHERENCE_MIN_N = 1 << 22; + +/** + * Where intermediate generations spill when they exceed the in-memory + * budget. `remove` deletes a spill file once its generation is consumed + * (optional; without it temp files are left behind). + */ +type DecimateSpill = { + writeFs: FileSystem; + readFs: ReadFileSystem; + scratchDir: string; + remove?: (path: string) => Promise; +}; + +type DecimateOptions = { + /** Exact number of gaussians to keep (≥ 1). */ + targetCount: number; + /** Optional GPU device factory; CPU fallback without it. */ + createDevice?: DeviceCreator; + /** Spill destination for over-budget intermediate generations. */ + spill?: DecimateSpill; + /** Resident-memory budget driving the candidate-K policy (default 24 GiB). */ + memoryBudgetBytes?: number; +}; + +// Candidate-K policy: keep 4 when the resident estimate fits the budget, +// else 2. Estimate: positions (12) + candidates (K*8) + selection counting +// sort (K*4) + memberGroup (4) per gaussian, plus a flat block-working fudge. +const chooseK = (n: number, budget: number): number => { + const estimate = (K: number) => n * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; + return estimate(4) <= budget ? 4 : 2; +}; + +// Read the position layer sequentially into resident columns (generation 1 +// only; later generations carry positions forward from the merge stream). +const extractPositions = async (source: ChunkSource, pool: ChunkDataPool): Promise => { + const { meta } = source; + const n = meta.numGaussians; + const out: ResidentPositions = { + x: new Float32Array(n), + y: new Float32Array(n), + z: new Float32Array(n) + }; + const bar = logger.bar('reading positions', meta.numChunks[0] ?? 0); + let base = 0; + for (let c = 0; c < (meta.numChunks[0] ?? 0); c++) { + const count = Math.min(meta.chunkSize, n - c * meta.chunkSize); + const cd = pool.acquire('position', meta.layouts.position!, count); + await source.read({ chunkIndex: c, position: cd }); + const p = new Float32Array(cd.data, 0, count * 3); + for (let i = 0; i < count; i++) { + out.x[base + i] = p[i * 3]; + out.y[base + i] = p[i * 3 + 1]; + out.z[base + i] = p[i * 3 + 2]; + } + base += count; + cd.release(); + bar.tick(); + } + bar.end(); + return out; +}; + +/** + * Chunk-native, memory-bounded decimation to an exact target count. + * + * Design: positions resident; KD blocks as an IO pattern only; per-block exact global 16-NN + + * edge costs (GPU when a device is supplied) reduced to K resident + * candidates; global bucketed greedy matching with chain closure; a second + * heavy pass moment-matches groups and streams the output. + * + * The returned source supports a single sequential pass (it computes the + * merge stream on demand) — the PLY-terminal consumption model. Its `close` + * releases the input source and any intermediate spill files. Deep targets + * run multiple generations; intermediates land in RAM when small enough, + * else in temp PLY spills under `opts.spill.scratchDir`. + * + * @param source - Input (consumed: the returned source owns it). Single LOD, gaussian layers required. + * @param pool - Chunk-data pool; its chunk size must match the source's. + * @param opts - Options. + * @returns The decimated stream-once source with exact metadata. + */ +const decimateSource = async ( + source: ChunkSource, + pool: ChunkDataPool, + opts: DecimateOptions +): Promise => { + const { targetCount } = opts; + const inputMeta = source.meta; + + if (inputMeta.numLods > 1) { + throw new Error( + `decimate requires a single-LOD source (got ${inputMeta.numLods} LODs); select a level first (--select-lod / selectLod)` + ); + } + for (const layer of ['position', 'geometric', 'color'] as const) { + if (!inputMeta.availableLayers.has(layer)) { + throw new Error(`decimate requires gaussian splat data (missing '${layer}' layer)`); + } + } + if (targetCount < 1) { + throw new Error(`decimate target must be at least 1 (got ${targetCount})`); + } + if (targetCount >= inputMeta.numGaussians) { + return source; + } + + const device: GraphicsDevice | undefined = opts.createDevice ? await opts.createDevice() : undefined; + const budget = opts.memoryBudgetBytes ?? DEFAULT_MEMORY_BUDGET; + const colorDim = inputMeta.layouts.color!.stride >> 2; + const otherStride = inputMeta.layouts.other?.stride ?? 0; + + // Bake the pending transform to PLY space up front (identity fast-path): + // intermediate spills go through writePlyStreaming (which bakes anyway) + // and carried-forward resident positions must match spilled values. + // Decimation is TRS-covariant, so bake timing cannot change the result. + let src: ChunkSource = bakeTransform(source, Transform.PLY); + let positions: ResidentPositions | null = null; + // Cleanup for the CURRENT generation's input (previous spill / RAM source). + let disposeCurrentInput: (() => Promise) | null = null; + + const totalGenerations = Math.max(1, Math.ceil(Math.log2(inputMeta.numGaussians / targetCount))); + + for (let generation = 1; ; generation++) { + const N = src.meta.numGaussians; + const gen = logger.group('Decimate generation', { + index: Math.min(generation, totalGenerations), + total: totalGenerations + }); + + positions ??= await extractPositions(src, pool); + + // Device binding-limit clamp: the largest per-binding buffer scales + // with block size, never scene size; halve the block size until it + // fits the adapter's storage-binding limit. + let blockSize = BLOCK_SIZE; + const bindingLimit = (device as unknown as { limits?: { maxStorageBufferBindingSize?: number } } | undefined) + ?.limits?.maxStorageBufferBindingSize; + if (typeof bindingLimit === 'number') { + const largestBinding = (bs: number) => bs * (1 + HALO_CAP) * Math.max(Math.min(APP_CHUNK, colorDim) * 4, 36); + while (blockSize > (1 << 16) && largestBinding(blockSize) > bindingLimit) { + blockSize >>= 1; + } + if (blockSize !== BLOCK_SIZE) { + logger.warn(`reducing decimate block size to ${fmtCount(blockSize)} to fit GPU binding limit ${fmtBytes(bindingLimit)}`); + } + } + + const partSub = logger.group('Partitioning'); + const { order, blocks } = kdPartition(positions, blockSize); + partSub.end(); + + if (generation === 1 && N >= COHERENCE_MIN_N) { + const runs = blocks.map(b => coherenceRuns(order, b.start, b.end, COHERENCE_GAP_ROWS)).sort((a, b) => a - b); + const median = runs[runs.length >> 1] ?? 0; + if (median > INCOHERENT_RUNS_PER_BLOCK) { + logger.warn( + 'input is spatially incoherent (scattered gathers expected); run a one-time --morton-order prepass for much faster IO' + ); + } + } + + const K = chooseK(N, budget); + const cand: CandidateArrays = { + idx: new Uint32Array(N * K).fill(0xFFFFFFFF), + cost: new Float32Array(N * K).fill(Infinity) + }; + + const priorityBar = logger.bar('computing merge priorities', N); + await runPriorityPass( + { source: src, pool, pos: positions, order, blocks, device, K, k: Math.min(KNN_K, Math.max(1, N - 1)) }, + cand, + n => priorityBar.tick(n) + ); + priorityBar.end(); + + const generationTarget = Math.max(targetCount, N - Math.floor(N / 2)); + const needed = N - generationTarget; + const selectSub = logger.group('Selecting merges'); + const selection = selectMerges(cand, N, K, needed); + selectSub.end(); + + if (selection.removed === 0) { + gen.end(); + const cause = device ? + 'the GPU step likely failed (e.g. out-of-memory) or produced non-finite costs' : + 'cost computation produced no finite merge candidates (e.g. non-finite inputs)'; + throw new Error( + `decimation found no valid merges at ${N} splats (target ${targetCount}) — ${cause}. ` + + 'Refusing to return an incompletely-decimated scene.' + ); + } + const removedFraction = selection.removed / N; + if (selection.removed < needed && removedFraction < MIN_ITERATION_PROGRESS) { + gen.end(); + throw new Error( + `decimation stalled at ${N} splats (target ${targetCount}): a generation removed only ` + + `${selection.removed} splat${selection.removed === 1 ? '' : 's'} (${(removedFraction * 100).toFixed(3)}% of ${N}) — ` + + 'the nearest-neighbour graph is too degenerate to merge further (e.g. many coincident splats). ' + + 'Refusing to grind toward the target.' + ); + } + + const outCount = N - selection.removed; + const outMeta: ChunkSourceMetadata = { + numGaussians: outCount, + numLods: 1, + lodCounts: [outCount], + chunkSize: src.meta.chunkSize, + numChunks: [Math.ceil(outCount / src.meta.chunkSize)], + shBands: src.meta.shBands, + extraColumns: src.meta.extraColumns, + transform: src.meta.transform, + availableLayers: src.meta.availableLayers, + layouts: src.meta.layouts + }; + + const isFinal = outCount <= targetCount; + const nextPositions: ResidentPositions | undefined = isFinal ? undefined : { + x: new Float32Array(outCount), + y: new Float32Array(outCount), + z: new Float32Array(outCount) + }; + + // `src` is reassigned each generation; capture this generation's + // values for the deferred producer closures. + const genSrc = src; + const genChunkSize = genSrc.meta.chunkSize; + const streamCtx = { source: genSrc, pool, pos: positions, order, blocks, selection, nextPositions }; + + if (isFinal) { + // The producer reads the input lazily while the consumer pulls + // chunks: the input chain (and any pending spill) is released on + // close. The merge bar lives outside the generation group since + // streaming happens after this function returns. + gen.end(); + const mergeBar = logger.bar('merging', N); + const producer = createBlockProducerSource(outMeta, () => mergeStream(streamCtx, genChunkSize, n => mergeBar.tick(n))); + const disposeSpill = disposeCurrentInput; + let closed = false; + return { + meta: producer.meta, + read: request => producer.read(request), + close: async () => { + if (closed) return; + closed = true; + mergeBar.end(); + await producer.close(); + await genSrc.close(); + await disposeSpill?.(); + } + }; + } + + const mergeBar = logger.bar('merging', N); + const producer = createBlockProducerSource(outMeta, () => mergeStream(streamCtx, genChunkSize, n => mergeBar.tick(n))); + + // Intermediate generation: materialize (RAM when comfortably within + // budget, else temp PLY spill), then advance the loop. + const estBytes = outCount * (12 + 32 + colorDim * 4 + otherStride); + let nextSrc: ChunkSource; + let disposeNext: (() => Promise) | null = null; + + if (estBytes <= budget / 4) { + nextSrc = await compact(producer, pool); + } else { + if (!opts.spill) { + throw new Error( + `decimation intermediate generation needs ${fmtBytes(estBytes)}, over the in-memory budget — ` + + 'a spill location is required (opts.spill / --scratch-dir)' + ); + } + const spill = opts.spill; + const filename = join(spill.scratchDir, `.decimate-gen${generation}.${Date.now().toString(36)}.tmp.ply`); + await writePlyStreaming(producer, pool, { filename }, spill.writeFs); + const readSource = await spill.readFs.createSource(filename); + const plySrc = await readPly(readSource, pool); + nextSrc = plySrc; + disposeNext = async () => { + await plySrc.close(); + await spill.remove?.(filename); + }; + } + mergeBar.end(); + await producer.close(); + + // The consumed input of THIS generation can now be released: for + // generation 1 that is the caller's source (we own it), for later + // generations the previous spill / RAM intermediate. + await src.close(); + await disposeCurrentInput?.(); + disposeCurrentInput = disposeNext; + + positions = nextPositions!; + src = nextSrc; + gen.end(); + } +}; + +export { decimateSource, type DecimateOptions, type DecimateSpill }; diff --git a/src/lib/decimate/edge-cost-legacy.ts b/src/lib/decimate-uniform/edge-cost-cpu.ts similarity index 67% rename from src/lib/decimate/edge-cost-legacy.ts rename to src/lib/decimate-uniform/edge-cost-cpu.ts index 47d8633a..8b46b900 100644 --- a/src/lib/decimate/edge-cost-legacy.ts +++ b/src/lib/decimate-uniform/edge-cost-cpu.ts @@ -20,24 +20,14 @@ import { gaussLogpdfDiagrot, type SplatView, type MergeScratch -} from './moment-match'; -import { type EdgeCostCacheLegacy } from '../gpu/gpu-edge-cost-legacy'; - -/** - * Appearance columns per storage chunk of the legacy GPU kernel. The kernel - * exposes three appearance bindings (appA/appB/appC), so the layout holds up - * to 3·APP_CHUNK columns; at 16 the widest chunk reaches the ~2 GB - * per-binding limit around ~33.5M splats. The kernel imports this same - * constant, so its strides and the host packing can't drift. - */ -export const APP_CHUNK = 16; +} from '../decimate/moment-match'; /** * Per-splat derived quantities for the cost function (legacy * `buildPerSplatCache`, forGpu = false). `mass` uses the cost-path epsilon * (+1e-12), matching legacy exactly. */ -type LegacyCostCache = { +type CostCache = { R: Float32Array; v: Float32Array; invdiag: Float32Array; @@ -46,7 +36,7 @@ type LegacyCostCache = { mass: Float32Array; }; -const buildCostCacheLegacy = (view: SplatView): LegacyCostCache => { +const buildCostCache = (view: SplatView): CostCache => { const { geo } = view; const n = geo.length / 8; const R = new Float32Array(n * 9); @@ -96,16 +86,16 @@ const buildCostCacheLegacy = (view: SplatView): LegacyCostCache => { * `computeEdgeCost`, verbatim. * * @param view - Splat columns. - * @param cache - Per-splat cache from {@link buildCostCacheLegacy}. + * @param cache - Per-splat cache from {@link buildCostCache}. * @param i - First splat (view row). * @param j - Second splat (view row). * @param Z - MC samples (legacy: one sample, seed 0). * @param scratch - Merge scratch (uses `sigm`). * @returns The edge cost. */ -const computeEdgeCostViewLegacy = ( +const computeEdgeCostView = ( view: SplatView, - cache: LegacyCostCache, + cache: CostCache, i: number, j: number, Z: Float64Array[], @@ -209,62 +199,4 @@ const computeEdgeCostViewLegacy = ( return geoCost + cSh; }; -// Pack the block view into the GpuEdgeCostLegacy cache layout (legacy packing: -// posScalars 8-wide, rotR from normalized quats, appearance in ≤APP_CHUNK -// column chunks with live-width strides). -const packGpuCacheLegacy = (view: SplatView): EdgeCostCacheLegacy => { - const { pos, geo, color, colorDim } = view; - const n = geo.length / 8; - const posScalars = new Float32Array(n * 8); - const rotR = new Float32Array(n * 9); - const rot = new Float32Array(9); - - for (let i = 0; i < n; i++) { - const i8 = i * 8; - const o = i * 8; - const linAlpha = sigmoid(geo[i8 + 7]); - const sx = Math.max(Math.exp(geo[i8 + 4]), 1e-12); - const sy = Math.max(Math.exp(geo[i8 + 5]), 1e-12); - const sz = Math.max(Math.exp(geo[i8 + 6]), 1e-12); - const vx = sx * sx + 1e-8; - const vy = sy * sy + 1e-8; - const vz = sz * sz + 1e-8; - posScalars[o] = pos[i * 3]; - posScalars[o + 1] = pos[i * 3 + 1]; - posScalars[o + 2] = pos[i * 3 + 2]; - posScalars[o + 3] = linAlpha * ellipsoidArea(sx, sy, sz) + 1e-12; - posScalars[o + 4] = Math.log(Math.max(vx, 1e-30)) + Math.log(Math.max(vy, 1e-30)) + Math.log(Math.max(vz, 1e-30)); - posScalars[o + 5] = vx; - posScalars[o + 6] = vy; - posScalars[o + 7] = vz; - - let qw = geo[i8], qx = geo[i8 + 1], qy = geo[i8 + 2], qz = geo[i8 + 3]; - const invq = 1 / Math.max(Math.hypot(qw, qx, qy, qz), 1e-12); - qw *= invq; qx *= invq; qy *= invq; qz *= invq; - const xx = qx * qx, yy = qy * qy, zz = qz * qz; - const wx = qw * qx, wy = qw * qy, wz = qw * qz; - const xy = qx * qy, xz = qx * qz, yz = qy * qz; - rot[0] = 1 - 2 * (yy + zz); rot[1] = 2 * (xy - wz); rot[2] = 2 * (xz + wy); - rot[3] = 2 * (xy + wz); rot[4] = 1 - 2 * (xx + zz); rot[5] = 2 * (yz - wx); - rot[6] = 2 * (xz - wy); rot[7] = 2 * (yz + wx); rot[8] = 1 - 2 * (xx + yy); - rotR.set(rot, i * 9); - } - - const numChunks = Math.ceil(colorDim / APP_CHUNK); - const appChunks: Float32Array[] = []; - for (let ch = 0; ch < numChunks; ch++) { - const kStart = ch * APP_CHUNK; - const width = Math.min(APP_CHUNK, colorDim - kStart); - const chunk = new Float32Array(n * width); - for (let s = 0; s < n; s++) { - const dst = s * width; - const src = s * colorDim + kStart; - for (let kk = 0; kk < width; kk++) chunk[dst + kk] = color[src + kk]; - } - appChunks.push(chunk); - } - - return { posScalars, rotR, appChunks, numAppCols: colorDim, numSplats: n }; -}; - -export { buildCostCacheLegacy, computeEdgeCostViewLegacy, packGpuCacheLegacy, type LegacyCostCache }; +export { buildCostCache, computeEdgeCostView, type CostCache }; diff --git a/src/lib/gpu/gpu-edge-cost-legacy.ts b/src/lib/decimate-uniform/gpu-edge-cost.ts similarity index 78% rename from src/lib/gpu/gpu-edge-cost-legacy.ts rename to src/lib/decimate-uniform/gpu-edge-cost.ts index a115b7b5..46928781 100644 --- a/src/lib/gpu/gpu-edge-cost-legacy.ts +++ b/src/lib/decimate-uniform/gpu-edge-cost.ts @@ -16,9 +16,14 @@ import { UniformFormat } from 'playcanvas'; -import { APP_CHUNK } from '../decimate/edge-cost-legacy'; - -export { APP_CHUNK }; +/** + * Appearance columns per storage chunk. The kernel exposes three appearance + * bindings (appA/appB/appC), so the layout holds up to 3·APP_CHUNK columns; at + * 16 the widest chunk reaches the ~2 GB per-binding limit around ~33.5M splats. + * The CPU-side packing in `decimate/priority.ts` imports this same constant, + * so the kernel strides and the host packing can't drift. + */ +export const APP_CHUNK = 16; /** * WGSL kernel: per-edge KL-style cost (matches `computeEdgeCostView` in @@ -30,47 +35,45 @@ export { APP_CHUNK }; * (the same `z` for both components, matching the CPU implementation), * and adds an L2 distance over the appearance (SH) coefficients. * - * @param k - Compile-time K, neighbour slots per owned row. * @param strideA - Live column count of appearance chunk A (0 if unused). * @param strideB - Live column count of appearance chunk B (0 if unused). * @param strideC - Live column count of appearance chunk C (0 if unused). * @returns WGSL source. */ -const edgeCostWgsl = (k: number, strideA: number, strideB: number, strideC: number) => /* wgsl */` +const edgeCostWgsl = (strideA: number, strideB: number, strideC: number) => /* wgsl */` struct Uniforms { - slotBase: u32, - slotCount: u32, + edgeCount: u32, z0: f32, z1: f32, z2: f32, } @group(0) @binding(0) var uniforms: Uniforms; -// Neighbour rows for the current dispatch batch (host uploads each batch's -// slice to offset 0): view-local row per slot, 0xFFFFFFFF for empty slots. -// Slot s belongs to owned row (slotBase + s) / K (dense-slot edge model). -@group(0) @binding(1) var nbRow: array; +// Edge list for the current dispatch batch only, split into two parallel +// arrays (avoids a host-side (i, j) interleave). The host uploads each batch's +// slice to offset 0, so we index edgesI/J[bid] directly — keeping these +// buffers batch-sized instead of N·k keeps them off the per-binding limit. +@group(0) @binding(1) var edgesI: array; +@group(0) @binding(2) var edgesJ: array; // Per-splat geometry, interleaved 8-wide: // posScalars[8s + 0..2] = position xyz // posScalars[8s + 3] = mass // posScalars[8s + 4] = logdet // posScalars[8s + 5..7] = variances (vx, vy, vz) -@group(0) @binding(2) var posScalars: array; +@group(0) @binding(3) var posScalars: array; // Row-major 3x3 rotation matrix per splat (9 floats per splat). -@group(0) @binding(3) var rotR: array; +@group(0) @binding(4) var rotR: array; // Appearance, split into up to three chunks (≤16 columns each) so no single // binding exceeds maxStorageBufferBindingSize (~2 GB). Each chunk's stride is // its live column count (STRIDE_A/B/C below); appA holds columns 0.., appB the // next span, appC the next. Unused chunks have stride 0, are bound to a dummy // buffer, and are never read. -@group(0) @binding(4) var appA: array; -@group(0) @binding(5) var appB: array; -@group(0) @binding(6) var appC: array; -// Output: cost per slot. -@group(0) @binding(7) var costs: array; - -const K: u32 = ${k}u; -const SENTINEL: u32 = 0xFFFFFFFFu; +@group(0) @binding(5) var appA: array; +@group(0) @binding(6) var appB: array; +@group(0) @binding(7) var appC: array; +// Output: cost per edge. +@group(0) @binding(8) var costs: array; + const EPS_COV: f32 = 1e-8; const LOG2PI: f32 = 1.8378770664093453; // Per-chunk appearance strides = live column count in each chunk (0 = unused, @@ -128,16 +131,10 @@ fn logAddExp(a: f32, b: f32) -> f32 { @compute @workgroup_size(64) fn main(@builtin(global_invocation_id) gid: vec3u) { let bid = gid.x; - if (bid >= uniforms.slotCount) { return; } - - // Empty slot: the reduction skips sentinel slots by id, so the cost value - // is never read. - let j = nbRow[bid]; - if (j == SENTINEL) { - costs[bid] = 0.0; - return; - } - let i = (uniforms.slotBase + bid) / K; + if (bid >= uniforms.edgeCount) { return; } + + let i = edgesI[bid]; + let j = edgesJ[bid]; let i8 = i * 8u; let j8 = j * 8u; @@ -260,7 +257,7 @@ fn main(@builtin(global_invocation_id) gid: vec3u) { * WebGPU per-stage storage-buffer count limit (8) and the per-binding size * limit (~2 GB) — appearance is split into 16-column chunks for the latter. */ -interface EdgeCostCacheLegacy { +interface EdgeCostCache { /** Per-splat geometry interleaved 8-wide: (x, y, z, mass, logdet, vx, vy, vz). */ posScalars: Float32Array; /** Row-major 3×3 rotation per splat (length 9N). */ @@ -287,17 +284,18 @@ interface EdgeCostCacheLegacy { * * Mirrors the CPU `computeEdgeCostView` in `decimate/edge-cost-cpu.ts`. */ -class GpuEdgeCostLegacy { +class GpuEdgeCost { /** * @param cache - Per-splat cache (uploaded once). - * @param nbRows - Dense neighbour slots (view-local row per slot, - * 0xFFFFFFFF sentinel for empty; slot s belongs to owned row s / k). + * @param edgeI - Edge u indices (length E). + * @param edgeJ - Edge v indices (length E). * @param z - Single Monte-Carlo sample (3 floats from N(0,1)). - * @param outCosts - Destination for per-slot costs (length = slots). + * @param outCosts - Destination for per-edge costs (length E). */ execute: ( - cache: EdgeCostCacheLegacy, - nbRows: Uint32Array, + cache: EdgeCostCache, + edgeI: Uint32Array, + edgeJ: Uint32Array, z: Float32Array, outCosts: Float32Array ) => Promise; @@ -306,13 +304,12 @@ class GpuEdgeCostLegacy { /** * @param device - PlayCanvas GraphicsDevice (WebGPU). * @param maxN - Maximum number of splats. - * @param k - Neighbour slots per owned row. + * @param maxE - Maximum number of edges in a single dispatch. * @param maxAppCols - Maximum appearance column count (over all bands). */ - constructor(device: GraphicsDevice, maxN: number, k: number, maxAppCols: number) { + constructor(device: GraphicsDevice, maxN: number, maxE: number, maxAppCols: number) { const workgroupSize = 64; - // Slots per dispatch: bounded by the 65,535 workgroups-per-dimension limit. - const slotsPerBatch = 65535 * workgroupSize; // 4,194,240 + const edgesPerBatch = 1024 * workgroupSize; // 65,536 // Appearance is split at fixed APP_CHUNK-column boundaries, but each // chunk's *stride* is its live column count — only the last non-empty // chunk is ever partial, so partial chunks neither allocate nor upload @@ -331,7 +328,8 @@ class GpuEdgeCostLegacy { const bindGroupFormat = new BindGroupFormat(device, [ new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE), - new BindStorageBufferFormat('nbRow', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('edgesI', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('edgesJ', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('posScalars', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('rotR', SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat('appA', SHADERSTAGE_COMPUTE, true), @@ -341,14 +339,13 @@ class GpuEdgeCostLegacy { ]); const shader = new Shader(device, { - name: 'compute-edge-cost-legacy', + name: 'compute-edge-cost', shaderLanguage: SHADERLANGUAGE_WGSL, - cshader: edgeCostWgsl(k, appStrides[0], appStrides[1], appStrides[2]), + cshader: edgeCostWgsl(appStrides[0], appStrides[1], appStrides[2]), // @ts-ignore computeUniformBufferFormats: { uniforms: new UniformBufferFormat(device, [ - new UniformFormat('slotBase', UNIFORMTYPE_UINT), - new UniformFormat('slotCount', UNIFORMTYPE_UINT), + new UniformFormat('edgeCount', UNIFORMTYPE_UINT), new UniformFormat('z0', UNIFORMTYPE_FLOAT), new UniformFormat('z1', UNIFORMTYPE_FLOAT), new UniformFormat('z2', UNIFORMTYPE_FLOAT) @@ -369,7 +366,7 @@ class GpuEdgeCostLegacy { const checkLimit = (label: string, bytes: number) => { if (bytes > maxStorage) { throw new Error( - `GpuEdgeCostLegacy: ${label} buffer (${bytes} bytes) exceeds device ` + + `GpuEdgeCost: ${label} buffer (${bytes} bytes) exceeds device ` + `maxStorageBufferBindingSize (${maxStorage})` ); } @@ -393,20 +390,24 @@ class GpuEdgeCostLegacy { appDummy; }); - // Neighbour-row buffer sized to a single dispatch batch (not the full - // N·k slot list): execute uploads each batch's slice before its - // dispatch, keeping it off the ~2 GB per-binding limit. - const nbRowBuf = new StorageBuffer(device, slotsPerBatch * 4, BUFFERUSAGE_COPY_DST); + // Two parallel u32 buffers, sized to a single dispatch batch (not the + // full N·k edge list): execute uploads each batch's slice before its + // dispatch. Batch-sizing keeps these ~256 KB instead of N·k·4 — off the + // ~2 GB per-binding limit (so edges never cap scene size) and ~1.6 GB + // less VRAM at 13M splats. Two parallel arrays avoid a host-side pack. + const edgesIBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); + const edgesJBuf = new StorageBuffer(device, edgesPerBatch * 4, BUFFERUSAGE_COPY_DST); const outBuf = new StorageBuffer( device, - slotsPerBatch * 4, + edgesPerBatch * 4, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST ); - const outScratch = new Float32Array(slotsPerBatch); + const outScratch = new Float32Array(edgesPerBatch); - const compute = new Compute(device, shader, 'compute-edge-cost-legacy'); - compute.setParameter('nbRow', nbRowBuf); + const compute = new Compute(device, shader, 'compute-edge-cost'); + compute.setParameter('edgesI', edgesIBuf); + compute.setParameter('edgesJ', edgesJBuf); compute.setParameter('posScalars', posScalarsBuf); compute.setParameter('rotR', rotRBuf); compute.setParameter('appA', appBufs[0]); @@ -415,27 +416,28 @@ class GpuEdgeCostLegacy { compute.setParameter('costs', outBuf); this.execute = async ( - cache: EdgeCostCacheLegacy, - nbRows: Uint32Array, + cache: EdgeCostCache, + edgeI: Uint32Array, + edgeJ: Uint32Array, z: Float32Array, outCosts: Float32Array ) => { const n = cache.numSplats; - const s = nbRows.length; + const e = edgeI.length; - if (n > maxN) throw new Error(`GpuEdgeCostLegacy: N=${n} exceeds maxN=${maxN}`); + if (n > maxN) throw new Error(`GpuEdgeCost: N=${n} exceeds maxN=${maxN}`); + if (e > maxE) throw new Error(`GpuEdgeCost: E=${e} exceeds maxE=${maxE}`); if (cache.numAppCols !== maxAppCols) { - throw new Error(`GpuEdgeCostLegacy: numAppCols=${cache.numAppCols} must equal maxAppCols=${maxAppCols} (baked into the kernel)`); + throw new Error(`GpuEdgeCost: numAppCols=${cache.numAppCols} must equal maxAppCols=${maxAppCols} (baked into the kernel)`); } if (cache.appChunks.length !== numAppChunks) { - throw new Error(`GpuEdgeCostLegacy: cache supplies ${cache.appChunks.length} appearance chunks but the kernel layout expects ${numAppChunks}`); + throw new Error(`GpuEdgeCost: cache supplies ${cache.appChunks.length} appearance chunks but the kernel layout expects ${numAppChunks}`); } - if (outCosts.length !== s) { - throw new Error('GpuEdgeCostLegacy: nbRows / outCosts must have same length'); + if (edgeJ.length !== e || outCosts.length !== e) { + throw new Error('GpuEdgeCost: edgeI / edgeJ / outCosts must have same length'); } - if (s % k !== 0) throw new Error(`GpuEdgeCostLegacy: slot count ${s} must be a multiple of k=${k}`); if (z.length < 3) { - throw new Error('GpuEdgeCostLegacy: z must have at least 3 elements'); + throw new Error('GpuEdgeCost: z must have at least 3 elements'); } // Upload per-splat cache. Each appearance chunk is row-major with @@ -450,23 +452,25 @@ class GpuEdgeCostLegacy { compute.setParameter('z1', z[1]); compute.setParameter('z2', z[2]); - const numBatches = Math.ceil(s / slotsPerBatch); + const numBatches = Math.ceil(e / edgesPerBatch); for (let batch = 0; batch < numBatches; batch++) { - const slotBase = batch * slotsPerBatch; - const slotCount = Math.min(slotsPerBatch, s - slotBase); - const groups = Math.ceil(slotCount / workgroupSize); + const edgeOffset = batch * edgesPerBatch; + const edgeCount = Math.min(edgesPerBatch, e - edgeOffset); + const groups = Math.ceil(edgeCount / workgroupSize); - nbRowBuf.write(0, nbRows, slotBase, slotCount); + // Upload just this batch's edges to offset 0; the kernel indexes + // edgesI/J[bid] within the batch. + edgesIBuf.write(0, edgeI, edgeOffset, edgeCount); + edgesJBuf.write(0, edgeJ, edgeOffset, edgeCount); - compute.setParameter('slotBase', slotBase); - compute.setParameter('slotCount', slotCount); + compute.setParameter('edgeCount', edgeCount); compute.setupDispatch(groups); device.computeDispatch([compute], `edge-cost-dispatch-${batch}`); - const readBytes = slotCount * 4; + const readBytes = edgeCount * 4; await outBuf.read(0, readBytes, outScratch, true); - outCosts.set(outScratch.subarray(0, slotCount), slotBase); + outCosts.set(outScratch.subarray(0, edgeCount), edgeOffset); } }; @@ -477,7 +481,8 @@ class GpuEdgeCostLegacy { if (buf !== appDummy) buf.destroy(); } appDummy.destroy(); - nbRowBuf.destroy(); + edgesIBuf.destroy(); + edgesJBuf.destroy(); outBuf.destroy(); shader.destroy(); bindGroupFormat.destroy(); @@ -485,4 +490,4 @@ class GpuEdgeCostLegacy { } } -export { GpuEdgeCostLegacy, type EdgeCostCacheLegacy }; +export { GpuEdgeCost, type EdgeCostCache }; diff --git a/src/lib/decimate-uniform/gpu-knn.ts b/src/lib/decimate-uniform/gpu-knn.ts new file mode 100644 index 00000000..11f49881 --- /dev/null +++ b/src/lib/decimate-uniform/gpu-knn.ts @@ -0,0 +1,345 @@ +import { + BUFFERUSAGE_COPY_DST, + BUFFERUSAGE_COPY_SRC, + SHADERLANGUAGE_WGSL, + SHADERSTAGE_COMPUTE, + UNIFORMTYPE_UINT, + BindGroupFormat, + BindStorageBufferFormat, + BindUniformBufferFormat, + Compute, + GraphicsDevice, + Shader, + StorageBuffer, + UniformBufferFormat, + UniformFormat +} from 'playcanvas'; + +import { type FlatKdTree } from '../spatial/kd-tree'; + +/** + * Block-local GPU KNN for `--decimate-uniform`. + * + * Separate from the quality path's `gpu/gpu-knn.ts` because the two have + * incompatible models: this one takes a max size and accepts a DIFFERENT tree + * per `execute` (the decimator re-uploads one block's tree at a time), while + * that one bakes each forest part's root and AABB in as compile-time constants + * at construction. + * + * The one deviation from 3.1.x here is the flat-tree layout it reads; see + * README.md in this directory before changing it. + */ + +/** + * WGSL kernel: iterative KD-tree K-nearest-neighbours. + * + * Each thread runs a depth-first traversal of the flattened KD-tree with a + * fixed-size per-thread stack. Visits at most `O(K · log N)` nodes per + * query thanks to the standard "skip the far subtree if its splitting plane + * is farther than the current K-th best" pruning. Top-K is maintained + * unsorted in per-thread storage with explicit worst-index tracking, so the + * common-case "candidate is rejected against worst" path is a single + * compare-and-branch (no dynamic-indexed shift). + * + * @param k - Compile-time K, the number of nearest neighbours per query. + * @param stackSize - Compile-time per-thread DFS stack depth. + * @returns WGSL source. + */ +const knnWgsl = (k: number, stackSize: number) => /* wgsl */` +struct Uniforms { + queryOffset: u32, + queryCount: u32, + rootIdx: u32, +} + +@group(0) @binding(0) var uniforms: Uniforms; +// Query positions interleaved xyz: positions[q*3 + 0/1/2]. +@group(0) @binding(1) var positions: array; +// Flattened KD-tree. Positions and children are interleaved so the kernel +// stays comfortably under the WebGPU per-stage storage-buffer minimum (8): +// nodePositions[t*3 + 0/1/2] for tree node t, nodeChildren[t*2 + 0/1] for +// (left, right). Kept separate from nodeSplatIdx to avoid mixing f32/u32. +@group(0) @binding(2) var nodeSplatIdx: array; +@group(0) @binding(3) var nodePositions: array; +@group(0) @binding(4) var nodeChildren: array; +// Output: per query, k neighbour splat indices (unsorted). +@group(0) @binding(5) var outIndices: array; + +const K: u32 = ${k}u; +const NULL_NODE: u32 = 0xFFFFFFFFu; +const F32_MAX: f32 = 3.4028234663852886e+38; +// log2(N) + slack — safe to ~2^40 nodes which is way past our limits. +const STACK_SIZE: u32 = ${stackSize}u; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3u) { + let bid = gid.x; + if (bid >= uniforms.queryCount) { return; } + let q = bid + uniforms.queryOffset; + + let q3 = q * 3u; + let qx = positions[q3 + 0u]; + let qy = positions[q3 + 1u]; + let qz = positions[q3 + 2u]; + + // Top-K state, unsorted. worstIdx points to the current K-th worst slot + // so accepts replace it in O(1) and we recompute worst via a fixed loop. + var topIdx: array; + var topDist: array; + var worst: f32 = F32_MAX; + var worstIdx: u32 = 0u; + for (var i: u32 = 0u; i < K; i++) { + topDist[i] = F32_MAX; + topIdx[i] = 0u; + } + + // Stack: (nodeIdx, axis) packed as u32. axis ∈ {0,1,2} in top 2 bits, + // nodeIdx in low 30 — supports up to ~1B nodes. + var stack: array; + var sp: u32 = 0u; + stack[0] = uniforms.rootIdx; // axis=0 → no axis bits set + sp = 1u; + + while (sp > 0u) { + sp = sp - 1u; + let packed = stack[sp]; + let nodeIdx = packed & 0x3FFFFFFFu; + let axis = packed >> 30u; + + // Read the node's position + splat id. + let np = nodeIdx * 3u; + let nx = nodePositions[np + 0u]; + let ny = nodePositions[np + 1u]; + let nz = nodePositions[np + 2u]; + let splatId = nodeSplatIdx[nodeIdx]; + + // Update top-K, skipping the query itself. + if (splatId != q) { + let dx = nx - qx; + let dy = ny - qy; + let dz = nz - qz; + let d2 = dx * dx + dy * dy + dz * dz; + if (d2 < worst) { + topDist[worstIdx] = d2; + topIdx[worstIdx] = splatId; + // Recompute worst with a constant-bound loop (compiler can + // unroll → all accesses to topDist resolve statically). + var w: f32 = topDist[0]; + var wi: u32 = 0u; + for (var i: u32 = 1u; i < K; i++) { + if (topDist[i] > w) { w = topDist[i]; wi = i; } + } + worst = w; + worstIdx = wi; + } + } + + // Choose near/far children based on which side of the splitting + // plane the query lies on. Walk near first (push far first so LIFO + // pops near first), with pruning on far. + var qAxisVal: f32; + var nAxisVal: f32; + if (axis == 0u) { qAxisVal = qx; nAxisVal = nx; } + else if (axis == 1u) { qAxisVal = qy; nAxisVal = ny; } + else { qAxisVal = qz; nAxisVal = nz; } + + let delta = qAxisVal - nAxisVal; + let nextAxis = select(axis + 1u, 0u, axis + 1u >= 3u); + let nextAxisPacked = nextAxis << 30u; + + let nc = nodeIdx * 2u; + let leftChild = nodeChildren[nc + 0u]; + let rightChild = nodeChildren[nc + 1u]; + let near = select(rightChild, leftChild, delta < 0.0); + let far = select(leftChild, rightChild, delta < 0.0); + + // Push far first iff its subtree could still hold a closer point + // than the current K-th best. + if (far != NULL_NODE && delta * delta < worst) { + stack[sp] = far | nextAxisPacked; + sp = sp + 1u; + } + if (near != NULL_NODE) { + stack[sp] = near | nextAxisPacked; + sp = sp + 1u; + } + } + + // Emit unsorted top-K (the decimator does not require sorted neighbours). + // Slots that never received a real candidate (n-1 < K) keep F32_MAX in + // topDist; emit the sentinel 0xFFFFFFFF for those so downstream + // edge-extraction can skip them, matching the CPU path. + let outBase = bid * K; + for (var i: u32 = 0u; i < K; i++) { + if (topDist[i] == F32_MAX) { + outIndices[outBase + i] = 0xFFFFFFFFu; + } else { + outIndices[outBase + i] = topIdx[i]; + } + } +} +`; + +/** + * GPU K-nearest-neighbours over a fixed point set using a flattened KD-tree. + * + * Algorithm: classic KD-tree DFS with bounded heap pruning, except the + * recursion is unrolled into an explicit per-thread stack and the top-K is + * maintained unsorted (with worst-index tracking) so the dominant + * candidate-rejection path is a single compare. Same O(N log N) total work + * as the CPU KD-tree the kernel mirrors, just parallelised across queries. + * + * The flattened tree is built by the caller (`buildFlatKdTree`, typically + * off-thread via the `flattenKdTree` worker task) — this class only uploads + * and traverses it. + * + * Memory footprint: ~24 N bytes for the flattened tree (3 floats + 3 + * u32 per node), plus query positions and the per-query output indices. + */ +class GpuKnn { + /** + * @param tree - Prebuilt flattened KD-tree over the `n` local points + * (see `buildFlatKdTree`; node splat ids are LOCAL indices). + * @param positions - Interleaved xyz for all `n` local points; queries + * are the first `queryCount` of them (owned-first ordering). + * @param n - Total local point count (tree size). + * @param queryCount - How many leading points to query. + * @param outNeighbours - destination for per-query K neighbour indices, + * length `queryCount * k`. `outNeighbours[i * k + j]` is one of the k + * nearest LOCAL neighbours of point i (UNSORTED). Excludes i itself; + * sentinel 0xFFFFFFFF fills surplus slots. + */ + execute: ( + tree: FlatKdTree, + positions: Float32Array, + n: number, + queryCount: number, + outNeighbours: Uint32Array + ) => Promise; + destroy: () => void; + + /** + * @param device - PlayCanvas GraphicsDevice (WebGPU). + * @param maxN - Maximum number of points the index will handle. + * @param k - Number of nearest neighbours per query. + */ + constructor(device: GraphicsDevice, maxN: number, k: number) { + const workgroupSize = 64; + const queriesPerBatch = 1024 * workgroupSize; // 65,536 + // Per-thread DFS stack depth: tree depth = log2(maxN) + slack. 48 is + // safe for any N within the 30-bit nodeIdx packing limit checked below. + const stackSize = 48; + if (maxN > 0x3FFFFFFF) { + throw new Error(`GpuKnn: maxN=${maxN} exceeds 30-bit nodeIdx packing limit (~1B nodes)`); + } + + // 5 storage buffers + 1 uniform — comfortably under the WebGPU + // per-stage minimum (8 storage buffers). Positions and KD-tree + // arrays are interleaved (see WGSL above) to keep the count down. + const bindGroupFormat = new BindGroupFormat(device, [ + new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE), + new BindStorageBufferFormat('positions', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('nodeSplatIdx', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('nodePositions', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('nodeChildren', SHADERSTAGE_COMPUTE, true), + new BindStorageBufferFormat('outIndices', SHADERSTAGE_COMPUTE) + ]); + + const shader = new Shader(device, { + name: 'compute-knn-kdtree', + shaderLanguage: SHADERLANGUAGE_WGSL, + cshader: knnWgsl(k, stackSize), + // @ts-ignore + computeUniformBufferFormats: { + uniforms: new UniformBufferFormat(device, [ + new UniformFormat('queryOffset', UNIFORMTYPE_UINT), + new UniformFormat('queryCount', UNIFORMTYPE_UINT), + new UniformFormat('rootIdx', UNIFORMTYPE_UINT) + ]) + }, + // @ts-ignore + computeBindGroupFormat: bindGroupFormat + }); + + const positionsBuf = new StorageBuffer(device, maxN * 3 * 4, BUFFERUSAGE_COPY_DST); + const nSplatIdxBuf = new StorageBuffer(device, maxN * 4, BUFFERUSAGE_COPY_DST); + const nPositionsBuf = new StorageBuffer(device, maxN * 3 * 4, BUFFERUSAGE_COPY_DST); + const nChildrenBuf = new StorageBuffer(device, maxN * 2 * 4, BUFFERUSAGE_COPY_DST); + + const outBatchBytes = queriesPerBatch * k * 4; + const outBuf = new StorageBuffer( + device, + outBatchBytes, + BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST + ); + const outScratch = new Uint32Array(queriesPerBatch * k); + + const compute = new Compute(device, shader, 'compute-knn-kdtree'); + compute.setParameter('positions', positionsBuf); + compute.setParameter('nodeSplatIdx', nSplatIdxBuf); + compute.setParameter('nodePositions', nPositionsBuf); + compute.setParameter('nodeChildren', nChildrenBuf); + compute.setParameter('outIndices', outBuf); + + this.execute = async ( + tree: FlatKdTree, + positions: Float32Array, + n: number, + queryCount: number, + outNeighbours: Uint32Array + ) => { + if (n > maxN) { + throw new Error(`GpuKnn: N=${n} exceeds maxN=${maxN}`); + } + if (positions.length < n * 3) { + throw new Error(`GpuKnn: positions length ${positions.length} must be at least N*3 = ${n * 3}`); + } + if (queryCount > n) { + throw new Error(`GpuKnn: queryCount=${queryCount} exceeds N=${n}`); + } + if (outNeighbours.length !== queryCount * k) { + throw new Error(`GpuKnn: outNeighbours length ${outNeighbours.length} must be queryCount*k = ${queryCount * k}`); + } + + // `FlatKdTree` already carries the interleaved layout this kernel + // wants (it is what this class used to pack for itself), so the + // uploads are direct. Bytes are unchanged from the packing loops + // this replaced — `buildFlatKdTree` is verified structurally + // identical to the pre-3.2 `KdTree.flatten()`. + positionsBuf.write(0, positions, 0, n * 3); + nSplatIdxBuf.write(0, tree.nodeSplatIdx, 0, n); + nPositionsBuf.write(0, tree.nodePositions, 0, n * 3); + nChildrenBuf.write(0, tree.nodeChildren, 0, n * 2); + compute.setParameter('rootIdx', tree.rootIdx); + + const numBatches = Math.ceil(queryCount / queriesPerBatch); + for (let batch = 0; batch < numBatches; batch++) { + const queryOffset = batch * queriesPerBatch; + const batchCount = Math.min(queriesPerBatch, queryCount - queryOffset); + const groups = Math.ceil(batchCount / workgroupSize); + + compute.setParameter('queryOffset', queryOffset); + compute.setParameter('queryCount', batchCount); + + compute.setupDispatch(groups); + device.computeDispatch([compute], `knn-dispatch-${batch}`); + + const readBytes = batchCount * k * 4; + await outBuf.read(0, readBytes, outScratch, true); + outNeighbours.set(outScratch.subarray(0, batchCount * k), queryOffset * k); + } + }; + + this.destroy = () => { + positionsBuf.destroy(); + nSplatIdxBuf.destroy(); + nPositionsBuf.destroy(); + nChildrenBuf.destroy(); + outBuf.destroy(); + shader.destroy(); + bindGroupFormat.destroy(); + }; + } +} + +export { GpuKnn }; diff --git a/src/lib/decimate-uniform/index.ts b/src/lib/decimate-uniform/index.ts new file mode 100644 index 00000000..017a2a64 --- /dev/null +++ b/src/lib/decimate-uniform/index.ts @@ -0,0 +1,8 @@ +// The pre-3.2 decimator — see README.md. Its `decimateSource` is exported +// under the `Uniform` name so the two paths can coexist without renaming +// anything inside this directory. +export { + decimateSource as decimateSourceUniform, + type DecimateOptions as DecimateUniformOptions, + type DecimateSpill as DecimateUniformSpill +} from './decimate-source'; diff --git a/src/lib/decimate-uniform/knn-blocks.ts b/src/lib/decimate-uniform/knn-blocks.ts new file mode 100644 index 00000000..c8874333 --- /dev/null +++ b/src/lib/decimate-uniform/knn-blocks.ts @@ -0,0 +1,332 @@ +import { knnQueryBlock, KNN_SENTINEL } from './knn-core'; +import { type BlockRange, type ResidentPositions } from './partition'; + +/** + * A block's local point set: owned gaussians first (in the block's sorted + * owned order), then halo members from neighbouring blocks. `ids` maps local + * index → global gaussian index; `positions` is interleaved xyz; `h` is the + * halo radius the set was collected with — `-Infinity` when no covering halo + * fits the cap (halo empty, every owned query takes the exact requery). + */ +type BlockLocals = { + ids: Uint32Array; + ownedCount: number; + positions: Float32Array; + h: number; +}; + +/** Local-slot marker: neighbour was fixed by verification; resolve via its global id. */ +const KNN_FIXED = 0xFFFFFFFE; + +// Density-based halo radius: haloFactor × the Poisson estimate of the k-NN +// radius from the block's AABB volume and count. Degenerate blocks (planar / +// coincident points) drive the estimate toward 0 — that's fine, the +// verification pass is the correctness backstop; h is only an efficiency hint. +const haloRadius = (block: BlockRange, k: number, haloFactor: number): number => { + const nOwned = block.end - block.start; + if (nOwned === 0) return 0; + const ex = Math.max(block.aabb[3] - block.aabb[0], 1e-12); + const ey = Math.max(block.aabb[4] - block.aabb[1], 1e-12); + const ez = Math.max(block.aabb[5] - block.aabb[2], 1e-12); + const lambda = nOwned / (ex * ey * ez); + const rk = Math.cbrt((k * 3) / (4 * Math.PI * lambda)); + return haloFactor * rk; +}; + +// Squared distance from a point to an AABB ([minx..z, maxx..z]). +const pointAabbDist2 = (px: number, py: number, pz: number, aabb: Float32Array): number => { + const dx = Math.max(0, aabb[0] - px, px - aabb[3]); + const dy = Math.max(0, aabb[1] - py, py - aabb[4]); + const dz = Math.max(0, aabb[2] - pz, pz - aabb[5]); + return dx * dx + dy * dy + dz * dz; +}; + +// Squared distance between two AABBs (0 when overlapping). +const aabbAabbDist2 = (a: Float32Array, b: Float32Array): number => { + let d2 = 0; + for (let c = 0; c < 3; c++) { + const gap = Math.max(0, b[c] - a[3 + c], a[c] - b[3 + c]); + d2 += gap * gap; + } + return d2; +}; + +/** + * Collect a block's local point set: its owned gaussians plus a halo of + * points from neighbouring blocks within `h` of the block AABB. + * + * @param pos - Resident positions. + * @param order - Partition index array. + * @param blocks - All block ranges. + * @param blockIdx - Which block to collect. + * @param k - Neighbours per query (drives the halo radius estimate). + * @param haloFactor - Multiplier on the k-NN radius estimate. + * @param haloCap - Maximum halo size as a multiple of the owned count (buffer-sizing bound; `h` shrinks until the full halo fits the cap — members are never dropped). + * @returns The block's locals. + */ +const collectBlock = ( + pos: ResidentPositions, + order: Uint32Array, + blocks: BlockRange[], + blockIdx: number, + k: number, + haloFactor: number, + haloCap = 1 +): BlockLocals => { + const block = blocks[blockIdx]; + const nOwned = block.end - block.start; + const maxHalo = Math.ceil(nOwned * haloCap); + + // The shrink loop only asks whether the halo exceeds the cap — bail as + // soon as that is known rather than counting the whole scene. + const countHalo = (h2: number): number => { + let count = 0; + for (let b2 = 0; b2 < blocks.length; b2++) { + if (b2 === blockIdx) continue; + const other = blocks[b2]; + if (aabbAabbDist2(block.aabb, other.aabb) > h2) continue; + for (let i = other.start; i < other.end; i++) { + const g = order[i]; + if (pointAabbDist2(pos.x[g], pos.y[g], pos.z[g], block.aabb) <= h2) { + if (++count > maxHalo) return count; + } + } + } + return count; + }; + + // The verification rule (`d_k ≤ depth + h`) is only sound when the halo + // FULLY covers AABB ⊕ h — so the size cap must never truncate members. + // Instead, shrink h until the full halo fits the cap; points beyond the + // reduced h are then legitimately outside the covered region and boundary + // queries fall through to the brute-force requery. + let h = haloRadius(block, k, haloFactor); + let fits = false; + for (let iter = 0; iter < 40 && h > 0; iter++) { + if (countHalo(h * h) <= maxHalo) { + fits = true; + break; + } + h *= 0.7; + if (h < 1e-12) h = 0; + } + if (!fits) { + if (countHalo(0) > maxHalo) { + // No h can fit the cap: other blocks' points sit INSIDE this + // block's AABB (e.g. an outlier-residual block enveloping the + // core), so the covering guarantee is unobtainable at any radius. + // Collect no halo and disable the guarantee (h = -Infinity): + // every owned query then takes the best-first requery, which is + // exact by construction. + const ids = new Uint32Array(nOwned); + const positions = new Float32Array(nOwned * 3); + for (let i = 0; i < nOwned; i++) { + const g = order[block.start + i]; + ids[i] = g; + positions[i * 3] = pos.x[g]; + positions[i * 3 + 1] = pos.y[g]; + positions[i * 3 + 2] = pos.z[g]; + } + return { ids, ownedCount: nOwned, positions, h: -Infinity }; + } + // The shrink iterations ran out, but a zero-radius halo fits. + h = 0; + } + const h2 = h * h; + + // Halo membership: points of other blocks within h of this block's AABB. + const haloIds: number[] = []; + for (let b2 = 0; b2 < blocks.length; b2++) { + if (b2 === blockIdx) continue; + const other = blocks[b2]; + if (aabbAabbDist2(block.aabb, other.aabb) > h2) continue; + for (let i = other.start; i < other.end; i++) { + const g = order[i]; + if (pointAabbDist2(pos.x[g], pos.y[g], pos.z[g], block.aabb) <= h2) { + haloIds.push(g); + } + } + } + + const n = nOwned + haloIds.length; + const ids = new Uint32Array(n); + const positions = new Float32Array(n * 3); + for (let i = 0; i < nOwned; i++) { + const g = order[block.start + i]; + ids[i] = g; + positions[i * 3] = pos.x[g]; + positions[i * 3 + 1] = pos.y[g]; + positions[i * 3 + 2] = pos.z[g]; + } + for (let i = 0; i < haloIds.length; i++) { + const g = haloIds[i]; + const l = nOwned + i; + ids[l] = g; + positions[l * 3] = pos.x[g]; + positions[l * 3 + 1] = pos.y[g]; + positions[l * 3 + 2] = pos.z[g]; + } + return { ids, ownedCount: nOwned, positions, h }; +}; + +/** + * CPU block KNN: exact k-NN of the owned points within block ∪ halo, as + * LOCAL indices (see {@link knnQueryBlock}). + * @param locals - The block's local point set. + * @param k - Neighbours per query. + * @returns Local neighbour indices, `ownedCount * k` long. + */ +const knnBlockCpu = (locals: BlockLocals, k: number): Uint32Array => { + return knnQueryBlock(locals.positions, locals.ownedCount, k); +}; + +/** + * Map local neighbour indices to global gaussian indices (sentinels pass + * through). + * @param locals - The block's local point set. + * @param nbLocal - Local neighbour indices. + * @returns A new array of global neighbour indices. + */ +const toGlobalNeighbors = (locals: BlockLocals, nbLocal: Uint32Array): Uint32Array => { + const out = new Uint32Array(nbLocal.length); + for (let s = 0; s < nbLocal.length; s++) { + out[s] = nbLocal[s] === KNN_SENTINEL ? KNN_SENTINEL : locals.ids[nbLocal[s]]; + } + return out; +}; + +// Insert (g, d2) into the k-best arrays (ascending by distance). +const kBestInsert = (bestIds: Uint32Array, bestD2: Float64Array, size: number, k: number, g: number, d2: number): number => { + if (size === k && d2 >= bestD2[k - 1]) return size; + let at = size < k ? size : k - 1; + while (at > 0 && bestD2[at - 1] > d2) { + bestD2[at] = bestD2[at - 1]; + bestIds[at] = bestIds[at - 1]; + at--; + } + bestD2[at] = d2; + bestIds[at] = g; + return Math.min(size + 1, k); +}; + +/** + * Exactness backstop for block KNN. A query's result is guaranteed correct + * when its k-th neighbour distance fits inside the halo-covered region + * (`d_k ≤ depth(q) + h`). Queries that fail — or that carry sentinel slots + * despite the scene having ≥ k other points — are re-queried best-first: + * blocks are visited in ascending point-to-AABB distance and the scan stops + * at the first block that can no longer improve on the k-th best so far. + * Same candidate set as a full scan, making the block KNN globally exact. + * + * Fixed entries are written into `nbGlobal`; the matching `nbLocal` slots + * (when provided) are marked {@link KNN_FIXED} so callers resolve those + * neighbours by global id. + * + * @param pos - Resident positions. + * @param order - Partition index array. + * @param blocks - All block ranges. + * @param blockIdx - Which block was queried. + * @param locals - The block's local point set. + * @param k - Neighbours per query. + * @param nbGlobal - Global neighbour indices (fixed in place). + * @param nbLocal - Optional parallel local indices to mark. + * @returns The number of re-queried gaussians. + */ +const verifyAndFixKnn = ( + pos: ResidentPositions, + order: Uint32Array, + blocks: BlockRange[], + blockIdx: number, + locals: BlockLocals, + k: number, + nbGlobal: Uint32Array, + nbLocal?: Uint32Array +): number => { + const block = blocks[blockIdx]; + const N = pos.x.length; + const { h } = locals; + let fixed = 0; + + const bestIds = new Uint32Array(k); + const bestD2 = new Float64Array(k); + const blockD2 = new Float64Array(blocks.length); + const blockOrd = new Uint32Array(blocks.length); + + for (let qi = 0; qi < locals.ownedCount; qi++) { + const g = locals.ids[qi]; + const qx = pos.x[g], qy = pos.y[g], qz = pos.z[g]; + + let dkSq = 0; + let sentinels = false; + for (let s = 0; s < k; s++) { + const nb = nbGlobal[qi * k + s]; + if (nb === KNN_SENTINEL) { + sentinels = true; + continue; + } + const dx = pos.x[nb] - qx, dy = pos.y[nb] - qy, dz = pos.z[nb] - qz; + const d2 = dx * dx + dy * dy + dz * dz; + if (d2 > dkSq) dkSq = d2; + } + + // Depth of q inside the block AABB (0 at/outside the boundary). + const depth = Math.max(0, Math.min( + qx - block.aabb[0], block.aabb[3] - qx, + qy - block.aabb[1], block.aabb[4] - qy, + qz - block.aabb[2], block.aabb[5] - qz + )); + + const needFix = (sentinels && N - 1 >= k) || Math.sqrt(dkSq) > depth + h; + if (!needFix) continue; + + // Best-first re-query: blocks in ascending point-to-AABB distance + // (insertion sort — block counts are small), stopping at the first + // block beyond the unverified k-th distance (a valid upper bound on + // the true one; unknown when sentinels) or beyond the k-th best so + // far, which tightens as candidates land. Blocks the loop never + // reaches provably contain no improving candidate. + const r2 = sentinels ? Infinity : dkSq; + for (let b2 = 0; b2 < blocks.length; b2++) { + const d2b = pointAabbDist2(qx, qy, qz, blocks[b2].aabb); + blockD2[b2] = d2b; + let at = b2; + while (at > 0 && blockD2[blockOrd[at - 1]] > d2b) { + blockOrd[at] = blockOrd[at - 1]; + at--; + } + blockOrd[at] = b2; + } + let size = 0; + for (let t = 0; t < blocks.length; t++) { + const b2 = blockOrd[t]; + const d2b = blockD2[b2]; + if (d2b > r2 || (size === k && d2b >= bestD2[k - 1])) break; + const other = blocks[b2]; + for (let i = other.start; i < other.end; i++) { + const cand = order[i]; + if (cand === g) continue; + const dx = pos.x[cand] - qx, dy = pos.y[cand] - qy, dz = pos.z[cand] - qz; + const d2 = dx * dx + dy * dy + dz * dz; + if (size < k || d2 < bestD2[size - 1]) { + size = kBestInsert(bestIds, bestD2, size, k, cand, d2); + } + } + } + for (let s = 0; s < k; s++) { + nbGlobal[qi * k + s] = s < size ? bestIds[s] : KNN_SENTINEL; + if (nbLocal) nbLocal[qi * k + s] = s < size ? KNN_FIXED : KNN_SENTINEL; + } + fixed++; + } + return fixed; +}; + +export { + collectBlock, + knnBlockCpu, + toGlobalNeighbors, + verifyAndFixKnn, + haloRadius, + KNN_FIXED, + type BlockLocals +}; diff --git a/src/lib/decimate-uniform/knn-core.ts b/src/lib/decimate-uniform/knn-core.ts new file mode 100644 index 00000000..7afb94bf --- /dev/null +++ b/src/lib/decimate-uniform/knn-core.ts @@ -0,0 +1,51 @@ +import { KdTree } from '../spatial/kd-tree'; + +/** Marks an unfilled neighbour slot (fewer than k non-self points available). */ +const KNN_SENTINEL = 0xFFFFFFFF; + +/** + * Exact k-nearest-neighbours for the owned prefix of a local point set. + * + * Engine-free (imported by worker tasks). Builds a {@link KdTree} over + * all `n` local points (owned first, then halo) and queries the first + * `ownedCount`. Output `out[q * k + s]` is a LOCAL index into `positions`, + * sorted ascending by distance, excluding the query itself, with + * {@link KNN_SENTINEL} filling surplus slots — the same contract as the + * legacy CPU KNN loop. + * + * @param positions - Interleaved xyz for all local points (owned + halo). + * @param ownedCount - Number of owned points at the front; only these are queried. + * @param k - Neighbours per query. + * @returns Local neighbour indices, `ownedCount * k` long. + */ +const knnQueryBlock = (positions: Float32Array, ownedCount: number, k: number): Uint32Array => { + const n = positions.length / 3; + const x = new Float32Array(n); + const y = new Float32Array(n); + const z = new Float32Array(n); + for (let i = 0; i < n; i++) { + x[i] = positions[i * 3]; + y[i] = positions[i * 3 + 1]; + z[i] = positions[i * 3 + 2]; + } + const tree = new KdTree([x, y, z]); + const out = new Uint32Array(ownedCount * k).fill(KNN_SENTINEL); + const q = new Float32Array(3); + for (let i = 0; i < ownedCount; i++) { + q[0] = x[i]; + q[1] = y[i]; + q[2] = z[i]; + // Request k+1 because the tree returns the query itself (distance 0). + const res = tree.findKNearest(q, k + 1); + let outPos = 0; + for (let m = 0; m < res.indices.length && outPos < k; m++) { + const j = res.indices[m]; + if (j === i) continue; + out[i * k + outPos] = j; + outPos++; + } + } + return out; +}; + +export { knnQueryBlock, KNN_SENTINEL }; diff --git a/src/lib/decimate-uniform/merge-stream.ts b/src/lib/decimate-uniform/merge-stream.ts new file mode 100644 index 00000000..12543c40 --- /dev/null +++ b/src/lib/decimate-uniform/merge-stream.ts @@ -0,0 +1,178 @@ +import { type ChunkPayload } from './block-producer'; +import { type ResidentPositions } from './partition'; +import { gatherBlockView, indexOfSorted, type PriorityContext } from './priority'; +import { type SelectionResult } from './select'; +import { WorkerQueue } from '../workers'; + +/** Context for the merge stream: the priority context plus the selection. */ +type MergeStreamContext = Pick & { + selection: SelectionResult; + /** When provided (sized to the output count), filled with output positions in emission order — the next generation's resident positions. */ + nextPositions?: ResidentPositions; +}; + +/** + * The merge stream (heavy read 2): walk blocks in partition order, gather + * geometric/color/(other) for owned rows + out-of-block group members, + * moment-match groups in workers, pass survivors through, and emit output + * rows in block order as chunk payloads of `chunkSize` rows (last partial). + * + * A group is emitted exactly once, at its minimum member's position; other + * members are consumed silently. Positions are never gathered — survivor + * positions and merged means come from the resident arrays / the merge. + * + * @param ctx - The stream context. + * @param chunkSize - Output rows per payload. + * @param tick - Optional progress callback (owned gaussians processed). + * @yields One {@link ChunkPayload} per output chunk, in order. + */ +async function *mergeStream( + ctx: MergeStreamContext, + chunkSize: number, + tick?: (n: number) => void +): AsyncGenerator { + const { source, pos, order, blocks, selection, nextPositions } = ctx; + const { memberGroup, groupMin, groupOffsets, groupMembers } = selection; + const { layouts, availableLayers } = source.meta; + + const colorDim = layouts.color!.stride >> 2; + const hasOther = availableLayers.has('other') && (layouts.other?.stride ?? 0) > 0; + const otherDim = hasOther ? layouts.other!.stride >> 2 : 0; + + // Rolling output buffers (reused across payloads: the consumer copies + // before pulling the next chunk). + const outPos = new Float32Array(chunkSize * 3); + const outGeo = new Float32Array(chunkSize * 8); + const outColor = new Float32Array(chunkSize * colorDim); + const outOther = hasOther ? new Uint32Array(chunkSize * otherDim) : undefined; + let rows = 0; + let emitted = 0; + + const payload = (): ChunkPayload => { + const p: ChunkPayload = { + count: rows, + position: outPos.subarray(0, rows * 3), + geometric: outGeo.subarray(0, rows * 8), + color: outColor.subarray(0, rows * colorDim) + }; + if (outOther) p.other = outOther.subarray(0, rows * otherDim); + return p; + }; + + for (let bi = 0; bi < blocks.length; bi++) { + const block = blocks[bi]; + const owned = order.subarray(block.start, block.end); + const nOwned = owned.length; + + // This block's emitted groups (min member owned here), in owned order, + // and the out-of-block members they pull in. + const blockGroups: number[] = []; + const extSet = new Map(); + for (let i = 0; i < nOwned; i++) { + const g = owned[i]; + const mg = memberGroup[g]; + if (mg === -1 || groupMin[mg] !== g) continue; + blockGroups.push(mg); + for (let m = groupOffsets[mg]; m < groupOffsets[mg + 1]; m++) { + const member = groupMembers[m]; + if (indexOfSorted(owned, member) < 0 && !extSet.has(member)) extSet.set(member, 0); + } + } + const extraGlobals = Uint32Array.from(extSet.keys()).sort(); + for (let i = 0; i < extraGlobals.length; i++) extSet.set(extraGlobals[i], nOwned + i); + + const { view, other } = await gatherBlockView(ctx, bi, extraGlobals, hasOther); + + // Merge this block's groups in a worker: pack member-major inputs. + let mergedPos: Float32Array | null = null; + let mergedGeo: Float32Array | null = null; + let mergedColor: Float32Array | null = null; + let mergedOther: Uint32Array | undefined; + if (blockGroups.length > 0) { + let totalMembers = 0; + for (const mg of blockGroups) totalMembers += groupOffsets[mg + 1] - groupOffsets[mg]; + const mPos = new Float32Array(totalMembers * 3); + const mGeo = new Float32Array(totalMembers * 8); + const mColor = new Float32Array(totalMembers * colorDim); + const mOther = hasOther ? new Uint32Array(totalMembers * otherDim) : undefined; + const sizes = new Uint32Array(blockGroups.length); + let mi = 0; + for (let gi = 0; gi < blockGroups.length; gi++) { + const mg = blockGroups[gi]; + sizes[gi] = groupOffsets[mg + 1] - groupOffsets[mg]; + for (let m = groupOffsets[mg]; m < groupOffsets[mg + 1]; m++) { + const member = groupMembers[m]; + const oi = indexOfSorted(owned, member); + const row = oi >= 0 ? oi : extSet.get(member)!; + mPos[mi * 3] = view.pos[row * 3]; + mPos[mi * 3 + 1] = view.pos[row * 3 + 1]; + mPos[mi * 3 + 2] = view.pos[row * 3 + 2]; + mGeo.set(view.geo.subarray(row * 8, row * 8 + 8), mi * 8); + mColor.set(view.color.subarray(row * colorDim, (row + 1) * colorDim), mi * colorDim); + if (mOther) mOther.set(other!.subarray(row * otherDim, (row + 1) * otherDim), mi * otherDim); + mi++; + } + } + const transfer: ArrayBuffer[] = [mPos.buffer as ArrayBuffer, mGeo.buffer as ArrayBuffer, mColor.buffer as ArrayBuffer]; + if (mOther) transfer.push(mOther.buffer as ArrayBuffer); + const merged = await WorkerQueue.run('mergeGroups', { + pos: mPos, + geo: mGeo, + color: mColor, + sizes, + colorDim, + other: mOther, + otherDim + }, transfer); + mergedPos = merged.pos; + mergedGeo = merged.geo; + mergedColor = merged.color; + mergedOther = merged.other; + } + + // Emit rows in owned order. + let nextMerged = 0; + for (let i = 0; i < nOwned; i++) { + const g = owned[i]; + const mg = memberGroup[g]; + if (mg !== -1 && groupMin[mg] !== g) continue; // consumed member + + if (mg === -1) { + // Survivor pass-through: position from resident arrays, + // geometric/color/other block-copied from the view. + outPos[rows * 3] = pos.x[g]; + outPos[rows * 3 + 1] = pos.y[g]; + outPos[rows * 3 + 2] = pos.z[g]; + outGeo.set(view.geo.subarray(i * 8, i * 8 + 8), rows * 8); + outColor.set(view.color.subarray(i * colorDim, (i + 1) * colorDim), rows * colorDim); + if (outOther) outOther.set(other!.subarray(i * otherDim, (i + 1) * otherDim), rows * otherDim); + } else { + const mi = nextMerged++; + outPos.set(mergedPos!.subarray(mi * 3, mi * 3 + 3), rows * 3); + outGeo.set(mergedGeo!.subarray(mi * 8, mi * 8 + 8), rows * 8); + outColor.set(mergedColor!.subarray(mi * colorDim, (mi + 1) * colorDim), rows * colorDim); + if (outOther) outOther.set(mergedOther!.subarray(mi * otherDim, (mi + 1) * otherDim), rows * otherDim); + } + + if (nextPositions) { + nextPositions.x[emitted] = outPos[rows * 3]; + nextPositions.y[emitted] = outPos[rows * 3 + 1]; + nextPositions.z[emitted] = outPos[rows * 3 + 2]; + } + rows++; + emitted++; + if (rows === chunkSize) { + yield payload(); + rows = 0; + } + } + tick?.(nOwned); + } + + if (rows > 0) { + yield payload(); + rows = 0; + } +} + +export { mergeStream, type MergeStreamContext }; diff --git a/src/lib/decimate-uniform/partition.ts b/src/lib/decimate-uniform/partition.ts new file mode 100644 index 00000000..07dfe5f3 --- /dev/null +++ b/src/lib/decimate-uniform/partition.ts @@ -0,0 +1,163 @@ +import { quickselect } from '../utils'; + +/** + * Resident per-gaussian position columns — the only whole-scene data + * decimation keeps in memory (12 B/gaussian). + */ +type ResidentPositions = { + x: Float32Array; + y: Float32Array; + z: Float32Array; +}; + +/** + * One spatial block of the KD partition: gaussians `order[start..end)`, + * sorted ascending (for gather coalescing), with the block's position AABB + * as `[minx, miny, minz, maxx, maxy, maxz]`. + */ +type BlockRange = { + start: number; + end: number; + aabb: Float32Array; +}; + +/** Outlier fence: expand the sampled per-axis quantile interval this much. */ +const OUTLIER_FENCE_FACTOR = 4; + +/** Treat out-of-fence points as flyaways only while they are rare. */ +const OUTLIER_MAX_FRACTION = 0.01; + +/** Position sample cap for the fence quantiles. */ +const OUTLIER_SAMPLE_CAP = 1 << 20; + +// Per-axis fence [lo, hi] from strided-sample quantiles: mid ± factor × the +// 0.1–99.9% half-spread. An axis with no spread stays unfenced (±Infinity). +const outlierFence = (pos: ResidentPositions): { lo: number[]; hi: number[] } => { + const n = pos.x.length; + const stride = Math.max(1, Math.ceil(n / OUTLIER_SAMPLE_CAP)); + const cols = [pos.x, pos.y, pos.z]; + const lo = [-Infinity, -Infinity, -Infinity]; + const hi = [Infinity, Infinity, Infinity]; + const samp = new Float32Array(Math.ceil(n / stride) || 1); + for (let c = 0; c < 3; c++) { + let m = 0; + for (let i = 0; i < n; i += stride) samp[m++] = cols[c][i]; + const s = samp.subarray(0, m).sort(); + const qlo = s[Math.min(m - 1, Math.floor(0.001 * m))]; + const qhi = s[Math.min(m - 1, Math.floor(0.999 * m))]; + const half = (qhi - qlo) / 2; + if (!(half > 0)) continue; + const mid = (qlo + qhi) / 2; + lo[c] = mid - OUTLIER_FENCE_FACTOR * half; + hi[c] = mid + OUTLIER_FENCE_FACTOR * half; + } + return { lo, hi }; +}; + +/** + * KD-partition the resident positions into spatial blocks of at most + * `blockSize` gaussians by recursive median splits on the largest AABB axis + * (quickselect, in place on one index array). Rare flyaway positions are set + * aside into trailing residual block(s) first, so core blocks keep tight + * AABBs — flyaways otherwise stretch AABBs scene-wide, which wrecks the + * density-based halo estimate and AABB-distance pruning downstream. With + * globally exact KNN, block boundaries cannot change which merges are + * possible or their costs — but they do set output row order, and selection + * tie-breaks between quantized-equal costs can resolve differently under a + * different partition. + * + * @param pos - Resident positions. + * @param blockSize - Maximum gaussians per block. + * @returns The permuted index array and the block ranges over it. + */ +const kdPartition = (pos: ResidentPositions, blockSize: number): { order: Uint32Array; blocks: BlockRange[] } => { + const n = pos.x.length; + const order = new Uint32Array(n); + for (let i = 0; i < n; i++) order[i] = i; + const blocks: BlockRange[] = []; + const cols = [pos.x, pos.y, pos.z]; + + const aabbOf = (start: number, end: number): Float32Array => { + const a = new Float32Array([Infinity, Infinity, Infinity, -Infinity, -Infinity, -Infinity]); + for (let i = start; i < end; i++) { + const g = order[i]; + for (let c = 0; c < 3; c++) { + const v = cols[c][g]; + if (v < a[c]) a[c] = v; + if (v > a[3 + c]) a[3 + c] = v; + } + } + return a; + }; + + const recurse = (start: number, end: number): void => { + const aabb = aabbOf(start, end); + if (end - start <= blockSize) { + order.subarray(start, end).sort(); + blocks.push({ start, end, aabb }); + return; + } + let axis = 0, ext = -Infinity; + for (let c = 0; c < 3; c++) { + const e = aabb[3 + c] - aabb[c]; + if (e > ext) { + ext = e; + axis = c; + } + } + const mid = start + ((end - start) >> 1); + quickselect(cols[axis], order.subarray(start, end), mid - start); + recurse(start, mid); + recurse(mid, end); + }; + + // Residual split: fence classification must stay rare — a scene that is + // mostly "outliers" is just sparse, and splitting it would recreate the + // stretched-AABB problem inside the residual. + let coreEnd = n; + if (n > 0) { + const { lo, hi } = outlierFence(pos); + let out = 0; + for (let i = 0; i < n; i++) { + if (cols[0][i] < lo[0] || cols[0][i] > hi[0] || + cols[1][i] < lo[1] || cols[1][i] > hi[1] || + cols[2][i] < lo[2] || cols[2][i] > hi[2]) out++; + } + if (out > 0 && out <= n * OUTLIER_MAX_FRACTION) { + coreEnd = n - out; + let c = 0, o = coreEnd; + for (let i = 0; i < n; i++) { + if (cols[0][i] < lo[0] || cols[0][i] > hi[0] || + cols[1][i] < lo[1] || cols[1][i] > hi[1] || + cols[2][i] < lo[2] || cols[2][i] > hi[2]) order[o++] = i; + else order[c++] = i; + } + } + } + if (coreEnd > 0) recurse(0, coreEnd); + if (coreEnd < n) recurse(coreEnd, n); + return { order, blocks }; +}; + +/** + * Count the coalesced runs a block's sorted source rows form under the + * reader's gap-merge threshold — the spatial-coherence signal. A coherent + * (Morton-ordered / block-ordered) file yields a handful of runs per block; + * a training-order file yields ~one run per row, which is the cue to + * recommend a one-time `--morton-order` prepass. + * + * @param sortedIndices - Row indices, ascending, typically `order`. + * @param start - Range start (inclusive). + * @param end - Range end (exclusive). + * @param mergeGapRows - Merge adjacent indices when the gap is at most this many rows. + * @returns The number of coalesced runs. + */ +const coherenceRuns = (sortedIndices: Uint32Array, start: number, end: number, mergeGapRows: number): number => { + let runs = end > start ? 1 : 0; + for (let i = start + 1; i < end; i++) { + if (sortedIndices[i] - sortedIndices[i - 1] > mergeGapRows) runs++; + } + return runs; +}; + +export { kdPartition, coherenceRuns, type BlockRange, type ResidentPositions }; diff --git a/src/lib/decimate-uniform/priority.ts b/src/lib/decimate-uniform/priority.ts new file mode 100644 index 00000000..f65e49f7 --- /dev/null +++ b/src/lib/decimate-uniform/priority.ts @@ -0,0 +1,387 @@ +import { type GraphicsDevice } from 'playcanvas'; + +import { buildCostCache, computeEdgeCostView } from './edge-cost-cpu'; +import { APP_CHUNK, GpuEdgeCost, type EdgeCostCache } from './gpu-edge-cost'; +import { GpuKnn } from './gpu-knn'; +import { collectBlock, verifyAndFixKnn, toGlobalNeighbors, KNN_FIXED, type BlockLocals } from './knn-blocks'; +import { KNN_SENTINEL } from './knn-core'; +import { type BlockRange, type ResidentPositions } from './partition'; +import { type ChunkData, type ChunkDataPool, type ChunkSource } from '../chunk'; +import { createMergeScratch, makeGaussianSamples, sigmoid, ellipsoidArea, type SplatView } from '../decimate/moment-match'; +import { WorkerQueue } from '../workers'; + +/** Halo radius multiplier on the density-estimated k-NN radius. */ +const HALO_FACTOR = 2.5; + +/** Halo size cap as a multiple of a block's owned count (buffer-sizing bound). */ +const HALO_CAP = 1; + +/** + * Per-gaussian best-K merge candidates, the resident output of the priority + * pass. `idx[g * K + s]` is the global index of gaussian g's s-th cheapest + * candidate (0xFFFFFFFF when absent); `cost[g * K + s]` its cost (+Inf when + * absent). + */ +type CandidateArrays = { + idx: Uint32Array; + cost: Float32Array; +}; + +/** Everything the block passes need: baked single-LOD source + resident state. */ +type PriorityContext = { + source: ChunkSource; + pool: ChunkDataPool; + pos: ResidentPositions; + order: Uint32Array; + blocks: BlockRange[]; + device?: GraphicsDevice; + /** Candidates kept per gaussian (K). */ + K: number; + /** Neighbours per query (16). */ + k: number; +}; + +/** + * A block's gathered splat columns: owned rows first (block order), then the + * requested extra globals. Positions come from the resident arrays, never + * from the source. + */ +type BlockView = { + view: SplatView; + /** u32 `other` columns (extraDim per row), when requested and present. */ + other?: Uint32Array; + otherDim: number; + ownedCount: number; +}; + +/** + * Gather geometric + color (and optionally `other`) for a block's owned rows + * plus `extraGlobals`, into tight column arrays. Reads are batched at the + * pool's chunk size; owned and extra index lists must be sorted ascending + * for gather coalescing. + * + * @param ctx - The pass context. + * @param blockIdx - Which block. + * @param extraGlobals - Sorted out-of-block rows to append after the owned rows. + * @param includeOther - Also gather the `other` layer (merge pass only). + * @returns The gathered block view. + */ +const gatherBlockView = async ( + ctx: Pick, + blockIdx: number, + extraGlobals: Uint32Array, + includeOther = false +): Promise => { + const { source, pool, pos, order, blocks } = ctx; + const block = blocks[blockIdx]; + const owned = order.subarray(block.start, block.end); + const nOwned = owned.length; + const n = nOwned + extraGlobals.length; + const { layouts, availableLayers } = source.meta; + + const colorDim = layouts.color!.stride >> 2; + const wantOther = includeOther && availableLayers.has('other') && (layouts.other?.stride ?? 0) > 0; + const otherDim = wantOther ? layouts.other!.stride >> 2 : 0; + + const view: SplatView = { + pos: new Float32Array(n * 3), + geo: new Float32Array(n * 8), + color: new Float32Array(n * colorDim), + colorDim + }; + const other = wantOther ? new Uint32Array(n * otherDim) : undefined; + + const readInto = async (indices: Uint32Array, rowBase: number): Promise => { + const batch = pool.chunkSize; + for (let off = 0; off < indices.length; off += batch) { + const count = Math.min(batch, indices.length - off); + const geoCd = pool.acquire('geometric', layouts.geometric!, count); + const colCd = pool.acquire('color', layouts.color!, count); + const othCd: ChunkData | undefined = wantOther ? pool.acquire('other', layouts.other!, count) : undefined; + await source.read({ + indices, + indexOffset: off, + count, + geometric: geoCd, + color: colCd, + other: othCd + }); + view.geo.set(new Float32Array(geoCd.data, 0, count * 8), (rowBase + off) * 8); + view.color.set(new Float32Array(colCd.data, 0, count * colorDim), (rowBase + off) * colorDim); + if (othCd) other!.set(new Uint32Array(othCd.data, 0, count * otherDim), (rowBase + off) * otherDim); + geoCd.release(); + colCd.release(); + othCd?.release(); + } + }; + + await readInto(owned, 0); + await readInto(extraGlobals, nOwned); + + for (let i = 0; i < nOwned; i++) { + const g = owned[i]; + view.pos[i * 3] = pos.x[g]; + view.pos[i * 3 + 1] = pos.y[g]; + view.pos[i * 3 + 2] = pos.z[g]; + } + for (let i = 0; i < extraGlobals.length; i++) { + const g = extraGlobals[i]; + const r = nOwned + i; + view.pos[r * 3] = pos.x[g]; + view.pos[r * 3 + 1] = pos.y[g]; + view.pos[r * 3 + 2] = pos.z[g]; + } + + return { view, other, otherDim, ownedCount: nOwned }; +}; + +// Binary search `g` in the sorted array; -1 when absent. +const indexOfSorted = (sorted: Uint32Array, g: number): number => { + let lo = 0, hi = sorted.length - 1; + while (lo <= hi) { + const mid = (lo + hi) >> 1; + const v = sorted[mid]; + if (v === g) return mid; + if (v < g) lo = mid + 1; + else hi = mid - 1; + } + return -1; +}; + +// Pack the block view into the GpuEdgeCost cache layout (legacy packing: +// posScalars 8-wide, rotR from normalized quats, appearance in ≤APP_CHUNK +// column chunks with live-width strides). +const packGpuCache = (view: SplatView): EdgeCostCache => { + const { pos, geo, color, colorDim } = view; + const n = geo.length / 8; + const posScalars = new Float32Array(n * 8); + const rotR = new Float32Array(n * 9); + const rot = new Float32Array(9); + + for (let i = 0; i < n; i++) { + const i8 = i * 8; + const o = i * 8; + const linAlpha = sigmoid(geo[i8 + 7]); + const sx = Math.max(Math.exp(geo[i8 + 4]), 1e-12); + const sy = Math.max(Math.exp(geo[i8 + 5]), 1e-12); + const sz = Math.max(Math.exp(geo[i8 + 6]), 1e-12); + const vx = sx * sx + 1e-8; + const vy = sy * sy + 1e-8; + const vz = sz * sz + 1e-8; + posScalars[o] = pos[i * 3]; + posScalars[o + 1] = pos[i * 3 + 1]; + posScalars[o + 2] = pos[i * 3 + 2]; + posScalars[o + 3] = linAlpha * ellipsoidArea(sx, sy, sz) + 1e-12; + posScalars[o + 4] = Math.log(Math.max(vx, 1e-30)) + Math.log(Math.max(vy, 1e-30)) + Math.log(Math.max(vz, 1e-30)); + posScalars[o + 5] = vx; + posScalars[o + 6] = vy; + posScalars[o + 7] = vz; + + let qw = geo[i8], qx = geo[i8 + 1], qy = geo[i8 + 2], qz = geo[i8 + 3]; + const invq = 1 / Math.max(Math.hypot(qw, qx, qy, qz), 1e-12); + qw *= invq; qx *= invq; qy *= invq; qz *= invq; + const xx = qx * qx, yy = qy * qy, zz = qz * qz; + const wx = qw * qx, wy = qw * qy, wz = qw * qz; + const xy = qx * qy, xz = qx * qz, yz = qy * qz; + rot[0] = 1 - 2 * (yy + zz); rot[1] = 2 * (xy - wz); rot[2] = 2 * (xz + wy); + rot[3] = 2 * (xy + wz); rot[4] = 1 - 2 * (xx + zz); rot[5] = 2 * (yz - wx); + rot[6] = 2 * (xz - wy); rot[7] = 2 * (yz + wx); rot[8] = 1 - 2 * (xx + yy); + rotR.set(rot, i * 9); + } + + const numChunks = Math.ceil(colorDim / APP_CHUNK); + const appChunks: Float32Array[] = []; + for (let ch = 0; ch < numChunks; ch++) { + const kStart = ch * APP_CHUNK; + const width = Math.min(APP_CHUNK, colorDim - kStart); + const chunk = new Float32Array(n * width); + for (let s = 0; s < n; s++) { + const dst = s * width; + const src = s * colorDim + kStart; + for (let kk = 0; kk < width; kk++) chunk[dst + kk] = color[src + kk]; + } + appChunks.push(chunk); + } + + return { posScalars, rotR, appChunks, numAppCols: colorDim, numSplats: n }; +}; + +/** + * The priority pass (heavy read 1): per block — exact global KNN, edge costs + * for each owned gaussian's k neighbours, reduction to the best K candidates + * — written into the resident candidate arrays. + * + * @param ctx - The pass context. + * @param cand - Preallocated candidate arrays (`N*K`), filled per block. + * @param tick - Optional progress callback (owned gaussians completed). + */ +const runPriorityPass = async ( + ctx: PriorityContext, + cand: CandidateArrays, + tick?: (n: number) => void +): Promise => { + const { pos, order, blocks, device, K, k } = ctx; + const Z = makeGaussianSamples(1, 0); + const z = new Float32Array([Z[0][0], Z[0][1], Z[0][2]]); + const colorDim = ctx.source.meta.layouts.color!.stride >> 2; + + let maxOwned = 0; + for (const b of blocks) maxOwned = Math.max(maxOwned, b.end - b.start); + const maxLocalN = maxOwned * (1 + HALO_CAP); + + let gpuKnn: GpuKnn | undefined; + let gpuCost: GpuEdgeCost | undefined; + let gpuCostCapacity = maxLocalN; + + // 1-deep prefetch: the next block's halo collection + tree build runs + // while the current block computes. GpuKnn executions share one set of + // buffers, so they are serialized through `gpuKnnQueue` — the prefetched + // block's KNN starts only after the current block's has finished. + let gpuKnnQueue: Promise = Promise.resolve(); + type Prepared = { locals: BlockLocals; nb: Promise }; + const prepare = (bi: number): Prepared => { + const locals = collectBlock(pos, order, blocks, bi, k, HALO_FACTOR, HALO_CAP); + const copy = locals.positions.slice(); + if (device) { + const treePromise = WorkerQueue.run('flattenKdTree', { positions: copy }, [copy.buffer as ArrayBuffer]); + const out = new Uint32Array(locals.ownedCount * k); + const run = Promise.all([treePromise, gpuKnnQueue]).then(([flat]) => { + return gpuKnn!.execute(flat, locals.positions, locals.ids.length, locals.ownedCount, out); + }); + gpuKnnQueue = run.catch(() => { /* surfaced by the awaiting block */ }); + return { locals, nb: run.then(() => out) }; + } + const nb = WorkerQueue.run('knnBlock', { positions: copy, ownedCount: locals.ownedCount, k }, [copy.buffer as ArrayBuffer]); + return { locals, nb }; + }; + + try { + if (device) { + gpuKnn = new GpuKnn(device, maxLocalN, k); + gpuCost = new GpuEdgeCost(device, maxLocalN, maxOwned * k, colorDim); + } + + let next: Prepared | null = blocks.length > 0 ? prepare(0) : null; + + for (let bi = 0; bi < blocks.length; bi++) { + const { locals, nb: nbPromise } = next!; + next = bi + 1 < blocks.length ? prepare(bi + 1) : null; + + const nOwned = locals.ownedCount; + const owned = order.subarray(blocks[bi].start, blocks[bi].end); + const nbLocal = await nbPromise; + const nbGlobal = toGlobalNeighbors(locals, nbLocal); + verifyAndFixKnn(pos, order, blocks, bi, locals, k, nbGlobal, nbLocal); + + // Externals: referenced rows outside the owned range (halo members + // and verification-fixed neighbours), sorted for the gather. + const extRow = new Map(); + for (let s = 0; s < nOwned * k; s++) { + const l = nbLocal[s]; + if (l === KNN_SENTINEL || l < nOwned) continue; + const g = nbGlobal[s]; + if (l !== KNN_FIXED) { + if (!extRow.has(g)) extRow.set(g, 0); + } else if (indexOfSorted(owned, g) < 0 && !extRow.has(g)) { + extRow.set(g, 0); + } + } + const extraGlobals = Uint32Array.from(extRow.keys()).sort(); + for (let i = 0; i < extraGlobals.length; i++) extRow.set(extraGlobals[i], nOwned + i); + + // Verification-fixed externals are not bounded by the halo cap, so + // a pathological block's view can exceed the preallocated cost + // buffers — grow them to the actual view size when that happens + // (rare; costs one reallocation). + const viewN = nOwned + extraGlobals.length; + if (gpuCost && viewN > gpuCostCapacity) { + gpuCost.destroy(); + gpuCostCapacity = Math.ceil(viewN * 1.1); + gpuCost = new GpuEdgeCost(device!, gpuCostCapacity, maxOwned * k, colorDim); + } + + const { view } = await gatherBlockView(ctx, bi, extraGlobals); + + // Edge lists in owned-major order (view-local endpoints). + const edgeI = new Uint32Array(nOwned * k); + const edgeJ = new Uint32Array(nOwned * k); + const edgeNb = new Uint32Array(nOwned * k); // global neighbour per edge + const edgeOf = new Uint32Array(nOwned + 1); // CSR into the edge list per owned row + let e = 0; + for (let qi = 0; qi < nOwned; qi++) { + edgeOf[qi] = e; + for (let s = 0; s < k; s++) { + const l = nbLocal[qi * k + s]; + if (l === KNN_SENTINEL) continue; + const g = nbGlobal[qi * k + s]; + let row: number; + if (l !== KNN_FIXED) { + row = l < nOwned ? l : extRow.get(g)!; + } else { + const oi = indexOfSorted(owned, g); + row = oi >= 0 ? oi : extRow.get(g)!; + } + edgeI[e] = qi; + edgeJ[e] = row; + edgeNb[e] = g; + e++; + } + } + edgeOf[nOwned] = e; + + const costs = new Float32Array(e); + if (device) { + await gpuCost!.execute(packGpuCache(view), edgeI.subarray(0, e), edgeJ.subarray(0, e), z, costs); + } else { + const cache = buildCostCache(view); + const scratch = createMergeScratch(); + for (let i = 0; i < e; i++) { + costs[i] = computeEdgeCostView(view, cache, edgeI[i], edgeJ[i], Z, scratch); + } + } + + // Reduce to best K candidates per owned gaussian (ascending by cost). + const bestIdx = new Uint32Array(K); + const bestCost = new Float64Array(K); + for (let qi = 0; qi < nOwned; qi++) { + let size = 0; + for (let s = edgeOf[qi]; s < edgeOf[qi + 1]; s++) { + const c = costs[s]; + if (!Number.isFinite(c)) continue; + if (size === K && c >= bestCost[K - 1]) continue; + let at = size < K ? size : K - 1; + while (at > 0 && bestCost[at - 1] > c) { + bestCost[at] = bestCost[at - 1]; + bestIdx[at] = bestIdx[at - 1]; + at--; + } + bestCost[at] = c; + bestIdx[at] = edgeNb[s]; + size = Math.min(size + 1, K); + } + const g = owned[qi]; + for (let s = 0; s < K; s++) { + cand.idx[g * K + s] = s < size ? bestIdx[s] : 0xFFFFFFFF; + cand.cost[g * K + s] = s < size ? bestCost[s] : Infinity; + } + } + + tick?.(nOwned); + } + } finally { + gpuKnn?.destroy(); + gpuCost?.destroy(); + } +}; + +export { + runPriorityPass, + gatherBlockView, + packGpuCache, + indexOfSorted, + HALO_FACTOR, + HALO_CAP, + type CandidateArrays, + type PriorityContext, + type BlockView +}; diff --git a/src/lib/decimate/select-legacy.ts b/src/lib/decimate-uniform/select.ts similarity index 96% rename from src/lib/decimate/select-legacy.ts rename to src/lib/decimate-uniform/select.ts index 9d6c5318..f7cfc72e 100644 --- a/src/lib/decimate/select-legacy.ts +++ b/src/lib/decimate-uniform/select.ts @@ -46,7 +46,7 @@ type SelectionResult = { * @param mergesNeeded - Target removal count for this generation. * @returns The selection. */ -const selectMergesLegacy = (cand: CandidateArrays, N: number, K: number, mergesNeeded: number): SelectionResult => { +const selectMerges = (cand: CandidateArrays, N: number, K: number, mergesNeeded: number): SelectionResult => { const E = N * K; // Pass 1: finite cost range. @@ -151,4 +151,4 @@ const selectMergesLegacy = (cand: CandidateArrays, N: number, K: number, mergesN return { groupOffsets, groupMembers, memberGroup, groupMin, mergedGroups: G, removed }; }; -export { selectMergesLegacy, SELECT_BUCKETS, type SelectionResult }; +export { selectMerges, SELECT_BUCKETS, type SelectionResult }; diff --git a/src/lib/decimate/block-plan.ts b/src/lib/decimate/block-plan.ts index 21b7d32d..135a68cd 100644 --- a/src/lib/decimate/block-plan.ts +++ b/src/lib/decimate/block-plan.ts @@ -217,7 +217,7 @@ const planBlockMerges = async (inputs: BlockPlanInputs): Promise => { const gpuFits = !!device && D * MAX_GROUP === 64 && GpuRecost.fits(device, N, D, WAVE, coreCount, true); if (requireGpu && !gpuFits) { throw new Error( - 'multi-block quality decimation requires a WebGPU block working set that fits the adapter limits ' + + 'multi-block adaptive decimation requires a WebGPU block working set that fits the adapter limits ' + `(core ${coreCount}, halo ${N - coreCount})` ); } diff --git a/src/lib/decimate/decimate-source.ts b/src/lib/decimate/decimate-source.ts index 316212eb..f676397c 100644 --- a/src/lib/decimate/decimate-source.ts +++ b/src/lib/decimate/decimate-source.ts @@ -10,13 +10,10 @@ import { blockPlanMergeStream } from './block-merge-stream'; import { planBlockMerges } from './block-plan'; import { prepareGpuBlock, type PreparedBlock } from './block-prepare'; import { createBlockProducerSource, type DestBuffers } from './block-producer'; -import { buildCostCacheLegacy, computeEdgeCostViewLegacy, packGpuCacheLegacy } from './edge-cost-legacy'; import { mergeStream } from './merge-stream'; -import { createMergeScratch, makeGaussianSamples } from './moment-match'; import { buildBlockHalo, kdPartition, coherenceRuns, type ResidentPositions } from './partition'; -import { runPriorityPass, VIEW_GROW, type CandidateArrays, type CostStrategy } from './priority'; +import { runPriorityPass, type CandidateArrays } from './priority'; import { selectMerges, type SelectionResult } from './select'; -import { selectMergesLegacy } from './select-legacy'; import { selectMergesRecosted, CACHE_STRIDE } from './select-recost'; import { compact, @@ -24,8 +21,6 @@ import { type ChunkSource, type ChunkSourceMetadata } from '../chunk'; -import { SPLAT_STRIDE } from '../gpu/gpu-edge-cost'; -import { GpuEdgeCostLegacy } from '../gpu/gpu-edge-cost-legacy'; import { type ReadFileSystem } from '../io/read'; import { type FileSystem } from '../io/write'; import { bakeTransform, permuteSource } from '../ops'; @@ -84,15 +79,6 @@ type DecimateOptions = { spill?: DecimateSpill; /** Resident-memory budget driving the candidate-K and re-costed-selection policies (default 48 GiB). */ memoryBudgetBytes?: number; - /** - * Decimation algorithm. `'quality'` (default): field-L2 cost with the - * scale-free colour term and re-costed selection — the quality-study - * winner (large gains on mixed-scale scenes, e.g. skies). `'legacy'`: the - * pre-study pipeline (KL-style cost with full-SH colour term, uniform - * matching) — faster, lower memory, and still measurably better on scenes - * of uniformly-sized gaussians (single objects, uniform texture). - */ - mode?: 'quality' | 'legacy'; }; // Candidate-K policy: keep 4 when the resident estimate fits the budget, @@ -131,26 +117,6 @@ const chooseBlockSize = ( return Math.max(1, Math.min(blockSize, n)); }; -// The pre-study KL-style cost kernel (--decimate-balanced): full-SH colour -// L2, single-Monte-Carlo geometric term, its own GPU cache/kernel layouts. -const createLegacyStrategy = (colorDim: number): CostStrategy => { - const Z = makeGaussianSamples(1, 0); - const z = new Float32Array([Z[0][0], Z[0][1], Z[0][2]]); - const scratch = createMergeScratch(); - return { - cacheForGpu: false, - buildCache: view => buildCostCacheLegacy(view), - createGpu(device, capacity, k) { - const gpu = new GpuEdgeCostLegacy(device, capacity, k, colorDim); - return { - execute: (view, _cache, nbRows, outCosts) => gpu.execute(packGpuCacheLegacy(view), nbRows, z, outCosts), - destroy: () => gpu.destroy() - }; - }, - cpuEdge: (view, cache, i, j) => computeEdgeCostViewLegacy(view, cache as ReturnType, i, j, Z, scratch) - }; -}; - // Read the position layer sequentially into resident columns (generation 1 // only; later generations carry positions forward from the merge stream). const extractPositions = async (source: ChunkSource, pool: ChunkDataPool): Promise => { @@ -256,28 +222,16 @@ const decimateSource = async ( positions ??= await extractPositions(src, pool); - // Quality multi-block working sets are sized from the actual host - // budget and adapter binding limits. Legacy retains its previous - // fixed-size batching policy. - const legacy = opts.mode === 'legacy'; + // Multi-block working sets are sized from the actual host budget + // and adapter binding limits. const nextCount = Math.max(targetCount, N - Math.floor(N / 2)); - let blockSize = legacy ? - BLOCK_SIZE : - chooseBlockSize(N, budget, residentInputBytes, nextCount * 12, device); - if (legacy) { - const bindingLimit = (device as unknown as { limits?: { maxStorageBufferBindingSize?: number } } | undefined) - ?.limits?.maxStorageBufferBindingSize; - if (typeof bindingLimit === 'number') { - const largestBinding = (bs: number) => Math.ceil(bs * VIEW_GROW) * SPLAT_STRIDE * 4; - while (blockSize > (1 << 16) && largestBinding(blockSize) > bindingLimit) blockSize >>= 1; - } - } + const blockSize = chooseBlockSize(N, budget, residentInputBytes, nextCount * 12, device); if (blockSize !== BLOCK_SIZE && N > blockSize) { logger.info(`decimate core size ${fmtCount(blockSize)} (memory/device working-set limit)`); } const partSub = logger.group('Partitioning'); - let partition = kdPartition(positions, blockSize, legacy ? null : generation); + let partition = kdPartition(positions, blockSize, generation); let { order, blocks } = partition; partSub.end(); @@ -285,10 +239,10 @@ const decimateSource = async ( const runs = blocks.map(b => coherenceRuns(order, b.start, b.end, COHERENCE_GAP_ROWS)).sort((a, b) => a - b); const median = runs[runs.length >> 1] ?? 0; if (median > INCOHERENT_RUNS_PER_BLOCK) { - if (!legacy && blocks.length > 1) { + if (blocks.length > 1) { if (!opts.spill) { throw new Error( - 'multi-block quality decimation needs scratch storage to stage spatially incoherent input; ' + + 'multi-block adaptive decimation needs scratch storage to stage spatially incoherent input; ' + 'provide opts.spill / --scratch-dir' ); } @@ -349,13 +303,12 @@ const decimateSource = async ( // Re-costed selection (exact within-generation greedy) when its // resident state fits the budget alongside the base state; one-shot // selection otherwise. Gated per generation, so large scenes regain - // re-costing as soon as the cascade shrinks under the budget. Legacy - // mode uses the pre-study pipeline throughout (no re-costing state). + // re-costing as soon as the cascade shrinks under the budget. const K = chooseK(N, budget); const k = Math.min(KNN_K, Math.max(1, N - 1)); const generationTarget = Math.max(targetCount, N - Math.floor(N / 2)); const needed = N - generationTarget; - const multiBlock = !legacy && blocks.length > 1; + const multiBlock = blocks.length > 1; let selection: SelectionResult | undefined; let storedPlans: StoredBlockPlan[] | undefined; let planPrefixes: Uint32Array | undefined; @@ -364,13 +317,13 @@ const decimateSource = async ( if (multiBlock) { if (!device) { throw new Error( - `multi-block quality decimation requires WebGPU (${fmtCount(N)} splats, ` + - `${fmtCount(blockSize)}-splat cores); provide a device, or use --decimate-balanced` + `multi-block adaptive decimation requires WebGPU (${fmtCount(N)} splats, ` + + `${fmtCount(blockSize)}-splat cores); provide a device, or use --decimate-uniform` ); } if (!opts.spill) { throw new Error( - 'multi-block quality decimation needs scratch storage for merge plans ' + + 'multi-block adaptive decimation needs scratch storage for merge plans ' + `(approximately ${fmtBytes(N * 12)} this generation); provide opts.spill / --scratch-dir` ); } @@ -479,10 +432,10 @@ const decimateSource = async ( ); } else { const baseBytes = residentInputBytes + N * (12 + K * 8 + K * 4 + 4) + 3 * 2 ** 30; - const recost = !legacy && baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; + const recost = baseBytes + N * RECOST_BYTES_PER_GAUSSIAN(k) <= budget; - // One-block quality deliberately follows the pre-existing path - // with no staging, halos, plan files, or k-way coordination. + // One block deliberately follows the pre-existing path with no + // staging, halos, plan files, or k-way coordination. const cand: CandidateArrays | undefined = recost ? undefined : { @@ -492,28 +445,17 @@ const decimateSource = async ( const cacheOut = recost ? new Float32Array(N * CACHE_STRIDE) : undefined; const neighborsOut = recost ? new Uint32Array(N * k) : undefined; const priorityBar = logger.bar('computing merge priorities', N); - if (legacy) { - await runPriorityPass( - { source: src, pool, pos: positions, order, blocks, device, K, k }, - cand!, - n => priorityBar.tick(n), - createLegacyStrategy(src.meta.layouts.color!.stride >> 2) - ); - } else { - await runPriorityPass( - { source: src, pool, pos: positions, order, blocks, device, K, k, cacheOut, neighborsOut }, - cand, - n => priorityBar.tick(n) - ); - } + await runPriorityPass( + { source: src, pool, pos: positions, order, blocks, device, K, k, cacheOut, neighborsOut }, + cand, + n => priorityBar.tick(n) + ); priorityBar.end(); const selectSub = logger.group(recost ? 'Selecting merges (re-costed)' : 'Selecting merges'); - selection = legacy ? - selectMergesLegacy(cand!, N, K, needed) : - cacheOut ? - await selectMergesRecosted({ splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed, device }) : - selectMerges(cand!, N, K, needed); + selection = cacheOut ? + await selectMergesRecosted({ splatCache: cacheOut, neighbors: neighborsOut!, D: k, N, mergesNeeded: needed, device }) : + selectMerges(cand!, N, K, needed); selectSub.end(); removed = selection.removed; } diff --git a/src/lib/decimate/priority.ts b/src/lib/decimate/priority.ts index 46885ad0..fcd9be2e 100644 --- a/src/lib/decimate/priority.ts +++ b/src/lib/decimate/priority.ts @@ -32,48 +32,13 @@ type CandidateArrays = { cost: Float32Array; }; -/** A per-slot GPU cost engine (one thread per dense neighbour slot). */ -type GpuSlotCost = { - execute(view: SplatView, cache: unknown, nbRows: Uint32Array, outCosts: Float32Array): Promise; - destroy(): void; -}; - /** - * The cost kernel the priority pass runs: quality (field-L2, DC-only - * colour) or legacy (KL-style, full-SH colour). Strategies own their cache - * shape (each implementation casts what it created). + * Colour components the pass gathers per row. The field-L2 cost reads only + * DC, so at SH band 3 this is 16× less colour RAM and copy traffic per block + * than gathering every band. (`--decimate-uniform` scores full SH and has its + * own pass — see lib/decimate-uniform/priority.ts.) */ -type CostStrategy = { - /** Colour components the block view needs (undefined = all). */ - colorComponents?: number; - /** Whether the CPU cache is also the GPU upload payload (quality's 16-f32 rows). */ - cacheForGpu: boolean; - /** Build the per-view CPU cost cache. */ - buildCache(view: SplatView): unknown; - /** Per-slot GPU engine factory (capacity = max view rows). */ - createGpu(device: GraphicsDevice, capacity: number, k: number): GpuSlotCost; - /** Evaluate one edge on the CPU. */ - cpuEdge(view: SplatView, cache: unknown, i: number, j: number): number; -}; - -/** The production field-L2 kernel (see edge-cost-cpu.ts / gpu-edge-cost.ts). */ -const qualityStrategy: CostStrategy = { - colorComponents: 3, - cacheForGpu: true, - buildCache(view) { - const cache = new Float32Array((view.geo.length / 8) * CACHE_STRIDE); - buildSplatCache(view, cache); - return cache; - }, - createGpu(device, capacity, k) { - const gpu = new GpuEdgeCost(device, capacity, k); - return { - execute: (view, cache, nbRows, outCosts) => gpu.execute(cache as Float32Array, view.geo.length / 8, nbRows, outCosts), - destroy: () => gpu.destroy() - }; - }, - cpuEdge: (view, cache, i, j) => computeEdgeCost(cache as Float32Array, i, j) -}; +const COLOR_COMPONENTS = 3; /** Everything the block passes need: baked single-LOD source + resident state. */ type PriorityContext = { @@ -121,7 +86,7 @@ type BlockView = { * @param extraGlobals - Sorted out-of-block rows to append after the owned rows. * @param includeOther - Also gather the `other` layer (merge pass only). * @param colorComponents - Colour components to copy per row (default: all). - * The quality cost reads only DC, so its pass gathers 3 — at SH band 3 that + * The adaptive cost reads only DC, so its pass gathers 3 — at SH band 3 that * is 16× less colour RAM and copy traffic per block. The read itself still * decodes whole rows (row-interleaved sources); only the view copy narrows. * @returns The gathered block view. @@ -326,13 +291,11 @@ const sortNeighborRows = ( * @param cand - Candidate arrays (`N*K`) to fill, or undefined to skip all * cost work (cacheOut/neighborsOut persistence only). * @param tick - Optional progress callback (owned gaussians completed). - * @param strategy - The cost kernel (default: the production field-L2). */ const runPriorityPass = async ( ctx: PriorityContext, cand: CandidateArrays | undefined, - tick?: (n: number) => void, - strategy: CostStrategy = qualityStrategy + tick?: (n: number) => void ): Promise => { const { pos, order, blocks, device, K, k } = ctx; @@ -340,7 +303,7 @@ const runPriorityPass = async ( for (const b of blocks) maxOwned = Math.max(maxOwned, b.end - b.start); let gpuKnn: GpuKnn | undefined; - let gpuCost: GpuSlotCost | undefined; + let gpuCost: GpuEdgeCost | undefined; let gpuCostCapacity = Math.ceil(maxOwned * VIEW_GROW); // The forest is built once per generation (its trees are exact and @@ -398,7 +361,7 @@ const runPriorityPass = async ( try { if (device) { gpuKnn = new GpuKnn(device, forest, k); - if (cand) gpuCost = strategy.createGpu(device, gpuCostCapacity, k); + if (cand) gpuCost = new GpuEdgeCost(device, gpuCostCapacity, k); } let next: Promise | null = blocks.length > 0 ? prepare(0) : null; @@ -451,15 +414,15 @@ const runPriorityPass = async ( if (gpuCost && viewN > gpuCostCapacity) { gpuCost.destroy(); gpuCostCapacity = Math.ceil(viewN * 1.1); - gpuCost = strategy.createGpu(device!, gpuCostCapacity, k); + gpuCost = new GpuEdgeCost(device!, gpuCostCapacity, k); } - const { view } = await gatherBlockView(ctx, bi, extraGlobals, false, strategy.colorComponents); + const { view } = await gatherBlockView(ctx, bi, extraGlobals, false, COLOR_COMPONENTS); - // Per-view cost cache: quality's 16-f32 rows double as the GPU - // upload and the re-costed selection's resident copy; legacy's is - // CPU-only (its GPU engine packs from the view internally). - const cache = (strategy.cacheForGpu || !device) ? strategy.buildCache(view) : undefined; + // Per-view cost cache: the 16-f32 rows double as the GPU upload + // and the re-costed selection's resident copy. + const cache = new Float32Array((view.geo.length / 8) * CACHE_STRIDE); + buildSplatCache(view, cache); if (cand) { // Translate neighbour slots: global ids → view rows @@ -478,13 +441,13 @@ const runPriorityPass = async ( const blockCosts = new Float32Array(slots); if (gpuCost) { - await gpuCost.execute(view, cache, nbRow, blockCosts); + await gpuCost.execute(cache, view.geo.length / 8, nbRow, blockCosts); } else { for (let s = 0; s < slots; s++) { const row = nbRow[s]; blockCosts[s] = row === KNN_SENTINEL ? 0 : - strategy.cpuEdge(view, cache, (s / k) | 0, row); + computeEdgeCost(cache, (s / k) | 0, row); } } @@ -519,15 +482,12 @@ const runPriorityPass = async ( } } - // Persist owned rows for re-costed selection: the splat cache - // (identical layout on both paths) and the global neighbour ids - // (sentinel-padded). Quality-only (cacheOut implies the quality - // strategy, whose cache is the 16-f32 rows). + // Persist owned rows for re-costed selection: the splat cache and + // the global neighbour ids (sentinel-padded). if (ctx.cacheOut) { const CO = ctx.cacheOut; - const c16 = cache as Float32Array; for (let qi = 0; qi < nOwned; qi++) { - CO.set(c16.subarray(qi * CACHE_STRIDE, (qi + 1) * CACHE_STRIDE), owned[qi] * CACHE_STRIDE); + CO.set(cache.subarray(qi * CACHE_STRIDE, (qi + 1) * CACHE_STRIDE), owned[qi] * CACHE_STRIDE); } } if (ctx.neighborsOut) { @@ -551,10 +511,8 @@ export { gatherBlockView, sortNeighborRows, indexOfSorted, - qualityStrategy, VIEW_GROW, type CandidateArrays, - type CostStrategy, type PriorityContext, type BlockView }; diff --git a/src/lib/index.ts b/src/lib/index.ts index 751ea5ac..aad13f4d 100644 --- a/src/lib/index.ts +++ b/src/lib/index.ts @@ -42,6 +42,11 @@ export type { export { decimateSource } from './decimate'; export type { DecimateOptions, DecimateSpill } from './decimate'; +// The frozen pre-3.2 decimator (--decimate-uniform); see +// lib/decimate-uniform/README.md +export { decimateSourceUniform } from './decimate-uniform'; +export type { DecimateUniformOptions, DecimateUniformSpill } from './decimate-uniform'; + // Statistics export { computeStats } from './stats'; export type { LodStats, LodStatsData, SourceStats } from './stats'; diff --git a/src/lib/workers/tasks.ts b/src/lib/workers/tasks.ts index 9aa818a7..61408e21 100644 --- a/src/lib/workers/tasks.ts +++ b/src/lib/workers/tasks.ts @@ -1,7 +1,8 @@ import type { TypedArray } from '../data-table/data-table'; import { knnForestQuery, type ForestPart } from '../decimate/knn-core'; import { mergeGroup, createMergeScratch, splatMass } from '../decimate/moment-match'; -import { buildFlatKdTree } from '../spatial/kd-tree'; +import { knnQueryBlock } from '../decimate-uniform/knn-core'; +import { buildFlatKdTree, type FlatKdTree } from '../spatial/kd-tree'; import { quantize1dColumns, type QuantizedColumns } from '../spatial/quantize-1d-core'; import { WebPCodec } from '../utils/webp-codec'; @@ -101,6 +102,36 @@ const taskHandlers = { return { result: out, transfer: [out.buffer as ArrayBuffer] }; }, + // Build + flatten a KD-tree over interleaved LOCAL positions (the + // `--decimate-uniform` GPU path: node splat ids stay local and the + // flattened arrays upload straight into that path's GpuKnn). Serves + // decimate-uniform/ — see its README before changing either handler. + flattenKdTree: (args: { positions: Float32Array }): TaskOutput => { + const n = args.positions.length / 3; + const x = new Float32Array(n); + const y = new Float32Array(n); + const z = new Float32Array(n); + for (let i = 0; i < n; i++) { + x[i] = args.positions[i * 3]; + y[i] = args.positions[i * 3 + 1]; + z[i] = args.positions[i * 3 + 2]; + } + const flat = buildFlatKdTree(x, y, z); + return { + result: flat, + transfer: [ + flat.nodeSplatIdx.buffer, flat.nodePositions.buffer, flat.nodeChildren.buffer + ] as ArrayBuffer[] + }; + }, + + // `--decimate-uniform` CPU-fallback block KNN: exact k-NN of the owned + // prefix within the local point set, as local indices. Frozen. + knnBlock: (args: { positions: Float32Array, ownedCount: number, k: number }): TaskOutput => { + const result = knnQueryBlock(args.positions, args.ownedCount, args.k); + return { result, transfer: [result.buffer as ArrayBuffer] }; + }, + // Decimation merge stream: n-ary moment match of packed member-major // groups. Inputs are member-major (pos 3 / geo 8 / color colorDim floats // per member, groups back to back per `sizes`); outputs are group-major. diff --git a/test/cli.test.mjs b/test/cli.test.mjs index 7b9e45fa..a1781970 100644 --- a/test/cli.test.mjs +++ b/test/cli.test.mjs @@ -185,17 +185,17 @@ describe('CLI decimate (terminal PLY restriction)', () => { assert.strictEqual(written, Math.round(inputCount / 2), `50% of ${inputCount}`); }); - it('--decimate-balanced runs the pre-3.2 algorithm', async () => { + it('--decimate-uniform runs the pre-3.2 algorithm', async () => { const { mkdtemp, readFile: readFileFs, rm } = await import('node:fs/promises'); const { tmpdir } = await import('node:os'); const { join } = await import('node:path'); - const dir = await mkdtemp(join(tmpdir(), 'st-decimate-balanced-cli-')); + const dir = await mkdtemp(join(tmpdir(), 'st-decimate-uniform-cli-')); const balancedPath = join(dir, 'balanced.ply'); const balanced = await runCli([ '--gpu', 'cpu', 'test/fixtures/splat/minimal.splat', - '--decimate-balanced', '50%', + '--decimate-uniform', '50%', balancedPath ]); assert.strictEqual(balanced.code, 0, `balanced CLI failed:\n${balanced.stderr}\n${balanced.stdout}`); diff --git a/test/decimate-multiblock.test.mjs b/test/decimate-multiblock.test.mjs index 1694c297..b0368488 100644 --- a/test/decimate-multiblock.test.mjs +++ b/test/decimate-multiblock.test.mjs @@ -22,12 +22,12 @@ after(() => { device?.destroy?.(); }); -describe('decimateSource multi-block quality path', () => { +describe('decimateSource multi-block adaptive path', () => { it('fails clearly when the memory budget requires multiple blocks without WebGPU', async () => { const { source, pool } = await makeSyntheticSource(65540, 0, 123, { chunkSize: 1024 }); await assert.rejects( decimateSource(source, pool, { targetCount: 65000, memoryBudgetBytes: 1 }), - /multi-block quality decimation requires WebGPU/ + /multi-block adaptive decimation requires WebGPU/ ); }); diff --git a/test/decimate-uniform-parity.test.mjs b/test/decimate-uniform-parity.test.mjs new file mode 100644 index 00000000..a78fd85f --- /dev/null +++ b/test/decimate-uniform-parity.test.mjs @@ -0,0 +1,143 @@ +/** + * Output-parity guards for the `--decimate-uniform` decimator in + * `src/lib/decimate-uniform/` — see that directory's README. + * + * That path is bit-for-bit output-compatible with the 3.1.6 binary, which is + * what makes it a usable reference baseline, so "it still works" is not the + * bar — "it still produces the same bytes" is. Equivalence was verified + * end to end on the two study scenes (fr-sky 5.81M and fr-snow 26.1M, six + * chained halvings, every level identical); these tests are the cheap + * in-suite tripwire for accidental drift away from that state. + * + * 1. Candidates come from the exact k-NN — the property the halo collection + * plus verify/requery backstop exists to guarantee. Integer-only, so it is + * immune to cross-platform float differences. + * 2. Candidate ids match a pinned digest. Also integer-only, and the digest + * was taken from the build verified byte-identical to 3.1.6. + * 3. The GPU pass agrees with the CPU pass — the only coverage of that + * directory's own gpu-knn.ts and gpu-edge-cost.ts. + */ + +import assert from 'node:assert'; +import { createHash } from 'node:crypto'; +import { after, before, describe, it } from 'node:test'; + +import { makeSyntheticSource } from './helpers/synthetic-source.mjs'; +import { kdPartition } from '../src/lib/decimate-uniform/partition.js'; +import { runPriorityPass } from '../src/lib/decimate-uniform/priority.js'; + +const N = 5000; +const K = 4; +const KNN_K = 16; +const SEED = 47; +// Small enough that the scene splits into several blocks, so halos, external +// rows and the verify/requery path all get exercised. +const BLOCK_SIZE = 1200; + +/** Digest of `cand.idx` from the build verified output-identical to 3.1.6. */ +const PINNED_IDX_DIGEST = 'aa4ffce1af6b16e0b02c46758891e078f352d98b23bab7da984adced8bca0e61'; + +let device = null; + +before(async () => { + try { + const { createDevice } = await import('../src/cli/node-device.js'); + device = await createDevice(); + } catch { + device = null; + } +}); + +after(() => { + device?.destroy?.(); +}); + +const runLegacy = async (dev) => { + const { source, pool, pos } = await makeSyntheticSource(N, 1, SEED, { chunkSize: 1024 }); + const { order, blocks } = kdPartition(pos, BLOCK_SIZE); + const cand = { + idx: new Uint32Array(N * K).fill(0xFFFFFFFF), + cost: new Float32Array(N * K).fill(Infinity) + }; + await runPriorityPass( + { source, pool, pos, order, blocks, device: dev, K, k: KNN_K }, + cand + ); + return { cand, pos }; +}; + +// Brute-force exact k-NN sets (integer sets, no distance ties to resolve). +const exactNeighbourSets = (pos, k) => { + const sets = []; + const d2 = new Float64Array(N); + const idx = new Uint32Array(N); + for (let i = 0; i < N; i++) { + for (let j = 0; j < N; j++) { + const dx = pos.x[j] - pos.x[i]; + const dy = pos.y[j] - pos.y[i]; + const dz = pos.z[j] - pos.z[i]; + d2[j] = j === i ? Infinity : dx * dx + dy * dy + dz * dz; + idx[j] = j; + } + const ordered = Array.from(idx).sort((a, b) => d2[a] - d2[b]); + sets.push(new Set(ordered.slice(0, k))); + } + return sets; +}; + +describe('uniform decimator parity', () => { + it('draws candidates from the exact k-NN (halo + verify backstop)', async () => { + const { cand, pos } = await runLegacy(undefined); + const exact = exactNeighbourSets(pos, KNN_K); + + let checked = 0; + for (let g = 0; g < N; g++) { + for (let s = 0; s < K; s++) { + const id = cand.idx[g * K + s]; + if (id === 0xFFFFFFFF) continue; + assert.ok( + exact[g].has(id), + `gaussian ${g} candidate ${s} = ${id} is not among its exact ${KNN_K}-NN ` + + '(halo collection or verify/requery regressed)' + ); + checked++; + } + } + assert.ok(checked > N, `expected most gaussians to have candidates, checked ${checked}`); + }); + + it('candidate ids match the pinned 3.1.6-parity digest', async () => { + const { cand } = await runLegacy(undefined); + const digest = createHash('sha256').update(Buffer.from(cand.idx.buffer)).digest('hex'); + assert.strictEqual( + digest, PINNED_IDX_DIGEST, + 'the uniform decimator changed its candidate selection — if that was ' + + 'deliberate, re-run the whole-scene comparison against the 3.1.6 binary and ' + + 're-baseline the `old` column before repinning (see the directory README)' + ); + }); + + it('GPU legacy pass agrees with the CPU legacy pass', async (t) => { + if (!device) return t.skip('no WebGPU adapter available'); + + const { cand: gpu } = await runLegacy(device); + const { cand: cpu } = await runLegacy(undefined); + + let idSetAgree = 0, costAgree = 0; + for (let g = 0; g < N; g++) { + const gpuIds = new Set(), cpuIds = new Set(); + let rowCostsAgree = true; + for (let s = 0; s < K; s++) { + const cg = gpu.cost[g * K + s], cc = cpu.cost[g * K + s]; + if (Math.abs(cg - cc) > Math.max(1e-3, Math.abs(cc) * 1e-3)) rowCostsAgree = false; + gpuIds.add(gpu.idx[g * K + s]); + cpuIds.add(cpu.idx[g * K + s]); + } + if (rowCostsAgree) costAgree++; + const inter = [...gpuIds].filter(x => cpuIds.has(x)).length; + if (inter >= K - 1) idSetAgree++; // allow one float-order swap at the K boundary + } + assert.ok(costAgree / N >= 0.99, `cost agreement ${(costAgree / N * 100).toFixed(2)}% (want >= 99%)`); + assert.ok(idSetAgree / N >= 0.95, `candidate-id agreement ${(idSetAgree / N * 100).toFixed(2)}% (want >= 95%)`); + }); +}); From 5fbec15e720c6850ecc6d95d41c572c4c826301c Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Thu, 30 Jul 2026 14:55:18 +0100 Subject: [PATCH 18/19] latest --- src/cli/index.ts | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/cli/index.ts b/src/cli/index.ts index af6d21f1..1467c43b 100644 --- a/src/cli/index.ts +++ b/src/cli/index.ts @@ -68,7 +68,6 @@ interface CliOptions extends LibOptions { listGpus: boolean; deviceIdx: number; // -1 = auto, -2 = CPU, 0+ = GPU index scratchDir: string | undefined; // decimation spill location (default: output directory) - decimateUniform: boolean; // --decimate-uniform: the frozen pre-3.2 decimator memoryBudgetBytes: number; // decimation residency policy ceiling (not an allocation, not user-facing) } @@ -127,6 +126,13 @@ const resolveInput = (arg: string): ResolvedInput => { // never dispatched as a data operation). type CliAction = ProcessAction | { kind: 'lod'; value: number }; +// `--decimate` and `--decimate-uniform` both produce a decimate action, so +// which decimator to run rides on the action itself rather than on global +// options — that way it always describes the action actually executed, with no +// dependence on flag ordering or on the "exactly one decimate action" check. +// The extra field is stripped before actions reach the library. +type CliDecimate = Extract & { uniform: boolean }; + // Strip the CLI-only lod tags, narrowing back to dispatchable actions. const stripLodTags = (actions: CliAction[]): ProcessAction[] => { return actions.filter((a): a is ProcessAction => a.kind !== 'lod'); @@ -522,7 +528,6 @@ const parseArguments = async () => { listGpus: v['list-gpus'], deviceIdx, scratchDir: v['scratch-dir'], - decimateUniform: false, // Residency policy ceiling for decimation (not an upfront allocation). // Half the machine's RAM, capped at 48 GiB — derived here because the // library is node-free and cannot read os.totalmem() itself. @@ -699,7 +704,6 @@ const parseArguments = async () => { break; case 'decimate': case 'decimate-uniform': { - if (t.name === 'decimate-uniform') options.decimateUniform = true; const value = t.value.trim(); let count: number | null = null; let percent: number | null = null; @@ -718,11 +722,13 @@ const parseArguments = async () => { } } - current.processActions.push({ + const decimate: CliDecimate = { kind: 'decimate', count, - percent - }); + percent, + uniform: t.name === 'decimate-uniform' + }; + current.processActions.push(decimate); break; } case 'filter-cluster': { @@ -1176,7 +1182,7 @@ const main = async () => { } } const decimateAction = decimateIdx.length === 1 ? - singleSceneActions[decimateIdx[0]] as Extract : + singleSceneActions[decimateIdx[0]] as CliDecimate : null; if ( @@ -1265,7 +1271,7 @@ const main = async () => { scratchDir: options.scratchDir ?? dirname(outputFilename), remove: (path: string) => unlink(path) }; - combined = options.decimateUniform ? + combined = decimateAction.uniform ? await decimateSourceUniform(combined, pool, { targetCount: keepCount, createDevice: deviceCreator, From 209dccf1229bb33bcf4f3651617b99ed5f83d549 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Thu, 30 Jul 2026 15:01:45 +0100 Subject: [PATCH 19/19] latest --- src/lib/decimate-uniform/README.md | 26 +- tools/decimate-exact.mjs | 1059 ++++++++++++++++++++++++++++ tools/decimate-parity.mjs | 196 +++++ tools/frustum-cull.mjs | 160 +++++ tools/sweep-fr.mjs | 113 +++ 5 files changed, 1546 insertions(+), 8 deletions(-) create mode 100644 tools/decimate-exact.mjs create mode 100644 tools/decimate-parity.mjs create mode 100644 tools/frustum-cull.mjs create mode 100644 tools/sweep-fr.mjs diff --git a/src/lib/decimate-uniform/README.md b/src/lib/decimate-uniform/README.md index e76c84fe..f671cc3b 100644 --- a/src/lib/decimate-uniform/README.md +++ b/src/lib/decimate-uniform/README.md @@ -63,19 +63,29 @@ Changes are fine — bug fixes, performance work, new capability — but they ar output changes to a path whose selling point is reproducibility, so they need to be deliberate rather than incidental. Before landing one: -- Re-run the whole-scene comparison against the 3.1.6 binary if you expect - output to be unchanged. Equivalence was last verified on both study scenes, - `fr-sky` (5.81M, 3 SH bands, multi-block) and `fr-snow` (26.1M, DC only, 13 - blocks), six chained halvings each: every level byte-identical, PSNR matching - the published `old` columns exactly. +- Re-run the whole-scene comparison if you expect output to be unchanged. That + is what `tools/decimate-parity.mjs` is for — it chains halvings through a + reference binary's `--decimate` and this tree's `--decimate-uniform`, compares + byte for byte, reports PSNR for both, and exits non-zero on any mismatch: + + ```bash + node tools/decimate-parity.mjs sky --ref splat-transform + node tools/decimate-parity.mjs snow --ref splat-transform + ``` + + Equivalence was last verified against 3.1.6 on both study scenes, `fr-sky` + (5.81M, 3 SH bands, multi-block) and `fr-snow` (26.1M, DC only, 13 blocks), + six chained halvings each: every level byte-identical, PSNR matching the + published `old` columns exactly. - If output *should* change, re-baseline `scenes/DECIMATION-RESULTS.md` — the `old` column is this path, and the study's conclusions are stated relative - to it. -- Repin the digest in `test/decimate-legacy-frozen.test.mjs`, which is the + to it. That document is local-only (`scenes/` is gitignored). +- Repin the digest in `test/decimate-uniform-parity.test.mjs`, which is the in-suite tripwire for accidental drift. ## If it is ever retired Nothing in `src/lib/decimate/` refers to this directory, so: `rm -rf` it, drop `decimateSourceUniform` from `src/lib/index.ts`, drop `--decimate-uniform` -from the CLI, and delete `test/decimate-legacy-frozen.test.mjs`. +from the CLI, and delete `test/decimate-uniform-parity.test.mjs` and +`tools/decimate-parity.mjs`. diff --git a/tools/decimate-exact.mjs b/tools/decimate-exact.mjs new file mode 100644 index 00000000..68eb1ba4 --- /dev/null +++ b/tools/decimate-exact.mjs @@ -0,0 +1,1059 @@ +#!/usr/bin/env node +/** + * Reference "quality ceiling" decimator — exact greedy agglomeration. + * + * The algorithm the production pipeline approximates, run without the + * approximations, to establish the best decimation quality achievable under + * the merge/moment-match model: + * + * - One merge at a time: pop the globally-cheapest candidate merge, commit + * it, re-evaluate the affected candidates, repeat. No batch selection, no + * stale costs — after every commit the merged cluster's candidates are + * re-costed against its *current* state. + * - Cost of a cluster C = exact L2 field error vs the ORIGINAL scene: + * E(C) = || Σ_{k∈C} f_k − f_m(C) ||² + * where f_k = α_k·c_k·G_k are the original member fields (DC colour) and + * f_m(C) is the single moment-matched Gaussian of ALL original members + * (same math as production mergeGroup). A candidate merge's cost is the + * marginal error ΔE = E(A∪B) − E(A) − E(B). Because E is always measured + * against the originals, chained/stretched clusters price in their full + * accumulated error — nothing is hidden by incremental approximations. + * - Nested LOD levels from ONE run: keep merging and snapshot a PLY at each + * target count. Every level approximates the ORIGINAL scene, not the + * previous level, so there is no compounding across levels. + * + * Deliberately not fast: everything resident, single-threaded greedy loop. + * Performance work comes after the quality ceiling is confirmed. + * + * Usage: + * node --import tsx --max-old-space-size=49152 tools/decimate-exact.mjs \ + * --input scenes/bad-sky.ply --out-prefix scenes/bad-sky.exact --halvings 3 + * node --import tsx tools/decimate-exact.mjs --selftest + * + * Emits 1.ply, 2.ply, ... (level counts = successive + * ceil(N/2), matching the production 50% cascade exactly). + */ + +import { openSync, writeSync, closeSync } from 'node:fs'; + +import { knnForestQuery, KNN_SENTINEL } from '../src/lib/decimate/knn-core.js'; +import { + EPS_COV, sigmoid, logit, ellipsoidArea, quatToRotmat, sigmaFromRotVar, + det3, eigenSymmetric3x3, rotmatToQuat, mergeGroup, createMergeScratch, + makeGaussianSamples, gaussLogpdfDiagrot, logAddExp, LOG2PI +} from '../src/lib/decimate/moment-match.js'; +import { buildSplatCache, computeEdgeCost, CACHE_STRIDE } from '../src/lib/decimate/edge-cost-cpu.js'; +import { createChunkDataPool } from '../src/lib/index.js'; +import { readPly } from '../src/lib/readers/read-ply.js'; +import { buildFlatKdTree } from '../src/lib/spatial/kd-tree.js'; +import { NodeReadFileSystem } from '../src/cli/node-file-system.js'; + +const C0 = 0.28209479177387814; +const PI_1_5 = Math.PI ** 1.5; +const TWO_PI_1_5 = (2 * Math.PI) ** 1.5; +const KNN_K = 16; +/** Skip Gaussian products whose exponent bound exceeds this (e^-60 ≈ 9e-27). */ +const CULL_QUAD = 120; + +const log = (msg) => console.log(`[${(performance.now() / 1000).toFixed(1)}s] ${msg}`); + +// --------------------------------------------------------------------------- +// Engine state (module-level typed arrays, sized once by initEngine). +// Clusters are union-find sets over original indices; all per-cluster state is +// indexed by the set's root. Original (per-splat) caches never change. +// --------------------------------------------------------------------------- + +let N = 0, colorDim = 3; + +// Per-original caches (immutable after init). +let px, py, pz; // f32 positions +let cs6; // f32 6N — Σ_k with EPS on the variance diagonal (product path) +let csd; // f32 N — √|Σ_k| (with-EPS) +let cal; // f32 N — α_k +let cmass; // f32 N — α·area + 1e-30 (merge weight) +let ctr; // f32 N — trace(Σ_k) for product culling +let cb; // f32 3N — DC base colour (0.5 + C0·f_dc) +let cbn2; // f32 N — |base|² +let knn; // u32 KNN_K·N — global neighbour ids (or KNN_SENTINEL) + +// Per-cluster state (valid at union-find roots). +let parent, ufsize; // u32 +let W; // f64 — total mass +let mx, my, mz; // f64 — mass-weighted mean +let M2; // f64 6N — central second moments Σ w(δδᵀ + Σ_noEPS) +let colorW; // f32 colorDim·N — mass-weighted raw colour sums +let Sself; // f64 — Σ_{k,l∈C}⟨f_k,f_l⟩ +let Err; // f64 — E(C) vs originals +let version, lastSeq; // u32 +let mHead, mTail, mNext; // u32 member chains (NIL = 0xFFFFFFFF) +const NIL = 0xFFFFFFFF; + +let liveCount = 0; + +/** Max original members per cluster (Infinity = uncapped reference). */ +let maxGroup = Infinity; + +/** + * Size-normalization exponent p: merge cost = ΔE / σ_merged^p, with σ the + * merged Gaussian's geometric-mean std (√|Σ|^{1/3}). p=0 is the pure field-L2 + * (volume, σ³-weighted); p=1 ≈ area (σ²) weighting; p=3 ≈ scale-free + * (uniform-rate) behavior. Interpolates the absolute↔relative spectrum. + */ +let sizeExponent = 0; + +/** + * Viewing-kernel dilation δ (world units): the field-L2 metric is evaluated + * after convolving both sides with N(0, δ²I) — every covariance in the + * products gains +δ²I and amplitudes dilute by √(|Σ|/|Σ+δ²I|) (mip-style). + * Prices merges as seen at resolution δ; output parameters are unaffected. + */ +let dilate = 0; + +/** Cost mode: 'l2' = field-L2 vs originals; 'kl' = legacy pairwise KL between cluster reps. */ +let costMode = 'l2'; + +/** + * Scale-free colour dissimilarity weight λ: cost += λ·Σ_allCoeffs(Δc)² between + * the two clusters' mass-weighted mean colours. Unlike the field-L2's own + * colour sensitivity (which vanishes ∝σ³ for faint splats), this term keeps + * light-vs-dark pairing selective at any scale (old's one good property). + * Ordering-only: Err/E bookkeeping stays pure field-L2. + */ +let colorWeight = 0; + +/** Restrict the λ colour term to the 3 DC coefficients (500M-residency probe). */ +let colorDcOnly = false; + +/** + * Needle-chaining guard: a merge is forbidden (cost = ∞) when the result is + * BOTH longer than either member (σmax > LEN_TOL × member max) AND still + * needle-like (σmax/σmid ≥ NEEDLE_TOL × member max needleness). End-to-end + * chaining of the scan's native thin splats (which otherwise compounds + * 0.2m source needles into metre-long artifacts) fails both tests; legitimate + * merges pass at least one: side-by-side joins thicken (needleness drops), + * and flat pancakes (sky clouds) have σmax/σmid ≈ 1 throughout. Relative + * everywhere — nothing is grandfathered, no scene-scale constants. + */ +let needleGuard = false; +const LEN_TOL = 1.15; +/** Absolute needle-shape threshold: σmax/σmid above this reads as a needle. */ +const NEEDLE_ABS = 6; +let clen; // f32[N] — max σmax over the cluster's ORIGINAL members (anchor, no ratchet) +let cneedle; // f32[N] — cluster rep σmax/σmid (diagnostics) + +const initEngine = (n, dim) => { + N = n; colorDim = dim; + px = new Float32Array(N); py = new Float32Array(N); pz = new Float32Array(N); + cs6 = new Float32Array(N * 6); + csd = new Float32Array(N); cal = new Float32Array(N); cmass = new Float32Array(N); + ctr = new Float32Array(N); cb = new Float32Array(N * 3); cbn2 = new Float32Array(N); + knn = new Uint32Array(N * KNN_K).fill(KNN_SENTINEL); + parent = new Uint32Array(N); ufsize = new Uint32Array(N).fill(1); + W = new Float64Array(N); + mx = new Float64Array(N); my = new Float64Array(N); mz = new Float64Array(N); + M2 = new Float64Array(N * 6); + colorW = new Float32Array(N * colorDim); + Sself = new Float64Array(N); Err = new Float64Array(N); + version = new Uint32Array(N).fill(1); lastSeq = new Uint32Array(N); + clen = new Float32Array(N); cneedle = new Float32Array(N).fill(1); + mHead = new Uint32Array(N); mTail = new Uint32Array(N); mNext = new Uint32Array(N).fill(NIL); + for (let i = 0; i < N; i++) { parent[i] = i; mHead[i] = i; mTail[i] = i; } + liveCount = N; +}; + +// Initialize original i (and its singleton cluster) from raw layer values: +// pos3 floats, geo8 = rot(4 wxyz) + log-scales(3) + logit-opacity, colour row. +const initRow = (i, x, y, z, geo, g8, color, cOff) => { + px[i] = x; py[i] = y; pz[i] = z; + + let qw = geo[g8], qx = geo[g8 + 1], qy = geo[g8 + 2], qz = geo[g8 + 3]; + const invq = 1 / Math.max(Math.hypot(qw, qx, qy, qz), 1e-12); + qw *= invq; qx *= invq; qy *= invq; qz *= invq; + const sx = Math.max(Math.exp(geo[g8 + 4]), 1e-12); + const sy = Math.max(Math.exp(geo[g8 + 5]), 1e-12); + const sz = Math.max(Math.exp(geo[g8 + 6]), 1e-12); + const alpha = sigmoid(geo[g8 + 7]); + + // Product-path caches carry the viewing-kernel dilation (Σ+δ²I with + // amplitude dilution √(|Σ|/|Σ+δ²I|)); moments/emission stay undilated. + const d2 = dilate * dilate; + const vx = sx * sx + EPS_COV, vy = sy * sy + EPS_COV, vz = sz * sz + EPS_COV; + const R = evR, S = evS; // scratch + quatToRotmat(qw, qx, qy, qz, R, 0); + sigmaFromRotVar(R, 0, vx + d2, vy + d2, vz + d2, S, 0); + const i6 = i * 6; + cs6[i6] = S[0]; cs6[i6 + 1] = S[1]; cs6[i6 + 2] = S[2]; + cs6[i6 + 3] = S[4]; cs6[i6 + 4] = S[5]; cs6[i6 + 5] = S[8]; + ctr[i] = S[0] + S[4] + S[8]; + const detU = vx * vy * vz; + const detD = (vx + d2) * (vy + d2) * (vz + d2); + csd[i] = Math.sqrt(Math.max(detD, 1e-60)); + cal[i] = alpha * Math.sqrt(detU / detD); + const mass = alpha * ellipsoidArea(sx, sy, sz) + 1e-30; + cmass[i] = mass; + const sHi = Math.max(sx, sy, sz), sLo = Math.min(sx, sy, sz); + const sMid = sx + sy + sz - sHi - sLo; + clen[i] = sHi; + cneedle[i] = sHi / Math.max(sMid, 1e-12); + + const b0 = 0.5 + C0 * color[cOff]; + const b1 = 0.5 + C0 * color[cOff + 1]; + const b2 = 0.5 + C0 * color[cOff + 2]; + cb[i * 3] = b0; cb[i * 3 + 1] = b1; cb[i * 3 + 2] = b2; + cbn2[i] = b0 * b0 + b1 * b1 + b2 * b2; + + // Singleton cluster: exact moments of one member (undilated, Σ WITHOUT the + // EPS used by the product path — matches production mergeGroup member math; + // the dilated cache S carries EPS+δ² on its diagonal). + W[i] = mass; + mx[i] = x; my[i] = y; mz[i] = z; + M2[i6] = mass * (S[0] - EPS_COV - d2); + M2[i6 + 1] = mass * S[1]; + M2[i6 + 2] = mass * S[2]; + M2[i6 + 3] = mass * (S[4] - EPS_COV - d2); + M2[i6 + 4] = mass * S[5]; + M2[i6 + 5] = mass * (S[8] - EPS_COV - d2); + for (let c = 0; c < colorDim; c++) colorW[i * colorDim + c] = mass * color[cOff + c]; + Sself[i] = cal[i] * cal[i] * cbn2[i] * PI_1_5 * csd[i]; + Err[i] = 0; +}; +const evR = new Float32Array(9), evS = new Float32Array(9); + +const find = (x) => { + while (parent[x] !== x) { parent[x] = parent[parent[x]]; x = parent[x]; } + return x; +}; + +// --------------------------------------------------------------------------- +// Cost evaluation. +// --------------------------------------------------------------------------- + +// ⟨G_a,G_b⟩ scaled by √|Σa|·√|Σb| for M = Σa+Σb (6 comps) and offset d. +const crossG = (sdAB, m0, m1, m2, m3, m4, m5, dx, dy, dz) => { + const c00 = m3 * m5 - m4 * m4; + const c01 = m2 * m4 - m1 * m5; + const c02 = m1 * m4 - m2 * m3; + const c11 = m0 * m5 - m2 * m2; + const c12 = m1 * m2 - m0 * m4; + const c22 = m0 * m3 - m1 * m1; + const det = Math.max(m0 * c00 + m1 * c01 + m2 * c02, 1e-60); + const quad = (c00 * dx * dx + c11 * dy * dy + c22 * dz * dz + + 2 * (c01 * dx * dy + c02 * dx * dz + c12 * dy * dz)) / det; + if (!(quad < CULL_QUAD)) return 0; + return TWO_PI_1_5 * sdAB / Math.sqrt(det) * Math.exp(-0.5 * quad); +}; + +// Smith closed-form eigenvalues of a symmetric 3×3 (6 comps) — for merged area. +const eig3 = (m0, m1, m2, m3, m4, m5, out) => { + const q = (m0 + m3 + m5) / 3; + const p1 = m1 * m1 + m2 * m2 + m4 * m4; + if (p1 <= 1e-30) { out[0] = m0; out[1] = m3; out[2] = m5; return; } + const p2 = (m0 - q) * (m0 - q) + (m3 - q) * (m3 - q) + (m5 - q) * (m5 - q) + 2 * p1; + const p = Math.sqrt(p2 / 6); + const ip = 1 / p; + const b00 = (m0 - q) * ip, b11 = (m3 - q) * ip, b22 = (m5 - q) * ip; + const b01 = m1 * ip, b02 = m2 * ip, b12 = m4 * ip; + const detB = b00 * (b11 * b22 - b12 * b12) - b01 * (b01 * b22 - b12 * b02) + b02 * (b01 * b12 - b11 * b02); + let r = detB / 2; + r = r < -1 ? -1 : (r > 1 ? 1 : r); + const phi = Math.acos(r) / 3; + const e0 = q + 2 * p * Math.cos(phi); + const e2 = q + 2 * p * Math.cos(phi + 2 * Math.PI / 3); + out[0] = e0; out[1] = 3 * q - e0 - e2; out[2] = e2; +}; +const eigOut = new Float64Array(3); + +// Member gather scratch (cluster A's members flattened for the Scross loop). +let gatherBuf = new Uint32Array(1 << 16); +const gatherMembers = (root) => { + let cnt = 0; + for (let m = mHead[root]; m !== NIL; m = mNext[m]) { + if (cnt === gatherBuf.length) { + const g = new Uint32Array(gatherBuf.length * 2); + g.set(gatherBuf); gatherBuf = g; + } + gatherBuf[cnt++] = m; + } + return cnt; +}; + +const evalOut = { E: 0, Scross: 0 }; + +// ---- Legacy-KL cost mode: pairwise cost between the two clusters' current +// representative gaussians (verbatim port of the legacy computeEdgeCost: +// KL-style geometric term with one MC sample + L2 over colour coefficients). +// O(1) per eval regardless of member count; incremental (not vs originals). +const klZ = makeGaussianSamples(1, 0)[0]; +const klS = { + SigA: new Float64Array(9), SigB: new Float64Array(9), + eigA: new Float64Array(9), eigV: new Float64Array(9), + RA: new Float64Array(9), RB: new Float64Array(9), + a: new Float64Array(8), b: new Float64Array(8), // v3, invd3... packed below + sigm: new Float64Array(9) +}; +// Fill side params for cluster C: Sig (9), R (9), out = [vx,vy,vz,ldet,mass]. +const klSide = (C, Sig, R, out) => { + const i6 = C * 6, iw = 1 / W[C]; + Sig[0] = M2[i6] * iw + EPS_COV; + Sig[1] = Sig[3] = M2[i6 + 1] * iw; + Sig[2] = Sig[6] = M2[i6 + 2] * iw; + Sig[4] = M2[i6 + 3] * iw + EPS_COV; + Sig[5] = Sig[7] = M2[i6 + 4] * iw; + Sig[8] = M2[i6 + 5] * iw + EPS_COV; + eigenSymmetric3x3(Sig, klS.eigA, klS.eigV); + R.set(klS.eigV); + const v0 = Math.max(klS.eigA[0], 1e-30); + const v1 = Math.max(klS.eigA[4], 1e-30); + const v2 = Math.max(klS.eigA[8], 1e-30); + out[0] = v0; out[1] = v1; out[2] = v2; + out[3] = Math.log(v0) + Math.log(v1) + Math.log(v2); + const s0 = Math.sqrt(v0), s1 = Math.sqrt(v1), s2 = Math.sqrt(v2); + const area = ellipsoidArea(s0, s1, s2); + const alphaM = Math.min(1, W[C] / Math.max(area, 1e-30)); + out[4] = alphaM * area + 1e-12; +}; + +const evalKl = (A, B) => { + evalOut.E = 0; evalOut.Scross = 0; + const { SigA, SigB, RA, RB, a, b, sigm } = klS; + klSide(A, SigA, RA, a); + klSide(B, SigB, RB, b); + + const mux = mx[A], muy = my[A], muz = mz[A]; + const mvx = mx[B], mvy = my[B], mvz = mz[B]; + const wi = a[4], wj = b[4]; + const Wsafe = wi + wj > 0 ? wi + wj : 1; + let pi_ = wi / Wsafe; + pi_ = Math.max(1e-12, Math.min(1 - 1e-12, pi_)); + const pj_ = 1 - pi_; + const logPi = Math.log(pi_), logPj = Math.log(pj_); + + const mmx = pi_ * mux + pj_ * mvx; + const mmy = pi_ * muy + pj_ * mvy; + const mmz = pi_ * muz + pj_ * mvz; + const dix = mux - mmx, diy = muy - mmy, diz = muz - mmz; + const djx = mvx - mmx, djy = mvy - mmy, djz = mvz - mmz; + + for (let t = 0; t < 9; t++) sigm[t] = pi_ * SigA[t] + pj_ * SigB[t]; + sigm[0] += pi_ * dix * dix + pj_ * djx * djx + EPS_COV; + sigm[1] += pi_ * dix * diy + pj_ * djx * djy; + sigm[2] += pi_ * dix * diz + pj_ * djx * djz; + sigm[3] = sigm[1]; + sigm[4] += pi_ * diy * diy + pj_ * djy * djy + EPS_COV; + sigm[5] += pi_ * diy * diz + pj_ * djy * djz; + sigm[6] = sigm[2]; sigm[7] = sigm[5]; + sigm[8] += pi_ * diz * diz + pj_ * djz * djz + EPS_COV; + const detm = Math.max(det3(sigm, 0), 1e-30); + const EpNegLogQ = 0.5 * (3 * LOG2PI + Math.log(detm) + 3); + + const z0 = klZ[0], z1 = klZ[1], z2 = klZ[2]; + const sia = Math.sqrt(a[0]), sib = Math.sqrt(a[1]), sic = Math.sqrt(a[2]); + const sja = Math.sqrt(b[0]), sjb = Math.sqrt(b[1]), sjc = Math.sqrt(b[2]); + const xix = mux + z0 * sia * RA[0] + z1 * sib * RA[1] + z2 * sic * RA[2]; + const xiy = muy + z0 * sia * RA[3] + z1 * sib * RA[4] + z2 * sic * RA[5]; + const xiz = muz + z0 * sia * RA[6] + z1 * sib * RA[7] + z2 * sic * RA[8]; + const xjx = mvx + z0 * sja * RB[0] + z1 * sjb * RB[1] + z2 * sjc * RB[2]; + const xjy = mvy + z0 * sja * RB[3] + z1 * sjb * RB[4] + z2 * sjc * RB[5]; + const xjz = mvz + z0 * sja * RB[6] + z1 * sjb * RB[7] + z2 * sjc * RB[8]; + + const ia = 1 / a[0], ib = 1 / a[1], ic = 1 / a[2]; + const ja = 1 / b[0], jb = 1 / b[1], jc = 1 / b[2]; + const logNiOnI = gaussLogpdfDiagrot(xix, xiy, xiz, mux, muy, muz, RA, 0, ia, ib, ic, a[3]); + const logNjOnI = gaussLogpdfDiagrot(xix, xiy, xiz, mvx, mvy, mvz, RB, 0, ja, jb, jc, b[3]); + const logNiOnJ = gaussLogpdfDiagrot(xjx, xjy, xjz, mux, muy, muz, RA, 0, ia, ib, ic, a[3]); + const logNjOnJ = gaussLogpdfDiagrot(xjx, xjy, xjz, mvx, mvy, mvz, RB, 0, ja, jb, jc, b[3]); + const Ei = logAddExp(logPi + logNiOnI, logPj + logNjOnI); + const Ej = logAddExp(logPi + logNiOnJ, logPj + logNjOnJ); + const geo = pi_ * Ei + pj_ * Ej + EpNegLogQ; + + let cSh = 0; + const cd = colorDim, iwA = 1 / W[A], iwB = 1 / W[B]; + for (let c = 0; c < cd; c++) { + const d = colorW[A * cd + c] * iwA - colorW[B * cd + c] * iwB; + cSh += d * d; + } + return geo + cSh; +}; + +// Marginal cost ΔE of merging clusters A and B (exact, vs originals). +// Fills evalOut with E(A∪B) and Scross(A,B) for reuse at commit. +const evalMerge = (A, B) => { + if (costMode === 'kl') return evalKl(A, B); + const WA = W[A], WB = W[B], WC = WA + WB; + const iw = 1 / WC; + const mcx = (WA * mx[A] + WB * mx[B]) * iw; + const mcy = (WA * my[A] + WB * my[B]) * iw; + const mcz = (WA * mz[A] + WB * mz[B]) * iw; + const dax = mx[A] - mcx, day = my[A] - mcy, daz = mz[A] - mcz; + const dbx = mx[B] - mcx, dby = my[B] - mcy, dbz = mz[B] - mcz; + const a6 = A * 6, b6 = B * 6; + + // Merged covariance Σm = (M2_A + M2_B + shift terms)/W + EPS·I. + const sm0 = (M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx) * iw + EPS_COV; + const sm1 = (M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby) * iw; + const sm2 = (M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz) * iw; + const sm3 = (M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby) * iw + EPS_COV; + const sm4 = (M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz) * iw; + const sm5 = (M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz) * iw + EPS_COV; + + const detm = Math.max( + sm0 * (sm3 * sm5 - sm4 * sm4) - sm1 * (sm1 * sm5 - sm4 * sm2) + sm2 * (sm1 * sm4 - sm3 * sm2), + 1e-60 + ); + const sdC = Math.sqrt(detm); + + eig3(sm0, sm1, sm2, sm3, sm4, sm5, eigOut); + const s0 = Math.sqrt(Math.max(eigOut[0], 1e-18)); + const s1 = Math.sqrt(Math.max(eigOut[1], 1e-18)); + const s2 = Math.sqrt(Math.max(eigOut[2], 1e-18)); + + // Needle-chaining guard: a needle-shaped result (σmax/σmid > NEEDLE_ABS) + // may never be longer than LEN_TOL × the longest ORIGINAL member — the + // original-length anchor kills the per-merge growth ratchet. Thickening is + // always allowed (aspect ≤ NEEDLE_ABS passes unconditionally). Checked + // before the expensive member loops. + if (needleGuard) { + const needleM = s0 / Math.max(s1, 1e-12); + if (needleM > NEEDLE_ABS && s0 > LEN_TOL * Math.max(clen[A], clen[B])) return Infinity; + } + + const alphaC = Math.min(1, WC / Math.max(ellipsoidArea(s0, s1, s2), 1e-30)); + + const cd = colorDim; + const bc0 = 0.5 + C0 * (colorW[A * cd] + colorW[B * cd]) * iw; + const bc1 = 0.5 + C0 * (colorW[A * cd + 1] + colorW[B * cd + 1]) * iw; + const bc2 = 0.5 + C0 * (colorW[A * cd + 2] + colorW[B * cd + 2]) * iw; + const bn2C = bc0 * bc0 + bc1 * bc1 + bc2 * bc2; + + // Viewing-kernel dilation of the merged gaussian for the product terms + // (member caches are already dilated); amplitude dilutes accordingly. + const dd = dilate * dilate; + const smD0 = sm0 + dd, smD3 = sm3 + dd, smD5 = sm5 + dd; + const detmD = dd === 0 ? detm : Math.max( + smD0 * (smD3 * smD5 - sm4 * sm4) - sm1 * (sm1 * smD5 - sm4 * sm2) + sm2 * (sm1 * sm4 - smD3 * sm2), + 1e-60 + ); + const sdCD = Math.sqrt(detmD); + const aCe = dd === 0 ? alphaC : alphaC * Math.sqrt(detm / detmD); + + const selfM = aCe * aCe * bn2C * PI_1_5 * sdCD; + + // ⟨Σ member fields, f_m⟩ over both chains (dilated space). + let memfm = 0; + for (let pass = 0; pass < 2; pass++) { + for (let m = pass === 0 ? mHead[A] : mHead[B]; m !== NIL; m = mNext[m]) { + const m6 = m * 6; + const wgt = cal[m] * aCe * (cb[m * 3] * bc0 + cb[m * 3 + 1] * bc1 + cb[m * 3 + 2] * bc2); + if (wgt === 0) continue; + memfm += wgt * crossG(csd[m] * sdCD, + cs6[m6] + smD0, cs6[m6 + 1] + sm1, cs6[m6 + 2] + sm2, + cs6[m6 + 3] + smD3, cs6[m6 + 4] + sm4, cs6[m6 + 5] + smD5, + px[m] - mcx, py[m] - mcy, pz[m] - mcz); + } + } + + // Scross(A,B) = Σ_{a∈A,b∈B}⟨f_a,f_b⟩ with distance culling. + const na = gatherMembers(A); + let scross = 0; + for (let b = mHead[B]; b !== NIL; b = mNext[b]) { + const b6i = b * 6, b3 = b * 3; + const bx = px[b], by = py[b], bz = pz[b]; + const trb = ctr[b], alb = cal[b], sdb = csd[b]; + const cb0 = cb[b3], cb1 = cb[b3 + 1], cb2 = cb[b3 + 2]; + for (let t = 0; t < na; t++) { + const a = gatherBuf[t]; + const dx = px[a] - bx, dy = py[a] - by, dz = pz[a] - bz; + const d2 = dx * dx + dy * dy + dz * dz; + if (d2 > CULL_QUAD * (ctr[a] + trb)) continue; + const a6i = a * 6, a3 = a * 3; + const wgt = cal[a] * alb * (cb[a3] * cb0 + cb[a3 + 1] * cb1 + cb[a3 + 2] * cb2); + if (wgt === 0) continue; + scross += wgt * crossG(csd[a] * sdb, + cs6[a6i] + cs6[b6i], cs6[a6i + 1] + cs6[b6i + 1], cs6[a6i + 2] + cs6[b6i + 2], + cs6[a6i + 3] + cs6[b6i + 3], cs6[a6i + 4] + cs6[b6i + 4], cs6[a6i + 5] + cs6[b6i + 5], + dx, dy, dz); + } + } + + const E = Sself[A] + Sself[B] + 2 * scross - 2 * memfm + selfM; + evalOut.E = E; + evalOut.Scross = scross; + const dE = E - Err[A] - Err[B]; + // Size normalization: σ_gm = (√|Σm|)^{1/3}; cost = ΔE / σ_gm^p. + let cost = sizeExponent === 0 ? dE : dE / Math.pow(sdC, sizeExponent / 3); + if (colorWeight > 0) { + const iwA = 1 / W[A], iwB = 1 / W[B]; + const lim = colorDcOnly ? Math.min(3, cd) : cd; + let cd2 = 0; + for (let c = 0; c < lim; c++) { + const d = colorW[A * cd + c] * iwA - colorW[B * cd + c] * iwB; + cd2 += d * d; + } + cost += colorWeight * cd2; + } + return cost; +}; + +// --------------------------------------------------------------------------- +// Candidate derivation: a cluster's candidates are the live clusters owning +// any original-KNN neighbour of any member. Reuses the static KNN graph. +// --------------------------------------------------------------------------- + +let stamp, stampGen = 0; +let candBuf = new Uint32Array(1 << 12); + +const deriveCandidates = (root) => { + stampGen++; + const gen = stampGen; + let cnt = 0; + for (let m = mHead[root]; m !== NIL; m = mNext[m]) { + const base = m * KNN_K; + for (let s = 0; s < KNN_K; s++) { + const nb = knn[base + s]; + if (nb === KNN_SENTINEL) continue; + const r = find(nb); + if (r === root || stamp[r] === gen) continue; + stamp[r] = gen; + if (cnt === candBuf.length) { + const g = new Uint32Array(candBuf.length * 2); + g.set(candBuf); candBuf = g; + } + candBuf[cnt++] = r; + } + } + return cnt; +}; + +// --------------------------------------------------------------------------- +// Binary min-heap of candidate edges (one live entry per cluster: its current +// best edge). Lazy invalidation via (seq, partner-version). +// --------------------------------------------------------------------------- + +let hCost, hA, hB, hSeq, hVb, heapSize = 0, heapCap = 0; +let seqCounter = 0; + +const heapInit = (cap) => { + heapCap = cap; + hCost = new Float64Array(cap); + hA = new Uint32Array(cap); hB = new Uint32Array(cap); + hSeq = new Uint32Array(cap); hVb = new Uint32Array(cap); + heapSize = 0; +}; + +const heapPush = (cost, a, b, seq, vb) => { + if (heapSize === heapCap) { + const nc = heapCap * 2; + const c2 = new Float64Array(nc); c2.set(hCost); hCost = c2; + const g = (old) => { const x = new Uint32Array(nc); x.set(old); return x; }; + hA = g(hA); hB = g(hB); hSeq = g(hSeq); hVb = g(hVb); + heapCap = nc; + } + let i = heapSize++; + hCost[i] = cost; hA[i] = a; hB[i] = b; hSeq[i] = seq; hVb[i] = vb; + while (i > 0) { + const p = (i - 1) >> 1; + if (hCost[p] <= hCost[i]) break; + swap(i, p); i = p; + } +}; + +const swap = (i, j) => { + let t; + t = hCost[i]; hCost[i] = hCost[j]; hCost[j] = t; + t = hA[i]; hA[i] = hA[j]; hA[j] = t; + t = hB[i]; hB[i] = hB[j]; hB[j] = t; + t = hSeq[i]; hSeq[i] = hSeq[j]; hSeq[j] = t; + t = hVb[i]; hVb[i] = hVb[j]; hVb[j] = t; +}; + +const popOut = { cost: 0, a: 0, b: 0, seq: 0, vb: 0 }; +const heapPop = () => { + if (heapSize === 0) return false; + popOut.cost = hCost[0]; popOut.a = hA[0]; popOut.b = hB[0]; popOut.seq = hSeq[0]; popOut.vb = hVb[0]; + heapSize--; + if (heapSize > 0) { + hCost[0] = hCost[heapSize]; hA[0] = hA[heapSize]; hB[0] = hB[heapSize]; + hSeq[0] = hSeq[heapSize]; hVb[0] = hVb[heapSize]; + let i = 0; + for (;;) { + const l = 2 * i + 1, r = l + 1; + let m = i; + if (l < heapSize && hCost[l] < hCost[m]) m = l; + if (r < heapSize && hCost[r] < hCost[m]) m = r; + if (m === i) break; + swap(i, m); i = m; + } + } + return true; +}; + +// Largest live cluster (diagnostic; O(N) but only called from throttled logs). +const maxLiveSize = () => { + let mx = 0; + for (let i = 0; i < N; i++) { + if (parent[i] === i && ufsize[i] > mx) mx = ufsize[i]; + } + return mx; +}; + +// Log2-bucketed live cluster-size histogram, e.g. "1:2041 2-3:511 4-7:88". +const sizeHistogram = () => { + const buckets = new Map(); + for (let i = 0; i < N; i++) { + if (parent[i] !== i) continue; + const b = Math.floor(Math.log2(ufsize[i])); + buckets.set(b, (buckets.get(b) ?? 0) + 1); + } + return [...buckets.entries()].sort((a, b) => a[0] - b[0]) + .map(([b, c]) => `${1 << b}${b > 0 ? `-${(1 << (b + 1)) - 1}` : ''}:${c}`).join(' '); +}; + +// Recompute a cluster's best edge and push it (bumps lastSeq — older entries +// for this cluster become stale). +const pushBestEdge = (root) => { + const cnt = deriveCandidates(root); + lastSeq[root] = ++seqCounter; + if (cnt === 0) return; + const sz = ufsize[root]; + let bc = Infinity, bp = -1, bv = 0; + for (let t = 0; t < cnt; t++) { + const cand = candBuf[t]; + if (sz + ufsize[cand] > maxGroup) continue; + const d = evalMerge(root, cand); + if (d < bc) { bc = d; bp = cand; bv = version[cand]; } + } + if (bp >= 0) heapPush(bc, root, bp, lastSeq[root], bv); +}; + +// Commit the merge of roots A and B. Returns the marginal error added. +const commitMerge = (A, B) => { + const dE = evalMerge(A, B); // exact recompute (E, Scross) + const E = evalOut.E, scross = evalOut.Scross; + + const keep = ufsize[A] >= ufsize[B] ? A : B; + const lose = keep === A ? B : A; + + // Moments (compute with locals before overwriting keep's slots). + const WA = W[A], WB = W[B], WC = WA + WB; + const iw = 1 / WC; + const mcx = (WA * mx[A] + WB * mx[B]) * iw; + const mcy = (WA * my[A] + WB * my[B]) * iw; + const mcz = (WA * mz[A] + WB * mz[B]) * iw; + const dax = mx[A] - mcx, day = my[A] - mcy, daz = mz[A] - mcz; + const dbx = mx[B] - mcx, dby = my[B] - mcy, dbz = mz[B] - mcz; + const a6 = A * 6, b6 = B * 6, k6 = keep * 6; + const n0 = M2[a6] + M2[b6] + WA * dax * dax + WB * dbx * dbx; + const n1 = M2[a6 + 1] + M2[b6 + 1] + WA * dax * day + WB * dbx * dby; + const n2 = M2[a6 + 2] + M2[b6 + 2] + WA * dax * daz + WB * dbx * dbz; + const n3 = M2[a6 + 3] + M2[b6 + 3] + WA * day * day + WB * dby * dby; + const n4 = M2[a6 + 4] + M2[b6 + 4] + WA * day * daz + WB * dby * dbz; + const n5 = M2[a6 + 5] + M2[b6 + 5] + WA * daz * daz + WB * dbz * dbz; + M2[k6] = n0; M2[k6 + 1] = n1; M2[k6 + 2] = n2; M2[k6 + 3] = n3; M2[k6 + 4] = n4; M2[k6 + 5] = n5; + W[keep] = WC; mx[keep] = mcx; my[keep] = mcy; mz[keep] = mcz; + + const cd = colorDim, kc = keep * cd, lc = lose * cd; + for (let c = 0; c < cd; c++) colorW[kc + c] += colorW[lc + c]; + + Sself[keep] = Sself[A] + Sself[B] + 2 * scross; + Err[keep] = E; + ufsize[keep] += ufsize[lose]; + mNext[mTail[keep]] = mHead[lose]; + mTail[keep] = mTail[lose]; + parent[lose] = keep; + version[keep]++; + liveCount--; + + // Maintain guard state: clen = max ORIGINAL member σmax (composes by max — + // the anchor that prevents compounding growth); cneedle = current rep + // needleness (diagnostics). + clen[keep] = Math.max(clen[A], clen[B]); + eig3(n0 / WC + EPS_COV, n1 / WC, n2 / WC, n3 / WC + EPS_COV, n4 / WC, n5 / WC + EPS_COV, eigOut); + cneedle[keep] = Math.sqrt(Math.max(eigOut[0], 1e-18) / Math.max(eigOut[1], 1e-18)); + + pushBestEdge(keep); + return dE; +}; + +// --------------------------------------------------------------------------- +// Main greedy loop with nested snapshots. +// --------------------------------------------------------------------------- + +const runGreedy = (targets, onSnapshot, progressEvery = 1_000_000) => { + log(`initial best edges (${liveCount} clusters)…`); + for (let i = 0; i < N; i++) { + pushBestEdge(i); + if ((i + 1) % 5_000_000 === 0) log(` seeded ${i + 1}/${N}`); + } + log(`greedy loop → targets [${targets.join(', ')}]`); + + let ti = 0; + let commits = 0, pops = 0, recomputes = 0; + let totalDE = 0; + const t0 = performance.now(); + let lastLogAt = t0; + + while (ti < targets.length && liveCount > targets[ti]) { + if (!heapPop()) { + log(`heap exhausted at ${liveCount} clusters (target ${targets[ti]}) — emitting partial level and stopping`); + onSnapshot(ti, liveCount); + ti++; + break; + } + pops++; + const a = popOut.a; + // Stale checks: a must still be a live root and this must be its + // latest entry (seq is globally monotonic, so no collisions); the + // partner must still be a live root with an unchanged version — + // find(b) is NOT sufficient (an absorbed b resolves to a different + // cluster whose independent version counter can coincidentally match, + // committing a stale cost and bypassing the cap). + if (parent[a] !== a || popOut.seq !== lastSeq[a]) continue; + const b = popOut.b; + if (parent[b] !== b || version[b] !== popOut.vb || + ufsize[a] + ufsize[b] > maxGroup /* defensive: cap must hold */) { + recomputes++; + pushBestEdge(a); + continue; + } + totalDE += commitMerge(a, b); + commits++; + + if (commits % progressEvery === 0 || (commits & 0xFFF) === 0) { + const now = performance.now(); + if (commits % progressEvery === 0 || now - lastLogAt > 30_000) { + lastLogAt = now; + const dt = (now - t0) / 1000; + log(` ${commits} merges (${(commits / dt).toFixed(0)}/s avg) · live ${liveCount} · heap ${heapSize} · pops ${pops} · recomputes ${recomputes} · maxSize ${maxLiveSize()} · ΣΔE ${totalDE.toExponential(3)}`); + } + } + if (liveCount === targets[ti]) { + log(`snapshot ${ti + 1}: ${liveCount} clusters · ΣΔE ${totalDE.toExponential(4)} · sizes ${sizeHistogram()}`); + onSnapshot(ti, liveCount); + ti++; + } + } + const dt = (performance.now() - t0) / 1000; + log(`greedy done: ${commits} merges in ${dt.toFixed(1)}s · ${recomputes} recomputes · ${pops} pops`); +}; + +// --------------------------------------------------------------------------- +// Emission: moment-match every live cluster (identical math to mergeGroup) and +// stream a binary-little-endian PLY. +// --------------------------------------------------------------------------- + +const emitPly = (path) => { + const props = ['x', 'y', 'z']; + props.push('f_dc_0', 'f_dc_1', 'f_dc_2'); + for (let r = 0; r < colorDim - 3; r++) props.push(`f_rest_${r}`); + props.push('opacity', 'scale_0', 'scale_1', 'scale_2', 'rot_0', 'rot_1', 'rot_2', 'rot_3'); + + const header = `ply\nformat binary_little_endian 1.0\nelement vertex ${liveCount}\n${props.map(p => `property float ${p}`).join('\n')}\nend_header\n`; + const stride = props.length * 4; + + const fd = openSync(path, 'w'); + writeSync(fd, Buffer.from(header, 'ascii')); + + const ROWS = 65536; + const buf = Buffer.allocUnsafe(ROWS * stride); + const f32 = new Float32Array(buf.buffer, buf.byteOffset, ROWS * (stride >> 2)); + let rows = 0, written = 0; + + const Sig = new Float64Array(9), eigA = new Float64Array(9), eigV = new Float64Array(9); + const Rm = new Float64Array(9), quat = new Float64Array(4); + + for (let i = 0; i < N; i++) { + if (parent[i] !== i) continue; + const i6 = i * 6, iw = 1 / W[i]; + // Σm = M2/W + EPS·I (same as mergeGroup's Σp(δδᵀ+Σ)+EPS). + Sig[0] = M2[i6] * iw + EPS_COV; + Sig[1] = Sig[3] = M2[i6 + 1] * iw; + Sig[2] = Sig[6] = M2[i6 + 2] * iw; + Sig[4] = M2[i6 + 3] * iw + EPS_COV; + Sig[5] = Sig[7] = M2[i6 + 4] * iw; + Sig[8] = M2[i6 + 5] * iw + EPS_COV; + eigenSymmetric3x3(Sig, eigA, eigV); + + // Order eigenpairs descending (mergeGroup's o0/o1/o2 logic). + const v0 = eigA[0], v1 = eigA[4], v2 = eigA[8]; + let o0, o1, o2; + if (v0 >= v1) { + if (v1 >= v2) { o0 = 0; o1 = 1; o2 = 2; } else if (v0 >= v2) { o0 = 0; o1 = 2; o2 = 1; } else { o0 = 2; o1 = 0; o2 = 1; } + } else if (v0 >= v2) { o0 = 1; o1 = 0; o2 = 2; } else if (v1 >= v2) { o0 = 1; o1 = 2; o2 = 0; } else { o0 = 2; o1 = 1; o2 = 0; } + const ev0 = Math.max(eigA[3 * o0 + o0], 1e-18); + const ev1 = Math.max(eigA[3 * o1 + o1], 1e-18); + const ev2 = Math.max(eigA[3 * o2 + o2], 1e-18); + const s0 = Math.sqrt(ev0), s1 = Math.sqrt(ev1), s2 = Math.sqrt(ev2); + const alphaM = Math.min(1, W[i] / Math.max(ellipsoidArea(s0, s1, s2), 1e-30)); + + Rm[0] = eigV[o0]; Rm[1] = eigV[o1]; Rm[2] = eigV[o2]; + Rm[3] = eigV[3 + o0]; Rm[4] = eigV[3 + o1]; Rm[5] = eigV[3 + o2]; + Rm[6] = eigV[6 + o0]; Rm[7] = eigV[6 + o1]; Rm[8] = eigV[6 + o2]; + if (det3(Rm, 0) < 0) { Rm[2] *= -1; Rm[5] *= -1; Rm[8] *= -1; } + rotmatToQuat(Rm, 0, quat, 0); + + const o = rows * (stride >> 2); + f32[o] = mx[i]; f32[o + 1] = my[i]; f32[o + 2] = mz[i]; + const cbase = i * colorDim; + for (let c = 0; c < colorDim; c++) f32[o + 3 + c] = colorW[cbase + c] * iw; + const oo = o + 3 + colorDim; + f32[oo] = logit(Math.max(0, Math.min(1, alphaM))); + f32[oo + 1] = Math.log(s0); f32[oo + 2] = Math.log(s1); f32[oo + 3] = Math.log(s2); + f32[oo + 4] = quat[0]; f32[oo + 5] = quat[1]; f32[oo + 6] = quat[2]; f32[oo + 7] = quat[3]; + + if (++rows === ROWS) { writeSync(fd, buf, 0, rows * stride); written += rows; rows = 0; } + } + if (rows > 0) { writeSync(fd, buf, 0, rows * stride); written += rows; } + closeSync(fd); + if (written !== liveCount) throw new Error(`emitted ${written} rows, expected ${liveCount}`); + log(`wrote ${path} (${written} gaussians)`); +}; + +// --------------------------------------------------------------------------- +// PLY load → engine init (single sequential pass; originals never stored raw). +// --------------------------------------------------------------------------- + +const loadPly = async (filename) => { + const pool = createChunkDataPool(); + const fs = new NodeReadFileSystem(); + const src = await readPly(await fs.createSource(filename), pool); + const { meta } = src; + if (meta.numLods !== 1) throw new Error('single-LOD input required'); + const dim = meta.layouts.color.stride >> 2; + log(`loading ${filename}: ${meta.numGaussians} gaussians · colorDim ${dim}`); + initEngine(meta.numGaussians, dim); + stamp = new Uint32Array(N); + + let base = 0; + const sigSample = []; + for (let c = 0; c < meta.numChunks[0]; c++) { + const count = Math.min(meta.chunkSize, meta.numGaussians - c * meta.chunkSize); + const pcd = pool.acquire('position', meta.layouts.position, count); + const gcd = pool.acquire('geometric', meta.layouts.geometric, count); + const ccd = pool.acquire('color', meta.layouts.color, count); + await src.read({ chunkIndex: c, position: pcd, geometric: gcd, color: ccd }); + const p = new Float32Array(pcd.data, 0, count * 3); + const g = new Float32Array(gcd.data, 0, count * 8); + const col = new Float32Array(ccd.data, 0, count * dim); + for (let i = 0; i < count; i++) { + initRow(base + i, p[i * 3], p[i * 3 + 1], p[i * 3 + 2], g, i * 8, col, i * dim); + if ((base + i) % 991 === 0) { + sigSample.push(Math.exp((g[i * 8 + 4] + g[i * 8 + 5] + g[i * 8 + 6]) / 3)); + } + } + base += count; + pcd.release(); gcd.release(); ccd.release(); + if ((c + 1) % 200 === 0) log(` loaded ${base}/${meta.numGaussians}`); + } + await src.close(); + sigSample.sort((x, y) => x - y); + const q = (f) => sigSample[Math.min(sigSample.length - 1, (sigSample.length * f) | 0)]; + log(`load complete (${base} rows) · σ_gm p10 ${q(0.1).toExponential(2)} · median ${q(0.5).toExponential(2)} · p90 ${q(0.9).toExponential(2)}`); +}; + +// Exact global KNN via the production forest (single part at crop scale). +const buildKnn = async () => { + const k = Math.min(KNN_K, Math.max(1, N - 1)); + log(`KNN: forest query (k=${k})`); + + const ids = new Uint32Array(N); + const aabb = new Float32Array([Infinity, Infinity, Infinity, -Infinity, -Infinity, -Infinity]); + for (let i = 0; i < N; i++) { + ids[i] = i; + aabb[0] = Math.min(aabb[0], px[i]); + aabb[1] = Math.min(aabb[1], py[i]); + aabb[2] = Math.min(aabb[2], pz[i]); + aabb[3] = Math.max(aabb[3], px[i]); + aabb[4] = Math.max(aabb[4], py[i]); + aabb[5] = Math.max(aabb[5], pz[i]); + } + const part = { ...buildFlatKdTree(px, py, pz), aabb }; // splat ids already global (identity) + const queryPos = new Float32Array(N * 3); + for (let i = 0; i < N; i++) { + queryPos[i * 3] = px[i]; + queryPos[i * 3 + 1] = py[i]; + queryPos[i * 3 + 2] = pz[i]; + } + const out = new Uint32Array(N * k); + knnForestQuery([part], queryPos, ids, N, k, out); + for (let g = 0; g < N; g++) { + for (let s = 0; s < k; s++) knn[g * KNN_K + s] = out[g * k + s]; + } + log(' KNN done'); +}; + +// --------------------------------------------------------------------------- +// Selftest: engine math vs the library (pair costs, moment composition, +// emission parameters) on synthetic data. +// --------------------------------------------------------------------------- + +const selftest = () => { + const n = 400, dim = 3; + const rand = (() => { let t = 12345; return () => { t = (t * 1103515245 + 12345) & 0x7fffffff; return t / 0x7fffffff; }; })(); + + const view = { + pos: new Float32Array(n * 3), + geo: new Float32Array(n * 8), + color: new Float32Array(n * dim), + colorDim: dim + }; + for (let i = 0; i < n; i++) { + view.pos[i * 3] = rand() * 4; view.pos[i * 3 + 1] = rand() * 4; view.pos[i * 3 + 2] = rand() * 4; + const qw = rand() - 0.5, qx = rand() - 0.5, qy = rand() - 0.5, qz = rand() - 0.5; + view.geo[i * 8] = qw; view.geo[i * 8 + 1] = qx; view.geo[i * 8 + 2] = qy; view.geo[i * 8 + 3] = qz; + view.geo[i * 8 + 4] = -2.5 + rand() * 2; + view.geo[i * 8 + 5] = -2.5 + rand() * 2; + view.geo[i * 8 + 6] = -2.5 + rand() * 2; + view.geo[i * 8 + 7] = -1 + rand() * 3; + view.color[i * dim] = rand() * 2 - 1; + view.color[i * dim + 1] = rand() * 2 - 1; + view.color[i * dim + 2] = rand() * 2 - 1; + } + + initEngine(n, dim); + stamp = new Uint32Array(n); + for (let i = 0; i < n; i++) { + initRow(i, view.pos[i * 3], view.pos[i * 3 + 1], view.pos[i * 3 + 2], view.geo, i * 8, view.color, i * dim); + } + // Brute-force KNN. + for (let i = 0; i < n; i++) { + const d2s = []; + for (let j = 0; j < n; j++) { + if (j === i) continue; + const dx = px[i] - px[j], dy = py[i] - py[j], dz = pz[i] - pz[j]; + d2s.push([dx * dx + dy * dy + dz * dz, j]); + } + d2s.sort((a, b) => a[0] - b[0]); + for (let s = 0; s < KNN_K; s++) knn[i * KNN_K + s] = d2s[s][1]; + } + + // 1. Singleton pair cost parity vs the library edge cost. + const cache = new Float32Array(n * CACHE_STRIDE); + buildSplatCache(view, cache); + let worst = 0; + for (let t = 0; t < 200; t++) { + const i = (rand() * n) | 0; + const j = knn[i * KNN_K + ((rand() * KNN_K) | 0)]; + const mine = evalMerge(i, j); + const lib = computeEdgeCost(cache, i, j); + const rel = Math.abs(mine - lib) / Math.max(1e-12, Math.abs(lib)); + if (rel > worst) worst = rel; + } + console.log(`selftest 1 — pair-cost parity vs lib: worst rel diff ${worst.toExponential(2)} ${worst < 1e-3 ? 'PASS' : 'FAIL'}`); + if (worst >= 1e-3) process.exit(1); + + // 2. Greedy run to 50%, then compare cluster moments vs n-ary mergeGroup. + heapInit(2 * n); + runGreedy([n >> 1], () => {}, 1e9); + const scratch = createMergeScratch(); + const out = { pos: new Float64Array(3), geo: new Float64Array(8), color: new Float64Array(dim) }; + let checked = 0, worstPos = 0, worstScale = 0, worstAlpha = 0, worstColor = 0, worstQuat = 1; + for (let i = 0; i < n && checked < 40; i++) { + if (parent[i] !== i || ufsize[i] < 2) continue; + const cnt = gatherMembers(i); + const members = Array.from(gatherBuf.subarray(0, cnt)); + mergeGroup(view, members, cnt, out, scratch); + + worstPos = Math.max(worstPos, Math.abs(out.pos[0] - mx[i]), Math.abs(out.pos[1] - my[i]), Math.abs(out.pos[2] - mz[i])); + + // Rebuild my emission params for this cluster. + const i6 = i * 6, iw = 1 / W[i]; + const Sig = new Float64Array([ + M2[i6] * iw + EPS_COV, M2[i6 + 1] * iw, M2[i6 + 2] * iw, + M2[i6 + 1] * iw, M2[i6 + 3] * iw + EPS_COV, M2[i6 + 4] * iw, + M2[i6 + 2] * iw, M2[i6 + 4] * iw, M2[i6 + 5] * iw + EPS_COV + ]); + const eigA = new Float64Array(9), eigV = new Float64Array(9); + eigenSymmetric3x3(Sig, eigA, eigV); + const evs = [eigA[0], eigA[4], eigA[8]].sort((x, y) => y - x); + const myScales = evs.map(v => Math.log(Math.sqrt(Math.max(v, 1e-18)))); + const libScales = [out.geo[4], out.geo[5], out.geo[6]]; + for (let s = 0; s < 3; s++) worstScale = Math.max(worstScale, Math.abs(myScales[s] - libScales[s])); + + const myAlpha = Math.min(1, W[i] / Math.max(ellipsoidArea(...evs.map(v => Math.sqrt(Math.max(v, 1e-18)))), 1e-30)); + worstAlpha = Math.max(worstAlpha, Math.abs(myAlpha - sigmoid(out.geo[7]))); + + for (let c = 0; c < dim; c++) { + worstColor = Math.max(worstColor, Math.abs(colorW[i * dim + c] * iw - out.color[c])); + } + checked++; + } + console.log(`selftest 2 — moment parity vs mergeGroup over ${checked} clusters: pos ${worstPos.toExponential(2)} scale ${worstScale.toExponential(2)} alpha ${worstAlpha.toExponential(2)} color ${worstColor.toExponential(2)}`); + const ok2 = worstPos < 1e-5 && worstScale < 1e-4 && worstAlpha < 1e-4 && worstColor < 1e-4; + console.log(ok2 ? 'PASS' : 'FAIL'); + if (!ok2) process.exit(1); + + // 3. Err invariants: every cluster error ≥ 0 (up to float noise), and + // liveCount bookkeeping is exact. + let minErr = Infinity, roots = 0; + for (let i = 0; i < n; i++) { + if (parent[i] !== i) continue; + roots++; + if (Err[i] < minErr) minErr = Err[i]; + } + console.log(`selftest 3 — ${roots} roots (expect ${n >> 1}) · min Err ${minErr.toExponential(2)} ${roots === n >> 1 && minErr > -1e-9 ? 'PASS' : 'FAIL'}`); + if (roots !== n >> 1 || minErr <= -1e-9) process.exit(1); + + console.log('selftest OK'); +}; + +// --------------------------------------------------------------------------- +// CLI +// --------------------------------------------------------------------------- + +const main = async () => { + const argv = process.argv.slice(2); + const argValue = (name, dflt) => { + const i = argv.indexOf(name); + return i >= 0 && i + 1 < argv.length ? argv[i + 1] : dflt; + }; + + if (argv.includes('--selftest')) { + selftest(); + process.exit(0); + } + + const input = argValue('--input'); + const outPrefix = argValue('--out-prefix'); + const halvings = parseInt(argValue('--halvings', '1'), 10); + const mg = argValue('--max-group'); + if (mg) maxGroup = parseInt(mg, 10); + const se = argValue('--size-exponent'); + if (se) sizeExponent = parseFloat(se); + const dl = argValue('--dilate'); + if (dl) dilate = parseFloat(dl); + const cm = argValue('--cost'); + if (cm) { + if (cm !== 'l2' && cm !== 'kl') throw new Error(`--cost must be l2|kl (got ${cm})`); + costMode = cm; + } + const cw = argValue('--color-weight'); + if (cw) colorWeight = parseFloat(cw); + if (argv.includes('--needle-guard')) needleGuard = true; + if (argv.includes('--color-dc-only')) colorDcOnly = true; + if (!input || !outPrefix) { + console.error('usage: decimate-exact.mjs --input --out-prefix --halvings [--max-group ] | --selftest'); + process.exit(1); + } + log(`max-group: ${maxGroup} · size-exponent: ${sizeExponent} · dilate: ${dilate} · cost: ${costMode} · color-weight: ${colorWeight} · needle-guard: ${needleGuard}`); + + await loadPly(input); + await buildKnn(); + + // Targets: successive ceil(count/2) — identical to the production cascade. + const targets = []; + let c = N; + for (let h = 0; h < halvings; h++) { + c = c - Math.floor(c / 2); + targets.push(c); + } + heapInit(Math.ceil(N * 1.25)); + runGreedy(targets, (ti) => emitPly(`${outPrefix}${ti + 1}.ply`)); + + log('done'); + process.exit(0); +}; + +main().catch((e) => { + console.error(e); + process.exit(1); +}); diff --git a/tools/decimate-parity.mjs b/tools/decimate-parity.mjs new file mode 100644 index 00000000..5cc445f6 --- /dev/null +++ b/tools/decimate-parity.mjs @@ -0,0 +1,196 @@ +#!/usr/bin/env node +/** + * Output-parity check for `--decimate-uniform` against a reference build. + * + * The uniform decimator is bit-for-bit output-compatible with the 3.1.x + * release, which is what makes it a usable reference baseline (see + * src/lib/decimate-uniform/README.md and the `old` column in + * scenes/DECIMATION-RESULTS.md). This script is how that claim is checked: + * run N chained halvings through a reference binary's `--decimate` and this + * working tree's `--decimate-uniform`, compare the outputs byte for byte, and + * report PSNR for both against the undecimated source. + * + * Exits non-zero if any level differs, so it can gate a change. + * + * Prerequisites: + * - npm run build (imports WebPCodec from ../dist) + * - a reference binary on PATH, or --ref . The reference must predate + * the split, i.e. its `--decimate` IS the uniform algorithm. + * + * Usage: + * node tools/decimate-parity.mjs [sky|snow] [options] + * + * --ref reference binary (default: splat-transform on PATH) + * --input override the preset's scene (cameras still come from it) + * --halvings chained 50% levels (default: 6) + * --out working directory (default: scenes/parity/) + * --skip-render byte comparison only; no renders, no PSNR + * + * The preset scenes are the two the study uses. They live under scenes/, which + * is gitignored — regenerate them with tools/frustum-cull.mjs if absent. + */ +import { execFileSync } from 'node:child_process'; +import { existsSync, mkdirSync, readFileSync } from 'node:fs'; +import { createHash } from 'node:crypto'; +import { dirname, resolve } from 'node:path'; +import { fileURLToPath } from 'node:url'; + +import { WebPCodec } from '../dist/index.mjs'; + +const ROOT = resolve(dirname(fileURLToPath(import.meta.url)), '..'); +const NODE = process.execPath; +const CLI = `${ROOT}/bin/cli.mjs`; + +const SCENES = { + sky: { + source: `${ROOT}/scenes/fr-sky.ply`, + P: [-5.632, 0.692, -2.550], + T: [-1.578, 1.265, -1.143] + }, + snow: { + source: `${ROOT}/scenes/fr-snow.ply`, + P: [-143.774, 29.955, -27.240], + T: [-108.901, 10.650, -46.631] + } +}; +const ANGLES = [-10, 0, 10]; + +const argv = process.argv.slice(2); +const flag = (name, fallback) => { + const i = argv.indexOf(`--${name}`); + return i >= 0 && argv[i + 1] !== undefined ? argv[i + 1] : fallback; +}; +const which = argv[0] && !argv[0].startsWith('--') ? argv[0] : 'sky'; +if (!SCENES[which]) { + console.error(`unknown scene "${which}" — expected one of: ${Object.keys(SCENES).join(', ')}`); + process.exit(2); +} + +const { P, T } = SCENES[which]; +const source = resolve(flag('input', SCENES[which].source)); +const ref = flag('ref', 'splat-transform'); +const halvings = parseInt(flag('halvings', '6'), 10); +const outDir = resolve(flag('out', `${ROOT}/scenes/parity/${which}`)); +const skipRender = argv.includes('--skip-render'); + +if (!existsSync(source)) { + console.error(`scene not found: ${source}\nscenes/ is gitignored — regenerate with tools/frustum-cull.mjs, or pass --input.`); + process.exit(2); +} +try { + execFileSync(ref, ['--version'], { stdio: 'ignore' }); +} catch { + console.error(`reference binary not runnable: ${ref}\nInstall the reference release or pass --ref .`); + process.exit(2); +} +mkdirSync(outDir, { recursive: true }); + +const sha256 = path => createHash('sha256').update(readFileSync(path)).digest('hex'); +const vertexCount = (path) => { + const head = readFileSync(path).subarray(0, 2048).toString('ascii'); + const m = head.match(/element vertex (\d+)/); + return m ? parseInt(m[1], 10) : -1; +}; + +// Chain `halvings` 50% steps, timing each. `ref` runs its own --decimate (the +// uniform algorithm at that revision); the working tree runs the explicit flag. +const chain = (label, argsFor) => { + const paths = []; + const secs = []; + let input = source; + for (let l = 1; l <= halvings; l++) { + const out = `${outDir}/${which}.${label}-${l}.ply`; + if (!existsSync(out)) { + const started = Date.now(); + const [bin, ...pre] = argsFor.bin; + execFileSync(bin, [...pre, input, ...argsFor.flags, '50%', out, '-w'], { stdio: 'ignore' }); + secs.push((Date.now() - started) / 1000); + } else { + secs.push(NaN); // cached from an earlier run + } + paths.push(out); + input = out; + } + return { paths, secs }; +}; + +console.log(`\n=== ${which}: ${ref} --decimate vs this tree --decimate-uniform ===`); +console.log(`source ${source} (${vertexCount(source)} splats), ${halvings} halvings\n`); + +const a = chain('ref', { bin: [ref], flags: ['--decimate'] }); +const b = chain('uni', { bin: [NODE, CLI], flags: ['--decimate-uniform'] }); + +let mismatches = 0; +console.log('level count ref s uni s identical'); +for (let l = 0; l < halvings; l++) { + const identical = sha256(a.paths[l]) === sha256(b.paths[l]); + if (!identical) mismatches++; + const ca = vertexCount(a.paths[l]); + const cb = vertexCount(b.paths[l]); + console.log( + `L${l + 1} ${String(ca).padStart(10)}${ca === cb ? '' : ` (uni ${cb})`} ` + + `${a.secs[l].toFixed(1).padStart(7)} ${b.secs[l].toFixed(1).padStart(7)} ${identical ? 'YES' : 'NO'}` + ); +} +const total = s => s.reduce((x, y) => x + (Number.isNaN(y) ? 0 : y), 0).toFixed(1); +console.log(`\ntotal cascade: ref ${total(a.secs)}s uniform ${total(b.secs)}s`); + +if (!skipRender) { + const codec = await WebPCodec.create(); + const render = (input, out, pos) => { + if (existsSync(out)) return; + execFileSync(NODE, [CLI, input, out, '--camera-pos', pos.join(','), '--camera-target', T.join(','), '-w'], { stdio: 'ignore' }); + }; + const rgbaOf = async (p) => { + const { rgba } = await codec.decodeRGBA(new Uint8Array(readFileSync(p))); + return rgba; + }; + const psnr = (x, y) => { + let sum = 0, n = 0; + for (let i = 0; i < x.length; i += 4) { + for (let c = 0; c < 3; c++) { + const d = x[i + c] - y[i + c]; + sum += d * d; + n++; + } + } + const mse = sum / n; + return mse === 0 ? Infinity : 10 * Math.log10(255 * 255 / mse); + }; + + const oy = P[1] - T[1]; + const r = Math.hypot(P[0] - T[0], P[2] - T[2]); + const theta0 = Math.atan2(P[2] - T[2], P[0] - T[0]); + const acc = { ref: new Array(halvings).fill(0), uni: new Array(halvings).fill(0) }; + + for (const ang of ANGLES) { + const th = theta0 + ang * Math.PI / 180; + const pos = [T[0] + r * Math.cos(th), T[1] + oy, T[2] + r * Math.sin(th)]; + const srcRender = `${outDir}/${which}.src_${ang}.webp`; + render(source, srcRender, pos); + const srcRgba = await rgbaOf(srcRender); + for (const [label, chained] of [['ref', a], ['uni', b]]) { + for (let l = 0; l < halvings; l++) { + const out = `${outDir}/${which}.${label}-${l + 1}_${ang}.webp`; + render(chained.paths[l], out, pos); + acc[label][l] += psnr(srcRgba, await rgbaOf(out)) / ANGLES.length; + } + } + } + + console.log(`\nPSNR dB vs source (mean of ${ANGLES.length} poses)\n`); + console.log('level ref uni delta'); + for (let l = 0; l < halvings; l++) { + const d = acc.uni[l] - acc.ref[l]; + console.log( + `L${l + 1} ${acc.ref[l].toFixed(2).padStart(7)} ${acc.uni[l].toFixed(2).padStart(7)} ` + + `${(d >= 0 ? '+' : '') + d.toFixed(2)}` + ); + } +} + +console.log(mismatches === 0 ? + `\nPASS — all ${halvings} levels byte-identical to ${ref}` : + `\nFAIL — ${mismatches} of ${halvings} level(s) differ`); +console.log(`artifacts: ${outDir}`); +process.exit(mismatches === 0 ? 0 : 1); diff --git a/tools/frustum-cull.mjs b/tools/frustum-cull.mjs new file mode 100644 index 00000000..ccd68053 --- /dev/null +++ b/tools/frustum-cull.mjs @@ -0,0 +1,160 @@ +#!/usr/bin/env node +/** + * Frustum-cull a binary splat PLY: keep only gaussians whose center (with a + * per-splat 2σ margin) lies inside the view frustum of the given render + * camera. Rows are passed through byte-for-byte (single streaming pass), so + * the output is the same PLY with fewer vertices. + * + * The camera is specified in RENDER space (the coordinates used by the CLI's + * --camera-pos/--camera-target). Raw PLY positions map to render space via + * Transform.PLY = rot-z 180°: (x, y, z) → (−x, −y, z). + * + * Usage: + * node tools/frustum-cull.mjs --input in.ply --output out.ply \ + * --camera-pos x,y,z --camera-target x,y,z \ + * [--fov 60] [--aspect 1.7778] [--widen-deg 12] [--near 0.2] + */ +import { openSync, readSync, writeSync, closeSync } from 'node:fs'; + +const argv = process.argv.slice(2); +const argValue = (name, dflt) => { + const i = argv.indexOf(name); + return i >= 0 && i + 1 < argv.length ? argv[i + 1] : dflt; +}; +const vec = (s) => s.split(',').map(Number); + +const input = argValue('--input'); +const output = argValue('--output'); +const P = vec(argValue('--camera-pos')); +const T = vec(argValue('--camera-target')); +const fovY = (parseFloat(argValue('--fov', '60')) * Math.PI) / 180; +const aspect = parseFloat(argValue('--aspect', String(1280 / 720))); +const widen = (parseFloat(argValue('--widen-deg', '12')) * Math.PI) / 180; +const near = parseFloat(argValue('--near', '0.2')); +if (!input || !output || P.length !== 3 || T.length !== 3) { + console.error('usage: frustum-cull.mjs --input in.ply --output out.ply --camera-pos x,y,z --camera-target x,y,z [--fov 60] [--aspect 1.7778] [--widen-deg 12] [--near 0.2]'); + process.exit(1); +} + +// Camera basis in render space (up = +y, standard lookAt). +const sub = (a, b) => [a[0] - b[0], a[1] - b[1], a[2] - b[2]]; +const norm = (a) => { const l = Math.hypot(...a); return [a[0] / l, a[1] / l, a[2] / l]; }; +const cross = (a, b) => [a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], a[0] * b[1] - a[1] * b[0]]; +const dot = (a, b) => a[0] * b[0] + a[1] * b[1] + a[2] * b[2]; + +const fwd = norm(sub(T, P)); +let right = cross(fwd, [0, 1, 0]); +const rl = Math.hypot(...right); +right = rl > 1e-6 ? [right[0] / rl, right[1] / rl, right[2] / rl] : [1, 0, 0]; +const up = cross(right, fwd); + +const tanY = Math.tan(fovY / 2 + widen); +const tanX = Math.tan(Math.atan(Math.tan(fovY / 2) * aspect) + widen); +const slopeNormX = Math.sqrt(1 + tanX * tanX); +const slopeNormY = Math.sqrt(1 + tanY * tanY); + +// ---- Parse header. +const fdIn = openSync(input, 'r'); +const head = Buffer.alloc(65536); +readSync(fdIn, head, 0, head.length, 0); +const headText = head.toString('latin1'); +const endIdx = headText.indexOf('end_header\n'); +if (endIdx < 0) throw new Error('no end_header'); +const headerLen = endIdx + 'end_header\n'.length; +const lines = headText.slice(0, endIdx).split('\n'); +let count = 0; +const props = []; +for (const l of lines) { + const mv = /^element vertex (\d+)/.exec(l); + if (mv) count = parseInt(mv[1], 10); + const mp = /^property (\w+) (\S+)/.exec(l); + if (mp && count > 0) props.push({ type: mp[1], name: mp[2] }); +} +const typeSize = { float: 4, float32: 4, double: 8, uchar: 1, uint8: 1, int: 4, uint: 4, uint32: 4, short: 2, ushort: 2 }; +let stride = 0; +const off = {}; +for (const p of props) { + off[p.name] = stride; + if (typeSize[p.type] !== 4 || (p.type !== 'float' && p.type !== 'float32')) { + if (['x', 'y', 'z', 'scale_0', 'scale_1', 'scale_2'].includes(p.name)) { + throw new Error(`property ${p.name} must be float32 (got ${p.type})`); + } + } + stride += typeSize[p.type]; +} +for (const n of ['x', 'y', 'z', 'scale_0', 'scale_1', 'scale_2']) { + if (!(n in off)) throw new Error(`missing property ${n}`); +} +console.log(`${input}: ${count} vertices, stride ${stride}`); + +// ---- Stream, test, write survivors. +const ROWS = 1 << 16; +const inBuf = Buffer.alloc(ROWS * stride); +const outBuf = Buffer.alloc(ROWS * stride); + +// First pass counts survivors while writing rows to a temp offset — instead, +// write rows after a placeholder header, then rewrite the header space with +// exact padding. Simpler: two passes would re-read 3GB; instead write header +// with the count later using a fixed-width count field. +const countField = String(count); // survivors <= count, pad to same width +const headerOut = headText.slice(0, endIdx).replace(/^element vertex \d+$/m, `element vertex COUNT_PLACEHOLDER`) + 'end_header\n'; + +const fdOut = openSync(output, 'w'); +// Reserve header space: replace placeholder with padded count at the end. +const headerTemplate = headerOut.replace('COUNT_PLACEHOLDER', countField); // max width +writeSync(fdOut, Buffer.from(headerTemplate, 'latin1')); +const headerOutLen = Buffer.byteLength(headerTemplate, 'latin1'); + +let kept = 0, read = 0, outRows = 0; +let inPos = headerLen; +while (read < count) { + const rows = Math.min(ROWS, count - read); + const bytes = rows * stride; + let got = 0; + while (got < bytes) { + const n = readSync(fdIn, inBuf, got, bytes - got, inPos + got); + if (n <= 0) throw new Error('short read'); + got += n; + } + inPos += bytes; + outRows = 0; + for (let r = 0; r < rows; r++) { + const base = r * stride; + // Raw → render space: rot-z 180°. + const qx = -inBuf.readFloatLE(base + off.x); + const qy = -inBuf.readFloatLE(base + off.y); + const qz = inBuf.readFloatLE(base + off.z); + const s0 = inBuf.readFloatLE(base + off.scale_0); + const s1 = inBuf.readFloatLE(base + off.scale_1); + const s2 = inBuf.readFloatLE(base + off.scale_2); + const sigma = Math.exp(Math.max(s0, s1, s2)); + const margin = 2 * (Number.isFinite(sigma) ? sigma : 0); + + const dx = qx - P[0], dy = qy - P[1], dz = qz - P[2]; + const zc = dx * fwd[0] + dy * fwd[1] + dz * fwd[2]; + if (zc < near - margin) continue; + const xc = dx * right[0] + dy * right[1] + dz * right[2]; + if (Math.abs(xc) > zc * tanX + margin * slopeNormX) continue; + const yc = dx * up[0] + dy * up[1] + dz * up[2]; + if (Math.abs(yc) > zc * tanY + margin * slopeNormY) continue; + + inBuf.copy(outBuf, outRows * stride, base, base + stride); + outRows++; + } + if (outRows > 0) writeSync(fdOut, outBuf, 0, outRows * stride); + kept += outRows; + read += rows; + if (read % (1 << 22) < ROWS) console.log(` ${read}/${count} scanned, ${kept} kept`); +} +closeSync(fdIn); + +// Rewrite header with the real count, zero-padded to the reserved width +// (leading zeros parse cleanly everywhere; trailing spaces may not). +const headerFinal = headerTemplate.replace( + `element vertex ${countField}`, + `element vertex ${String(kept).padStart(countField.length, '0')}` +); +if (Buffer.byteLength(headerFinal, 'latin1') !== headerOutLen) throw new Error('header size drift'); +writeSync(fdOut, Buffer.from(headerFinal, 'latin1'), 0, headerOutLen, 0); +closeSync(fdOut); +console.log(`${output}: kept ${kept} / ${count} (${((kept / count) * 100).toFixed(1)}%)`); diff --git a/tools/sweep-fr.mjs b/tools/sweep-fr.mjs new file mode 100644 index 00000000..3eccadf7 --- /dev/null +++ b/tools/sweep-fr.mjs @@ -0,0 +1,113 @@ +#!/usr/bin/env node +// Full-scale confirmation on frustum-culled scenes: old vs new4 vs b4, +// 6 levels, poses −10°/0°/+10° (inside the +12°-widened culling frustum). +// Renders under scenes/sweep-fr//; PSNR vs the culled source. +// +// Prerequisite: npm run build (imports WebPCodec from ../dist). +// Usage: node tools/sweep-fr.mjs [sky|snow] +import { execFileSync } from 'node:child_process'; +import { readFileSync, mkdirSync, existsSync } from 'node:fs'; +import { dirname, resolve } from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { WebPCodec } from '../dist/index.mjs'; + +const ROOT = resolve(dirname(fileURLToPath(import.meta.url)), '..'); +const NODE = process.execPath; +const CLI = `${ROOT}/bin/cli.mjs`; + +const L = [1, 2, 3, 4, 5, 6]; +const SCENES = { + sky: { + source: `${ROOT}/scenes/fr-sky.ply`, + P: [-5.632, 0.692, -2.550], + T: [-1.578, 1.265, -1.143], + methods: { + old: L.map(l => `${ROOT}/scenes/fr-sky.old-${l}.ply`), + new4: L.map(l => `${ROOT}/scenes/fr-sky.new4-${l}.ply`), + b4: L.map(l => `${ROOT}/scenes/fr-sky.b4-${l}1.ply`), + cw6: L.map(l => `${ROOT}/scenes/fr-sky.cw6-${l}1.ply`), + new4r: L.map(l => `${ROOT}/scenes/fr-sky.new4r-${l}.ply`) + } + }, + snow: { + source: `${ROOT}/scenes/fr-snow.ply`, + P: [-143.774, 29.955, -27.240], + T: [-108.901, 10.650, -46.631], + methods: { + old: L.map(l => `${ROOT}/scenes/fr-snow.old-${l}.ply`), + new4: L.map(l => `${ROOT}/scenes/fr-snow.new4-${l}.ply`), + b4: L.map(l => `${ROOT}/scenes/fr-snow.b4-${l}1.ply`), + cw6: L.map(l => `${ROOT}/scenes/fr-snow.cw6-${l}1.ply`), + new4r: L.map(l => `${ROOT}/scenes/fr-snow.new4r-${l}.ply`) + } + } +}; +const ANGLES = [-10, 0, 10]; + +const which = process.argv[2] || 'sky'; +const cfg = SCENES[which]; +const OUT = `${ROOT}/scenes/sweep-fr/${which}`; +mkdirSync(OUT, { recursive: true }); + +const { P, T } = cfg; +const oy = P[1] - T[1]; +const r = Math.hypot(P[0] - T[0], P[2] - T[2]); +const theta0 = Math.atan2(P[2] - T[2], P[0] - T[0]); + +const render = (input, out, pos) => { + if (existsSync(out)) return; + execFileSync(NODE, [CLI, input, out, '--camera-pos', pos.join(','), '--camera-target', T.join(',')], { stdio: 'ignore' }); +}; + +const codec = await WebPCodec.create(); +const psnr = (aPath, bPath) => { + const a = codec.decodeRGBA(readFileSync(aPath)); + const b = codec.decodeRGBA(readFileSync(bPath)); + const n = a.width * a.height; + let se = 0; + for (let i = 0; i < n; i++) { + for (let c = 0; c < 3; c++) { + const d = a.rgba[i * 4 + c] - b.rgba[i * 4 + c]; + se += d * d; + } + } + const mse = se / (n * 3); + return mse === 0 ? Infinity : 10 * Math.log10((255 * 255) / mse); +}; + +const methods = Object.keys(cfg.methods); +const perPose = {}; +for (const ang of ANGLES) { + const th = theta0 + (ang * Math.PI) / 180; + const pos = [T[0] + r * Math.cos(th), T[1] + oy, T[2] + r * Math.sin(th)]; + const ref = `${OUT}/src_${ang}.webp`; + render(cfg.source, ref, pos); + perPose[ang] = {}; + for (const l of L) { + perPose[ang][l] = {}; + for (const m of methods) { + const ply = cfg.methods[m][l - 1]; + if (!existsSync(ply)) { perPose[ang][l][m] = null; continue; } + const out = `${OUT}/${m}${l}_${ang}.webp`; + render(ply, out, pos); + perPose[ang][l][m] = psnr(ref, out); + } + } +} + +console.log(`\n=== fr-${which} === PSNR dB vs culled source (mean over ${ANGLES.length} poses; pose 0 = user camera)`); +console.log(` level ${methods.map(m => m.padStart(8)).join(' ')}`); +for (const l of L) { + const cells = methods.map((m) => { + const vals = ANGLES.map(a => perPose[a][l][m]).filter(v => v != null); + return (vals.length ? (vals.reduce((x, y) => x + y, 0) / vals.length).toFixed(2) : ' —').padStart(8); + }); + console.log(` L${l} ${cells.join(' ')}`); +} +for (const ang of ANGLES) { + console.log(`pose ${ang}°`); + for (const l of L) { + const cells = methods.map(m => (perPose[ang][l][m] == null ? ' —' : perPose[ang][l][m].toFixed(2)).padStart(8)); + console.log(` L${l} ${cells.join(' ')}`); + } +}