playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
452 lines (451 loc) • 19.1 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";
import gsplatModifyDefaultSource from "../shader-lib/wgsl/chunks/gsplat/vert/gsplatModify.js";
import gsplatHelpersSource from "../shader-lib/wgsl/chunks/gsplat/vert/gsplatHelpers.js";
const INDEX_COUNT = 6 * GSplatResourceBase.instanceSize;
const PROJECTOR_WORKGROUP_SIZE = 256;
const PROJECTOR_INTERNAL_DEFINES = /* @__PURE__ */ new Set([
"{CACHE_STRIDE}",
"RADIAL_SORT",
"PICK_MODE",
"GSPLAT_FISHEYE",
"GSPLAT_AA",
"GSPLAT_COLOR_FLOAT",
"GSPLAT_XR"
]);
const _cameraDir = new Vec3();
const _dispatchSize = new Vec2();
const _viewProjMat = new Mat4();
const _viewProjData = new Float32Array(16);
const _viewProj1Data = 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;
_projectorUniformBufferFormatStereo = null;
_writeIndirectArgsCompute = null;
_writeArgsBindGroupFormat = null;
_writeArgsUniformBufferFormat = null;
_formatVersion = -1;
_materialKey = "";
_userModifySource = null;
_userVaryingsSource = null;
_userCacheWriteSource = null;
_userDefines = null;
_userCacheWords = 0;
_allocatedCacheCount = 0;
_allocatedCacheStride = 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._projectorUniformBufferFormatStereo = 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),
new UniformFormat("foveationStrength", UNIFORMTYPE_FLOAT),
new UniformFormat("foveationCenter", 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._projectorUniformBufferFormatStereo = new UniformBufferFormat(device, [
...baseFields.map((f) => new UniformFormat(f.name, f.type)),
new UniformFormat("viewProj1", UNIFORMTYPE_MAT4)
]);
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, stereo) {
return `${radialSort ? "r" : "l"}${pickMode ? "p" : ""}${fisheyeMode ? "f" : ""}${antiAlias ? "a" : ""}${stereo ? "s" : ""}`;
}
_createProjectorCompute(workBuffer, radialSort, pickMode, fisheyeMode, antiAlias, stereo) {
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("gsplatHelpersVS", gsplatHelpersSource);
cincludes.set("gsplatModifyVS", this._userModifySource ?? gsplatModifyDefaultSource);
cincludes.set("gsplatUserVaryingsCS", this._userVaryingsSource ?? "");
cincludes.set("gsplatUserCacheWriteCS", this._userCacheWriteSource ?? "");
cincludes.set("gsplatProjectCommonCS", computeGsplatProjectCommonSource);
const cdefines = /* @__PURE__ */ new Map();
cdefines.set("{CACHE_STRIDE}", (CACHE_STRIDE + this._userCacheWords).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", "");
}
if (stereo) {
cdefines.set("GSPLAT_XR", "");
}
const colorStream = wbFormat.getStream("dataColor");
if (colorStream && colorStream.format !== PIXELFORMAT_RGBA16U) {
cdefines.set("GSPLAT_COLOR_FLOAT", "");
}
if (this._userDefines) {
this._userDefines.forEach((value, key) => {
if (!PROJECTOR_INTERNAL_DEFINES.has(key)) {
cdefines.set(key, value);
}
});
}
const name = `GSplatProjector${radialSort ? "Radial" : "Linear"}${pickMode ? "Pick" : ""}${fisheyeMode ? "Fisheye" : ""}${antiAlias ? "Aa" : ""}${stereo ? "Stereo" : ""}`;
const ubFormat = stereo ? this._projectorUniformBufferFormatStereo : 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, stereo = false) {
const wbFormat = workBuffer.format;
if (this._formatVersion !== wbFormat.extraStreamsVersion) {
this._destroyProjectorComputes();
this._formatVersion = wbFormat.extraStreamsVersion;
}
const key = this._projectorKey(radialSort, pickMode, fisheyeMode, antiAlias, stereo);
let compute = this._projectorComputes.get(key);
if (!compute) {
compute = this._createProjectorCompute(workBuffer, radialSort, pickMode, fisheyeMode, antiAlias, stereo);
this._projectorComputes.set(key, compute);
}
return compute;
}
_updateMaterial(material, userCacheWords = 0) {
const chunksKey = material?.shaderChunks?.key ?? "";
const definesKey = material?.definesKey ?? "";
const materialKey = `${chunksKey}|${definesKey}|${userCacheWords}`;
if (materialKey !== this._materialKey) {
this._materialKey = materialKey;
this._userDefines = material?.defines ?? null;
this._userCacheWords = userCacheWords;
const wgslChunks = material?.getShaderChunks?.(SHADERLANGUAGE_WGSL);
this._userModifySource = wgslChunks?.get("gsplatModifyVS") ?? null;
this._userVaryingsSource = wgslChunks?.get("gsplatUserVaryingsCS") ?? null;
this._userCacheWriteSource = wgslChunks?.get("gsplatUserCacheWriteCS") ?? null;
this._destroyProjectorComputes();
}
}
_ensureCapacity(capacity) {
const cacheStride = CACHE_STRIDE + this._userCacheWords;
if (capacity > this._allocatedCacheCount || cacheStride !== this._allocatedCacheStride) {
this.projCache?.destroy();
this.sortKeys?.destroy();
this._allocatedCacheCount = capacity;
this._allocatedCacheStride = cacheStride;
this.projCache = new StorageBuffer(this.device, capacity * cacheStride * 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,
foveationStrength = 0,
foveationCenter = 0.3,
viewportWidth,
viewportHeight,
flipY,
pickMode = false,
fisheyeProj,
antiAlias = false,
isStereo = false,
material,
userCacheWords = 0
} = params;
const fisheyeMode = !!fisheyeProj?.enabled;
const stereoMode = !!isStereo && !pickMode && !fisheyeMode;
const aaMode = antiAlias && !pickMode;
this._updateMaterial(material, userCacheWords);
this._ensureCapacity(totalCapacity);
this.renderCounter.clear();
const compute = this._getProjectorCompute(workBuffer, radialSort, pickMode, fisheyeMode, aaMode, stereoMode);
if (material) {
const srcParams = material.parameters;
for (const name in srcParams) {
if (srcParams.hasOwnProperty(name)) {
compute.setParameter(name, srcParams[name].data);
}
}
}
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 webgpu = this.device.isWebGPU;
let focal;
if (stereoMode) {
const views = cam.xrViews;
cam.updateViewTransforms();
_viewProjData.set(views[0].projViewOffMat.data);
_viewProj1Data.set(views[1].projViewOffMat.data);
_viewData.set(views[0].viewOffMat.data);
focal = viewportWidth * views[0].projMat.data[0];
} else {
const view = cam.viewMatrix;
_viewProjMat.mul2(Camera.applyShaderProjectionTransform(cam.projectionMatrix, _shaderProjMat, flipY, webgpu), view);
_viewProjData.set(_viewProjMat.data);
_viewData.set(view.data);
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);
if (stereoMode) {
compute.setParameter("viewProj1", _viewProj1Data);
}
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("foveationStrength", foveationStrength);
compute.setParameter("foveationCenter", foveationCenter);
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
};