gaussian-splat-lite 1.5.0 → 1.5.2

package/bin/rad/rad.js

@@ -16,7 +16,7 @@ import { HALF, half, halfBits, opacityRadius } from "./math.js";
 import { WorkerGroup, encoders } from "./parallel.js";
 import { Progress } from "./progress.js";
 import { SH_DIMS } from "./splats.js";
-import { readAll, removePath, writeAll } from "./storage.js";
+import { PAGE_BYTES, readAll, removePath } from "./storage.js";
 
 const CHUNK = 65536;
 // Fixed s8 range with SPZ's 1/128 SH step; -1 clips to -127/128.
@@ -158,27 +158,34 @@ const unorm = (value, min, max) =>
   Math.round(Math.max(0, Math.min(255, ((value - min) / (max - min)) * 255)));
 
 // The columns encodeChunk() reads. With a codebook, labels replace the SH.
-function gather(splats, base, count, codebook) {
+// The records before `kept` come through the cache, where counting left
+// them; the rest of them and the other columns, read once, pass it by.
+function gather(splats, base, count, codebook, kept) {
   return {
     base,
     count,
     maxSh: splats.maxSh,
-    records: splats.records.read(base * 16, new Float32Array(count * 16)),
+    records: splats.records.scan(
+      base * 16,
+      new Float32Array(count * 16),
+      count * 16,
+      kept * 16,
+    ),
     sh: codebook
       ? []
       : splats.sh
           .slice(0, splats.maxSh)
           .map((column, band) =>
-            column.read(
+            column.stream(
               base * SH_DIMS[band],
               new Uint16Array(count * SH_DIMS[band]),
             ),
           ),
     ...(codebook && {
-      labels: codebook.labels.read(base, new Uint16Array(count)),
+      labels: codebook.labels.stream(base, new Uint16Array(count)),
     }),
-    counts: splats.childCount.read(base, new Uint32Array(count)),
-    starts: splats.childStart.read(base, new Float64Array(count)),
+    counts: splats.childCount.stream(base, new Uint32Array(count)),
+    starts: splats.childStart.stream(base, new Float64Array(count)),
   };
 }
 
@@ -447,15 +454,16 @@ if (!isMainThread && workerData?.coder) {
 
 // With a codebook from clusterSh(), splats store SH as its labels.
 export async function writeRad(splats, output, chunked, lodBase, codebook) {
+  const { workspace } = splats;
   const chunks = [];
   const stem = chunkStem(output);
-  const spool = chunked ? undefined : splats.workspace.file();
+  const spool = chunked ? undefined : workspace.file();
   // Write every file completely before publishing chunks, then the header.
   // Track published names too, so cancellation rolls back this conversion.
   const staged = [];
   const created = [];
   const stage = (file, write) => {
-    splats.workspace.check();
+    workspace.check();
     const temporary = path.join(
       path.dirname(file),
       `.${path.basename(file)}.${process.pid}.${randomBytes(6).toString("hex")}.tmp`,
@@ -464,14 +472,15 @@ export async function writeRad(splats, output, chunked, lodBase, codebook) {
     staged.push({ temporary, file });
     try {
       write(fd);
+      workspace.check();
       fsyncSync(fd);
-    } catch (error) {
+    } finally {
       try {
+        workspace.writeback.synced(fd);
+      } finally {
         closeSync(fd);
-      } catch {}
-      throw error;
+      }
     }
-    closeSync(fd);
   };
   let total = 0;
   let progress;
@@ -479,11 +488,11 @@ export async function writeRad(splats, output, chunked, lodBase, codebook) {
     if (chunked) {
       const filename = `${stem}-${chunks.length}.radc`;
       stage(path.join(path.dirname(output), filename), (fd) =>
-        writeAll(fd, bytes, 0),
+        workspace.write(fd, bytes, 0),
       );
       chunks.push({ offset: 0, bytes: bytes.length, filename });
     } else {
-      writeAll(spool, bytes, total);
+      workspace.write(spool, bytes, total);
       chunks.push({ offset: total, bytes: bytes.length });
     }
     total += bytes.length;
@@ -493,21 +502,29 @@ export async function writeRad(splats, output, chunked, lodBase, codebook) {
   // order. A chunk's failure is thrown when its turn comes, not as an
   // unhandled rejection while an earlier chunk is awaited.
   const pending = [];
-  const encoder = new Encoder(
-    splats.workspace,
-    Math.ceil(splats.length / CHUNK),
-  );
+  const encoder = new Encoder(workspace, Math.ceil(splats.length / CHUNK));
   try {
-    // Count the encoding ranges in a first pass.
+    // Count the encoding ranges in a first pass. The records that the cache
+    // holds stay there for encoding; see gather(). Recompute the prefix
+    // after measuring worker growth so it fits a cache that has shrunk.
+    const keep = () => {
+      workspace.adjust();
+      return Math.floor(Math.max(0, workspace.capacity - PAGE_BYTES) / 64);
+    };
     const chunkCount = Math.ceil(splats.length / CHUNK);
     progress = new Progress("Analyzing RAD encoding", chunkCount);
     let countedChunks = 0;
     for (let base = 0; base < splats.length; base += CHUNK) {
-      splats.workspace.check();
+      workspace.check();
+      // The encoders grow as they work, and reading past the cache loads
+      // no pages, which would measure them.
+      const kept = keep();
       const count = Math.min(CHUNK, splats.length - base);
-      const records = splats.records.read(
+      const records = splats.records.scan(
         base * 16,
         new Float32Array(count * 16),
+        count * 16,
+        kept * 16,
       );
       const counted = encoder
         .count(records)
@@ -522,12 +539,14 @@ export async function writeRad(splats, output, chunked, lodBase, codebook) {
     progress.finish();
     progress = new Progress("Encoding RAD chunks", chunkCount);
     for (let base = 0; base < splats.length; base += CHUNK) {
-      splats.workspace.check();
+      workspace.check();
+      const kept = keep();
       const chunk = gather(
         splats,
         base,
         Math.min(CHUNK, splats.length - base),
         codebook,
+        kept,
       );
       if (codebook && base === 0)
         chunk.codes = encodeCodes(codebook, splats.maxSh);
@@ -556,25 +575,28 @@ export async function writeRad(splats, output, chunked, lodBase, codebook) {
       meta.length + (spool === undefined ? 0 : total),
     );
     stage(output, (fd) => {
-      writeAll(fd, meta, 0);
+      workspace.write(fd, meta, 0);
       if (spool !== undefined) {
         const buffer = Buffer.allocUnsafe(1024 * 1024);
         for (let offset = 0; offset < total; offset += buffer.length) {
-          splats.workspace.check();
+          workspace.check();
           const bytes = buffer.subarray(
             0,
             Math.min(buffer.length, total - offset),
           );
           readAll(spool, bytes, offset);
-          writeAll(fd, bytes, meta.length + offset);
+          workspace.write(fd, bytes, meta.length + offset);
           writing.update(meta.length + offset + bytes.length);
         }
       }
     });
     writing.finish();
+    // Closing an output waits only for a sync that includes it; a temporary
+    // file's failed writeback must surface before anything is published.
+    workspace.writeback.wait();
     const publishing = new Progress("Publishing RAD files", staged.length);
     for (const { temporary, file } of staged) {
-      splats.workspace.check();
+      workspace.check();
       // rename can replace an existing file, so reserve the name exclusively.
       const fd = openSync(file, "wx");
       created.push(file);
@@ -582,7 +604,7 @@ export async function writeRad(splats, output, chunked, lodBase, codebook) {
       renameSync(temporary, file);
       publishing.update(created.length);
     }
-    splats.workspace.check();
+    workspace.check();
     publishing.finish();
     return meta.length + total;
   } catch (error) {

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();