playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
254 lines (251 loc) • 10.9 kB
JavaScript
import { Texture } from '../../platform/graphics/texture.js';
import { StorageBuffer } from '../../platform/graphics/storage-buffer.js';
import { Compute } from '../../platform/graphics/compute.js';
import { Shader } from '../../platform/graphics/shader.js';
import { BindUniformBufferFormat, BindStorageBufferFormat, BindStorageTextureFormat, BindTextureFormat, BindGroupFormat } from '../../platform/graphics/bind-group-format.js';
import { UniformFormat, UniformBufferFormat } from '../../platform/graphics/uniform-buffer-format.js';
import { BUFFERUSAGE_COPY_DST, BUFFERUSAGE_COPY_SRC, BUFFERUSAGE_INDIRECT, UNIFORMTYPE_UINT, UNIFORMTYPE_FLOAT, UNIFORMTYPE_VEC3, SHADERSTAGE_COMPUTE, PIXELFORMAT_R32U, PIXELFORMAT_RGBA16F, TEXTUREDIMENSION_2D, SAMPLETYPE_UNFILTERABLE_FLOAT, SHADERLANGUAGE_WGSL, FILTER_NEAREST } from '../../platform/graphics/constants.js';
import { PrefixSumKernel } from '../graphics/prefix-sum-kernel.js';
import { shaderChunksWGSL } from '../shader-lib/wgsl/collections/shader-chunks-wgsl.js';
import { computeGsplatLocalRasterizeSource } from '../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-rasterize.js';
const MAX_CHUNKS_PER_TILE = 8;
class GSplatLocalDispatchSet {
resizeOutputTextures(width, height) {
if (this.pickMode) {
this.pickIdTexture?.resize(width, height);
this.pickDepthTexture?.resize(width, height);
} else {
this.outputTexture?.resize(width, height);
}
}
ensureTileBuffers(numTiles) {
const requiredTileSlots = numTiles + 1;
if (requiredTileSlots <= this._allocatedTileCapacity) return;
this._tileSplatCountsBuffer?.destroy();
this._smallTileListBuffer?.destroy();
this._largeTileListBuffer?.destroy();
this._largeTileOverflowBasesBuffer?.destroy();
this._rasterizeTileListBuffer?.destroy();
this._tileListCountsBuffer?.destroy();
this._chunkRangesBuffer?.destroy();
this._totalChunksBuffer?.destroy();
this._chunkSortIndirectBuffer?.destroy();
this._allocatedTileCapacity = requiredTileSlots;
this._tileSplatCountsBuffer = new StorageBuffer(this.device, requiredTileSlots * 4, BUFFERUSAGE_COPY_DST | BUFFERUSAGE_COPY_SRC);
this._smallTileListBuffer = new StorageBuffer(this.device, numTiles * 4);
this._largeTileListBuffer = new StorageBuffer(this.device, numTiles * 4);
this._largeTileOverflowBasesBuffer = new StorageBuffer(this.device, numTiles * 4);
this._rasterizeTileListBuffer = new StorageBuffer(this.device, numTiles * 4);
this._tileListCountsBuffer = new StorageBuffer(this.device, 4 * 4, BUFFERUSAGE_COPY_DST | BUFFERUSAGE_COPY_SRC);
const maxChunks = numTiles * MAX_CHUNKS_PER_TILE;
this._chunkRangesBuffer = new StorageBuffer(this.device, maxChunks * 8);
this._totalChunksBuffer = new StorageBuffer(this.device, 1 * 4, BUFFERUSAGE_COPY_DST);
this._chunkSortIndirectBuffer = new StorageBuffer(this.device, 3 * 4, BUFFERUSAGE_COPY_DST | BUFFERUSAGE_INDIRECT);
this.prefixSumKernel.destroyPasses();
}
getCountCompute(fisheyeEnabled, createShaderAndFormat) {
if (fisheyeEnabled) {
if (!this._countComputeFisheye) {
const { shader, bindGroupFormat } = createShaderAndFormat(this.pickMode, true);
this._countShaderFisheye = shader;
this._countBindGroupFormatFisheye = bindGroupFormat;
const label = this.pickMode ? 'GSplatPickTileCountFisheye' : 'GSplatLocalTileCountFisheye';
this._countComputeFisheye = new Compute(this.device, shader, label);
}
return this._countComputeFisheye;
}
if (!this._countCompute) {
const { shader, bindGroupFormat } = createShaderAndFormat(this.pickMode, false);
this._countShader = shader;
this._countBindGroupFormat = bindGroupFormat;
const label = this.pickMode ? 'GSplatPickTileCount' : 'GSplatLocalTileCount';
this._countCompute = new Compute(this.device, shader, label);
}
return this._countCompute;
}
destroyCountResources() {
this._countShader?.destroy();
this._countBindGroupFormat?.destroy();
this._countCompute?.destroy();
this._countShaderFisheye?.destroy();
this._countBindGroupFormatFisheye?.destroy();
this._countComputeFisheye?.destroy();
this._countShader = null;
this._countBindGroupFormat = null;
this._countCompute = null;
this._countShaderFisheye = null;
this._countBindGroupFormatFisheye = null;
this._countComputeFisheye = null;
}
getRasterizeCompute(pickMode, depthTest, fogType = 'none', heatmap = false) {
let key = pickMode ? 'pick' : 'color';
if (depthTest) key += '-depth';
if (fogType !== 'none') key += `-fog-${fogType}`;
if (heatmap) key += '-heatmap';
let variant = this._rasterizeVariants.get(key);
if (!variant) {
const { shader, bindGroupFormat } = this._createRasterizeShaderAndFormat(pickMode, depthTest, fogType, heatmap);
const compute = new Compute(this.device, shader, `GSplatRasterize-${key}`);
variant = {
shader,
bindGroupFormat,
compute
};
this._rasterizeVariants.set(key, variant);
}
return variant.compute;
}
_createRasterizeShaderAndFormat(pickMode, depthTest = false, fogType = 'none', heatmap = false) {
const device = this.device;
const hasFog = fogType !== 'none';
const uniforms = [
new UniformFormat('screenWidth', UNIFORMTYPE_UINT),
new UniformFormat('screenHeight', UNIFORMTYPE_UINT),
new UniformFormat('numTilesX', UNIFORMTYPE_UINT),
new UniformFormat('nearClip', UNIFORMTYPE_FLOAT),
new UniformFormat('farClip', UNIFORMTYPE_FLOAT),
new UniformFormat('alphaClip', UNIFORMTYPE_FLOAT)
];
if (hasFog) {
uniforms.push(new UniformFormat('fog_color', UNIFORMTYPE_VEC3), new UniformFormat('fog_start', UNIFORMTYPE_FLOAT), new UniformFormat('fog_end', UNIFORMTYPE_FLOAT), new UniformFormat('fog_density', UNIFORMTYPE_FLOAT));
}
const ubf = new UniformBufferFormat(device, uniforms);
const sharedBindings = [
new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE),
new BindStorageBufferFormat('tileEntries', SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat('tileSplatCounts', SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat('projCache', SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat('rasterizeTileList', SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat('tileListCounts', SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat('depthBuffer', SHADERSTAGE_COMPUTE, true)
];
const outputBindings = pickMode ? [
new BindStorageTextureFormat('pickIdTexture', PIXELFORMAT_R32U),
new BindStorageTextureFormat('pickDepthTexture', PIXELFORMAT_RGBA16F)
] : [
new BindStorageTextureFormat('outputTexture', PIXELFORMAT_RGBA16F)
];
const depthBindings = depthTest ? [
new BindTextureFormat('sceneDepthMap', SHADERSTAGE_COMPUTE, TEXTUREDIMENSION_2D, SAMPLETYPE_UNFILTERABLE_FLOAT, false)
] : [];
const bgf = new BindGroupFormat(device, [
...sharedBindings,
...outputBindings,
...depthBindings
]);
const cdefines = new Map();
if (pickMode) cdefines.set('PICK_MODE', '');
if (depthTest) cdefines.set('DEPTH_TEST', '');
if (heatmap) cdefines.set('HEATMAP_MODE', '');
cdefines.set('GAMMA', 'SRGB');
cdefines.set('FOG', hasFog ? fogType.toUpperCase() : 'NONE');
const cincludes = pickMode ? undefined : new Map([
[
'decodePS',
shaderChunksWGSL.decodePS
]
]);
if (hasFog && cincludes) {
cincludes.set('fogMathPS', shaderChunksWGSL.fogMathPS);
cincludes.set('gammaPS', shaderChunksWGSL.gammaPS);
}
let name = 'GSplatLocalRasterize';
if (pickMode) name = 'GSplatLocalRasterizePick';
else if (depthTest) name = 'GSplatLocalRasterizeDepth';
const shader = new Shader(device, {
name,
shaderLanguage: SHADERLANGUAGE_WGSL,
cshader: computeGsplatLocalRasterizeSource,
cdefines: cdefines.size > 0 ? cdefines : undefined,
cincludes,
computeBindGroupFormat: bgf,
computeUniformBufferFormats: {
uniforms: ubf
}
});
return {
shader,
bindGroupFormat: bgf
};
}
destroy() {
this.destroyCountResources();
for (const { shader, bindGroupFormat, compute } of this._rasterizeVariants.values()){
compute.destroy();
shader.destroy();
bindGroupFormat.destroy();
}
this._tileSplatCountsBuffer?.destroy();
this._smallTileListBuffer?.destroy();
this._largeTileListBuffer?.destroy();
this._largeTileOverflowBasesBuffer?.destroy();
this._rasterizeTileListBuffer?.destroy();
this._tileListCountsBuffer?.destroy();
this._chunkRangesBuffer?.destroy();
this._totalChunksBuffer?.destroy();
this._chunkSortIndirectBuffer?.destroy();
this.prefixSumKernel.destroy();
this.outputTexture?.destroy();
this.pickIdTexture?.destroy();
this.pickDepthTexture?.destroy();
}
constructor(device, pickMode){
this._countShader = null;
this._countBindGroupFormat = null;
this._countCompute = null;
this._countShaderFisheye = null;
this._countBindGroupFormatFisheye = null;
this._countComputeFisheye = null;
this._rasterizeVariants = new Map();
this._tileSplatCountsBuffer = null;
this._smallTileListBuffer = null;
this._largeTileListBuffer = null;
this._largeTileOverflowBasesBuffer = null;
this._rasterizeTileListBuffer = null;
this._tileListCountsBuffer = null;
this._chunkRangesBuffer = null;
this._totalChunksBuffer = null;
this._chunkSortIndirectBuffer = null;
this._allocatedTileCapacity = 0;
this.outputTexture = null;
this.pickIdTexture = null;
this.pickDepthTexture = null;
this.device = device;
this.pickMode = pickMode;
this.prefixSumKernel = new PrefixSumKernel(device);
if (pickMode) {
this.pickIdTexture = new Texture(device, {
name: 'GSplatLocalPickId',
width: 4,
height: 4,
format: PIXELFORMAT_R32U,
mipmaps: false,
minFilter: FILTER_NEAREST,
magFilter: FILTER_NEAREST,
storage: true
});
this.pickDepthTexture = new Texture(device, {
name: 'GSplatLocalPickDepth',
width: 4,
height: 4,
format: PIXELFORMAT_RGBA16F,
mipmaps: false,
minFilter: FILTER_NEAREST,
magFilter: FILTER_NEAREST,
storage: true
});
} else {
this.outputTexture = new Texture(device, {
name: 'GSplatLocalComputeOutput',
width: 4,
height: 4,
format: PIXELFORMAT_RGBA16F,
mipmaps: false,
minFilter: FILTER_NEAREST,
magFilter: FILTER_NEAREST,
storage: true
});
}
}
}
export { GSplatLocalDispatchSet };