playcanvas/build/playcanvas.prf/src/scene/gsplat-unified/gsplat-work-buffer.js

PlayCanvas WebGL game engine

import { SEMANTIC_POSITION, PIXELFORMAT_RGBA16U, PIXELFORMAT_RGBA16F, PIXELFORMAT_RGBA32U, PIXELFORMAT_RG32U, BUFFERUSAGE_COPY_DST, PIXELFORMAT_R32U, ADDRESS_CLAMP_TO_EDGE, FILTER_NEAREST } from '../../platform/graphics/constants.js';
import { RenderTarget } from '../../platform/graphics/render-target.js';
import { StorageBuffer } from '../../platform/graphics/storage-buffer.js';
import { Texture } from '../../platform/graphics/texture.js';
import { UploadStream } from '../../platform/graphics/upload-stream.js';
import { QuadRender } from '../graphics/quad-render.js';
import { ShaderUtils } from '../shader-lib/shader-utils.js';
import glslGsplatCopyToWorkBufferPS from '../shader-lib/glsl/chunks/gsplat/frag/gsplatCopyToWorkbuffer.js';
import wgslGsplatCopyToWorkBufferPS from '../shader-lib/wgsl/chunks/gsplat/frag/gsplatCopyToWorkbuffer.js';
import { GSplatWorkBufferRenderPass } from './gsplat-work-buffer-render-pass.js';

let id = 0;
class WorkBufferRenderInfo {
		constructor(device, key, material, colorTextureFormat, colorOnly){
				this.device = device;
				this.material = material;
				const clonedDefines = new Map(material.defines);
				const isColorUint = colorTextureFormat === PIXELFORMAT_RGBA16U;
				const colorOutputType = isColorUint ? 'uvec4' : 'vec4';
				if (isColorUint) {
						clonedDefines.set('GSPLAT_COLOR_UINT', '');
				}
				if (colorOnly) {
						clonedDefines.set('GSPLAT_COLOR_ONLY', '');
				}
				const shader = ShaderUtils.createShader(this.device, {
						uniqueName: `SplatCopyToWorkBuffer:${key}`,
						attributes: {
								vertex_position: SEMANTIC_POSITION
						},
						vertexDefines: clonedDefines,
						fragmentDefines: clonedDefines,
						vertexChunk: 'fullscreenQuadVS',
						fragmentGLSL: glslGsplatCopyToWorkBufferPS,
						fragmentWGSL: wgslGsplatCopyToWorkBufferPS,
						fragmentOutputTypes: colorOnly ? [
								colorOutputType
						] : [
								colorOutputType,
								'uvec4',
								'uvec2'
						]
				});
				this.quadRender = new QuadRender(shader);
		}
		destroy() {
				this.material?.destroy();
				this.quadRender?.destroy();
		}
}
class GSplatWorkBuffer {
		constructor(device){
				this.id = id++;
				this._textureSize = 1;
				this.device = device;
				this.colorTextureFormat = device.getRenderableHdrFormat([
						PIXELFORMAT_RGBA16F
				]) || PIXELFORMAT_RGBA16U;
				this.colorTexture = this.createTexture('splatColor', this.colorTextureFormat, 1, 1);
				this.splatTexture0 = this.createTexture('splatTexture0', PIXELFORMAT_RGBA32U, 1, 1);
				this.splatTexture1 = this.createTexture('splatTexture1', PIXELFORMAT_RG32U, 1, 1);
				this.renderTarget = new RenderTarget({
						name: `GsplatWorkBuffer-MRT-${this.id}`,
						colorBuffers: [
								this.colorTexture,
								this.splatTexture0,
								this.splatTexture1
						],
						depth: false,
						flipY: true
				});
				this.colorRenderTarget = new RenderTarget({
						name: `GsplatWorkBuffer-Color-${this.id}`,
						colorBuffer: this.colorTexture,
						depth: false,
						flipY: true
				});
				this.uploadStream = new UploadStream(device);
				if (device.isWebGPU) {
						this.orderBuffer = new StorageBuffer(device, 4, BUFFERUSAGE_COPY_DST);
				} else {
						this.orderTexture = this.createTexture('SplatGlobalOrder', PIXELFORMAT_R32U, 1, 1);
				}
				this.renderPass = new GSplatWorkBufferRenderPass(device, this);
				this.renderPass.init(this.renderTarget);
				this.colorRenderPass = new GSplatWorkBufferRenderPass(device, this, true);
				this.colorRenderPass.init(this.colorRenderTarget);
		}
		destroy() {
				this.renderPass?.destroy();
				this.colorRenderPass?.destroy();
				this.colorTexture?.destroy();
				this.splatTexture0?.destroy();
				this.splatTexture1?.destroy();
				this.orderTexture?.destroy();
				this.orderBuffer?.destroy();
				this.renderTarget?.destroy();
				this.colorRenderTarget?.destroy();
				this.uploadStream.destroy();
		}
		get textureSize() {
				return this._textureSize;
		}
		setOrderData(data) {
				if (this.device.isWebGPU) {
						this.uploadStream.upload(data, this.orderBuffer, 0, data.length);
				} else {
						this.uploadStream.upload(data, this.orderTexture, 0, data.length);
				}
		}
		createTexture(name, format, w, h) {
				return new Texture(this.device, {
						name: name,
						width: w,
						height: h,
						format: format,
						cubemap: false,
						mipmaps: false,
						minFilter: FILTER_NEAREST,
						magFilter: FILTER_NEAREST,
						addressU: ADDRESS_CLAMP_TO_EDGE,
						addressV: ADDRESS_CLAMP_TO_EDGE
				});
		}
		resize(textureSize) {
				this.renderTarget.resize(textureSize, textureSize);
				this.colorRenderTarget.resize(textureSize, textureSize);
				this._textureSize = textureSize;
				if (this.device.isWebGPU) {
						const newByteSize = textureSize * textureSize * 4;
						if (this.orderBuffer.byteSize < newByteSize) {
								this.orderBuffer.destroy();
								this.orderBuffer = new StorageBuffer(this.device, newByteSize, BUFFERUSAGE_COPY_DST);
						}
				} else {
						this.orderTexture.resize(textureSize, textureSize);
				}
		}
		render(splats, cameraNode, colorsByLod) {
				if (this.renderPass.update(splats, cameraNode, colorsByLod)) {
						this.renderPass.render();
				}
		}
		renderColor(splats, cameraNode, colorsByLod) {
				if (this.colorRenderPass.update(splats, cameraNode, colorsByLod)) {
						this.colorRenderPass.render();
				}
		}
}

export { GSplatWorkBuffer, WorkBufferRenderInfo };