gaussian-splat-lite 1.5.1 → 1.5.2

package/bin/rad/splats.js

@@ -12,7 +12,7 @@ import {
   quaternion,
 } from "./math.js";
 import { Progress } from "./progress.js";
-import { IndexArray, PagedArray, Permutation } from "./storage.js";
+import { IndexArray, PAGE_BYTES, PagedArray, Permutation } from "./storage.js";
 
 export const SH_DIMS = [9, 15, 21];
 
@@ -75,8 +75,15 @@ export class Splats {
     // Each splat's merge weight, area × opacity, which merge() reads instead
     // of two records.
     this.weight = new PagedArray(workspace, Float64Array);
+    // Each splat's feature size, for the passes that read nothing else of a
+    // record, and the bounds of the centers; see reorder().
+    this.sizes = new PagedArray(workspace, Float32Array);
+    this.bounds = undefined;
     // The two child weights of each merged splat, for mergeSh().
     this.blend = new PagedArray(workspace, Float64Array);
+    // Each decoded splat's row in the SH columns: as decoded, and from
+    // mergeSh() on, in the order merges read them; see shRow().
+    this.shOrder = undefined;
     this.merged = new Float64Array(16);
     this.bits = new Uint16Array(0);
     this.merger = new Merger();
@@ -133,9 +140,6 @@ export class Splats {
   read(index, output) {
     return this.records.read(index * 16, output, 16);
   }
-  feature(index) {
-    return this.records.get(index * 16 + 14);
-  }
 
   // Skips empty splats and counts invalid ones in `invalid`: NaN or infinite
   // values, including a zero rotation, which the decoder normalizes to NaN.
@@ -192,6 +196,8 @@ export class Splats {
             this.sh[band].write(count * dims, rows[band]);
           }
         }
+        // The record's page is cached.
+        this.sizes.set(count, this.records.get(count * 16 + 14));
         count++;
       }
       progress.update(base + size);
@@ -226,6 +232,7 @@ export class Splats {
     const index = this.length++;
     this.records.write(index * 16, out);
     this.weight.set(index, out[15] * out[3]);
+    this.sizes.set(index, out[14]);
     this.children.push(a);
     this.children.push(b);
     // SH coefficients do not affect merging; mergeSh() blends them later.
@@ -236,46 +243,105 @@ export class Splats {
     return index;
   }
 
-  // Blends the SH coefficients of the splats merge() created from `initial`
-  // on, in creation order, so children are done before their parents. Going
-  // band by band keeps the random reads of children within one column.
-  // Workers blend batches of splats whose children precede the batch; the
-  // rest wait for their children in the batch.
+  // Blend each SH band in creation order, children before parents. Reorder
+  // leaves for sequential reads; only merged children need random reads.
+  // Batches use at most half the cache for child pages. Workers handle
+  // earlier children; children within the batch are blended below in order.
   mergeSh(initial, pool) {
+    if (!this.maxSh) return;
     const { nodes, entries, later, blends, blended, results } = pool;
+    // Each leaf occurs once, as a merge child or the lone root. Replace its
+    // decoded SH row with its position in merge-read order.
+    const reads = new IndexArray(this.workspace);
+    const collecting = new Progress("Collecting SH order", this.length);
+    for (let k = 0; k < (this.length - initial) * 2; k++) {
+      const child = this.children.get(k);
+      if (child < initial) {
+        reads.push(this.shOrder.get(child));
+        this.shOrder.set(child, reads.length - 1);
+      }
+      if ((k & 4095) === 0) collecting.update(k / 2);
+    }
+    if (this.length === initial) {
+      reads.push(this.shOrder.get(0));
+      this.shOrder.set(0, 0);
+    }
+    collecting.finish();
+    const leafRows = new Permutation(
+      this.workspace,
+      reads,
+      initial,
+      "Preparing SH order",
+    );
+    reads.dispose();
     // Splats per batch.
     const size = later.length;
     const rows = [this.shLeft, this.shRight];
     for (let band = 0; band < this.maxSh; band++) {
+      const dims = SH_DIMS[band];
+      const column = leafRows.apply(
+        this.sh[band],
+        dims,
+        pool,
+        false,
+        `Reordering SH${band + 1}`,
+      );
+      this.sh[band] = column;
       const progress = new Progress(
         `Blending SH${band + 1}`,
         this.length - initial,
       );
-      const dims = SH_DIMS[band];
-      const column = this.sh[band].share();
-      for (let start = initial; start < this.length; start += size) {
-        const count = Math.min(size, this.length - start);
+      // The batch that last counted each page of the column.
+      const counted = new Float64Array(
+        Math.ceil((this.length * dims) / column.items),
+      );
+      let batch = 0;
+      // The row of the next decoded splat that a merge reads.
+      let next = 0;
+      for (let start = initial; start < this.length; ) {
+        const limit = Math.min(size, this.length - start);
+        const pages = Math.max(
+          1,
+          Math.floor(this.workspace.capacity / PAGE_BYTES / 2),
+        );
+        batch++;
+        let taken = 0;
+        let count = 0;
         // Load the pages of the children's rows for the workers. Consecutive
         // splats' children often share pages; a page evicted meanwhile is a
         // miss that the main thread blends.
         let left = -1;
         let right = -1;
-        for (let i = 0; i < count; i++) {
-          const a = this.children.get((start + i - initial) * 2);
-          const b = this.children.get((start + i - initial) * 2 + 1);
-          nodes[i] = a;
-          entries[i] = b;
-          blends[i * 2] = this.blend.get((start + i - initial) * 2);
-          blends[i * 2 + 1] = this.blend.get((start + i - initial) * 2 + 1);
-          later[i] = Number(a >= start || b >= start);
+        while (count < limit) {
+          const k = (start + count - initial) * 2;
+          const a = this.children.get(k);
+          const b = this.children.get(k + 1);
+          const ra = a < initial ? next : a;
+          const rb = b < initial ? next + Number(a < initial) : b;
+          const pa = Math.floor((ra * dims) / column.items);
+          const pb = Math.floor((rb * dims) / column.items);
+          let added = Number(ra < start && counted[pa] !== batch);
+          if (rb < start && counted[pb] !== batch && !(added && pb === pa))
+            added++;
+          if (count && taken + added > pages) break;
+          taken += added;
+          if (ra < start) counted[pa] = batch;
+          if (rb < start) counted[pb] = batch;
+          next += Number(a < initial) + Number(b < initial);
+          const i = count++;
+          nodes[i] = ra;
+          entries[i] = rb;
+          blends[i * 2] = this.blend.get(k);
+          blends[i * 2 + 1] = this.blend.get(k + 1);
+          later[i] = Number(ra >= start || rb >= start);
           if (later[i]) continue;
-          if (Math.floor((a * dims) / column.items) !== left) {
-            left = Math.floor((a * dims) / column.items);
-            column.page(a * dims);
+          if (pa !== left) {
+            left = pa;
+            column.page(ra * dims);
           }
-          if (Math.floor((b * dims) / column.items) !== right) {
-            right = Math.floor((b * dims) / column.items);
-            column.page(b * dims);
+          if (pb !== right) {
+            right = pb;
+            column.page(rb * dims);
           }
         }
         pool.blend(column, count, dims);
@@ -297,14 +363,19 @@ export class Splats {
             );
         }
         column.write(start * dims, blended, count * dims);
-        progress.update(start + count - initial);
+        start += count;
+        progress.update(start - initial);
       }
       progress.finish();
     }
+    leafRows.dispose();
     this.blend.dispose();
   }
 
-  // Moves row indices[i] of the records and SH to row i.
+  // Moves row indices[i] of the records to row i, and keeps `indices` for
+  // the SH rows, which move once merging has built the tree; see mergeSh().
+  // The splats' sizes and weights, and the bounds of their centers, come
+  // from the records while each window of them is cached.
   reorder(indices, pool) {
     const permutation = new Permutation(
       this.workspace,
@@ -312,23 +383,43 @@ export class Splats {
       this.length,
       "Preparing splat order",
     );
+    this.sizes.dispose();
+    const sizes = this.sizes;
+    const bounds = {
+      min: [0, 1, 2].map(() => Number.POSITIVE_INFINITY),
+      max: [0, 1, 2].map(() => Number.NEGATIVE_INFINITY),
+    };
+    const record = new Float64Array(16);
     this.records = permutation.apply(
       this.records,
       16,
       pool,
       false,
       "Reordering splats",
+      (records, start, end) => {
+        for (let i = start; i < end; i++) {
+          records.read(i * 16, record, 16);
+          for (let d = 0; d < 3; d++) {
+            bounds.min[d] = Math.min(bounds.min[d], record[d]);
+            bounds.max[d] = Math.max(bounds.max[d], record[d]);
+          }
+          this.weight.set(i, record[15] * record[3]);
+          sizes.set(i, record[14]);
+        }
+      },
     );
-    for (let band = 0; band < this.maxSh; band++)
-      this.sh[band] = permutation.apply(
-        this.sh[band],
-        SH_DIMS[band],
-        pool,
-        false,
-        `Reordering SH${band + 1}`,
-      );
+    this.bounds = bounds;
     permutation.dispose();
     this.length = indices.length;
+    if (this.maxSh) this.shOrder = indices;
+    else indices.dispose();
+  }
+
+  // The row of splat `index` in the SH columns after mergeSh(): merged
+  // splats keep theirs, and the decoded ones are in the order merges read
+  // them.
+  shRow(index) {
+    return index < this.shOrder.length ? this.shOrder.get(index) : index;
   }
 
   dispose() {
@@ -339,6 +430,7 @@ export class Splats {
       this.childCount,
       this.children,
       this.weight,
+      this.sizes,
       this.blend,
     ])
       column.dispose();

package/bin/rad/storage.js

@@ -12,6 +12,7 @@ import { join } from "node:path";
 import v8 from "node:v8";
 import { runInNewContext } from "node:vm";
 import { Progress } from "./progress.js";
+import { Writeback } from "./writeback.js";
 
 export const PAGE_BYTES = 65536;
 export const GIB = 1024 ** 3;
@@ -39,7 +40,7 @@ export function collect() {
 // which a SharedArrayBuffer does not on every platform. The cache takes what
 // the budget leaves after the rest of the process, measured as it runs.
 export class Workspace {
-  constructor(budget, parent = tmpdir(), cancelled = null) {
+  constructor(budget, sync, parent = tmpdir(), cancelled = null) {
     this.budget = budget;
     this.parent = parent;
     this.cancelled = cancelled;
@@ -74,6 +75,9 @@ export class Workspace {
     this.nextFile = 0;
     this.readBytes = 0;
     this.writeBytes = 0;
+    // File writes can remain in the system's cache after writeSync returns.
+    // Reserve room for pending writeback inside the budget on every platform.
+    this.writeback = new Writeback(64 * 1024 ** 2, sync, () => this.check());
     // An earlier conversion's arenas return once collected.
     this.settle();
   }
@@ -87,13 +91,28 @@ export class Workspace {
   get capacity() {
     return Math.max(
       this.arenaBytes,
-      this.budget - MARGIN - Math.max(this.outside, this.floor),
+      this.budget -
+        MARGIN -
+        this.writeback.limit -
+        Math.max(this.outside, this.floor),
     );
   }
 
   // Bytes the rest of the process may still grow by.
   get room() {
-    return this.budget - MARGIN - this.arenaBytes - this.outside;
+    return (
+      this.budget -
+      MARGIN -
+      this.writeback.limit -
+      this.arenaBytes -
+      this.outside
+    );
+  }
+
+  // Rebalance the cache every 64 allocations and between streaming passes.
+  adjust() {
+    this.measure();
+    this.shrink();
   }
 
   // Measures the memory of the rest of the process: all of it but the
@@ -117,7 +136,10 @@ export class Workspace {
   expect(bytes) {
     this.measure();
     this.floor = this.outside + bytes;
-    if (this.budget - MARGIN - this.floor < this.arenaBytes)
+    if (
+      this.budget - MARGIN - this.writeback.limit - this.floor <
+      this.arenaBytes
+    )
       throw new Error("Memory budget too small; increase --memory-budget");
     this.shrink();
   }
@@ -161,9 +183,32 @@ export class Workspace {
     const file = this.files.get(fd);
     this.files.delete(fd);
     try {
-      closeSync(fd);
-      rmSync(file);
-    } catch {}
+      this.writeback.synced(fd);
+    } finally {
+      try {
+        closeSync(fd);
+        rmSync(file);
+      } catch {}
+    }
+  }
+
+  // Reserve writeback room before each write, including large output chunks.
+  // Counting whole 64 KiB pages also covers partial pages on systems whose
+  // native pages are larger than 4 KiB.
+  write(fd, bytes, position) {
+    for (let offset = 0; offset < bytes.byteLength; offset += 1024 * 1024) {
+      const part =
+        bytes.byteLength <= 1024 * 1024
+          ? bytes
+          : bytes.subarray(offset, offset + 1024 * 1024);
+      const at = position + offset;
+      const size =
+        Math.ceil(((at % PAGE_BYTES) + part.byteLength) / PAGE_BYTES) *
+        PAGE_BYTES;
+      this.writeback.prepare(size);
+      writeAll(fd, part, at);
+      this.writeback.wrote(fd, size);
+    }
   }
 
   evict() {
@@ -198,10 +243,7 @@ export class Workspace {
   }
 
   allocate(array, index) {
-    if (++this.allocations % 64 === 0) {
-      this.measure();
-      this.shrink();
-    }
+    if (++this.allocations % 64 === 0) this.adjust();
     if (this.used + PAGE_BYTES > this.capacity && this.used) this.evict();
     let slot = this.free.pop();
     if (slot === undefined) {
@@ -278,25 +320,29 @@ export class Workspace {
   }
 
   dispose() {
-    const progress = this.directory
-      ? new Progress("Cleaning temporary files", this.files.size + 1)
-      : undefined;
-    let closed = 0;
-    for (const fd of this.files.keys()) {
-      this.close(fd);
-      progress?.update(++closed);
-    }
-    if (this.directory) removePath(this.directory, { recursive: true });
-    this.directory = undefined;
-    this.slots.length = 0;
-    this.reuse.length = 0;
-    this.free.length = 0;
-    this.arenas.length = 0;
-    progress?.finish();
-    if (this.readBytes || this.writeBytes) {
-      console.log(
-        `Temporary array I/O: ${(this.readBytes / 1048576).toFixed(1)} MiB read, ${(this.writeBytes / 1048576).toFixed(1)} MiB written`,
-      );
+    try {
+      this.writeback.wait(false);
+    } finally {
+      const progress = this.directory
+        ? new Progress("Cleaning temporary files", this.files.size + 1)
+        : undefined;
+      let closed = 0;
+      for (const fd of this.files.keys()) {
+        this.close(fd);
+        progress?.update(++closed);
+      }
+      if (this.directory) removePath(this.directory, { recursive: true });
+      this.directory = undefined;
+      this.slots.length = 0;
+      this.reuse.length = 0;
+      this.free.length = 0;
+      this.arenas.length = 0;
+      progress?.finish();
+      if (this.readBytes || this.writeBytes) {
+        console.log(
+          `Temporary array I/O: ${(this.readBytes / 1048576).toFixed(1)} MiB read, ${(this.writeBytes / 1048576).toFixed(1)} MiB written`,
+        );
+      }
     }
   }
 }
@@ -348,7 +394,7 @@ export class PagedArray {
     const stripe = block % this.stripes;
     const position = Math.floor(block / this.stripes) * PAGE_BYTES;
     this.fds[stripe] ??= this.workspace.file();
-    writeAll(this.fds[stripe], bytes, position);
+    this.workspace.write(this.fds[stripe], bytes, position);
     this.sizes[stripe] = Math.max(
       this.sizes[stripe] ?? 0,
       position + PAGE_BYTES,
@@ -427,6 +473,47 @@ export class PagedArray {
     return output;
   }
 
+  // Read into the same array type, bypassing the cache on misses so a scan
+  // cannot evict cached pages.
+  stream(base, output, count = output.length, at = 0) {
+    const bytes = this.Type.BYTES_PER_ELEMENT;
+    let offset = 0;
+    while (offset < count) {
+      if ((++this.workspace.accesses & 4095) === 0) this.workspace.check();
+      const index = base + offset;
+      const block = Math.floor(index * this.inverse);
+      const within = index & this.mask;
+      const size = Math.min(count - offset, this.items - within);
+      const page = this.pages[block];
+      if (page)
+        output.set(page.values.subarray(within, within + size), at + offset);
+      else {
+        const stripe = block % this.stripes;
+        const position = Math.floor(block / this.stripes) * PAGE_BYTES;
+        const part = new Uint8Array(
+          output.buffer,
+          output.byteOffset + (at + offset) * bytes,
+          size * bytes,
+        );
+        // As in allocate(), a page that no file holds is zeros.
+        if (position >= (this.sizes[stripe] ?? 0)) part.fill(0);
+        else {
+          readAll(this.fds[stripe], part, position + within * bytes);
+          this.workspace.readBytes += part.byteLength;
+        }
+      }
+      offset += size;
+    }
+    return output;
+  }
+
+  // Cache the prefix before `kept`; stream the rest to preserve it for reuse.
+  scan(base, output, count, kept) {
+    const cached = Math.max(0, Math.min(count, kept - base));
+    this.read(base, output, cached);
+    return this.stream(base + cached, output, count - cached, cached);
+  }
+
   write(base, values, count = values.length) {
     let offset = 0;
     while (offset < count) {
@@ -657,8 +744,17 @@ export class Permutation {
   // rows into it; rows across pages are left to the main thread. The output
   // spreads over as many files as `source`.
   // With `words`, rows are copied as 32-bit words, which keeps any bits of
-  // floats, and the output has word views; see PagedArray.at().
-  apply(source, stride, pool, words = false, label = undefined) {
+  // floats, and the output has word views; see PagedArray.at(). `visit` is
+  // called with the output and each window's rows [start, end) once they are
+  // in place, while their pages are still cached.
+  apply(
+    source,
+    stride,
+    pool,
+    words = false,
+    label = undefined,
+    visit = undefined,
+  ) {
     const output = new PagedArray(
       this.workspace,
       source.Type,
@@ -679,6 +775,11 @@ export class Permutation {
     );
     const progress =
       label && new Progress(label, Math.ceil(this.length / window) * this.rows);
+    // The lowest and highest target + 1 of each block of source rows, found
+    // as the first window reads them: a later window that none of a block's
+    // rows land in skips the block.
+    const low = new Float64Array(Math.ceil(this.rows / block));
+    const high = new Float64Array(low.length);
     for (let start = 0; start < this.length; start += window) {
       const end = Math.min(this.length, start + window);
       output.pinned = [
@@ -687,8 +788,19 @@ export class Permutation {
       ];
       output.load(start * stride, end * stride, true);
       for (let base = 0; base < this.rows; base += block) {
+        const b = base / block;
+        if (start && (high[b] <= start || low[b] > end)) continue;
         const count = Math.min(block, this.rows - base);
         this.targets.read(base, targets, count);
+        if (!start) {
+          low[b] = Number.POSITIVE_INFINITY;
+          for (let i = 0; i < count; i++)
+            if (targets[i]) {
+              low[b] = Math.min(low[b], targets[i]);
+              high[b] = Math.max(high[b], targets[i]);
+            }
+          if (low[b] > end) continue;
+        }
         source.load(base * stride, (base + count) * stride);
         pool.copy(source, output, stride, count, base, start, end, words);
         for (let i = 0; i < count; i++) {
@@ -710,6 +822,7 @@ export class Permutation {
         }
         progress?.update(Math.floor(start / window) * this.rows + base + count);
       }
+      visit?.(output, start, end);
     }
     output.pinned = [0, 0];
     source.dispose();
@@ -763,7 +876,7 @@ export function readAll(fd, bytes, position) {
   return bytes;
 }
 
-export function writeAll(fd, bytes, position) {
+function writeAll(fd, bytes, position) {
   let offset = 0;
   while (offset < bytes.byteLength) {
     const count = writeSync(

package/bin/rad/writeback.js

@@ -0,0 +1,100 @@
+import { fdatasync } from "node:fs";
+import { getSystemErrorName } from "node:util";
+
+// The CLI thread can receive these requests while conversion runs synchronous
+// code. Completion is shared so the converter need not yield its event loop.
+export function syncFiles(files, state, cancelled) {
+  // A batch can still be in the message queue when cancellation arrives.
+  // Acknowledge it without submitting any more native I/O.
+  if (Atomics.load(cancelled, 0)) {
+    Atomics.store(state, 0, 0);
+    Atomics.notify(state, 0);
+    return;
+  }
+  for (const fd of files)
+    fdatasync(fd, (error) => {
+      if (error) Atomics.compareExchange(state, 1, 0, error.errno);
+      Atomics.sub(state, 0, 1);
+      Atomics.notify(state, 0);
+    });
+}
+
+// Bounds buffered writes that have not yet been confirmed written back. One
+// batch syncs while the next fills; only the older batch is acknowledged when
+// its sync completes, even if the same files have been written again.
+export class Writeback {
+  // `check` throws once the conversion is cancelled.
+  constructor(limit, submit, check) {
+    this.limit = limit;
+    this.submit = submit;
+    this.check = check;
+    this.files = new Map();
+    this.bytes = 0;
+    this.pending = undefined;
+    // Outstanding syncs, then the first negative system error number.
+    this.state = new Int32Array(new SharedArrayBuffer(8));
+  }
+
+  complete() {
+    if (!this.pending || Atomics.load(this.state, 0)) return;
+    this.pending = undefined;
+    const error = Atomics.load(this.state, 1);
+    if (error)
+      throw new Error(
+        `Unable to write back file data: ${getSystemErrorName(error)}`,
+      );
+  }
+
+  // Cancellation interrupts conversion, but keeps the batch tracked. Cleanup
+  // waits without cancellation so every fd stays open until its sync returns.
+  wait(cancellable = true) {
+    if (!this.pending) return;
+    for (;;) {
+      if (cancellable) this.check();
+      const remaining = Atomics.load(this.state, 0);
+      if (!remaining) break;
+      Atomics.wait(this.state, 0, remaining, cancellable ? 100 : undefined);
+    }
+    this.complete();
+  }
+
+  // Called before writing: the write itself must fit, including while an
+  // earlier batch is still being written back.
+  prepare(bytes) {
+    this.check();
+    this.complete();
+    if (this.bytes + (this.pending?.bytes ?? 0) + bytes <= this.limit) return;
+    this.wait();
+    if (this.bytes + bytes > this.limit) {
+      this.flush();
+      this.wait();
+    }
+  }
+
+  wrote(fd, bytes) {
+    this.files.set(fd, (this.files.get(fd) ?? 0) + bytes);
+    this.bytes += bytes;
+    // Each sync costs a call per file, so batches are large; the last
+    // quarter of the limit takes writes while the older batch syncs.
+    if (!this.pending && this.bytes >= (this.limit * 3) / 4) this.flush();
+  }
+
+  flush() {
+    this.check();
+    const files = [...this.files.keys()];
+    Atomics.store(this.state, 0, files.length);
+    Atomics.store(this.state, 1, 0);
+    this.pending = { files: this.files, bytes: this.bytes };
+    this.submit(files, this.state);
+    this.files = new Map();
+    this.bytes = 0;
+  }
+
+  // Before a file is closed, its outstanding syncs must finish: otherwise a
+  // queued sync could operate on a different file that reuses its descriptor.
+  synced(fd) {
+    if (this.pending?.files.has(fd)) this.wait(false);
+    this.bytes -= this.files.get(fd) ?? 0;
+    this.files.delete(fd);
+  }
+}