/
rayon.rs
106 lines (89 loc) · 3.21 KB
/
rayon.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
use bevy::{
prelude::*,
asset::LoadState,
utils::Instant,
};
use rayon::prelude::*;
use crate::{
GaussianCloud,
GaussianCloudSettings,
sort::{
SortedEntries,
SortMode,
},
};
#[derive(Default)]
pub struct RayonSortPlugin;
impl Plugin for RayonSortPlugin {
fn build(&self, app: &mut App) {
app.add_systems(Update, rayon_sort);
}
}
pub fn rayon_sort(
asset_server: Res<AssetServer>,
gaussian_clouds_res: Res<Assets<GaussianCloud>>,
mut sorted_entries_res: ResMut<Assets<SortedEntries>>,
gaussian_clouds: Query<(
&Handle<GaussianCloud>,
&Handle<SortedEntries>,
&GaussianCloudSettings,
)>,
cameras: Query<(
&GlobalTransform,
&Camera3d,
)>,
mut last_camera_position: Local<Vec3>,
mut last_sort_time: Local<Option<Instant>>,
) {
let period = std::time::Duration::from_millis(100);
if let Some(last_sort_time) = last_sort_time.as_ref() {
if last_sort_time.elapsed() < period {
return;
}
}
// TODO: move sort to render world, use extracted views and update the existing buffer instead of creating new
for (
camera_transform,
_camera,
) in cameras.iter() {
let camera_position = camera_transform.compute_transform().translation;
if *last_camera_position == camera_position {
return;
}
for (
gaussian_cloud_handle,
sorted_entries_handle,
settings,
) in gaussian_clouds.iter() {
if settings.sort_mode != SortMode::Rayon {
continue;
}
if Some(LoadState::Loading) == asset_server.get_load_state(gaussian_cloud_handle) {
continue;
}
if Some(LoadState::Loading) == asset_server.get_load_state(sorted_entries_handle) {
continue;
}
if let Some(gaussian_cloud) = gaussian_clouds_res.get(gaussian_cloud_handle) {
if let Some(sorted_entries) = sorted_entries_res.get_mut(sorted_entries_handle) {
assert_eq!(gaussian_cloud.gaussians.len(), sorted_entries.sorted.len());
*last_camera_position = camera_position;
*last_sort_time = Some(Instant::now());
gaussian_cloud.gaussians.par_iter()
.zip(sorted_entries.sorted.par_iter_mut())
.enumerate()
.for_each(|(idx, (gaussian, sort_entry))| {
let position = Vec3::from_slice(gaussian.position.as_ref());
let delta = camera_position - position;
sort_entry.key = bytemuck::cast(delta.length_squared());
sort_entry.index = idx as u32;
});
sorted_entries.sorted.par_sort_unstable_by(|a, b| {
bytemuck::cast::<u32, f32>(b.key).partial_cmp(&bytemuck::cast::<u32, f32>(a.key)).unwrap()
});
// TODO: update DrawIndirect buffer during sort phase (GPU sort will override default DrawIndirect)
}
}
}
}
}