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,
14}
15#import bevy_gaussian_splatting::transform::{
16    world_to_clip,
17    in_frustum,
18}
19
20#ifdef PACKED_F32
21#import bevy_gaussian_splatting::packed::get_position
22#else
23
24#ifdef BUFFER_STORAGE
25#import bevy_gaussian_splatting::planar::get_position
26#endif
27
28#endif
29
30#ifdef BUFFER_TEXTURE
31#import bevy_gaussian_splatting::texture::get_position
32#endif
33
34struct SortingGlobal {
35    digit_histogram: array<array<atomic<u32>, #{RADIX_BASE}>, #{RADIX_DIGIT_PLACES}>,
36    assignment_counter: atomic<u32>,
37    digit_tile_head: array<atomic<u32>, #{RADIX_BASE}>,
38}
39
40@group(3) @binding(0) var<uniform> sorting_pass_index: u32;
41@group(3) @binding(1) var<storage, read_write> sorting: SortingGlobal;
42// NOTE: status_counters at binding(2) is NO LONGER USED by the corrected shader.
43// It can be removed from the Rust host code.
44@group(3) @binding(2) var<storage, read_write> status_counters: array<array<atomic<u32>, #{RADIX_BASE}>>;
45@group(3) @binding(3) var<storage, read_write> draw_indirect: DrawIndirect;
46@group(3) @binding(4) var<storage, read_write> input_entries: array<Entry>;
47@group(3) @binding(5) var<storage, read_write> output_entries: array<Entry>;
48
49//
50// The following three functions (`radix_reset`, `radix_sort_a`, `radix_sort_b`)
51// form a standard three-phase GPU sort setup and were already correct.
52// They are included here without changes.
53//
54
55@compute @workgroup_size(#{RADIX_BASE}, #{RADIX_DIGIT_PLACES})
56fn radix_reset(
57    @builtin(local_invocation_id) local_id: vec3<u32>,
58    @builtin(global_invocation_id) global_id: vec3<u32>,
59){
60    let b = local_id.x;
61    let p = local_id.y;
62    atomicStore(&sorting.digit_histogram[p][b], 0u);
63    if (p == 0u) {
64        atomicStore(&sorting.digit_tile_head[b], 0u);
65    }
66    if (global_id.x == 0u && global_id.y == 0u) {
67        atomicStore(&sorting.assignment_counter, 0u);
68        draw_indirect.instance_count = 0u;
69    }
70}
71
72@compute @workgroup_size(#{RADIX_BASE}, #{RADIX_DIGIT_PLACES})
73fn radix_sort_a(
74    @builtin(local_invocation_id) gl_LocalInvocationID: vec3<u32>,
75    @builtin(global_invocation_id) gl_GlobalInvocationID: vec3<u32>,
76) {
77    if (gl_LocalInvocationID.x == 0u && gl_LocalInvocationID.y == 0u && gl_GlobalInvocationID.x == 0u) {
78        draw_indirect.vertex_count = 4u;
79        atomicStore(&draw_indirect.instance_count, gaussian_uniforms.count);
80    }
81    workgroupBarrier();
82
83    let thread_index = gl_GlobalInvocationID.x * #{RADIX_DIGIT_PLACES}u + gl_GlobalInvocationID.y;
84    let start_entry_index = thread_index * #{ENTRIES_PER_INVOCATION_A}u;
85    let end_entry_index = start_entry_index + #{ENTRIES_PER_INVOCATION_A}u;
86
87    for (var entry_index = start_entry_index; entry_index < end_entry_index; entry_index += 1u) {
88        if (entry_index >= gaussian_uniforms.count) { continue; }
89        var key: u32 = 0xFFFFFFFFu;
90        let position = vec4<f32>(get_position(entry_index), 1.0);
91        let transformed_position = (gaussian_uniforms.transform * position).xyz;
92        let clip_space_pos = world_to_clip(transformed_position);
93        let diff = transformed_position - view.world_position;
94        let dist2 = dot(diff, diff);
95        let dist_bits = bitcast<u32>(dist2);
96        let key_distance = 0xFFFFFFFFu - dist_bits;
97        if (in_frustum(clip_space_pos.xyz)) {
98            key = key_distance;
99        }
100        input_entries[entry_index].key = key;
101        input_entries[entry_index].value = entry_index;
102        for(var shift = 0u; shift < #{RADIX_DIGIT_PLACES}u; shift += 1u) {
103            let digit = (key >> (shift * #{RADIX_BITS_PER_DIGIT}u)) & (#{RADIX_BASE}u - 1u);
104            atomicAdd(&sorting.digit_histogram[shift][digit], 1u);
105        }
106    }
107}
108
109@compute @workgroup_size(1)
110fn radix_sort_b(
111    @builtin(global_invocation_id) gl_GlobalInvocationID: vec3<u32>,
112) {
113    var sum = 0u;
114    for(var digit = 0u; digit < #{RADIX_BASE}u; digit += 1u) {
115        let tmp = atomicLoad(&sorting.digit_histogram[gl_GlobalInvocationID.y][digit]);
116        atomicStore(&sorting.digit_histogram[gl_GlobalInvocationID.y][digit], sum);
117        sum += tmp;
118    }
119}
120
121// --- SHARED MEMORY for the final, stable `radix_sort_c` ---
122var<workgroup> tile_input_entries: array<Entry, #{WORKGROUP_ENTRIES_C}>;
123var<workgroup> sorted_tile_entries: array<Entry, #{WORKGROUP_ENTRIES_C}>;
124var<workgroup> local_digit_counts: array<u32, #{RADIX_BASE}>;
125var<workgroup> local_digit_offsets: array<u32, #{RADIX_BASE}>;
126var<workgroup> digit_global_base_ws: array<u32, #{RADIX_BASE}>;
127var<workgroup> tile_entry_count_ws: u32;
128const INVALID_KEY: u32 = 0xFFFFFFFFu;
129
130//
131// Pass C (REWRITTEN): A fully stable implementation that discards the faulty spin-lock.
132//
133@compute @workgroup_size(#{WORKGROUP_INVOCATIONS_C})
134fn radix_sort_c(
135    @builtin(local_invocation_id) local_id: vec3<u32>,
136    @builtin(workgroup_id) workgroup_id: vec3<u32>,
137) {
138    let tid = local_id.x;
139    let tile_size = #{WORKGROUP_ENTRIES_C}u;
140    let threads = #{WORKGROUP_INVOCATIONS_C}u;
141    let global_entry_offset = workgroup_id.y * tile_size;
142
143    // Clear per-digit base cache so stale values are never reused across tiles.
144    if (tid < #{RADIX_BASE}u) {
145        digit_global_base_ws[tid] = 0u;
146    }
147    workgroupBarrier();
148
149    // --- Step 1: Parallel load ---
150    for (var i = tid; i < tile_size; i += threads) {
151        let idx = global_entry_offset + i;
152        if (idx < gaussian_uniforms.count) {
153            tile_input_entries[i] = input_entries[idx];
154        } else {
155            tile_input_entries[i] = Entry(INVALID_KEY, INVALID_KEY);
156        }
157    }
158    workgroupBarrier();
159
160    // --- Step 2: Serial, stable sort within the tile by a single thread ---
161    if (tid == 0u) {
162        for (var i = 0u; i < #{RADIX_BASE}u; i+=1u) { local_digit_counts[i] = 0u; }
163
164        var entries_in_tile = 0u;
165        for (var i = 0u; i < tile_size; i+=1u) {
166            let entry = tile_input_entries[i];
167            if (entry.value == INVALID_KEY) { continue; } // value sentinel marks padding
168
169            let digit = (entry.key >> (sorting_pass_index * #{RADIX_BITS_PER_DIGIT}u)) & (#{RADIX_BASE}u - 1u);
170            local_digit_counts[digit] += 1u;
171            entries_in_tile += 1u;
172        }
173        tile_entry_count_ws = entries_in_tile;
174
175        var sum = 0u;
176        for (var i = 0u; i < #{RADIX_BASE}u; i+=1u) {
177            local_digit_offsets[i] = sum;
178            sum += local_digit_counts[i];
179        }
180
181        for (var i = 0u; i < tile_size; i+=1u) {
182            let entry = tile_input_entries[i];
183            if (entry.value == INVALID_KEY) { continue; } // value sentinel marks padding
184
185            let digit = (entry.key >> (sorting_pass_index * #{RADIX_BITS_PER_DIGIT}u)) & (#{RADIX_BASE}u - 1u);
186            let dest_idx = local_digit_offsets[digit];
187            local_digit_offsets[digit] = dest_idx + 1u;
188            sorted_tile_entries[dest_idx] = entry;
189        }
190    }
191    workgroupBarrier();
192
193    // --- Step 3: Determine deterministic global base for each digit ---
194    let tile_count = max((gaussian_uniforms.count + tile_size - 1u) / tile_size, 1u);
195    let expected_ticket = sorting_pass_index * tile_count + workgroup_id.y;
196
197    // Acquire a per-tile ticket so we only serialize once per tile instead of once per digit.
198    if (tid == 0u) {
199        loop {
200            let head = atomicLoad(&sorting.assignment_counter);
201            if (head == expected_ticket) {
202                let exchange = atomicCompareExchangeWeak(&sorting.assignment_counter, expected_ticket, expected_ticket + 1u);
203                if (exchange.exchanged) { break; }
204            }
205        }
206    }
207    workgroupBarrier();
208
209    if (tid < #{RADIX_BASE}u) {
210        let count = local_digit_counts[tid];
211        let base = atomicAdd(&sorting.digit_histogram[sorting_pass_index][tid], count);
212        digit_global_base_ws[tid] = base;
213    }
214    workgroupBarrier();
215
216    // --- Step 4: Parallel write from the locally-sorted tile to global memory ---
217    if (tid == 0u) {
218        var sum = 0u;
219        for (var i = 0u; i < #{RADIX_BASE}u; i += 1u) {
220            local_digit_offsets[i] = sum;
221            sum += local_digit_counts[i];
222        }
223    }
224    workgroupBarrier();
225
226    for (var i = tid; i < tile_size; i += threads) {
227        if (i < tile_entry_count_ws) {
228            let entry = sorted_tile_entries[i];
229            let digit = (entry.key >> (sorting_pass_index * #{RADIX_BITS_PER_DIGIT}u)) & (#{RADIX_BASE}u - 1u);
230
231            let bin_start_offset = local_digit_offsets[digit];
232            let rank_in_bin = i - bin_start_offset;
233            let global_base = digit_global_base_ws[digit];
234            let dst = global_base + rank_in_bin;
235
236            if (dst < gaussian_uniforms.count) {
237                output_entries[dst] = entry;
238            }
239        }
240    }
241    if (sorting_pass_index == #{RADIX_DIGIT_PLACES}u - 1u && tid == 0u) {
242        atomicStore(&draw_indirect.instance_count, gaussian_uniforms.count);
243    }
244}