UNPKG

playcanvas

Version:

Open-source WebGL/WebGPU 3D engine for the web

867 lines (866 loc) 43.2 kB
import { Compute } from "../../platform/graphics/compute.js"; import { Shader } from "../../platform/graphics/shader.js"; import { StorageBuffer } from "../../platform/graphics/storage-buffer.js"; import { BindGroupFormat, BindStorageBufferFormat, BindUniformBufferFormat } from "../../platform/graphics/bind-group-format.js"; import { UniformBufferFormat, UniformFormat } from "../../platform/graphics/uniform-buffer-format.js"; import { BUFFERUSAGE_COPY_DST, BUFFERUSAGE_COPY_SRC, BUFFERUSAGE_INDIRECT, PIXELFORMAT_RGBA16U, SHADERLANGUAGE_WGSL, SHADERSTAGE_COMPUTE, UNIFORMTYPE_FLOAT, UNIFORMTYPE_MAT4, UNIFORMTYPE_UINT } from "../../platform/graphics/constants.js"; import { GSPLAT_FORWARD, PROJECTION_ORTHOGRAPHIC, FOG_NONE, GSPLAT_DEBUG_HEATMAP } from "../constants.js"; import { Color } from "../../core/math/color.js"; import { Mat4 } from "../../core/math/mat4.js"; import { Camera } from "../camera.js"; import { GSplatRenderer } from "./gsplat-renderer.js"; import { FramePassGSplatComputeLocal } from "./frame-pass-gsplat-compute-local.js"; import { computeGsplatLocalDispatchPrepSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-dispatch-prep.js"; import { computeGsplatLocalDispatchPrepLargeSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-dispatch-prep-large.js"; import { computeGsplatLocalTileCountSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-tile-count.js"; import { computeGsplatLocalTileCountLargeSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-tile-count-large.js"; import { computeGsplatLocalPlaceEntriesSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-place-entries.js"; import { computeGsplatLocalPlaceEntriesLargeSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-place-entries-large.js"; import { computeGsplatLocalTileSortSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-tile-sort.js"; import { computeGsplatLocalClassifySource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-classify.js"; import { computeGsplatLocalBucketSortSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-bucket-sort.js"; import { computeGsplatLocalChunkSortSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-chunk-sort.js"; import { computeGsplatLocalCopySource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-copy.js"; import { computeGsplatLocalBitonicSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-bitonic.js"; import { computeGsplatProjectCommonSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-project-common.js"; import { computeGsplatCommonSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-common.js"; import { computeGsplatTileIntersectSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-tile-intersect.js"; import { GSplatTileComposite } from "./gsplat-tile-composite.js"; import { GSplatLocalDispatchSet } from "./gsplat-local-dispatch-set.js"; import { ALPHA_VISIBILITY_THRESHOLD, CACHE_STRIDE } from "./constants.js"; import computeSplatSource from "../shader-lib/wgsl/chunks/gsplat/vert/gsplatComputeSplat.js"; const INITIAL_TILE_ENTRY_MULTIPLIER = 2.5; const ENTRY_HEADROOM_MULTIPLIER = 2; const SHRINK_THRESHOLD = 200; const INITIAL_LARGE_SPLAT_CAPACITY = 16384; const TILE_SIZE = 16; const MAX_TILES = 65535; const MAX_CHUNKS_PER_TILE = 8; const _viewProjMat = new Mat4(); const _viewProjData = new Float32Array(16); const _viewData = new Float32Array(16); const _shaderProjMat = new Mat4(); const _fogColorLinear = new Color(); const _fogColorArray = new Float32Array(3); class GSplatComputeLocalRenderer extends GSplatRenderer { _mainSet; _pickSet = null; framePass; tileComposite; _needsFramePassRegister = false; _textureSize = 0; _minPixelSize = 2; _minContribution = 3; _alphaClip = 0.3; _alphaClipForward = ALPHA_VISIBILITY_THRESHOLD; _exposure = 1; _numSplats = 0; _compactedSplatIds = null; _sortElementCountBuffer = null; // --- Shared splat/entry-dependent buffers --- _projCacheBuffer = null; _tileEntriesBuffer = null; // Pair buffer for scatter-free tile binning: stores packed (tileIdx << 16 | localOffset) // per splat-tile intersection. Same size as tileEntries (one u32 per pair). _pairBuffer = null; _countersBuffer = null; _splatPairStartBuffer = null; _splatPairCountBuffer = null; // --- Large-splat deferred processing buffers --- _largeSplatIdsBuffer = null; _largeSplatIdsCapacity = INITIAL_LARGE_SPLAT_CAPACITY; _allocatedLargeSplatCapacity = 0; _allocatedSplatCapacity = 0; _allocatedEntryCapacity = 0; _tileEntryMultiplier = INITIAL_TILE_ENTRY_MULTIPLIER; _lastBufferSubmitVersion = -1; _lastReadbackEntryCount = -1; _shrinkFrameCount = 0; _fisheye = 0; _dataSource; _placeEntriesShader; _placeEntriesBindGroupFormat; _placeEntryPrepShader; _placeEntryPrepBindGroupFormat; _placeEntryPrepDispatchBuffer = null; _placeEntryPrepCompute = null; _largeSplatShader; _largeSplatBindGroupFormat; _largeSplatPrepShader; _largeSplatPrepBindGroupFormat; _largeSplatDispatchBuffer = null; _largeSplatPrepCompute = null; _largePlaceEntriesShader; _largePlaceEntriesBindGroupFormat; _classifyShader; _classifyBindGroupFormat; _sortShader; _sortBindGroupFormat; _bucketSortShader; _bucketSortBindGroupFormat; _copyShader; _copyBindGroupFormat; _chunkSortShader; _chunkSortBindGroupFormat; constructor(device, node, cameraNode, layer, workBuffer) { super(device, node, cameraNode, layer, workBuffer); this._dataSource = workBuffer; this._createSharedShaders(); this._mainSet = this._createDispatchSet(false); this.framePass = new FramePassGSplatComputeLocal(this); const thisCamera = cameraNode.camera; this.tileComposite = new GSplatTileComposite(device, node, (camera) => { const renderMode = this.renderMode ?? 0; return thisCamera.camera === camera && (renderMode & GSPLAT_FORWARD) !== 0 && this._mainSet._rasterizeTileListBuffer !== null; }); } setDataSource(source) { this._dataSource = source; } destroy() { this._unregisterFramePass(); if (this.renderMode) { if (this.renderMode & GSPLAT_FORWARD) { this.layer.removeMeshInstances([this.tileComposite.meshInstance], true); } } this._mainSet.destroy(); this._pickSet?.destroy(); this._placeEntriesShader.destroy(); this._placeEntriesBindGroupFormat.destroy(); this._placeEntryPrepShader.destroy(); this._placeEntryPrepBindGroupFormat.destroy(); this._placeEntryPrepDispatchBuffer?.destroy(); this._largeSplatShader.destroy(); this._largeSplatBindGroupFormat.destroy(); this._largeSplatPrepShader.destroy(); this._largeSplatPrepBindGroupFormat.destroy(); this._largeSplatDispatchBuffer?.destroy(); this._largePlaceEntriesShader.destroy(); this._largePlaceEntriesBindGroupFormat.destroy(); this._classifyShader.destroy(); this._classifyBindGroupFormat.destroy(); this._sortShader.destroy(); this._sortBindGroupFormat.destroy(); this._bucketSortShader.destroy(); this._bucketSortBindGroupFormat.destroy(); this._copyShader.destroy(); this._copyBindGroupFormat.destroy(); this._chunkSortShader.destroy(); this._chunkSortBindGroupFormat.destroy(); this._projCacheBuffer?.destroy(); this._depthBuffer?.destroy(); this._tileEntriesBuffer?.destroy(); this._pairBuffer?.destroy(); this._countersBuffer?.destroy(); this._splatPairStartBuffer?.destroy(); this._splatPairCountBuffer?.destroy(); this._largeSplatIdsBuffer?.destroy(); this.tileComposite.destroy(); super.destroy(); } get material() { return this.tileComposite.material; } setRenderMode(renderMode) { const oldRenderMode = this.renderMode ?? 0; const wasForward = (oldRenderMode & GSPLAT_FORWARD) !== 0; const isForward = (renderMode & GSPLAT_FORWARD) !== 0; if (!wasForward && isForward) { this.layer.addMeshInstances([this.tileComposite.meshInstance], true); this._registerFramePass(); } if (wasForward && !isForward) { this.layer.removeMeshInstances([this.tileComposite.meshInstance], true); this._unregisterFramePass(); } super.setRenderMode(renderMode); } frameUpdate(gsplat, exposure, fogParams) { if (this._needsFramePassRegister) { this._registerFramePass(); } this._minPixelSize = gsplat.minPixelSize; this._minContribution = gsplat.minContribution; this._alphaClip = gsplat.alphaClip; this._alphaClipForward = gsplat.alphaClipForward; this._exposure = exposure ?? 1; this._fisheye = gsplat.fisheye; this._fogParams = fogParams ?? null; this._debugMode = gsplat.debug; const formatHash = this._dataSource.format.hash; if (formatHash !== this._formatHash) { this._formatHash = formatHash; this._invalidateCountCompute(); } } setCompactedData(compactedSplatIds, sortElementCountBuffer, textureSize, numSplats) { this._compactedSplatIds = compactedSplatIds; this._sortElementCountBuffer = sortElementCountBuffer; this._textureSize = textureSize; this._numSplats = numSplats; } _registerFramePass() { const camera = this.cameraNode.camera?.camera; if (camera) { if (!camera.beforePasses.includes(this.framePass)) { camera.beforePasses.push(this.framePass); } this._needsFramePassRegister = false; } else { this._needsFramePassRegister = true; } } _unregisterFramePass() { this._needsFramePassRegister = false; const camera = this.cameraNode.camera?.camera; if (camera) { const idx = camera.beforePasses.indexOf(this.framePass); if (idx !== -1) { camera.beforePasses.splice(idx, 1); } } } resizeOutputTexture(width, height) { this._mainSet.resizeOutputTextures(width, height); } _ensureSharedBuffers(numSplats) { const device = this.device; const canResize = device.submitVersion !== this._lastBufferSubmitVersion; const readbackValue = this._lastReadbackEntryCount; if (readbackValue !== -1) { this._lastReadbackEntryCount = -1; const currentCapacity = this._allocatedEntryCapacity; if (currentCapacity > 0 && readbackValue < currentCapacity / 2) { this._shrinkFrameCount++; } else { this._shrinkFrameCount = 0; } if (this._shrinkFrameCount >= SHRINK_THRESHOLD && numSplats > 0) { const target = Math.max(INITIAL_TILE_ENTRY_MULTIPLIER, readbackValue / numSplats * ENTRY_HEADROOM_MULTIPLIER); this._tileEntryMultiplier = Math.max(target, this._tileEntryMultiplier * 0.9); this._shrinkFrameCount = 0; } } if (!canResize) return; if (numSplats > this._allocatedSplatCapacity) { this._projCacheBuffer?.destroy(); this._depthBuffer?.destroy(); this._splatPairStartBuffer?.destroy(); this._splatPairCountBuffer?.destroy(); this._allocatedSplatCapacity = numSplats; this._projCacheBuffer = new StorageBuffer(device, numSplats * CACHE_STRIDE * 4); this._depthBuffer = new StorageBuffer(device, numSplats * 4); this._splatPairStartBuffer = new StorageBuffer(device, numSplats * 4); this._splatPairCountBuffer = new StorageBuffer(device, numSplats * 4); } if (!this._countersBuffer) { this._countersBuffer = new StorageBuffer(device, 8, BUFFERUSAGE_COPY_DST | BUFFERUSAGE_COPY_SRC); } if (this._largeSplatIdsCapacity > this._allocatedLargeSplatCapacity) { this._largeSplatIdsBuffer?.destroy(); this._allocatedLargeSplatCapacity = this._largeSplatIdsCapacity; this._largeSplatIdsBuffer = new StorageBuffer(device, this._largeSplatIdsCapacity * 4); } const requiredEntryCapacity = Math.ceil(numSplats * this._tileEntryMultiplier); const needsGrow = requiredEntryCapacity > this._allocatedEntryCapacity; const needsShrink = this._allocatedEntryCapacity > 0 && requiredEntryCapacity * 2 < this._allocatedEntryCapacity; if (needsGrow || needsShrink) { this._tileEntriesBuffer?.destroy(); this._pairBuffer?.destroy(); this._allocatedEntryCapacity = requiredEntryCapacity; this._tileEntriesBuffer = new StorageBuffer(device, requiredEntryCapacity * 4, BUFFERUSAGE_COPY_DST); this._pairBuffer = new StorageBuffer(device, requiredEntryCapacity * 4); } this._lastBufferSubmitVersion = device.submitVersion; } dispatch() { const set = this._mainSet; const outputTex = set.outputTexture; if (!outputTex) return; const width = outputTex.width; const height = outputTex.height; this._dispatchPipeline(set, width, height, false); this.tileComposite.update( this._lastDrawSlot, outputTex, set._rasterizeTileListBuffer, this._lastNumTilesX, width, height ); } dispatchPick(cam, width, height) { if (!this._pickSet) { this._pickSet = this._createDispatchSet(true); } const set = this._pickSet; set.resizeOutputTextures(width, height); this._dispatchPipeline(set, width, height, true); return this.tileComposite.prepareForPicking( this._lastDrawSlot, set.pickIdTexture, set.pickDepthTexture, set._rasterizeTileListBuffer, this._lastNumTilesX, width, height ); } _lastDrawSlot = 0; _lastNumTilesX = 0; _dispatchPipeline(set, width, height, pickMode) { const numSplats = this._numSplats; if (!this._compactedSplatIds || !this._sortElementCountBuffer || numSplats === 0) return; const device = this.device; let numTilesX = Math.ceil(width / TILE_SIZE); let numTilesY = Math.ceil(height / TILE_SIZE); if (numTilesX * numTilesY > MAX_TILES) { const scale = Math.sqrt(MAX_TILES / (numTilesX * numTilesY)); numTilesX = Math.max(1, Math.floor(numTilesX * scale)); numTilesY = Math.max(1, Math.floor(numTilesY * scale)); } const numTiles = numTilesX * numTilesY; this._ensureSharedBuffers(numSplats); set.ensureTileBuffers(numTiles); const maxEntries = this._allocatedEntryCapacity; const ds = this._dataSource; const camera = this.cameraNode.camera; const cam = camera.camera; const view = cam.viewMatrix; const flipY = !!camera.renderTarget?.flipY; const webgpu = device.isWebGPU; _viewProjMat.mul2(Camera.applyShaderProjectionTransform(cam.projectionMatrix, _shaderProjMat, flipY, webgpu), view); _viewProjData.set(_viewProjMat.data); _viewData.set(view.data); const focal = width * _shaderProjMat.data[0]; const alphaClip = pickMode ? this._alphaClip : Math.max(ALPHA_VISIBILITY_THRESHOLD, this._alphaClipForward); this.fisheyeProj.update(this._fisheye, camera.fov, cam.projectionMatrix); const fisheyeEnabled = this.fisheyeProj.enabled; const createCountShader = (pick, fisheye) => this._createCountShaderAndFormat(pick, fisheye); const countCompute = set.getCountCompute(fisheyeEnabled, createCountShader); this._placeEntryPrepCompute.setParameter("sortElementCount", this._sortElementCountBuffer); this._placeEntryPrepCompute.setParameter("dispatchArgs", this._placeEntryPrepDispatchBuffer); this._placeEntryPrepCompute.setupDispatch(1, 1, 1); device.computeDispatch([this._placeEntryPrepCompute], "GSplatLocalDispatchPrep"); set._tileSplatCountsBuffer.clear(); this._countersBuffer.clear(); countCompute.setParameter("compactedSplatIds", this._compactedSplatIds); countCompute.setParameter("sortElementCount", this._sortElementCountBuffer); countCompute.setParameter("projCache", this._projCacheBuffer); countCompute.setParameter("tileSplatCounts", set._tileSplatCountsBuffer); countCompute.setParameter("pairBuffer", this._pairBuffer); countCompute.setParameter("countersBuffer", this._countersBuffer); countCompute.setParameter("splatPairStart", this._splatPairStartBuffer); countCompute.setParameter("splatPairCount", this._splatPairCountBuffer); countCompute.setParameter("largeSplatIds", this._largeSplatIdsBuffer); countCompute.setParameter("depthBuffer", this._depthBuffer); for (const stream of ds.format.streams) { countCompute.setParameter(stream.name, ds.getTexture(stream.name)); } for (const stream of ds.format.extraStreams) { countCompute.setParameter(stream.name, ds.getTexture(stream.name)); } countCompute.setParameter("splatTextureSize", this._textureSize); countCompute.setParameter("numTilesX", numTilesX); countCompute.setParameter("numTilesY", numTilesY); countCompute.setParameter("viewProj", _viewProjData); countCompute.setParameter("viewMatrix", _viewData); countCompute.setParameter("focal", focal); countCompute.setParameter("viewportWidth", width); countCompute.setParameter("viewportHeight", height); countCompute.setParameter("nearClip", cam.nearClip); countCompute.setParameter("farClip", cam.farClip); countCompute.setParameter("minPixelSize", this._minPixelSize * 0.5); countCompute.setParameter("isOrtho", cam.projection === PROJECTION_ORTHOGRAPHIC ? 1 : 0); countCompute.setParameter("exposure", this._exposure); countCompute.setParameter("alphaClip", alphaClip); countCompute.setParameter("minContribution", this._minContribution); if (fisheyeEnabled) { const fp = this.fisheyeProj; countCompute.setParameter("fisheye_k", fp.k); countCompute.setParameter("fisheye_inv_k", fp.invK); countCompute.setParameter("fisheye_projMat00", fp.projMat00); countCompute.setParameter("fisheye_projMat11", fp.projMat11); } countCompute.setupIndirectDispatch(0, this._placeEntryPrepDispatchBuffer); device.computeDispatch([countCompute], pickMode ? "GSplatPickTileCount" : "GSplatLocalTileCount"); this._largeSplatPrepCompute.setParameter("countersBuffer", this._countersBuffer); this._largeSplatPrepCompute.setParameter("dispatchArgs", this._largeSplatDispatchBuffer); this._largeSplatPrepCompute.setParameter("largeSplatIds", this._largeSplatIdsBuffer); this._largeSplatPrepCompute.setupDispatch(1, 1, 1); device.computeDispatch([this._largeSplatPrepCompute], pickMode ? "GSplatPickLargeSplatPrep" : "GSplatLocalLargeSplatPrep"); set.largeSplatCompute.setParameter("projCache", this._projCacheBuffer); set.largeSplatCompute.setParameter("tileSplatCounts", set._tileSplatCountsBuffer); set.largeSplatCompute.setParameter("pairBuffer", this._pairBuffer); set.largeSplatCompute.setParameter("countersBuffer", this._countersBuffer); set.largeSplatCompute.setParameter("splatPairStart", this._splatPairStartBuffer); set.largeSplatCompute.setParameter("splatPairCount", this._splatPairCountBuffer); set.largeSplatCompute.setParameter("largeSplatIds", this._largeSplatIdsBuffer); set.largeSplatCompute.setParameter("numTilesX", numTilesX); set.largeSplatCompute.setParameter("numTilesY", numTilesY); set.largeSplatCompute.setParameter("viewportWidth", width); set.largeSplatCompute.setParameter("viewportHeight", height); set.largeSplatCompute.setParameter("alphaClip", alphaClip); set.largeSplatCompute.setupIndirectDispatch(0, this._largeSplatDispatchBuffer); device.computeDispatch([set.largeSplatCompute], pickMode ? "GSplatPickLargeTileCount" : "GSplatLocalLargeTileCount"); set.prefixSumKernel.resize(set._tileSplatCountsBuffer, numTiles + 1); set.prefixSumKernel.dispatch(device); set.placeEntriesCompute.setParameter("pairBuffer", this._pairBuffer); set.placeEntriesCompute.setParameter("splatPairStart", this._splatPairStartBuffer); set.placeEntriesCompute.setParameter("splatPairCount", this._splatPairCountBuffer); set.placeEntriesCompute.setParameter("tileSplatCounts", set._tileSplatCountsBuffer); set.placeEntriesCompute.setParameter("tileEntries", this._tileEntriesBuffer); set.placeEntriesCompute.setParameter("sortElementCount", this._sortElementCountBuffer); set.placeEntriesCompute.setupIndirectDispatch(1, this._placeEntryPrepDispatchBuffer); device.computeDispatch([set.placeEntriesCompute], pickMode ? "GSplatPickPlaceEntries" : "GSplatLocalPlaceEntries"); set.largePlaceEntriesCompute.setParameter("pairBuffer", this._pairBuffer); set.largePlaceEntriesCompute.setParameter("splatPairStart", this._splatPairStartBuffer); set.largePlaceEntriesCompute.setParameter("splatPairCount", this._splatPairCountBuffer); set.largePlaceEntriesCompute.setParameter("tileSplatCounts", set._tileSplatCountsBuffer); set.largePlaceEntriesCompute.setParameter("tileEntries", this._tileEntriesBuffer); set.largePlaceEntriesCompute.setParameter("largeSplatIds", this._largeSplatIdsBuffer); set.largePlaceEntriesCompute.setParameter("countersBuffer", this._countersBuffer); set.largePlaceEntriesCompute.setupIndirectDispatch(0, this._largeSplatDispatchBuffer); device.computeDispatch([set.largePlaceEntriesCompute], pickMode ? "GSplatPickLargePlaceEntries" : "GSplatLocalLargePlaceEntries"); set._tileListCountsBuffer.clear(); set._totalChunksBuffer.clear(); set._chunkSortIndirectBuffer.clear(); const indirectSlot = device.getIndirectDispatchSlot(3); const drawSlot = device.getIndirectDrawSlot(1); set.classifyCompute.setParameter("tileSplatCounts", set._tileSplatCountsBuffer); set.classifyCompute.setParameter("smallTileList", set._smallTileListBuffer); set.classifyCompute.setParameter("largeTileList", set._largeTileListBuffer); set.classifyCompute.setParameter("rasterizeTileList", set._rasterizeTileListBuffer); set.classifyCompute.setParameter("tileListCounts", set._tileListCountsBuffer); set.classifyCompute.setParameter("indirectDispatchArgs", device.indirectDispatchBuffer); set.classifyCompute.setParameter("largeTileOverflowBases", set._largeTileOverflowBasesBuffer); set.classifyCompute.setParameter("indirectDrawArgs", device.indirectDrawBuffer); set.classifyCompute.setParameter("numTiles", numTiles); set.classifyCompute.setParameter("dispatchSlotOffset", indirectSlot * 3); set.classifyCompute.setParameter("bufferCapacity", maxEntries); set.classifyCompute.setParameter("maxWorkgroupsPerDim", device.limits.maxComputeWorkgroupsPerDimension || 65535); set.classifyCompute.setParameter("drawSlot", drawSlot); set.classifyCompute.setupDispatch(1, 1, 1); device.computeDispatch([set.classifyCompute], pickMode ? "GSplatPickClassify" : "GSplatLocalClassify"); set.bucketSortCompute.setParameter("tileEntries", this._tileEntriesBuffer); set.bucketSortCompute.setParameter("largeTileOverflowBases", set._largeTileOverflowBasesBuffer); set.bucketSortCompute.setParameter("tileSplatCounts", set._tileSplatCountsBuffer); set.bucketSortCompute.setParameter("depthBuffer", this._depthBuffer); set.bucketSortCompute.setParameter("largeTileList", set._largeTileListBuffer); set.bucketSortCompute.setParameter("chunkRanges", set._chunkRangesBuffer); set.bucketSortCompute.setParameter("totalChunks", set._totalChunksBuffer); set.bucketSortCompute.setParameter("tileListCounts", set._tileListCountsBuffer); set.bucketSortCompute.setParameter("bufferCapacity", maxEntries); set.bucketSortCompute.setParameter("maxChunks", numTiles * MAX_CHUNKS_PER_TILE); set.bucketSortCompute.setupIndirectDispatch(indirectSlot + 1); device.computeDispatch([set.bucketSortCompute], pickMode ? "GSplatPickBucketSort" : "GSplatLocalBucketSort"); set.copyCompute.setParameter("totalChunks", set._totalChunksBuffer); set.copyCompute.setParameter("chunkSortIndirect", set._chunkSortIndirectBuffer); set.copyCompute.setParameter("maxChunks", numTiles * MAX_CHUNKS_PER_TILE); set.copyCompute.setParameter("maxWorkgroupsPerDim", device.limits.maxComputeWorkgroupsPerDimension || 65535); set.copyCompute.setupDispatch(1, 1, 1); device.computeDispatch([set.copyCompute], pickMode ? "GSplatPickCopy" : "GSplatLocalCopy"); set.sortCompute.setParameter("tileEntries", this._tileEntriesBuffer); set.sortCompute.setParameter("tileSplatCounts", set._tileSplatCountsBuffer); set.sortCompute.setParameter("depthBuffer", this._depthBuffer); set.sortCompute.setParameter("smallTileList", set._smallTileListBuffer); set.sortCompute.setParameter("tileListCounts", set._tileListCountsBuffer); set.sortCompute.setupIndirectDispatch(indirectSlot); device.computeDispatch([set.sortCompute], pickMode ? "GSplatPickTileSort" : "GSplatLocalTileSort"); set.chunkSortCompute.setParameter("tileEntries", this._tileEntriesBuffer); set.chunkSortCompute.setParameter("depthBuffer", this._depthBuffer); set.chunkSortCompute.setParameter("chunkRanges", set._chunkRangesBuffer); set.chunkSortCompute.setParameter("totalChunks", set._totalChunksBuffer); set.chunkSortCompute.setParameter("maxChunks", numTiles * MAX_CHUNKS_PER_TILE); set.chunkSortCompute.setupIndirectDispatch(0, set._chunkSortIndirectBuffer); device.computeDispatch([set.chunkSortCompute], pickMode ? "GSplatPickChunkSort" : "GSplatLocalChunkSort"); const hasLinearDepth = cam.shaderParams.sceneDepthMapLinear; const sceneDepthMap = hasLinearDepth ? device.scope.resolve("uSceneDepthMap").value : null; const useDepth = !pickMode && sceneDepthMap; const fogParams = this._fogParams; const fogType = fogParams && fogParams.type !== FOG_NONE ? fogParams.type : "none"; const heatmap = !pickMode && this._debugMode === GSPLAT_DEBUG_HEATMAP; const rasterizeCompute = set.getRasterizeCompute(pickMode, useDepth, fogType, heatmap); rasterizeCompute.setParameter("screenWidth", width); rasterizeCompute.setParameter("screenHeight", height); rasterizeCompute.setParameter("numTilesX", numTilesX); rasterizeCompute.setParameter("nearClip", cam.nearClip); rasterizeCompute.setParameter("farClip", cam.farClip); rasterizeCompute.setParameter("alphaClip", alphaClip); if (fogType !== "none") { _fogColorLinear.linear(fogParams.color); _fogColorArray[0] = _fogColorLinear.r; _fogColorArray[1] = _fogColorLinear.g; _fogColorArray[2] = _fogColorLinear.b; rasterizeCompute.setParameter("fog_color", _fogColorArray); rasterizeCompute.setParameter("fog_start", fogParams.start); rasterizeCompute.setParameter("fog_end", fogParams.end); rasterizeCompute.setParameter("fog_density", fogParams.density); } rasterizeCompute.setParameter("tileEntries", this._tileEntriesBuffer); rasterizeCompute.setParameter("tileSplatCounts", set._tileSplatCountsBuffer); rasterizeCompute.setParameter("projCache", this._projCacheBuffer); rasterizeCompute.setParameter("rasterizeTileList", set._rasterizeTileListBuffer); rasterizeCompute.setParameter("tileListCounts", set._tileListCountsBuffer); rasterizeCompute.setParameter("depthBuffer", this._depthBuffer); if (pickMode) { rasterizeCompute.setParameter("pickIdTexture", set.pickIdTexture); rasterizeCompute.setParameter("pickDepthTexture", set.pickDepthTexture); } else { rasterizeCompute.setParameter("outputTexture", set.outputTexture); if (useDepth) { rasterizeCompute.setParameter("sceneDepthMap", sceneDepthMap); } } rasterizeCompute.setupIndirectDispatch(indirectSlot + 2); device.computeDispatch([rasterizeCompute], pickMode ? "GSplatPickRasterize" : "GSplatLocalRasterize"); this._lastDrawSlot = drawSlot; this._lastNumTilesX = numTilesX; this._scheduleReadback(numSplats, numTiles, set); } _scheduleReadback(numSplats, numTiles, set) { if (numSplats === 0) return; const capturedNumSplats = numSplats; const readback1 = set._tileSplatCountsBuffer.read(numTiles * 4, 4); const readback2 = set._tileListCountsBuffer.read(3 * 4, 4); const readback3 = this._countersBuffer.read(4, 4); Promise.all([readback1, readback2, readback3]).then(([r1, r2, r3]) => { const totalEntries = new Uint32Array(r1.buffer, r1.byteOffset, 1)[0]; const totalOverflow = new Uint32Array(r2.buffer, r2.byteOffset, 1)[0]; const needed = totalEntries + totalOverflow; if (capturedNumSplats > 0) { this._tileEntryMultiplier = Math.max(this._tileEntryMultiplier, needed / capturedNumSplats * ENTRY_HEADROOM_MULTIPLIER); this._tileEntryMultiplier = Math.max(this._tileEntryMultiplier, INITIAL_TILE_ENTRY_MULTIPLIER); } this._lastReadbackEntryCount = needed; const largeSplatDemand = new Uint32Array(r3.buffer, r3.byteOffset, 1)[0]; if (largeSplatDemand > this._largeSplatIdsCapacity) { this._largeSplatIdsCapacity = Math.ceil(largeSplatDemand * 1.2); } }).catch(() => { }); } _invalidateCountCompute() { this._mainSet.destroyCountResources(); this._pickSet?.destroyCountResources(); } // ---- Shader / BindGroupFormat creation ---- _createCommonIncludes() { const cincludes = /* @__PURE__ */ new Map(); cincludes.set("gsplatCommonCS", computeGsplatCommonSource); cincludes.set("gsplatTileIntersectCS", computeGsplatTileIntersectSource); return cincludes; } _createBitonicIncludes() { const cincludes = /* @__PURE__ */ new Map(); cincludes.set("gsplatLocalBitonicCS", computeGsplatLocalBitonicSource); return cincludes; } _createSharedShaders() { const device = this.device; this._placeEntriesBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("pairBuffer", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("splatPairStart", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("splatPairCount", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("tileSplatCounts", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("tileEntries", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("sortElementCount", SHADERSTAGE_COMPUTE, true) ]); this._placeEntriesShader = new Shader(device, { name: "GSplatLocalPlaceEntries", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalPlaceEntriesSource, computeBindGroupFormat: this._placeEntriesBindGroupFormat }); { const maxDim = device.limits.maxComputeWorkgroupsPerDimension || 65535; const prepDefines = /* @__PURE__ */ new Map(); prepDefines.set("{MAX_DIM}", maxDim.toString()); prepDefines.set("{SPLATS_PER_WG}", "256"); prepDefines.set("{SPLATS_PER_WG_MINUS_1}", "255"); this._placeEntryPrepBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("sortElementCount", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("dispatchArgs", SHADERSTAGE_COMPUTE) ]); this._placeEntryPrepShader = new Shader(device, { name: "GSplatLocalPlaceEntryPrep", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalDispatchPrepSource, cdefines: prepDefines, computeBindGroupFormat: this._placeEntryPrepBindGroupFormat }); this._placeEntryPrepDispatchBuffer = new StorageBuffer(device, 6 * 4, BUFFERUSAGE_INDIRECT); this._placeEntryPrepCompute = new Compute(device, this._placeEntryPrepShader); } { const cincludes = this._createCommonIncludes(); const ubf = new UniformBufferFormat(device, [ new UniformFormat("numTilesX", UNIFORMTYPE_UINT), new UniformFormat("numTilesY", UNIFORMTYPE_UINT), new UniformFormat("viewportWidth", UNIFORMTYPE_FLOAT), new UniformFormat("viewportHeight", UNIFORMTYPE_FLOAT), new UniformFormat("alphaClip", UNIFORMTYPE_FLOAT) ]); this._largeSplatBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("projCache", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("tileSplatCounts", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("pairBuffer", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("countersBuffer", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("splatPairStart", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("splatPairCount", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("largeSplatIds", SHADERSTAGE_COMPUTE, true), new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE) ]); const cdefines = /* @__PURE__ */ new Map(); cdefines.set("{CACHE_STRIDE}", CACHE_STRIDE.toString()); this._largeSplatShader = new Shader(device, { name: "GSplatLocalTileCountLarge", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalTileCountLargeSource, cincludes, cdefines, computeBindGroupFormat: this._largeSplatBindGroupFormat, computeUniformBufferFormats: { uniforms: ubf } }); } { const maxDim = device.limits.maxComputeWorkgroupsPerDimension || 65535; const largePrepDefines = /* @__PURE__ */ new Map(); largePrepDefines.set("{MAX_DIM}", maxDim.toString()); this._largeSplatPrepBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("countersBuffer", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("dispatchArgs", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("largeSplatIds", SHADERSTAGE_COMPUTE, true) ]); this._largeSplatPrepShader = new Shader(device, { name: "GSplatLocalLargeSplatPrep", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalDispatchPrepLargeSource, cdefines: largePrepDefines, computeBindGroupFormat: this._largeSplatPrepBindGroupFormat }); this._largeSplatDispatchBuffer = new StorageBuffer(device, 3 * 4, BUFFERUSAGE_INDIRECT); this._largeSplatPrepCompute = new Compute(device, this._largeSplatPrepShader); } this._largePlaceEntriesBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("pairBuffer", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("splatPairStart", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("splatPairCount", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("tileSplatCounts", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("tileEntries", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("largeSplatIds", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("countersBuffer", SHADERSTAGE_COMPUTE, true) ]); this._largePlaceEntriesShader = new Shader(device, { name: "GSplatLocalPlaceEntriesLarge", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalPlaceEntriesLargeSource, computeBindGroupFormat: this._largePlaceEntriesBindGroupFormat }); { const ubf = new UniformBufferFormat(device, [ new UniformFormat("numTiles", UNIFORMTYPE_UINT), new UniformFormat("dispatchSlotOffset", UNIFORMTYPE_UINT), new UniformFormat("bufferCapacity", UNIFORMTYPE_UINT), new UniformFormat("maxWorkgroupsPerDim", UNIFORMTYPE_UINT), new UniformFormat("drawSlot", UNIFORMTYPE_UINT) ]); this._classifyBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("tileSplatCounts", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("smallTileList", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("largeTileList", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("rasterizeTileList", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("tileListCounts", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("indirectDispatchArgs", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("largeTileOverflowBases", SHADERSTAGE_COMPUTE), new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("indirectDrawArgs", SHADERSTAGE_COMPUTE) ]); this._classifyShader = new Shader(device, { name: "GSplatLocalClassify", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalClassifySource, computeBindGroupFormat: this._classifyBindGroupFormat, computeUniformBufferFormats: { uniforms: ubf } }); } this._sortBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("tileEntries", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("tileSplatCounts", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("depthBuffer", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("smallTileList", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("tileListCounts", SHADERSTAGE_COMPUTE, true) ]); this._sortShader = new Shader(device, { name: "GSplatLocalTileSort", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalTileSortSource, cincludes: this._createBitonicIncludes(), computeBindGroupFormat: this._sortBindGroupFormat }); { const ubf = new UniformBufferFormat(device, [ new UniformFormat("bufferCapacity", UNIFORMTYPE_UINT), new UniformFormat("maxChunks", UNIFORMTYPE_UINT) ]); this._bucketSortBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("tileEntries", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("largeTileOverflowBases", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("tileSplatCounts", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("depthBuffer", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("largeTileList", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("chunkRanges", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("totalChunks", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("tileListCounts", SHADERSTAGE_COMPUTE, true), new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE) ]); this._bucketSortShader = new Shader(device, { name: "GSplatLocalBucketSort", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalBucketSortSource, computeBindGroupFormat: this._bucketSortBindGroupFormat, computeUniformBufferFormats: { uniforms: ubf } }); } { const ubf = new UniformBufferFormat(device, [ new UniformFormat("maxChunks", UNIFORMTYPE_UINT), new UniformFormat("maxWorkgroupsPerDim", UNIFORMTYPE_UINT) ]); this._copyBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("totalChunks", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("chunkSortIndirect", SHADERSTAGE_COMPUTE), new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE) ]); this._copyShader = new Shader(device, { name: "GSplatLocalCopy", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalCopySource, computeBindGroupFormat: this._copyBindGroupFormat, computeUniformBufferFormats: { uniforms: ubf } }); } { const ubf = new UniformBufferFormat(device, [ new UniformFormat("maxChunks", UNIFORMTYPE_UINT) ]); this._chunkSortBindGroupFormat = new BindGroupFormat(device, [ new BindStorageBufferFormat("tileEntries", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("depthBuffer", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("chunkRanges", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("totalChunks", SHADERSTAGE_COMPUTE, true), new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE) ]); this._chunkSortShader = new Shader(device, { name: "GSplatLocalChunkSort", shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalChunkSortSource, cincludes: this._createBitonicIncludes(), computeBindGroupFormat: this._chunkSortBindGroupFormat, computeUniformBufferFormats: { uniforms: ubf } }); } } _createCountShaderAndFormat(pickMode, fisheyeEnabled) { const device = this.device; const uniforms = [ new UniformFormat("splatTextureSize", UNIFORMTYPE_UINT), new UniformFormat("numTilesX", UNIFORMTYPE_UINT), new UniformFormat("numTilesY", UNIFORMTYPE_UINT), new UniformFormat("viewProj", UNIFORMTYPE_MAT4), new UniformFormat("viewMatrix", UNIFORMTYPE_MAT4), new UniformFormat("focal", UNIFORMTYPE_FLOAT), new UniformFormat("viewportWidth", UNIFORMTYPE_FLOAT), new UniformFormat("viewportHeight", UNIFORMTYPE_FLOAT), new UniformFormat("nearClip", UNIFORMTYPE_FLOAT), new UniformFormat("farClip", UNIFORMTYPE_FLOAT), new UniformFormat("minPixelSize", UNIFORMTYPE_FLOAT), new UniformFormat("isOrtho", UNIFORMTYPE_UINT), new UniformFormat("exposure", UNIFORMTYPE_FLOAT), new UniformFormat("alphaClip", UNIFORMTYPE_FLOAT), new UniformFormat("minContribution", UNIFORMTYPE_FLOAT) ]; if (fisheyeEnabled) { uniforms.push( new UniformFormat("fisheye_k", UNIFORMTYPE_FLOAT), new UniformFormat("fisheye_inv_k", UNIFORMTYPE_FLOAT), new UniformFormat("fisheye_projMat00", UNIFORMTYPE_FLOAT), new UniformFormat("fisheye_projMat11", UNIFORMTYPE_FLOAT) ); } const uniformBufferFormat = new UniformBufferFormat(device, uniforms); const wbFormat = this._dataSource.format; const fixedBindings = [ new BindStorageBufferFormat("compactedSplatIds", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("sortElementCount", SHADERSTAGE_COMPUTE, true), new BindStorageBufferFormat("projCache", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("tileSplatCounts", SHADERSTAGE_COMPUTE), new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("pairBuffer", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("countersBuffer", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("splatPairStart", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("splatPairCount", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("largeSplatIds", SHADERSTAGE_COMPUTE), new BindStorageBufferFormat("depthBuffer", SHADERSTAGE_COMPUTE) ]; const bindGroupFormat = new BindGroupFormat(device, [ ...fixedBindings, ...wbFormat.getComputeBindFormats() ]); const cincludes = this._createCommonIncludes(); cincludes.set("gsplatComputeSplatCS", computeSplatSource); cincludes.set("gsplatFormatDeclCS", wbFormat.getComputeInputDeclarations(fixedBindings.length)); cincludes.set("gsplatFormatReadCS", wbFormat.getReadCode()); cincludes.set("gsplatProjectCommonCS", computeGsplatProjectCommonSource); const cdefines = /* @__PURE__ */ new Map(); cdefines.set("{CACHE_STRIDE}", CACHE_STRIDE.toString()); if (fisheyeEnabled) { cdefines.set("GSPLAT_FISHEYE", ""); } if (pickMode) { cdefines.set("PICK_MODE", ""); } const colorStream = wbFormat.getStream("dataColor"); if (colorStream && colorStream.format !== PIXELFORMAT_RGBA16U) { cdefines.set("GSPLAT_COLOR_FLOAT", ""); } const suffix = fisheyeEnabled ? "Fisheye" : ""; const shader = new Shader(device, { name: pickMode ? `GSplatLocalTileCountPick${suffix}` : `GSplatLocalTileCount${suffix}`, shaderLanguage: SHADERLANGUAGE_WGSL, cshader: computeGsplatLocalTileCountSource, cincludes, cdefines, computeBindGroupFormat: bindGroupFormat, computeUniformBufferFormats: { uniforms: uniformBufferFormat } }); return { shader, bindGroupFormat }; } _createDispatchSet(pickMode) { const device = this.device; const set = new GSplatLocalDispatchSet(device, pickMode); set.placeEntriesCompute = new Compute(device, this._placeEntriesShader, pickMode ? "GSplatPickPlaceEntries" : "GSplatLocalPlaceEntries"); set.largeSplatCompute = new Compute(device, this._largeSplatShader, pickMode ? "GSplatPickLargeTileCount" : "GSplatLocalLargeTileCount"); set.largePlaceEntriesCompute = new Compute(device, this._largePlaceEntriesShader, pickMode ? "GSplatPickLargePlaceEntries" : "GSplatLocalLargePlaceEntries"); set.classifyCompute = new Compute(device, this._classifyShader, pickMode ? "GSplatPickClassify" : "GSplatLocalClassify"); set.sortCompute = new Compute(device, this._sortShader, pickMode ? "GSplatPickTileSort" : "GSplatLocalTileSort"); set.bucketSortCompute = new Compute(device, this._bucketSortShader, pickMode ? "GSplatPickBucketSort" : "GSplatLocalBucketSort"); set.copyCompute = new Compute(device, this._copyShader, pickMode ? "GSplatPickCopy" : "GSplatLocalCopy"); set.chunkSortCompute = new Compute(device, this._chunkSortShader, pickMode ? "GSplatPickChunkSort" : "GSplatLocalChunkSort"); return set; } } export { GSplatComputeLocalRenderer };