UNPKG

playcanvas

Version:

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

playcanvas.com

playcanvas/engine

3 lines (2 loc) • 3.33 kB
1
2
3
export const prefixSumSource: "\n\n@group(0) @binding(0) var<storage, read_write> items: array<u32>;\n@group(0) @binding(1) var<storage, read_write> blockSums: array<u32>;\n\n// Uniform for runtime element count (changes per-sort without shader recompilation)\nstruct PrefixSumUniforms {\n elementCount: u32\n};\n@group(0) @binding(2) var<uniform> uniforms: PrefixSumUniforms;\n\n// Compile-time constants\nconst WORKGROUP_SIZE_X: u32 = {WORKGROUP_SIZE_X}u;\nconst WORKGROUP_SIZE_Y: u32 = {WORKGROUP_SIZE_Y}u;\nconst THREADS_PER_WORKGROUP: u32 = {THREADS_PER_WORKGROUP}u;\nconst ITEMS_PER_WORKGROUP: u32 = {ITEMS_PER_WORKGROUP}u;\n\nvar<workgroup> temp: array<u32, ITEMS_PER_WORKGROUP * 2>;\n\n@compute @workgroup_size(WORKGROUP_SIZE_X, WORKGROUP_SIZE_Y, 1)\nfn reduce_downsweep(\n @builtin(workgroup_id) w_id: vec3<u32>,\n @builtin(num_workgroups) w_dim: vec3<u32>,\n @builtin(local_invocation_index) TID: u32,\n) {\n let WORKGROUP_ID = w_id.x + w_id.y * w_dim.x;\n let WID = WORKGROUP_ID * THREADS_PER_WORKGROUP;\n let GID = WID + TID;\n \n let ELM_TID = TID * 2;\n let ELM_GID = GID * 2;\n \n // Load input to shared memory\n temp[ELM_TID] = select(items[ELM_GID], 0u, ELM_GID >= uniforms.elementCount);\n temp[ELM_TID + 1u] = select(items[ELM_GID + 1u], 0u, ELM_GID + 1u >= uniforms.elementCount);\n\n var offset: u32 = 1u;\n\n // Up-sweep (reduce) phase\n for (var d: u32 = ITEMS_PER_WORKGROUP >> 1u; d > 0u; d >>= 1u) {\n workgroupBarrier();\n\n if (TID < d) {\n var ai: u32 = offset * (ELM_TID + 1u) - 1u;\n var bi: u32 = offset * (ELM_TID + 2u) - 1u;\n temp[bi] += temp[ai];\n }\n\n offset *= 2u;\n }\n\n // Save workgroup sum and clear last element\n if (TID == 0u) {\n let last_offset = ITEMS_PER_WORKGROUP - 1u;\n blockSums[WORKGROUP_ID] = temp[last_offset];\n temp[last_offset] = 0u;\n }\n\n // Down-sweep phase\n for (var d: u32 = 1u; d < ITEMS_PER_WORKGROUP; d *= 2u) {\n offset >>= 1u;\n workgroupBarrier();\n\n if (TID < d) {\n var ai: u32 = offset * (ELM_TID + 1u) - 1u;\n var bi: u32 = offset * (ELM_TID + 2u) - 1u;\n\n let t: u32 = temp[ai];\n temp[ai] = temp[bi];\n temp[bi] += t;\n }\n }\n workgroupBarrier();\n\n // Copy result from shared memory to global memory\n if (ELM_GID < uniforms.elementCount) {\n items[ELM_GID] = temp[ELM_TID];\n }\n\n if (ELM_GID + 1u < uniforms.elementCount) {\n items[ELM_GID + 1u] = temp[ELM_TID + 1u];\n }\n}\n\n@compute @workgroup_size(WORKGROUP_SIZE_X, WORKGROUP_SIZE_Y, 1)\nfn add_block_sums(\n @builtin(workgroup_id) w_id: vec3<u32>,\n @builtin(num_workgroups) w_dim: vec3<u32>,\n @builtin(local_invocation_index) TID: u32,\n) {\n let WORKGROUP_ID = w_id.x + w_id.y * w_dim.x;\n let WID = WORKGROUP_ID * THREADS_PER_WORKGROUP;\n let GID = WID + TID;\n\n let ELM_ID = GID * 2u;\n\n if (ELM_ID >= uniforms.elementCount) {\n return;\n }\n\n let blockSum = blockSums[WORKGROUP_ID];\n\n items[ELM_ID] += blockSum;\n\n if (ELM_ID + 1u >= uniforms.elementCount) {\n return;\n }\n\n items[ELM_ID + 1u] += blockSum;\n}\n"; export default prefixSumSource;