UNPKG

playcanvas

Version:

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

playcanvas.com

playcanvas/engine

2 lines (1 loc) • 17.6 kB
1
2
export const computeGsplatLocalRasterizeSource: "\n\n#include \"halfTypesCS\"\n#ifndef PICK_MODE\n #include \"decodePS\"\n #if FOG != NONE\n #include \"fogMathPS\"\n #include \"gammaPS\"\n #endif\n#endif\n\nconst CACHE_STRIDE: u32 = 7u;\nconst BATCH_SIZE: u32 = 64u;\nconst ALPHA_THRESHOLD: half = half(1.0) / half(255.0);\nconst EXP4: half = exp(half(-4.0));\nconst INV_EXP4: half = half(1.0) / (half(1.0) - EXP4);\n\n// Shared uniforms and storage buffers (same layout for both color and pick modes).\n// Pick-only uniforms (nearClip, farClip, alphaClip) are present but unused in color mode.\nstruct Uniforms {\n screenWidth: u32,\n screenHeight: u32,\n numTilesX: u32,\n nearClip: f32,\n farClip: f32,\n alphaClip: f32,\n #if FOG != NONE\n fog_color: vec3f,\n fog_start: f32,\n fog_end: f32,\n fog_density: f32,\n #endif\n}\n@group(0) @binding(0) var<uniform> uniforms: Uniforms;\n@group(0) @binding(1) var<storage, read> tileEntries: array<u32>;\n@group(0) @binding(2) var<storage, read> tileSplatCounts: array<u32>;\n@group(0) @binding(3) var<storage, read> projCache: array<u32>;\n@group(0) @binding(4) var<storage, read> rasterizeTileList: array<u32>;\n@group(0) @binding(5) var<storage, read> tileListCounts: array<u32>;\n@group(0) @binding(6) var<storage, read> depthBuffer: array<u32>;\n\n// Mode-specific output textures appended after the shared bindings.\n#ifdef PICK_MODE\n @group(0) @binding(7) var pickIdTexture: texture_storage_2d<r32uint, write>;\n @group(0) @binding(8) var pickDepthTexture: texture_storage_2d<rgba16float, write>;\n#else\n @group(0) @binding(7) var outputTexture: texture_storage_2d<rgba16float, write>;\n #ifdef DEPTH_TEST\n @group(0) @binding(8) var sceneDepthMap: texture_2d<f32>;\n #endif\n#endif\n\nvar<workgroup> sharedCenterScreen: array<vec2f, 64>;\nvar<workgroup> sharedCoeffs: array<vec3f, 64>;\n#ifdef HEATMAP_MODE\n var<workgroup> sharedHeatCount: atomic<u32>;\n#endif\n\n// Pick mode stores per-splat opacity, ID and depth; color mode stores packed RGBA.\n// Depth test mode also needs per-splat view depth for occlusion against scene geometry.\n#ifdef PICK_MODE\n var<workgroup> sharedOpacity: array<half, 64>;\n var<workgroup> sharedPickId: array<u32, 64>;\n var<workgroup> sharedViewDepth: array<f32, 64>;\n#else\n var<workgroup> sharedColor: array<half4, 64>;\n #ifdef DEPTH_TEST\n var<workgroup> sharedViewDepth: array<f32, 64>;\n #endif\n#endif\n\n// Evaluate a single splat for picking. Records the front-most pick ID (first splat above\n// alphaClip) and accumulates alpha-weighted depth for sub-pixel depth reconstruction.\n#ifdef PICK_MODE\nfn evalSplatPick(pixelCoord: vec2f, center: vec2f, coeffX: f32, coeffY: f32, coeffXY: f32,\n opacity: half, pickId: u32, viewDepth: f32, alphaClip: half,\n bestPickId: ptr<function, u32>, depthAccum: ptr<function, f32>,\n weightAccum: ptr<function, f32>, T: ptr<function, half>) {\n let dx = pixelCoord - center;\n let power = coeffX * dx.x * dx.x + coeffXY * dx.x * dx.y + coeffY * dx.y * dx.y;\n let gauss = (half(exp(power)) - EXP4) * INV_EXP4;\n let alpha = half(min(half(0.99), opacity * gauss));\n let newT = *T * (half(1.0) - alpha);\n let visible = power > -4.0 && alpha > ALPHA_THRESHOLD && *T >= ALPHA_THRESHOLD;\n if (!visible) { return; }\n\n // Per-pixel alphaClip: only solid splat centers contribute to pick\n if (alpha >= alphaClip) {\n if (*bestPickId == 0xFFFFFFFFu) {\n *bestPickId = pickId;\n }\n let normalizedDepth = saturate((viewDepth - uniforms.nearClip) / (uniforms.farClip - uniforms.nearClip));\n let w = f32(alpha) * f32(*T);\n *depthAccum += w * normalizedDepth;\n *weightAccum += w;\n }\n\n *T = newT;\n}\n#endif\n\n#ifdef HEATMAP_MODE\nfn heatmapColor(v: f32) -> vec3f {\n let t = saturate(v / 2000.0);\n if (t < 0.2) {\n return mix(vec3f(0.0, 0.0, 1.0), vec3f(0.0, 1.0, 1.0), t * 5.0);\n } else if (t < 0.4) {\n return mix(vec3f(0.0, 1.0, 1.0), vec3f(1.0, 1.0, 0.0), (t - 0.2) * 5.0);\n } else if (t < 0.6) {\n return mix(vec3f(1.0, 1.0, 0.0), vec3f(1.0, 0.0, 0.0), (t - 0.4) * 5.0);\n }\n return mix(vec3f(1.0, 0.0, 0.0), vec3f(0.15, 0.0, 0.0), (t - 0.6) * 2.5);\n}\n#endif\n\n@compute @workgroup_size(8, 8)\nfn main(\n @builtin(local_invocation_id) lid: vec3u,\n @builtin(local_invocation_index) localIdx: u32,\n @builtin(workgroup_id) wid: vec3u,\n @builtin(num_workgroups) numWorkgroups: vec3u\n) {\n let workgroupIdx = wid.y * numWorkgroups.x + wid.x;\n if (workgroupIdx >= tileListCounts[2]) {\n return;\n }\n let tileIdx = rasterizeTileList[workgroupIdx];\n let tileX = tileIdx % uniforms.numTilesX;\n let tileY = tileIdx / uniforms.numTilesX;\n let tStart = tileSplatCounts[tileIdx];\n let tEnd = tileSplatCounts[tileIdx + 1u];\n\n let basePixel = vec2u(tileX * 16u + lid.x * 2u, tileY * 16u + lid.y * 2u);\n let p00 = vec2f(f32(basePixel.x) + 0.5, f32(basePixel.y) + 0.5);\n let p10 = p00 + vec2f(1.0, 0.0);\n let p01 = p00 + vec2f(0.0, 1.0);\n let p11 = p00 + vec2f(1.0, 1.0);\n\n // Per-pixel state for the 2x2 quad.\n\n #ifdef PICK_MODE\n var T00: half = half(1.0); var T10: half = half(1.0);\n var T01: half = half(1.0); var T11: half = half(1.0);\n\n // front-most pick ID per pixel\n var pickId00: u32 = 0xFFFFFFFFu; var pickId10: u32 = 0xFFFFFFFFu;\n var pickId01: u32 = 0xFFFFFFFFu; var pickId11: u32 = 0xFFFFFFFFu;\n\n // alpha-weighted depth accumulators\n var dAcc00: f32 = 0.0; var dAcc10: f32 = 0.0;\n var dAcc01: f32 = 0.0; var dAcc11: f32 = 0.0;\n\n // alpha-weighted weight accumulators\n var wAcc00: f32 = 0.0; var wAcc10: f32 = 0.0;\n var wAcc01: f32 = 0.0; var wAcc11: f32 = 0.0;\n let clipH = half(uniforms.alphaClip);\n #else\n // Transmittance for the 2x2 quad packed as vec4<half> (x=00, y=10, z=01, w=11).\n // Tracking how much light passes through the splat stack at that pixel. Packing\n // them into a single vec4 lets the compiler use vector ALU for the branchless\n // alpha-blending update and the all-saturated early-out test.\n var T = half4(1.0);\n\n // accumulated color per pixel\n var c00 = half3(0.0); var c10 = half3(0.0);\n var c01 = half3(0.0); var c11 = half3(0.0);\n\n #ifdef DEPTH_TEST\n // Load per-pixel linear scene depth for the 2x2 quad (x=00, y=10, z=01, w=11).\n // Splats behind this depth are skipped during rasterization.\n var sceneDepth = vec4f(1e20);\n let depthY0 = uniforms.screenHeight - 1u - basePixel.y;\n let depthY1 = depthY0 - 1u;\n if (basePixel.x < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {\n sceneDepth.x = textureLoad(sceneDepthMap, vec2i(vec2u(basePixel.x, depthY0)), 0).r;\n }\n if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {\n sceneDepth.y = textureLoad(sceneDepthMap, vec2i(vec2u(basePixel.x + 1u, depthY0)), 0).r;\n }\n if (basePixel.x < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {\n sceneDepth.z = textureLoad(sceneDepthMap, vec2i(vec2u(basePixel.x, depthY1)), 0).r;\n }\n if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {\n sceneDepth.w = textureLoad(sceneDepthMap, vec2i(vec2u(basePixel.x + 1u, depthY1)), 0).r;\n }\n #endif\n #endif\n\n let tileCount = tEnd - tStart;\n\n #ifdef HEATMAP_MODE\n if (localIdx == 0u) { atomicStore(&sharedHeatCount, 0u); }\n workgroupBarrier();\n var processedCount: u32 = 0u;\n #endif\n\n let numBatches = (tileCount + BATCH_SIZE - 1u) / BATCH_SIZE;\n var threadDone = false;\n\n for (var batch: u32 = 0u; batch < numBatches; batch++) {\n\n let batchOffset = batch * BATCH_SIZE + localIdx;\n if (batchOffset < tileCount) {\n let cacheIdx = tileEntries[tStart + batchOffset];\n let base = cacheIdx * CACHE_STRIDE;\n sharedCenterScreen[localIdx] = vec2f(\n bitcast<f32>(projCache[base + 0u]),\n bitcast<f32>(projCache[base + 1u])\n );\n // Conic values cx/cy/cz stored as f32; convert to evaluation coefficients.\n let cx = bitcast<f32>(projCache[base + 2u]);\n let cy = bitcast<f32>(projCache[base + 3u]);\n let cz = bitcast<f32>(projCache[base + 4u]);\n sharedCoeffs[localIdx] = vec3f(cx * -0.5, cz * -0.5, -cy);\n\n #ifdef PICK_MODE\n sharedPickId[localIdx] = projCache[base + 5u];\n sharedOpacity[localIdx] = half(unpack2x16float(projCache[base + 6u]).y);\n sharedViewDepth[localIdx] = bitcast<f32>(depthBuffer[cacheIdx]);\n #else\n let rg = unpack2x16float(projCache[base + 5u]);\n let ba = unpack2x16float(projCache[base + 6u]);\n\n #if FOG != NONE\n let viewDepth = bitcast<f32>(depthBuffer[cacheIdx]);\n #if (FOG == LINEAR)\n let fogFactor = evaluateFogFactorLinear(viewDepth, uniforms.fog_start, uniforms.fog_end);\n #elif (FOG == EXP)\n let fogFactor = evaluateFogFactorExp(viewDepth, uniforms.fog_density);\n #elif (FOG == EXP2)\n let fogFactor = evaluateFogFactorExp2(viewDepth, uniforms.fog_density);\n #endif\n var foggedColor = decodeGamma3(vec3f(rg.x, rg.y, ba.x));\n foggedColor = mix(uniforms.fog_color, foggedColor, fogFactor);\n sharedColor[localIdx] = half4(half3(gammaCorrectOutput(foggedColor)), half(ba.y));\n #else\n sharedColor[localIdx] = half4(half(rg.x), half(rg.y), half(ba.x), half(ba.y));\n #endif\n\n #ifdef DEPTH_TEST\n sharedViewDepth[localIdx] = bitcast<f32>(depthBuffer[cacheIdx]);\n #endif\n #endif\n }\n\n workgroupBarrier();\n\n if (!threadDone) {\n let batchCount = min(BATCH_SIZE, tileCount - batch * BATCH_SIZE);\n\n for (var i: u32 = 0u; i < batchCount; i++) {\n let center = sharedCenterScreen[i];\n let coeffs = sharedCoeffs[i];\n\n #ifdef PICK_MODE\n let splatOpacity = sharedOpacity[i];\n let splatPickId = sharedPickId[i];\n let splatDepth = sharedViewDepth[i];\n\n evalSplatPick(p00, center, coeffs.x, coeffs.y, coeffs.z, splatOpacity, splatPickId, splatDepth, clipH, &pickId00, &dAcc00, &wAcc00, &T00);\n evalSplatPick(p10, center, coeffs.x, coeffs.y, coeffs.z, splatOpacity, splatPickId, splatDepth, clipH, &pickId10, &dAcc10, &wAcc10, &T10);\n evalSplatPick(p01, center, coeffs.x, coeffs.y, coeffs.z, splatOpacity, splatPickId, splatDepth, clipH, &pickId01, &dAcc01, &wAcc01, &T01);\n evalSplatPick(p11, center, coeffs.x, coeffs.y, coeffs.z, splatOpacity, splatPickId, splatDepth, clipH, &pickId11, &dAcc11, &wAcc11, &T11);\n\n if (all(vec4<half>(T00, T10, T01, T11) < half4(ALPHA_THRESHOLD))) {\n threadDone = true;\n break;\n }\n #else\n let splatColor = sharedColor[i];\n\n #ifdef DEPTH_TEST\n let splatDepth = sharedViewDepth[i];\n\n // Splats are front-to-back; if behind all four depth samples, all remaining splats will be too.\n if (all(vec4f(splatDepth) > sceneDepth)) {\n threadDone = true;\n break;\n }\n #endif\n\n // Vectorized Gaussian evaluation for the 2x2 pixel quad. Compute dx\n // once from p00, build the four pixel offsets as vec4f (exploiting the\n // regular +1 grid), and evaluate power/gauss/alpha/transmittance as\n // vec4 operations to share ALU across the quad.\n let d = p00 - center;\n let dxV = vec4f(d.x, d.x + 1.0, d.x, d.x + 1.0);\n let dyV = vec4f(d.y, d.y, d.y + 1.0, d.y + 1.0);\n let power4 = coeffs.x * dxV * dxV + coeffs.z * dxV * dyV + coeffs.y * dyV * dyV;\n let gauss4 = (half4(exp(power4)) - half4(EXP4)) * half4(INV_EXP4);\n let alpha4 = min(half4(0.99), half4(splatColor.a) * gauss4);\n let newT = T * (half4(1.0) - alpha4);\n\n var valid = (power4 > vec4f(-4.0)) & (alpha4 > half4(ALPHA_THRESHOLD)) & (T >= half4(ALPHA_THRESHOLD));\n #ifdef DEPTH_TEST\n valid = valid & (vec4f(splatDepth) <= sceneDepth);\n #endif\n\n let weight = alpha4 * T * select(half4(0.0), half4(1.0), valid);\n c00 += splatColor.rgb * weight.x;\n c10 += splatColor.rgb * weight.y;\n c01 += splatColor.rgb * weight.z;\n c11 += splatColor.rgb * weight.w;\n T = select(T, newT, valid);\n\n #ifdef HEATMAP_MODE\n processedCount += 1u;\n #endif\n\n if (all(T < half4(ALPHA_THRESHOLD))) {\n threadDone = true;\n break;\n }\n #endif\n }\n }\n\n workgroupBarrier();\n }\n\n #ifdef HEATMAP_MODE\n atomicAdd(&sharedHeatCount, processedCount);\n workgroupBarrier();\n let avgCount = f32(atomicLoad(&sharedHeatCount)) / 64.0;\n let heatColor = vec4f(heatmapColor(avgCount), 1.0);\n if (basePixel.x < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {\n textureStore(outputTexture, basePixel, heatColor);\n }\n if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {\n textureStore(outputTexture, vec2u(basePixel.x + 1u, basePixel.y), heatColor);\n }\n if (basePixel.x < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {\n textureStore(outputTexture, vec2u(basePixel.x, basePixel.y + 1u), heatColor);\n }\n if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {\n textureStore(outputTexture, vec2u(basePixel.x + 1u, basePixel.y + 1u), heatColor);\n }\n #else\n\n // Write results for the 2x2 pixel quad owned by this thread.\n // Pick mode: store the front-most pick ID and (accumulated depth, weight) per pixel.\n // Color mode: convert accumulated gamma-space color to linear via decodeGamma3 and store\n // to the rgba16float output texture; alpha holds total opacity (1 - transmittance).\n if (basePixel.x < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {\n #ifdef PICK_MODE\n textureStore(pickIdTexture, basePixel, vec4u(pickId00, 0u, 0u, 0u));\n textureStore(pickDepthTexture, basePixel, vec4f(dAcc00, wAcc00, 0.0, 0.0));\n #else\n textureStore(outputTexture, basePixel, vec4f(decodeGamma3(vec3f(c00)), f32(half(1.0) - T.x)));\n #endif\n }\n if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {\n let px10 = vec2u(basePixel.x + 1u, basePixel.y);\n #ifdef PICK_MODE\n textureStore(pickIdTexture, px10, vec4u(pickId10, 0u, 0u, 0u));\n textureStore(pickDepthTexture, px10, vec4f(dAcc10, wAcc10, 0.0, 0.0));\n #else\n textureStore(outputTexture, px10, vec4f(decodeGamma3(vec3f(c10)), f32(half(1.0) - T.y)));\n #endif\n }\n if (basePixel.x < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {\n let px01 = vec2u(basePixel.x, basePixel.y + 1u);\n #ifdef PICK_MODE\n textureStore(pickIdTexture, px01, vec4u(pickId01, 0u, 0u, 0u));\n textureStore(pickDepthTexture, px01, vec4f(dAcc01, wAcc01, 0.0, 0.0));\n #else\n textureStore(outputTexture, px01, vec4f(decodeGamma3(vec3f(c01)), f32(half(1.0) - T.z)));\n #endif\n }\n if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {\n let px11 = vec2u(basePixel.x + 1u, basePixel.y + 1u);\n #ifdef PICK_MODE\n textureStore(pickIdTexture, px11, vec4u(pickId11, 0u, 0u, 0u));\n textureStore(pickDepthTexture, px11, vec4f(dAcc11, wAcc11, 0.0, 0.0));\n #else\n textureStore(outputTexture, px11, vec4f(decodeGamma3(vec3f(c11)), f32(half(1.0) - T.w)));\n #endif\n }\n\n #endif\n}\n";