playcanvas

Version:
Open-source WebGL/WebGPU 3D engine for the web
86 lines (76 loc) • 2.43 kB
JavaScript
const computeGsplatSortKeySource = `
@group(0) @binding(0) var dataTransformA: texture_2d<u32>;
@group(0) @binding(1) var<storage, read_write> sortKeys: array<u32>;
struct SortKeyUniforms {
	cameraPosition: vec3f,
	elementCount: u32,
	cameraDirection: vec3f,
	numBits: u32,
	textureSize: u32,
	minDist: f32,
	invRange: f32,
	numWorkgroupsX: u32,
	numBins: u32
};
@group(0) @binding(2) var<uniform> uniforms: SortKeyUniforms;
struct BinWeight {
	base: f32,
	divider: f32
};
@group(0) @binding(3) var<storage, read> binWeights: array<BinWeight>;
#ifdef USE_INDIRECT_SORT
	@group(0) @binding(4) var<storage, read> compactedSplatIds: array<u32>;
	@group(0) @binding(5) var<storage, read> sortElementCountBuf: array<u32>;
#endif
@compute @workgroup_size({WORKGROUP_SIZE_X}, {WORKGROUP_SIZE_Y}, 1)
fn main(
	@builtin(global_invocation_id) global_id: vec3u,
	@builtin(workgroup_id) w_id: vec3u,
	@builtin(num_workgroups) w_dim: vec3u,
	@builtin(local_invocation_index) TID: u32
) {
	#ifdef USE_INDIRECT_SORT
		let WORKGROUP_ID = w_id.x + w_id.y * w_dim.x;
		let gid = WORKGROUP_ID * ({WORKGROUP_SIZE_X}u * {WORKGROUP_SIZE_Y}u) + TID;
	#else
		let gid = global_id.x + global_id.y * ({WORKGROUP_SIZE_X} * uniforms.numWorkgroupsX);
	#endif
	
	if (gid >= uniforms.elementCount) {
		return;
	}
	#ifdef USE_INDIRECT_SORT
		let visibleCount = sortElementCountBuf[0];
		if (gid >= visibleCount) {
			return;
		}
		let splatId = compactedSplatIds[gid];
	#else
		let splatId = gid;
	#endif
	
	let textureSize = uniforms.textureSize;
	let uv = vec2i(i32(splatId % textureSize), i32(splatId / textureSize));
	
	let packed = textureLoad(dataTransformA, uv, 0);
	let worldCenter = vec3f(
		bitcast<f32>(packed.r),
		bitcast<f32>(packed.g),
		bitcast<f32>(packed.b)
	);
	
	var dist: f32;
	
	#ifdef RADIAL_SORT
		let delta = worldCenter - uniforms.cameraPosition;
		let radialDist = length(delta);
		dist = (1.0 / uniforms.invRange) - radialDist - uniforms.minDist;
	#else
		let toSplat = worldCenter - uniforms.cameraPosition;
		dist = dot(toSplat, uniforms.cameraDirection) - uniforms.minDist;
	#endif
	
	let numBins = uniforms.numBins;
	let d = dist * uniforms.invRange * f32(numBins);
	let binFloat = clamp(d, 0.0, f32(numBins) - 0.001);
	let bin = u32(binFloat);
	let binFrac = binFloat - f32(bin);
	
	let sortKey = u32(binWeights[bin].base + binWeights[bin].divider * binFrac);
	
	sortKeys[gid] = sortKey;
}
`;

export { computeGsplatSortKeySource };