playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
3 lines (2 loc) • 8.19 kB
TypeScript
export const computeGsplatProjectorSource: "\n\n#include \"gsplatCommonCS\"\n#include \"gsplatTileIntersectCS\"\n\n@group(0) @binding(0) var<storage, read> compactedSplatIds: array<u32>;\n@group(0) @binding(1) var<storage, read> sortElementCount: array<u32>;\n@group(0) @binding(2) var<storage, read_write> projCache: array<u32>;\n@group(0) @binding(3) var<storage, read_write> sortKeys: array<u32>;\n@group(0) @binding(4) var<storage, read_write> renderCounter: array<atomic<u32>>;\n\n// Camera-relative bin weighting for sort key computation (32 entries of {base, divider}).\nstruct BinWeight {\n base: f32,\n divider: f32\n}\n@group(0) @binding(5) var<storage, read> binWeights: array<BinWeight>;\n\nstruct ProjectorUniforms {\n splatTextureSize: u32,\n numBins: u32,\n isOrtho: u32,\n pad0: u32,\n viewProj: mat4x4f,\n viewMatrix: mat4x4f,\n cameraPosition: vec3f,\n minPixelSize: f32,\n cameraDirection: vec3f,\n focal: f32,\n viewportWidth: f32,\n viewportHeight: f32,\n nearClip: f32,\n farClip: f32,\n alphaClip: f32,\n minContribution: f32,\n minDist: f32,\n invRange: f32,\n #ifdef GSPLAT_FISHEYE\n fisheye_k: f32,\n fisheye_inv_k: f32,\n fisheye_projMat00: f32,\n fisheye_projMat11: f32,\n #endif\n}\n@group(0) @binding(6) var<uniform> uniforms: ProjectorUniforms;\n\n#include \"gsplatComputeSplatCS\"\n#include \"gsplatFormatDeclCS\"\n#include \"gsplatFormatReadCS\"\n#include \"gsplatProjectCommonCS\"\n\n// One global atomicAdd per workgroup (256 threads) \u2014 drastically lowers contention\n// vs a per-thread atomic on the global counter without needing subgroup ops.\nvar<workgroup> wgCount: atomic<u32>;\nvar<workgroup> wgBase: u32;\n\n@compute @workgroup_size(256)\nfn main(\n @builtin(global_invocation_id) gid: vec3u,\n @builtin(num_workgroups) numWorkgroups: vec3u,\n @builtin(local_invocation_index) localIdx: u32\n) {\n if (localIdx == 0u) {\n atomicStore(&wgCount, 0u);\n }\n workgroupBarrier();\n\n // Match the indirect dispatch linearisation used by compute-gsplat-local-tile-count.js:\n // a 2D grid expanded into a flat thread index.\n let threadIdx = gid.y * (numWorkgroups.x * 256u) + gid.x;\n let numVisible = sortElementCount[0];\n\n var valid = false;\n var clipPos: vec4f = vec4f(0.0);\n var v1: vec2f = vec2f(0.0);\n var v2: vec2f = vec2f(0.0);\n var rgb: vec3f = vec3f(0.0);\n var alpha: f32 = 0.0;\n var pcId: u32 = 0u;\n var sortKey: u32 = 0u;\n\n let projected = projectSplatCommon(\n threadIdx,\n numVisible,\n uniforms.alphaClip,\n uniforms.minPixelSize,\n uniforms.minContribution,\n uniforms.viewMatrix,\n uniforms.viewProj,\n uniforms.focal,\n uniforms.viewportWidth,\n uniforms.viewportHeight,\n uniforms.nearClip,\n uniforms.farClip,\n uniforms.isOrtho,\n #ifdef GSPLAT_FISHEYE\n uniforms.fisheye_k,\n uniforms.fisheye_inv_k,\n uniforms.fisheye_projMat00,\n uniforms.fisheye_projMat11,\n #endif\n );\n\n if (projected.valid) {\n let center = projected.center;\n let opacity = projected.opacity;\n let proj = projected.proj;\n\n // Eigen-decomposition matches gsplatCorner.js initCornerCov: derives the two\n // screen-pixel eigen-vectors of the 2D screen-space covariance.\n let mid = 0.5 * (proj.a + proj.c);\n let radius = length(vec2f(0.5 * (proj.a - proj.c), proj.b));\n let lambda1 = mid + radius;\n let lambda2 = max(mid - radius, 0.1);\n\n // capScale was already applied to a, b, c inside computeSplatCov so the\n // eigenvalues here are already in the radius-capped space. The vmin\n // saturation mirrors gsplatCorner.js (line 89-90) for hard parity with\n // the existing rasterizer when capScale > 1.\n let vmin = min(1024.0, min(uniforms.viewportWidth, uniforms.viewportHeight));\n let l1 = 2.0 * min(sqrt(2.0 * lambda1), vmin);\n let l2 = 2.0 * min(sqrt(2.0 * lambda2), vmin);\n\n let dir = normalize(vec2f(proj.b, lambda1 - proj.a));\n v1 = l1 * dir;\n v2 = l2 * vec2f(dir.y, -dir.x);\n\n // Clip-space center. The hybrid VS uses clipPos.w directly for both\n // output.position.w (rasterizer correctness) and corner scale, and\n // reconstructs linear view depth from clipPos via the clipToViewZ uniform.\n #ifdef GSPLAT_FISHEYE\n // Fisheye: invert the screen\u2192pixel mapping done by computeSplatCov to recover\n // NDC, then store NDC + linear depthNdc with w=1.0 (matches gsplatCenter.js).\n // The rasterizer's perspective divide is a no-op (w=1), so output.position is\n // NDC directly. The VS recovers linear -view.z via clipToViewZ = (0, 0, far-near, near).\n let viewCenter = uniforms.viewMatrix * vec4f(center, 1.0);\n let neg_z = -viewCenter.z;\n let ndcX = proj.screen.x / uniforms.viewportWidth * 2.0 - 1.0;\n let ndcY = proj.screen.y / uniforms.viewportHeight * 2.0 - 1.0;\n let depthNdc = clamp(\n (neg_z - uniforms.nearClip) / (uniforms.farClip - uniforms.nearClip),\n 0.0, 1.0\n );\n clipPos = vec4f(ndcX, ndcY, depthNdc, 1.0);\n #else\n clipPos = uniforms.viewProj * vec4f(center, 1.0);\n clipPos.z = clamp(clipPos.z, 0.0, abs(clipPos.w));\n #endif\n\n // Sort key \u2014 shared depth-bin weighting (same as CPU worker).\n #ifdef RADIAL_SORT\n let delta = center - uniforms.cameraPosition;\n let radialDist = length(delta);\n let dist = (1.0 / uniforms.invRange) - radialDist - uniforms.minDist;\n #else\n let toSplat = center - uniforms.cameraPosition;\n let dist = dot(toSplat, uniforms.cameraDirection) - uniforms.minDist;\n #endif\n let d = dist * uniforms.invRange * f32(uniforms.numBins);\n let binFloat = clamp(d, 0.0, f32(uniforms.numBins) - 0.001);\n let bin = u32(binFloat);\n let binFrac = binFloat - f32(bin);\n sortKey = u32(binWeights[bin].base + binWeights[bin].divider * binFrac);\n\n #ifdef PICK_MODE\n pcId = loadPcId().r;\n alpha = opacity;\n #else\n let color = getColor();\n rgb = max(color, vec3f(0.0));\n #if GSPLAT_AA\n // Bake the AA opacity compensation into the cached alpha (the hybrid VS\n // reads it pre-modulated). Only compiled for the non-pick variant.\n alpha = opacity * proj.aaFactor;\n #else\n alpha = opacity;\n #endif\n #endif\n\n valid = true;\n }\n\n // Reserve a per-workgroup slot for this thread.\n var localDst: u32 = 0u;\n if (valid) {\n localDst = atomicAdd(&wgCount, 1u);\n }\n workgroupBarrier();\n\n // Workgroup leader reserves a contiguous output range in renderCounter[0].\n if (localIdx == 0u) {\n let total = atomicLoad(&wgCount);\n wgBase = atomicAdd(&renderCounter[0], total);\n }\n workgroupBarrier();\n\n if (valid) {\n let dst = wgBase + localDst;\n let base = dst * {CACHE_STRIDE}u;\n\n projCache[base + 0u] = bitcast<u32>(clipPos.x);\n projCache[base + 1u] = bitcast<u32>(clipPos.y);\n projCache[base + 2u] = bitcast<u32>(clipPos.z);\n projCache[base + 3u] = bitcast<u32>(clipPos.w);\n projCache[base + 4u] = pack2x16float(v1);\n projCache[base + 5u] = pack2x16float(v2);\n\n #ifdef PICK_MODE\n projCache[base + 6u] = pcId;\n projCache[base + 7u] = pack2x16float(vec2f(0.0, alpha));\n #else\n projCache[base + 6u] = pack2x16float(vec2f(rgb.x, rgb.y));\n projCache[base + 7u] = pack2x16float(vec2f(rgb.z, alpha));\n #endif\n\n sortKeys[dst] = sortKey;\n }\n}\n";
export default computeGsplatProjectorSource;