playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
269 lines (268 loc) • 10.5 kB
JavaScript
import { Texture } from "../../platform/graphics/texture.js";
import { StorageBuffer } from "../../platform/graphics/storage-buffer.js";
import { Compute } from "../../platform/graphics/compute.js";
import { Shader } from "../../platform/graphics/shader.js";
import { BindGroupFormat, BindStorageBufferFormat, BindStorageTextureFormat, BindTextureFormat, 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,
BUFFERUSAGE_INDIRECT,
FILTER_NEAREST,
PIXELFORMAT_R32U,
PIXELFORMAT_RGBA16F,
SAMPLETYPE_UNFILTERABLE_FLOAT,
SHADERLANGUAGE_WGSL,
SHADERSTAGE_COMPUTE,
TEXTUREDIMENSION_2D,
UNIFORMTYPE_FLOAT,
UNIFORMTYPE_UINT,
UNIFORMTYPE_VEC3
} from "../../platform/graphics/constants.js";
import { PrefixSumKernel } from "../graphics/prefix-sum-kernel.js";
import { shaderChunksWGSL } from "../shader-lib/wgsl/collections/shader-chunks-wgsl.js";
import { computeGsplatLocalRasterizeSource } from "../shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-rasterize.js";
import { CACHE_STRIDE } from "./constants.js";
const MAX_CHUNKS_PER_TILE = 8;
class GSplatLocalDispatchSet {
device;
pickMode;
// Count compute caching: standard and fisheye variants, created lazily
_countShader = null;
_countBindGroupFormat = null;
_countCompute = null;
_countShaderFisheye = null;
_countBindGroupFormatFisheye = null;
_countComputeFisheye = null;
placeEntriesCompute;
largeSplatCompute;
largePlaceEntriesCompute;
classifyCompute;
sortCompute;
bucketSortCompute;
copyCompute;
chunkSortCompute;
_rasterizeVariants = /* @__PURE__ */ new Map();
prefixSumKernel;
_tileSplatCountsBuffer = null;
_smallTileListBuffer = null;
_largeTileListBuffer = null;
_largeTileOverflowBasesBuffer = null;
_rasterizeTileListBuffer = null;
_tileListCountsBuffer = null;
_chunkRangesBuffer = null;
_totalChunksBuffer = null;
_chunkSortIndirectBuffer = null;
_allocatedTileCapacity = 0;
outputTexture = null;
pickIdTexture = null;
pickDepthTexture = null;
constructor(device, pickMode) {
this.device = device;
this.pickMode = pickMode;
this.prefixSumKernel = new PrefixSumKernel(device);
if (pickMode) {
this.pickIdTexture = new Texture(device, {
name: "GSplatLocalPickId",
width: 4,
height: 4,
format: PIXELFORMAT_R32U,
mipmaps: false,
minFilter: FILTER_NEAREST,
magFilter: FILTER_NEAREST,
storage: true
});
this.pickDepthTexture = new Texture(device, {
name: "GSplatLocalPickDepth",
width: 4,
height: 4,
format: PIXELFORMAT_RGBA16F,
mipmaps: false,
minFilter: FILTER_NEAREST,
magFilter: FILTER_NEAREST,
storage: true
});
} else {
this.outputTexture = new Texture(device, {
name: "GSplatLocalComputeOutput",
width: 4,
height: 4,
format: PIXELFORMAT_RGBA16F,
mipmaps: false,
minFilter: FILTER_NEAREST,
magFilter: FILTER_NEAREST,
storage: true
});
}
}
resizeOutputTextures(width, height) {
if (this.pickMode) {
this.pickIdTexture?.resize(width, height);
this.pickDepthTexture?.resize(width, height);
} else {
this.outputTexture?.resize(width, height);
}
}
ensureTileBuffers(numTiles) {
const requiredTileSlots = numTiles + 1;
if (requiredTileSlots <= this._allocatedTileCapacity) return;
this._tileSplatCountsBuffer?.destroy();
this._smallTileListBuffer?.destroy();
this._largeTileListBuffer?.destroy();
this._largeTileOverflowBasesBuffer?.destroy();
this._rasterizeTileListBuffer?.destroy();
this._tileListCountsBuffer?.destroy();
this._chunkRangesBuffer?.destroy();
this._totalChunksBuffer?.destroy();
this._chunkSortIndirectBuffer?.destroy();
this._allocatedTileCapacity = requiredTileSlots;
this._tileSplatCountsBuffer = new StorageBuffer(this.device, requiredTileSlots * 4, BUFFERUSAGE_COPY_DST | BUFFERUSAGE_COPY_SRC);
this._smallTileListBuffer = new StorageBuffer(this.device, numTiles * 4);
this._largeTileListBuffer = new StorageBuffer(this.device, numTiles * 4);
this._largeTileOverflowBasesBuffer = new StorageBuffer(this.device, numTiles * 4);
this._rasterizeTileListBuffer = new StorageBuffer(this.device, numTiles * 4);
this._tileListCountsBuffer = new StorageBuffer(this.device, 4 * 4, BUFFERUSAGE_COPY_DST | BUFFERUSAGE_COPY_SRC);
const maxChunks = numTiles * MAX_CHUNKS_PER_TILE;
this._chunkRangesBuffer = new StorageBuffer(this.device, maxChunks * 8);
this._totalChunksBuffer = new StorageBuffer(this.device, 1 * 4, BUFFERUSAGE_COPY_DST);
this._chunkSortIndirectBuffer = new StorageBuffer(this.device, 3 * 4, BUFFERUSAGE_COPY_DST | BUFFERUSAGE_INDIRECT);
this.prefixSumKernel.destroyPasses();
}
getCountCompute(fisheyeEnabled, createShaderAndFormat) {
if (fisheyeEnabled) {
if (!this._countComputeFisheye) {
const { shader, bindGroupFormat } = createShaderAndFormat(this.pickMode, true);
this._countShaderFisheye = shader;
this._countBindGroupFormatFisheye = bindGroupFormat;
const label = this.pickMode ? "GSplatPickTileCountFisheye" : "GSplatLocalTileCountFisheye";
this._countComputeFisheye = new Compute(this.device, shader, label);
}
return this._countComputeFisheye;
}
if (!this._countCompute) {
const { shader, bindGroupFormat } = createShaderAndFormat(this.pickMode, false);
this._countShader = shader;
this._countBindGroupFormat = bindGroupFormat;
const label = this.pickMode ? "GSplatPickTileCount" : "GSplatLocalTileCount";
this._countCompute = new Compute(this.device, shader, label);
}
return this._countCompute;
}
destroyCountResources() {
this._countShader?.destroy();
this._countBindGroupFormat?.destroy();
this._countCompute?.destroy();
this._countShaderFisheye?.destroy();
this._countBindGroupFormatFisheye?.destroy();
this._countComputeFisheye?.destroy();
this._countShader = null;
this._countBindGroupFormat = null;
this._countCompute = null;
this._countShaderFisheye = null;
this._countBindGroupFormatFisheye = null;
this._countComputeFisheye = null;
}
getRasterizeCompute(pickMode, depthTest, fogType = "none", heatmap = false) {
let key = pickMode ? "pick" : "color";
if (depthTest) key += "-depth";
if (fogType !== "none") key += `-fog-${fogType}`;
if (heatmap) key += "-heatmap";
let variant = this._rasterizeVariants.get(key);
if (!variant) {
const { shader, bindGroupFormat } = this._createRasterizeShaderAndFormat(pickMode, depthTest, fogType, heatmap);
const compute = new Compute(this.device, shader, `GSplatRasterize-${key}`);
variant = { shader, bindGroupFormat, compute };
this._rasterizeVariants.set(key, variant);
}
return variant.compute;
}
_createRasterizeShaderAndFormat(pickMode, depthTest = false, fogType = "none", heatmap = false) {
const device = this.device;
const hasFog = fogType !== "none";
const uniforms = [
new UniformFormat("screenWidth", UNIFORMTYPE_UINT),
new UniformFormat("screenHeight", UNIFORMTYPE_UINT),
new UniformFormat("numTilesX", UNIFORMTYPE_UINT),
new UniformFormat("nearClip", UNIFORMTYPE_FLOAT),
new UniformFormat("farClip", UNIFORMTYPE_FLOAT),
new UniformFormat("alphaClip", UNIFORMTYPE_FLOAT)
];
if (hasFog) {
uniforms.push(
new UniformFormat("fog_color", UNIFORMTYPE_VEC3),
new UniformFormat("fog_start", UNIFORMTYPE_FLOAT),
new UniformFormat("fog_end", UNIFORMTYPE_FLOAT),
new UniformFormat("fog_density", UNIFORMTYPE_FLOAT)
);
}
const ubf = new UniformBufferFormat(device, uniforms);
const sharedBindings = [
new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE),
new BindStorageBufferFormat("tileEntries", SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat("tileSplatCounts", SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat("projCache", SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat("rasterizeTileList", SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat("tileListCounts", SHADERSTAGE_COMPUTE, true),
new BindStorageBufferFormat("depthBuffer", SHADERSTAGE_COMPUTE, true)
];
const outputBindings = pickMode ? [
new BindStorageTextureFormat("pickIdTexture", PIXELFORMAT_R32U),
new BindStorageTextureFormat("pickDepthTexture", PIXELFORMAT_RGBA16F)
] : [
new BindStorageTextureFormat("outputTexture", PIXELFORMAT_RGBA16F)
];
const depthBindings = depthTest ? [
new BindTextureFormat("sceneDepthMap", SHADERSTAGE_COMPUTE, TEXTUREDIMENSION_2D, SAMPLETYPE_UNFILTERABLE_FLOAT, false)
] : [];
const bgf = new BindGroupFormat(device, [...sharedBindings, ...outputBindings, ...depthBindings]);
const cdefines = /* @__PURE__ */ new Map();
cdefines.set("{CACHE_STRIDE}", CACHE_STRIDE.toString());
if (pickMode) cdefines.set("PICK_MODE", "");
if (depthTest) cdefines.set("DEPTH_TEST", "");
if (heatmap) cdefines.set("HEATMAP_MODE", "");
cdefines.set("GAMMA", "SRGB");
cdefines.set("FOG", hasFog ? fogType.toUpperCase() : "NONE");
const cincludes = pickMode ? void 0 : /* @__PURE__ */ new Map([["decodePS", shaderChunksWGSL.decodePS]]);
if (hasFog && cincludes) {
cincludes.set("fogMathPS", shaderChunksWGSL.fogMathPS);
cincludes.set("gammaPS", shaderChunksWGSL.gammaPS);
}
let name = "GSplatLocalRasterize";
if (pickMode) name = "GSplatLocalRasterizePick";
else if (depthTest) name = "GSplatLocalRasterizeDepth";
const shader = new Shader(device, {
name,
shaderLanguage: SHADERLANGUAGE_WGSL,
cshader: computeGsplatLocalRasterizeSource,
cdefines: cdefines.size > 0 ? cdefines : void 0,
cincludes,
computeBindGroupFormat: bgf,
computeUniformBufferFormats: { uniforms: ubf }
});
return { shader, bindGroupFormat: bgf };
}
destroy() {
this.destroyCountResources();
for (const { shader, bindGroupFormat, compute } of this._rasterizeVariants.values()) {
compute.destroy();
shader.destroy();
bindGroupFormat.destroy();
}
this._tileSplatCountsBuffer?.destroy();
this._smallTileListBuffer?.destroy();
this._largeTileListBuffer?.destroy();
this._largeTileOverflowBasesBuffer?.destroy();
this._rasterizeTileListBuffer?.destroy();
this._tileListCountsBuffer?.destroy();
this._chunkRangesBuffer?.destroy();
this._totalChunksBuffer?.destroy();
this._chunkSortIndirectBuffer?.destroy();
this.prefixSumKernel.destroy();
this.outputTexture?.destroy();
this.pickIdTexture?.destroy();
this.pickDepthTexture?.destroy();
}
}
export {
GSplatLocalDispatchSet
};