playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
364 lines (363 loc) • 15.6 kB
JavaScript
import { Mat4 } from "../../core/math/mat4.js";
import { Vec2 } from "../../core/math/vec2.js";
import { Vec3 } from "../../core/math/vec3.js";
import { Compute } from "../../platform/graphics/compute.js";
import { Shader } from "../../platform/graphics/shader.js";
import { StorageBuffer } from "../../platform/graphics/storage-buffer.js";
import {
BindGroupFormat,
BindStorageBufferFormat,
BindUniformBufferFormat
} from "../../platform/graphics/bind-group-format.js";
import {
UniformBufferFormat,
UniformFormat
} from "../../platform/graphics/uniform-buffer-format.js";
import {
BUFFERUSAGE_COPY_DST,
BUFFERUSAGE_COPY_SRC,
PIXELFORMAT_RGBA16U,
SHADERLANGUAGE_WGSL,
SHADERSTAGE_COMPUTE,
UNIFORMTYPE_FLOAT,
UNIFORMTYPE_MAT4,
UNIFORMTYPE_UINT,
UNIFORMTYPE_UVEC4,
UNIFORMTYPE_VEC3
} from "../../platform/graphics/constants.js";
import { PROJECTION_ORTHOGRAPHIC } from "../constants.js";
import { Camera } from "../camera.js";
import { GSplatResourceBase } from "../gsplat/gsplat-resource-base.js";
import { GSplatSortBinWeights } from "./gsplat-sort-bin-weights.js";
import { CACHE_STRIDE } from "./gsplat-projector-constants.js";
import { computeGsplatProjectorSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-projector.js";
import { computeGsplatProjectorWriteIndirectArgsSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-projector-write-indirect-args.js";
import { computeGsplatProjectCommonSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-project-common.js";
import { computeGsplatCommonSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-common.js";
import { computeGsplatTileIntersectSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-tile-intersect.js";
import computeSplatSource from "../shader-lib/wgsl/chunks/gsplat/vert/gsplatComputeSplat.js";
const INDEX_COUNT = 6 * GSplatResourceBase.instanceSize;
const PROJECTOR_WORKGROUP_SIZE = 256;
const _cameraDir = new Vec3();
const _dispatchSize = new Vec2();
const _viewProjMat = new Mat4();
const _viewProjData = new Float32Array(16);
const _viewData = new Float32Array(16);
const _shaderProjMat = new Mat4();
class GSplatProjector {
device;
projCache = null;
sortKeys = null;
renderCounter = null;
binWeightsBuffer = null;
binWeightsUtil;
_projectorComputes = /* @__PURE__ */ new Map();
_projectorBindGroupFormat = null;
_projectorUniformBufferFormat = null;
_projectorUniformBufferFormatFisheye = null;
_writeIndirectArgsCompute = null;
_writeArgsBindGroupFormat = null;
_writeArgsUniformBufferFormat = null;
_formatVersion = -1;
_allocatedCacheCount = 0;
cameraPositionData = new Float32Array(3);
cameraDirectionData = new Float32Array(3);
constructor(device) {
this.device = device;
this.binWeightsUtil = new GSplatSortBinWeights();
this.binWeightsBuffer = new StorageBuffer(
device,
GSplatSortBinWeights.NUM_BINS * 2 * 4,
BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST
);
this.renderCounter = new StorageBuffer(device, 4, BUFFERUSAGE_COPY_SRC | BUFFERUSAGE_COPY_DST);
this._createUniformBufferFormats();
this._createWriteIndirectArgsCompute();
}
destroy() {
this.projCache?.destroy();
this.sortKeys?.destroy();
this.renderCounter?.destroy();
this.binWeightsBuffer?.destroy();
for (const compute of this._projectorComputes.values()) {
compute.shader?.destroy();
}
this._projectorComputes.clear();
this._projectorBindGroupFormat?.destroy();
this._writeIndirectArgsCompute?.shader?.destroy();
this._writeArgsBindGroupFormat?.destroy();
this.projCache = null;
this.sortKeys = null;
this.renderCounter = null;
this.binWeightsBuffer = null;
this._projectorBindGroupFormat = null;
this._projectorUniformBufferFormat = null;
this._projectorUniformBufferFormatFisheye = null;
this._writeIndirectArgsCompute = null;
this._writeArgsBindGroupFormat = null;
this._writeArgsUniformBufferFormat = null;
}
_createUniformBufferFormats() {
const device = this.device;
const baseFields = [
new UniformFormat("splatTextureSize", UNIFORMTYPE_UINT),
new UniformFormat("numBins", UNIFORMTYPE_UINT),
new UniformFormat("isOrtho", UNIFORMTYPE_UINT),
new UniformFormat("pad0", UNIFORMTYPE_UINT),
new UniformFormat("viewProj", UNIFORMTYPE_MAT4),
new UniformFormat("viewMatrix", UNIFORMTYPE_MAT4),
new UniformFormat("cameraPosition", UNIFORMTYPE_VEC3),
new UniformFormat("minPixelSize", UNIFORMTYPE_FLOAT),
new UniformFormat("cameraDirection", UNIFORMTYPE_VEC3),
new UniformFormat("focal", UNIFORMTYPE_FLOAT),
new UniformFormat("viewportWidth", UNIFORMTYPE_FLOAT),
new UniformFormat("viewportHeight", UNIFORMTYPE_FLOAT),
new UniformFormat("nearClip", UNIFORMTYPE_FLOAT),
new UniformFormat("farClip", UNIFORMTYPE_FLOAT),
new UniformFormat("alphaClip", UNIFORMTYPE_FLOAT),
new UniformFormat("minContribution", UNIFORMTYPE_FLOAT),
new UniformFormat("minDist", UNIFORMTYPE_FLOAT),
new UniformFormat("invRange", UNIFORMTYPE_FLOAT)
];
this._projectorUniformBufferFormat = new UniformBufferFormat(device, baseFields);
this._projectorUniformBufferFormatFisheye = new UniformBufferFormat(device, [
...baseFields.map((f) => new UniformFormat(f.name, f.type)),
new UniformFormat("fisheye_k", UNIFORMTYPE_FLOAT),
new UniformFormat("fisheye_inv_k", UNIFORMTYPE_FLOAT),
new UniformFormat("fisheye_projMat00", UNIFORMTYPE_FLOAT),
new UniformFormat("fisheye_projMat11", UNIFORMTYPE_FLOAT)
]);
this._writeArgsUniformBufferFormat = new UniformBufferFormat(device, [
new UniformFormat("drawSlot", UNIFORMTYPE_UINT),
new UniformFormat("indexCount", UNIFORMTYPE_UINT),
new UniformFormat("sortSlotBase", UNIFORMTYPE_UINT),
new UniformFormat("pad0", UNIFORMTYPE_UINT),
new UniformFormat("sortIndirectInfo", UNIFORMTYPE_UVEC4)
]);
}
_createWriteIndirectArgsCompute() {
const device = this.device;
this._writeArgsBindGroupFormat = new BindGroupFormat(device, [
new BindStorageBufferFormat("renderCounter", SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat("indirectDrawArgs", SHADERSTAGE_COMPUTE, false),
new BindStorageBufferFormat("numSplatsBuf", SHADERSTAGE_COMPUTE, false),
new BindStorageBufferFormat("indirectDispatchArgs", SHADERSTAGE_COMPUTE, false),
new BindStorageBufferFormat("sortElementCountBuf", SHADERSTAGE_COMPUTE, false),
new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE)
]);
const cdefines = /* @__PURE__ */ new Map([
["{INSTANCE_SIZE}", GSplatResourceBase.instanceSize.toString()]
]);
const shader = new Shader(device, {
name: "GSplatProjectorWriteIndirectArgs",
shaderLanguage: SHADERLANGUAGE_WGSL,
cshader: computeGsplatProjectorWriteIndirectArgsSource,
cdefines,
computeBindGroupFormat: this._writeArgsBindGroupFormat,
computeUniformBufferFormats: { uniforms: this._writeArgsUniformBufferFormat }
});
this._writeIndirectArgsCompute = new Compute(device, shader, "GSplatProjectorWriteIndirectArgs");
}
_destroyProjectorComputes() {
for (const compute of this._projectorComputes.values()) {
compute.shader?.destroy();
}
this._projectorComputes.clear();
this._projectorBindGroupFormat?.destroy();
this._projectorBindGroupFormat = null;
}
_projectorKey(radialSort, pickMode, fisheyeMode, antiAlias) {
return `${radialSort ? "r" : "l"}${pickMode ? "p" : ""}${fisheyeMode ? "f" : ""}${antiAlias ? "a" : ""}`;
}
_createProjectorCompute(workBuffer, radialSort, pickMode, fisheyeMode, antiAlias) {
const device = this.device;
const wbFormat = workBuffer.format;
const fixedBindings = [
new BindStorageBufferFormat("compactedSplatIds", SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat("sortElementCount", SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat("projCache", SHADERSTAGE_COMPUTE),
new BindStorageBufferFormat("sortKeys", SHADERSTAGE_COMPUTE),
new BindStorageBufferFormat("renderCounter", SHADERSTAGE_COMPUTE),
new BindStorageBufferFormat("binWeights", SHADERSTAGE_COMPUTE, true),
new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE)
];
if (!this._projectorBindGroupFormat) {
this._projectorBindGroupFormat = new BindGroupFormat(device, [
...fixedBindings,
...wbFormat.getComputeBindFormats()
]);
}
const cincludes = /* @__PURE__ */ new Map();
cincludes.set("gsplatCommonCS", computeGsplatCommonSource);
cincludes.set("gsplatTileIntersectCS", computeGsplatTileIntersectSource);
cincludes.set("gsplatComputeSplatCS", computeSplatSource);
cincludes.set("gsplatFormatDeclCS", wbFormat.getComputeInputDeclarations(fixedBindings.length));
cincludes.set("gsplatFormatReadCS", wbFormat.getReadCode());
cincludes.set("gsplatProjectCommonCS", computeGsplatProjectCommonSource);
const cdefines = /* @__PURE__ */ new Map();
cdefines.set("{CACHE_STRIDE}", CACHE_STRIDE.toString());
if (radialSort) {
cdefines.set("RADIAL_SORT", "");
}
if (pickMode) {
cdefines.set("PICK_MODE", "");
}
if (fisheyeMode) {
cdefines.set("GSPLAT_FISHEYE", "");
}
if (antiAlias) {
cdefines.set("GSPLAT_AA", "");
}
const colorStream = wbFormat.getStream("dataColor");
if (colorStream && colorStream.format !== PIXELFORMAT_RGBA16U) {
cdefines.set("GSPLAT_COLOR_FLOAT", "");
}
const name = `GSplatProjector${radialSort ? "Radial" : "Linear"}${pickMode ? "Pick" : ""}${fisheyeMode ? "Fisheye" : ""}${antiAlias ? "Aa" : ""}`;
const ubFormat = fisheyeMode ? this._projectorUniformBufferFormatFisheye : this._projectorUniformBufferFormat;
const shader = new Shader(device, {
name,
shaderLanguage: SHADERLANGUAGE_WGSL,
cshader: computeGsplatProjectorSource,
cincludes,
cdefines,
computeBindGroupFormat: this._projectorBindGroupFormat,
computeUniformBufferFormats: { uniforms: ubFormat }
});
return new Compute(device, shader, name);
}
_getProjectorCompute(workBuffer, radialSort, pickMode = false, fisheyeMode = false, antiAlias = false) {
const wbFormat = workBuffer.format;
if (this._formatVersion !== wbFormat.extraStreamsVersion) {
this._destroyProjectorComputes();
this._formatVersion = wbFormat.extraStreamsVersion;
}
const key = this._projectorKey(radialSort, pickMode, fisheyeMode, antiAlias);
let compute = this._projectorComputes.get(key);
if (!compute) {
compute = this._createProjectorCompute(workBuffer, radialSort, pickMode, fisheyeMode, antiAlias);
this._projectorComputes.set(key, compute);
}
return compute;
}
_ensureCapacity(capacity) {
if (capacity > this._allocatedCacheCount) {
this.projCache?.destroy();
this.sortKeys?.destroy();
this._allocatedCacheCount = capacity;
this.projCache = new StorageBuffer(this.device, capacity * CACHE_STRIDE * 4);
this.sortKeys = new StorageBuffer(this.device, capacity * 4, BUFFERUSAGE_COPY_SRC);
}
}
dispatch(params) {
const {
workBuffer,
cameraNode,
compactedSplatIds,
sortElementCountBuffer,
totalCapacity,
radialSort,
numBits,
minDist,
maxDist,
alphaClip,
minPixelSize,
minContribution,
viewportWidth,
viewportHeight,
flipY,
pickMode = false,
fisheyeProj,
antiAlias = false
} = params;
const fisheyeMode = !!fisheyeProj?.enabled;
const aaMode = antiAlias && !pickMode;
this._ensureCapacity(totalCapacity);
this.renderCounter.clear();
const compute = this._getProjectorCompute(workBuffer, radialSort, pickMode, fisheyeMode, aaMode);
const cameraPos = cameraNode.getPosition();
const cameraMat = cameraNode.getWorldTransform();
const cameraDir = cameraMat.getZ(_cameraDir).normalize();
const range = maxDist - minDist;
const invRange = range > 0 ? 1 / range : 1;
const bucketCount = 1 << numBits;
const cameraBin = GSplatSortBinWeights.computeCameraBin(radialSort, minDist, range);
const binWeights = this.binWeightsUtil.compute(cameraBin, bucketCount);
this.binWeightsBuffer.write(0, binWeights);
compute.setParameter("compactedSplatIds", compactedSplatIds);
compute.setParameter("sortElementCount", sortElementCountBuffer);
compute.setParameter("projCache", this.projCache);
compute.setParameter("sortKeys", this.sortKeys);
compute.setParameter("renderCounter", this.renderCounter);
compute.setParameter("binWeights", this.binWeightsBuffer);
for (const stream of workBuffer.format.resourceStreams) {
const texture = workBuffer.getTexture(stream.name);
if (texture) {
compute.setParameter(stream.name, texture);
}
}
const cameraComponent = cameraNode.camera;
const cam = cameraComponent.camera;
const view = cam.viewMatrix;
const webgpu = this.device.isWebGPU;
_viewProjMat.mul2(Camera.applyShaderProjectionTransform(cam.projectionMatrix, _shaderProjMat, flipY, webgpu), view);
_viewProjData.set(_viewProjMat.data);
_viewData.set(view.data);
const focal = viewportWidth * _shaderProjMat.data[0];
this.cameraPositionData[0] = cameraPos.x;
this.cameraPositionData[1] = cameraPos.y;
this.cameraPositionData[2] = cameraPos.z;
compute.setParameter("cameraPosition", this.cameraPositionData);
this.cameraDirectionData[0] = cameraDir.x;
this.cameraDirectionData[1] = cameraDir.y;
this.cameraDirectionData[2] = cameraDir.z;
compute.setParameter("cameraDirection", this.cameraDirectionData);
compute.setParameter("viewMatrix", _viewData);
compute.setParameter("viewProj", _viewProjData);
compute.setParameter("focal", focal);
compute.setParameter("viewportWidth", viewportWidth);
compute.setParameter("viewportHeight", viewportHeight);
compute.setParameter("nearClip", cam.nearClip);
compute.setParameter("farClip", cam.farClip);
compute.setParameter("alphaClip", alphaClip);
compute.setParameter("minPixelSize", minPixelSize);
compute.setParameter("minContribution", minContribution);
compute.setParameter("isOrtho", cam.projection === PROJECTION_ORTHOGRAPHIC ? 1 : 0);
compute.setParameter("splatTextureSize", workBuffer.textureSize);
compute.setParameter("numBins", GSplatSortBinWeights.NUM_BINS);
compute.setParameter("minDist", minDist);
compute.setParameter("invRange", invRange);
compute.setParameter("pad0", 0);
if (fisheyeMode) {
compute.setParameter("fisheye_k", fisheyeProj.k);
compute.setParameter("fisheye_inv_k", fisheyeProj.invK);
compute.setParameter("fisheye_projMat00", fisheyeProj.projMat00);
compute.setParameter("fisheye_projMat11", fisheyeProj.projMat11);
}
const workgroupCount = Math.ceil(totalCapacity / PROJECTOR_WORKGROUP_SIZE);
Compute.calcDispatchSize(
workgroupCount,
_dispatchSize,
this.device.limits.maxComputeWorkgroupsPerDimension || 65535
);
compute.setupDispatch(_dispatchSize.x, _dispatchSize.y, 1);
this.device.computeDispatch([compute], "GSplatProjector");
}
writeIndirectArgs(drawSlot, sortSlotBase, numSplatsBuffer, sortElementCountBuffer, sortIndirectInfo) {
const compute = this._writeIndirectArgsCompute;
compute.setParameter("renderCounter", this.renderCounter);
compute.setParameter("indirectDrawArgs", this.device.indirectDrawBuffer);
compute.setParameter("numSplatsBuf", numSplatsBuffer);
compute.setParameter("indirectDispatchArgs", this.device.indirectDispatchBuffer);
compute.setParameter("sortElementCountBuf", sortElementCountBuffer);
compute.setParameter("drawSlot", drawSlot);
compute.setParameter("indexCount", INDEX_COUNT);
compute.setParameter("sortSlotBase", sortSlotBase);
compute.setParameter("pad0", 0);
compute.setParameter("sortIndirectInfo", sortIndirectInfo);
compute.setupDispatch(1);
this.device.computeDispatch([compute], "GSplatProjectorWriteIndirectArgs");
}
}
export {
GSplatProjector
};