UNPKG

playcanvas

Version:

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

playcanvas.com

playcanvas/engine

3 lines (2 loc) • 6.92 kB
1
2
3
export const onesweepGlobalHistSource: "\n\n@group(0) @binding(0) var<storage, read> b_sort: array<vec4<u32>>;\n@group(0) @binding(1) var<storage, read_write> b_globalHist: array<atomic<u32>>;\n\nstruct OneSweepUniforms {\n numKeys: u32,\n threadBlocks: u32, // number of DigitBinningPass workgroups (unused here)\n numPasses: u32, // 1..MAX_PASSES\n _pad: u32\n};\n@group(0) @binding(2) var<uniform> uniforms: OneSweepUniforms;\n\n#ifdef USE_INDIRECT_SORT\n// Indirect dispatch: numKeys is GPU-written. uniforms.numKeys is ignored.\n@group(0) @binding(3) var<storage, read> b_sortElementCount: array<u32>;\n#endif\n\nconst RADIX: u32 = 256u;\nconst MAX_PASSES: u32 = 4u;\nconst G_HIST_DIM: u32 = {G_HIST_DIM}u;\nconst G_HIST_PART_SIZE: u32 = {G_HIST_PART_SIZE}u;\nconst G_HIST_PART_SIZE_VEC: u32 = G_HIST_PART_SIZE / 4u; // partition in vec4 units\nconst SHARED_HIST_SIZE: u32 = 2u * MAX_PASSES * RADIX; // 2048\n\n// 2 rows \u00D7 MAX_PASSES \u00D7 RADIX u32 atomics. For NUM_PASSES < MAX_PASSES the\n// trailing rows are unused but harmless.\nvar<workgroup> g_gHist: array<atomic<u32>, SHARED_HIST_SIZE>;\n\nfn histOffset(row: u32, pass_: u32) -> u32 {\n return row * RADIX + pass_ * 2u * RADIX;\n}\n\n@compute @workgroup_size(G_HIST_DIM, 1, 1)\nfn main(\n @builtin(local_invocation_index) gtid: u32,\n @builtin(workgroup_id) gid: vec3<u32>,\n @builtin(num_workgroups) nwg: vec3<u32>,\n) {\n let flatGid = gid.x + gid.y * nwg.x;\n let numPasses = uniforms.numPasses;\n\n #ifdef USE_INDIRECT_SORT\n let numKeys = b_sortElementCount[0];\n #else\n let numKeys = uniforms.numKeys;\n #endif\n\n // Clear shared histogram (only the rows we'll actually use).\n let sharedEnd = 2u * numPasses * RADIX;\n for (var i = gtid; i < sharedEnd; i = i + G_HIST_DIM) {\n atomicStore(&g_gHist[i], 0u);\n }\n workgroupBarrier();\n\n // Process this workgroup's partition tile (vec4 units). The loop bound\n // is clamped to FULL vec4s (numKeys >> 2u); any partial trailing vec4\n // is handled out-of-loop by a single thread below.\n let row = gtid / 64u; // 0 or 1\n let numKeysVecFull = numKeys >> 2u;\n let partitionStartVec = flatGid * G_HIST_PART_SIZE_VEC;\n let partitionEndVec = min(partitionStartVec + G_HIST_PART_SIZE_VEC, numKeysVecFull);\n\n // Fast path: all 4 lanes always valid, no per-lane bounds checks.\n // Unrolled per-lane, per-pass digit extraction; numPasses is known at\n // runtime but MAX_PASSES is compile-time, so naga strips the guarded\n // blocks for NUM_PASSES < 4.\n for (var i = partitionStartVec + gtid; i < partitionEndVec; i = i + G_HIST_DIM) {\n let q = b_sort[i];\n if (numPasses >= 1u) {\n let off = histOffset(row, 0u);\n atomicAdd(&g_gHist[(q.x & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[(q.y & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[(q.z & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[(q.w & 0xFFu) + off], 1u);\n }\n if (numPasses >= 2u) {\n let off = histOffset(row, 1u);\n atomicAdd(&g_gHist[((q.x >> 8u) & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[((q.y >> 8u) & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[((q.z >> 8u) & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[((q.w >> 8u) & 0xFFu) + off], 1u);\n }\n if (numPasses >= 3u) {\n let off = histOffset(row, 2u);\n atomicAdd(&g_gHist[((q.x >> 16u) & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[((q.y >> 16u) & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[((q.z >> 16u) & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[((q.w >> 16u) & 0xFFu) + off], 1u);\n }\n if (numPasses >= 4u) {\n let off = histOffset(row, 3u);\n atomicAdd(&g_gHist[((q.x >> 24u) & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[((q.y >> 24u) & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[((q.z >> 24u) & 0xFFu) + off], 1u);\n atomicAdd(&g_gHist[((q.w >> 24u) & 0xFFu) + off], 1u);\n }\n }\n\n // Ragged-tail path: at most ONE vec4 at index (numKeys >> 2u) has 1-3\n // valid lanes (never 4 \u2014 if it were, the fast loop would own it). The\n // thread whose position in the strided loop would have landed on\n // tailIdx handles it; everyone else falls through.\n let tailIdx = numKeysVecFull;\n if ((numKeys & 3u) != 0u &&\n tailIdx >= partitionStartVec &&\n tailIdx < partitionStartVec + G_HIST_PART_SIZE_VEC &&\n (tailIdx - partitionStartVec) % G_HIST_DIM == gtid) {\n let q = b_sort[tailIdx];\n let base = tailIdx << 2u;\n // Lane 3 is always invalid here (numKeys % 4 == 1/2/3), so it is\n // unconditionally dropped. Lanes 0..2 are guarded.\n if (numPasses >= 1u) {\n let off = histOffset(row, 0u);\n if (base + 0u < numKeys) { atomicAdd(&g_gHist[(q.x & 0xFFu) + off], 1u); }\n if (base + 1u < numKeys) { atomicAdd(&g_gHist[(q.y & 0xFFu) + off], 1u); }\n if (base + 2u < numKeys) { atomicAdd(&g_gHist[(q.z & 0xFFu) + off], 1u); }\n }\n if (numPasses >= 2u) {\n let off = histOffset(row, 1u);\n if (base + 0u < numKeys) { atomicAdd(&g_gHist[((q.x >> 8u) & 0xFFu) + off], 1u); }\n if (base + 1u < numKeys) { atomicAdd(&g_gHist[((q.y >> 8u) & 0xFFu) + off], 1u); }\n if (base + 2u < numKeys) { atomicAdd(&g_gHist[((q.z >> 8u) & 0xFFu) + off], 1u); }\n }\n if (numPasses >= 3u) {\n let off = histOffset(row, 2u);\n if (base + 0u < numKeys) { atomicAdd(&g_gHist[((q.x >> 16u) & 0xFFu) + off], 1u); }\n if (base + 1u < numKeys) { atomicAdd(&g_gHist[((q.y >> 16u) & 0xFFu) + off], 1u); }\n if (base + 2u < numKeys) { atomicAdd(&g_gHist[((q.z >> 16u) & 0xFFu) + off], 1u); }\n }\n if (numPasses >= 4u) {\n let off = histOffset(row, 3u);\n if (base + 0u < numKeys) { atomicAdd(&g_gHist[((q.x >> 24u) & 0xFFu) + off], 1u); }\n if (base + 1u < numKeys) { atomicAdd(&g_gHist[((q.y >> 24u) & 0xFFu) + off], 1u); }\n if (base + 2u < numKeys) { atomicAdd(&g_gHist[((q.z >> 24u) & 0xFFu) + off], 1u); }\n }\n }\n workgroupBarrier();\n\n // Reduce rows and atomically add into global histogram.\n for (var i = gtid; i < RADIX; i = i + G_HIST_DIM) {\n for (var p = 0u; p < numPasses; p = p + 1u) {\n let row0 = atomicLoad(&g_gHist[i + histOffset(0u, p)]);\n let row1 = atomicLoad(&g_gHist[i + histOffset(1u, p)]);\n let total = row0 + row1;\n if (total != 0u) {\n atomicAdd(&b_globalHist[i + p * RADIX], total);\n }\n }\n }\n}\n"; export default onesweepGlobalHistSource;