playcanvas/build/playcanvas/src/scene/graphics/frame-pass-radix-sort.d.ts

Open-source WebGL/WebGPU 3D engine for the web

/**
 * A frame pass that performs GPU-based radix sort using mipmap-based prefix sums.
 *
 * This implementation is based on:
 * - VRChat Gaussian Splatting by MichaelMoroz: https://github.com/MichaelMoroz/VRChatGaussianSplatting
 * - Mipmap prefix sum trick by d4rkpl4y3r: https://github.com/d4rkc0d3r/CompactSparseTextureDemo
 *
 * ## Algorithm Overview
 *
 * The sort uses a 4-bit radix (16 buckets) and processes keys in multiple passes,
 * one pass per 4-bit chunk. Each pass consists of:
 *
 * 1. **Count Pass**: For each digit (0-15), count how many keys in each group have that digit.
 *    Output is an R32F texture where each pixel stores a count. Groups are 16 elements.
 *
 * 2. **Mipmap Generation**: Generate mipmaps for the count texture using hardware mipmap
 *    generation. This creates a quadtree of counts that enables efficient binary search.
 *
 * 3. **Reorder Pass**: For each output position, binary search through the mipmap hierarchy
 *    to find which source element maps to it. The mipmap structure enables O(log N) lookup
 *    per element instead of O(N) linear scan.
 *
 * ## Mipmap Prefix Sum Trick
 *
 * The key insight is that mipmaps naturally form a quadtree of averages. By writing counts
 * (e.g., 1.0 for active pixels) into an R32F texture with auto-generated mipmaps:
 *
 * - Each mip level stores the average of the 4 pixels below it
 * - To reconstruct actual counts, multiply by 4^level (i.e., `1 << (level * 2)`)
 * - This gives us a hierarchical prefix sum structure
 *
 * Binary search traversal:
 * - Start at maxMipLevel and work down to level 0
 * - At each level, check 3 quadrants (can skip 4th - if not in first 3, must be in 4th)
 * - Order: bottom-left → bottom-right → top-left → top-right (Z-order/Morton curve)
 * - Accumulate prefix sums while descending to find the target element
 *
 * The Z-order traversal ensures stable sorting: if element A comes before B in the input,
 * it remains before B in the output.
 *
 * ## Internal Data Layout
 *
 * - Internal keys/indices use Morton order (Z-order curve) for better texture cache locality
 * - Source keys texture uses linear (row-major) layout
 * - Output sorted indices use linear layout for simple consumer access
 *
 * ## Complexity
 *
 * - Time: O(N log N) per pass due to mipmap binary search
 * - Passes: ceil(numBits / 4) passes for numBits-bit keys
 * - Memory: 2x keys textures + 2x indices textures + 1x prefix sums texture (all power-of-2)
 *
 * @category Graphics
 * @ignore
 */
export class FramePassRadixSort extends FramePass {
    /**
     * The current sorted indices texture (R32U). Access sorted indices using Morton lookup.
     *
     * @type {Texture|null}
     */
    _currentIndices: Texture | null;
    /**
     * Current number of radix passes.
     */
    _numPasses: number;
    /**
     * Current internal texture size (power of 2).
     */
    _internalSize: number;
    /**
     * Internal keys texture 0 (ping-pong buffer).
     *
     * @type {Texture|null}
     */
    _keys0: Texture | null;
    /**
     * Internal keys texture 1 (ping-pong buffer).
     *
     * @type {Texture|null}
     */
    _keys1: Texture | null;
    /**
     * Internal indices texture 0 (ping-pong buffer).
     *
     * @type {Texture|null}
     */
    _indices0: Texture | null;
    /**
     * Internal indices texture 1 (ping-pong buffer).
     *
     * @type {Texture|null}
     */
    _indices1: Texture | null;
    /**
     * Prefix sums texture (R32F with mipmaps).
     *
     * @type {Texture|null}
     */
    _prefixSums: Texture | null;
    /**
     * Sort render target 0 (MRT for keys + indices).
     *
     * @type {RenderTarget|null}
     */
    _sortRT0: RenderTarget | null;
    /**
     * Sort render target 1 (MRT for keys + indices).
     *
     * @type {RenderTarget|null}
     */
    _sortRT1: RenderTarget | null;
    /**
     * Prefix sums render target.
     *
     * @type {RenderTarget|null}
     */
    _prefixSumsRT: RenderTarget | null;
    /**
     * Count passes for each radix iteration.
     *
     * @type {RenderPassRadixSortCount[]}
     */
    _countPasses: RenderPassRadixSortCount[];
    /**
     * Reorder passes for each radix iteration.
     *
     * @type {RenderPassRadixSortReorder[]}
     */
    _reorderPasses: RenderPassRadixSortReorder[];
    /**
     * Number of elements to sort (set by setup()).
     */
    _elementCount: number;
    /**
     * The source keys texture (set by setup()).
     *
     * @type {Texture|null}
     */
    _keysTexture: Texture | null;
    /**
     * 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(): Texture | null;
    /**
     * 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: Texture, elementCount: number, numBits?: number): void;
    /**
     * 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
     */
    private _calculateInternalSize;
    /**
     * Creates or resizes internal textures.
     *
     * @param {number} size - Power-of-2 size for textures.
     * @private
     */
    private _resizeInternalTextures;
    /**
     * 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
     */
    private _createTexture;
    /**
     * Destroys internal textures and render targets.
     *
     * @private
     */
    private _destroyInternalTextures;
    /**
     * Creates the sort passes based on numBits.
     * Sets up beforePasses with the complete pass sequence (count, mipmap, reorder for each iteration).
     *
     * @private
     */
    private _createPasses;
    /**
     * Destroys all sort passes.
     *
     * @private
     */
    private _destroyPasses;
    /**
     * 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: Texture, elementCount: number, numBits?: number): Texture;
}
import { FramePass } from '../../platform/graphics/frame-pass.js';
import { Texture } from '../../platform/graphics/texture.js';
import { RenderTarget } from '../../platform/graphics/render-target.js';
import { RenderPassRadixSortCount } from './render-pass-radix-sort-count.js';
import { RenderPassRadixSortReorder } from './render-pass-radix-sort-reorder.js';