From 5e095219b4dc22fd6e90cdc9e61df2b2fd839245 Mon Sep 17 00:00:00 2001 From: Cedric Guillemet <1312968+CedricGuillemet@users.noreply.github.com> Date: Thu, 2 Jul 2026 09:25:44 +0200 Subject: [PATCH 1/9] GS WebGPU fastpath --- .../gaussianSplattingGpuSorter.pure.ts | 122 +++++++++++++++ .../gaussianSplattingGpuSorter.ts | 3 + .../gaussianSplattingMeshBase.pure.ts | 141 ++++++++++++++++++ .../gaussianSplattingMeshBase.ts | 1 + .../core/src/Meshes/GaussianSplatting/pure.ts | 1 + .../core/src/Meshes/thinInstanceMesh.pure.ts | 18 +++ .../core/src/Meshes/thinInstanceMesh.types.ts | 4 + ...gaussianSplattingSplatIndexInit.compute.fx | 25 ++++ ...abylon.gaussianSplatting.gpuSorter.test.ts | 18 +++ packages/public/@babylonjs/core/package.json | 1 + .../side-effects-manifest/core/Meshes.json | 1 + .../core/ShadersWGSL.json | 1 + 12 files changed, 336 insertions(+) create mode 100644 packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts create mode 100644 packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts create mode 100644 packages/dev/core/src/ShadersWGSL/gaussianSplattingSplatIndexInit.compute.fx create mode 100644 packages/dev/core/test/unit/Meshes/babylon.gaussianSplatting.gpuSorter.test.ts diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts new file mode 100644 index 000000000000..97d64335e087 --- /dev/null +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts @@ -0,0 +1,122 @@ +/** This file must only contain pure code and pure imports */ + +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"; + +/** + * @internal + * WebGPU compute driver that produces the per-splat `splatIndex` order buffer for a + * {@link GaussianSplattingMeshBase} on the GPU. + * + * Phase 0 scaffold: it uploads a packed (interval-aware) source-index list and normalizes it + * through a compute pass, so the resulting buffer is produced on the GPU and can be bound directly + * as the `splatIndex` instanced vertex buffer with no CPU readback. This validates the + * compute to storage buffer to instanced vertex buffer path that the full fast path depends on. + * Phase 1 replaces the seed copy with GPU depth-key generation followed by a radix sort. + */ +export class GaussianSplattingGpuSorter { + private readonly _engine: AbstractEngine; + private _compute: Nullable = null; + private _params: Nullable = null; + private _seedBuffer: Nullable = null; + private _sortedIndexBuffer: Nullable = null; + private _capacity = 0; + + /** + * 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; + } + + /** + * Creates a new GPU sorter. + * @param engine the (WebGPU) engine to run the compute passes on + */ + public constructor(engine: AbstractEngine) { + this._engine = engine; + } + + private _ensureResources(paddedCount: number): void { + if (!this._compute) { + const bindingsMapping: ComputeBindingMapping = { + seedBuffer: { group: 0, binding: 0 }, + sortedIndexBuffer: { group: 0, binding: 1 }, + params: { group: 0, binding: 2 }, + }; + this._compute = new ComputeShader("gaussianSplattingSplatIndexInit", this._engine, "gaussianSplattingSplatIndexInit", { bindingsMapping }); + this._params = new UniformBuffer(this._engine, undefined, undefined, "GaussianSplattingGpuSorterParams"); + this._params.addUniform("count", 1); + } + + if (paddedCount > this._capacity) { + this._seedBuffer?.dispose(); + this._sortedIndexBuffer?.dispose(); + const byteLength = paddedCount * Float32Array.BYTES_PER_ELEMENT; + this._seedBuffer = new StorageBuffer(this._engine as WebGPUEngine, byteLength, Constants.BUFFER_CREATIONFLAG_READWRITE, "GaussianSplattingSplatIndexSeed"); + // The sorted-index buffer is consumed as an instanced vertex buffer, so it needs the VERTEX usage flag + // in addition to READWRITE (written by the compute pass). + this._sortedIndexBuffer = new StorageBuffer( + this._engine as WebGPUEngine, + byteLength, + Constants.BUFFER_CREATIONFLAG_READWRITE | Constants.BUFFER_CREATIONFLAG_VERTEX, + "GaussianSplattingSplatIndexSorted" + ); + this._capacity = paddedCount; + } + } + + /** + * Generates the sorted-index buffer from a packed source-index list. + * @param seedIndices packed (interval-aware) source indices, one f32 per rendered slot (already padded) + * @returns true when the compute pass was dispatched (the shader was ready); false when it must be retried + */ + public build(seedIndices: Float32Array): boolean { + const paddedCount = seedIndices.length; + if (paddedCount === 0) { + return false; + } + this._ensureResources(paddedCount); + + this._seedBuffer!.update(seedIndices); + this._params!.updateUInt("count", paddedCount); + this._params!.update(); + + this._compute!.setStorageBuffer("seedBuffer", this._seedBuffer!); + this._compute!.setStorageBuffer("sortedIndexBuffer", this._sortedIndexBuffer!); + this._compute!.setUniformBuffer("params", this._params!); + // dispatch returns false when the compute effect (or a bound resource) is not ready yet; the caller + // keeps the "dirty" state and retries on a later frame. + return this._compute!.dispatch(Math.ceil(paddedCount / 256)); + } + + /** + * The GPU data buffer to bind as the `splatIndex` instanced vertex buffer, or null before the first build. + */ + public get sortedDataBuffer(): Nullable { + return this._sortedIndexBuffer ? this._sortedIndexBuffer.getBuffer() : null; + } + + /** + * Releases all GPU resources held by the sorter. + */ + public dispose(): void { + this._seedBuffer?.dispose(); + this._sortedIndexBuffer?.dispose(); + this._params?.dispose(); + this._seedBuffer = null; + this._sortedIndexBuffer = null; + this._params = null; + this._compute = null; + this._capacity = 0; + } +} 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..4c6be800a728 --- /dev/null +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts @@ -0,0 +1,3 @@ +export * from "./gaussianSplattingGpuSorter.pure"; + +import "../../ShadersWGSL/gaussianSplattingSplatIndexInit.compute"; diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts index 956cb1361422..e3ee700da449 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.pure"; +import { type DataBuffer } from "core/Buffers/dataBuffer"; // eslint-disable-next-line @typescript-eslint/naming-convention declare const _native: INative; @@ -58,6 +60,20 @@ 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 keeps the CPU worker / native path until GPU sorting is complete. */ + 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 +501,26 @@ 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; + /** @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; private _modelViewProjectionMatrix = Matrix.Identity(); private _depthMix: BigInt64Array; protected _canPostToWorker = true; @@ -1320,6 +1353,86 @@ 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 keeps the worker path until GPU depth sorting is implemented; only an explicit Gpu request opts in. + return mode === GaussianSplattingGpuSortMode.Gpu; + } + + /** + * Generates (when needed) the GPU sorted-index buffer and binds it to each active camera mesh. + * @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) { + return; + } + + if (this._gpuSortDirty || !this._gpuSortedDataBuffer) { + // The packed source-index seed (identity / interval-aware order) is produced on the CPU by + // _updateSplatIndexBuffer; the compute pass copies it into the GPU sorted-index buffer. + const ready = sorter.build(this._splatIndex); + const dataBuffer = sorter.sortedDataBuffer; + if (dataBuffer !== this._gpuSortedDataBuffer) { + // The underlying buffer was (re)allocated, so every camera mesh must rebind against it. + this._gpuSortedDataBuffer = dataBuffer; + this._cameraViewInfos.forEach((info) => (info.splatIndexBufferSet = false)); + } + // Keep retrying while the compute shader is still compiling; clear the flag only once it ran. + if (ready) { + this._gpuSortDirty = false; + } + } + + const dataBuffer = this._gpuSortedDataBuffer; + if (!dataBuffer) { + return; + } + + const instanceCount = Math.max(this._splatIndex.length >> 4, 1); + const worldMatrix = this.computeWorldMatrix(true); + activeViewInfos.forEach((info) => { + if (!info.splatIndexBufferSet) { + info.mesh._thinInstanceSetSplatIndexBuffer(dataBuffer, instanceCount); + info.splatIndexBufferSet = true; + } + // Keep the per-view sort bookkeeping current so isReady() / _isSortStateDirty consider this view's + // ordering applied (the GPU path bypasses the worker-posting block that normally updates these). + const camera = info.camera; + const cameraViewMatrix = camera.getViewMatrix(); + info.cameraDirection.copyFrom(this._getCameraDirection(camera)); + info.sortWorldMatrix.copyFrom(worldMatrix); + info.sortCameraForward.set(cameraViewMatrix.m[2], cameraViewMatrix.m[6], cameraViewMatrix.m[10]); + info.sortCameraPosition.copyFrom(camera.globalPosition); + info.sortAppliedId = info.sortRequestId; + info.frameIdLastUpdate = frameId; + }); + this._readyToDisplay = true; + this._canPostToWorker = true; + } + /** @internal */ public _postToWorker(forced = false): void { const scene = this._scene; @@ -1391,6 +1504,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 +2418,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(); @@ -2956,6 +3080,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 +3186,20 @@ 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; + } + // no worker in native if (IsNative) { return; diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts index 5062e4eb9f9d..14db97466ff3 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts @@ -1,3 +1,4 @@ export * from "./gaussianSplattingMeshBase.pure"; import "core/Meshes/thinInstanceMesh"; +import "./gaussianSplattingGpuSorter"; diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/pure.ts index ab7f89df1939..aefea68e2324 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/pure.ts @@ -1,6 +1,7 @@ /** Pure barrel — re-exports only side-effect-free modules */ export * from "./gaussianSplattingCompoundMesh.pure"; export * from "./gaussianSplattingDebugger.pure"; +export * from "./gaussianSplattingGpuSorter.pure"; export * from "./gaussianSplattingMesh.pure"; export * from "./gaussianSplattingMeshBase.pure"; export * from "./gaussianSplattingPartProxyMesh.pure"; diff --git a/packages/dev/core/src/Meshes/thinInstanceMesh.pure.ts b/packages/dev/core/src/Meshes/thinInstanceMesh.pure.ts index 00055b1c1763..fbb9884c1104 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,23 @@ 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); + 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/gaussianSplattingSplatIndexInit.compute.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSplatIndexInit.compute.fx new file mode 100644 index 000000000000..b273faa19123 --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSplatIndexInit.compute.fx @@ -0,0 +1,25 @@ +// Gaussian Splatting - splat index generation (WebGPU compute). +// +// Phase 0 scaffold for the WebGPU Gaussian Splatting fast path: copies a packed +// (interval-aware) source-index list into the sorted-index buffer that is bound +// directly as the `splatIndex` instanced vertex buffer. This proves the +// compute -> storage buffer -> instanced vertex buffer data path end to end, +// without any CPU readback. Phase 1 replaces the seed copy with GPU depth-key +// generation followed by a radix sort. + +struct Params { + count : u32, +}; + +@group(0) @binding(0) var seedBuffer : array; +@group(0) @binding(1) var sortedIndexBuffer : array; +@group(0) @binding(2) var params : Params; + +@compute @workgroup_size(256, 1, 1) +fn main(@builtin(global_invocation_id) global_id : vec3) { + let index = global_id.x; + if (index >= params.count) { + return; + } + sortedIndexBuffer[index] = seedBuffer[index]; +} 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/packages/public/@babylonjs/core/package.json b/packages/public/@babylonjs/core/package.json index b3d83f51ee4f..3e442d08e2a1 100644 --- a/packages/public/@babylonjs/core/package.json +++ b/packages/public/@babylonjs/core/package.json @@ -457,6 +457,7 @@ "Meshes/Builders/tubeBuilder.js", "Meshes/GaussianSplatting/gaussianSplattingCompoundMesh.js", "Meshes/GaussianSplatting/gaussianSplattingDebugger.js", + "Meshes/GaussianSplatting/gaussianSplattingGpuSorter.js", "Meshes/GaussianSplatting/gaussianSplattingMesh.js", "Meshes/GaussianSplatting/gaussianSplattingMeshBase.js", "Meshes/GaussianSplatting/gaussianSplattingPartProxyMesh.js", diff --git a/scripts/treeshaking/side-effects-manifest/core/Meshes.json b/scripts/treeshaking/side-effects-manifest/core/Meshes.json index 7d442d2ec47c..8624970c35fc 100644 --- a/scripts/treeshaking/side-effects-manifest/core/Meshes.json +++ b/scripts/treeshaking/side-effects-manifest/core/Meshes.json @@ -22,6 +22,7 @@ "Meshes/Builders/tubeBuilder.ts": ["top-level-call"], "Meshes/GaussianSplatting/gaussianSplattingCompoundMesh.ts": ["top-level-call"], "Meshes/GaussianSplatting/gaussianSplattingDebugger.ts": ["bare-import"], + "Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts": ["bare-import"], "Meshes/GaussianSplatting/gaussianSplattingMesh.ts": ["bare-import", "top-level-call"], "Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts": ["bare-import"], "Meshes/GaussianSplatting/gaussianSplattingPartProxyMesh.ts": ["top-level-call"], diff --git a/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json b/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json index 1846c83ec41b..b889c11e4f17 100644 --- a/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json +++ b/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json @@ -196,6 +196,7 @@ "ShadersWGSL/gaussianSplatting.vertex.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplattingDepth.fragment.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplattingDepth.vertex.ts": ["shader-store-write"], + "ShadersWGSL/gaussianSplattingSplatIndexInit.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"], From ddbccd1ce9433d0dfed06b7952e51b7355022aa1 Mon Sep 17 00:00:00 2001 From: Cedric Guillemet <1312968+CedricGuillemet@users.noreply.github.com> Date: Thu, 2 Jul 2026 10:22:17 +0200 Subject: [PATCH 2/9] Phase 1 with GPU sorting --- .../core/src/Compute/prefixSumCompute.pure.ts | 159 ++++++++++ .../dev/core/src/Compute/prefixSumCompute.ts | 4 + packages/dev/core/src/Compute/pure.ts | 1 + .../gaussianSplattingGpuSorter.pure.ts | 280 ++++++++++++++---- .../gaussianSplattingGpuSorter.ts | 6 +- .../gaussianSplattingMeshBase.pure.ts | 83 ++++-- .../gaussianSplattingSortClear.compute.fx | 20 ++ .../gaussianSplattingSortDepth.compute.fx | 62 ++++ .../gaussianSplattingSortHistogram.compute.fx | 46 +++ .../gaussianSplattingSortScatter.compute.fx | 35 +++ ...gaussianSplattingSplatIndexInit.compute.fx | 25 -- .../prefixSumAddOffsets.compute.fx | 20 ++ .../ShadersWGSL/prefixSumScanBlock.compute.fx | 46 +++ packages/public/@babylonjs/core/package.json | 1 + .../side-effects-manifest/core/Compute.json | 3 +- .../core/ShadersWGSL.json | 7 +- 16 files changed, 689 insertions(+), 109 deletions(-) create mode 100644 packages/dev/core/src/Compute/prefixSumCompute.pure.ts create mode 100644 packages/dev/core/src/Compute/prefixSumCompute.ts create mode 100644 packages/dev/core/src/ShadersWGSL/gaussianSplattingSortClear.compute.fx create mode 100644 packages/dev/core/src/ShadersWGSL/gaussianSplattingSortDepth.compute.fx create mode 100644 packages/dev/core/src/ShadersWGSL/gaussianSplattingSortHistogram.compute.fx create mode 100644 packages/dev/core/src/ShadersWGSL/gaussianSplattingSortScatter.compute.fx delete mode 100644 packages/dev/core/src/ShadersWGSL/gaussianSplattingSplatIndexInit.compute.fx create mode 100644 packages/dev/core/src/ShadersWGSL/prefixSumAddOffsets.compute.fx create mode 100644 packages/dev/core/src/ShadersWGSL/prefixSumScanBlock.compute.fx diff --git a/packages/dev/core/src/Compute/prefixSumCompute.pure.ts b/packages/dev/core/src/Compute/prefixSumCompute.pure.ts new file mode 100644 index 000000000000..3307dcb59b05 --- /dev/null +++ b/packages/dev/core/src/Compute/prefixSumCompute.pure.ts @@ -0,0 +1,159 @@ +/** This file must only contain pure code and pure imports */ + +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 is intended to be + * reused by the interval culling pass in a later phase. + */ +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[] = []; + + /** + * Creates a new prefix-sum compute helper. + * @param engine the (WebGPU) engine to run the compute passes on + */ + public constructor(engine: AbstractEngine) { + this._engine = engine; + } + + private _ensureShaders(): void { + if (this._scanBlock) { + 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"); + 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 compiled and ready to dispatch. + * @returns true when the scan can run + */ + public isReady(): boolean { + 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 { + if (count <= 1) { + 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/Compute/prefixSumCompute.ts b/packages/dev/core/src/Compute/prefixSumCompute.ts new file mode 100644 index 000000000000..e76765a2d7c1 --- /dev/null +++ b/packages/dev/core/src/Compute/prefixSumCompute.ts @@ -0,0 +1,4 @@ +export * from "./prefixSumCompute.pure"; + +import "../ShadersWGSL/prefixSumScanBlock.compute"; +import "../ShadersWGSL/prefixSumAddOffsets.compute"; diff --git a/packages/dev/core/src/Compute/pure.ts b/packages/dev/core/src/Compute/pure.ts index 6ae86a736bb1..aacba9a1d7ec 100644 --- a/packages/dev/core/src/Compute/pure.ts +++ b/packages/dev/core/src/Compute/pure.ts @@ -1,3 +1,4 @@ /** Pure barrel — re-exports only side-effect-free modules */ export * from "./computeEffect"; export * from "./computeShader.pure"; +export * from "./prefixSumCompute.pure"; diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts index 97d64335e087..f95704665ea5 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts @@ -4,30 +4,56 @@ 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"; +import { PrefixSumCompute } from "core/Compute/prefixSumCompute.pure"; + +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 per-splat `splatIndex` order buffer for a - * {@link GaussianSplattingMeshBase} on the GPU. + * 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 * - * Phase 0 scaffold: it uploads a packed (interval-aware) source-index list and normalizes it - * through a compute pass, so the resulting buffer is produced on the GPU and can be bound directly - * as the `splatIndex` instanced vertex buffer with no CPU readback. This validates the - * compute to storage buffer to instanced vertex buffer path that the full fast path depends on. - * Phase 1 replaces the seed copy with GPU depth-key generation followed by a radix sort. + * 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. */ export class GaussianSplattingGpuSorter { private readonly _engine: AbstractEngine; - private _compute: Nullable = null; + + private _clear: Nullable = null; + private _depth: Nullable = null; + private _histogram: Nullable = null; + private _scatter: Nullable = null; + private _prefixSum: Nullable = null; private _params: Nullable = null; - private _seedBuffer: 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; - private _capacity = 0; + + private _positionCapacity = 0; + private _paddedCapacity = 0; + private _histCapacity = 0; + private _renderedCount = 0; + private _paddedCount = 0; + private _numBuckets = 0; + private _hasSource = false; /** * Whether the running engine supports the WebGPU compute sort path. @@ -38,6 +64,17 @@ export class GaussianSplattingGpuSorter { 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. * @param engine the (WebGPU) engine to run the compute passes on @@ -46,61 +83,178 @@ export class GaussianSplattingGpuSorter { this._engine = engine; } - private _ensureResources(paddedCount: number): void { - if (!this._compute) { - const bindingsMapping: ComputeBindingMapping = { - seedBuffer: { group: 0, binding: 0 }, - sortedIndexBuffer: { group: 0, binding: 1 }, - params: { group: 0, binding: 2 }, - }; - this._compute = new ComputeShader("gaussianSplattingSplatIndexInit", this._engine, "gaussianSplattingSplatIndexInit", { bindingsMapping }); - this._params = new UniformBuffer(this._engine, undefined, undefined, "GaussianSplattingGpuSorterParams"); - this._params.addUniform("count", 1); + 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); + this._params = new UniformBuffer(this._engine, undefined, undefined, "GaussianSplattingGpuSorterParams"); + 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); + } - if (paddedCount > this._capacity) { - this._seedBuffer?.dispose(); + 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 byteLength = paddedCount * Float32Array.BYTES_PER_ELEMENT; - this._seedBuffer = new StorageBuffer(this._engine as WebGPUEngine, byteLength, Constants.BUFFER_CREATIONFLAG_READWRITE, "GaussianSplattingSplatIndexSeed"); - // The sorted-index buffer is consumed as an instanced vertex buffer, so it needs the VERTEX usage flag - // in addition to READWRITE (written by the compute pass). - this._sortedIndexBuffer = new StorageBuffer( - this._engine as WebGPUEngine, - byteLength, - Constants.BUFFER_CREATIONFLAG_READWRITE | Constants.BUFFER_CREATIONFLAG_VERTEX, - "GaussianSplattingSplatIndexSorted" - ); - this._capacity = paddedCount; + 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"); } + + this._positions!.update(positions); + this._seed!.update(seedIndices); + this._hasSource = true; } /** - * Generates the sorted-index buffer from a packed source-index list. - * @param seedIndices packed (interval-aware) source indices, one f32 per rendered slot (already padded) - * @returns true when the compute pass was dispatched (the shader was ready); false when it must be retried + * Whether all compute shaders are compiled and ready to dispatch. + * @returns true when a sort can run */ - public build(seedIndices: Float32Array): boolean { - const paddedCount = seedIndices.length; - if (paddedCount === 0) { + public isReady(): boolean { + if (!this._clear) { return false; } - this._ensureResources(paddedCount); + return !!this._clear.isReady() && !!this._depth!.isReady() && !!this._histogram!.isReady() && !!this._scatter!.isReady() && !!this._prefixSum!.isReady(); + } - this._seedBuffer!.update(seedIndices); - this._params!.updateUInt("count", paddedCount); + /** + * 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(); - this._compute!.setStorageBuffer("seedBuffer", this._seedBuffer!); - this._compute!.setStorageBuffer("sortedIndexBuffer", this._sortedIndexBuffer!); - this._compute!.setUniformBuffer("params", this._params!); - // dispatch returns false when the compute effect (or a bound resource) is not ready yet; the caller - // keeps the "dirty" state and retries on a later frame. - return this._compute!.dispatch(Math.ceil(paddedCount / 256)); + // 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 the first build. + * 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; @@ -110,13 +264,31 @@ export class GaussianSplattingGpuSorter { * Releases all GPU resources held by the sorter. */ public dispose(): void { - this._seedBuffer?.dispose(); + this._positions?.dispose(); + this._seed?.dispose(); + this._depthBuffer?.dispose(); + this._bucket?.dispose(); + this._range?.dispose(); + this._histogramBuffer?.dispose(); this._sortedIndexBuffer?.dispose(); this._params?.dispose(); - this._seedBuffer = null; + this._prefixSum?.dispose(); + this._positions = null; + this._seed = null; + this._depthBuffer = null; + this._bucket = null; + this._range = null; + this._histogramBuffer = null; this._sortedIndexBuffer = null; this._params = null; - this._compute = null; - this._capacity = 0; + this._prefixSum = null; + this._clear = null; + this._depth = null; + this._histogram = null; + this._scatter = null; + this._positionCapacity = 0; + this._paddedCapacity = 0; + this._histCapacity = 0; + this._hasSource = false; } } diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts index 4c6be800a728..7096a18cae13 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts @@ -1,3 +1,7 @@ export * from "./gaussianSplattingGpuSorter.pure"; -import "../../ShadersWGSL/gaussianSplattingSplatIndexInit.compute"; +import "core/Compute/prefixSumCompute"; +import "../../ShadersWGSL/gaussianSplattingSortClear.compute"; +import "../../ShadersWGSL/gaussianSplattingSortDepth.compute"; +import "../../ShadersWGSL/gaussianSplattingSortHistogram.compute"; +import "../../ShadersWGSL/gaussianSplattingSortScatter.compute"; diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts index e3ee700da449..aecf5efb5395 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts @@ -1375,58 +1375,82 @@ export class GaussianSplattingMeshBase extends Mesh { if (!GaussianSplattingGpuSorter.IsSupported(this._scene.getEngine())) { return false; } - // Auto keeps the worker path until GPU depth sorting is implemented; only an explicit Gpu request opts in. + // 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 sorted-index buffer and binds it to each active camera mesh. + * 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) { + if (!sorter || !this._splatIndex || !this._splatPositions || !activeViewInfos.length) { return; } - if (this._gpuSortDirty || !this._gpuSortedDataBuffer) { - // The packed source-index seed (identity / interval-aware order) is produced on the CPU by - // _updateSplatIndexBuffer; the compute pass copies it into the GPU sorted-index buffer. - const ready = sorter.build(this._splatIndex); - const dataBuffer = sorter.sortedDataBuffer; - if (dataBuffer !== this._gpuSortedDataBuffer) { - // The underlying buffer was (re)allocated, so every camera mesh must rebind against it. - this._gpuSortedDataBuffer = dataBuffer; - this._cameraViewInfos.forEach((info) => (info.splatIndexBufferSet = false)); - } - // Keep retrying while the compute shader is still compiling; clear the flag only once it ran. - if (ready) { - this._gpuSortDirty = false; + // 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. + const activeCamera = this._scene.activeCamera; + const primary = (activeCamera && activeViewInfos.find((info) => info.camera === activeCamera)) || activeViewInfos[0]; + 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 needSort = this._gpuSortDirty || !this._readyToDisplay || this._isSortStateDirty(primary, worldMatrix, camera); + 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; } - const dataBuffer = this._gpuSortedDataBuffer; + const dataBuffer = sorter.sortedDataBuffer; if (!dataBuffer) { return; } + if (dataBuffer !== this._gpuSortedDataBuffer) { + // The underlying buffer was (re)allocated, so every camera mesh must rebind against it. + this._gpuSortedDataBuffer = dataBuffer; + this._cameraViewInfos.forEach((info) => (info.splatIndexBufferSet = false)); + } const instanceCount = Math.max(this._splatIndex.length >> 4, 1); - const worldMatrix = this.computeWorldMatrix(true); activeViewInfos.forEach((info) => { if (!info.splatIndexBufferSet) { info.mesh._thinInstanceSetSplatIndexBuffer(dataBuffer, instanceCount); info.splatIndexBufferSet = true; } - // Keep the per-view sort bookkeeping current so isReady() / _isSortStateDirty consider this view's - // ordering applied (the GPU path bypasses the worker-posting block that normally updates these). - const camera = info.camera; - const cameraViewMatrix = camera.getViewMatrix(); - info.cameraDirection.copyFrom(this._getCameraDirection(camera)); - info.sortWorldMatrix.copyFrom(worldMatrix); - info.sortCameraForward.set(cameraViewMatrix.m[2], cameraViewMatrix.m[6], cameraViewMatrix.m[10]); - info.sortCameraPosition.copyFrom(camera.globalPosition); - info.sortAppliedId = info.sortRequestId; info.frameIdLastUpdate = frameId; }); this._readyToDisplay = true; @@ -2775,6 +2799,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); } } @@ -2818,6 +2844,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; } /** @@ -2838,6 +2866,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( 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/gaussianSplattingSplatIndexInit.compute.fx b/packages/dev/core/src/ShadersWGSL/gaussianSplattingSplatIndexInit.compute.fx deleted file mode 100644 index b273faa19123..000000000000 --- a/packages/dev/core/src/ShadersWGSL/gaussianSplattingSplatIndexInit.compute.fx +++ /dev/null @@ -1,25 +0,0 @@ -// Gaussian Splatting - splat index generation (WebGPU compute). -// -// Phase 0 scaffold for the WebGPU Gaussian Splatting fast path: copies a packed -// (interval-aware) source-index list into the sorted-index buffer that is bound -// directly as the `splatIndex` instanced vertex buffer. This proves the -// compute -> storage buffer -> instanced vertex buffer data path end to end, -// without any CPU readback. Phase 1 replaces the seed copy with GPU depth-key -// generation followed by a radix sort. - -struct Params { - count : u32, -}; - -@group(0) @binding(0) var seedBuffer : array; -@group(0) @binding(1) var sortedIndexBuffer : array; -@group(0) @binding(2) var params : Params; - -@compute @workgroup_size(256, 1, 1) -fn main(@builtin(global_invocation_id) global_id : vec3) { - let index = global_id.x; - if (index >= params.count) { - return; - } - sortedIndexBuffer[index] = seedBuffer[index]; -} 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..6ece483d6119 --- /dev/null +++ b/packages/dev/core/src/ShadersWGSL/prefixSumScanBlock.compute.fx @@ -0,0 +1,46 @@ +// 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; + + let v = select(0u, data[g], g < params.count); + 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/public/@babylonjs/core/package.json b/packages/public/@babylonjs/core/package.json index 3e442d08e2a1..3155bd19a4e7 100644 --- a/packages/public/@babylonjs/core/package.json +++ b/packages/public/@babylonjs/core/package.json @@ -122,6 +122,7 @@ "Cameras/virtualJoysticksCamera.js", "Collisions/collisionCoordinator.js", "Compute/computeShader.types.js", + "Compute/prefixSumCompute.js", "Culling/Helper/computeShaderBoundingHelper.js", "Culling/Helper/transformFeedbackBoundingHelper.js", "Culling/Octrees/octreeSceneComponent.js", diff --git a/scripts/treeshaking/side-effects-manifest/core/Compute.json b/scripts/treeshaking/side-effects-manifest/core/Compute.json index 383e36a2ea46..a15b090c9a99 100644 --- a/scripts/treeshaking/side-effects-manifest/core/Compute.json +++ b/scripts/treeshaking/side-effects-manifest/core/Compute.json @@ -2,6 +2,7 @@ "version": 1, "files": { "Compute/computeShader.ts": ["top-level-call"], - "Compute/computeShader.types.ts": ["declare-module"] + "Compute/computeShader.types.ts": ["declare-module"], + "Compute/prefixSumCompute.ts": ["bare-import"] } } diff --git a/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json b/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json index b889c11e4f17..f95ef1f67a7e 100644 --- a/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json +++ b/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json @@ -196,7 +196,10 @@ "ShadersWGSL/gaussianSplatting.vertex.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplattingDepth.fragment.ts": ["shader-store-write"], "ShadersWGSL/gaussianSplattingDepth.vertex.ts": ["shader-store-write"], - "ShadersWGSL/gaussianSplattingSplatIndexInit.compute.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"], @@ -274,6 +277,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"], From 2dc578811635ad6457bca1ba20ef15844d471b65 Mon Sep 17 00:00:00 2001 From: Cedric Guillemet <1312968+CedricGuillemet@users.noreply.github.com> Date: Thu, 2 Jul 2026 15:54:59 +0200 Subject: [PATCH 3/9] webgpu culling --- .../core/src/Compute/prefixSumCompute.pure.ts | 2 +- .../core/src/Engines/abstractEngine.pure.ts | 13 ++ .../dev/core/src/Engines/webgpuEngine.pure.ts | 34 ++- .../gaussianSplattingGpuSorter.pure.ts | 212 +++++++++++++++++- .../gaussianSplattingGpuSorter.ts | 3 + .../gaussianSplattingMeshBase.pure.ts | 67 +++++- packages/dev/core/src/Meshes/mesh.pure.ts | 11 + .../core/src/Meshes/thinInstanceMesh.pure.ts | 4 + .../ShadersWGSL/gaussianSplatting.vertex.fx | 11 +- .../gaussianSplattingCull.compute.fx | 43 ++++ .../gaussianSplattingCullCompact.compute.fx | 25 +++ .../gaussianSplattingCullFinalize.compute.fx | 46 ++++ .../devHost/src/testScene/createScene.ts | 68 +++--- packages/tools/devHost/src/testScene/main.ts | 20 +- .../core/ShadersWGSL.json | 3 + 15 files changed, 508 insertions(+), 54 deletions(-) create mode 100644 packages/dev/core/src/ShadersWGSL/gaussianSplattingCull.compute.fx create mode 100644 packages/dev/core/src/ShadersWGSL/gaussianSplattingCullCompact.compute.fx create mode 100644 packages/dev/core/src/ShadersWGSL/gaussianSplattingCullFinalize.compute.fx diff --git a/packages/dev/core/src/Compute/prefixSumCompute.pure.ts b/packages/dev/core/src/Compute/prefixSumCompute.pure.ts index 3307dcb59b05..86e8b1789ef7 100644 --- a/packages/dev/core/src/Compute/prefixSumCompute.pure.ts +++ b/packages/dev/core/src/Compute/prefixSumCompute.pure.ts @@ -61,7 +61,7 @@ export class PrefixSumCompute { private _getUbo(count: number): UniformBuffer { if (this._uboIndex >= this._ubos.length) { - const ubo = new UniformBuffer(this._engine, undefined, undefined, "PrefixSumComputeParams"); + const ubo = new UniformBuffer(this._engine, undefined, undefined, "PrefixSumComputeParams", false, false); ubo.addUniform("count", 1); this._ubos.push(ubo); } 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..9448c1afff00 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; @@ -3764,7 +3764,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 +3841,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); @@ -3882,6 +3893,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.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts index f95704665ea5..e62a13321a1d 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts @@ -1,6 +1,7 @@ /** This file must only contain pure code and pure imports */ 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"; @@ -36,8 +37,15 @@ export class GaussianSplattingGpuSorter { 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; @@ -46,6 +54,16 @@ export class GaussianSplattingGpuSorter { 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; @@ -118,12 +136,89 @@ export class GaussianSplattingGpuSorter { }, }); this._prefixSum = new PrefixSumCompute(this._engine); - this._params = new UniformBuffer(this._engine, undefined, undefined, "GaussianSplattingGpuSorterParams"); + // 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 { @@ -175,11 +270,104 @@ export class GaussianSplattingGpuSorter { 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 { + 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 @@ -271,8 +459,17 @@ export class GaussianSplattingGpuSorter { 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; @@ -280,15 +477,28 @@ export class GaussianSplattingGpuSorter { 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/gaussianSplattingGpuSorter.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts index 7096a18cae13..6352e0104e68 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts @@ -5,3 +5,6 @@ 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"; diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts index aecf5efb5395..663b10096fa4 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts @@ -509,6 +509,30 @@ export class GaussianSplattingMeshBase extends Mesh { */ 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 per-frame frustum cull + * on the GPU and draws only the visible splats via an indirect draw. The cull test matches the vertex shader + * exactly, so it never removes a splat the shader would have shown. Set false to draw all sorted splats. + * + * Currently defaults to false: the cull path is complete but blocked by a one-dispatch UniformBuffer-upload + * lag for compute dispatched inside the mesh render pass (the cull reads the previous dispatch's clip matrix). + * The GPU sort + non-culled indirect draw path is unaffected and remains the shipping default. + */ + /** + * When true (and {@link EnableGpuIndirectDraw} is on), the WebGPU GPU-sort path runs a per-frame frustum cull + * on the GPU 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; + /** @internal */ public _vertexCount = 0; protected _worker: Nullable = null; @@ -521,6 +545,11 @@ export class GaussianSplattingMeshBase extends Mesh { 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; + // 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; @@ -1404,7 +1433,16 @@ export class GaussianSplattingMeshBase extends Mesh { 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 needSort = this._gpuSortDirty || !this._readyToDisplay || this._isSortStateDirty(primary, worldMatrix, camera); + // Also keep re-running while culling is enabled but not yet active (its compute shaders may still be + // compiling on the first frames — without this a static camera would never start culling). + const cullingWanted = GaussianSplattingMeshBase.EnableGpuIndirectDraw && GaussianSplattingMeshBase.EnableGpuCulling; + const needSort = this._gpuSortDirty || !this._readyToDisplay || (cullingWanted && !this._gpuCulling) || this._isSortStateDirty(primary, worldMatrix, camera); + + // Local->clip matrix (world * view * projection) for the frustum cull pass. Recomputed every frame so the + // per-frame cull tracks the camera even when the (more expensive) depth sort is skipped. + camera.getViewMatrix().multiplyToRef(camera.getProjectionMatrix(), this._gpuViewProjMatrix); + worldMatrix.multiplyToRef(this._gpuViewProjMatrix, this._gpuClipMatrix); + if (needSort) { const cameraViewMatrix = camera.getViewMatrix(); const world = worldMatrix.m; @@ -1435,22 +1473,45 @@ export class GaussianSplattingMeshBase extends Mesh { primary.sortAppliedId = primary.sortRequestId; } - const dataBuffer = sorter.sortedDataBuffer; + // Post-sort GPU frustum cull runs EVERY frame (not only when the sort re-runs) so it tracks the camera: + // once the sort has produced a sorted-index buffer, the cull compacts the visible splats and writes the + // indirect instance count using the same 1.2*w bounds the vertex shader uses. + if (cullingWanted && !this._gpuSortDirty) { + this._gpuCulling = sorter.cull(GaussianSplattingMeshBase._BatchSize * 6, this._gpuClipMatrix); + } else if (!cullingWanted) { + this._gpuCulling = false; + } + + // When culling is active the draw reads the compacted visible list and the GPU-written indirect args; + // otherwise it reads the full sorted list with a CPU-written full instance count. + const cullingActive = GaussianSplattingMeshBase.EnableGpuIndirectDraw && GaussianSplattingMeshBase.EnableGpuCulling && this._gpuCulling; + const dataBuffer = cullingActive ? sorter.culledDrawBuffer : sorter.sortedDataBuffer; if (!dataBuffer) { return; } if (dataBuffer !== this._gpuSortedDataBuffer) { - // The underlying buffer was (re)allocated, so every camera mesh must rebind against it. + // 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; + } activeViewInfos.forEach((info) => { if (!info.splatIndexBufferSet) { info.mesh._thinInstanceSetSplatIndexBuffer(dataBuffer, instanceCount); info.splatIndexBufferSet = true; } + info.mesh._indirectDrawBuffer = indirectBuffer; info.frameIdLastUpdate = frameId; }); this._readyToDisplay = true; diff --git a/packages/dev/core/src/Meshes/mesh.pure.ts b/packages/dev/core/src/Meshes/mesh.pure.ts index d78533e17241..354fdc04626a 100644 --- a/packages/dev/core/src/Meshes/mesh.pure.ts +++ b/packages/dev/core/src/Meshes/mesh.pure.ts @@ -539,6 +539,14 @@ 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). + */ + public _indirectDrawBuffer: Nullable = null; + /** @internal */ public _instanceDataStorage: _InstanceDataStorage; @@ -2141,6 +2149,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 fbb9884c1104..cc93d3382444 100644 --- a/packages/dev/core/src/Meshes/thinInstanceMesh.pure.ts +++ b/packages/dev/core/src/Meshes/thinInstanceMesh.pure.ts @@ -390,6 +390,10 @@ export function RegisterThinInstanceMesh(): void { 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)); 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/tools/devHost/src/testScene/createScene.ts b/packages/tools/devHost/src/testScene/createScene.ts index d2c02ab55cf7..5b698a7c3c23 100644 --- a/packages/tools/devHost/src/testScene/createScene.ts +++ b/packages/tools/devHost/src/testScene/createScene.ts @@ -1,44 +1,50 @@ import { type Engine } from "core/Engines/engine"; import { Scene } from "core/scene"; -import { FreeCamera } from "core/Cameras/freeCamera"; +import { ArcRotateCamera } from "core/Cameras/arcRotateCamera"; import { HemisphericLight } from "core/Lights/hemisphericLight"; -import { StandardMaterial } from "core/Materials/standardMaterial"; -import { MeshBuilder } from "core/Meshes/meshBuilder"; import { Vector3 } from "core/Maths/math.vector"; +import { GaussianSplattingMesh } from "core/Meshes/GaussianSplatting/gaussianSplattingMesh"; +import { GaussianSplattingMeshBase, GaussianSplattingGpuSortMode } from "core/Meshes/GaussianSplatting/gaussianSplattingMeshBase"; +// Register the .splat file loader (side-effect). +import "loaders/SPLAT/splatFileLoader"; + +const SplatUrl = "https://raw.githubusercontent.com/CedricGuillemet/dump/master/Halo_Believe.splat"; // eslint-disable-next-line @typescript-eslint/naming-convention, no-restricted-syntax export const createScene = async function (engine: Engine, canvas: HTMLCanvasElement): Promise { - // This creates a basic Babylon Scene object (non-mesh) const scene = new Scene(engine); - // This creates and positions a free camera (non-mesh) - const camera = new FreeCamera("camera1", new Vector3(0, 5, -10), scene); - - // This targets the camera to scene origin - camera.setTarget(Vector3.Zero()); - - // This attaches the camera to the canvas + const camera = new ArcRotateCamera("camera1", 0.6, 1.3, 8, Vector3.Zero(), scene); camera.attachControl(canvas, true); - - // This creates a light, aiming 0,1,0 - to the sky (non-mesh) - const light = new HemisphericLight("light1", new Vector3(0, 1, 0), scene); - - // Default intensity is 1. Let's dim the light a small amount - light.intensity = 0.7; - - // Default material for all meshes - const material = new StandardMaterial("sceneMaterial", scene); - - // Our built-in 'sphere' shape. Params: name, options, scene - const sphere = MeshBuilder.CreateSphere("sphere", {diameter: 2, segments: 32}, scene); - sphere.material = material; - - // Move the sphere upward 1/2 its height - sphere.position.y = 1; - - // Our built-in 'ground' shape. Params: name, options, scene - const ground = MeshBuilder.CreateGround("ground", {width: 6, height: 6}, scene); - ground.material = material; + camera.wheelPrecision = 50; + camera.minZ = 0.01; + new HemisphericLight("light1", new Vector3(0, 1, 0), scene); + + // Force the WebGPU GPU sort + cull path (bypasses the CPU worker fallback). + GaussianSplattingMeshBase.GpuSortMode = GaussianSplattingGpuSortMode.Gpu; + + const gs = new GaussianSplattingMesh("gs", null, scene); + await gs.loadFileAsync(SplatUrl); + + gs.computeWorldMatrix(true); + const bounds = gs.getBoundingInfo().boundingSphere; + camera.setTarget(bounds.centerWorld.clone()); + camera.radius = bounds.radiusWorld * 2.5 + 1; + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (window as any).__gsScene = scene; + scene.onBeforeRenderObservable.add(() => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const win = window as any; + const anyGs = gs as unknown as { _useGpuSort: boolean; _worker: unknown; _gpuSorter: unknown }; + win.__gsVerify = { + found: true, + useGpuSort: anyGs._useGpuSort, + hasWorker: !!anyGs._worker, + hasGpuSorter: !!anyGs._gpuSorter, + renderedSplatCount: gs.renderedSplatCount, + }; + }); return scene; }; diff --git a/packages/tools/devHost/src/testScene/main.ts b/packages/tools/devHost/src/testScene/main.ts index 0fc09372fd7a..47a74f5f58d0 100644 --- a/packages/tools/devHost/src/testScene/main.ts +++ b/packages/tools/devHost/src/testScene/main.ts @@ -1,8 +1,8 @@ import { type Scene } from "core/scene"; -import { Engine } from "core/Engines/engine"; +import { WebGPUEngine } from "core/Engines/webgpuEngine"; +import "core/Engines/WebGPU/Extensions/engine.computeShader"; import { createScene as createSceneTs } from "./createScene"; -import { createScene as createSceneJs } from "./createSceneJS.js"; /** * Main entry point for the default scene for the devhost @@ -15,19 +15,11 @@ export async function Main(searchParams: URLSearchParams): Promise { canvas.id = "babylon-canvas"; mainDiv.appendChild(canvas); - // Whether to use the TS or JS scene files, default to TS - const useTsParam = searchParams.get("usets"); - const useTs = useTsParam !== "false"; // Default to true if not specified + // Setup a WebGPU engine (required for the Gaussian Splatting GPU sort/cull compute path) + const engine = new WebGPUEngine(canvas, { antialias: true }); + await engine.initAsync(); - // Setup the engine and create the scene - const engine = new Engine(canvas, true); - - let scene: Scene | undefined = undefined; - if (useTs) { - scene = await createSceneTs(engine, canvas); - } else { - scene = await createSceneJs(engine, canvas); - } + const scene: Scene = await createSceneTs(engine, canvas); // Register a render loop to repeatedly render the scene engine.runRenderLoop(function () { diff --git a/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json b/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json index f95ef1f67a7e..2f6fad96cb7a 100644 --- a/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json +++ b/scripts/treeshaking/side-effects-manifest/core/ShadersWGSL.json @@ -194,6 +194,9 @@ "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"], From 8e35713aeb094e51669b6f7c815873c7012ddb4d Mon Sep 17 00:00:00 2001 From: Cedric Guillemet <1312968+CedricGuillemet@users.noreply.github.com> Date: Thu, 2 Jul 2026 16:55:55 +0200 Subject: [PATCH 4/9] fix build --- packages/dev/core/src/Meshes/mesh.pure.ts | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/packages/dev/core/src/Meshes/mesh.pure.ts b/packages/dev/core/src/Meshes/mesh.pure.ts index 354fdc04626a..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 */ @@ -543,9 +546,16 @@ export class Mesh extends AbstractMesh implements IGetSetVerticesData { * @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). + * 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 _indirectDrawBuffer: Nullable = null; + public get _indirectDrawBuffer(): Nullable { + return this._internalMeshDataInfo._indirectDrawBuffer; + } + + public set _indirectDrawBuffer(value: Nullable) { + this._internalMeshDataInfo._indirectDrawBuffer = value; + } /** @internal */ public _instanceDataStorage: _InstanceDataStorage; From 2fd440f976f5cbf729157839a9fc99c66c1abf37 Mon Sep 17 00:00:00 2001 From: Cedric Guillemet <1312968+CedricGuillemet@users.noreply.github.com> Date: Thu, 2 Jul 2026 17:52:56 +0200 Subject: [PATCH 5/9] culling with static camera --- .../gaussianSplattingMeshBase.pure.ts | 57 +++++++++++-------- 1 file changed, 34 insertions(+), 23 deletions(-) diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts index 663b10096fa4..cf3b853b3141 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts @@ -517,21 +517,21 @@ export class GaussianSplattingMeshBase extends Mesh { public static EnableGpuIndirectDraw = true; /** - * When true (and {@link EnableGpuIndirectDraw} is on), the WebGPU GPU-sort path runs a per-frame frustum cull - * on the GPU and draws only the visible splats via an indirect draw. The cull test matches the vertex shader - * exactly, so it never removes a splat the shader would have shown. Set false to draw all sorted splats. - * - * Currently defaults to false: the cull path is complete but blocked by a one-dispatch UniformBuffer-upload - * lag for compute dispatched inside the mesh render pass (the cull reads the previous dispatch's clip matrix). - * The GPU sort + non-culled indirect draw path is unaffected and remains the shipping default. + * 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; + /** - * When true (and {@link EnableGpuIndirectDraw} is on), the WebGPU GPU-sort path runs a per-frame frustum cull - * on the GPU 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. + * 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 EnableGpuCulling = true; + public static GpuCullResumeDelayFrames = 2; /** @internal */ public _vertexCount = 0; @@ -547,6 +547,9 @@ export class GaussianSplattingMeshBase extends Mesh { 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(); @@ -1433,13 +1436,19 @@ export class GaussianSplattingMeshBase extends Mesh { 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. - // Also keep re-running while culling is enabled but not yet active (its compute shaders may still be - // compiling on the first frames — without this a static camera would never start culling). const cullingWanted = GaussianSplattingMeshBase.EnableGpuIndirectDraw && GaussianSplattingMeshBase.EnableGpuCulling; - const needSort = this._gpuSortDirty || !this._readyToDisplay || (cullingWanted && !this._gpuCulling) || this._isSortStateDirty(primary, worldMatrix, camera); + 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 - // per-frame cull tracks the camera even when the (more expensive) depth sort is skipped. + // cull tracks the camera when it re-engages. camera.getViewMatrix().multiplyToRef(camera.getProjectionMatrix(), this._gpuViewProjMatrix); worldMatrix.multiplyToRef(this._gpuViewProjMatrix, this._gpuClipMatrix); @@ -1473,18 +1482,20 @@ export class GaussianSplattingMeshBase extends Mesh { primary.sortAppliedId = primary.sortRequestId; } - // Post-sort GPU frustum cull runs EVERY frame (not only when the sort re-runs) so it tracks the camera: - // once the sort has produced a sorted-index buffer, the cull compacts the visible splats and writes the - // indirect instance count using the same 1.2*w bounds the vertex shader uses. - if (cullingWanted && !this._gpuSortDirty) { + // 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 if (!cullingWanted) { + } else { this._gpuCulling = false; } - // When culling is active the draw reads the compacted visible list and the GPU-written indirect args; + // 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 = GaussianSplattingMeshBase.EnableGpuIndirectDraw && GaussianSplattingMeshBase.EnableGpuCulling && this._gpuCulling; + const cullingActive = this._gpuCulling && this._gpuCullSettleFrames >= cullDelay; const dataBuffer = cullingActive ? sorter.culledDrawBuffer : sorter.sortedDataBuffer; if (!dataBuffer) { return; From bd9d758e53042503c121891bce450c8c7bd03ee1 Mon Sep 17 00:00:00 2001 From: Cedric Guillemet <1312968+CedricGuillemet@users.noreply.github.com> Date: Thu, 2 Jul 2026 18:23:25 +0200 Subject: [PATCH 6/9] revert --- .../devHost/src/testScene/createScene.ts | 68 +++++++++---------- packages/tools/devHost/src/testScene/main.ts | 20 ++++-- 2 files changed, 45 insertions(+), 43 deletions(-) diff --git a/packages/tools/devHost/src/testScene/createScene.ts b/packages/tools/devHost/src/testScene/createScene.ts index 5b698a7c3c23..d2c02ab55cf7 100644 --- a/packages/tools/devHost/src/testScene/createScene.ts +++ b/packages/tools/devHost/src/testScene/createScene.ts @@ -1,50 +1,44 @@ import { type Engine } from "core/Engines/engine"; import { Scene } from "core/scene"; -import { ArcRotateCamera } from "core/Cameras/arcRotateCamera"; +import { FreeCamera } from "core/Cameras/freeCamera"; import { HemisphericLight } from "core/Lights/hemisphericLight"; +import { StandardMaterial } from "core/Materials/standardMaterial"; +import { MeshBuilder } from "core/Meshes/meshBuilder"; import { Vector3 } from "core/Maths/math.vector"; -import { GaussianSplattingMesh } from "core/Meshes/GaussianSplatting/gaussianSplattingMesh"; -import { GaussianSplattingMeshBase, GaussianSplattingGpuSortMode } from "core/Meshes/GaussianSplatting/gaussianSplattingMeshBase"; -// Register the .splat file loader (side-effect). -import "loaders/SPLAT/splatFileLoader"; - -const SplatUrl = "https://raw.githubusercontent.com/CedricGuillemet/dump/master/Halo_Believe.splat"; // eslint-disable-next-line @typescript-eslint/naming-convention, no-restricted-syntax export const createScene = async function (engine: Engine, canvas: HTMLCanvasElement): Promise { + // This creates a basic Babylon Scene object (non-mesh) const scene = new Scene(engine); - const camera = new ArcRotateCamera("camera1", 0.6, 1.3, 8, Vector3.Zero(), scene); + // This creates and positions a free camera (non-mesh) + const camera = new FreeCamera("camera1", new Vector3(0, 5, -10), scene); + + // This targets the camera to scene origin + camera.setTarget(Vector3.Zero()); + + // This attaches the camera to the canvas camera.attachControl(canvas, true); - camera.wheelPrecision = 50; - camera.minZ = 0.01; - new HemisphericLight("light1", new Vector3(0, 1, 0), scene); - - // Force the WebGPU GPU sort + cull path (bypasses the CPU worker fallback). - GaussianSplattingMeshBase.GpuSortMode = GaussianSplattingGpuSortMode.Gpu; - - const gs = new GaussianSplattingMesh("gs", null, scene); - await gs.loadFileAsync(SplatUrl); - - gs.computeWorldMatrix(true); - const bounds = gs.getBoundingInfo().boundingSphere; - camera.setTarget(bounds.centerWorld.clone()); - camera.radius = bounds.radiusWorld * 2.5 + 1; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (window as any).__gsScene = scene; - scene.onBeforeRenderObservable.add(() => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const win = window as any; - const anyGs = gs as unknown as { _useGpuSort: boolean; _worker: unknown; _gpuSorter: unknown }; - win.__gsVerify = { - found: true, - useGpuSort: anyGs._useGpuSort, - hasWorker: !!anyGs._worker, - hasGpuSorter: !!anyGs._gpuSorter, - renderedSplatCount: gs.renderedSplatCount, - }; - }); + + // This creates a light, aiming 0,1,0 - to the sky (non-mesh) + const light = new HemisphericLight("light1", new Vector3(0, 1, 0), scene); + + // Default intensity is 1. Let's dim the light a small amount + light.intensity = 0.7; + + // Default material for all meshes + const material = new StandardMaterial("sceneMaterial", scene); + + // Our built-in 'sphere' shape. Params: name, options, scene + const sphere = MeshBuilder.CreateSphere("sphere", {diameter: 2, segments: 32}, scene); + sphere.material = material; + + // Move the sphere upward 1/2 its height + sphere.position.y = 1; + + // Our built-in 'ground' shape. Params: name, options, scene + const ground = MeshBuilder.CreateGround("ground", {width: 6, height: 6}, scene); + ground.material = material; return scene; }; diff --git a/packages/tools/devHost/src/testScene/main.ts b/packages/tools/devHost/src/testScene/main.ts index 47a74f5f58d0..0fc09372fd7a 100644 --- a/packages/tools/devHost/src/testScene/main.ts +++ b/packages/tools/devHost/src/testScene/main.ts @@ -1,8 +1,8 @@ import { type Scene } from "core/scene"; -import { WebGPUEngine } from "core/Engines/webgpuEngine"; -import "core/Engines/WebGPU/Extensions/engine.computeShader"; +import { Engine } from "core/Engines/engine"; import { createScene as createSceneTs } from "./createScene"; +import { createScene as createSceneJs } from "./createSceneJS.js"; /** * Main entry point for the default scene for the devhost @@ -15,11 +15,19 @@ export async function Main(searchParams: URLSearchParams): Promise { canvas.id = "babylon-canvas"; mainDiv.appendChild(canvas); - // Setup a WebGPU engine (required for the Gaussian Splatting GPU sort/cull compute path) - const engine = new WebGPUEngine(canvas, { antialias: true }); - await engine.initAsync(); + // Whether to use the TS or JS scene files, default to TS + const useTsParam = searchParams.get("usets"); + const useTs = useTsParam !== "false"; // Default to true if not specified - const scene: Scene = await createSceneTs(engine, canvas); + // Setup the engine and create the scene + const engine = new Engine(canvas, true); + + let scene: Scene | undefined = undefined; + if (useTs) { + scene = await createSceneTs(engine, canvas); + } else { + scene = await createSceneJs(engine, canvas); + } // Register a render loop to repeatedly render the scene engine.runRenderLoop(function () { From 13e2df73cc32364b9274f73f7d0e71adc5fbf945 Mon Sep 17 00:00:00 2001 From: Cedric Guillemet <1312968+CedricGuillemet@users.noreply.github.com> Date: Fri, 3 Jul 2026 11:15:46 +0200 Subject: [PATCH 7/9] feedback --- .../core/src/Compute/prefixSumCompute.pure.ts | 4 +++- .../src/Engines/WebGPU/webgpuDrawContext.ts | 10 ++++++++++ .../dev/core/src/Engines/webgpuEngine.pure.ts | 16 +++++++++++++++- .../gaussianSplattingMeshBase.pure.ts | 17 ++++++++++++++++- .../ShadersWGSL/prefixSumScanBlock.compute.fx | 8 +++++++- 5 files changed, 51 insertions(+), 4 deletions(-) diff --git a/packages/dev/core/src/Compute/prefixSumCompute.pure.ts b/packages/dev/core/src/Compute/prefixSumCompute.pure.ts index 86e8b1789ef7..7ee1ba6a5be7 100644 --- a/packages/dev/core/src/Compute/prefixSumCompute.pure.ts +++ b/packages/dev/core/src/Compute/prefixSumCompute.pure.ts @@ -117,7 +117,9 @@ export class PrefixSumCompute { } private _scanRecursive(data: DataBuffer, count: number, level: number): void { - if (count <= 1) { + // 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); 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/webgpuEngine.pure.ts b/packages/dev/core/src/Engines/webgpuEngine.pure.ts index 9448c1afff00..7511bed68b21 100644 --- a/packages/dev/core/src/Engines/webgpuEngine.pure.ts +++ b/packages/dev/core/src/Engines/webgpuEngine.pure.ts @@ -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; } @@ -3863,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); } diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts index cf3b853b3141..8fe6d1855ee7 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts @@ -66,7 +66,10 @@ const Native = IsNative ? _native : null; * @see GaussianSplattingMeshBase.GpuSortMode */ export enum GaussianSplattingGpuSortMode { - /** Pick automatically. Currently keeps the CPU worker / native path until GPU sorting is complete. */ + /** + * 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, @@ -3301,6 +3304,18 @@ export class GaussianSplattingMeshBase extends Mesh { 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/ShadersWGSL/prefixSumScanBlock.compute.fx b/packages/dev/core/src/ShadersWGSL/prefixSumScanBlock.compute.fx index 6ece483d6119..7532dc00870c 100644 --- a/packages/dev/core/src/ShadersWGSL/prefixSumScanBlock.compute.fx +++ b/packages/dev/core/src/ShadersWGSL/prefixSumScanBlock.compute.fx @@ -19,7 +19,13 @@ fn main(@builtin(global_invocation_id) gid : vec3u, @builtin(local_invocation_id let g = gid.x; let t = lid.x; - let v = select(0u, data[g], g < params.count); + // 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(); From 9544c4d9359a465a0275bd482011f1a1c8c1ae93 Mon Sep 17 00:00:00 2001 From: Cedric Guillemet <1312968+CedricGuillemet@users.noreply.github.com> Date: Fri, 3 Jul 2026 14:10:53 +0200 Subject: [PATCH 8/9] shader loading --- .../core/src/Compute/prefixSumCompute.pure.ts | 161 ------ .../dev/core/src/Compute/prefixSumCompute.ts | 180 +++++- packages/dev/core/src/Compute/pure.ts | 1 - .../gaussianSplattingGpuSorter.pure.ts | 504 ---------------- .../gaussianSplattingGpuSorter.ts | 541 +++++++++++++++++- .../gaussianSplattingMeshBase.pure.ts | 2 +- .../gaussianSplattingMeshBase.ts | 1 - .../core/src/Meshes/GaussianSplatting/pure.ts | 1 - packages/public/@babylonjs/core/package.json | 2 - .../side-effects-manifest/core/Compute.json | 3 +- .../side-effects-manifest/core/Meshes.json | 1 - 11 files changed, 710 insertions(+), 687 deletions(-) delete mode 100644 packages/dev/core/src/Compute/prefixSumCompute.pure.ts delete mode 100644 packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts diff --git a/packages/dev/core/src/Compute/prefixSumCompute.pure.ts b/packages/dev/core/src/Compute/prefixSumCompute.pure.ts deleted file mode 100644 index 7ee1ba6a5be7..000000000000 --- a/packages/dev/core/src/Compute/prefixSumCompute.pure.ts +++ /dev/null @@ -1,161 +0,0 @@ -/** This file must only contain pure code and pure imports */ - -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 is intended to be - * reused by the interval culling pass in a later phase. - */ -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[] = []; - - /** - * Creates a new prefix-sum compute helper. - * @param engine the (WebGPU) engine to run the compute passes on - */ - public constructor(engine: AbstractEngine) { - this._engine = engine; - } - - private _ensureShaders(): void { - if (this._scanBlock) { - 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 compiled and ready to dispatch. - * @returns true when the scan can run - */ - public isReady(): boolean { - 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/Compute/prefixSumCompute.ts b/packages/dev/core/src/Compute/prefixSumCompute.ts index e76765a2d7c1..40923108e357 100644 --- a/packages/dev/core/src/Compute/prefixSumCompute.ts +++ b/packages/dev/core/src/Compute/prefixSumCompute.ts @@ -1,4 +1,178 @@ -export * from "./prefixSumCompute.pure"; +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"; -import "../ShadersWGSL/prefixSumScanBlock.compute"; -import "../ShadersWGSL/prefixSumAddOffsets.compute"; +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/Compute/pure.ts b/packages/dev/core/src/Compute/pure.ts index aacba9a1d7ec..6ae86a736bb1 100644 --- a/packages/dev/core/src/Compute/pure.ts +++ b/packages/dev/core/src/Compute/pure.ts @@ -1,4 +1,3 @@ /** Pure barrel — re-exports only side-effect-free modules */ export * from "./computeEffect"; export * from "./computeShader.pure"; -export * from "./prefixSumCompute.pure"; diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts deleted file mode 100644 index e62a13321a1d..000000000000 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.pure.ts +++ /dev/null @@ -1,504 +0,0 @@ -/** This file must only contain pure code and pure imports */ - -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.pure"; - -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. - */ -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; - - /** - * 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. - * @param engine the (WebGPU) engine to run the compute passes on - */ - public constructor(engine: AbstractEngine) { - this._engine = engine; - } - - 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 { - 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._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/gaussianSplattingGpuSorter.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts index 6352e0104e68..e81d670399c0 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts @@ -1,10 +1,531 @@ -export * from "./gaussianSplattingGpuSorter.pure"; - -import "core/Compute/prefixSumCompute"; -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"; +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 8fe6d1855ee7..94a178f43de9 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts @@ -27,7 +27,7 @@ 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.pure"; +import { GaussianSplattingGpuSorter } from "./gaussianSplattingGpuSorter"; import { type DataBuffer } from "core/Buffers/dataBuffer"; // eslint-disable-next-line @typescript-eslint/naming-convention diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts index 14db97466ff3..5062e4eb9f9d 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts @@ -1,4 +1,3 @@ export * from "./gaussianSplattingMeshBase.pure"; import "core/Meshes/thinInstanceMesh"; -import "./gaussianSplattingGpuSorter"; diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/pure.ts index aefea68e2324..ab7f89df1939 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/pure.ts @@ -1,7 +1,6 @@ /** Pure barrel — re-exports only side-effect-free modules */ export * from "./gaussianSplattingCompoundMesh.pure"; export * from "./gaussianSplattingDebugger.pure"; -export * from "./gaussianSplattingGpuSorter.pure"; export * from "./gaussianSplattingMesh.pure"; export * from "./gaussianSplattingMeshBase.pure"; export * from "./gaussianSplattingPartProxyMesh.pure"; diff --git a/packages/public/@babylonjs/core/package.json b/packages/public/@babylonjs/core/package.json index 3155bd19a4e7..b3d83f51ee4f 100644 --- a/packages/public/@babylonjs/core/package.json +++ b/packages/public/@babylonjs/core/package.json @@ -122,7 +122,6 @@ "Cameras/virtualJoysticksCamera.js", "Collisions/collisionCoordinator.js", "Compute/computeShader.types.js", - "Compute/prefixSumCompute.js", "Culling/Helper/computeShaderBoundingHelper.js", "Culling/Helper/transformFeedbackBoundingHelper.js", "Culling/Octrees/octreeSceneComponent.js", @@ -458,7 +457,6 @@ "Meshes/Builders/tubeBuilder.js", "Meshes/GaussianSplatting/gaussianSplattingCompoundMesh.js", "Meshes/GaussianSplatting/gaussianSplattingDebugger.js", - "Meshes/GaussianSplatting/gaussianSplattingGpuSorter.js", "Meshes/GaussianSplatting/gaussianSplattingMesh.js", "Meshes/GaussianSplatting/gaussianSplattingMeshBase.js", "Meshes/GaussianSplatting/gaussianSplattingPartProxyMesh.js", diff --git a/scripts/treeshaking/side-effects-manifest/core/Compute.json b/scripts/treeshaking/side-effects-manifest/core/Compute.json index a15b090c9a99..383e36a2ea46 100644 --- a/scripts/treeshaking/side-effects-manifest/core/Compute.json +++ b/scripts/treeshaking/side-effects-manifest/core/Compute.json @@ -2,7 +2,6 @@ "version": 1, "files": { "Compute/computeShader.ts": ["top-level-call"], - "Compute/computeShader.types.ts": ["declare-module"], - "Compute/prefixSumCompute.ts": ["bare-import"] + "Compute/computeShader.types.ts": ["declare-module"] } } diff --git a/scripts/treeshaking/side-effects-manifest/core/Meshes.json b/scripts/treeshaking/side-effects-manifest/core/Meshes.json index 8624970c35fc..7d442d2ec47c 100644 --- a/scripts/treeshaking/side-effects-manifest/core/Meshes.json +++ b/scripts/treeshaking/side-effects-manifest/core/Meshes.json @@ -22,7 +22,6 @@ "Meshes/Builders/tubeBuilder.ts": ["top-level-call"], "Meshes/GaussianSplatting/gaussianSplattingCompoundMesh.ts": ["top-level-call"], "Meshes/GaussianSplatting/gaussianSplattingDebugger.ts": ["bare-import"], - "Meshes/GaussianSplatting/gaussianSplattingGpuSorter.ts": ["bare-import"], "Meshes/GaussianSplatting/gaussianSplattingMesh.ts": ["bare-import", "top-level-call"], "Meshes/GaussianSplatting/gaussianSplattingMeshBase.ts": ["bare-import"], "Meshes/GaussianSplatting/gaussianSplattingPartProxyMesh.ts": ["top-level-call"], From 532adbe462f18973ce150a9ee14048a3a10b3a53 Mon Sep 17 00:00:00 2001 From: Cedric Guillemet <1312968+CedricGuillemet@users.noreply.github.com> Date: Fri, 3 Jul 2026 14:24:06 +0200 Subject: [PATCH 9/9] release gc pressure --- .../gaussianSplattingMeshBase.pure.ts | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts index 94a178f43de9..8c4ba2b1048f 100644 --- a/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts +++ b/packages/dev/core/src/Meshes/GaussianSplatting/gaussianSplattingMeshBase.pure.ts @@ -1433,8 +1433,17 @@ export class GaussianSplattingMeshBase extends Mesh { } // 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; - const primary = (activeCamera && activeViewInfos.find((info) => info.camera === activeCamera)) || activeViewInfos[0]; + 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); @@ -1520,14 +1529,16 @@ export class GaussianSplattingMeshBase extends Mesh { // When culling is active the cull's finalize pass already wrote the indirect args on the GPU. indirectBuffer = sorter.indirectArgsBuffer; } - activeViewInfos.forEach((info) => { + // 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; }