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

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

export const onesweepScanSource: "\n\n@group(0) @binding(0) var<storage, read> b_globalHist: array<u32>;\n@group(0) @binding(1) var<storage, read_write> b_passHist: array<atomic<u32>>;\n\nstruct OneSweepScanUniforms {\n    threadBlocks: u32,   // number of DigitBinningPass workgroups per pass (ignored in indirect mode)\n    _pad0: u32,\n    _pad1: u32,\n    _pad2: u32\n};\n@group(0) @binding(2) var<uniform> uniforms: OneSweepScanUniforms;\n\n#ifdef USE_INDIRECT_SORT\n// Indirect dispatch: threadBlocks is derived from a GPU-written element count.\n@group(0) @binding(3) var<storage, read> b_sortElementCount: array<u32>;\n#endif\n\nconst RADIX: u32 = 256u;\nconst FLAG_INCLUSIVE: u32 = 2u;\nconst PART_SIZE: u32 = {PART_SIZE}u;\n\n// Parametrized by the host from device.maxSubgroupSize (256 / sgSize).\nconst MAX_SUBGROUPS: u32 = {MAX_SUBGROUPS}u;\n\n// Scratch for the hierarchical exclusive scan. sg_totals holds one entry per\n// subgroup; lane 0 of the workgroup scans it serially.\nvar<workgroup> g_scan: array<u32, RADIX>;\nvar<workgroup> sg_totals: array<u32, MAX_SUBGROUPS>;\n\n@compute @workgroup_size(RADIX, 1, 1)\nfn main(\n    @builtin(local_invocation_index) gtid: u32,\n    @builtin(workgroup_id) gid: vec3<u32>,\n    @builtin(subgroup_invocation_id) sgInvId: u32,\n    @builtin(subgroup_size) sgSize: u32,\n) {\n    let pass_ = gid.x;\n    #ifdef USE_INDIRECT_SORT\n    let numKeys = b_sortElementCount[0];\n    let threadBlocks = (numKeys + PART_SIZE - 1u) / PART_SIZE;\n    #else\n    let threadBlocks = uniforms.threadBlocks;\n    #endif\n    let waveIndex = gtid / sgSize;\n\n    // Load this pass's digit counts.\n    let t = b_globalHist[gtid + pass_ * RADIX];\n\n    // Phase 1: subgroup-level exclusive scan.\n    let sgExcl = subgroupExclusiveAdd(t);\n    let sgTotal = subgroupAdd(t);\n\n    if (sgInvId == 0u) {\n        sg_totals[waveIndex] = sgTotal;\n    }\n    workgroupBarrier();\n\n    // Phase 2: scan the subgroup totals (serially in thread 0; MAX_SUBGROUPS entries).\n    if (gtid == 0u) {\n        var acc: u32 = 0u;\n        for (var i = 0u; i < MAX_SUBGROUPS; i = i + 1u) {\n            let v = sg_totals[i];\n            sg_totals[i] = acc;\n            acc = acc + v;\n        }\n    }\n    workgroupBarrier();\n\n    // Phase 3: combine subgroup-local prefix with the subgroup base.\n    let excl = sgExcl + sg_totals[waveIndex];\n    g_scan[gtid] = excl;\n\n    // Publish to block-0 slot of passHist with FLAG_INCLUSIVE.\n    // Layout: b_passHist[pass * threadBlocks * RADIX + block * RADIX + digit].\n    let dst = pass_ * threadBlocks * RADIX + gtid;\n    atomicStore(&b_passHist[dst], (excl << 2u) | FLAG_INCLUSIVE);\n}\n";
export default onesweepScanSource;