spatialrust_registration/
gicp_reduction.rs

1use super::{delta_from_solution, update_magnitude, GicpRegistration};
2use crate::RegistrationResult;
3use spatialrust_core::{
4    HasPositions3, PointCloud, SpatialError, SpatialResult, TransferDirection, TransferStats,
5};
6use spatialrust_gpu::{GpuGicpEquations, GpuGicpInput, WgpuRuntime};
7use spatialrust_math::{solve_linear_system, Isometry3, LeastSquaresResult};
8use spatialrust_search::KdTree;
9
10/// One completed update in the explicit resident GICP precision path.
11#[derive(Clone, Debug, PartialEq)]
12pub struct GpuGicpIteration {
13    /// One-based completed update number.
14    pub iteration: usize,
15    /// Gated pairs used by the normal-equation solve.
16    pub correspondences: u32,
17    /// Gated pairs used for post-update Euclidean fitness.
18    pub evaluated_correspondences: u32,
19    /// Post-update Euclidean mean squared distance, f64::MAX when unmatched.
20    pub fitness: f64,
21    /// Norm of the world rotation/translation linearized solution, as in CPU GICP.
22    pub update_magnitude: f64,
23}
24
25/// Explicit GICP resident reduction result, trace and full transfer receipt.
26#[derive(Clone, Debug, PartialEq)]
27pub struct GpuGicpReductionDiagnostics {
28    /// Estimated pose, Euclidean fitness and stopping result.
29    pub result: RegistrationResult,
30    /// Completed update rows from this device precision path.
31    pub history: Vec<GpuGicpIteration>,
32    /// Position/tree/covariance upload, optional radius covariance transfers and equations readback.
33    pub transfers: TransferStats,
34    /// Completed resident evaluations (updates plus one).
35    pub evaluation_batches: usize,
36    /// Explicit GPU covariance estimates, zero or two (radius or resident k-NN).
37    pub covariance_batches: usize,
38}
39
40fn invalid() -> SpatialError {
41    SpatialError::InvalidArgument("GPU GICP reduction requires at least one iteration, minimum six correspondences, finite positive covariance epsilon/distance square, finite nonnegative stopping thresholds and a finite unit initial pose".to_owned())
42}
43fn valid_pose(pose: Isometry3<f32>) -> bool {
44    let q = pose.rotation();
45    let t = pose.translation();
46    let norm = f64::from(q.x).hypot(f64::from(q.y)).hypot(f64::from(q.z)).hypot(f64::from(q.w));
47    t.x.is_finite()
48        && t.y.is_finite()
49        && t.z.is_finite()
50        && norm.is_finite()
51        && (norm - 1.0).abs() <= 1e-4
52}
53fn record(total: &mut TransferStats, part: TransferStats) {
54    total.record(TransferDirection::HostToDevice, part.host_to_device_bytes());
55    total.record(TransferDirection::DeviceToHost, part.device_to_host_bytes());
56}
57fn fitness(equations: &GpuGicpEquations) -> f64 {
58    if equations.fitness_correspondences == 0 {
59        f64::MAX
60    } else {
61        equations.squared_error_sum / f64::from(equations.fitness_correspondences)
62    }
63}
64
65#[derive(Clone, Copy)]
66enum CovariancePreparation {
67    Legacy,
68    ResidentKnn,
69    ResidentRadius(u32),
70    ResidentGrid(u32),
71}
72
73type TreeBuilder = fn(&[f32], &[f32], &[f32]) -> KdTree;
74
75#[path = "gicp_covariance_preparation.rs"]
76mod covariance_preparation;
77
78impl GicpRegistration {
79    /// Retains source/target/covariances on GPU for transform, nearest matching,
80    /// rotated covariance inverse, Hessian/gradient and Euclidean fitness passes.
81    /// CPU solves normalized 6x6 equations and converts the linearized step to
82    /// the existing world update/stopping convention. Each evaluation explicitly
83    /// downloads 256 bytes; post-update equations feed the next solve/final fitness.
84    /// Covariance preparation remains CPU k-NN or explicit radius-GPU readback
85    /// and re-upload with full receipts. CPU k-NN rows are packed directly with
86    /// bounded scoped workers for large native inputs; small/WASM inputs stay
87    /// serial. Per-point neighbor and f64 accumulation order is preserved.
88    /// No automatic CPU fallback occurs.
89    /// Configuration/conditioning restrictions are stricter than legacy GICP;
90    /// f32 weights, transforms and sums can change pose/matching/stopping.
91    pub fn align_with_gpu_reduction(
92        &self,
93        runtime: &WgpuRuntime,
94        source: &PointCloud,
95        target: &PointCloud,
96    ) -> SpatialResult<GpuGicpReductionDiagnostics> {
97        self.align_reduction_impl(runtime, source, target, CovariancePreparation::Legacy)
98    }
99
100    /// Computes k-NN neighborhoods and plane covariances on GPU, retaining the
101    /// covariance buffers for the existing resident equation/update loop. No
102    /// neighborhood/covariance download or re-upload occurs. CPU builds both
103    /// trees, solves normalized 6x6 equations and composes/stops the pose.
104    /// Requires covariance_radius=None, k_neighbors=3..=64, epsilon (as f32)
105    /// in 1e-4..=1, and the same valid configuration as align_with_gpu_reduction.
106    /// Centered/scaled f32 covariance/eigenvectors may differ from CPU raw
107    /// products, particularly at repeated eigenvalues and search boundaries.
108    /// Invalid device preparation rejects equations without CPU fallback.
109    pub fn align_with_gpu_resident_covariances(
110        &self,
111        runtime: &WgpuRuntime,
112        source: &PointCloud,
113        target: &PointCloud,
114    ) -> SpatialResult<GpuGicpReductionDiagnostics> {
115        self.align_reduction_impl(runtime, source, target, CovariancePreparation::ResidentKnn)
116    }
117
118    /// Computes complete inclusive radius neighborhoods and covariances on GPU
119    /// and retains them for the resident equation/update loop. Requires configured
120    /// covariance_radius=Some(finite positive radius) with a normal f32 square,
121    /// epsilon (as f32) in 1e-4..=1 and positive `neighbor_capacity` per point.
122    /// Capacity includes self/duplicates; overflow fails without truncation or
123    /// fallback. Fewer than three neighbors use epsilon*I as in the grid route.
124    /// Centered/scaled compensated f32 covariance and max-pivot Jacobi can differ
125    /// from grid accumulation/sweeps, especially at repeated eigenvalues/gates.
126    /// CPU builds both trees, solves 6x6 equations and composes/stops the pose.
127    /// Upload is 32*(S+T+source_nodes+target_nodes) bytes; downloads are only
128    /// 256 bytes per evaluation. No neighborhood/covariance payload crosses CPU.
129    pub fn align_with_gpu_resident_radius_covariances(
130        &self,
131        runtime: &WgpuRuntime,
132        source: &PointCloud,
133        target: &PointCloud,
134        neighbor_capacity: u32,
135    ) -> SpatialResult<GpuGicpReductionDiagnostics> {
136        self.align_reduction_impl(
137            runtime,
138            source,
139            target,
140            CovariancePreparation::ResidentRadius(neighbor_capacity),
141        )
142    }
143
144    /// Prepares radius covariances with CPU counting-sort grids and GPU gather,
145    /// retaining the canonical grid's raw f32 means/central sums/Jacobi sweeps.
146    /// No source KD-tree, per-point neighborhood buffer or covariance readback
147    /// is needed. Positive capacity includes self/duplicates; overflow rejects
148    /// all equations. Configured radius/epsilon and normal-equation constraints
149    /// match the resident-radius API. Dense grids must fit existing cell/device
150    /// limits; no automatic sparse-path/CPU fallback occurs. CPU prepares grids
151    /// and the target tree and solves 6x6 equations. Downloads remain 256 bytes
152    /// per evaluation; all grid-index uploads are included in the receipt.
153    pub fn align_with_gpu_resident_grid_covariances(
154        &self,
155        runtime: &WgpuRuntime,
156        source: &PointCloud,
157        target: &PointCloud,
158        neighbor_capacity: u32,
159    ) -> SpatialResult<GpuGicpReductionDiagnostics> {
160        self.align_reduction_impl(
161            runtime,
162            source,
163            target,
164            CovariancePreparation::ResidentGrid(neighbor_capacity),
165        )
166    }
167
168    fn align_reduction_impl(
169        &self,
170        runtime: &WgpuRuntime,
171        source: &PointCloud,
172        target: &PointCloud,
173        preparation: CovariancePreparation,
174    ) -> SpatialResult<GpuGicpReductionDiagnostics> {
175        self.align_reduction_with_tree_builder(
176            runtime,
177            source,
178            target,
179            preparation,
180            KdTree::from_slices_parallel,
181        )
182    }
183
184    fn align_reduction_with_tree_builder(
185        &self,
186        runtime: &WgpuRuntime,
187        source: &PointCloud,
188        target: &PointCloud,
189        preparation: CovariancePreparation,
190        build_tree: TreeBuilder,
191    ) -> SpatialResult<GpuGicpReductionDiagnostics> {
192        self.validate_inputs(source, target)?;
193        let c = self.config;
194        let distance_squared = c.max_correspondence_distance * c.max_correspondence_distance;
195        let epsilon = c.epsilon as f32;
196        if matches!(preparation, CovariancePreparation::ResidentKnn)
197            && (c.covariance_radius.is_some()
198                || c.k_neighbors > spatialrust_gpu::MAX_GPU_K_NEIGHBORS as usize
199                || !(1e-4..=1.0).contains(&epsilon))
200        {
201            return Err(SpatialError::InvalidArgument("resident GPU GICP covariances require covariance_radius=None, k_neighbors=3..=64 and epsilon (as f32) in 1e-4..=1".to_owned()));
202        }
203        if let CovariancePreparation::ResidentRadius(capacity)
204        | CovariancePreparation::ResidentGrid(capacity) = preparation
205        {
206            if capacity == 0
207                || !(1e-4..=1.0).contains(&epsilon)
208                || c.covariance_radius
209                    .map_or(true, |r| !r.is_finite() || r <= 0.0 || !(r * r).is_normal())
210            {
211                return Err(SpatialError::InvalidArgument("resident GPU GICP radius covariances require a positive capacity including self, covariance_radius=Some(finite positive radius) with a normal f32 square and epsilon (as f32) in 1e-4..=1".to_owned()));
212            }
213        }
214        if c.max_iterations == 0
215            || c.min_correspondences < 6
216            || !c.epsilon.is_finite()
217            || epsilon <= 0.0
218            || !epsilon.is_finite()
219            || !c.max_correspondence_distance.is_finite()
220            || c.max_correspondence_distance <= 0.0
221            || !distance_squared.is_finite()
222            || distance_squared == 0.0
223            || !c.transformation_epsilon.is_finite()
224            || c.transformation_epsilon < 0.0
225            || !c.fitness_epsilon.is_finite()
226            || c.fitness_epsilon < 0.0
227            || !valid_pose(c.initial_guess)
228        {
229            return Err(invalid());
230        }
231        let (sx, sy, sz) = source.positions3()?;
232        let (tx, ty, tz) = target.positions3()?;
233        if sx.iter().chain(sy).chain(sz).chain(tx).chain(ty).chain(tz).any(|v| !v.is_finite()) {
234            return Err(SpatialError::InvalidArgument("GPU GICP requires finite XYZ".to_owned()));
235        }
236        let tree = build_tree(tx, ty, tz);
237        let mut transfers = TransferStats::default();
238        let mut covariance_batches = 0;
239        let input = if let CovariancePreparation::ResidentGrid(capacity) = preparation {
240            covariance_batches = 2;
241            GpuGicpInput::upload_with_gpu_grid_covariances(
242                runtime,
243                (sx, sy, sz),
244                &tree,
245                c.covariance_radius.unwrap(),
246                capacity,
247                epsilon,
248            )?
249        } else if !matches!(preparation, CovariancePreparation::Legacy) {
250            let source_tree = build_tree(sx, sy, sz);
251            covariance_batches = 2;
252            match preparation {
253                CovariancePreparation::ResidentKnn => {
254                    let k = u32::try_from(c.k_neighbors).map_err(|_| invalid())?;
255                    GpuGicpInput::upload_with_gpu_covariances(
256                        runtime,
257                        &source_tree,
258                        &tree,
259                        k,
260                        epsilon,
261                    )?
262                }
263                CovariancePreparation::ResidentRadius(capacity) => {
264                    GpuGicpInput::upload_with_gpu_radius_covariances(
265                        runtime,
266                        &source_tree,
267                        &tree,
268                        c.covariance_radius.unwrap(),
269                        capacity,
270                        epsilon,
271                    )?
272                }
273                CovariancePreparation::Legacy | CovariancePreparation::ResidentGrid(_) => {
274                    unreachable!()
275                }
276            }
277        } else {
278            let mut prepare = |x: &[f32],
279                               y: &[f32],
280                               z: &[f32],
281                               tree: Option<&KdTree>|
282             -> SpatialResult<Vec<[f32; 6]>> {
283                if let Some(radius) = c.covariance_radius {
284                    let (covariance, receipt) =
285                        spatialrust_gpu::estimate_plane_covariances_grid_gpu_with_receipt(
286                            runtime, x, y, z, radius, epsilon,
287                        )?;
288                    record(&mut transfers, receipt);
289                    covariance_batches += 1;
290                    Ok(covariance)
291                } else {
292                    let built;
293                    let tree = if let Some(tree) = tree {
294                        tree
295                    } else {
296                        built = build_tree(x, y, z);
297                        &built
298                    };
299                    Ok(covariance_preparation::prepare(x, y, z, tree, c.k_neighbors, c.epsilon))
300                }
301            };
302            let source_cov = prepare(sx, sy, sz, None)?;
303            let target_cov = prepare(tx, ty, tz, Some(&tree))?;
304            let positions: Vec<_> = (0..source.len()).map(|i| [sx[i], sy[i], sz[i]]).collect();
305            GpuGicpInput::upload(runtime, &positions, &source_cov, &tree, &target_cov)?
306        };
307        record(&mut transfers, input.upload_transfers());
308        let mut evaluation_batches = 0;
309        let mut evaluate = |pose| -> SpatialResult<GpuGicpEquations> {
310            let equations = input
311                .equations_resident(runtime, pose, c.max_correspondence_distance)?
312                .download(runtime)?;
313            record(&mut transfers, equations.transfers);
314            evaluation_batches += 1;
315            Ok(equations)
316        };
317        let mut transform = c.initial_guess;
318        let mut equations = evaluate(transform)?;
319        let mut converged = false;
320        let mut history = Vec::new();
321        for iteration in 1..=c.max_iterations {
322            let count = equations.correspondences;
323            if (count as usize) < c.min_correspondences {
324                return Err(SpatialError::InvalidArgument(format!(
325                    "GICP found only {count} correspondences, minimum is {}",
326                    c.min_correspondences
327                )));
328            }
329            let rows = equations.hessian.iter().map(|row| row.to_vec()).collect();
330            let gradient = equations.gradient.iter().map(|g| -g).collect();
331            let LeastSquaresResult::Solved(solution) = solve_linear_system(rows, gradient) else {
332                return Err(SpatialError::InvalidArgument(
333                    "GICP normal equations were singular".to_owned(),
334                ));
335            };
336            let solution: [f64; 6] = solution.try_into().map_err(|_| invalid())?;
337            let world = equations.world_step(solution)?;
338            transform = delta_from_solution(&world).compose(transform);
339            if !valid_pose(transform) {
340                return Err(invalid());
341            }
342            equations = evaluate(transform)?;
343            let fitness = fitness(&equations);
344            let update_magnitude = update_magnitude(&world);
345            history.push(GpuGicpIteration {
346                iteration,
347                correspondences: count,
348                evaluated_correspondences: equations.fitness_correspondences,
349                fitness,
350                update_magnitude,
351            });
352            if update_magnitude < c.transformation_epsilon || fitness < c.fitness_epsilon {
353                converged = true;
354                break;
355            }
356        }
357        Ok(GpuGicpReductionDiagnostics {
358            result: RegistrationResult {
359                transform,
360                fitness: fitness(&equations),
361                iterations: history.len(),
362                converged,
363            },
364            history,
365            transfers,
366            evaluation_batches,
367            covariance_batches,
368        })
369    }
370}
371
372#[cfg(all(test, not(target_arch = "wasm32")))]
373#[path = "gicp_tree_build_tests.rs"]
374mod tree_build_tests;
spatialrust_registration/gicp_reduction.rs

spatialrust_registration/
gicp_reduction.rs
