playcanvas

Version:
Open-source WebGL/WebGPU 3D engine for the web
73 lines (72 loc) • 2.92 kB
JavaScript
import indirectCoreCS from "../common/comp/indirect-core.js";
import dispatchCoreCS from "../common/comp/dispatch-core.js";
const computeGsplatLocalClassifySource = `
${indirectCoreCS}
${dispatchCoreCS}
const MAX_TILE_ENTRIES: u32 = 4096u;
const CLASSIFY_WORKGROUP: u32 = 256u;
@group(0) @binding(0) var<storage, read> tileSplatCounts: array<u32>;
@group(0) @binding(1) var<storage, read_write> smallTileList: array<u32>;
@group(0) @binding(2) var<storage, read_write> largeTileList: array<u32>;
@group(0) @binding(3) var<storage, read_write> rasterizeTileList: array<u32>;
@group(0) @binding(4) var<storage, read_write> tileListCounts: array<atomic<u32>>;
@group(0) @binding(5) var<storage, read_write> indirectDispatchArgs: array<u32>;
@group(0) @binding(6) var<storage, read_write> largeTileOverflowBases: array<u32>;
@group(0) @binding(8) var<storage, read_write> indirectDrawArgs: array<DrawIndirectArgs>;
struct Uniforms {
	numTiles: u32,
	dispatchSlotOffset: u32,
	bufferCapacity: u32,
	maxWorkgroupsPerDim: u32,
	drawSlot: u32,
}
@group(0) @binding(7) var<uniform> uniforms: Uniforms;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_index) localIdx: u32) {
	let numTiles = uniforms.numTiles;
	let totalEntries = tileSplatCounts[numTiles];
	for (var i: u32 = localIdx; i < numTiles; i += CLASSIFY_WORKGROUP) {
		let tStart = tileSplatCounts[i];
		let tEnd = tileSplatCounts[i + 1u];
		let count = tEnd - tStart;
		if (count == 0u || tEnd > uniforms.bufferCapacity) {
			continue;
		}
		let rIdx = atomicAdd(&tileListCounts[2], 1u);
		rasterizeTileList[rIdx] = i;
		if (count <= MAX_TILE_ENTRIES) {
			let sIdx = atomicAdd(&tileListCounts[0], 1u);
			smallTileList[sIdx] = i;
		} else {
			let overflowOffset = atomicAdd(&tileListCounts[3], count);
			let lIdx = atomicAdd(&tileListCounts[1], 1u);
			largeTileList[lIdx] = i;
			largeTileOverflowBases[lIdx] = totalEntries + overflowOffset;
		}
	}
	workgroupBarrier();
	if (localIdx == 0u) {
		let smallCount = atomicLoad(&tileListCounts[0]);
		let largeCount = atomicLoad(&tileListCounts[1]);
		let rasterizeCount = atomicLoad(&tileListCounts[2]);
		let off = uniforms.dispatchSlotOffset;
		let maxDim = uniforms.maxWorkgroupsPerDim;
		let smallDim = calcDispatch2D(smallCount, maxDim);
		indirectDispatchArgs[off + 0u] = smallDim.x;
		indirectDispatchArgs[off + 1u] = smallDim.y;
		indirectDispatchArgs[off + 2u] = 1u;
		let largeDim = calcDispatch2D(largeCount, maxDim);
		indirectDispatchArgs[off + 3u] = largeDim.x;
		indirectDispatchArgs[off + 4u] = largeDim.y;
		indirectDispatchArgs[off + 5u] = 1u;
		let rasterDim = calcDispatch2D(rasterizeCount, maxDim);
		indirectDispatchArgs[off + 6u] = rasterDim.x;
		indirectDispatchArgs[off + 7u] = rasterDim.y;
		indirectDispatchArgs[off + 8u] = 1u;
		indirectDrawArgs[uniforms.drawSlot] = DrawIndirectArgs(rasterizeCount * 6u, 1u, 0u, 0u, 0u);
	}
}
`;
export {
	computeGsplatLocalClassifySource
};