playcanvas/build/playcanvas/src/scene/gsplat-unified/gsplat-sort-key-compute.js

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

import { Vec2 } from '../../core/math/vec2.js';
import { Vec3 } from '../../core/math/vec3.js';
import { Compute } from '../../platform/graphics/compute.js';
import { Shader } from '../../platform/graphics/shader.js';
import { StorageBuffer } from '../../platform/graphics/storage-buffer.js';
import { BindGroupFormat, BindTextureFormat, BindStorageBufferFormat, BindUniformBufferFormat } from '../../platform/graphics/bind-group-format.js';
import { UniformBufferFormat, UniformFormat } from '../../platform/graphics/uniform-buffer-format.js';
import { SHADERLANGUAGE_WGSL, UNIFORMTYPE_VEC3, UNIFORMTYPE_UINT, UNIFORMTYPE_FLOAT, SHADERSTAGE_COMPUTE, SAMPLETYPE_UINT, BUFFERUSAGE_COPY_SRC, BUFFERUSAGE_COPY_DST } from '../../platform/graphics/constants.js';
import { computeGsplatSortKeySource } from '../shader-lib/wgsl/chunks/gsplat/compute-gsplat-sort-key.js';
import { GSplatSortBinWeights } from './gsplat-sort-bin-weights.js';

const WORKGROUP_SIZE_X = 16;
const WORKGROUP_SIZE_Y = 16;
const THREADS_PER_WORKGROUP = WORKGROUP_SIZE_X * WORKGROUP_SIZE_Y;
const _cameraDir = new Vec3();
const _dispatchSize = new Vec2();
class GSplatSortKeyCompute {
		destroy() {
				this.keysBuffer?.destroy();
				this.binWeightsBuffer?.destroy();
				this.compute?.shader?.destroy();
				this.bindGroupFormat?.destroy();
				this.bindGroupFormatIndirect?.destroy();
				this.keysBuffer = null;
				this.binWeightsBuffer = null;
				this.compute = null;
				this.bindGroupFormat = null;
				this.bindGroupFormatIndirect = null;
				this.uniformBufferFormat = null;
		}
		_getCompute(computeRadialSort, computeUseIndirectSort = false) {
				if (!this.compute || this.computeRadialSort !== computeRadialSort || this.computeUseIndirectSort !== computeUseIndirectSort) {
						this.compute?.shader?.destroy();
						const modeName = computeRadialSort ? 'Radial' : 'Linear';
						const name = `GSplatSortKeyCompute-${modeName}${computeUseIndirectSort ? '-Indirect' : ''}`;
						const cdefines = new Map([
								[
										'{WORKGROUP_SIZE_X}',
										`${WORKGROUP_SIZE_X}`
								],
								[
										'{WORKGROUP_SIZE_Y}',
										`${WORKGROUP_SIZE_Y}`
								]
						]);
						if (computeRadialSort) {
								cdefines.set('RADIAL_SORT', '');
						}
						if (computeUseIndirectSort) {
								cdefines.set('USE_INDIRECT_SORT', '');
						}
						const bgFormat = computeUseIndirectSort ? this.bindGroupFormatIndirect : this.bindGroupFormat;
						const shader = new Shader(this.device, {
								name: name,
								shaderLanguage: SHADERLANGUAGE_WGSL,
								cshader: computeGsplatSortKeySource,
								cdefines: cdefines,
								computeBindGroupFormat: bgFormat,
								computeUniformBufferFormats: {
										uniforms: this.uniformBufferFormat
								}
						});
						this.compute = new Compute(this.device, shader, name);
						this.computeRadialSort = computeRadialSort;
						this.computeUseIndirectSort = computeUseIndirectSort;
				}
				return this.compute;
		}
		_createBindGroupFormat() {
				const device = this.device;
				this.uniformBufferFormat = new UniformBufferFormat(device, [
						new UniformFormat('cameraPosition', UNIFORMTYPE_VEC3),
						new UniformFormat('elementCount', UNIFORMTYPE_UINT),
						new UniformFormat('cameraDirection', UNIFORMTYPE_VEC3),
						new UniformFormat('numBits', UNIFORMTYPE_UINT),
						new UniformFormat('textureSize', UNIFORMTYPE_UINT),
						new UniformFormat('minDist', UNIFORMTYPE_FLOAT),
						new UniformFormat('invRange', UNIFORMTYPE_FLOAT),
						new UniformFormat('numWorkgroupsX', UNIFORMTYPE_UINT),
						new UniformFormat('numBins', UNIFORMTYPE_UINT)
				]);
				this.bindGroupFormat = new BindGroupFormat(device, [
						new BindTextureFormat('dataTransformA', SHADERSTAGE_COMPUTE, undefined, SAMPLETYPE_UINT, false),
						new BindStorageBufferFormat('sortKeys', SHADERSTAGE_COMPUTE, false),
						new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE),
						new BindStorageBufferFormat('binWeights', SHADERSTAGE_COMPUTE, true)
				]);
				this.bindGroupFormatIndirect = new BindGroupFormat(device, [
						new BindTextureFormat('dataTransformA', SHADERSTAGE_COMPUTE, undefined, SAMPLETYPE_UINT, false),
						new BindStorageBufferFormat('sortKeys', SHADERSTAGE_COMPUTE, false),
						new BindUniformBufferFormat('uniforms', SHADERSTAGE_COMPUTE),
						new BindStorageBufferFormat('binWeights', SHADERSTAGE_COMPUTE, true),
						new BindStorageBufferFormat('compactedSplatIds', SHADERSTAGE_COMPUTE, true),
						new BindStorageBufferFormat('sortElementCountBuf', SHADERSTAGE_COMPUTE, true)
				]);
		}
		_ensureCapacity(elementCount) {
				if (elementCount > this.allocatedCount) {
						this.keysBuffer?.destroy();
						this.allocatedCount = elementCount;
						this.keysBuffer = new StorageBuffer(this.device, elementCount * 4, BUFFERUSAGE_COPY_SRC);
				}
		}
		generate(workBuffer, cameraNode, computeRadialSort, elementCount, numBits, minDist, maxDist) {
				this._ensureCapacity(elementCount);
				const workgroupCount = Math.ceil(elementCount / THREADS_PER_WORKGROUP);
				Compute.calcDispatchSize(workgroupCount, _dispatchSize, this.device.limits.maxComputeWorkgroupsPerDimension || 65535);
				const compute = this._getCompute(computeRadialSort);
				const cameraPos = cameraNode.getPosition();
				const cameraMat = cameraNode.getWorldTransform();
				const cameraDir = cameraMat.getZ(_cameraDir).normalize();
				const range = maxDist - minDist;
				const invRange = range > 0 ? 1.0 / range : 1.0;
				const bucketCount = 1 << numBits;
				const cameraBin = GSplatSortBinWeights.computeCameraBin(computeRadialSort, minDist, range);
				const binWeights = this.binWeightsUtil.compute(cameraBin, bucketCount);
				this.binWeightsBuffer.write(0, binWeights);
				compute.setParameter('dataTransformA', workBuffer.getTexture('dataTransformA'));
				compute.setParameter('sortKeys', this.keysBuffer);
				compute.setParameter('binWeights', this.binWeightsBuffer);
				this.cameraPositionData[0] = cameraPos.x;
				this.cameraPositionData[1] = cameraPos.y;
				this.cameraPositionData[2] = cameraPos.z;
				compute.setParameter('cameraPosition', this.cameraPositionData);
				this.cameraDirectionData[0] = cameraDir.x;
				this.cameraDirectionData[1] = cameraDir.y;
				this.cameraDirectionData[2] = cameraDir.z;
				compute.setParameter('cameraDirection', this.cameraDirectionData);
				compute.setParameter('elementCount', elementCount);
				compute.setParameter('numBits', numBits);
				compute.setParameter('textureSize', workBuffer.textureSize);
				compute.setParameter('minDist', minDist);
				compute.setParameter('invRange', invRange);
				compute.setParameter('numWorkgroupsX', _dispatchSize.x);
				compute.setParameter('numBins', GSplatSortBinWeights.NUM_BINS);
				compute.setupDispatch(_dispatchSize.x, _dispatchSize.y, 1);
				this.device.computeDispatch([
						compute
				], 'GSplatSortKeyCompute');
				return this.keysBuffer;
		}
		generateIndirect(workBuffer, cameraNode, computeRadialSort, maxElementCount, numBits, minDist, maxDist, compactedSplatIds, sortElementCountBuffer, dispatchSlot) {
				this._ensureCapacity(maxElementCount);
				const compute = this._getCompute(computeRadialSort, true);
				const cameraPos = cameraNode.getPosition();
				const cameraMat = cameraNode.getWorldTransform();
				const cameraDir = cameraMat.getZ(_cameraDir).normalize();
				const range = maxDist - minDist;
				const invRange = range > 0 ? 1.0 / range : 1.0;
				const bucketCount = 1 << numBits;
				const cameraBin = GSplatSortBinWeights.computeCameraBin(computeRadialSort, minDist, range);
				const binWeights = this.binWeightsUtil.compute(cameraBin, bucketCount);
				this.binWeightsBuffer.write(0, binWeights);
				compute.setParameter('dataTransformA', workBuffer.getTexture('dataTransformA'));
				compute.setParameter('sortKeys', this.keysBuffer);
				compute.setParameter('binWeights', this.binWeightsBuffer);
				compute.setParameter('compactedSplatIds', compactedSplatIds);
				compute.setParameter('sortElementCountBuf', sortElementCountBuffer);
				this.cameraPositionData[0] = cameraPos.x;
				this.cameraPositionData[1] = cameraPos.y;
				this.cameraPositionData[2] = cameraPos.z;
				compute.setParameter('cameraPosition', this.cameraPositionData);
				this.cameraDirectionData[0] = cameraDir.x;
				this.cameraDirectionData[1] = cameraDir.y;
				this.cameraDirectionData[2] = cameraDir.z;
				compute.setParameter('cameraDirection', this.cameraDirectionData);
				compute.setParameter('elementCount', maxElementCount);
				compute.setParameter('numBits', numBits);
				compute.setParameter('textureSize', workBuffer.textureSize);
				compute.setParameter('minDist', minDist);
				compute.setParameter('invRange', invRange);
				const workgroupCount = Math.ceil(maxElementCount / THREADS_PER_WORKGROUP);
				Compute.calcDispatchSize(workgroupCount, _dispatchSize, this.device.limits.maxComputeWorkgroupsPerDimension || 65535);
				compute.setParameter('numWorkgroupsX', _dispatchSize.x);
				compute.setParameter('numBins', GSplatSortBinWeights.NUM_BINS);
				compute.setupIndirectDispatch(dispatchSlot);
				this.device.computeDispatch([
						compute
				], 'GSplatSortKeyCompute-Indirect');
				return this.keysBuffer;
		}
		constructor(device){
				this.allocatedCount = 0;
				this.keysBuffer = null;
				this.binWeightsBuffer = null;
				this.compute = null;
				this.computeRadialSort = false;
				this.computeUseIndirectSort = false;
				this.bindGroupFormat = null;
				this.bindGroupFormatIndirect = null;
				this.uniformBufferFormat = null;
				this.cameraPositionData = new Float32Array(3);
				this.cameraDirectionData = new Float32Array(3);
				this.device = device;
				this.binWeightsUtil = new GSplatSortBinWeights();
				this.binWeightsBuffer = new StorageBuffer(device, GSplatSortBinWeights.NUM_BINS * 2 * 4, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST);
				this._createBindGroupFormat();
		}
}

export { GSplatSortKeyCompute };