spatialrust_gpu/kernels/
ransac_plane.rs

1use bytemuck::{Pod, Zeroable};
2use spatialrust_core::{SpatialError, SpatialResult};
3use wgpu::util::DeviceExt;
4
5use crate::runtime::WgpuRuntime;
6
7const WORKGROUP_SIZE: u32 = 256;
8
9#[repr(C)]
10#[derive(Clone, Copy, Debug, Pod, Zeroable)]
11struct RansacPlaneUniform {
12    point_count: u32,
13    hypothesis_count: u32,
14    distance_threshold: f32,
15    _pad: u32,
16}
17
18#[repr(C)]
19#[derive(Clone, Copy, Debug, Pod, Zeroable)]
20struct RansacHypothesisPod {
21    i0: u32,
22    i1: u32,
23    i2: u32,
24    _pad: u32,
25}
26
27/// GPU score for one RANSAC plane hypothesis.
28#[repr(C)]
29#[derive(Clone, Copy, Debug, Pod, Zeroable)]
30pub struct GpuPlaneScore {
31    /// Number of inlier points for this hypothesis.
32    pub inlier_count: u32,
33    /// Unit plane normal `(nx, ny, nz)`.
34    pub normal: [f32; 3],
35    /// Plane offset term in Hessian form.
36    pub d: f32,
37}
38
39const RANSAC_PLANE_WGSL: &str = r#"
40struct Params {
41    point_count: u32,
42    hypothesis_count: u32,
43    distance_threshold: f32,
44    pad: f32,
45};
46
47struct Hypothesis {
48    i0: u32,
49    i1: u32,
50    i2: u32,
51    pad: u32,
52};
53
54struct Score {
55    inlier_count: u32,
56    nx: f32,
57    ny: f32,
58    nz: f32,
59    d: f32,
60};
61
62@group(0) @binding(0) var<uniform> params: Params;
63@group(0) @binding(1) var<storage, read> xs: array<f32>;
64@group(0) @binding(2) var<storage, read> ys: array<f32>;
65@group(0) @binding(3) var<storage, read> zs: array<f32>;
66@group(0) @binding(4) var<storage, read> hypotheses: array<Hypothesis>;
67@group(0) @binding(5) var<storage, read_write> scores: array<Score>;
68
69fn plane_from_indices(i0: u32, i1: u32, i2: u32) -> Score {
70    let p0 = vec3<f32>(xs[i0], ys[i0], zs[i0]);
71    let p1 = vec3<f32>(xs[i1], ys[i1], zs[i1]);
72    let p2 = vec3<f32>(xs[i2], ys[i2], zs[i2]);
73    let v1 = p1 - p0;
74    let v2 = p2 - p0;
75    var normal = cross(v1, v2);
76    if (dot(normal, normal) < 1e-12) {
77        return Score(0u, 0.0, 0.0, 0.0, 0.0);
78    }
79    normal = normalize(normal);
80    let d = -dot(normal, p0);
81    var count = 0u;
82    for (var i: u32 = 0u; i < params.point_count; i = i + 1u) {
83        let dist = abs(normal.x * xs[i] + normal.y * ys[i] + normal.z * zs[i] + d);
84        if (dist <= params.distance_threshold) {
85            count = count + 1u;
86        }
87    }
88    return Score(count, normal.x, normal.y, normal.z, d);
89}
90
91@compute @workgroup_size(256)
92fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
93    let h = gid.x;
94    if (h >= params.hypothesis_count) {
95        return;
96    }
97    let hyp = hypotheses[h];
98    scores[h] = plane_from_indices(hyp.i0, hyp.i1, hyp.i2);
99}
100"#;
101
102/// Scores RANSAC plane hypotheses in parallel on the GPU.
103pub fn score_ransac_plane_hypotheses_gpu(
104    runtime: &WgpuRuntime,
105    x: &[f32],
106    y: &[f32],
107    z: &[f32],
108    hypotheses: &[[u32; 3]],
109    distance_threshold: f32,
110) -> SpatialResult<Vec<GpuPlaneScore>> {
111    if x.len() != y.len() || x.len() != z.len() {
112        return Err(SpatialError::InvalidArgument("xyz arrays must have equal length".to_owned()));
113    }
114    if hypotheses.is_empty() {
115        return Ok(Vec::new());
116    }
117
118    let device = runtime.device();
119    let queue = runtime.queue();
120    let point_count = x.len();
121
122    let x_buffer = runtime.upload_f32_storage("ransac-plane-x", x)?;
123    let y_buffer = runtime.upload_f32_storage("ransac-plane-y", y)?;
124    let z_buffer = runtime.upload_f32_storage("ransac-plane-z", z)?;
125
126    let hypothesis_pods: Vec<RansacHypothesisPod> = hypotheses
127        .iter()
128        .map(|indices| RansacHypothesisPod {
129            i0: indices[0],
130            i1: indices[1],
131            i2: indices[2],
132            _pad: 0,
133        })
134        .collect();
135    let hypothesis_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
136        label: Some("ransac-plane-hypotheses"),
137        contents: bytemuck::cast_slice(&hypothesis_pods),
138        usage: wgpu::BufferUsages::STORAGE,
139    });
140
141    let hypothesis_count = hypotheses.len();
142    let output_len = (hypothesis_count * std::mem::size_of::<GpuPlaneScore>()) as u64;
143    let output_buffer = device.create_buffer(&wgpu::BufferDescriptor {
144        label: Some("ransac-plane-scores"),
145        size: output_len,
146        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
147        mapped_at_creation: false,
148    });
149
150    let uniform = RansacPlaneUniform {
151        point_count: point_count as u32,
152        hypothesis_count: hypothesis_count as u32,
153        distance_threshold,
154        _pad: 0,
155    };
156    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
157        label: Some("ransac-plane-uniform"),
158        contents: bytemuck::bytes_of(&uniform),
159        usage: wgpu::BufferUsages::UNIFORM,
160    });
161
162    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
163        label: Some("ransac-plane-shader"),
164        source: wgpu::ShaderSource::Wgsl(RANSAC_PLANE_WGSL.into()),
165    });
166    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
167        label: Some("ransac-plane-pipeline"),
168        layout: None,
169        module: &module,
170        entry_point: Some("main"),
171        compilation_options: wgpu::PipelineCompilationOptions::default(),
172        cache: None,
173    });
174
175    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
176        label: Some("ransac-plane-bind-group"),
177        layout: &pipeline.get_bind_group_layout(0),
178        entries: &[
179            wgpu::BindGroupEntry { binding: 0, resource: uniform_buffer.as_entire_binding() },
180            wgpu::BindGroupEntry { binding: 1, resource: x_buffer.as_entire_binding() },
181            wgpu::BindGroupEntry { binding: 2, resource: y_buffer.as_entire_binding() },
182            wgpu::BindGroupEntry { binding: 3, resource: z_buffer.as_entire_binding() },
183            wgpu::BindGroupEntry { binding: 4, resource: hypothesis_buffer.as_entire_binding() },
184            wgpu::BindGroupEntry { binding: 5, resource: output_buffer.as_entire_binding() },
185        ],
186    });
187
188    let mut encoder = device
189        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("ransac-plane") });
190    {
191        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
192            label: Some("ransac-plane-pass"),
193            timestamp_writes: None,
194        });
195        pass.set_pipeline(&pipeline);
196        pass.set_bind_group(0, &bind_group, &[]);
197        pass.dispatch_workgroups(hypothesis_count.div_ceil(WORKGROUP_SIZE as usize) as u32, 1, 1);
198    }
199    queue.submit(Some(encoder.finish()));
200
201    let staging = device.create_buffer(&wgpu::BufferDescriptor {
202        label: Some("ransac-plane-staging"),
203        size: output_len,
204        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
205        mapped_at_creation: false,
206    });
207    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
208        label: Some("ransac-plane-readback"),
209    });
210    encoder.copy_buffer_to_buffer(&output_buffer, 0, &staging, 0, output_len);
211    queue.submit(Some(encoder.finish()));
212
213    let slice = staging.slice(..);
214    let (sender, receiver) = std::sync::mpsc::channel();
215    slice.map_async(wgpu::MapMode::Read, move |result| {
216        let _ = sender.send(result);
217    });
218    device.poll(wgpu::Maintain::Wait);
219    receiver
220        .recv()
221        .map_err(|_| SpatialError::InvalidArgument("failed to receive wgpu map result".to_owned()))?
222        .map_err(|error| {
223            SpatialError::InvalidArgument(format!("failed to map wgpu buffer: {error}"))
224        })?;
225
226    let data = slice.get_mapped_range();
227    let scores: Vec<GpuPlaneScore> = bytemuck::cast_slice(&data).to_vec();
228    drop(data);
229    staging.unmap();
230
231    runtime.recycle_storage(std::mem::size_of_val(x) as u64, x_buffer);
232    runtime.recycle_storage(std::mem::size_of_val(y) as u64, y_buffer);
233    runtime.recycle_storage(std::mem::size_of_val(z) as u64, z_buffer);
234
235    Ok(scores)
236}
237
238#[cfg(test)]
239mod tests {
240    use super::{score_ransac_plane_hypotheses_gpu, GpuPlaneScore};
241    use crate::runtime::WgpuRuntime;
242
243    #[test]
244    fn scores_planar_patch_hypothesis() {
245        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
246        let mut x = Vec::new();
247        let mut y = Vec::new();
248        let mut z = Vec::new();
249        for i in 0..10 {
250            for j in 0..10 {
251                x.push(i as f32);
252                y.push(j as f32);
253                z.push(0.0);
254            }
255        }
256        let hypotheses = [[0u32, 1, 10], [5, 15, 50]];
257        let scores = score_ransac_plane_hypotheses_gpu(&runtime, &x, &y, &z, &hypotheses, 0.05)
258            .expect("gpu scores");
259        assert_eq!(scores.len(), 2);
260        assert!(scores.iter().all(|score: &GpuPlaneScore| score.inlier_count >= 90));
261        assert!(scores[0].normal[2].abs() > 0.9);
262    }
263}
spatialrust_gpu/kernels/ransac_plane.rs

spatialrust_gpu/kernels/
ransac_plane.rs
