playcanvas

Version:
Open-source WebGL/WebGPU 3D engine for the web
47 lines (44 loc) • 1.98 kB
JavaScript
import indirectCoreCS from '../common/comp/indirect-core.js';
import dispatchCoreCS from '../common/comp/dispatch-core.js';

const computeGsplatWriteIndirectArgsSource = `
${indirectCoreCS}
${dispatchCoreCS}
@group(0) @binding(0) var<storage, read> prefixSumBuffer: array<u32>;
@group(0) @binding(1) var<storage, read_write> indirectDrawArgs: array<DrawIndexedIndirectArgs>;
@group(0) @binding(2) var<storage, read_write> numSplatsBuf: array<u32>;
@group(0) @binding(3) var<storage, read_write> indirectDispatchArgs: array<u32>;
@group(0) @binding(4) var<storage, read_write> sortElementCountBuf: array<u32>;
struct WriteArgsUniforms {
	drawSlot: u32,
	indexCount: u32,
	dispatchSlotOffset: u32,
	totalSplats: u32
};
@group(0) @binding(5) var<uniform> uniforms: WriteArgsUniforms;
@compute @workgroup_size(1)
fn main(@builtin(global_invocation_id) gid: vec3u) {
	let count = prefixSumBuffer[uniforms.totalSplats];
	let instanceCount = (count + {INSTANCE_SIZE}u - 1u) / {INSTANCE_SIZE}u;
	indirectDrawArgs[uniforms.drawSlot] = DrawIndexedIndirectArgs(
		uniforms.indexCount,
		instanceCount,
		0u,
		0,
		0u
	);
	numSplatsBuf[0] = count;
	let dispatchOffset = uniforms.dispatchSlotOffset;
	let keygenWorkgroupCount = (count + {KEYGEN_THREADS_PER_WORKGROUP}u - 1u) / {KEYGEN_THREADS_PER_WORKGROUP}u;
	let keygenDim = calcDispatch2D(keygenWorkgroupCount, {MAX_WORKGROUPS_PER_DIM}u);
	indirectDispatchArgs[dispatchOffset + 0u] = keygenDim.x;
	indirectDispatchArgs[dispatchOffset + 1u] = keygenDim.y;
	indirectDispatchArgs[dispatchOffset + 2u] = 1u;
	let sortWorkgroupCount = (count + {SORT_ELEMENTS_PER_WORKGROUP}u - 1u) / {SORT_ELEMENTS_PER_WORKGROUP}u;
	let sortDim = calcDispatch2D(sortWorkgroupCount, {MAX_WORKGROUPS_PER_DIM}u);
	indirectDispatchArgs[dispatchOffset + 3u] = sortDim.x;
	indirectDispatchArgs[dispatchOffset + 4u] = sortDim.y;
	indirectDispatchArgs[dispatchOffset + 5u] = 1u;
	sortElementCountBuf[0] = count;
}
`;

export { computeGsplatWriteIndirectArgsSource };