playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
813 lines (810 loc) • 43.6 kB
JavaScript
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, SHADERSTAGE_COMPUTE, SHADERLANGUAGE_WGSL, BUFFERUSAGE_INDIRECT, UNIFORMTYPE_UINT, UNIFORMTYPE_FLOAT, UNIFORMTYPE_MAT4, PIXELFORMAT_RGBA16U } 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 { 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 { 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 computeSplatSource from '../shader-lib/wgsl/chunks/gsplat/vert/gsplatComputeSplat.js';
const INITIAL_TILE_ENTRY_MULTIPLIER = 2.5;
const ENTRY_HEADROOM_MULTIPLIER = 2.0;
const SHRINK_THRESHOLD = 200;
const INITIAL_LARGE_SPLAT_CAPACITY = 16384;
const TILE_SIZE = 16;
const MAX_TILES = 65535;
const CACHE_STRIDE = 7;
const MAX_CHUNKS_PER_TILE = 8;
const _viewProjMat = new Mat4();
const _viewProjData = new Float32Array(16);
const _viewData = new Float32Array(16);
const _fogColorLinear = new Color();
const _fogColorArray = new Float32Array(3);
class GSplatComputeLocalRenderer extends GSplatRenderer {
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._exposure = exposure ?? 1.0;
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);
}
_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 proj = cam.projectionMatrix;
_viewProjMat.mul2(proj, view);
_viewProjData.set(_viewProjMat.data);
_viewData.set(view.data);
const focal = width * proj.data[0];
const alphaClip = pickMode ? this._alphaClip : 1.0 / 255.0;
this.fisheyeProj.update(this._fisheye, camera.fov, proj);
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();
}
_createCommonIncludes() {
const cincludes = new Map();
cincludes.set('gsplatCommonCS', computeGsplatCommonSource);
cincludes.set('gsplatTileIntersectCS', computeGsplatTileIntersectSource);
return cincludes;
}
_createBitonicIncludes() {
const cincludes = 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 = 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)
]);
this._largeSplatShader = new Shader(device, {
name: 'GSplatLocalTileCountLarge',
shaderLanguage: SHADERLANGUAGE_WGSL,
cshader: computeGsplatLocalTileCountLargeSource,
cincludes,
computeBindGroupFormat: this._largeSplatBindGroupFormat,
computeUniformBufferFormats: {
uniforms: ubf
}
});
}
{
const maxDim = device.limits.maxComputeWorkgroupsPerDimension || 65535;
const largePrepDefines = 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());
const cdefines = new Map();
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;
}
constructor(device, node, cameraNode, layer, workBuffer){
super(device, node, cameraNode, layer, workBuffer), this._pickSet = null, this._needsFramePassRegister = false, this._textureSize = 0, this._minPixelSize = 2.0, this._minContribution = 3.0, this._alphaClip = 0.3, this._exposure = 1.0, this._numSplats = 0, this._compactedSplatIds = null, this._sortElementCountBuffer = null, this._projCacheBuffer = null, this._tileEntriesBuffer = null, this._pairBuffer = null, this._countersBuffer = null, this._splatPairStartBuffer = null, this._splatPairCountBuffer = null, this._largeSplatIdsBuffer = null, this._largeSplatIdsCapacity = INITIAL_LARGE_SPLAT_CAPACITY, this._allocatedLargeSplatCapacity = 0, this._allocatedSplatCapacity = 0, this._allocatedEntryCapacity = 0, this._tileEntryMultiplier = INITIAL_TILE_ENTRY_MULTIPLIER, this._lastBufferSubmitVersion = -1, this._lastReadbackEntryCount = -1, this._shrinkFrameCount = 0, this._fisheye = 0, this._placeEntryPrepDispatchBuffer = null, this._placeEntryPrepCompute = null, this._largeSplatDispatchBuffer = null, this._largeSplatPrepCompute = null, this._lastDrawSlot = 0, this._lastNumTilesX = 0;
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;
});
}
}
export { GSplatComputeLocalRenderer };