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

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

import { Vec2 } from "../../core/math/vec2.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, BindStorageBufferFormat, BindUniformBufferFormat } from "../../platform/graphics/bind-group-format.js";
import { UniformBufferFormat, UniformFormat } from "../../platform/graphics/uniform-buffer-format.js";
import {
	BUFFERUSAGE_COPY_DST,
	CULLFACE_NONE,
	PIXELFORMAT_RGBA16U,
	SEMANTIC_POSITION,
	SHADERLANGUAGE_WGSL,
	SHADERSTAGE_COMPUTE,
	UNIFORMTYPE_FLOAT,
	UNIFORMTYPE_UINT,
	UNIFORMTYPE_VEC4
} from "../../platform/graphics/constants.js";
import { BLEND_PREMULTIPLIED, LIGHTTYPE_DIRECTIONAL } from "../constants.js";
import { ShaderMaterial } from "../materials/shader-material.js";
import { MeshInstance } from "../mesh-instance.js";
import { GSplatResourceBase } from "../gsplat/gsplat-resource-base.js";
import { computeGsplatShadowCullSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-shadow-cull.js";
import { computeGsplatShadowIndirectArgsSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-shadow-indirect-args.js";
import computeSplatSource from "../shader-lib/wgsl/chunks/gsplat/vert/gsplatComputeSplat.js";
import gsplatModifyDefaultSource from "../shader-lib/wgsl/chunks/gsplat/vert/gsplatModify.js";
import gsplatHelpersSource from "../shader-lib/wgsl/chunks/gsplat/vert/gsplatHelpers.js";
import { GSplatIntervalCompaction } from "./gsplat-interval-compaction.js";
const WORKGROUP_SIZE = 256;
const INDEX_COUNT = 6 * GSplatResourceBase.instanceSize;
class GSplatShadowRenderer {
	device;
	node;
	cameraNode;
	layer;
	world;
	entries = /* @__PURE__ */ new Map();
	_desiredLights = /* @__PURE__ */ new Set();
	_compaction = null;
	_cullDispatchSize = new Vec2(1, 1);
	_frustumPlanes = new Float32Array(24);
	_userChunksKey = "";
	_userModifyWgsl = null;
	// The cull/args shaders are shared, but each light entry gets its OWN Compute instances
	// (created in _createEntry). A Compute owns a persistent uniform buffer, so a single shared
	// Compute dispatched once per light per frame would have all dispatches read the last-written
	// uniforms (frustum planes / draw slot) — making all but one light's shadow draw empty.
	_cullShader = null;
	_cullBindGroupFormat = null;
	_cullFormatVersion = -1;
	_cullBuiltChunksKey = null;
	_cullShaderGen = 0;
	_argsShader = null;
	_argsBindGroupFormat = null;
	constructor(device, node, cameraNode, layer, world, scratch = null) {
		this.device = device;
		this.node = node;
		this.cameraNode = cameraNode;
		this.layer = layer;
		this.world = world;
		this._compaction = new GSplatIntervalCompaction(device, scratch);
		this._createArgsShader();
	}
	destroy() {
		this.entries.forEach((entry) => this._destroyEntry(entry));
		this.entries.clear();
		this._compaction?.destroy();
		this._compaction = null;
		this._cullShader?.destroy();
		this._cullBindGroupFormat?.destroy();
		this._argsShader?.destroy();
		this._argsBindGroupFormat?.destroy();
		this._cullShader = null;
		this._argsShader = null;
	}
	_ensureCullShader() {
		const wbFormat = this.world.workBuffer.format;
		const version = wbFormat.extraStreamsVersion;
		if (!this._cullShader || version !== this._cullFormatVersion || this._userChunksKey !== this._cullBuiltChunksKey) {
			this._cullFormatVersion = version;
			this._cullBuiltChunksKey = this._userChunksKey;
			this._buildCullShader();
		}
	}
	_buildCullShader() {
		const device = this.device;
		const wbFormat = this.world.workBuffer.format;
		const fixedBindings = [
			new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE),
			// pass-1 outputs (read): the shared candidate list + its count at [numIntervals]
			new BindStorageBufferFormat("compactedSplatIds", SHADERSTAGE_COMPUTE, true),
			new BindStorageBufferFormat("candidateCountBuffer", SHADERSTAGE_COMPUTE, true),
			// per-light outputs (read/write): the final visible list + atomic count
			new BindStorageBufferFormat("outputIndices", SHADERSTAGE_COMPUTE, false),
			new BindStorageBufferFormat("globalCount", SHADERSTAGE_COMPUTE, false)
		];
		this._cullBindGroupFormat?.destroy();
		this._cullBindGroupFormat = new BindGroupFormat(device, [
			...fixedBindings,
			...wbFormat.getComputeBindFormats()
		]);
		const uniformBufferFormat = new UniformBufferFormat(device, [
			new UniformFormat("frustumPlanes", UNIFORMTYPE_VEC4, 6),
			new UniformFormat("numIntervals", UNIFORMTYPE_UINT),
			new UniformFormat("splatTextureSize", UNIFORMTYPE_UINT),
			new UniformFormat("alphaClip", UNIFORMTYPE_FLOAT),
			new UniformFormat("worldSizeThreshold", UNIFORMTYPE_FLOAT)
		]);
		const cincludes = /* @__PURE__ */ new Map();
		cincludes.set("gsplatComputeSplatCS", computeSplatSource);
		cincludes.set("gsplatFormatDeclCS", wbFormat.getComputeInputDeclarations(fixedBindings.length));
		cincludes.set("gsplatFormatReadCS", wbFormat.getReadCode());
		cincludes.set("gsplatHelpersVS", gsplatHelpersSource);
		cincludes.set("gsplatModifyVS", this._userModifyWgsl ?? gsplatModifyDefaultSource);
		const cdefines = /* @__PURE__ */ new Map([["{WORKGROUP_SIZE}", WORKGROUP_SIZE.toString()]]);
		const colorStream = wbFormat.getStream("dataColor");
		if (colorStream && colorStream.format !== PIXELFORMAT_RGBA16U) {
			cdefines.set("GSPLAT_COLOR_FLOAT", "");
		}
		this._cullShader?.destroy();
		this._cullShader = new Shader(device, {
			name: "GSplatShadowCull",
			shaderLanguage: SHADERLANGUAGE_WGSL,
			cshader: computeGsplatShadowCullSource,
			cincludes,
			cdefines,
			computeBindGroupFormat: this._cullBindGroupFormat,
			computeUniformBufferFormats: { uniforms: uniformBufferFormat }
		});
		this._cullShaderGen++;
	}
	_createArgsShader() {
		const device = this.device;
		this._argsBindGroupFormat = new BindGroupFormat(device, [
			new BindStorageBufferFormat("countBuffer", SHADERSTAGE_COMPUTE, true),
			new BindStorageBufferFormat("indirectDrawArgs", SHADERSTAGE_COMPUTE, false),
			new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE)
		]);
		const uniformBufferFormat = new UniformBufferFormat(device, [
			new UniformFormat("drawSlot", UNIFORMTYPE_UINT),
			new UniformFormat("indexCount", UNIFORMTYPE_UINT),
			new UniformFormat("pad0", UNIFORMTYPE_UINT),
			new UniformFormat("pad1", UNIFORMTYPE_UINT)
		]);
		this._argsShader = new Shader(device, {
			name: "GSplatShadowIndirectArgs",
			shaderLanguage: SHADERLANGUAGE_WGSL,
			cshader: computeGsplatShadowIndirectArgsSource,
			cdefines: /* @__PURE__ */ new Map([["{INSTANCE_SIZE}", GSplatResourceBase.instanceSize.toString()]]),
			computeBindGroupFormat: this._argsBindGroupFormat,
			computeUniformBufferFormats: { uniforms: uniformBufferFormat }
		});
	}
	setDataSource(workBuffer) {
		this._compaction?.invalidateUpload();
		this._cullFormatVersion = -1;
		this.entries.forEach((entry) => {
			this._configureMaterialWorkBuffer(entry.material);
			entry.material.update();
		});
	}
	setCastersAabb(aabb) {
		if (!aabb) return;
		this.entries.forEach((entry) => {
			entry.meshInstance.setCustomAabb(aabb);
		});
	}
	syncLights() {
		const lights = this.layer.splitLights[LIGHTTYPE_DIRECTIONAL];
		const desired = this._desiredLights;
		desired.clear();
		for (let i = 0; i < lights.length; i++) {
			const light = lights[i];
			if (!light.enabled || !light.castShadows) continue;
			if (light.numCascades !== 1) {
				continue;
			}
			desired.add(light);
		}
		this.entries.forEach((entry, light) => {
			if (!desired.has(light)) {
				this._destroyEntry(entry);
				this.entries.delete(light);
			}
		});
		desired.forEach((light) => {
			if (!this.entries.has(light)) {
				this.entries.set(light, this._createEntry(light));
			}
		});
	}
	cull(gsplatParams) {
		const worldState = this.world.getState(this.world.currentVersion);
		const ready = worldState && worldState.sortedBefore && worldState.totalActiveSplats > 0;
		if (!ready) {
			this.entries.forEach((entry) => {
				entry.meshInstance.visible = false;
			});
			return;
		}
		this._compaction.uploadIntervals(worldState);
		this.world.workBuffer.frustumCuller.updateTransformsData(worldState.boundsGroups);
		this._syncUserModify(gsplatParams);
		this._ensureCullShader();
		const numIntervals = worldState.totalIntervals;
		const totalActiveSplats = worldState.totalActiveSplats;
		const textureSize = this.world.workBuffer.textureSize;
		this.entries.forEach((entry) => {
			this._cullEntry(entry, numIntervals, totalActiveSplats, textureSize, gsplatParams);
		});
	}
	_syncUserModify(gsplatParams) {
		const userMat = gsplatParams.material;
		if (!userMat) return;
		const chunksKey = userMat.shaderChunks?.key ?? "";
		if (chunksKey !== this._userChunksKey) {
			this._userChunksKey = chunksKey;
			this._userModifyWgsl = userMat.getShaderChunks?.("wgsl")?.get("gsplatModifyVS") ?? null;
			this.entries.forEach((entry) => this._applyUserModify(entry));
		}
		const params = userMat.parameters;
		this.entries.forEach((entry) => {
			for (const name in params) {
				if (params.hasOwnProperty(name)) {
					entry.material.setParameter(name, params[name].data);
				}
			}
		});
	}
	_applyUserModify(entry) {
		const wgsl = entry.material.shaderChunks.wgsl;
		if (this._userModifyWgsl) {
			wgsl.set("gsplatModifyVS", this._userModifyWgsl);
		} else {
			wgsl.delete("gsplatModifyVS");
		}
		entry.material.update();
	}
	_cullEntry(entry, numIntervals, totalActiveSplats, textureSize, gsplatParams) {
		const device = this.device;
		const sceneCamera = this.cameraNode.camera?.camera;
		const shadowCamera = sceneCamera && entry.light.getRenderData(sceneCamera, 0).shadowCamera;
		const frustum = shadowCamera && shadowCamera.frustum;
		const frustumCuller = this.world.workBuffer.frustumCuller;
		if (!frustum || !frustumCuller?.boundsBuffer || !frustumCuller?.transformsBuffer) {
			entry.meshInstance.visible = false;
			return;
		}
		this._fillFrustumPlanes(frustum);
		const orthoHeight = shadowCamera.orthoHeight;
		const shadowRes = entry.light._shadowResolution;
		const minPixelSize = gsplatParams.minPixelSize;
		const focal = orthoHeight > 0 && shadowRes > 0 ? shadowRes / orthoHeight : 0;
		const t2 = minPixelSize * minPixelSize * 0.5 - 0.3;
		const worldSizeThreshold = focal > 0 && t2 > 0 ? Math.sqrt(t2) / focal : 0;
		if (totalActiveSplats > entry.allocatedIndexCount) {
			entry.indexBuffer?.destroy();
			entry.allocatedIndexCount = totalActiveSplats;
			entry.indexBuffer = new StorageBuffer(device, totalActiveSplats * 4);
		}
		const compaction = this._compaction;
		compaction.dispatchCompact({
			boundsBuffer: frustumCuller.boundsBuffer,
			transformsBuffer: frustumCuller.transformsBuffer,
			frustumPlanes: this._frustumPlanes
		}, numIntervals, totalActiveSplats, false);
		entry.countBuffer.clear();
		if (!entry.cullCompute || entry.cullComputeGen !== this._cullShaderGen) {
			entry.cullCompute?.destroy();
			entry.cullCompute = new Compute(device, this._cullShader, "GSplatShadowCull");
			entry.cullComputeGen = this._cullShaderGen;
		}
		const cull = entry.cullCompute;
		const userMat = gsplatParams.material;
		if (userMat) {
			const srcParams = userMat.parameters;
			for (const name in srcParams) {
				if (srcParams.hasOwnProperty(name)) {
					cull.setParameter(name, srcParams[name].data);
				}
			}
		}
		cull.setParameter("compactedSplatIds", compaction.compactedSplatIds);
		cull.setParameter("candidateCountBuffer", compaction.countBuffer);
		cull.setParameter("outputIndices", entry.indexBuffer);
		cull.setParameter("globalCount", entry.countBuffer);
		cull.setParameter("frustumPlanes[0]", this._frustumPlanes);
		cull.setParameter("numIntervals", numIntervals);
		cull.setParameter("splatTextureSize", textureSize);
		cull.setParameter("alphaClip", gsplatParams.alphaClip);
		cull.setParameter("worldSizeThreshold", worldSizeThreshold);
		const workBuffer = this.world.workBuffer;
		for (const stream of workBuffer.format.resourceStreams) {
			const texture = workBuffer.getTexture(stream.name);
			if (texture) {
				cull.setParameter(stream.name, texture);
			}
		}
		const workgroupCount = Math.ceil(totalActiveSplats / WORKGROUP_SIZE);
		Compute.calcDispatchSize(workgroupCount, this._cullDispatchSize, device.limits.maxComputeWorkgroupsPerDimension || 65535);
		cull.setupDispatch(this._cullDispatchSize.x, this._cullDispatchSize.y, 1);
		device.computeDispatch([cull], "GSplatShadowCull");
		const drawSlot = device.getIndirectDrawSlot(1);
		const args = entry.argsCompute;
		args.setParameter("countBuffer", entry.countBuffer);
		args.setParameter("indirectDrawArgs", device.indirectDrawBuffer);
		args.setParameter("drawSlot", drawSlot);
		args.setParameter("indexCount", INDEX_COUNT);
		args.setParameter("pad0", 0);
		args.setParameter("pad1", 0);
		args.setupDispatch(1);
		device.computeDispatch([args], "GSplatShadowIndirectArgs");
		const material = entry.material;
		entry.meshInstance.setIndirect(null, drawSlot, 1);
		material.setParameter("compactedSplatIds", entry.indexBuffer);
		material.setParameter("numSplatsStorage", entry.countBuffer);
		material.setParameter("splatTextureSize", textureSize);
		material.setParameter("alphaClip", gsplatParams.alphaClip);
		entry.meshInstance.visible = true;
		if (entry.meshInstance.instancingCount <= 0) {
			entry.meshInstance.instancingCount = 1;
		}
	}
	_fillFrustumPlanes(frustum) {
		const p = this._frustumPlanes;
		for (let i = 0; i < 6; i++) {
			const plane = frustum.planes[i];
			p[i * 4 + 0] = plane.normal.x;
			p[i * 4 + 1] = plane.normal.y;
			p[i * 4 + 2] = plane.normal.z;
			p[i * 4 + 3] = plane.distance;
		}
	}
	_createEntry(light) {
		const device = this.device;
		const material = this._createMaterial();
		const meshInstance = this._createMeshInstance(light, material);
		meshInstance.castShadow = true;
		this.layer.addShadowCasters([meshInstance]);
		const countBuffer = new StorageBuffer(device, 4, BUFFERUSAGE_COPY_DST);
		const argsCompute = new Compute(device, this._argsShader, "GSplatShadowIndirectArgs");
		const entry = {
			light,
			material,
			meshInstance,
			indexBuffer: null,
			allocatedIndexCount: 0,
			countBuffer,
			cullCompute: null,
			cullComputeGen: -1,
			argsCompute
		};
		this._applyUserModify(entry);
		return entry;
	}
	_destroyEntry(entry) {
		this.layer.removeShadowCasters([entry.meshInstance]);
		entry.meshInstance.destroy();
		entry.material.destroy();
		entry.indexBuffer?.destroy();
		entry.countBuffer.destroy();
		entry.cullCompute?.destroy();
		entry.argsCompute.destroy();
	}
	_createMaterial() {
		const material = new ShaderMaterial({
			uniqueName: "GSplatShadowMaterial",
			vertexGLSL: '#include "gsplatVS"',
			fragmentGLSL: '#include "gsplatPS"',
			vertexWGSL: '#include "gsplatVS"',
			fragmentWGSL: '#include "gsplatPS"',
			attributes: {
				vertex_position: SEMANTIC_POSITION
			}
		});
		material.setDefine("{GSPLAT_INSTANCE_SIZE}", GSplatResourceBase.instanceSize);
		material.setDefine("SH_BANDS", "0");
		material.setDefine("GSPLAT_SEPARATE_OPACITY", "");
		material.setDefine("DITHER_NONE", "");
		material.setDefine("GSPLAT_INDIRECT_DRAW", true);
		this._configureMaterialWorkBuffer(material);
		material.cull = CULLFACE_NONE;
		material.blendType = BLEND_PREMULTIPLIED;
		material.depthWrite = false;
		material.update();
		return material;
	}
	_configureMaterialWorkBuffer(material) {
		const workBuffer = this.world.workBuffer;
		const wbFormat = workBuffer.format;
		const chunks = this.device.isWebGPU ? material.shaderChunks.wgsl : material.shaderChunks.glsl;
		chunks.set("gsplatDeclarationsVS", wbFormat.getInputDeclarations());
		chunks.set("gsplatReadVS", wbFormat.getReadCode());
		const colorStream = wbFormat.getStream("dataColor");
		if (colorStream && colorStream.format !== PIXELFORMAT_RGBA16U) {
			material.setDefine("GSPLAT_COLOR_FLOAT", "");
		}
		const hasPcId = !!wbFormat.getStream("pcId");
		material.setDefine("GSPLAT_UNIFIED_ID", hasPcId);
		material.setDefine("PICK_CUSTOM_ID", hasPcId);
		for (const stream of wbFormat.resourceStreams) {
			const texture = workBuffer.getTexture(stream.name);
			if (texture) {
				material.setParameter(stream.name, texture);
			}
		}
	}
	_createMeshInstance(light, material) {
		const mesh = GSplatResourceBase.createMesh(this.device);
		const meshInstance = new MeshInstance(mesh, material);
		meshInstance.node = this.node;
		meshInstance.setInstancing(true, true);
		meshInstance.instancingCount = 0;
		meshInstance.pick = false;
		const cameraNode = this.cameraNode;
		meshInstance.isVisibleFunc = (camera) => {
			const sceneCamera = cameraNode.camera?.camera;
			if (!sceneCamera) return false;
			return camera === light.getRenderData(sceneCamera, 0).shadowCamera;
		};
		return meshInstance;
	}
}
export {
	GSplatShadowRenderer
};