gaussian-splat-lite 1.0.1 → 1.0.2

package/dist/gaussian-splat-lite.module.js

@@ -4913,11 +4913,8 @@ const ELEMENTS_PER_THREAD = 8;
 const ELEMENTS_PER_WORKGROUP = WORKGROUP_SIZE$1 * ELEMENTS_PER_THREAD;
 const PREFIX_ITEMS_PER_WORKGROUP = WORKGROUP_SIZE$1 * 2;
 const PREFIX_LEVELS = 3;
-function storage(buffer2, type, name, readOnly = false) {
-  const getBuffer = () => buffer2.value;
-  const node = N$2.storage(getBuffer(), type).onObjectUpdate(getBuffer);
-  node.name = name;
-  return readOnly ? node.toReadOnly() : node;
+function storage(buffer2, name) {
+  return N$2.storage(buffer2.value, "uint").setName(name).onObjectUpdate(() => buffer2.value);
 }
 function mutableUniform(value, type) {
   const holder = { value };
@@ -4942,8 +4939,8 @@ function makeHistogramTask({
   bitOffset,
   workgroupCount
 }) {
-  const inputKeys = storage(input, "uint", "gslRadixInputKeys", true);
-  const sums = storage(blockSums, "uint", "gslRadixBlockSums");
+  const inputKeys = storage(input, "gslRadixInputKeys").toReadOnly();
+  const sums = storage(blockSums, "gslRadixBlockSums");
   const bit = N$2.uint(bitOffset);
   const histogram = N$2.workgroupArray("uint", RADIX_BUCKETS).toAtomic().setName("gslRadixHistogram");
   return N$2.Fn(() => {
@@ -4976,8 +4973,8 @@ function makeHistogramTask({
   })().compute([1, 1, 1], [WORKGROUP_SIZE$1]).setName("Splat radix histogram");
 }
 function makePrefixScanTask(itemsBuffer, blockSumsBuffer, elementCount) {
-  const items = storage(itemsBuffer, "uint", "gslPrefixItems");
-  const blockSums = storage(blockSumsBuffer, "uint", "gslPrefixBlockSums");
+  const items = storage(itemsBuffer, "gslPrefixItems");
+  const blockSums = storage(blockSumsBuffer, "gslPrefixBlockSums");
   const temp = N$2.workgroupArray("uint", PREFIX_ITEMS_PER_WORKGROUP).setName(
     "gslPrefixTemp"
   );
@@ -5049,13 +5046,8 @@ function makePrefixScanTask(itemsBuffer, blockSumsBuffer, elementCount) {
   return compute;
 }
 function makePrefixAddTask(itemsBuffer, blockSumsBuffer, elementCount) {
-  const items = storage(itemsBuffer, "uint", "gslPrefixItems");
-  const blockSums = storage(
-    blockSumsBuffer,
-    "uint",
-    "gslPrefixBlockSums",
-    true
-  );
+  const items = storage(itemsBuffer, "gslPrefixItems");
+  const blockSums = storage(blockSumsBuffer, "gslPrefixBlockSums").toReadOnly();
   const compute = N$2.Fn(() => {
     const workgroup = workgroupIndex();
     const first = workgroup.mul(PREFIX_ITEMS_PER_WORKGROUP).add(N$2.invocationLocalIndex.mul(2));
@@ -5085,23 +5077,15 @@ function makeReorderTask({
 }) {
   const inputKeys = storage(
     inputKeysAttribute,
-    "uint",
-    "gslRadixInputKeys",
-    true
-  );
-  const outputKeys = lastPass ? null : storage(outputKeysAttribute, "uint", "gslRadixOutputKeys");
+    "gslRadixInputKeys"
+  ).toReadOnly();
+  const outputKeys = lastPass ? null : storage(outputKeysAttribute, "gslRadixOutputKeys");
   const inputValues = storage(
     inputValuesAttribute,
-    "uint",
-    "gslRadixInputValues",
-    true
-  );
-  const outputValues = storage(
-    outputValuesAttribute,
-    "uint",
-    "gslRadixOutputValues"
-  );
-  const prefix = storage(prefixAttribute, "uint", "gslRadixPrefix", true);
+    "gslRadixInputValues"
+  ).toReadOnly();
+  const outputValues = storage(outputValuesAttribute, "gslRadixOutputValues");
+  const prefix = storage(prefixAttribute, "gslRadixPrefix").toReadOnly();
   const bitOffsetNode = N$2.uint(bitOffset);
   const digitMasks = N$2.workgroupArray("uint", RADIX_BUCKETS * 8).toAtomic().setName("gslRadixDigitMasks");
   const digitOffsets = N$2.workgroupArray("uint", RADIX_BUCKETS).setName(
@@ -5224,7 +5208,7 @@ class WebGPURadixSort {
     this.values = [makeBufferRef(safeCapacity), makeBufferRef(safeCapacity)];
     const histogramCount = RADIX_BUCKETS * this.paddedWorkgroups(safeCapacity, ELEMENTS_PER_WORKGROUP);
     this.blockSums = makeBufferRef(histogramCount);
-    const counts = storage(this.counts, "uint", "gslRadixCounts", true);
+    const counts = storage(this.counts, "gslRadixCounts").toReadOnly();
     this.createPrefixLevels(histogramCount, counts);
     const setup = this.makeSetupTask(options.count);
     const prefixNodes = this.prefixLevels.map((_, lastLevel) => {
@@ -5291,13 +5275,11 @@ class WebGPURadixSort {
     return Math.max(1, width * Math.ceil(groups / this.maxWorkgroups));
   }
   makeSetupTask(gpuCount) {
-    const counts = storage(this.counts, "uint", "gslRadixCounts");
+    const counts = storage(this.counts, "gslRadixCounts");
     const dispatchBuffers = [
       this.sortDispatch,
       ...this.prefixLevels.map((level) => level.dispatch)
-    ].map(
-      (value, index2) => storage({ value }, "uint", `gslRadixDispatch${index2}`)
-    );
+    ].map((value, index2) => storage({ value }, `gslRadixDispatch${index2}`));
     const maxWorkgroups = N$2.uint(this.maxWorkgroups);
     return N$2.Fn(() => {
       const count = N$2.min(
@@ -5409,29 +5391,31 @@ class WebGPURadixSort {
 const N$1 = TSL;
 const WORKGROUP_SIZE = 256;
 const PROJECT_SLOTS = 8;
-function buffer(count = 1) {
-  return { value: new StorageBufferAttribute(new Uint32Array(count), 1) };
+function buffer(name) {
+  return { value: new StorageBufferAttribute(new Uint32Array(1), 1), name };
 }
-function bindBuffer(ref, readOnly = false) {
-  const node = N$1.storage(ref.value, "uint").onObjectUpdate(() => ref.value);
-  return readOnly ? node.toReadOnly() : node;
+function bindBuffer(ref) {
+  return N$1.storage(ref.value, "uint").setName(ref.name).onObjectUpdate(() => ref.value);
 }
 class ProjectedSplats {
   constructor(renderer, uniforms) {
     this.renderer = renderer;
     this.uniforms = uniforms;
-    this.indirect = new IndirectStorageBufferAttribute(new Uint32Array(5), 1);
-    this.pending = true;
+    this.indirect = new IndirectStorageBufferAttribute(
+      new Uint32Array([WEBGPU_SPLATS_PER_INSTANCE * 6, 0, 0, 0, 0]),
+      1
+    );
     this.error = null;
     this.cache = new ProjectionCache();
     this.visibleCount = new StorageBufferAttribute(
       new Uint32Array(1),
       1
     );
-    this.keys = buffer();
-    this.seeds = buffer();
-    this.counts = buffer();
+    this.keys = buffer("gslProjectionKeys");
+    this.seeds = buffer("gslProjectionSeeds");
+    this.counts = buffer("gslProjectionCounts");
     this.slots = [];
+    this.compiledSlots = 0;
     this.matrix = new THREE.Matrix4();
     this.translation = new THREE.Matrix4();
     this.scale = new THREE.Vector3();
@@ -5461,8 +5445,8 @@ class ProjectedSplats {
     const viewBase = uniformBinding(this.state, "viewBase", "uint");
     const multiView = uniformBinding(this.state, "multiView", "bool");
     const stochastic = uniformBinding(uniforms, "stochastic", "bool");
-    const count = N$1.storage(this.visibleCount, "uint").toReadOnly().element(0);
-    const seeds = bindBuffer(this.seeds, true);
+    const count = N$1.storage(this.visibleCount, "uint").setName("gslVisibleCount").toReadOnly().element(0);
+    const seeds = bindBuffer(this.seeds).toReadOnly();
     this.sorter = new WebGPURadixSort(1, this.keys, {
       count,
       storeOrder: (index2, value) => {
@@ -5479,15 +5463,8 @@ class ProjectedSplats {
         this.limits.maxBufferSize
       )
     });
-    const draw = N$1.storage(this.indirect, "uint");
-    this.resetDraw = N$1.Fn(() => {
-      draw.element(0).assign(WEBGPU_SPLATS_PER_INSTANCE * 6);
-      draw.element(1).assign(0);
-      draw.element(2).assign(0);
-      draw.element(3).assign(0);
-      draw.element(4).assign(0);
-    })().compute(1).setName("Splat reset draw");
-    const counter = N$1.storage(this.visibleCount, "uint").toAtomic();
+    const drawCount = N$1.storage(this.indirect, "uint").setName("gslDrawIndirect").element(1);
+    const counter = N$1.storage(this.visibleCount, "uint").setName("gslVisibleCount").toAtomic();
     this.resetCount = N$1.Fn(() => {
       N$1.atomicStore(counter.element(0), N$1.uint(0));
     })().compute(1, [1]).setName("Splat reset visible count");
@@ -5498,7 +5475,9 @@ class ProjectedSplats {
       N$1.If(index2.equal(0), () => {
         counts.element(viewIndex).assign(count);
         const instances = count.add(WEBGPU_SPLATS_PER_INSTANCE - 1).div(N$1.uint(WEBGPU_SPLATS_PER_INSTANCE));
-        draw.element(1).assign(draw.element(1).max(instances));
+        drawCount.assign(
+          N$1.select(viewIndex.equal(0), instances, drawCount.max(instances))
+        );
       });
       N$1.If(multiView.and(stochastic).and(index2.lessThan(count)), () => {
         this.cache.storeOrder(
@@ -5509,17 +5488,30 @@ class ProjectedSplats {
     })().compute(1, [WORKGROUP_SIZE]).setName("Splat visible draw arguments");
     for (let i = 0; i < PROJECT_SLOTS; i++) this.slots.push(this.createSlot());
     this.ready = computeRenderer.compileComputeAsync([
-      this.resetDraw,
       this.resetCount,
       this.finish,
-      ...this.slots.map((slot) => slot.node),
-      ...this.sorter.nodes
-    ]).catch((error) => {
+      ...this.sorter.nodes,
+      this.slots[0].node
+    ]).then(() => {
+      this.compiledSlots = 1;
+    }).catch((error) => {
       this.error = error;
-    }).finally(() => {
-      this.pending = false;
-      if (this.disposed) this.disposeResources();
     });
+    this.compilation = this.ready.then(() => this.compileRemainingSlots());
+  }
+  async compileRemainingSlots() {
+    if (this.disposed || this.compiledSlots === 0) return;
+    try {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      while (!this.disposed && this.compiledSlots < this.slots.length) {
+        await this.renderer.compileComputeAsync([
+          this.slots[this.compiledSlots].node
+        ]);
+        this.compiledSlots++;
+      }
+    } catch (error) {
+      this.error = error;
+    }
   }
   createSlot() {
     const uniforms = { ...this.state, ...makeGenerateUniforms() };
@@ -5543,7 +5535,7 @@ class ProjectedSplats {
     const pixelScale = u("renderSize", "vec2").mul(u("focalAdjustment", "float")).mul(0.5);
     const keys = bindBuffer(this.keys);
     const seeds = bindBuffer(this.seeds);
-    const counter = N$1.storage(this.visibleCount, "uint").toAtomic();
+    const counter = N$1.storage(this.visibleCount, "uint").setName("gslVisibleCount").toAtomic();
     const node = N$1.Fn(() => {
       const index2 = N$1.uint(N$1.instanceIndex);
       const generated = generate.prepare(index2);
@@ -5596,7 +5588,7 @@ class ProjectedSplats {
     const stride = uniformBinding(this.state, "viewStride", "uint");
     const base = eye.mul(stride).toVar();
     const i = N$1.uint(N$1.instanceIndex).mul(WEBGPU_SPLATS_PER_INSTANCE).add(N$1.uint(N$1.positionGeometry.z));
-    const counts = bindBuffer(this.counts, true);
+    const counts = bindBuffer(this.counts).toReadOnly();
     const clipPosition = N$1.vec4(0, 0, 2, 1).toVar();
     const rgba = N$1.vec4(0).toVar();
     const splatUv = N$1.vec2(0).toVar();
@@ -5612,7 +5604,7 @@ class ProjectedSplats {
       const cacheIndex = i.toVar();
       const assignSeed = () => {
         stochasticSeed.assign(
-          multiView ? this.cache.readOrder(base.add(i)).y : bindBuffer(this.seeds, true).element(i)
+          multiView ? this.cache.readOrder(base.add(i)).y : bindBuffer(this.seeds).toReadOnly().element(i)
         );
       };
       if (depthOnly) {
@@ -5691,7 +5683,7 @@ class ProjectedSplats {
   }
   render(accumulator, camera, geometry, radial, shrink = false) {
     if (this.error) throw this.error;
-    if (this.pending || this.disposed) {
+    if (this.compiledSlots === 0 || this.disposed) {
       geometry.instanceCount = 0;
       return;
     }
@@ -5725,12 +5717,11 @@ class ProjectedSplats {
         this.scale
       );
       this.state.renderToViewScale.value = (this.scale.x + this.scale.y + this.scale.z) / 3;
-      const pending = eye === 0 ? [this.resetDraw] : [];
-      pending.push(this.resetCount);
+      const pending = [this.resetCount];
       let slotCount = 0;
       for (const { node, count } of accumulator.mapping) {
         if (!view.layers.test(node.layers)) continue;
-        if (slotCount === PROJECT_SLOTS) {
+        if (slotCount === this.compiledSlots) {
           this.renderer.compute(pending);
           pending.length = 0;
           slotCount = 0;
@@ -5751,12 +5742,11 @@ class ProjectedSplats {
   dispose() {
     if (this.disposed) return;
     this.disposed = true;
-    if (!this.pending) this.disposeResources();
+    void this.compilation.then(() => this.disposeResources());
   }
   disposeResources() {
     for (const slot of this.slots) slot.node.dispose();
-    for (const node of [this.resetDraw, this.resetCount, this.finish])
-      node.dispose();
+    for (const node of [this.resetCount, this.finish]) node.dispose();
     this.sorter.dispose();
     this.cache.dispose();
     this.visibleCount.dispose();