diff --git a/packages/dev/core/src/Compute/prefixSumCompute.ts b/packages/dev/core/src/Compute/prefixSumCompute.ts new file mode 100644 index 000000000000..40923108e357 --- /dev/null +++ b/packages/dev/core/src/Compute/prefixSumCompute.ts @@ -0,0 +1,178 @@ +import { type Nullable } from "core/types"; +import { type AbstractEngine } from "core/Engines/abstractEngine.pure"; +import { type WebGPUEngine } from "core/Engines/webgpuEngine.pure"; +import { type DataBuffer } from "core/Buffers/dataBuffer"; +import { type ComputeBindingMapping } from "core/Engines/Extensions/engine.computeShader.pure"; +import { ComputeShader } from "core/Compute/computeShader.pure"; +import { StorageBuffer } from "core/Buffers/storageBuffer"; +import { UniformBuffer } from "core/Materials/uniformBuffer"; +import { Constants } from "core/Engines/constants"; + +const BlockSize = 256; + +/** + * @internal + * WebGPU compute utility that performs an in-place hierarchical EXCLUSIVE prefix sum (scan) over a + * `StorageBuffer` of u32 values. + * + * It scans the array in blocks of {@link BlockSize}, recursively scans the per-block totals, and adds + * the resulting offsets back, so it handles arrays much larger than a single workgroup. It is used by + * the Gaussian Splatting GPU depth sort (histogram into per-bucket start offsets) and its interval + * culling pass. + * + * The WGSL compute shaders are loaded asynchronously during construction (they self-register into the + * ShaderStore), so this module stays side-effect free: no top-level shader imports and therefore no + * `.pure` / side-effect-wrapper split. {@link isReady} returns false until the shaders have loaded and + * the compute pipelines have compiled. + */ +export class PrefixSumCompute { + private readonly _engine: AbstractEngine; + private _scanBlock: Nullable = null; + private _addOffsets: Nullable = null; + // One UBO per dispatch, cycled per scan, to avoid a later dispatch's parameters overwriting an + // earlier (still-pending) dispatch's uniform buffer (the update/dispatch hazard). + private _ubos: UniformBuffer[] = []; + private _uboIndex = 0; + // Per-recursion-level scratch buffers holding the block totals; grown/cached by level. + private _levelBuffers: StorageBuffer[] = []; + // Set once the WGSL shader modules have been dynamically imported (and thus registered in the ShaderStore). + private _shadersLoaded = false; + + /** + * Creates a new prefix-sum compute helper and kicks off asynchronous loading of its WGSL shaders. + * @param engine the (WebGPU) engine to run the compute passes on + */ + public constructor(engine: AbstractEngine) { + this._engine = engine; + // Async-load the compute shaders during construction. Importing the generated shader modules registers + // them into the ShaderStore as a side effect; isReady() stays false until this resolves. + // eslint-disable-next-line @typescript-eslint/no-floating-promises + this._loadShadersAsync(); + } + + private async _loadShadersAsync(): Promise { + await Promise.all([import("../ShadersWGSL/prefixSumScanBlock.compute"), import("../ShadersWGSL/prefixSumAddOffsets.compute")]); + this._shadersLoaded = true; + } + + private _ensureShaders(): void { + if (this._scanBlock || !this._shadersLoaded) { + return; + } + const scanBlockBindings: ComputeBindingMapping = { + data: { group: 0, binding: 0 }, + blockSums: { group: 0, binding: 1 }, + params: { group: 0, binding: 2 }, + }; + this._scanBlock = new ComputeShader("prefixSumScanBlock", this._engine, "prefixSumScanBlock", { bindingsMapping: scanBlockBindings }); + const addOffsetsBindings: ComputeBindingMapping = { + data: { group: 0, binding: 0 }, + blockOffsets: { group: 0, binding: 1 }, + params: { group: 0, binding: 2 }, + }; + this._addOffsets = new ComputeShader("prefixSumAddOffsets", this._engine, "prefixSumAddOffsets", { bindingsMapping: addOffsetsBindings }); + } + + private _getUbo(count: number): UniformBuffer { + if (this._uboIndex >= this._ubos.length) { + const ubo = new UniformBuffer(this._engine, undefined, undefined, "PrefixSumComputeParams", false, false); + ubo.addUniform("count", 1); + this._ubos.push(ubo); + } + const ubo = this._ubos[this._uboIndex++]; + ubo.updateUInt("count", count); + ubo.update(); + return ubo; + } + + private _getLevelBuffer(level: number, numBlocks: number): StorageBuffer { + // Padded to a multiple of BlockSize so the next level's block scan never reads out of bounds. + const capacity = (Math.ceil(numBlocks / BlockSize) * BlockSize) | 0; + const existing = this._levelBuffers[level]; + if (existing && existing.getBuffer().capacity >= capacity * Uint32Array.BYTES_PER_ELEMENT) { + return existing; + } + existing?.dispose(); + const buffer = new StorageBuffer( + this._engine as WebGPUEngine, + Math.max(capacity, BlockSize) * Uint32Array.BYTES_PER_ELEMENT, + Constants.BUFFER_CREATIONFLAG_READWRITE, + "PrefixSumLevel" + level + ); + this._levelBuffers[level] = buffer; + return buffer; + } + + /** + * Whether both compute shaders are loaded, compiled and ready to dispatch. + * @returns true when the scan can run + */ + public isReady(): boolean { + if (!this._shadersLoaded) { + return false; + } + this._ensureShaders(); + return !!this._scanBlock!.isReady() && !!this._addOffsets!.isReady(); + } + + /** + * Runs an in-place exclusive prefix sum over the first `count` entries of `buffer`. + * Call {@link resetForFrame} once before the sequence of scans issued in a frame. + * @param buffer the u32 storage buffer to scan in place + * @param count number of valid entries to scan + */ + public scanExclusive(buffer: StorageBuffer, count: number): void { + this._ensureShaders(); + this._scanRecursive(buffer.getBuffer(), count, 0); + } + + /** + * Resets the internal UBO ring. Must be called once at the start of each frame's scan sequence. + */ + public resetForFrame(): void { + this._uboIndex = 0; + } + + private _scanRecursive(data: DataBuffer, count: number, level: number): void { + // count === 1 must still run the block scan so the single element is written to its exclusive value (0); + // returning here would leave it unchanged. Only an empty scan is a true no-op. + if (count <= 0) { + return; + } + const numBlocks = Math.ceil(count / BlockSize); + const blockSums = this._getLevelBuffer(level, numBlocks); + + const scanUbo = this._getUbo(count); + this._scanBlock!.setStorageBuffer("data", data); + this._scanBlock!.setStorageBuffer("blockSums", blockSums); + this._scanBlock!.setUniformBuffer("params", scanUbo); + this._scanBlock!.dispatch(numBlocks); + + if (numBlocks > 1) { + this._scanRecursive(blockSums.getBuffer(), numBlocks, level + 1); + + const addUbo = this._getUbo(count); + this._addOffsets!.setStorageBuffer("data", data); + this._addOffsets!.setStorageBuffer("blockOffsets", blockSums); + this._addOffsets!.setUniformBuffer("params", addUbo); + this._addOffsets!.dispatch(numBlocks); + } + } + + /** + * Releases all GPU resources held by the helper. + */ + public dispose(): void { + for (const ubo of this._ubos) { + ubo.dispose(); + } + for (const buffer of this._levelBuffers) { + buffer.dispose(); + } + this._ubos = []; + this._levelBuffers = []; + this._uboIndex = 0; + this._scanBlock = null; + this._addOffsets = null; + } +} diff --git a/packages/dev/core/src/Engines/WebGPU/webgpuDrawContext.ts b/packages/dev/core/src/Engines/WebGPU/webgpuDrawContext.ts index d8484291762b..7d8cc6692792 100644 --- a/packages/dev/core/src/Engines/WebGPU/webgpuDrawContext.ts +++ b/packages/dev/core/src/Engines/WebGPU/webgpuDrawContext.ts @@ -39,6 +39,16 @@ export class WebGPUDrawContext implements IDrawContext { public indirectDrawBuffer?: GPUBuffer; + /** + * @internal + * The caller-owned (external) indirect buffer and offset the cached {@link fastBundle} was recorded against, + * if any. Used to invalidate the bundle when the caller passes a different external indirect buffer/offset + * (the bundle records drawIndirect against a specific buffer, which is otherwise not part of the dirty key). + */ + public _externalIndirectBuffer?: GPUBuffer; + /** @internal */ + public _externalIndirectOffset = 0; + private _materialContextUpdateId: number; private _bufferManager: WebGPUBufferManager; private _useInstancing: boolean; diff --git a/packages/dev/core/src/Engines/abstractEngine.pure.ts b/packages/dev/core/src/Engines/abstractEngine.pure.ts index b79ce9bd2286..07fe45a64713 100644 --- a/packages/dev/core/src/Engines/abstractEngine.pure.ts +++ b/packages/dev/core/src/Engines/abstractEngine.pure.ts @@ -1345,6 +1345,19 @@ export abstract class AbstractEngine { */ public abstract drawElementsType(fillMode: number, indexStart: number, indexCount: number, instancesCount?: number): void; + /** + * Draws indexed instanced primitives where the draw arguments (including the instance count) are read from a + * caller-owned GPU indirect buffer. Only supported on WebGPU; other engines throw. + * @param _fillMode defines the primitive to use + * @param _indexStart defines the starting index + * @param _indexCount defines the number of indices per instance + * @param _indirectBuffer the GPU buffer holding the draw-indexed-indirect arguments + * @param _indirectByteOffset byte offset of the arguments within the buffer (default 0) + */ + public drawElementsInstancedIndirect(_fillMode: number, _indexStart: number, _indexCount: number, _indirectBuffer: DataBuffer, _indirectByteOffset = 0): void { + throw new Error("drawElementsInstancedIndirect is only supported on WebGPU engines."); + } + /** * Unbind the current render target texture from the webGL context * @param texture defines the render target wrapper to unbind diff --git a/packages/dev/core/src/Engines/webgpuEngine.pure.ts b/packages/dev/core/src/Engines/webgpuEngine.pure.ts index 548f2f16089d..7511bed68b21 100644 --- a/packages/dev/core/src/Engines/webgpuEngine.pure.ts +++ b/packages/dev/core/src/Engines/webgpuEngine.pure.ts @@ -3722,7 +3722,7 @@ export class WebGPUEngine extends ThinWebGPUEngine { } } - private _draw(drawType: number, fillMode: number, start: number, count: number, instancesCount: number): void { + private _draw(drawType: number, fillMode: number, start: number, count: number, instancesCount: number, externalIndirect?: { buffer: GPUBuffer; offset: number }): void { const renderPass = this._getCurrentRenderPass(); const bundleList = this._bundleList; @@ -3750,9 +3750,19 @@ export class WebGPUEngine extends ThinWebGPUEngine { return; } + // The cached fast bundle records drawIndirect against a specific external buffer/offset, which is not part + // of the draw/material context dirty key. Invalidate it when the caller supplies a different external + // indirect buffer/offset (or stops/starts using one), so a stale bundle can't draw from the wrong buffer. + const externalIndirectChanged = externalIndirect + ? this._currentDrawContext._externalIndirectBuffer !== externalIndirect.buffer || this._currentDrawContext._externalIndirectOffset !== externalIndirect.offset + : this._currentDrawContext._externalIndirectBuffer !== undefined; + if ( !this.compatibilityMode && - (this._currentDrawContext.isDirty(this._currentMaterialContext.updateId) || this._currentMaterialContext.isDirty || this._currentMaterialContext.forceBindGroupCreation) + (this._currentDrawContext.isDirty(this._currentMaterialContext.updateId) || + this._currentMaterialContext.isDirty || + this._currentMaterialContext.forceBindGroupCreation || + externalIndirectChanged) ) { this._currentDrawContext.fastBundle = undefined; } @@ -3764,7 +3774,10 @@ export class WebGPUEngine extends ThinWebGPUEngine { this._applyRenderPassChanges(bundleList); if (!this._snapshotRendering.record) { this._counters.numBundleReuseNonCompatMode++; - if (this._currentDrawContext.indirectDrawBuffer) { + if (externalIndirect) { + // The instance count lives in an app-owned indirect buffer written by a compute shader; the + // cached bundle already records drawIndexedIndirect against it, so nothing to update here. + } else if (this._currentDrawContext.indirectDrawBuffer) { this._currentDrawContext.setIndirectData(count, instancesCount || 1, start); } bundleList.addBundle(this._currentDrawContext.fastBundle); @@ -3838,7 +3851,15 @@ export class WebGPUEngine extends ThinWebGPUEngine { // draw const nonCompatMode = !this.compatibilityMode && !this._snapshotRendering.record; - if ((nonCompatMode || this._currentDrawContext._enableIndirectDrawInCompatMode) && this._currentDrawContext.indirectDrawBuffer) { + if (externalIndirect) { + // App-driven indirect draw: the draw arguments (including a GPU-computed instance count) live in a + // caller-owned buffer. Used by the Gaussian Splatting GPU culling path. + if (drawType === 0) { + renderPass2.drawIndexedIndirect(externalIndirect.buffer, externalIndirect.offset); + } else { + renderPass2.drawIndirect(externalIndirect.buffer, externalIndirect.offset); + } + } else if ((nonCompatMode || this._currentDrawContext._enableIndirectDrawInCompatMode) && this._currentDrawContext.indirectDrawBuffer) { this._currentDrawContext.setIndirectData(count, instancesCount || 1, start); if (drawType === 0) { renderPass2.drawIndexedIndirect(this._currentDrawContext.indirectDrawBuffer, 0); @@ -3852,6 +3873,10 @@ export class WebGPUEngine extends ThinWebGPUEngine { } if (nonCompatMode) { + // Remember which external indirect buffer/offset this bundle draws from, so it is invalidated above if + // the caller changes it on a later draw sharing this draw context. + this._currentDrawContext._externalIndirectBuffer = externalIndirect?.buffer; + this._currentDrawContext._externalIndirectOffset = externalIndirect?.offset ?? 0; this._currentDrawContext.fastBundle = (renderPass2 as GPURenderBundleEncoder).finish(); bundleList.addBundle(this._currentDrawContext.fastBundle); } @@ -3882,6 +3907,23 @@ export class WebGPUEngine extends ThinWebGPUEngine { this._draw(1, fillMode, verticesStart, verticesCount, instancesCount); } + /** + * Draws indexed instanced primitives where the draw arguments (including the instance count) are read + * from a caller-owned GPU buffer (drawIndexedIndirect). The buffer must contain 5 consecutive u32: + * [indexCount, instanceCount, firstIndex, baseVertex, firstInstance] and be created with the + * {@link Constants.BUFFER_CREATIONFLAG_INDIRECT} usage. Used by the Gaussian Splatting GPU culling path so + * that a compute-computed visible count drives the draw without any CPU readback. + * @param fillMode defines the primitive to use + * @param indexStart defines the starting index (used for pipeline/index-buffer binding) + * @param indexCount defines the number of indices per instance (used for pipeline/index-buffer binding) + * @param indirectBuffer the GPU buffer holding the draw-indexed-indirect arguments + * @param indirectByteOffset byte offset of the arguments within the buffer (default 0) + */ + public override drawElementsInstancedIndirect(fillMode: number, indexStart: number, indexCount: number, indirectBuffer: DataBuffer, indirectByteOffset: number = 0): void { + const buffer = (indirectBuffer as WebGPUDataBuffer).underlyingResource as GPUBuffer; + this._draw(0, fillMode, indexStart, indexCount, 1, { buffer, offset: indirectByteOffset }); + } + //------------------------------------------------------------------------------ // Async Pipeline Pre-Warming //------------------------------------------------------------------------------ diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts new file mode 100644 index 000000000000..e81d670399c0 --- /dev/null +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts @@ -0,0 +1,531 @@ +import { type Nullable } from "core/types"; +import { type Matrix } from "core/Maths/math.vector.pure"; +import { type AbstractEngine } from "core/Engines/abstractEngine.pure"; +import { type WebGPUEngine } from "core/Engines/webgpuEngine.pure"; +import { type DataBuffer } from "core/Buffers/dataBuffer"; +import { ComputeShader } from "core/Compute/computeShader.pure"; +import { StorageBuffer } from "core/Buffers/storageBuffer"; +import { UniformBuffer } from "core/Materials/uniformBuffer"; +import { Constants } from "core/Engines/constants"; +import { PrefixSumCompute } from "core/Compute/prefixSumCompute"; + +const WorkgroupSize = 256; +// Bit patterns of +Infinity / -Infinity used to seed the atomic depth-range reduction each sort. +const RangeInit = /* #__PURE__ */ new Int32Array(new Float32Array([Number.POSITIVE_INFINITY, Number.NEGATIVE_INFINITY]).buffer); + +/** + * @internal + * WebGPU compute driver that produces the depth-sorted per-splat `splatIndex` order buffer for a + * {@link GaussianSplattingMeshBase} entirely on the GPU, with no CPU readback. + * + * The sort is a single-pass counting sort by camera-forward depth (matching the CPU worker's algorithm): + * 1. depth pass - compute each active splat's signed depth and reduce the min/max range + * 2. histogram pass - map depth to a bucket and count per bucket + * 3. prefix sum - exclusive-scan the histogram into per-bucket start offsets + * 4. scatter pass - place each source index at the next free slot of its bucket + * + * The resulting sorted-index buffer is bound directly as the `splatIndex` instanced vertex buffer. + * The farthest splat gets bucket 0, so the ascending scatter yields back-to-front order for correct + * alpha blending. + * + * The WGSL compute shaders are loaded asynchronously during construction (they self-register into the + * ShaderStore), so this module stays side-effect free: no top-level shader imports and therefore no + * `.pure` / side-effect-wrapper split. {@link isReady} / {@link isCullReady} return false until the + * shaders have loaded and the compute pipelines have compiled. + */ +export class GaussianSplattingGpuSorter { + private readonly _engine: AbstractEngine; + + private _clear: Nullable = null; + private _depth: Nullable = null; + private _histogram: Nullable = null; + private _scatter: Nullable = null; + private _cull: Nullable = null; + private _compact: Nullable = null; + private _finalize: Nullable = null; + private _prefixSum: Nullable = null; + private _cullPrefixSum: Nullable = null; + private _params: Nullable = null; + private _cullParams: Nullable = null; + private _compactParams: Nullable = null; + private _finalizeParams: Nullable = null; + + private _positions: Nullable = null; + private _seed: Nullable = null; + private _depthBuffer: Nullable = null; + private _bucket: Nullable = null; + private _range: Nullable = null; + private _histogramBuffer: Nullable = null; + private _sortedIndexBuffer: Nullable = null; + // Culling buffers. + private _flag: Nullable = null; + private _offsets: Nullable = null; + private _drawList: Nullable = null; + private _count: Nullable = null; + private _cullCapacity = 0; + // Draw-indexed-indirect arguments [indexCount, instanceCount, firstIndex, baseVertex, firstInstance]. + // Fixed 5-u32 buffer (never grows), so cached render bundles referencing it stay valid while its contents update. + private _indirectArgs: Nullable = null; + private _indirectArgsData = new Uint32Array(5); + + private _positionCapacity = 0; + private _paddedCapacity = 0; + private _histCapacity = 0; + private _renderedCount = 0; + private _paddedCount = 0; + private _numBuckets = 0; + private _hasSource = false; + // Set once the WGSL shader modules have been dynamically imported (and thus registered in the ShaderStore). + // Until then isReady()/isCullReady() return false so the child ComputeShaders are never compiled early (which + // would fall through to a .fx network fetch for a shader that is not yet in the store). + private _shadersLoaded = false; + + /** + * Whether the running engine supports the WebGPU compute sort path. + * @param engine the engine to test + * @returns true when compute shaders are available + */ + public static IsSupported(engine: AbstractEngine): boolean { + return !!engine.getCaps().supportComputeShaders; + } + + /** + * Number of depth buckets to use for a given active splat count. + * Mirrors the CPU worker: round(log2(n/4)) clamped to [10, 20]. + * @param renderedCount active splat count + * @returns the bucket count (a power of two) + */ + private static _BucketCount(renderedCount: number): number { + const bits = Math.max(10, Math.min(20, Math.round(Math.log2(Math.max(1, renderedCount) / 4)))); + return 2 ** bits; + } + + /** + * Creates a new GPU sorter and kicks off asynchronous loading of its WGSL compute shaders. + * @param engine the (WebGPU) engine to run the compute passes on + */ + public constructor(engine: AbstractEngine) { + this._engine = engine; + // Async-load the compute shaders during construction. Importing the generated shader modules registers + // them into the ShaderStore as a side effect; isReady()/isCullReady() stay false until this resolves. + // eslint-disable-next-line @typescript-eslint/no-floating-promises + this._loadShadersAsync(); + } + + private async _loadShadersAsync(): Promise { + await Promise.all([ + import("../../ShadersWGSL/gaussianSplattingSortClear.compute"), + import("../../ShadersWGSL/gaussianSplattingSortDepth.compute"), + import("../../ShadersWGSL/gaussianSplattingSortHistogram.compute"), + import("../../ShadersWGSL/gaussianSplattingSortScatter.compute"), + import("../../ShadersWGSL/gaussianSplattingCull.compute"), + import("../../ShadersWGSL/gaussianSplattingCullCompact.compute"), + import("../../ShadersWGSL/gaussianSplattingCullFinalize.compute"), + ]); + this._shadersLoaded = true; + } + + private _ensureShaders(): void { + if (this._clear) { + return; + } + this._clear = new ComputeShader("gaussianSplattingSortClear", this._engine, "gaussianSplattingSortClear", { + bindingsMapping: { histogram: { group: 0, binding: 0 }, params: { group: 0, binding: 1 } }, + }); + this._depth = new ComputeShader("gaussianSplattingSortDepth", this._engine, "gaussianSplattingSortDepth", { + bindingsMapping: { + positions: { group: 0, binding: 0 }, + seed: { group: 0, binding: 1 }, + depth: { group: 0, binding: 2 }, + range: { group: 0, binding: 3 }, + params: { group: 0, binding: 4 }, + }, + }); + this._histogram = new ComputeShader("gaussianSplattingSortHistogram", this._engine, "gaussianSplattingSortHistogram", { + bindingsMapping: { + depth: { group: 0, binding: 0 }, + range: { group: 0, binding: 1 }, + bucket: { group: 0, binding: 2 }, + histogram: { group: 0, binding: 3 }, + params: { group: 0, binding: 4 }, + }, + }); + this._scatter = new ComputeShader("gaussianSplattingSortScatter", this._engine, "gaussianSplattingSortScatter", { + bindingsMapping: { + bucket: { group: 0, binding: 0 }, + seed: { group: 0, binding: 1 }, + offsets: { group: 0, binding: 2 }, + sortedIndex: { group: 0, binding: 3 }, + params: { group: 0, binding: 4 }, + }, + }); + this._prefixSum = new PrefixSumCompute(this._engine); + // The cull owns a separate prefix-sum instance from the sort. Sharing one ring made the cull's scan reuse a + // UBO the sort had just written with a different element count; combined with the one-dispatch compute-UBO + // upload lag, the cull then scanned with the wrong count. A dedicated ring always holds the cull's count. + this._cullPrefixSum = new PrefixSumCompute(this._engine); + // trackUBOsInFrame=false: these UBOs are updated once and dispatched immediately during mesh render; the + // default per-frame buffer rotation left the compute reading a stale buffer, freezing the sort/cull data. + this._params = new UniformBuffer(this._engine, undefined, undefined, "GaussianSplattingGpuSorterParams", false, false); + this._params.addUniform("count", 1); + this._params.addUniform("paddedCount", 1); + this._params.addUniform("numBuckets", 1); + this._params.addUniform("pad", 1); + this._params.addUniform("coeff", 4); + + this._cull = new ComputeShader("gaussianSplattingCull", this._engine, "gaussianSplattingCull", { + bindingsMapping: { + positions: { group: 0, binding: 0 }, + sortedIndex: { group: 0, binding: 1 }, + flag: { group: 0, binding: 2 }, + offsets: { group: 0, binding: 3 }, + params: { group: 0, binding: 4 }, + }, + }); + this._compact = new ComputeShader("gaussianSplattingCullCompact", this._engine, "gaussianSplattingCullCompact", { + bindingsMapping: { + sortedIndex: { group: 0, binding: 0 }, + flag: { group: 0, binding: 1 }, + offsets: { group: 0, binding: 2 }, + drawList: { group: 0, binding: 3 }, + params: { group: 0, binding: 4 }, + }, + }); + this._finalize = new ComputeShader("gaussianSplattingCullFinalize", this._engine, "gaussianSplattingCullFinalize", { + bindingsMapping: { + flag: { group: 0, binding: 0 }, + offsets: { group: 0, binding: 1 }, + countOut: { group: 0, binding: 2 }, + indirectArgs: { group: 0, binding: 3 }, + drawList: { group: 0, binding: 4 }, + params: { group: 0, binding: 5 }, + }, + }); + this._compactParams = new UniformBuffer(this._engine, undefined, undefined, "GaussianSplattingCompactParams", false, false); + this._compactParams.addUniform("count", 1); + this._cullParams = new UniformBuffer(this._engine, undefined, undefined, "GaussianSplattingCullParams", false, false); + this._cullParams.addUniform("count", 1); + this._cullParams.addUniform("pad0", 1); + this._cullParams.addUniform("pad1", 1); + this._cullParams.addUniform("pad2", 1); + this._cullParams.addUniform("clip0", 4); + this._cullParams.addUniform("clip1", 4); + this._cullParams.addUniform("clip2", 4); + this._cullParams.addUniform("clip3", 4); + this._finalizeParams = new UniformBuffer(this._engine, undefined, undefined, "GaussianSplattingCullFinalizeParams", false, false); + this._finalizeParams.addUniform("count", 1); + this._finalizeParams.addUniform("indexCount", 1); + this._finalizeParams.addUniform("pad0", 1); + this._finalizeParams.addUniform("pad1", 1); + + // Fixed-size indirect draw-args buffer (also compute-writable in the culling path). + this._indirectArgs = this._createStorage(5 * Uint32Array.BYTES_PER_ELEMENT, Constants.BUFFER_CREATIONFLAG_INDIRECT, "GaussianSplattingSortIndirectArgs"); + this._count = this._createStorage(Uint32Array.BYTES_PER_ELEMENT, 0, "GaussianSplattingCullCount"); + } + + /** + * The GPU indirect draw-args buffer to bind as the mesh's indirect draw buffer. + */ + public get indirectArgsBuffer(): Nullable { + return this._indirectArgs ? this._indirectArgs.getBuffer() : null; + } + + /** + * Writes the draw-indexed-indirect arguments from the CPU. + * @param indexCount indices per instance (subMesh.indexCount) + * @param instanceCount number of instances to draw + */ + public setIndirectArgs(indexCount: number, instanceCount: number): void { + this._ensureShaders(); + this._indirectArgsData[0] = indexCount; + this._indirectArgsData[1] = instanceCount; + this._indirectArgsData[2] = 0; + this._indirectArgsData[3] = 0; + this._indirectArgsData[4] = 0; + this._indirectArgs!.update(this._indirectArgsData); + } + + private _createStorage(byteLength: number, extraFlags: number, label: string): StorageBuffer { + return new StorageBuffer(this._engine as WebGPUEngine, byteLength, Constants.BUFFER_CREATIONFLAG_READWRITE | extraFlags, label); + } + + /** + * Uploads the source splat data. Call whenever the positions or the active source-index set change. + * @param positions splat centers, stride 4 (xyz + 1) + * @param seedIndices packed active source indices, one f32 per slot, padded to a multiple of 16 + * @param renderedCount number of active (non-padding) slots + */ + public setSource(positions: Float32Array, seedIndices: Float32Array, renderedCount: number): void { + this._ensureShaders(); + + this._renderedCount = renderedCount; + this._paddedCount = seedIndices.length; + this._numBuckets = GaussianSplattingGpuSorter._BucketCount(renderedCount); + + if (positions.length > this._positionCapacity) { + this._positions?.dispose(); + this._positions = this._createStorage(positions.length * Float32Array.BYTES_PER_ELEMENT, 0, "GaussianSplattingSortPositions"); + this._positionCapacity = positions.length; + } + + if (this._paddedCount > this._paddedCapacity) { + this._seed?.dispose(); + this._depthBuffer?.dispose(); + this._bucket?.dispose(); + this._sortedIndexBuffer?.dispose(); + const f32Bytes = this._paddedCount * Float32Array.BYTES_PER_ELEMENT; + this._seed = this._createStorage(f32Bytes, 0, "GaussianSplattingSortSeed"); + this._depthBuffer = this._createStorage(f32Bytes, 0, "GaussianSplattingSortDepth"); + this._bucket = this._createStorage(this._paddedCount * Uint32Array.BYTES_PER_ELEMENT, 0, "GaussianSplattingSortBucket"); + // The sorted-index buffer is consumed as an instanced vertex buffer, so it also needs the VERTEX usage flag. + this._sortedIndexBuffer = this._createStorage(f32Bytes, Constants.BUFFER_CREATIONFLAG_VERTEX, "GaussianSplattingSortSorted"); + this._paddedCapacity = this._paddedCount; + } + + // Histogram/offsets buffer, padded to a workgroup multiple so the scan never reads out of bounds. + const histCapacity = Math.ceil(this._numBuckets / WorkgroupSize) * WorkgroupSize; + if (histCapacity > this._histCapacity) { + this._histogramBuffer?.dispose(); + this._histogramBuffer = this._createStorage(histCapacity * Uint32Array.BYTES_PER_ELEMENT, 0, "GaussianSplattingSortHistogram"); + this._histCapacity = histCapacity; + } + + if (!this._range) { + this._range = this._createStorage(2 * Int32Array.BYTES_PER_ELEMENT, 0, "GaussianSplattingSortRange"); + } + + // Culling buffers: flag/offsets padded to a workgroup multiple so the compaction scan stays in bounds; + // drawList is the compacted (visible, sorted) index buffer consumed as the instanced vertex buffer. + const cullCapacity = Math.ceil(this._paddedCount / WorkgroupSize) * WorkgroupSize; + if (cullCapacity > this._cullCapacity) { + this._flag?.dispose(); + this._offsets?.dispose(); + this._drawList?.dispose(); + this._flag = this._createStorage(cullCapacity * Uint32Array.BYTES_PER_ELEMENT, 0, "GaussianSplattingCullFlag"); + this._offsets = this._createStorage(cullCapacity * Uint32Array.BYTES_PER_ELEMENT, 0, "GaussianSplattingCullOffsets"); + this._drawList = this._createStorage(cullCapacity * Float32Array.BYTES_PER_ELEMENT, Constants.BUFFER_CREATIONFLAG_VERTEX, "GaussianSplattingCullDrawList"); + this._cullCapacity = cullCapacity; + } + + this._positions!.update(positions); + this._seed!.update(seedIndices); + this._hasSource = true; + } + + /** + * Whether the culling compute shaders are compiled and ready to dispatch. + * @returns true when a cull pass can run + */ + public isCullReady(): boolean { + if (!this._shadersLoaded) { + return false; + } + return !!this._cull && !!this._cull.isReady() && !!this._compact!.isReady() && !!this._finalize!.isReady() && !!this._cullPrefixSum!.isReady(); + } + + /** + * Runs a post-sort frustum cull + compaction. Uses a dedicated clip-matrix UBO and prefix-sum instance so it + * can run every frame independently of {@link sort}. Produces the compacted visible draw list + * ({@link culledDrawBuffer}) and writes the draw-indexed-indirect arguments with a GPU-computed instance count. + * @param indexCount indices per instance (subMesh.indexCount) + * @param clipMatrix local->clip transform (world * view * projection) used for the frustum test + * @returns true when the cull passes were dispatched; false when shaders are not ready + */ + public cull(indexCount: number, clipMatrix: Matrix): boolean { + if (!this._hasSource || this._renderedCount === 0 || !this.isCullReady()) { + return false; + } + + // The cull owns its clip-matrix UBO and uploads it immediately before dispatching (mirroring sort()'s + // update-then-dispatch adjacency). Sharing sort()'s params UBO left the cull reading a frame-1 snapshot. + const cm = clipMatrix.m; + this._cullParams!.updateUInt("count", this._renderedCount); + this._cullParams!.updateUInt("pad0", 0); + this._cullParams!.updateUInt("pad1", 0); + this._cullParams!.updateUInt("pad2", 0); + this._cullParams!.updateFloat4("clip0", cm[0], cm[1], cm[2], cm[3]); + this._cullParams!.updateFloat4("clip1", cm[4], cm[5], cm[6], cm[7]); + this._cullParams!.updateFloat4("clip2", cm[8], cm[9], cm[10], cm[11]); + this._cullParams!.updateFloat4("clip3", cm[12], cm[13], cm[14], cm[15]); + this._cullParams!.update(); + + this._compactParams!.updateUInt("count", this._renderedCount); + this._compactParams!.update(); + + this._finalizeParams!.updateUInt("count", this._renderedCount); + this._finalizeParams!.updateUInt("indexCount", indexCount); + this._finalizeParams!.update(); + + const groups = Math.ceil(this._renderedCount / WorkgroupSize); + + this._cull!.setStorageBuffer("positions", this._positions!); + this._cull!.setStorageBuffer("sortedIndex", this._sortedIndexBuffer!); + this._cull!.setStorageBuffer("flag", this._flag!); + this._cull!.setStorageBuffer("offsets", this._offsets!); + this._cull!.setUniformBuffer("params", this._cullParams!); + this._cull!.dispatch(groups); + + // Cull-owned prefix-sum ring (see constructor): reset and scan on the dedicated instance so it never + // collides with the sort's ring and always scans the cull's element count. + this._cullPrefixSum!.resetForFrame(); + this._cullPrefixSum!.scanExclusive(this._offsets!, this._renderedCount); + + this._compact!.setStorageBuffer("sortedIndex", this._sortedIndexBuffer!); + this._compact!.setStorageBuffer("flag", this._flag!); + this._compact!.setStorageBuffer("offsets", this._offsets!); + this._compact!.setStorageBuffer("drawList", this._drawList!); + this._compact!.setUniformBuffer("params", this._compactParams!); + this._compact!.dispatch(groups); + + this._finalize!.setStorageBuffer("flag", this._flag!); + this._finalize!.setStorageBuffer("offsets", this._offsets!); + this._finalize!.setStorageBuffer("countOut", this._count!); + this._finalize!.setStorageBuffer("indirectArgs", this._indirectArgs!); + this._finalize!.setStorageBuffer("drawList", this._drawList!); + this._finalize!.setUniformBuffer("params", this._finalizeParams!); + this._finalize!.dispatch(1); + + return true; + } + + /** + * The compacted visible draw list to bind as the `splatIndex` instanced vertex buffer when culling is active. + */ + public get culledDrawBuffer(): Nullable { + return this._drawList ? this._drawList.getBuffer() : null; + } + + /** + * Whether all compute shaders are compiled and ready to dispatch. + * @returns true when a sort can run + */ + public isReady(): boolean { + if (!this._shadersLoaded || !this._clear) { + return false; + } + return !!this._clear.isReady() && !!this._depth!.isReady() && !!this._histogram!.isReady() && !!this._scatter!.isReady() && !!this._prefixSum!.isReady(); + } + + /** + * Runs a depth sort for a camera view. The depth coefficients (a,b,c,d) must already fold in the + * world matrix, camera forward, camera position, and the right-handed depth sign, so that + * depth = a*x + b*y + c*z + d matches the CPU worker's convention. + * @param a depth coefficient for x + * @param b depth coefficient for y + * @param c depth coefficient for z + * @param d constant depth offset + * @returns true when the sort was dispatched; false when shaders are not ready or no source is set + */ + public sort(a: number, b: number, c: number, d: number): boolean { + if (!this._hasSource || this._renderedCount === 0) { + return false; + } + if (!this.isReady()) { + return false; + } + + this._params!.updateUInt("count", this._renderedCount); + this._params!.updateUInt("paddedCount", this._paddedCount); + this._params!.updateUInt("numBuckets", this._numBuckets); + this._params!.updateUInt("pad", 0); + this._params!.updateFloat4("coeff", a, b, c, d); + this._params!.update(); + + // Seed the atomic depth range with (+Inf, -Inf) before the reduction. + this._range!.update(RangeInit); + + const renderGroups = Math.ceil(this._renderedCount / WorkgroupSize); + const paddedGroups = Math.ceil(this._paddedCount / WorkgroupSize); + + this._clear!.setStorageBuffer("histogram", this._histogramBuffer!); + this._clear!.setUniformBuffer("params", this._params!); + this._clear!.dispatch(Math.ceil(this._numBuckets / WorkgroupSize)); + + this._depth!.setStorageBuffer("positions", this._positions!); + this._depth!.setStorageBuffer("seed", this._seed!); + this._depth!.setStorageBuffer("depth", this._depthBuffer!); + this._depth!.setStorageBuffer("range", this._range!); + this._depth!.setUniformBuffer("params", this._params!); + this._depth!.dispatch(renderGroups); + + this._histogram!.setStorageBuffer("depth", this._depthBuffer!); + this._histogram!.setStorageBuffer("range", this._range!); + this._histogram!.setStorageBuffer("bucket", this._bucket!); + this._histogram!.setStorageBuffer("histogram", this._histogramBuffer!); + this._histogram!.setUniformBuffer("params", this._params!); + this._histogram!.dispatch(renderGroups); + + this._prefixSum!.resetForFrame(); + this._prefixSum!.scanExclusive(this._histogramBuffer!, this._numBuckets); + + this._scatter!.setStorageBuffer("bucket", this._bucket!); + this._scatter!.setStorageBuffer("seed", this._seed!); + this._scatter!.setStorageBuffer("offsets", this._histogramBuffer!); + this._scatter!.setStorageBuffer("sortedIndex", this._sortedIndexBuffer!); + this._scatter!.setUniformBuffer("params", this._params!); + this._scatter!.dispatch(paddedGroups); + + return true; + } + + /** + * The GPU data buffer to bind as the `splatIndex` instanced vertex buffer, or null before allocation. + */ + public get sortedDataBuffer(): Nullable { + return this._sortedIndexBuffer ? this._sortedIndexBuffer.getBuffer() : null; + } + + /** + * Releases all GPU resources held by the sorter. + */ + public dispose(): void { + this._positions?.dispose(); + this._seed?.dispose(); + this._depthBuffer?.dispose(); + this._bucket?.dispose(); + this._range?.dispose(); + this._histogramBuffer?.dispose(); + this._sortedIndexBuffer?.dispose(); + this._indirectArgs?.dispose(); + this._flag?.dispose(); + this._offsets?.dispose(); + this._drawList?.dispose(); + this._count?.dispose(); + this._params?.dispose(); + this._compactParams?.dispose(); + this._cullParams?.dispose(); + this._finalizeParams?.dispose(); + this._prefixSum?.dispose(); + this._cullPrefixSum?.dispose(); + this._positions = null; + this._seed = null; + this._depthBuffer = null; + this._bucket = null; + this._range = null; + this._histogramBuffer = null; + this._sortedIndexBuffer = null; + this._indirectArgs = null; + this._flag = null; + this._offsets = null; + this._drawList = null; + this._count = null; + this._params = null; + this._compactParams = null; + this._cullParams = null; + this._finalizeParams = null; + this._prefixSum = null; + this._cullPrefixSum = null; + this._clear = null; + this._depth = null; + this._histogram = null; + this._scatter = null; + this._cull = null; + this._compact = null; + this._finalize = null; + this._positionCapacity = 0; + this._paddedCapacity = 0; + this._histCapacity = 0; + this._cullCapacity = 0; + this._hasSource = false; + } +} diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts index 956cb1361422..8c4ba2b1048f 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts @@ -27,6 +27,8 @@ import { type INative } from "core/Engines/Native/nativeInterfaces"; import { type AbstractEngine } from "core/Engines/abstractEngine.pure"; import { type ICustomAnimationFrameRequester } from "core/Misc/customAnimationFrameRequester"; import { GaussianSplattingSortWorker, GaussianSplattingSortWorkerCommand } from "./gaussianSplattingSortWorker"; +import { GaussianSplattingGpuSorter } from "./gaussianSplattingGpuSorter"; +import { type DataBuffer } from "core/Buffers/dataBuffer"; // eslint-disable-next-line @typescript-eslint/naming-convention declare const _native: INative; @@ -58,6 +60,23 @@ interface ISogPackInternal { const IsNative = typeof _native !== "undefined"; const Native = IsNative ? _native : null; + +/** + * Depth-sort backend selection for Gaussian Splatting meshes. + * @see GaussianSplattingMeshBase.GpuSortMode + */ +export enum GaussianSplattingGpuSortMode { + /** + * Pick automatically. Currently uses the CPU worker / Babylon Native path; the WebGPU compute sort is opt-in + * via {@link GaussianSplattingGpuSortMode.Gpu} while it gains broader validation. + */ + Auto, + /** Force the WebGPU compute sort path. Falls back to the worker when compute shaders are unavailable. */ + Gpu, + /** Force the legacy CPU worker (or Babylon Native) sort path. */ + Worker, +} + interface IDelayedTextureUpdate { covA: Uint16Array; covB: Uint16Array; @@ -485,9 +504,58 @@ export class GaussianSplattingMeshBase extends Mesh { */ public static LogSortPerformance = false; + /** + * Selects the depth-sort backend used by every Gaussian Splatting mesh. + * Defaults to {@link GaussianSplattingGpuSortMode.Auto}. The CPU worker / native path is always kept as a + * fallback: when the WebGPU compute path is requested but unsupported (or not applicable, e.g. compound + * meshes), the mesh transparently falls back to the worker. + */ + public static GpuSortMode: GaussianSplattingGpuSortMode = GaussianSplattingGpuSortMode.Auto; + + /** + * When true, the WebGPU GPU-sort path drives the draw via a GPU indirect-args buffer (drawIndexedIndirect), + * which the culling pipeline uses to reduce the drawn instance count. Set false to fall back to a CPU-set + * instance count (draws all active splats). Only affects meshes using the GPU sort path. + */ + public static EnableGpuIndirectDraw = true; + + /** + * When true (and {@link EnableGpuIndirectDraw} is on), the WebGPU GPU-sort path runs a GPU frustum cull and + * draws only the visible splats via an indirect draw. The cull test matches the vertex shader exactly (same + * local->clip transform and 1.2*w bounds), so it never removes a splat the shader would have shown. Set false + * to draw all sorted splats. + */ + public static EnableGpuCulling = true; + + /** + * Number of frames the camera/world must stay stationary before the GPU frustum cull re-engages after a move. + * While the camera is moving (and for this many frames after it stops) the full sorted splat list is drawn + * instead of the culled list: this keeps culling off the expensive motion frames (where the depth sort also + * re-runs) and avoids edge popping, and gives the per-frame cull compute a frame to converge (its UBO upload + * reaches the shader one dispatch late). Set to 0 to cull every frame regardless of motion. + */ + public static GpuCullResumeDelayFrames = 2; + /** @internal */ public _vertexCount = 0; protected _worker: Nullable = null; + // WebGPU compute sort driver, created only when the GPU path is active for this mesh. When null the mesh + // uses the CPU worker (or Babylon Native) sort path. + private _gpuSorter: Nullable = null; + // Resolved once per (re)build in _instantiateWorker: true when this mesh currently uses the GPU sort path. + private _useGpuSort = false; + // Set when the packed source-index seed changed and the GPU sorted-index buffer must be regenerated. + private _gpuSortDirty = false; + // Last GPU data buffer bound as the splatIndex vertex buffer; its identity changes when the buffer is grown. + private _gpuSortedDataBuffer: Nullable = null; + // Whether the GPU frustum-cull pass ran for the current sort (culling active for this mesh this frame). + private _gpuCulling = false; + // Frames the camera/world has been stationary. Reset to 0 on any move; gates when culling re-engages + // (see GpuCullResumeDelayFrames). + private _gpuCullSettleFrames = 0; + // Reused matrices for the GPU cull local->clip transform (avoids per-frame allocation). + private _gpuClipMatrix = Matrix.Identity(); + private _gpuViewProjMatrix = Matrix.Identity(); private _modelViewProjectionMatrix = Matrix.Identity(); private _depthMix: BigInt64Array; protected _canPostToWorker = true; @@ -1320,6 +1388,161 @@ export class GaussianSplattingMeshBase extends Mesh { ); } + /** + * Resolves whether this mesh should use the WebGPU compute sort path for the current build. + * The CPU worker / native path is always available as a fallback, so this returns false whenever the GPU + * path is not applicable (unsupported engine, native, depth sort disabled, or compound meshes which still + * sort on the worker in this phase). + * @returns true when the GPU compute sort should be used + */ + protected _resolveUseGpuSort(): boolean { + const mode = GaussianSplattingMeshBase.GpuSortMode; + if (mode === GaussianSplattingGpuSortMode.Worker) { + return false; + } + if (IsNative || this._disableDepthSort) { + return false; + } + // Compound (rig) meshes need per-part depth coefficients that the GPU scaffold does not handle yet. + if (this.isCompound) { + return false; + } + if (!GaussianSplattingGpuSorter.IsSupported(this._scene.getEngine())) { + return false; + } + // Auto currently keeps the CPU worker as the default; the GPU depth sort is opt-in via GpuSortMode.Gpu + // while it gains broader validation. The worker path remains the fallback in every case above. + return mode === GaussianSplattingGpuSortMode.Gpu; + } + + /** + * Generates (when needed) the GPU depth-sorted `splatIndex` buffer and binds it to each active camera + * mesh. The sort is re-run when the source data changed or when the active camera / world transform moved. + * @param activeViewInfos the per-camera view infos to bind the buffer to + * @param frameId the current frame id + */ + private _updateGpuSortBuffers(activeViewInfos: ICameraViewInfo[], frameId: number): void { + const sorter = this._gpuSorter; + if (!sorter || !this._splatIndex || !this._splatPositions || !activeViewInfos.length) { + return; + } + + // Re-upload the source (positions + packed active source indices) whenever the data/seed changed. + if (this._gpuSortDirty) { + sorter.setSource(this._splatPositions, this._splatIndex, this.renderedSplatCount); + } + + // The GPU produces a single ordering per frame, so sort for the camera being rendered. + // Plain loop (no find/closure) to avoid a per-frame allocation in this render-loop path. + const activeCamera = this._scene.activeCamera; + let primary = activeViewInfos[0]; + if (activeCamera) { + for (let i = 0; i < activeViewInfos.length; i++) { + if (activeViewInfos[i].camera === activeCamera) { + primary = activeViewInfos[i]; + break; + } + } + } + const camera = primary.camera; + const worldMatrix = this.computeWorldMatrix(true); + + // Re-sort when the data changed, when this view has not been displayed yet, or when the camera/world moved. + const cullingWanted = GaussianSplattingMeshBase.EnableGpuIndirectDraw && GaussianSplattingMeshBase.EnableGpuCulling; + const cameraMoved = this._isSortStateDirty(primary, worldMatrix, camera); + const needSort = this._gpuSortDirty || !this._readyToDisplay || cameraMoved; + + // Track how long the camera has been stationary so the frustum cull only engages once it settles. + if (cameraMoved) { + this._gpuCullSettleFrames = 0; + } else if (this._gpuCullSettleFrames < 0xffff) { + this._gpuCullSettleFrames++; + } + + // Local->clip matrix (world * view * projection) for the frustum cull pass. Recomputed every frame so the + // cull tracks the camera when it re-engages. + camera.getViewMatrix().multiplyToRef(camera.getProjectionMatrix(), this._gpuViewProjMatrix); + worldMatrix.multiplyToRef(this._gpuViewProjMatrix, this._gpuClipMatrix); + + if (needSort) { + const cameraViewMatrix = camera.getViewMatrix(); + const world = worldMatrix.m; + const fx = cameraViewMatrix.m[2]; + const fy = cameraViewMatrix.m[6]; + const fz = cameraViewMatrix.m[10]; + const cameraPosition = camera.globalPosition; + const camDot = fx * cameraPosition.x + fy * cameraPosition.y + fz * cameraPosition.z; + // Fold the world matrix and camera forward into depth coefficients so depth = a*x + b*y + c*z + d. + // Negate in right-handed scenes so "larger depth = farther" holds (matching the worker's counting sort). + const depthSign = this._scene.useRightHandedSystem ? -1 : 1; + const a = (fx * world[0] + fy * world[1] + fz * world[2]) * depthSign; + const b = (fx * world[4] + fy * world[5] + fz * world[6]) * depthSign; + const c = (fx * world[8] + fy * world[9] + fz * world[10]) * depthSign; + const d = (fx * world[12] + fy * world[13] + fz * world[14] - camDot) * depthSign; + + if (!sorter.sort(a, b, c, d)) { + // Shaders still compiling (or no source yet): keep the dirty state and retry next frame. + return; + } + this._gpuSortDirty = false; + + // Record the sort state on the view we sorted for so it is not re-sorted until the camera/world moves. + primary.sortWorldMatrix.copyFrom(worldMatrix); + primary.sortCameraForward.set(fx, fy, fz); + primary.sortCameraPosition.copyFrom(cameraPosition); + primary.cameraDirection.copyFrom(this._getCameraDirection(camera)); + primary.sortAppliedId = primary.sortRequestId; + } + + // GPU frustum cull. Skipped while the camera is moving (the full sorted list is drawn instead, avoiding + // cull cost on the expensive motion frames and any edge popping). The cull compute is dispatched one frame + // before its result is used so the one-dispatch UBO upload lag has converged (see GpuCullResumeDelayFrames). + const cullDelay = GaussianSplattingMeshBase.GpuCullResumeDelayFrames; + const cullWarming = cullingWanted && !this._gpuSortDirty && this._gpuCullSettleFrames >= Math.max(0, cullDelay - 1); + if (cullWarming) { + this._gpuCulling = sorter.cull(GaussianSplattingMeshBase._BatchSize * 6, this._gpuClipMatrix); + } else { + this._gpuCulling = false; + } + + // The draw uses the compacted visible list only once the cull has settled (its result is converged); + // otherwise it reads the full sorted list with a CPU-written full instance count. + const cullingActive = this._gpuCulling && this._gpuCullSettleFrames >= cullDelay; + const dataBuffer = cullingActive ? sorter.culledDrawBuffer : sorter.sortedDataBuffer; + if (!dataBuffer) { + return; + } + if (dataBuffer !== this._gpuSortedDataBuffer) { + // The bound buffer changed (realloc, or culling toggled on/off), so every camera mesh must rebind. + this._gpuSortedDataBuffer = dataBuffer; + this._cameraViewInfos.forEach((info) => (info.splatIndexBufferSet = false)); + } + + const instanceCount = Math.max(this._splatIndex.length >> 4, 1); + const indexCount = GaussianSplattingMeshBase._BatchSize * 6; + let indirectBuffer: Nullable = null; + if (GaussianSplattingMeshBase.EnableGpuIndirectDraw) { + if (!cullingActive) { + // Non-culled indirect draw: write the full instance count from the CPU. + sorter.setIndirectArgs(indexCount, instanceCount); + } + // When culling is active the cull's finalize pass already wrote the indirect args on the GPU. + indirectBuffer = sorter.indirectArgsBuffer; + } + // Plain loop (no forEach closure) to avoid a per-frame allocation in this render-loop path. + for (let i = 0; i < activeViewInfos.length; i++) { + const info = activeViewInfos[i]; + if (!info.splatIndexBufferSet) { + info.mesh._thinInstanceSetSplatIndexBuffer(dataBuffer, instanceCount); + info.splatIndexBufferSet = true; + } + info.mesh._indirectDrawBuffer = indirectBuffer; + info.frameIdLastUpdate = frameId; + } + this._readyToDisplay = true; + this._canPostToWorker = true; + } + /** @internal */ public _postToWorker(forced = false): void { const scene = this._scene; @@ -1391,6 +1614,14 @@ export class GaussianSplattingMeshBase extends Mesh { return a.frameIdLastUpdate - b.frameIdLastUpdate; }); + // WebGPU compute sort path: generate the sorted-index buffer on the GPU and bind it to each active + // camera mesh. This bypasses the CPU worker / native branches entirely (they remain the fallback and + // are used whenever _useGpuSort is false). + if (this._useGpuSort && this._gpuSorter) { + this._updateGpuSortBuffers(activeViewInfos, frameId); + return; + } + const hasSortFunction = this._worker || Native?.sortSplats || this._disableDepthSort; if ((forced || outdated) && hasSortFunction && (this._scene.activeCameras?.length || this._scene.activeCamera) && this._canPostToWorker) { const worldMatrix = this.computeWorldMatrix(true); @@ -2297,6 +2528,9 @@ export class GaussianSplattingMeshBase extends Mesh { _ReleaseGsInterFrameYield(this.getEngine()); this._worker?.terminate(); this._worker = null; + this._gpuSorter?.dispose(); + this._gpuSorter = null; + this._gpuSortedDataBuffer = null; this.onPartCountChangedObservable.clear(); this.onPartRemovedObservable.clear(); @@ -2651,6 +2885,8 @@ export class GaussianSplattingMeshBase extends Mesh { // Re-sync the active interval set in case the source splat count changed. this._postIntervalsToWorker(); } + // The GPU sort path (when active) re-uploads its source positions/indices on the next sort. + this._gpuSortDirty = true; this._postToWorker(true); } } @@ -2694,6 +2930,8 @@ export class GaussianSplattingMeshBase extends Mesh { this._postIntervalsToWorker(); } this._sortIsDirty = true; + // The GPU sort path (when active) re-uploads its source positions/indices on the next sort. + this._gpuSortDirty = true; } /** @@ -2714,6 +2952,7 @@ export class GaussianSplattingMeshBase extends Mesh { const data = this._splatPositions.slice(floatOffset, floatOffset + splatCount * 4); this._worker.postMessage({ command: GaussianSplattingSortWorkerCommand.POSITIONS_UPDATE, offset: floatOffset, data }, [data.buffer]); this._sortIsDirty = true; + this._gpuSortDirty = true; } private *_updateData( @@ -2956,6 +3195,9 @@ export class GaussianSplattingMeshBase extends Mesh { } this.forcedInstanceCount = Math.max(paddedVertexCount >> 4, 1); + + // The packed source-index seed changed: the GPU sorted-index buffer (when active) must be regenerated. + this._gpuSortDirty = true; } protected _updateTextureFromData = (texture: BaseTexture, data: ArrayBufferView, width: number, lineStart: number, lineCount: number) => { @@ -3059,6 +3301,32 @@ export class GaussianSplattingMeshBase extends Mesh { } this._updateSplatIndexBuffer(this._vertexCount); + // Resolve the sort backend. When the WebGPU compute path is active, no CPU worker is created: the + // sorted-index buffer is produced on the GPU and bound directly as the splatIndex vertex buffer. + this._useGpuSort = this._resolveUseGpuSort(); + if (this._useGpuSort) { + this._worker?.terminate(); + this._worker = null; + this._canPostToWorker = true; + if (!this._gpuSorter) { + this._gpuSorter = new GaussianSplattingGpuSorter(this._scene.getEngine()); + } + this._gpuSortDirty = true; + return; + } + + // Falling back to the CPU worker / Babylon Native path. If this mesh was previously on the GPU sort path, + // clear that draw state so the worker rebinds a CPU splatIndex buffer (the GPU buffer wrapper is not + // updatable, so thinInstanceBufferUpdated would silently fail) and Mesh._draw stops issuing GPU indirect + // draws from the now-unused sorter's indirect-args buffer. + if (this._gpuSortedDataBuffer) { + this._gpuSortedDataBuffer = null; + this._cameraViewInfos.forEach((info) => { + info.splatIndexBufferSet = false; + info.mesh._indirectDrawBuffer = null; + }); + } + // no worker in native if (IsNative) { return; diff --git a/packages/dev/core/src/Meshes/mesh.pure.ts b/packages/dev/core/src/Meshes/mesh.pure.ts index d78533e17241..1d6b1c9b5b8a 100644 --- a/packages/dev/core/src/Meshes/mesh.pure.ts +++ b/packages/dev/core/src/Meshes/mesh.pure.ts @@ -205,6 +205,9 @@ class _InternalMeshDataInfo { public meshMap: Nullable<{ [id: string]: Mesh | undefined }> = null; public _preActivateId: number = -1; + // Optional GPU indirect draw-args buffer (Gaussian Splatting GPU culling path). Stored here rather than as a + // direct Mesh field to keep Mesh under the V8 fast-property own-property threshold (see dictionaryMode test). + public _indirectDrawBuffer: Nullable = null; // eslint-disable-next-line @typescript-eslint/naming-convention public _LODLevels = new Array(); /** Alternative definition of LOD level, using screen coverage instead of distance */ @@ -539,6 +542,21 @@ export class Mesh extends AbstractMesh implements IGetSetVerticesData { this._internalMeshDataInfo._forcedInstanceCount = count; } + /** + * @internal + * Optional GPU indirect-args buffer (5 u32: indexCount, instanceCount, firstIndex, baseVertex, firstInstance). + * When set, indexed draws for this mesh issue drawIndexedIndirect from this buffer instead of a CPU instance + * count. Used by the Gaussian Splatting GPU culling path (WebGPU only). Backed by `_internalMeshDataInfo` (via + * an accessor) so it does not add an own property to every Mesh instance (V8 fast-property threshold). + */ + public get _indirectDrawBuffer(): Nullable { + return this._internalMeshDataInfo._indirectDrawBuffer; + } + + public set _indirectDrawBuffer(value: Nullable) { + this._internalMeshDataInfo._indirectDrawBuffer = value; + } + /** @internal */ public _instanceDataStorage: _InstanceDataStorage; @@ -2141,6 +2159,9 @@ export class Mesh extends AbstractMesh implements IGetSetVerticesData { } else if (useVertexPulling) { // We're rendering the number of indices in the index buffer but the vertex shader is handling the data itself. engine.drawArraysType(fillMode, subMesh.indexStart, subMesh.indexCount, this.forcedInstanceCount || instancesCount); + } else if (this._indirectDrawBuffer) { + // GPU-driven indirect draw: the instance count is computed on the GPU (e.g. Gaussian Splatting culling). + engine.drawElementsInstancedIndirect(fillMode, subMesh.indexStart, subMesh.indexCount, this._indirectDrawBuffer); } else { engine.drawElementsType(fillMode, subMesh.indexStart, subMesh.indexCount, this.forcedInstanceCount || instancesCount); } diff --git a/packages/dev/core/src/Meshes/thinInstanceMesh.pure.ts b/packages/dev/core/src/Meshes/thinInstanceMesh.pure.ts index 00055b1c1763..cc93d3382444 100644 --- a/packages/dev/core/src/Meshes/thinInstanceMesh.pure.ts +++ b/packages/dev/core/src/Meshes/thinInstanceMesh.pure.ts @@ -6,6 +6,7 @@ import { Logger } from "../Misc/logger"; import { BoundingInfo } from "core/Culling/boundingInfo"; import { Mesh } from "../Meshes/mesh.pure"; import { VertexBuffer, Buffer } from "../Buffers/buffer.pure"; +import { type DataBuffer } from "../Buffers/dataBuffer"; const BakedVertexAnimationSettingsInstancedKind = "bakedVertexAnimationSettingsInstanced"; @@ -378,6 +379,27 @@ export function RegisterThinInstanceMesh(): void { } }; + Mesh.prototype._thinInstanceSetSplatIndexBuffer = function (dataBuffer: Nullable, instanceCount: number): void { + if (!dataBuffer) { + return; + } + // GPU-produced counterpart of the `splatIndex` branch in thinInstanceSetBuffer: instead of wrapping a CPU + // Float32Array, it binds an existing GPU DataBuffer (typically a StorageBuffer written by a compute pass) + // as the four `splatIndex0..3` instanced vertex attributes, avoiding any CPU readback. + this._thinInstanceInitializeUserStorage(); + this._thinInstanceDataStorage.instancesCount = instanceCount; + this._thinInstanceDataStorage.matrixBuffer?.dispose(); + const splatInstancesBuffer = new Buffer(this.getEngine(), dataBuffer, false, 16, false, true); + // The passed DataBuffer is owned and reused across frames by its producer (e.g. the GPU splat sorter), + // not by this wrapper. Take a reference so the wrapper's dispose() (on the next rebind) only releases our + // borrow instead of destroying the shared buffer while the producer's compute passes still read it. + dataBuffer.references++; + this._thinInstanceDataStorage.matrixBuffer = splatInstancesBuffer; + for (let i = 0; i < 4; i++) { + this.setVerticesBuffer(splatInstancesBuffer.createVertexBuffer("splatIndex" + i, i * 4, 4)); + } + }; + Mesh.prototype.thinInstanceBufferUpdated = function (kind: string): void { if (kind === "matrix") { if (this.thinInstanceAllowAutomaticStaticBufferRecreation && this._thinInstanceDataStorage.matrixBuffer && !this._thinInstanceDataStorage.matrixBuffer.isUpdatable()) { diff --git a/packages/dev/core/src/Meshes/thinInstanceMesh.types.ts b/packages/dev/core/src/Meshes/thinInstanceMesh.types.ts index 696628ed9d07..5d08af3e8a2b 100644 --- a/packages/dev/core/src/Meshes/thinInstanceMesh.types.ts +++ b/packages/dev/core/src/Meshes/thinInstanceMesh.types.ts @@ -1,6 +1,7 @@ import { type Nullable, type DeepImmutableObject } from "../types"; import { type Matrix } from "../Maths/math.vector"; import { type VertexBuffer, type Buffer } from "../Buffers/buffer"; +import { type DataBuffer } from "../Buffers/dataBuffer"; declare module "./mesh.pure" { /** @internal */ // eslint-disable-next-line @typescript-eslint/naming-convention @@ -108,6 +109,9 @@ declare module "./mesh.pure" { /** @internal */ _thinInstanceInitializeUserStorage(): void; + /** @internal */ + _thinInstanceSetSplatIndexBuffer(dataBuffer: Nullable, instanceCount: number): void; + /** @internal */ _thinInstanceUpdateBufferSize(kind: string, numInstances?: number): void; diff --git a/packages/dev/core/src/ShadersWGSL/gaussianSplatting.vertex.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplatting.vertex.fx index 526ce932261f..1e29086e3f64 100644 --- a/packages/dev/core/src/ShadersWGSL/gaussianSplatting.vertex.fx +++ b/packages/dev/core/src/ShadersWGSL/gaussianSplatting.vertex.fx @@ -84,7 +84,11 @@ fn main(input : VertexInputs) -> FragmentInputs { let splatIndex: f32 = getSplatIndex(i32(vertexInputs.position.z + 0.5), vertexInputs.splatIndex0, vertexInputs.splatIndex1, vertexInputs.splatIndex2, vertexInputs.splatIndex3); - var splat: Splat = readSplat(splatIndex, uniforms.dataTextureSize); + // A negative index is the GPU-culling padding sentinel: the last indirect-draw instance has < 16 real splats, + // and the remaining slots are set to -1 so this vertex is clipped and renders nothing. + let isPaddingSplat: bool = splatIndex < 0.0; + + var splat: Splat = readSplat(max(splatIndex, 0.0), uniforms.dataTextureSize); var covA: vec3f = splat.covA.xyz; var covB: vec3f = vec3f(splat.covA.w, splat.covB.xy); @@ -143,6 +147,11 @@ fn main(input : VertexInputs) -> FragmentInputs { vertexOutputs.position = gaussianSplatting(vertexInputs.position.xy, worldPos.xyz, scale, covA, covB, splatWorld, scene.view, scene.projection, uniforms.focal, uniforms.invViewport, uniforms.kernelSize, uniforms.minPixelSize); + // Clip GPU-culling padding vertices (sentinel index) so the last partial instance renders nothing. + if (isPaddingSplat) { + vertexOutputs.position = vec4f(0.0, 0.0, 2.0, 1.0); + } + #include #include #include diff --git a/packages/dev/core/src/ShadersWGSL/gaussianSplattingCull.compute.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplattingCull.compute.fx new file mode 100644 index 000000000000..c9c38b294f23 --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/gaussianSplattingCull.compute.fx @@ -0,0 +1,43 @@ +// Gaussian Splatting GPU culling - flag pass (post-sort). +// +// For each already-sorted render slot, fetches the splat center, projects it to clip space with the same +// local->clip matrix and the same 1.2*w frustum bounds the vertex shader uses (gaussianSplatting.fx), and +// writes a visibility flag. Because the test matches the vertex shader exactly, culled splats are precisely +// those the shader would clip anyway (no edge popping). The flag is written to two buffers: one preserved for +// the compaction pass, one to be exclusive-scanned into output offsets. + +struct Params { + count : u32, + pad0 : u32, + pad1 : u32, + pad2 : u32, + clip0 : vec4f, + clip1 : vec4f, + clip2 : vec4f, + clip3 : vec4f, +}; + +@group(0) @binding(0) var positions : array; +@group(0) @binding(1) var sortedIndex : array; +@group(0) @binding(2) var flag : array; +@group(0) @binding(3) var offsets : array; +@group(0) @binding(4) var params : Params; + +@compute @workgroup_size(256, 1, 1) +fn main(@builtin(global_invocation_id) gid : vec3u) { + let i = gid.x; + if (i >= params.count) { + return; + } + let src = u32(sortedIndex[i] + 0.5); + let o = src * 4u; + // clip = M * (x,y,z,1) built from the local->clip matrix columns. + let clipPos = params.clip0 * positions[o] + params.clip1 * positions[o + 1u] + params.clip2 * positions[o + 2u] + params.clip3; + let bounds = 1.2 * clipPos.w; + var visible = 1u; + if (clipPos.z < 0.0 || clipPos.x < -bounds || clipPos.x > bounds || clipPos.y < -bounds || clipPos.y > bounds) { + visible = 0u; + } + flag[i] = visible; + offsets[i] = visible; +} diff --git a/packages/dev/core/src/ShadersWGSL/gaussianSplattingCullCompact.compute.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplattingCullCompact.compute.fx new file mode 100644 index 000000000000..07a1a5a02c8f --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/gaussianSplattingCullCompact.compute.fx @@ -0,0 +1,25 @@ +// Gaussian Splatting GPU culling - compaction pass (post-sort). +// +// Stream-compacts the visible sorted indices into a dense draw list, preserving back-to-front order. `offsets` +// holds the exclusive prefix sum of the visibility flags, so each visible slot writes to its packed position. + +struct Params { + count : u32, +}; + +@group(0) @binding(0) var sortedIndex : array; +@group(0) @binding(1) var flag : array; +@group(0) @binding(2) var offsets : array; +@group(0) @binding(3) var drawList : array; +@group(0) @binding(4) var params : Params; + +@compute @workgroup_size(256, 1, 1) +fn main(@builtin(global_invocation_id) gid : vec3u) { + let i = gid.x; + if (i >= params.count) { + return; + } + if (flag[i] == 1u) { + drawList[offsets[i]] = sortedIndex[i]; + } +} diff --git a/packages/dev/core/src/ShadersWGSL/gaussianSplattingCullFinalize.compute.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplattingCullFinalize.compute.fx new file mode 100644 index 000000000000..16605611b33b --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/gaussianSplattingCullFinalize.compute.fx @@ -0,0 +1,46 @@ +// Gaussian Splatting GPU culling - finalize pass (post-sort). +// +// Computes the visible splat count from the scanned flags, writes the draw-indexed-indirect arguments +// (instanceCount = ceil(visible / 16)) and the visible count, then zero-pads the draw list's final partial +// instance so the indirect draw never renders stale indices in the last (< 16 splats) instance. + +struct Params { + count : u32, + indexCount : u32, + pad0 : u32, + pad1 : u32, +}; + +@group(0) @binding(0) var flag : array; +@group(0) @binding(1) var offsets : array; +@group(0) @binding(2) var countOut : array; +@group(0) @binding(3) var indirectArgs : array; +@group(0) @binding(4) var drawList : array; +@group(0) @binding(5) var params : Params; + +@compute @workgroup_size(64, 1, 1) +fn main(@builtin(global_invocation_id) gid : vec3u) { + let n = params.count; + var visible = 0u; + if (n > 0u) { + visible = offsets[n - 1u] + flag[n - 1u]; + } + let padded = (visible + 15u) & ~15u; + + let t = gid.x; + if (t == 0u) { + countOut[0] = visible; + indirectArgs[0] = params.indexCount; + indirectArgs[1] = padded / 16u; + indirectArgs[2] = 0u; + indirectArgs[3] = 0u; + indirectArgs[4] = 0u; + } + // Point the last partial instance's padding slots at a sentinel index (-1) that the vertex shader clips, + // so they render nothing. (Index 0 — as the non-culled path uses — would render a real splat here, which is + // wrong when that splat is culled: it would appear as an orphan in the middle of the view.) + let idx = visible + t; + if (idx < padded) { + drawList[idx] = -1.0; + } +} diff --git a/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortClear.compute.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortClear.compute.fx new file mode 100644 index 000000000000..c6d6aa7c8179 --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortClear.compute.fx @@ -0,0 +1,20 @@ +// Gaussian Splatting GPU depth sort - clear pass. +// Zeroes the per-bucket histogram before the histogram pass accumulates into it. + +struct Params { + count : u32, + paddedCount : u32, + numBuckets : u32, + pad : u32, + coeff : vec4f, +}; + +@group(0) @binding(0) var histogram : array>; +@group(0) @binding(1) var params : Params; + +@compute @workgroup_size(256, 1, 1) +fn main(@builtin(global_invocation_id) gid : vec3u) { + if (gid.x < params.numBuckets) { + atomicStore(&histogram[gid.x], 0u); + } +} diff --git a/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortDepth.compute.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortDepth.compute.fx new file mode 100644 index 000000000000..1fb3f2b56f03 --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortDepth.compute.fx @@ -0,0 +1,62 @@ +// Gaussian Splatting GPU depth sort - depth + range pass. +// +// For each active render slot, reads its source splat index from `seed`, fetches the splat center from +// `positions` (stride 4: xyz + 1), and computes the signed camera-forward depth using precomputed +// coefficients (coeff = a,b,c,d, already multiplied by the right-handed sign on the CPU): +// depth = a*x + b*y + c*z + d +// It stores the depth per slot and atomically tracks the min/max depth across all active slots so the +// histogram pass can map depths to buckets. + +struct Params { + count : u32, + paddedCount : u32, + numBuckets : u32, + pad : u32, + coeff : vec4f, +}; + +@group(0) @binding(0) var positions : array; +@group(0) @binding(1) var seed : array; +@group(0) @binding(2) var depth : array; +@group(0) @binding(3) var range : array>; +@group(0) @binding(4) var params : Params; + +fn atomicMinFloat(atomicVar : ptr, read_write>, value : f32) { + let intValue = bitcast(value); + loop { + let oldIntValue = atomicLoad(atomicVar); + if (value >= bitcast(oldIntValue)) { + break; + } + if (atomicCompareExchangeWeak(atomicVar, oldIntValue, intValue).old_value == oldIntValue) { + break; + } + } +} + +fn atomicMaxFloat(atomicVar : ptr, read_write>, value : f32) { + let intValue = bitcast(value); + loop { + let oldIntValue = atomicLoad(atomicVar); + if (value <= bitcast(oldIntValue)) { + break; + } + if (atomicCompareExchangeWeak(atomicVar, oldIntValue, intValue).old_value == oldIntValue) { + break; + } + } +} + +@compute @workgroup_size(256, 1, 1) +fn main(@builtin(global_invocation_id) gid : vec3u) { + let i = gid.x; + if (i >= params.count) { + return; + } + let src = u32(seed[i] + 0.5); + let o = src * 4u; + let d = params.coeff.x * positions[o] + params.coeff.y * positions[o + 1u] + params.coeff.z * positions[o + 2u] + params.coeff.w; + depth[i] = d; + atomicMinFloat(&range[0], d); + atomicMaxFloat(&range[1], d); +} diff --git a/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortHistogram.compute.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortHistogram.compute.fx new file mode 100644 index 000000000000..b170d04dbdef --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortHistogram.compute.fx @@ -0,0 +1,46 @@ +// Gaussian Splatting GPU depth sort - histogram pass. +// +// Maps each active slot's depth to a bucket in [0, numBuckets) and accumulates the per-bucket count. +// The farthest splat (max depth) maps to bucket 0 so the subsequent ascending scatter produces +// back-to-front order for correct alpha blending (matching the CPU worker's counting sort). + +struct Params { + count : u32, + paddedCount : u32, + numBuckets : u32, + pad : u32, + coeff : vec4f, +}; + +@group(0) @binding(0) var depth : array; +@group(0) @binding(1) var range : array; +@group(0) @binding(2) var bucket : array; +@group(0) @binding(3) var histogram : array>; +@group(0) @binding(4) var params : Params; + +@compute @workgroup_size(256, 1, 1) +fn main(@builtin(global_invocation_id) gid : vec3u) { + let i = gid.x; + if (i >= params.count) { + return; + } + let minDepth = bitcast(range[0]); + let maxDepth = bitcast(range[1]); + let span = maxDepth - minDepth; + + var b : u32 = 0u; + if (span > 1e-12) { + let scale = f32(params.numBuckets - 1u) / span; + var k = i32((maxDepth - depth[i]) * scale); + if (k < 0) { + k = 0; + } + let maxBucket = i32(params.numBuckets - 1u); + if (k > maxBucket) { + k = maxBucket; + } + b = u32(k); + } + bucket[i] = b; + atomicAdd(&histogram[b], 1u); +} diff --git a/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortScatter.compute.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortScatter.compute.fx new file mode 100644 index 000000000000..7b08f723b187 --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSortScatter.compute.fx @@ -0,0 +1,35 @@ +// Gaussian Splatting GPU depth sort - scatter pass. +// +// After the histogram has been exclusive-scanned into per-bucket start offsets, each active slot claims +// the next free position within its bucket via an atomic increment and writes its source index (as f32) +// into the sorted-index buffer. Trailing padded slots are set to the reserved (invisible) index 0. + +struct Params { + count : u32, + paddedCount : u32, + numBuckets : u32, + pad : u32, + coeff : vec4f, +}; + +@group(0) @binding(0) var bucket : array; +@group(0) @binding(1) var seed : array; +@group(0) @binding(2) var offsets : array>; +@group(0) @binding(3) var sortedIndex : array; +@group(0) @binding(4) var params : Params; + +@compute @workgroup_size(256, 1, 1) +fn main(@builtin(global_invocation_id) gid : vec3u) { + let i = gid.x; + if (i >= params.paddedCount) { + return; + } + if (i >= params.count) { + // Padding beyond the active count renders the reserved invisible splat 0. + sortedIndex[i] = 0.0; + return; + } + let b = bucket[i]; + let dst = atomicAdd(&offsets[b], 1u); + sortedIndex[dst] = seed[i]; +} diff --git a/packages/dev/core/src/ShadersWGSL/prefixSumAddOffsets.compute.fx b/packages/dev/core/src/ShadersWGSL/prefixSumAddOffsets.compute.fx new file mode 100644 index 000000000000..71f10ba6d430 --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/prefixSumAddOffsets.compute.fx @@ -0,0 +1,20 @@ +// Adds each block's scanned offset back onto its elements, completing a hierarchical exclusive scan. +// Pairs with prefixSumScanBlock: after the per-block exclusive sums are written and the block totals +// have themselves been exclusive-scanned into `blockOffsets`, this adds blockOffsets[block] to every +// element of that block. + +struct Params { + count : u32, +}; + +@group(0) @binding(0) var data : array; +@group(0) @binding(1) var blockOffsets : array; +@group(0) @binding(2) var params : Params; + +@compute @workgroup_size(256, 1, 1) +fn main(@builtin(global_invocation_id) gid : vec3u, @builtin(workgroup_id) wid : vec3u) { + let g = gid.x; + if (g < params.count) { + data[g] = data[g] + blockOffsets[wid.x]; + } +} diff --git a/packages/dev/core/src/ShadersWGSL/prefixSumScanBlock.compute.fx b/packages/dev/core/src/ShadersWGSL/prefixSumScanBlock.compute.fx new file mode 100644 index 000000000000..7532dc00870c --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/prefixSumScanBlock.compute.fx @@ -0,0 +1,52 @@ +// Blelloch/Hillis-Steele block scan (WebGPU compute). +// +// Scans `count` u32 values in `data` in place into per-block EXCLUSIVE prefix sums, and writes each +// block's total into `blockSums`. A hierarchical driver scans `blockSums` and adds the offsets back +// (see prefixSumAddOffsets), producing a full exclusive scan over arrays larger than one workgroup. + +struct Params { + count : u32, +}; + +@group(0) @binding(0) var data : array; +@group(0) @binding(1) var blockSums : array; +@group(0) @binding(2) var params : Params; + +var temp : array; + +@compute @workgroup_size(256, 1, 1) +fn main(@builtin(global_invocation_id) gid : vec3u, @builtin(local_invocation_id) lid : vec3u, @builtin(workgroup_id) wid : vec3u) { + let g = gid.x; + let t = lid.x; + + // Explicit bounds check: `select` would still evaluate `data[g]` for out-of-range lanes, and the recursive + // block-sums buffer is sized to the exact block count (not padded to the workgroup size), so that read can be + // out of bounds. Guard the load so padded lanes contribute 0. + var v = 0u; + if (g < params.count) { + v = data[g]; + } + temp[t] = v; + workgroupBarrier(); + + // Hillis-Steele inclusive scan. Each step reads old values into a register before the shared + // array is updated, so no thread observes a partially-updated neighbour. + for (var d = 1u; d < 256u; d = d * 2u) { + var val = temp[t]; + if (t >= d) { + val = val + temp[t - d]; + } + workgroupBarrier(); + temp[t] = val; + workgroupBarrier(); + } + + // Exclusive result = inclusive - own value. + if (g < params.count) { + data[g] = temp[t] - v; + } + // The last lane holds the inclusive total of the whole block (padded lanes contribute 0). + if (t == 255u) { + blockSums[wid.x] = temp[t]; + } +} diff --git a/packages/dev/core/test/unit/Meshes/babylon.gaussianSplatting.gpuSorter.test.ts b/packages/dev/core/test/unit/Meshes/babylon.gaussianSplatting.gpuSorter.test.ts new file mode 100644 index 000000000000..8ca61c69061d --- /dev/null +++ b/packages/dev/core/test/unit/Meshes/babylon.gaussianSplatting.gpuSorter.test.ts @@ -0,0 +1,18 @@ +import { NullEngine } from "core/Engines/nullEngine"; +import { type AbstractEngine } from "core/Engines/abstractEngine"; +import { GaussianSplattingGpuSorter } from "core/Meshes/GaussianSplatting/gaussianSplattingGpuSorter"; +import { describe, expect, it } from "vitest"; + +describe("GaussianSplattingGpuSorter.IsSupported", () => { + it("returns false when the engine does not support compute shaders", () => { + const engine = new NullEngine(); + expect(GaussianSplattingGpuSorter.IsSupported(engine)).toBe(false); + }); + + it("reflects the engine compute capability flag", () => { + const makeEngine = (supportComputeShaders: boolean) => ({ getCaps: () => ({ supportComputeShaders }) }) as unknown as AbstractEngine; + + expect(GaussianSplattingGpuSorter.IsSupported(makeEngine(true))).toBe(true); + expect(GaussianSplattingGpuSorter.IsSupported(makeEngine(false))).toBe(false); + }); +}); diff --git a/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json b/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json index 1846c83ec41b..2f6fad96cb7a 100644 --- a/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json +++ b/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json @@ -194,8 +194,15 @@ "ShadersWGSL/fxaa.vertex.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplatting.fragment.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplatting.vertex.ts": ["shader-store-write"], + "ShadersWGSL/gaussianSplattingCull.compute.ts": ["shader-store-write"], + "ShadersWGSL/gaussianSplattingCullCompact.compute.ts": ["shader-store-write"], + "ShadersWGSL/gaussianSplattingCullFinalize.compute.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplattingDepth.fragment.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplattingDepth.vertex.ts": ["shader-store-write"], + "ShadersWGSL/gaussianSplattingSortClear.compute.ts": ["shader-store-write"], + "ShadersWGSL/gaussianSplattingSortDepth.compute.ts": ["shader-store-write"], + "ShadersWGSL/gaussianSplattingSortHistogram.compute.ts": ["shader-store-write"], + "ShadersWGSL/gaussianSplattingSortScatter.compute.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplattingVoxel.fragment.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplattingVoxel.vertex.ts": ["shader-store-write"], "ShadersWGSL/geometry.fragment.ts": ["shader-store-write"], @@ -273,6 +280,8 @@ "ShadersWGSL/picking.fragment.ts": ["shader-store-write"], "ShadersWGSL/picking.vertex.ts": ["shader-store-write"], "ShadersWGSL/postprocess.vertex.ts": ["shader-store-write"], + "ShadersWGSL/prefixSumAddOffsets.compute.ts": ["shader-store-write"], + "ShadersWGSL/prefixSumScanBlock.compute.ts": ["shader-store-write"], "ShadersWGSL/procedural.vertex.ts": ["shader-store-write"], "ShadersWGSL/rgbdDecode.fragment.ts": ["shader-store-write"], "ShadersWGSL/rgbdEncode.fragment.ts": ["shader-store-write"],