UNPKG

playcanvas

Version:

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

playcanvas.com

playcanvas/engine

2 lines (1 loc) • 6.33 kB
1
2
export const computeGsplatLocalBucketSortSource: "\n\nconst NUM_BUCKETS: u32 = 128u;\nconst MAX_CHUNK_SIZE: u32 = 4096u;\nconst WG_SIZE: u32 = 256u;\n@group(0) @binding(0) var<storage, read_write> tileEntries: array<u32>;\n@group(0) @binding(1) var<storage, read> largeTileOverflowBases: array<u32>;\n@group(0) @binding(2) var<storage, read> tileSplatCounts: array<u32>;\n@group(0) @binding(3) var<storage, read> depthBuffer: array<u32>;\n@group(0) @binding(4) var<storage, read> largeTileList: array<u32>;\n@group(0) @binding(5) var<storage, read_write> chunkRanges: array<u32>;\n@group(0) @binding(6) var<storage, read_write> totalChunks: array<atomic<u32>>;\n\n@group(0) @binding(7) var<storage, read> tileListCounts: array<u32>;\n\nstruct Uniforms {\n bufferCapacity: u32,\n maxChunks: u32,\n}\n@group(0) @binding(8) var<uniform> uniforms: Uniforms;\n\nvar<workgroup> sDepthMin: atomic<u32>;\nvar<workgroup> sDepthMax: atomic<u32>;\nvar<workgroup> sBucketCounts: array<atomic<u32>, NUM_BUCKETS>;\nvar<workgroup> sBucketOffsets: array<u32, NUM_BUCKETS + 1>;\nvar<workgroup> sBucketCursors: array<atomic<u32>, NUM_BUCKETS>;\n\n@compute @workgroup_size(256)\nfn main(\n @builtin(local_invocation_index) localIdx: u32,\n @builtin(workgroup_id) wid: vec3u,\n @builtin(num_workgroups) numWorkgroups: vec3u\n) {\n let largeTileIdx = wid.y * numWorkgroups.x + wid.x;\n if (largeTileIdx >= tileListCounts[1]) {\n return;\n }\n let tileIdx = largeTileList[largeTileIdx];\n let tStart = tileSplatCounts[tileIdx];\n let tEnd = tileSplatCounts[tileIdx + 1u];\n let count = tEnd - tStart;\n\n // Overflow scratch base for this tile (assigned by classify pass).\n // If it extends beyond buffer capacity, skip — tile renders unsorted.\n let overflowBase = largeTileOverflowBases[largeTileIdx];\n if (overflowBase + count > uniforms.bufferCapacity) {\n return;\n }\n\n // --- Phase 1: Find depth min/max ---\n if (localIdx == 0u) {\n atomicStore(&sDepthMin, 0xFFFFFFFFu);\n atomicStore(&sDepthMax, 0u);\n }\n if (localIdx < NUM_BUCKETS) {\n atomicStore(&sBucketCounts[localIdx], 0u);\n atomicStore(&sBucketCursors[localIdx], 0u);\n }\n\n workgroupBarrier();\n\n for (var i: u32 = localIdx; i < count; i += WG_SIZE) {\n let entryIdx = tileEntries[tStart + i];\n let depthU = depthBuffer[entryIdx];\n atomicMin(&sDepthMin, depthU);\n atomicMax(&sDepthMax, depthU);\n }\n\n workgroupBarrier();\n\n let depthMinU = atomicLoad(&sDepthMin);\n let depthMaxU = atomicLoad(&sDepthMax);\n let depthMin = bitcast<f32>(depthMinU);\n let depthMax = bitcast<f32>(depthMaxU);\n\n // Logarithmic bucketing: more precision for near depths, less for far.\n // Avoids oversized buckets when a dense surface (wall) occupies a narrow depth range.\n let logMin = log(max(depthMin, 1e-6));\n let logRange = log(max(depthMax, 1e-6)) - logMin;\n let bucketScale = select(f32(NUM_BUCKETS) / logRange, 0.0, logRange < 1e-10);\n\n // --- Phase 2: Histogram + save entry indices to overflow scratch ---\n // Write entryIdx to the overflow region so Phase 4 can read without aliasing\n // the main tileEntries range (which Phase 4 writes to).\n for (var i: u32 = localIdx; i < count; i += WG_SIZE) {\n let entryIdx = tileEntries[tStart + i];\n let depth = bitcast<f32>(depthBuffer[entryIdx]);\n let bucket = min(u32((log(max(depth, 1e-6)) - logMin) * bucketScale), NUM_BUCKETS - 1u);\n atomicAdd(&sBucketCounts[bucket], 1u);\n tileEntries[overflowBase + i] = entryIdx;\n }\n\n workgroupBarrier();\n\n // --- Phase 3: Prefix sum on bucket counts (thread 0, serial) ---\n if (localIdx == 0u) {\n sBucketOffsets[0] = 0u;\n for (var b: u32 = 0u; b < NUM_BUCKETS; b++) {\n sBucketOffsets[b + 1u] = sBucketOffsets[b] + atomicLoad(&sBucketCounts[b]);\n }\n }\n\n workgroupBarrier();\n\n // --- Phase 4: Scatter entries to tileEntries in bucket order ---\n // Read from overflow scratch, recompute bucket, scatter to tileEntries main range.\n for (var i: u32 = localIdx; i < count; i += WG_SIZE) {\n let entryIdx = tileEntries[overflowBase + i];\n let depth = bitcast<f32>(depthBuffer[entryIdx]);\n let bucket = min(u32((log(max(depth, 1e-6)) - logMin) * bucketScale), NUM_BUCKETS - 1u);\n let writePos = sBucketOffsets[bucket] + atomicAdd(&sBucketCursors[bucket], 1u);\n tileEntries[tStart + writePos] = entryIdx;\n }\n\n workgroupBarrier();\n\n // --- Phase 5: Thread 0 greedy-packs whole buckets into chunks ---\n // Buckets larger than MAX_CHUNK_SIZE are split into multiple pieces.\n // Chunk emission is bounds-checked against maxChunks; excess chunks are dropped\n // (those entries retain bucket-level ordering but skip the bitonic sort pass).\n if (localIdx == 0u) {\n var chunkStart: u32 = 0u;\n var currentSize: u32 = 0u;\n let maxChunks = uniforms.maxChunks;\n\n for (var b: u32 = 0u; b < NUM_BUCKETS; b++) {\n var bRemaining = sBucketOffsets[b + 1u] - sBucketOffsets[b];\n if (bRemaining == 0u) {\n continue;\n }\n\n // Split oversized buckets into MAX_CHUNK_SIZE pieces\n while (bRemaining > 0u) {\n let space = MAX_CHUNK_SIZE - currentSize;\n let take = min(bRemaining, space);\n currentSize += take;\n bRemaining -= take;\n\n if (currentSize == MAX_CHUNK_SIZE) {\n let cIdx = atomicAdd(&totalChunks[0], 1u);\n if (cIdx < maxChunks) {\n chunkRanges[cIdx * 2u] = tStart + chunkStart;\n chunkRanges[cIdx * 2u + 1u] = currentSize;\n }\n chunkStart += currentSize;\n currentSize = 0u;\n }\n }\n }\n\n if (currentSize > 0u) {\n let cIdx = atomicAdd(&totalChunks[0], 1u);\n if (cIdx < maxChunks) {\n chunkRanges[cIdx * 2u] = tStart + chunkStart;\n chunkRanges[cIdx * 2u + 1u] = currentSize;\n }\n }\n }\n}\n";