playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
195 lines (194 loc) • 6.35 kB
JavaScript
import { Color } from "../../core/math/color.js";
import { Texture } from "../../platform/graphics/texture.js";
import { RenderTarget } from "../../platform/graphics/render-target.js";
import { FramePass } from "../../platform/graphics/frame-pass.js";
import {
ADDRESS_CLAMP_TO_EDGE,
FILTER_NEAREST,
FILTER_NEAREST_MIPMAP_NEAREST,
PIXELFORMAT_R32F,
PIXELFORMAT_R32U
} from "../../platform/graphics/constants.js";
import { RenderPassRadixSortCount } from "./render-pass-radix-sort-count.js";
import { RenderPassRadixSortReorder } from "./render-pass-radix-sort-reorder.js";
const BITS_PER_STEP = 4;
const GROUP_SIZE = 4;
class FramePassRadixSort extends FramePass {
_currentIndices = null;
_numPasses = 0;
_internalSize = 0;
_keys0 = null;
_keys1 = null;
_indices0 = null;
_indices1 = null;
_prefixSums = null;
_sortRT0 = null;
_sortRT1 = null;
_prefixSumsRT = null;
_countPasses = [];
_reorderPasses = [];
_elementCount = 0;
_keysTexture = null;
// eslint-disable-next-line no-useless-constructor
constructor(device) {
super(device);
}
destroy() {
this._destroyPasses();
this._destroyInternalTextures();
super.destroy();
}
get sortedIndices() {
return this._currentIndices;
}
setup(keysTexture, elementCount, numBits = 16) {
this._keysTexture = keysTexture;
this._elementCount = elementCount;
const numPasses = Math.ceil(numBits / BITS_PER_STEP);
if (numPasses !== this._numPasses) {
this._destroyPasses();
this._numPasses = numPasses;
}
const requiredSize = this._calculateInternalSize(elementCount);
if (requiredSize !== this._internalSize) {
this._destroyPasses();
this._resizeInternalTextures(requiredSize);
this._internalSize = requiredSize;
}
if (this._countPasses.length === 0) {
this._createPasses();
}
}
_calculateInternalSize(elementCount) {
const side = Math.ceil(Math.sqrt(elementCount));
return Math.pow(2, Math.ceil(Math.log2(side)));
}
_resizeInternalTextures(size) {
this._destroyInternalTextures();
this._keys0 = this._createTexture("RadixSortKeys0", size, PIXELFORMAT_R32U);
this._keys1 = this._createTexture("RadixSortKeys1", size, PIXELFORMAT_R32U);
this._indices0 = this._createTexture("RadixSortIndices0", size, PIXELFORMAT_R32U);
this._indices1 = this._createTexture("RadixSortIndices1", size, PIXELFORMAT_R32U);
const prefixSize = size * Math.pow(2, BITS_PER_STEP / 2) / Math.pow(2, GROUP_SIZE / 2);
this._prefixSums = this._createTexture("RadixSortPrefixSums", prefixSize, PIXELFORMAT_R32F, true);
this._sortRT0 = new RenderTarget({
name: "RadixSortRT0",
colorBuffers: [this._keys0, this._indices0],
depth: false
});
this._sortRT1 = new RenderTarget({
name: "RadixSortRT1",
colorBuffers: [this._keys1, this._indices1],
depth: false
});
this._prefixSumsRT = new RenderTarget({
name: "RadixSortPrefixSumsRT",
colorBuffer: this._prefixSums,
depth: false
});
}
_createTexture(name, size, format, mipmaps = false) {
return new Texture(this.device, {
name,
width: size,
height: size,
format,
mipmaps,
minFilter: mipmaps ? FILTER_NEAREST_MIPMAP_NEAREST : FILTER_NEAREST,
magFilter: FILTER_NEAREST,
addressU: ADDRESS_CLAMP_TO_EDGE,
addressV: ADDRESS_CLAMP_TO_EDGE
});
}
_destroyInternalTextures() {
this._sortRT0?.destroy();
this._sortRT1?.destroy();
this._prefixSumsRT?.destroy();
this._keys0?.destroy();
this._keys1?.destroy();
this._indices0?.destroy();
this._indices1?.destroy();
this._prefixSums?.destroy();
this._sortRT0 = null;
this._sortRT1 = null;
this._prefixSumsRT = null;
this._keys0 = null;
this._keys1 = null;
this._indices0 = null;
this._indices1 = null;
this._prefixSums = null;
}
_createPasses() {
const device = this.device;
const numPasses = this._numPasses;
let nextRT = this._sortRT1;
for (let i = 0; i < numPasses; i++) {
const sourceLinear = i === 0;
const outputLinear = i === numPasses - 1;
const currentBit = i * BITS_PER_STEP;
const countPass = new RenderPassRadixSortCount(device, sourceLinear, BITS_PER_STEP, GROUP_SIZE, currentBit);
countPass.init(this._prefixSumsRT);
countPass.setClearColor(new Color(0, 0, 0, 0));
this._countPasses.push(countPass);
this.beforePasses.push(countPass);
const reorderPass = new RenderPassRadixSortReorder(device, sourceLinear, outputLinear, BITS_PER_STEP, GROUP_SIZE, currentBit);
reorderPass.setPrefixSumsTexture(this._prefixSums);
reorderPass.init(nextRT);
this._reorderPasses.push(reorderPass);
this.beforePasses.push(reorderPass);
nextRT = nextRT === this._sortRT1 ? this._sortRT0 : this._sortRT1;
}
this._currentIndices = numPasses % 2 === 1 ? this._indices1 : this._indices0;
}
_destroyPasses() {
for (const pass of this.beforePasses) {
pass.destroy();
}
this.beforePasses.length = 0;
this._countPasses.length = 0;
this._reorderPasses.length = 0;
}
frameUpdate() {
super.frameUpdate();
if (!this._keysTexture || this._countPasses.length === 0) {
return;
}
const numPasses = this._countPasses.length;
const elementCount = this._elementCount;
const imageElementsLog2 = Math.log2(this._internalSize * this._internalSize);
const imageSize = this._internalSize;
let currentKeys = this._keys0;
let currentIndices = this._indices0;
for (let i = 0; i < numPasses; i++) {
const sourceLinear = i === 0;
const countPass = this._countPasses[i];
const reorderPass = this._reorderPasses[i];
if (sourceLinear) {
countPass.setKeysTexture(this._keysTexture);
} else {
countPass.setKeysTexture(currentKeys);
}
countPass.setDynamicParams(elementCount, imageElementsLog2);
if (sourceLinear) {
reorderPass.setKeysTexture(this._keysTexture);
} else {
reorderPass.setKeysTexture(currentKeys);
reorderPass.setIndicesTexture(currentIndices);
}
reorderPass.setDynamicParams(elementCount, imageElementsLog2, imageSize);
currentKeys = currentKeys === this._keys0 ? this._keys1 : this._keys0;
currentIndices = currentIndices === this._indices0 ? this._indices1 : this._indices0;
}
}
sort(keysTexture, elementCount, numBits = 16) {
this.setup(keysTexture, elementCount, numBits);
this.frameUpdate();
for (const pass of this.beforePasses) {
pass.render();
}
return this.sortedIndices;
}
}
export {
FramePassRadixSort
};