src/stg_prepare.rs
与本次数学页面一起发布的源码快照。SHA-256: 7f9d9ae9e00637cd37636ee99598d8ed006ec9bb3f4aa66d3f3b1c244ea5b6ca
1//! Reproducible CPU preparation variants. Variant 0 preserves the native baseline.
2//! Source IDs and the 1/255 visibility threshold never change between variants.
3use crate::{parse_stg_ply, sigmoid, ResearchSplat};
4use glam::{Mat4, Vec3};
5
6const CUT: f32 = 1.0 / 255.0;
7const BUCKETS: usize = 64;
8
9pub struct StgPreparation {
10 source: Vec<ResearchSplat>,
11 opacity: Vec<f32>,
12 radius: Vec<f32>,
13 inner: Vec<f32>,
14 buckets: Vec<Vec<u32>>,
15 ordered: Vec<(f32, u32)>,
16 scratch: Vec<(f32, u32)>,
17 indices: Vec<u32>,
18 sampled: Vec<(Vec3, u32)>,
19 last_sample: Option<f32>,
20 last_frame: Option<(f32, Mat4, Mat4, u8)>,
21 pub candidates: usize,
22 pub cache_hit: bool,
23 pub sort_ms: f64,
24}
25
26impl StgPreparation {
27 pub fn from_ply(bytes: &[u8]) -> Result<Self, String> {
28 Ok(Self::build(parse_stg_ply(bytes)?.source, true))
29 }
30
31 pub(crate) fn new(source: Vec<ResearchSplat>) -> Self {
32 Self::build(source, false)
33 }
34 pub fn from_ply_adaptive(bytes: &[u8]) -> Result<Self, String> {
35 Ok(Self::new(parse_stg_ply(bytes)?.source))
36 }
37 pub fn temporal_index_bytes(&self) -> usize {
38 self.buckets.iter().map(|b| b.len() * 4).sum()
39 }
40 fn build(source: Vec<ResearchSplat>, force_index: bool) -> Self {
41 let opacity: Vec<_> = source.iter().map(|s| sigmoid(s.opacity_logit)).collect();
42 // Expand the analytic support slightly. The exact baseline test is still
43 // evaluated afterwards; this index must only exclude definitely invisible IDs.
44 let radius: Vec<_> = source
45 .iter()
46 .zip(&opacity)
47 .map(|(s, &a)| {
48 if a < CUT {
49 -1.0
50 } else {
51 s.temporal_scale * (a / (CUT * (1.0 - 1e-5))).ln().max(0.0).sqrt() + 1e-5
52 }
53 })
54 .collect();
55 let inner: Vec<_> = source
56 .iter()
57 .zip(&opacity)
58 .map(|(s, &a)| {
59 if a <= CUT * (1.0 + 1e-5) {
60 -1.0
61 } else {
62 (s.temporal_scale * (a / (CUT * (1.0 + 1e-5))).ln().sqrt() - 1e-5).max(-1.0)
63 }
64 })
65 .collect();
66 let coverage: f64 = radius
67 .iter()
68 .zip(&source)
69 .filter(|(r, _)| **r >= 0.0)
70 .map(|(&r, s)| {
71 ((s.temporal_center + r).min(1.0) - (s.temporal_center - r).max(0.0)).max(0.0)
72 as f64
73 })
74 .sum::<f64>()
75 / source.len().max(1) as f64;
76 let indexed = force_index || coverage < 0.65;
77 let mut buckets = if indexed {
78 vec![Vec::new(); BUCKETS]
79 } else {
80 Vec::new()
81 };
82 for (i, (s, &r)) in source.iter().zip(&radius).enumerate() {
83 if r < 0.0 {
84 continue;
85 }
86 for (b, ids) in buckets.iter_mut().enumerate() {
87 if s.temporal_center + r >= b as f32 / BUCKETS as f32
88 && s.temporal_center - r <= (b + 1) as f32 / BUCKETS as f32
89 {
90 ids.push(i as u32);
91 }
92 }
93 }
94 let n = source.len();
95 Self {
96 source,
97 opacity,
98 radius,
99 inner,
100 buckets,
101 ordered: Vec::with_capacity(n),
102 scratch: Vec::with_capacity(n),
103 indices: Vec::with_capacity(n),
104 sampled: Vec::with_capacity(n),
105 last_sample: None,
106 last_frame: None,
107 candidates: 0,
108 cache_hit: false,
109 sort_ms: 0.0,
110 }
111 }
112
113 pub fn len(&self) -> usize {
114 self.source.len()
115 }
116 pub fn is_empty(&self) -> bool {
117 self.source.is_empty()
118 }
119 pub fn indices(&self) -> &[u32] {
120 &self.indices
121 }
122 pub fn ordered(&self) -> &[(f32, u32)] {
123 &self.ordered
124 }
125 pub fn invalidate(&mut self) {
126 self.last_frame = None;
127 }
128
129 /// 0 baseline; 1 cached sigmoid; 2 temporal support rejection;
130 /// 3 temporal buckets; 4 exact radix order; 5 repeated-frame cache;
131 /// 6 reuse temporal samples when only the camera changes.
132 pub fn prepare(&mut self, time: f32, view: Mat4, projection: Mat4, variant: u8) -> &[u32] {
133 let key = (time, view, projection, variant);
134 self.cache_hit = variant >= 5 && self.last_frame == Some(key);
135 if self.cache_hit {
136 self.sort_ms = 0.0;
137 return &self.indices;
138 }
139 let vp = projection * view;
140 self.ordered.clear();
141 let bucket = ((time.clamp(0.0, 1.0) * BUCKETS as f32) as usize).min(BUCKETS - 1);
142 let capture_samples = variant >= 6 && (variant < 9 || self.last_sample == Some(time));
143 let indexed = variant >= 3 && !self.buckets.is_empty();
144 let reuse = variant >= 6 && self.last_sample == Some(time) && !self.sampled.is_empty();
145 if !reuse {
146 self.sampled.clear();
147 self.candidates = if indexed {
148 self.buckets[bucket].len()
149 } else {
150 self.source.len()
151 };
152 for j in 0..self.candidates {
153 let i = if indexed {
154 self.buckets[bucket][j] as usize
155 } else {
156 j
157 };
158 let s = self.source[i];
159 let dt = time - s.temporal_center;
160 if variant >= 2 && (self.radius[i] < 0.0 || dt.abs() > self.radius[i]) {
161 continue;
162 }
163 let (center, a) = if variant == 0 {
164 s.sample(time)
165 } else {
166 let dt2 = dt * dt;
167 let center = s.base
168 + s.motion_linear * dt
169 + s.motion_quadratic * dt2
170 + s.motion_cubic * (dt2 * dt);
171 let a = if variant >= 7 && dt.abs() < self.inner[i] {
172 1.0
173 } else {
174 self.opacity[i] * (-(dt / s.temporal_scale).powi(2)).exp()
175 };
176 (center, a)
177 };
178 if a < CUT {
179 continue;
180 }
181 if capture_samples {
182 self.sampled.push((center, i as u32));
183 }
184 project(&mut self.ordered, center, i as u32, view, vp);
185 }
186 } else {
187 for &(center, id) in &self.sampled {
188 project(&mut self.ordered, center, id, view, vp);
189 }
190 }
191 self.last_sample = if variant >= 6 { Some(time) } else { None };
192 let sort_started = web_time::Instant::now();
193 if variant >= 4 {
194 if variant >= 8 {
195 radix11(&mut self.ordered, &mut self.scratch);
196 } else {
197 radix(&mut self.ordered, &mut self.scratch);
198 }
199 } else {
200 self.ordered
201 .sort_unstable_by(|a, b| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1)));
202 }
203 self.sort_ms = sort_started.elapsed().as_secs_f64() * 1000.0;
204 self.indices.clear();
205 self.indices.extend(self.ordered.iter().map(|x| x.1));
206 self.last_frame = Some(key);
207 &self.indices
208 }
209}
210
211fn project(out: &mut Vec<(f32, u32)>, center: Vec3, id: u32, view: Mat4, vp: Mat4) {
212 let clip = vp * center.extend(1.0);
213 if clip.w <= 0.0 {
214 return;
215 }
216 let ndc = clip.truncate() / clip.w;
217 if ndc.x.abs() > 1.35 || ndc.y.abs() > 1.35 || ndc.z < 0.0 || ndc.z > 1.0 {
218 return;
219 }
220 out.push((-(view * center.extend(1.0)).z, id));
221}
222
223// Stable LSD sort. Input IDs are ascending, so equal depths retain the baseline
224// ID tie-break. Float key also handles signed depths and IEEE total-order values.
225fn radix(values: &mut Vec<(f32, u32)>, temp: &mut Vec<(f32, u32)>) {
226 temp.resize(values.len(), (0.0, 0));
227 let key = |x: f32| {
228 let bits = x.to_bits();
229 !(bits
230 ^ if bits >> 31 != 0 {
231 u32::MAX
232 } else {
233 0x8000_0000
234 })
235 };
236 for shift in [0, 8, 16, 24] {
237 let mut counts = [0usize; 256];
238 for &(depth, _) in values.iter() {
239 counts[((key(depth) >> shift) & 255) as usize] += 1;
240 }
241 let mut sum = 0;
242 for count in &mut counts {
243 let n = *count;
244 *count = sum;
245 sum += n;
246 }
247 for &value in values.iter() {
248 let bucket = ((key(value.0) >> shift) & 255) as usize;
249 temp[counts[bucket]] = value;
250 counts[bucket] += 1;
251 }
252 std::mem::swap(values, temp);
253 }
254}
255
256#[cfg(test)]
257mod tests {
258 use super::*;
259 #[test]
260 fn radix_matches_total_order_and_ties() {
261 let mut values = vec![
262 (1.0, 0),
263 (-3.0, 1),
264 (1.0, 2),
265 (0.0, 3),
266 (-0.0, 4),
267 (f32::INFINITY, 5),
268 ];
269 let mut expected = values.clone();
270 expected.sort_unstable_by(|a, b| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1)));
271 radix(&mut values, &mut Vec::new());
272 assert_eq!(
273 values.iter().map(|x| x.1).collect::<Vec<_>>(),
274 expected.iter().map(|x| x.1).collect::<Vec<_>>()
275 );
276 }
277}
278
279fn radix11(values: &mut Vec<(f32, u32)>, temp: &mut Vec<(f32, u32)>) {
280 temp.resize(values.len(), (0.0, 0));
281 let key = |x: f32| {
282 let b = x.to_bits();
283 !(b ^ if b >> 31 != 0 { u32::MAX } else { 0x8000_0000 })
284 };
285 for shift in [0, 11, 22] {
286 let mut counts = [0usize; 2048];
287 for &(depth, _) in values.iter() {
288 counts[((key(depth) >> shift) & 2047) as usize] += 1;
289 }
290 let mut sum = 0;
291 for c in &mut counts {
292 let n = *c;
293 *c = sum;
294 sum += n;
295 }
296 for &v in values.iter() {
297 let b = ((key(v.0) >> shift) & 2047) as usize;
298 temp[counts[b]] = v;
299 counts[b] += 1;
300 }
301 std::mem::swap(values, temp);
302 }
303}
304