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