playcanvas/build/playcanvas/src/scene/gsplat-unified/gsplat-hybrid-renderer.js

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

import { Mat4 } from "../../core/math/mat4.js";
import { Vec3 } from "../../core/math/vec3.js";
import { SEMANTIC_POSITION, CULLFACE_NONE } from "../../platform/graphics/constants.js";
import {
	BLEND_NONE,
	BLEND_PREMULTIPLIED,
	BLEND_ADDITIVE,
	GSPLAT_FORWARD,
	SHADOWCAMERA_NAME
} from "../constants.js";
import { ShaderMaterial } from "../materials/shader-material.js";
import { GSplatResourceBase } from "../gsplat/gsplat-resource-base.js";
import { MeshInstance } from "../mesh-instance.js";
import { GSplatRenderer } from "./gsplat-renderer.js";
import { GSplatProjector } from "./gsplat-projector.js";
import { GSplatIntervalCompaction } from "./gsplat-interval-compaction.js";
import { ComputeRadixSort } from "../graphics/radix-sort/compute-radix-sort.js";
import { CACHE_STRIDE } from "./gsplat-projector-constants.js";
import { ALPHA_VISIBILITY_THRESHOLD } from "./constants.js";
import { Camera } from "../camera.js";
const _invProjMat = new Mat4();
const _shaderProjMat = new Mat4();
const _camPos = new Vec3();
const _camDir = new Vec3();
const _tmpV = new Vec3();
class GSplatHybridRenderer extends GSplatRenderer {
	_material;
	meshInstance;
	_pickMaterial = null;
	_pickMeshInstance = null;
	_clipToViewZ = new Float32Array(4);
	_clipToViewZPick = null;
	originalBlendType = BLEND_ADDITIVE;
	_internalDefines = /* @__PURE__ */ new Set();
	forceCopyMaterial = true;
	_lastSourceChunksKey = "";
	_cacheStride = CACHE_STRIDE;
	gpuSorter = null;
	projector = null;
	intervalCompaction = null;
	_scratch = null;
	indirectDrawSlot = -1;
	indirectDispatchSlot = -1;
	lastCompactedNumIntervals = 0;
	constructor(device, node, cameraNode, layer, workBuffer, scratch = null) {
		super(device, node, cameraNode, layer, workBuffer);
		this._scratch = scratch;
		this._material = new ShaderMaterial({
			uniqueName: "UnifiedSplatHybridMaterial",
			vertexWGSL: '#include "gsplatHybridVS"',
			fragmentWGSL: '#include "gsplatPS"',
			attributes: {
				vertex_position: SEMANTIC_POSITION
			}
		});
		this._material.setDefine("{GSPLAT_INSTANCE_SIZE}", GSplatResourceBase.instanceSize);
		this._material.setDefine("{CACHE_STRIDE}", CACHE_STRIDE);
		this.configureMaterial();
		this._material.defines.forEach((value, key) => {
			this._internalDefines.add(key);
		});
		this._internalDefines.add("{GSPLAT_INSTANCE_SIZE}");
		this._internalDefines.add("{CACHE_STRIDE}");
		this._internalDefines.add("GSPLAT_UNIFIED_ID");
		this._internalDefines.add("PICK_CUSTOM_ID");
		this._internalDefines.add("GSPLAT_OVERDRAW");
		this._internalDefines.add("GSPLAT_NO_FOG");
		this._internalDefines.add("GSPLAT_XR");
		this.meshInstance = this.createMeshInstance();
	}
	setRenderMode(renderMode) {
		const oldRenderMode = this.renderMode ?? 0;
		const wasForward = (oldRenderMode & GSPLAT_FORWARD) !== 0;
		const isForward = (renderMode & GSPLAT_FORWARD) !== 0;
		if (wasForward && !isForward) {
			this.layer.removeMeshInstances([this.meshInstance], true);
		}
		if (!wasForward && isForward) {
			this.layer.addMeshInstances([this.meshInstance], true);
		}
		super.setRenderMode(renderMode);
	}
	destroy() {
		if (this.renderMode && this.renderMode & GSPLAT_FORWARD) {
			this.layer.removeMeshInstances([this.meshInstance], true);
		}
		this.gpuSorter?.destroy();
		this.gpuSorter = null;
		this.projector?.destroy();
		this.projector = null;
		this.intervalCompaction?.destroy();
		this.intervalCompaction = null;
		this._material.destroy();
		this._pickMaterial?.destroy();
		this.meshInstance.destroy();
		this._pickMeshInstance?.destroy();
		super.destroy();
	}
	get material() {
		return this._material;
	}
	get usesGpuSort() {
		return true;
	}
	get requiresBounds() {
		return true;
	}
	onWorkBufferFormatChanged() {
		this.configureMaterial();
	}
	configureMaterial() {
		this._material.setDefine("SH_BANDS", "0");
		this._material.setDefine("GSPLAT_INDIRECT_DRAW", true);
		this._updateIdDefines(this._material);
		const dither = false;
		this._material.setDefine(`DITHER_${dither ? "BLUENOISE" : "NONE"}`, "");
		this._material.cull = CULLFACE_NONE;
		this._material.blendType = dither ? BLEND_NONE : BLEND_PREMULTIPLIED;
		this._material.depthWrite = !!dither;
		this._material.update();
	}
	setStereo(enabled) {
		if (this._material.getDefine("GSPLAT_XR") !== enabled) {
			this._material.setDefine("GSPLAT_XR", enabled);
			this._material.update();
		}
	}
	update(count, textureSize) {
		if (this.meshInstance.instancingCount <= 0) {
			this.meshInstance.instancingCount = 1;
		}
		this.meshInstance.visible = count > 0;
	}
	invalidateCullUpload() {
		this.intervalCompaction?.invalidateUpload();
	}
	_ensureGpuPipeline() {
		if (!this.gpuSorter) this.gpuSorter = new ComputeRadixSort(this.device, { indirect: true });
		if (!this.projector) this.projector = new GSplatProjector(this.device);
		if (!this.intervalCompaction) this.intervalCompaction = new GSplatIntervalCompaction(this.device, this._scratch);
	}
	prepareRenderView(world, worldState, params) {
		const cameraNode = params.cameraNode;
		const cam = cameraNode.camera;
		const sceneCam = cam.camera;
		const rt = cam.renderTarget;
		const rect = cam.rect;
		const xrView = sceneCam.xrActive ? sceneCam.xrViews[0] ?? null : null;
		const viewportWidth = Math.floor((xrView ? xrView.viewport.z : rt ? rt.width : this.device.width) * rect.z);
		const viewportHeight = Math.floor((xrView ? xrView.viewport.w : rt ? rt.height : this.device.height) * rect.w);
		const xrViewCount = sceneCam.xrActive ? sceneCam.xrViews.length : 0;
		if (xrViewCount > 2) {
		}
		const isStereo = xrViewCount === 2;
		this.setStereo(isStereo);
		const sortedIndices = this.sortAndProjectForCamera(
			world,
			worldState,
			cameraNode,
			viewportWidth,
			viewportHeight,
			Math.max(ALPHA_VISIBILITY_THRESHOLD, params.alphaClipForward),
			false,
			isStereo,
			params
		);
		if (!sortedIndices) return false;
		this.setHybridSortedRendering(
			this.indirectDrawSlot,
			sortedIndices,
			this.projector.projCache,
			this.intervalCompaction.numSplatsBuffer
		);
		return true;
	}
	preparePickingView(world, worldState, pickParams) {
		const pickMode = !!world.workBuffer.format.getStream("pcId");
		const sortedIndices = this.sortAndProjectForCamera(
			world,
			worldState,
			pickParams.cameraNode,
			pickParams.width,
			pickParams.height,
			Math.max(ALPHA_VISIBILITY_THRESHOLD, pickParams.alphaClip),
			pickMode,
			false,
			pickParams
		);
		if (!sortedIndices) return null;
		return this.prepareForPicking(
			this.indirectDrawSlot,
			sortedIndices,
			this.projector.projCache,
			this.intervalCompaction.numSplatsBuffer,
			pickParams.alphaClip,
			pickParams.alphaClipForward,
			pickParams.cameraNode
		);
	}
	sortAndProjectForCamera(world, worldState, cameraNode, viewportWidth, viewportHeight, alphaClip, pickMode, isStereo, params) {
		const elementCount = worldState.totalActiveSplats;
		if (elementCount === 0) return null;
		this._ensureGpuPipeline();
		const gpuSorter = this.gpuSorter;
		const projector = this.projector;
		this.intervalCompaction.uploadIntervals(worldState);
		if (world.hasBounds) {
			const state = world.getState(world.currentVersion);
			if (state) {
				this._runFrustumCulling(world, state, cameraNode, params);
			}
		}
		const fisheyeProj = this.fisheyeProj;
		const numIntervals = worldState.totalIntervals;
		const totalActiveSplats = worldState.totalActiveSplats;
		this.intervalCompaction.dispatchCompact(world.workBuffer.frustumCuller, numIntervals, totalActiveSplats, fisheyeProj.enabled);
		this.allocateAndWriteIntervalIndirectArgs(numIntervals);
		const ic = this.intervalCompaction;
		const compactedSplatIds = ic.compactedSplatIds;
		const numBits = Math.max(10, Math.min(20, Math.round(Math.log2(elementCount / 4))));
		const radixBits = gpuSorter.radixBits;
		const roundedNumBits = Math.ceil(numBits / radixBits) * radixBits;
		const { minDist, maxDist } = this.computeDistanceRange(worldState, cameraNode, params.radialSorting);
		const sortIndirectInfo = gpuSorter.prepareIndirect();
		projector.dispatch({
			workBuffer: world.workBuffer,
			cameraNode,
			compactedSplatIds,
			sortElementCountBuffer: ic.sortElementCountBuffer,
			totalCapacity: elementCount,
			radialSort: params.radialSorting,
			numBits: roundedNumBits,
			minDist,
			maxDist,
			alphaClip,
			minPixelSize: params.minPixelSize * 0.5,
			minContribution: params.minContribution,
			foveationStrength: params.foveationStrength,
			foveationCenter: params.foveationCenter,
			viewportWidth,
			viewportHeight,
			flipY: !!cameraNode.camera.renderTarget?.flipY,
			pickMode,
			fisheyeProj,
			antiAlias: params.antiAlias,
			isStereo,
			material: params.material,
			userCacheWords: params.varyings.words
		});
		projector.writeIndirectArgs(
			this.indirectDrawSlot,
			this.indirectDispatchSlot + 1,
			ic.numSplatsBuffer,
			ic.sortElementCountBuffer,
			sortIndirectInfo
		);
		if (pickMode) {
			this.device.submit();
		}
		return gpuSorter.sortIndirect(
			projector.sortKeys,
			elementCount,
			roundedNumBits,
			this.indirectDispatchSlot + 1,
			ic.sortElementCountBuffer,
			void 0,
			false,
			true
			// destructiveKeys: projector overwrites sortKeys each frame before the sort
		);
	}
	allocateAndWriteIntervalIndirectArgs(numIntervals) {
		const gpuSorter = this.gpuSorter;
		const sortInfo = gpuSorter.prepareIndirect();
		const sortSlotCount = sortInfo[0];
		this.indirectDrawSlot = this.device.getIndirectDrawSlot(1);
		this.indirectDispatchSlot = this.device.getIndirectDispatchSlot(1 + sortSlotCount);
		const ic = this.intervalCompaction;
		ic.writeIndirectArgs(this.indirectDrawSlot, this.indirectDispatchSlot, numIntervals, sortInfo);
		this.lastCompactedNumIntervals = numIntervals;
	}
	_runFrustumCulling(world, worldState, cameraNode, params) {
		world.workBuffer.frustumCuller.updateTransformsData(worldState.boundsGroups);
		const cam = cameraNode.camera;
		const sceneCamera = cam.camera;
		const xrViews = sceneCamera.xrViews;
		if (xrViews?.length) {
			sceneCamera.updateViewTransforms();
			sceneCamera.updateXrFrustum();
			world.workBuffer.frustumCuller.setFrustumPlanes(sceneCamera.frustum);
		} else {
			world.workBuffer.frustumCuller.computeFrustumPlanes(cam.projectionMatrix, cam.viewMatrix);
		}
		const fp = this.fisheyeProj;
		fp.update(this.resolveFisheye(params.fisheye), cam.fov, cam.projectionMatrix);
		if (fp.enabled) {
			world.workBuffer.frustumCuller.setFisheyeData(
				cameraNode.getPosition(),
				cameraNode.forward,
				fp.maxTheta
			);
		}
	}
	computeDistanceRange(worldState, cameraNode, radialSort) {
		const cameraMat = cameraNode.getWorldTransform();
		cameraMat.getTranslation(_camPos);
		cameraMat.getZ(_camDir).normalize();
		let minDist = radialSort ? 0 : Infinity;
		let maxDist = radialSort ? 0 : -Infinity;
		for (const splat of worldState.splats) {
			const modelMat = splat.node.getWorldTransform();
			const aabbMin = splat.aabb.getMin();
			const aabbMax = splat.aabb.getMax();
			for (let i = 0; i < 8; i++) {
				_tmpV.x = i & 1 ? aabbMax.x : aabbMin.x;
				_tmpV.y = i & 2 ? aabbMax.y : aabbMin.y;
				_tmpV.z = i & 4 ? aabbMax.z : aabbMin.z;
				modelMat.transformPoint(_tmpV, _tmpV);
				if (radialSort) {
					const dist = _tmpV.distance(_camPos);
					if (dist > maxDist) maxDist = dist;
				} else {
					const dist = _tmpV.sub(_camPos).dot(_camDir);
					if (dist < minDist) minDist = dist;
					if (dist > maxDist) maxDist = dist;
				}
			}
		}
		if (maxDist === 0 || maxDist === -Infinity) {
			return { minDist: 0, maxDist: 1 };
		}
		return { minDist, maxDist };
	}
	setHybridSortedRendering(drawSlot, sortedIndices, projCache, numSplatsBuffer) {
		this.meshInstance.setIndirect(null, drawSlot, 1);
		this._material.setParameter("sortedIndices", sortedIndices);
		this._material.setParameter("projCache", projCache);
		this._material.setParameter("numSplatsStorage", numSplatsBuffer);
		this._computeClipToViewZ(this.cameraNode, this._clipToViewZ);
		this._material.setParameter("clipToViewZ", this._clipToViewZ);
		this.meshInstance.visible = true;
		if (this.meshInstance.instancingCount <= 0) {
			this.meshInstance.instancingCount = 1;
		}
	}
	prepareForPicking(drawSlot, sortedIndices, projCache, numSplatsBuffer, alphaClip, alphaClipForward, cameraNode) {
		if (!this._pickMaterial) {
			this._pickMaterial = new ShaderMaterial({
				uniqueName: "UnifiedSplatHybridPickMaterial",
				vertexWGSL: '#include "gsplatHybridVS"',
				fragmentWGSL: '#include "gsplatPS"',
				attributes: {
					vertex_position: SEMANTIC_POSITION
				}
			});
			this._pickMaterial.setDefine("{GSPLAT_INSTANCE_SIZE}", GSplatResourceBase.instanceSize);
			this._pickMaterial.setDefine("{CACHE_STRIDE}", this._cacheStride);
			this._pickMaterial.setDefine("SH_BANDS", "0");
			this._pickMaterial.setDefine("GSPLAT_INDIRECT_DRAW", true);
			this._pickMaterial.setDefine("DITHER_NONE", "");
			this._updateIdDefines(this._pickMaterial);
			this._pickMaterial.cull = CULLFACE_NONE;
			this._pickMaterial.blendType = BLEND_NONE;
			this._pickMaterial.depthWrite = false;
			this._pickMaterial.update();
			const mesh = GSplatResourceBase.createMesh(this.device);
			this._pickMeshInstance = new MeshInstance(mesh, this._pickMaterial);
			this._pickMeshInstance.node = this.node;
			this._pickMeshInstance.setInstancing(true, true);
			this._pickMeshInstance.instancingCount = 1;
		} else {
			if (this._updateIdDefines(this._pickMaterial)) {
				this._pickMaterial.update();
			}
		}
		const pickMaterial = this._pickMaterial;
		const pickMeshInstance = this._pickMeshInstance;
		pickMeshInstance.setIndirect(null, drawSlot, 1);
		pickMaterial.setParameter("sortedIndices", sortedIndices);
		pickMaterial.setParameter("projCache", projCache);
		pickMaterial.setParameter("numSplatsStorage", numSplatsBuffer);
		pickMaterial.setParameter("alphaClip", alphaClip);
		pickMaterial.setParameter("alphaClipForward", alphaClipForward);
		this._clipToViewZPick ?? (this._clipToViewZPick = new Float32Array(4));
		this._computeClipToViewZ(cameraNode, this._clipToViewZPick);
		pickMaterial.setParameter("clipToViewZ", this._clipToViewZPick);
		return pickMeshInstance;
	}
	_computeClipToViewZ(cameraNode, dst) {
		const camComp = cameraNode.camera;
		const cam = camComp.camera;
		if (this.fisheyeProj.enabled) {
			const near = cam.nearClip;
			const far = cam.farClip;
			dst[0] = 0;
			dst[1] = 0;
			dst[2] = far - near;
			dst[3] = near;
			return;
		}
		const flipY = !!camComp.renderTarget?.flipY;
		_invProjMat.copy(Camera.applyShaderProjectionTransform(cam.projectionMatrix, _shaderProjMat, flipY, this.device.isWebGPU)).invert();
		const d = _invProjMat.data;
		dst[0] = -d[2];
		dst[1] = -d[6];
		dst[2] = -d[10];
		dst[3] = -d[14];
	}
	setCpuSortedRendering() {
		this.meshInstance.setIndirect(null, -1);
		this.meshInstance.visible = false;
	}
	setOrderData() {
	}
	frameUpdate(params) {
		this._material.setParameter("alphaClip", params.alphaClip);
		this._material.setParameter("alphaClipForward", params.alphaClipForward);
		this._pickMaterial?.setParameter("alphaClip", params.alphaClip);
		this._pickMaterial?.setParameter("alphaClipForward", params.alphaClipForward);
		if (params.colorRamp) {
			this._material.setParameter("colorRampIntensity", params.colorRampIntensity);
		}
		const noFog = !params.useFog;
		if (noFog !== this._lastNoFog) {
			this._lastNoFog = noFog;
			this._material.setDefine("GSPLAT_NO_FOG", noFog);
			this._material.update();
		}
		const cacheStride = CACHE_STRIDE + params.varyings.words;
		if (cacheStride !== this._cacheStride) {
			this._cacheStride = cacheStride;
			this._material.setDefine("{CACHE_STRIDE}", cacheStride);
			this._material.update();
			if (this._pickMaterial) {
				this._pickMaterial.setDefine("{CACHE_STRIDE}", cacheStride);
				this._pickMaterial.update();
			}
		}
		if (this.forceCopyMaterial || params.material.dirty) {
			this.copyMaterialSettings(params.material);
			this.forceCopyMaterial = false;
		}
	}
	copyMaterialSettings(sourceMaterial) {
		const keysToDelete = [];
		this._material.defines.forEach((value, key) => {
			if (!this._internalDefines.has(key) && !sourceMaterial.defines.has(key)) {
				keysToDelete.push(key);
			}
		});
		keysToDelete.forEach((key) => this._material.setDefine(key, void 0));
		sourceMaterial.defines.forEach((value, key) => {
			this._material.setDefine(key, value);
		});
		const srcParams = sourceMaterial.parameters;
		for (const paramName in srcParams) {
			if (srcParams.hasOwnProperty(paramName)) {
				this._material.setParameter(paramName, srcParams[paramName].data);
			}
		}
		if (sourceMaterial.hasShaderChunks) {
			const sourceChunksKey = sourceMaterial.shaderChunks.key;
			if (sourceChunksKey !== this._lastSourceChunksKey) {
				this._material.shaderChunks.copy(sourceMaterial.shaderChunks);
				this._lastSourceChunksKey = sourceChunksKey;
			}
		}
		this._material.update();
	}
	_updateIdDefines(material) {
		const hasPcId = !!this.workBuffer.format.getStream("pcId");
		const changed = material.getDefine("GSPLAT_UNIFIED_ID") !== hasPcId || material.getDefine("PICK_CUSTOM_ID") !== hasPcId;
		material.setDefine("GSPLAT_UNIFIED_ID", hasPcId);
		material.setDefine("PICK_CUSTOM_ID", hasPcId);
		return changed;
	}
	updateOverdrawMode(params) {
		const overdrawEnabled = !!params.colorRamp;
		const wasOverdrawEnabled = this._material.getDefine("GSPLAT_OVERDRAW");
		if (overdrawEnabled) {
			this._material.setParameter("colorRamp", params.colorRamp);
			this._material.setParameter("colorRampIntensity", params.colorRampIntensity);
		}
		if (overdrawEnabled !== wasOverdrawEnabled) {
			this._material.setDefine("GSPLAT_OVERDRAW", overdrawEnabled);
			if (overdrawEnabled) {
				this.originalBlendType = this._material.blendType;
				this._material.blendType = BLEND_ADDITIVE;
			} else {
				this._material.blendType = this.originalBlendType;
			}
			this._material.update();
		}
	}
	createMeshInstance() {
		const mesh = GSplatResourceBase.createMesh(this.device);
		const meshInstance = new MeshInstance(mesh, this._material);
		meshInstance.node = this.node;
		meshInstance.setInstancing(true, true);
		meshInstance.instancingCount = 0;
		meshInstance.pick = false;
		const thisCamera = this.cameraNode.camera;
		meshInstance.isVisibleFunc = (camera) => {
			const renderMode = this.renderMode ?? 0;
			if (thisCamera.camera === camera && renderMode & GSPLAT_FORWARD) {
				return true;
			}
			if (camera.node?.name === SHADOWCAMERA_NAME) {
				return false;
			}
			return false;
		};
		return meshInstance;
	}
}
export {
	GSplatHybridRenderer
};