Bevy Shader Explorer
Search across packages, versions, files, declarations, and comments. Use multiple terms to narrow results (AND), and add @vertex, @fragment, or @compute to filter by shader stage.
Packages › bevy_gaussian_splatting › radix.wgsl › source
← Back to declarations 
1#import bevy_gaussian_splatting::bindings::{
2    view,
3    globals,
4    gaussian_uniforms,
5    sorting_pass_index,
6    sorting,
7    status_counters,
8    draw_indirect,
9    input_entries,
10    output_entries,
11    sorted_entries,
12    DrawIndirect,
13    Entry,
15#import bevy_gaussian_splatting::transform::{
16    world_to_clip,
17    in_frustum,
20#ifdef PACKED_F32
21#import bevy_gaussian_splatting::packed::get_position
22#else
24#ifdef BUFFER_STORAGE
25#import bevy_gaussian_splatting::planar::get_position
26#endif
28#endif
30#ifdef BUFFER_TEXTURE
31#import bevy_gaussian_splatting::texture::get_position
32#endif
35struct SortingGlobal {
36    digit_histogram: array<array<atomic<u32>, #{RADIX_BASE}>, #{RADIX_DIGIT_PLACES}>,
37    assignment_counter: atomic<u32>,
40@group(3) @binding(0) var<uniform> sorting_pass_index: u32;
41@group(3) @binding(1) var<storage, read_write> sorting: SortingGlobal;
42@group(3) @binding(2) var<storage, read_write> status_counters: array<array<atomic<u32>, #{RADIX_BASE}>>;
43@group(3) @binding(3) var<storage, read_write> draw_indirect: DrawIndirect;
44@group(3) @binding(4) var<storage, read_write> input_entries: array<Entry>;
45@group(3) @binding(5) var<storage, read_write> output_entries: array<Entry>;
48struct SortingSharedA {
49    digit_histogram: array<array<atomic<u32>, #{RADIX_BASE}>, #{RADIX_DIGIT_PLACES}>,
51var<workgroup> sorting_shared_a: SortingSharedA;
53@compute @workgroup_size(#{RADIX_BASE}, #{RADIX_DIGIT_PLACES})
54fn radix_sort_a(
55    @builtin(local_invocation_id) gl_LocalInvocationID: vec3<u32>,
56    @builtin(global_invocation_id) gl_GlobalInvocationID: vec3<u32>,
57) {
58    sorting_shared_a.digit_histogram[gl_LocalInvocationID.y][gl_LocalInvocationID.x] = 0u;
59    workgroupBarrier();
61    let thread_index = gl_GlobalInvocationID.x * #{RADIX_DIGIT_PLACES}u + gl_GlobalInvocationID.y;
62    let start_entry_index = thread_index * #{ENTRIES_PER_INVOCATION_A}u;
63    let end_entry_index = start_entry_index + #{ENTRIES_PER_INVOCATION_A}u;
64    for(var entry_index = start_entry_index; entry_index < end_entry_index; entry_index += 1u) {
65        if(entry_index >= gaussian_uniforms.count) {
66            continue;
67        }
68        var key: u32 = 0xFFFFFFFFu; // Stream compaction for frustum culling
69        let position = vec4<f32>(get_position(entry_index), 1.0);
70        let transformed_position = (gaussian_uniforms.transform * position).xyz;
71        let clip_space_pos = world_to_clip(transformed_position);
72        if(in_frustum(clip_space_pos.xyz)) {
73            // key = bitcast<u32>(1.0 - clip_space_pos.z);
74            key = u32(clip_space_pos.z * 0xFFFF.0) << 16u;
75            key |= u32((clip_space_pos.x * 0.5 + 0.5) * 0xFF.0) << 8u;
76            key |= u32((clip_space_pos.y * 0.5 + 0.5) * 0xFF.0);
77        }
78        output_entries[entry_index].key = key;
79        output_entries[entry_index].value = entry_index;
80        for(var shift = 0u; shift < #{RADIX_DIGIT_PLACES}u; shift += 1u) {
81            let digit = (key >> (shift * #{RADIX_BITS_PER_DIGIT}u)) & (#{RADIX_BASE}u - 1u);
82            atomicAdd(&sorting_shared_a.digit_histogram[shift][digit], 1u);
83        }
84    }
85    workgroupBarrier();
87    atomicAdd(&sorting.digit_histogram[gl_LocalInvocationID.y][gl_LocalInvocationID.x], sorting_shared_a.digit_histogram[gl_LocalInvocationID.y][gl_LocalInvocationID.x]);
90@compute @workgroup_size(1)
91fn radix_sort_b(
92    @builtin(global_invocation_id) gl_GlobalInvocationID: vec3<u32>,
93) {
94    var sum = 0u;
95    for(var digit = 0u; digit < #{RADIX_BASE}u; digit += 1u) {
96        let tmp = sorting.digit_histogram[gl_GlobalInvocationID.y][digit];
97        sorting.digit_histogram[gl_GlobalInvocationID.y][digit] = sum;
98        sum += tmp;
99    }
102struct SortingSharedC {
103    entries: array<atomic<u32>, #{WORKGROUP_ENTRIES_C}>,
104    gather_sources: array<atomic<u32>, #{WORKGROUP_ENTRIES_C}>,
105    scan: array<atomic<u32>, #{WORKGROUP_INVOCATIONS_C}>,
106    total: u32,
108var<workgroup> sorting_shared_c: SortingSharedC;
110const NUM_BANKS: u32 = 16u;
111const LOG_NUM_BANKS: u32 = 4u;
112fn conflict_free_offset(n: u32) -> u32 {
113    return 0u;//n >> NUM_BANKS + n >> (2u * LOG_NUM_BANKS);
116fn exclusive_scan(local_invocation_index: u32, value: u32) -> u32 {
117    sorting_shared_c.scan[local_invocation_index + conflict_free_offset(local_invocation_index)] = value;
119    var offset = 1u;
120    for (var d = #{WORKGROUP_INVOCATIONS_C}u >> 1u; d > 0u; d >>= 1u) {
121        workgroupBarrier();
122        if(local_invocation_index < d) {
123            var ai = offset * (2u * local_invocation_index + 1u) - 1u;
124            var bi = offset * (2u * local_invocation_index + 2u) - 1u;
125            ai += conflict_free_offset(ai);
126            bi += conflict_free_offset(bi);
127            sorting_shared_c.scan[bi] += sorting_shared_c.scan[ai];
128        }
130        offset <<= 1u;
131    }
133    if (local_invocation_index == 0u) {
134      var i = #{WORKGROUP_INVOCATIONS_C}u - 1u;
135      i += conflict_free_offset(i);
136      sorting_shared_c.total = sorting_shared_c.scan[i];
137      sorting_shared_c.scan[i] = 0u;
138    }
140    for (var d = 1u; d < #{WORKGROUP_INVOCATIONS_C}u; d <<= 1u) {
141        workgroupBarrier();
142        offset >>= 1u;
143        if(local_invocation_index < d) {
144            var ai = offset * (2u * local_invocation_index + 1u) - 1u;
145            var bi = offset * (2u * local_invocation_index + 2u) - 1u;
146            ai += conflict_free_offset(ai);
147            bi += conflict_free_offset(bi);
148            let t = sorting_shared_c.scan[ai];
149            sorting_shared_c.scan[ai] = sorting_shared_c.scan[bi];
150            sorting_shared_c.scan[bi] += t;
151        }
152    }
154    workgroupBarrier();
155    return sorting_shared_c.scan[local_invocation_index + conflict_free_offset(local_invocation_index)];
158@compute @workgroup_size(#{WORKGROUP_INVOCATIONS_C})
159fn radix_sort_c(
160    @builtin(local_invocation_id) gl_LocalInvocationID: vec3<u32>,
161    @builtin(global_invocation_id) gl_GlobalInvocationID: vec3<u32>,
162) {
163    // Draw an assignment number
164    if(gl_LocalInvocationID.x == 0u) {
165        sorting_shared_c.entries[0] = atomicAdd(&sorting.assignment_counter, 1u);
166    }
168    // Reset histogram
169    sorting_shared_c.scan[gl_LocalInvocationID.x + conflict_free_offset(gl_LocalInvocationID.x)] = 0u;
170    workgroupBarrier();
172    let assignment = sorting_shared_c.entries[0];
173    let global_entry_offset = assignment * #{WORKGROUP_ENTRIES_C}u;
174    // TODO: Specialize end shader
175    if(gl_LocalInvocationID.x == 0u && assignment * #{WORKGROUP_ENTRIES_C}u + #{WORKGROUP_ENTRIES_C}u >= gaussian_uniforms.count) {
176        // Last workgroup resets the assignment number for the next pass
177        sorting.assignment_counter = 0u;
178    }
180    // Load keys from global memory into registers and rank them
181    var keys: array<u32, #{ENTRIES_PER_INVOCATION_C}>;
182    var ranks: array<u32, #{ENTRIES_PER_INVOCATION_C}>;
183    for(var entry_index = 0u; entry_index < #{ENTRIES_PER_INVOCATION_C}u; entry_index += 1u) {
184        keys[entry_index] = input_entries[global_entry_offset + #{WORKGROUP_INVOCATIONS_C}u * entry_index + gl_LocalInvocationID.x].key;
185        let digit = (keys[entry_index] >> (sorting_pass_index * #{RADIX_BITS_PER_DIGIT}u)) & (#{RADIX_BASE}u - 1u);
186        // TODO: Implement warp-level multi-split (WLMS) once WebGPU supports subgroup operations
187        ranks[entry_index] = atomicAdd(&sorting_shared_c.scan[digit + conflict_free_offset(digit)], 1u);
188    }
189    workgroupBarrier();
191    // Cumulate histogram
192    let local_digit_count = sorting_shared_c.scan[gl_LocalInvocationID.x + conflict_free_offset(gl_LocalInvocationID.x)];
193    let local_digit_offset = exclusive_scan(gl_LocalInvocationID.x, local_digit_count);
194    sorting_shared_c.scan[gl_LocalInvocationID.x + conflict_free_offset(gl_LocalInvocationID.x)] = local_digit_offset;
196    // Chained decoupling lookback
197    atomicStore(&status_counters[assignment][gl_LocalInvocationID.x], 0x40000000u | local_digit_count);
198    var global_digit_count = 0u;
199    var previous_tile = assignment;
200    while true {
201        if(previous_tile == 0u) {
202            global_digit_count += sorting.digit_histogram[sorting_pass_index][gl_LocalInvocationID.x];
203            break;
204        }
205        previous_tile -= 1u;
206        var status_counter = 0u;
207        while((status_counter & 0xC0000000u) == 0u) {
208            status_counter = atomicLoad(&status_counters[previous_tile][gl_LocalInvocationID.x]);
209        }
210        global_digit_count += status_counter & 0x3FFFFFFFu;
211        if((status_counter & 0x80000000u) != 0u) {
212            break;
213        }
214    }
215    atomicStore(&status_counters[assignment][gl_LocalInvocationID.x], 0x80000000u | (global_digit_count + local_digit_count));
216    if(sorting_pass_index == #{RADIX_DIGIT_PLACES}u - 1u && gl_LocalInvocationID.x == #{WORKGROUP_INVOCATIONS_C}u - 2u && global_entry_offset + #{WORKGROUP_ENTRIES_C}u >= gaussian_uniforms.count) {
217        draw_indirect.vertex_count = 4u;
218        draw_indirect.instance_count = global_digit_count + local_digit_count;
219    }
221    // Scatter keys inside shared memory
222    for(var entry_index = 0u; entry_index < #{ENTRIES_PER_INVOCATION_C}u; entry_index += 1u) {
223        let key = keys[entry_index];
224        let digit = (key >> (sorting_pass_index * #{RADIX_BITS_PER_DIGIT}u)) & (#{RADIX_BASE}u - 1u);
225        ranks[entry_index] += sorting_shared_c.scan[digit + conflict_free_offset(digit)];
226        sorting_shared_c.entries[ranks[entry_index]] = key;
227    }
228    workgroupBarrier();
230    // Add global offset
231    sorting_shared_c.scan[gl_LocalInvocationID.x + conflict_free_offset(gl_LocalInvocationID.x)] = global_digit_count - local_digit_offset;
232    workgroupBarrier();
234    // Store keys from shared memory into global memory
235    for(var entry_index = 0u; entry_index < #{ENTRIES_PER_INVOCATION_C}u; entry_index += 1u) {
236        let key = sorting_shared_c.entries[#{WORKGROUP_INVOCATIONS_C}u * entry_index + gl_LocalInvocationID.x];
237        let digit = (key >> (sorting_pass_index * #{RADIX_BITS_PER_DIGIT}u)) & (#{RADIX_BASE}u - 1u);
238        keys[entry_index] = digit;
239        output_entries[sorting_shared_c.scan[digit + conflict_free_offset(digit)] + #{WORKGROUP_INVOCATIONS_C}u * entry_index + gl_LocalInvocationID.x].key = key;
240    }
241    workgroupBarrier();
243    // Load values from global memory and scatter them inside shared memory
244    for(var entry_index = 0u; entry_index < #{ENTRIES_PER_INVOCATION_C}u; entry_index += 1u) {
245        let value = input_entries[global_entry_offset + #{WORKGROUP_INVOCATIONS_C}u * entry_index + gl_LocalInvocationID.x].value;
246        sorting_shared_c.entries[ranks[entry_index]] = value;
247    }
248    workgroupBarrier();
250    // Store values from shared memory into global memory
251    for(var entry_index = 0u; entry_index < #{ENTRIES_PER_INVOCATION_C}u; entry_index += 1u) {
252        let value = sorting_shared_c.entries[#{WORKGROUP_INVOCATIONS_C}u * entry_index + gl_LocalInvocationID.x];
253        let digit = keys[entry_index];
254        output_entries[sorting_shared_c.scan[digit + conflict_free_offset(digit)] + #{WORKGROUP_INVOCATIONS_C}u * entry_index + gl_LocalInvocationID.x][1] = value;
255    }
Contents