UNPKG

playcanvas

Version:

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

playcanvas.com

playcanvas/engine

3 lines (2 loc) • 4.24 kB
1
2
3
export const computeGsplatSortKeySource: "\n\n// Work buffer texture containing world-space centers (RGBA32U: xyz as floatBitsToUint)\n@group(0) @binding(0) var dataTransformA: texture_2d<u32>;\n\n// Output sort keys (one u32 per splat)\n@group(0) @binding(1) var<storage, read_write> sortKeys: array<u32>;\n\n// Uniforms\nstruct SortKeyUniforms {\n cameraPosition: vec3f,\n elementCount: u32,\n cameraDirection: vec3f,\n numBits: u32,\n textureSize: u32,\n minDist: f32,\n invRange: f32,\n numWorkgroupsX: u32,\n numBins: u32\n};\n@group(0) @binding(2) var<uniform> uniforms: SortKeyUniforms;\n\n// Camera-relative bin weighting (entries with base and divider)\nstruct BinWeight {\n base: f32,\n divider: f32\n};\n@group(0) @binding(3) var<storage, read> binWeights: array<BinWeight>;\n\n#ifdef USE_INDIRECT_SORT\n // Compacted visible splat IDs from stream compaction\n @group(0) @binding(4) var<storage, read> compactedSplatIds: array<u32>;\n\n // sortElementCount from write-indirect-args (same buffer the radix sort reads)\n @group(0) @binding(5) var<storage, read> sortElementCountBuf: array<u32>;\n#endif\n\n@compute @workgroup_size({WORKGROUP_SIZE_X}, {WORKGROUP_SIZE_Y}, 1)\nfn main(\n @builtin(global_invocation_id) global_id: vec3u,\n @builtin(workgroup_id) w_id: vec3u,\n @builtin(num_workgroups) w_dim: vec3u,\n @builtin(local_invocation_index) TID: u32\n) {\n // Compute flat GID: for indirect dispatch use workgroup/local builtins (dispatch\n // dimensions are GPU-written so uniforms.numWorkgroupsX won't match). For direct\n // dispatch the 2D global_invocation_id linearization is used.\n #ifdef USE_INDIRECT_SORT\n let WORKGROUP_ID = w_id.x + w_id.y * w_dim.x;\n let gid = WORKGROUP_ID * ({WORKGROUP_SIZE_X}u * {WORKGROUP_SIZE_Y}u) + TID;\n #else\n let gid = global_id.x + global_id.y * ({WORKGROUP_SIZE_X} * uniforms.numWorkgroupsX);\n #endif\n \n // Early exit for out-of-bounds threads\n if (gid >= uniforms.elementCount) {\n return;\n }\n\n #ifdef USE_INDIRECT_SORT\n // With indirect dispatch, only visibleCount threads are launched (plus up to\n // 255 padding threads from workgroup rounding). Early-out for padding threads.\n let visibleCount = sortElementCountBuf[0];\n if (gid >= visibleCount) {\n return;\n }\n\n // Read actual splat ID from compacted buffer\n let splatId = compactedSplatIds[gid];\n #else\n let splatId = gid;\n #endif\n \n // Calculate texture UV from splat ID\n let textureSize = uniforms.textureSize;\n let uv = vec2i(i32(splatId % textureSize), i32(splatId / textureSize));\n \n // Load world-space center from work buffer (stored as floatBitsToUint)\n let packed = textureLoad(dataTransformA, uv, 0);\n let worldCenter = vec3f(\n bitcast<f32>(packed.r),\n bitcast<f32>(packed.g),\n bitcast<f32>(packed.b)\n );\n \n // Calculate distance based on sort mode\n var dist: f32;\n \n #ifdef RADIAL_SORT\n // Radial mode: distance from camera (inverted so far objects get small keys)\n let delta = worldCenter - uniforms.cameraPosition;\n let radialDist = length(delta);\n // Invert distance so far objects get small keys (rendered first, back-to-front)\n dist = (1.0 / uniforms.invRange) - radialDist - uniforms.minDist;\n #else\n // Linear mode: distance along camera forward vector\n let toSplat = worldCenter - uniforms.cameraPosition;\n dist = dot(toSplat, uniforms.cameraDirection) - uniforms.minDist;\n #endif\n \n // Apply bin-based mapping for camera-relative precision weighting\n let numBins = uniforms.numBins;\n let d = dist * uniforms.invRange * f32(numBins);\n let binFloat = clamp(d, 0.0, f32(numBins) - 0.001);\n let bin = u32(binFloat);\n let binFrac = binFloat - f32(bin);\n \n // Calculate final sort key using pre-computed bin weighting\n let sortKey = u32(binWeights[bin].base + binWeights[bin].divider * binFrac);\n \n // Write sort key\n sortKeys[gid] = sortKey;\n}\n"; export default computeGsplatSortKeySource;