playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
179 lines (176 loc) • 6.78 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 { PIXELFORMAT_R32U, PIXELFORMAT_R32F, ADDRESS_CLAMP_TO_EDGE, FILTER_NEAREST, FILTER_NEAREST_MIPMAP_NEAREST } 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 {
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: name,
width: size,
height: size,
format: format,
mipmaps: 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;
}
constructor(device){
super(device), this._currentIndices = null, this._numPasses = 0, this._internalSize = 0, this._keys0 = null, this._keys1 = null, this._indices0 = null, this._indices1 = null, this._prefixSums = null, this._sortRT0 = null, this._sortRT1 = null, this._prefixSumsRT = null, this._countPasses = [], this._reorderPasses = [], this._elementCount = 0, this._keysTexture = null;
}
}
export { FramePassRadixSort };