playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
178 lines (177 loc) • 5.34 kB
JavaScript
const computeGsplatProjectorSource = `
@group(0) @binding(0) var<storage, read> compactedSplatIds: array<u32>;
@group(0) @binding(1) var<storage, read> sortElementCount: array<u32>;
@group(0) @binding(2) var<storage, read_write> projCache: array<u32>;
@group(0) @binding(3) var<storage, read_write> sortKeys: array<u32>;
@group(0) @binding(4) var<storage, read_write> renderCounter: array<atomic<u32>>;
struct BinWeight {
base: f32,
divider: f32
}
@group(0) @binding(5) var<storage, read> binWeights: array<BinWeight>;
struct ProjectorUniforms {
splatTextureSize: u32,
numBins: u32,
isOrtho: u32,
pad0: u32,
viewProj: mat4x4f,
viewMatrix: mat4x4f,
cameraPosition: vec3f,
minPixelSize: f32,
cameraDirection: vec3f,
focal: f32,
viewportWidth: f32,
viewportHeight: f32,
nearClip: f32,
farClip: f32,
alphaClip: f32,
minContribution: f32,
minDist: f32,
invRange: f32,
fisheye_k: f32,
fisheye_inv_k: f32,
fisheye_projMat00: f32,
fisheye_projMat11: f32,
}
@group(0) @binding(6) var<uniform> uniforms: ProjectorUniforms;
var<workgroup> wgCount: atomic<u32>;
var<workgroup> wgBase: u32;
@compute @workgroup_size(256)
fn main(
@builtin(global_invocation_id) gid: vec3u,
@builtin(num_workgroups) numWorkgroups: vec3u,
@builtin(local_invocation_index) localIdx: u32
) {
if (localIdx == 0u) {
atomicStore(&wgCount, 0u);
}
workgroupBarrier();
let threadIdx = gid.y * (numWorkgroups.x * 256u) + gid.x;
let numVisible = sortElementCount[0];
var valid = false;
var clipPos: vec4f = vec4f(0.0);
var v1: vec2f = vec2f(0.0);
var v2: vec2f = vec2f(0.0);
var rgb: vec3f = vec3f(0.0);
var alpha: f32 = 0.0;
var pcId: u32 = 0u;
var sortKey: u32 = 0u;
let projected = projectSplatCommon(
threadIdx,
numVisible,
uniforms.alphaClip,
uniforms.minPixelSize,
uniforms.minContribution,
uniforms.viewMatrix,
uniforms.viewProj,
uniforms.focal,
uniforms.viewportWidth,
uniforms.viewportHeight,
uniforms.nearClip,
uniforms.farClip,
uniforms.isOrtho,
uniforms.fisheye_k,
uniforms.fisheye_inv_k,
uniforms.fisheye_projMat00,
uniforms.fisheye_projMat11,
);
if (projected.valid) {
let center = projected.center;
let opacity = projected.opacity;
let proj = projected.proj;
let mid = 0.5 * (proj.a + proj.c);
let radius = length(vec2f(0.5 * (proj.a - proj.c), proj.b));
let lambda1 = mid + radius;
let lambda2 = max(mid - radius, 0.1);
let vmin = min(1024.0, min(uniforms.viewportWidth, uniforms.viewportHeight));
let l1 = 2.0 * min(sqrt(2.0 * lambda1), vmin);
let l2 = 2.0 * min(sqrt(2.0 * lambda2), vmin);
let dir = normalize(vec2f(proj.b, lambda1 - proj.a));
v1 = l1 * dir;
v2 = l2 * vec2f(dir.y, -dir.x);
let viewCenter = uniforms.viewMatrix * vec4f(center, 1.0);
let neg_z = -viewCenter.z;
let ndcX = proj.screen.x / uniforms.viewportWidth * 2.0 - 1.0;
let ndcY = proj.screen.y / uniforms.viewportHeight * 2.0 - 1.0;
let depthNdc = clamp(
(neg_z - uniforms.nearClip) / (uniforms.farClip - uniforms.nearClip),
0.0, 1.0
);
clipPos = vec4f(ndcX, ndcY, depthNdc, 1.0);
clipPos = uniforms.viewProj * vec4f(center, 1.0);
clipPos.z = clamp(clipPos.z, 0.0, abs(clipPos.w));
let delta = center - uniforms.cameraPosition;
let radialDist = length(delta);
let dist = (1.0 / uniforms.invRange) - radialDist - uniforms.minDist;
let toSplat = center - uniforms.cameraPosition;
let dist = dot(toSplat, uniforms.cameraDirection) - uniforms.minDist;
let d = dist * uniforms.invRange * f32(uniforms.numBins);
let binFloat = clamp(d, 0.0, f32(uniforms.numBins) - 0.001);
let bin = u32(binFloat);
let binFrac = binFloat - f32(bin);
sortKey = u32(binWeights[bin].base + binWeights[bin].divider * binFrac);
pcId = loadPcId().r;
alpha = opacity;
let color = getColor();
rgb = max(color, vec3f(0.0));
alpha = opacity * proj.aaFactor;
alpha = opacity;
valid = true;
}
var localDst: u32 = 0u;
if (valid) {
localDst = atomicAdd(&wgCount, 1u);
}
workgroupBarrier();
if (localIdx == 0u) {
let total = atomicLoad(&wgCount);
wgBase = atomicAdd(&renderCounter[0], total);
}
workgroupBarrier();
if (valid) {
let dst = wgBase + localDst;
let base = dst * {CACHE_STRIDE}u;
projCache[base + 0u] = bitcast<u32>(clipPos.x);
projCache[base + 1u] = bitcast<u32>(clipPos.y);
projCache[base + 2u] = bitcast<u32>(clipPos.z);
projCache[base + 3u] = bitcast<u32>(clipPos.w);
projCache[base + 4u] = pack2x16float(v1);
projCache[base + 5u] = pack2x16float(v2);
projCache[base + 6u] = pcId;
projCache[base + 7u] = pack2x16float(vec2f(0.0, alpha));
projCache[base + 6u] = pack2x16float(vec2f(rgb.x, rgb.y));
projCache[base + 7u] = pack2x16float(vec2f(rgb.z, alpha));
sortKeys[dst] = sortKey;
}
}
`;
var compute_gsplat_projector_default = computeGsplatProjectorSource;
export {
computeGsplatProjectorSource,
compute_gsplat_projector_default as default
};