playcanvas/build/playcanvas/src/scene/shader-lib/wgsl/chunks/radix-sort/onesweep-binning.d.ts

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

export const onesweepBinningSource: "\n\n@group(0) @binding(0) var<storage, read> inputKeys: array<u32>;\n@group(0) @binding(1) var<storage, read_write> outputKeys: array<u32>;\n@group(0) @binding(2) var<storage, read> inputValues: array<u32>;\n@group(0) @binding(3) var<storage, read_write> outputValues: array<u32>;\n@group(0) @binding(4) var<storage, read_write> b_passHist: array<atomic<u32>>;\n@group(0) @binding(5) var<storage, read_write> b_index: array<atomic<u32>>;\n\nstruct OneSweepBinningUniforms {\n    numKeys: u32,        // ignored in indirect mode\n    threadBlocks: u32,   // DigitBinningPass workgroup count per pass (ignored in indirect mode)\n    pass_: u32,          // 0..NUM_PASSES-1\n    flags: u32           // bit 0: isFirstPass, bit 1: isLastPass (skip key write)\n};\n@group(0) @binding(6) var<uniform> uniforms: OneSweepBinningUniforms;\n\n#ifdef USE_INDIRECT_SORT\n// Indirect dispatch: numKeys/threadBlocks are derived from a GPU-written\n// element count. The uniform fields are ignored.\n@group(0) @binding(7) var<storage, read> b_sortElementCount: array<u32>;\n#endif\n\nconst RADIX: u32 = 256u;\nconst RADIX_MASK: u32 = 255u;\nconst RADIX_LOG: u32 = 8u;\n\nconst D_DIM: u32 = {D_DIM}u;\nconst KEYS_PER_THREAD: u32 = {KEYS_PER_THREAD}u;\nconst PART_SIZE: u32 = D_DIM * KEYS_PER_THREAD; // 3840 for D_DIM=256, KEYS=15\n// Parametrized by the host from device.maxSubgroupSize:\n//  sgSize=32 (NVIDIA / Apple M-series / Intel / desktop AMD): MAX_SUBGROUPS = 8,  WAVE_HISTS_SIZE = 2048.\n//  sgSize=16 (Mali / Pixel / some Imagination):                MAX_SUBGROUPS = 16, WAVE_HISTS_SIZE = 4096.\n//  sgSize=64/128 (Adreno):                                     MAX_SUBGROUPS = 4,  WAVE_HISTS_SIZE = 1024.\nconst MAX_SUBGROUPS: u32 = {MAX_SUBGROUPS}u;\nconst WAVE_HISTS_SIZE: u32 = MAX_SUBGROUPS * RADIX;\n\n// g_d must be large enough for both the ranking phase (WAVE_HISTS_SIZE slots)\n// and the staging phase (PART_SIZE slots). For D_DIM=256, KEYS_PER_THREAD=15\n// this is:\n//  sgSize=32 (MAX_SUBGROUPS=8 ): max(3840, 2048) = 3840\n//  sgSize=16 (MAX_SUBGROUPS=16): max(3840, 4096) = 4096  (+1 KiB vs sgSize=32)\n//  sgSize=64 (MAX_SUBGROUPS=4 ): max(3840, 1024) = 3840\n// Sizing g_d to PART_SIZE alone (as in the original port) corrupts waves 15..\n// on sgSize=16 hardware because their per-warp histogram slots fall out of\n// bounds.\nconst G_D_SIZE: u32 = max(PART_SIZE, WAVE_HISTS_SIZE);\n\nconst FLAG_NOT_READY: u32 = 0u;\nconst FLAG_REDUCTION: u32 = 1u;\nconst FLAG_INCLUSIVE: u32 = 2u;\nconst FLAG_MASK: u32 = 3u;\n\n// Staging memory. Reused across phases:\n//   phase A (ranking):          per-warp histograms (MAX_SUBGROUPS \u00D7 256 u32) in slots 0..WAVE_HISTS_SIZE\n//   phase E (key staging):      sorted keys at block-local offset, slots 0..PART_SIZE\n//   phase G (value staging):    values at block-local offset, slots 0..PART_SIZE\n// Declared atomic to satisfy atomicAdd in phase A. Other phases use\n// atomicStore/atomicLoad which behave like plain stores/loads on modern GPUs.\nvar<workgroup> g_d: array<atomic<u32>, G_D_SIZE>;\n\n// After phase B: per-digit block-local base.\n// After phase D: per-digit GLOBAL base (minus block-local exclusive prefix).\n// Adding a linear staging index to digit_base[digit] gives the global output\n// position for any key at that staging slot with that digit.\nvar<workgroup> digit_base: array<u32, RADIX>;\n\n// Scratch for the 2-level exclusive scan of per-digit block totals.\nvar<workgroup> sg_totals: array<u32, MAX_SUBGROUPS>;\n\n// Broadcast slot for the atomically-acquired partition tile id.\nvar<workgroup> wg_partIndex: u32;\n\n// 'passHistOffset' needs 'threadBlocks' to compute the per-pass row stride.\n// 'threadBlocks' is a local at the top of 'main' (from either the uniform or\n// the GPU-side element count); we plumb it through as an explicit argument\n// rather than a module-scope 'var' so the value stays in registers.\nfn passHistOffset(tb: u32, pass_: u32, partitionIdx: u32) -> u32 {\n    return pass_ * tb * RADIX + partitionIdx * RADIX;\n}\n\n@compute @workgroup_size(D_DIM, 1, 1)\nfn main(\n    @builtin(local_invocation_index) TID: u32,\n    @builtin(subgroup_invocation_id) sgInvId: u32,\n    @builtin(subgroup_size) sgSize: u32,\n) {\n    let waveIndex = TID / sgSize;\n    let ltMask = (1u << sgInvId) - 1u;\n    // Active-lane mask for the match-any ballot below. WGSL says inactive-lane\n    // bits of subgroupBallot are 0, but drivers (notably Mali / Imagination\n    // at sgSize<32) don't always honour this for subgroupBallot(is_valid).\n    // Initialising waveFlag to only cover active lanes makes the per-bit\n    // AND-chain correct regardless of driver behaviour. 1u << 32u is UB so\n    // branch on sgSize < 32.\n    let activeMask = select(0xFFFFFFFFu, (1u << sgSize) - 1u, sgSize < 32u);\n    let pass_ = uniforms.pass_;\n    let currentBit = pass_ << 3u;\n    #ifdef USE_INDIRECT_SORT\n    let numKeys = b_sortElementCount[0];\n    let threadBlocks = (numKeys + PART_SIZE - 1u) / PART_SIZE;\n    #else\n    let numKeys = uniforms.numKeys;\n    let threadBlocks = uniforms.threadBlocks;\n    #endif\n    let isFirstPass = (uniforms.flags & 1u) != 0u;\n    let isLastPass = (uniforms.flags & 2u) != 0u;\n\n    // ---- Phase 0: assign partition tile ----\n    if (TID == 0u) {\n        wg_partIndex = atomicAdd(&b_index[pass_], 1u);\n    }\n    let partitionIndex = workgroupUniformLoad(&wg_partIndex);\n\n    // ---- Phase A.1: clear per-warp histograms ----\n    // Only the first 2048 slots hold wave hists during ranking. We clear and\n    // re-use them; later phases (staging) overwrite beyond slot 2048 too.\n    for (var i = TID; i < WAVE_HISTS_SIZE; i = i + D_DIM) {\n        atomicStore(&g_d[i], 0u);\n    }\n    workgroupBarrier();\n\n    let tileStart = partitionIndex * PART_SIZE;\n    let validInBlock = select(\n        0u,\n        min(PART_SIZE, numKeys - tileStart),\n        tileStart < numKeys\n    );\n\n    // ---- Phase A.2: wave-interleaved load of keys + values ----\n    // Values are loaded up-front (into registers) alongside keys so that phase G\n    // does not need a second full-bandwidth read of inputValues. This saves one\n    // pass over the value buffer per radix pass, at the cost of keeping\n    // KEYS_PER_THREAD extra u32 live across the ranking loop. With\n    // KEYS_PER_THREAD a compile-time constant and static indexing preserved,\n    // the compiler register-allocates values[] the same way it does keys[].\n    let subPartSize = sgSize * KEYS_PER_THREAD;\n    let waveBase = tileStart + waveIndex * subPartSize;\n\n    var keys: array<u32, {KEYS_PER_THREAD}>;\n    var values: array<u32, {KEYS_PER_THREAD}>;\n    var validMask: u32 = 0u;\n    for (var i = 0u; i < KEYS_PER_THREAD; i = i + 1u) {\n        let gid = waveBase + sgInvId + i * sgSize;\n        let is_valid = gid < numKeys;\n        // Dummy 0xFFFFFFFF for invalid lanes: validBallot drops them from\n        // any real digit's run.\n        keys[i] = select(0xFFFFFFFFu, inputKeys[gid], is_valid);\n        // On the first pass, values are synthesised as the original index\n        // (identity permutation), so we skip the value-buffer load entirely.\n        // On subsequent passes, values are the permutation from the previous\n        // pass; read once here and reuse in phase G.\n        values[i] = select(\n            select(0u, inputValues[gid], is_valid),\n            gid,\n            isFirstPass\n        );\n        if (is_valid) {\n            validMask = validMask | (1u << i);\n        }\n    }\n\n    // ---- Phase A.3: rank keys (RankKeysWGE16) ----\n    var offsets: array<u32, {KEYS_PER_THREAD}>;\n    for (var i = 0u; i < KEYS_PER_THREAD; i = i + 1u) {\n        let k = keys[i];\n        let isValid = ((validMask >> i) & 1u) == 1u;\n        let digit = (k >> currentBit) & RADIX_MASK;\n\n        var waveFlag: u32 = activeMask;\n        for (var b = 0u; b < 8u; b = b + 1u) {\n            let t = ((digit >> b) & 1u) == 1u;\n            let ballot = subgroupBallot(t).x;\n            waveFlag = waveFlag & select(~ballot, ballot, t);\n        }\n        let validBallot = subgroupBallot(isValid).x;\n        waveFlag = waveFlag & validBallot;\n\n        let peerBits = countOneBits(waveFlag & ltMask);\n        let totalBits = countOneBits(waveFlag);\n        let lowestRankPeer = firstTrailingBit(waveFlag);\n\n        var preIncrementVal: u32 = 0u;\n        if (isValid && peerBits == 0u) {\n            preIncrementVal = atomicAdd(&g_d[waveIndex * RADIX + digit], totalBits);\n        }\n        offsets[i] = subgroupShuffle(preIncrementVal, lowestRankPeer) + peerBits;\n\n        // Force lane reconvergence before the next iteration. Without this,\n        // NVIDIA Turing+ Independent Thread Scheduling can let two different\n        // rounds of this loop interleave within a single warp, corrupting the\n        // atomicAdd/subgroupShuffle pairing. WGSL has no subgroupBarrier;\n        // workgroupBarrier is the cheapest portable substitute.\n        workgroupBarrier();\n    }\n\n    // ---- Phase A.4: circular-shift inclusive scan across warps ----\n    // After this loop, for digit TID (TID < RADIX):\n    //   - myHistRed = total count of digit TID across all warps in this block.\n    //   - g_d[TID + w*RADIX] for w >= 1 holds the exclusive per-warp prefix.\n    var myHistRed: u32 = 0u;\n    {\n        var histReduction = atomicLoad(&g_d[TID]);\n        for (var w = 1u; w < MAX_SUBGROUPS; w = w + 1u) {\n            let idx = TID + w * RADIX;\n            let cnt = atomicLoad(&g_d[idx]);\n            histReduction = histReduction + cnt;\n            atomicStore(&g_d[idx], histReduction - cnt);\n        }\n        myHistRed = histReduction;\n    }\n\n    // ---- Phase A.5: publish this block's per-digit totals ----\n    // DeviceBroadcastReductionsWGE16: the block at partitionIndex writes to\n    // slot partitionIndex+1 of passHist (i.e. its successor's inbox). The\n    // last block has no successor and skips this step.\n    if (partitionIndex + 1u < threadBlocks) {\n        let dst = passHistOffset(threadBlocks, pass_, partitionIndex + 1u) + TID;\n        atomicAdd(&b_passHist[dst], FLAG_REDUCTION | (myHistRed << 2u));\n    }\n\n    // ---- Phase B: per-digit exclusive scan (hierarchical) ----\n    let warpExcl = subgroupExclusiveAdd(myHistRed);\n    let warpTotal = subgroupAdd(myHistRed);\n\n    if (sgInvId == 0u) {\n        sg_totals[waveIndex] = warpTotal;\n    }\n    workgroupBarrier();\n\n    if (TID == 0u) {\n        var acc: u32 = 0u;\n        for (var w = 0u; w < MAX_SUBGROUPS; w = w + 1u) {\n            let t = sg_totals[w];\n            sg_totals[w] = acc;\n            acc = acc + t;\n        }\n    }\n    workgroupBarrier();\n\n    let myDigitBase = warpExcl + sg_totals[waveIndex];\n\n    // ---- Phase C: per-key scatter positions (block-local) ----\n    // scatterPos[i] = intra-warp rank + per-warp base for this digit\n    //               + block-local base across earlier digits.\n    var scatterPos: array<u32, {KEYS_PER_THREAD}>;\n    for (var i = 0u; i < KEYS_PER_THREAD; i = i + 1u) {\n        let k = keys[i];\n        let digit = (k >> currentBit) & RADIX_MASK;\n        // digit_base not yet populated; we need warp base + myDigitBase for\n        // the key's digit (which is DIFFERENT from TID's digit). We read\n        // myDigitBase for any digit by publishing digit_base[TID] = myDigitBase\n        // below, which means we need a staging step. Alternatively, compute\n        // the combined base on-the-fly using the per-warp prefix stored in\n        // g_d and reading myDigitBase-equivalent via a shared array.\n        // We stage myDigitBase into digit_base first.\n        let warpBase = select(0u, atomicLoad(&g_d[waveIndex * RADIX + digit]), waveIndex > 0u);\n        scatterPos[i] = offsets[i] + warpBase; // add per-digit block base below\n    }\n\n    // Publish myDigitBase so each thread can look up the base for its keys'\n    // digits (which usually differ from TID for most keys).\n    digit_base[TID] = myDigitBase;\n    workgroupBarrier();\n\n    for (var i = 0u; i < KEYS_PER_THREAD; i = i + 1u) {\n        let k = keys[i];\n        let digit = (k >> currentBit) & RADIX_MASK;\n        scatterPos[i] = scatterPos[i] + digit_base[digit];\n    }\n    workgroupBarrier();\n\n    // ---- Phase D: Lookback + global base resolution ----\n    // Plain decoupled lookback: each digit-owning thread walks backward\n    // through passHist for its digit until it finds FLAG_INCLUSIVE,\n    // accumulating reductions on the way. Requires forward-thread-progress\n    // guarantees (NVIDIA Turing+, recent AMD, Intel Gen9+). On devices\n    // without those guarantees (Apple Silicon, Mali, Adreno) this may\n    // deadlock; callers should use {@link ComputeRadixSort} on those\n    // architectures instead.\n    // On finding FLAG_INCLUSIVE, it atomically upgrades its own (partition+1)\n    // slot to FLAG_INCLUSIVE so later blocks terminate faster.\n    // No workgroupBarrier inside the loop: the spin is per-thread/digit and\n    // we only need sync *after* all threads complete, before scatter uses\n    // digit_base in its new form.\n    if (TID < RADIX) {\n        var lookbackReduction: u32 = 0u;\n        var k: u32 = partitionIndex;\n        var done: bool = false;\n        loop {\n            if (done) { break; }\n            let flagPayload = atomicLoad(&b_passHist[passHistOffset(threadBlocks, pass_, k) + TID]);\n            let flag = flagPayload & FLAG_MASK;\n\n            if (flag == FLAG_INCLUSIVE) {\n                lookbackReduction = lookbackReduction + (flagPayload >> 2u);\n                if (partitionIndex + 1u < threadBlocks) {\n                    // Flip FLAG_REDUCTION (01) to FLAG_INCLUSIVE (10) by adding\n                    // 1, and fold in the full exclusive prefix so downstream\n                    // blocks can terminate their lookback on this slot.\n                    let dst = passHistOffset(threadBlocks, pass_, partitionIndex + 1u) + TID;\n                    atomicAdd(&b_passHist[dst], 1u | (lookbackReduction << 2u));\n                }\n                // Convert digit_base[TID] from block-local base to the value\n                // needed during scatter: globalPrefix - blockLocalExclusive.\n                digit_base[TID] = lookbackReduction - myDigitBase;\n                done = true;\n            } else if (flag == FLAG_REDUCTION) {\n                lookbackReduction = lookbackReduction + (flagPayload >> 2u);\n                // Scan kernel writes block 0's slot as FLAG_INCLUSIVE, so we\n                // must see that before underflowing. Guard anyway.\n                if (k == 0u) { done = true; }\n                else { k = k - 1u; }\n            }\n            // FLAG_NOT_READY: spin on the same slot.\n        }\n    }\n\n    // ---- Phase E: scatter keys into shared-memory staging ----\n    for (var i = 0u; i < KEYS_PER_THREAD; i = i + 1u) {\n        if (((validMask >> i) & 1u) == 1u) {\n            atomicStore(&g_d[scatterPos[i]], keys[i]);\n        }\n    }\n    workgroupBarrier();\n\n    // ---- Phase F: linear key read \u2192 coalesced global write ----\n    var linearDigits: array<u32, {KEYS_PER_THREAD}>;\n    for (var r = 0u; r < KEYS_PER_THREAD; r = r + 1u) {\n        let linearIdx = TID + r * D_DIM;\n        if (linearIdx < validInBlock) {\n            let k = atomicLoad(&g_d[linearIdx]);\n            let digit = (k >> currentBit) & RADIX_MASK;\n            linearDigits[r] = digit;\n            let globalPos = digit_base[digit] + linearIdx;\n            if (!isLastPass) {\n                outputKeys[globalPos] = k;\n            }\n        }\n    }\n    workgroupBarrier();\n\n    // ---- Phase G: scatter values into staging ----\n    // Values were loaded into registers in phase A.2; reuse them here.\n    for (var i = 0u; i < KEYS_PER_THREAD; i = i + 1u) {\n        if (((validMask >> i) & 1u) == 1u) {\n            atomicStore(&g_d[scatterPos[i]], values[i]);\n        }\n    }\n    workgroupBarrier();\n\n    // ---- Phase H: linear value read \u2192 coalesced global write ----\n    for (var r = 0u; r < KEYS_PER_THREAD; r = r + 1u) {\n        let linearIdx = TID + r * D_DIM;\n        if (linearIdx < validInBlock) {\n            let v = atomicLoad(&g_d[linearIdx]);\n            let digit = linearDigits[r];\n            let globalPos = digit_base[digit] + linearIdx;\n            outputValues[globalPos] = v;\n        }\n    }\n}\n";
export default onesweepBinningSource;