spatialrust_gpu/kernels/
voxel_pipeline.rs

1use spatialrust_core::SpatialResult;
2
3use crate::kernels::gpu_segments::GpuVoxelSegments;
4use crate::kernels::voxel_gather::gather_voxel_first_xyz_gpu_buffers;
5use crate::kernels::voxel_keys::{compute_voxel_keys_gpu_buffers, GpuVoxelKeyBuffers};
6use crate::kernels::voxel_reduce::reduce_voxel_centroids_xyz_gpu_buffers;
7use crate::kernels::voxel_sort::build_voxel_segments_gpu_from_keys_buffer;
8use crate::runtime::WgpuRuntime;
9
10/// GPU-resident centroid downsample result.
11pub struct VoxelCentroidGpuResult {
12    /// Averaged x coordinates per voxel cell.
13    pub out_x: Vec<f32>,
14    /// Averaged y coordinates per voxel cell.
15    pub out_y: Vec<f32>,
16    /// Averaged z coordinates per voxel cell.
17    pub out_z: Vec<f32>,
18    /// Segment metadata kept on the GPU for follow-up reductions.
19    pub segments: GpuVoxelSegments,
20    /// Source position buffers kept on the GPU.
21    pub positions: GpuVoxelKeyBuffers,
22}
23
24/// Runs the chained GPU voxel centroid pipeline without intermediate key/segment readbacks.
25pub fn downsample_voxel_centroid_gpu(
26    runtime: &WgpuRuntime,
27    x: &[f32],
28    y: &[f32],
29    z: &[f32],
30    origin: [f32; 3],
31    inv_leaf: f32,
32) -> SpatialResult<VoxelCentroidGpuResult> {
33    let positions = compute_voxel_keys_gpu_buffers(runtime, x, y, z, origin, inv_leaf)?;
34    let point_count = positions.point_count();
35    let padded_count = point_count.next_power_of_two();
36    let segments = build_voxel_segments_gpu_from_keys_buffer(
37        runtime,
38        positions.keys_buffer(),
39        point_count,
40        padded_count,
41    )?;
42    let (out_x, out_y, out_z) = reduce_voxel_centroids_xyz_gpu_buffers(
43        runtime,
44        positions.x_buffer(),
45        positions.y_buffer(),
46        positions.z_buffer(),
47        &segments,
48    )?;
49
50    Ok(VoxelCentroidGpuResult { out_x, out_y, out_z, segments, positions })
51}
52
53/// GPU-resident approximate-first downsample result.
54pub struct VoxelApproximateFirstGpuResult {
55    /// First-point x coordinates per voxel cell.
56    pub out_x: Vec<f32>,
57    /// First-point y coordinates per voxel cell.
58    pub out_y: Vec<f32>,
59    /// First-point z coordinates per voxel cell.
60    pub out_z: Vec<f32>,
61    /// Segment metadata kept on the GPU for follow-up gathers/reductions.
62    pub segments: GpuVoxelSegments,
63    /// Source position buffers kept on the GPU.
64    pub positions: GpuVoxelKeyBuffers,
65}
66
67/// Runs the chained GPU approximate-first pipeline without intermediate readbacks.
68pub fn downsample_voxel_approximate_first_gpu(
69    runtime: &WgpuRuntime,
70    x: &[f32],
71    y: &[f32],
72    z: &[f32],
73    origin: [f32; 3],
74    inv_leaf: f32,
75) -> SpatialResult<VoxelApproximateFirstGpuResult> {
76    let positions = compute_voxel_keys_gpu_buffers(runtime, x, y, z, origin, inv_leaf)?;
77    let point_count = positions.point_count();
78    let padded_count = point_count.next_power_of_two();
79    let segments = build_voxel_segments_gpu_from_keys_buffer(
80        runtime,
81        positions.keys_buffer(),
82        point_count,
83        padded_count,
84    )?;
85    let (out_x, out_y, out_z) = gather_voxel_first_xyz_gpu_buffers(
86        runtime,
87        positions.x_buffer(),
88        positions.y_buffer(),
89        positions.z_buffer(),
90        &segments,
91    )?;
92
93    Ok(VoxelApproximateFirstGpuResult { out_x, out_y, out_z, segments, positions })
94}
95
96#[cfg(test)]
97mod tests {
98    use super::{downsample_voxel_approximate_first_gpu, downsample_voxel_centroid_gpu};
99    use crate::kernels::voxel_reduce::reduce_voxel_centroids_xyz;
100    use crate::kernels::voxel_segments::build_voxel_segments;
101    use crate::runtime::WgpuRuntime;
102
103    #[test]
104    fn chained_pipeline_matches_staged_gpu_reference() {
105        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
106        let x = [0.0_f32, 0.1, 1.0, 1.1];
107        let y = [0.0_f32, 0.0, 0.0, 0.0];
108        let z = [0.0_f32, 0.0, 0.0, 0.0];
109        let origin = [0.0_f32, 0.0, 0.0];
110        let inv_leaf = 2.0_f32;
111
112        let chained = downsample_voxel_centroid_gpu(&runtime, &x, &y, &z, origin, inv_leaf)
113            .expect("chained pipeline");
114
115        let keys: Vec<(i64, i64, i64)> = x
116            .iter()
117            .zip(y.iter())
118            .zip(z.iter())
119            .map(|((x, y), z)| {
120                (
121                    ((x - origin[0]) * inv_leaf).floor() as i64,
122                    ((y - origin[1]) * inv_leaf).floor() as i64,
123                    ((z - origin[2]) * inv_leaf).floor() as i64,
124                )
125            })
126            .collect();
127        let segments = build_voxel_segments(&keys);
128        let (ref_x, ref_y, ref_z) =
129            reduce_voxel_centroids_xyz(&runtime, &x, &y, &z, &segments).expect("staged reduce");
130
131        assert!((chained.out_x[0] - ref_x[0]).abs() < 1e-5);
132        assert!((chained.out_x[1] - ref_x[1]).abs() < 1e-5);
133        assert_eq!(chained.out_y, ref_y);
134        assert_eq!(chained.out_z, ref_z);
135    }
136
137    #[test]
138    fn approximate_first_pipeline_keeps_first_point() {
139        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
140        let x = [0.0_f32, 0.1];
141        let y = [0.0_f32, 0.0];
142        let z = [0.0_f32, 0.0];
143        let origin = [0.0_f32, 0.0, 0.0];
144        let inv_leaf = 1.0_f32;
145
146        let result = downsample_voxel_approximate_first_gpu(&runtime, &x, &y, &z, origin, inv_leaf)
147            .expect("approximate-first pipeline");
148
149        assert_eq!(result.out_x.len(), 1);
150        assert!((result.out_x[0] - 0.0).abs() < 1e-5);
151    }
152
153    #[test]
154    fn centroid_pipeline_handles_non_pot_point_count() {
155        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
156        let point_count = 100_000_usize;
157        let mut x = Vec::with_capacity(point_count);
158        let mut y = Vec::with_capacity(point_count);
159        let mut z = Vec::with_capacity(point_count);
160        for index in 0..point_count {
161            let t = index as f32;
162            x.push((t * 0.013).fract() * 100.0);
163            y.push(((index % 97) as f32) * 0.017);
164            z.push(((index % 53) as f32) * 0.019);
165        }
166
167        downsample_voxel_centroid_gpu(&runtime, &x, &y, &z, [0.0; 3], 0.25)
168            .expect("centroid pipeline for 100k points");
169    }
170}
spatialrust_gpu/kernels/voxel_pipeline.rs

spatialrust_gpu/kernels/
voxel_pipeline.rs
