playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
340 lines (339 loc) • 11.5 kB
JavaScript
var __defProp = Object.defineProperty;
var __defNormalProp = (obj, key, value) => key in obj ? __defProp(obj, key, { enumerable: true, configurable: true, writable: true, value }) : obj[key] = value;
var __publicField = (obj, key, value) => __defNormalProp(obj, typeof key !== "symbol" ? key + "" : key, value);
import { Debug } from "../../core/debug.js";
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 {
/**
* Creates a new FramePassRadixSort instance.
*
* @param {GraphicsDevice} device - The graphics device.
*/
// eslint-disable-next-line no-useless-constructor
constructor(device) {
super(device);
/**
* The current sorted indices texture (R32U). Access sorted indices using Morton lookup.
*
* @type {Texture|null}
*/
__publicField(this, "_currentIndices", null);
/**
* Current number of radix passes.
*/
__publicField(this, "_numPasses", 0);
/**
* Current internal texture size (power of 2).
*/
__publicField(this, "_internalSize", 0);
/**
* Internal keys texture 0 (ping-pong buffer).
*
* @type {Texture|null}
*/
__publicField(this, "_keys0", null);
/**
* Internal keys texture 1 (ping-pong buffer).
*
* @type {Texture|null}
*/
__publicField(this, "_keys1", null);
/**
* Internal indices texture 0 (ping-pong buffer).
*
* @type {Texture|null}
*/
__publicField(this, "_indices0", null);
/**
* Internal indices texture 1 (ping-pong buffer).
*
* @type {Texture|null}
*/
__publicField(this, "_indices1", null);
/**
* Prefix sums texture (R32F with mipmaps).
*
* @type {Texture|null}
*/
__publicField(this, "_prefixSums", null);
/**
* Sort render target 0 (MRT for keys + indices).
*
* @type {RenderTarget|null}
*/
__publicField(this, "_sortRT0", null);
/**
* Sort render target 1 (MRT for keys + indices).
*
* @type {RenderTarget|null}
*/
__publicField(this, "_sortRT1", null);
/**
* Prefix sums render target.
*
* @type {RenderTarget|null}
*/
__publicField(this, "_prefixSumsRT", null);
/**
* Count passes for each radix iteration.
*
* @type {RenderPassRadixSortCount[]}
*/
__publicField(this, "_countPasses", []);
/**
* Reorder passes for each radix iteration.
*
* @type {RenderPassRadixSortReorder[]}
*/
__publicField(this, "_reorderPasses", []);
/**
* Number of elements to sort (set by setup()).
*/
__publicField(this, "_elementCount", 0);
/**
* The source keys texture (set by setup()).
*
* @type {Texture|null}
*/
__publicField(this, "_keysTexture", null);
}
destroy() {
this._destroyPasses();
this._destroyInternalTextures();
super.destroy();
}
/**
* Gets the sorted indices texture (R32U, linear layout). Use `.width` for texture dimensions.
* Access with: `texelFetch(texture, ivec2(index % width, index / width), 0).r`
*
* @type {Texture|null}
*/
get sortedIndices() {
return this._currentIndices;
}
/**
* Sets up the sort for the current frame.
*
* Note: The source keys texture is read-only and can be any size.
* The sorted indices will be in a separate power-of-2 texture.
*
* @param {Texture} keysTexture - R32U texture containing sort keys (linear layout, any size).
* @param {number} elementCount - Number of elements to sort.
* @param {number} [numBits] - Number of bits to sort (1-24). More bits = more passes.
*/
setup(keysTexture, elementCount, numBits = 16) {
Debug.assert(keysTexture, "FramePassRadixSort.setup: keysTexture is required");
Debug.assert(elementCount > 0, "FramePassRadixSort.setup: elementCount must be > 0");
Debug.assert(numBits >= 1 && numBits <= 24, "FramePassRadixSort.setup: numBits must be 1-24");
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();
}
}
/**
* Calculates the required power-of-2 texture size for the given element count.
*
* @param {number} elementCount - Number of elements.
* @returns {number} Power-of-2 size.
* @private
*/
_calculateInternalSize(elementCount) {
const side = Math.ceil(Math.sqrt(elementCount));
return Math.pow(2, Math.ceil(Math.log2(side)));
}
/**
* Creates or resizes internal textures.
*
* @param {number} size - Power-of-2 size for textures.
* @private
*/
_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
});
}
/**
* Creates a texture for radix sort.
*
* @param {string} name - Texture name.
* @param {number} size - Texture size.
* @param {number} format - Pixel format (PIXELFORMAT_R32U or PIXELFORMAT_R32F).
* @param {boolean} [mipmaps] - Whether to generate mipmaps. Defaults to false.
* @returns {Texture} The created texture.
* @private
*/
_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
});
}
/**
* Destroys internal textures and render targets.
*
* @private
*/
_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;
}
/**
* Creates the sort passes based on numBits.
* Sets up beforePasses with the complete pass sequence (count, mipmap, reorder for each iteration).
*
* @private
*/
_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;
}
/**
* Destroys all sort passes.
*
* @private
*/
_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;
}
}
/**
* Executes the GPU radix sort. This is a convenience method that combines setup, frameUpdate,
* and rendering all passes in one call.
*
* @param {Texture} keysTexture - R32U texture containing sort keys (linear layout, any size).
* @param {number} elementCount - Number of elements to sort.
* @param {number} [numBits] - Number of bits to sort (1-24). More bits = more passes. Defaults to 16.
* @returns {Texture} The sorted indices texture (R32U, linear layout).
*/
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
};