spatialrust_gpu/kernels/
normals_grid.rs

1use bytemuck::{Pod, Zeroable};
2use spatialrust_core::{SpatialError, SpatialResult};
3use wgpu::util::DeviceExt;
4
5use crate::kernels::normals::GpuNormal;
6use crate::runtime::WgpuRuntime;
7
8use spatialrust_search::{build_grid, grid_bounds};
9
10pub use spatialrust_search::uniform_grid_fits;
11
12const WORKGROUP_SIZE: u32 = 256;
13
14#[repr(C)]
15#[derive(Clone, Copy, Debug, Pod, Zeroable)]
16struct GridUniform {
17    origin: [f32; 4],
18    dims: [u32; 4], // dimx, dimy, dimz, point_count
19    inv_cell: f32,
20    radius_sq: f32,
21    _pad0: f32,
22    _pad1: f32,
23}
24
25pub(crate) const NORMALS_GRID_WGSL: &str = r#"
26struct Params {
27    origin: vec4<f32>,
28    dims: vec4<u32>,
29    inv_cell: f32,
30    radius_sq: f32,
31    pad0: f32,
32    pad1: f32,
33};
34
35@group(0) @binding(0) var<uniform> params: Params;
36@group(0) @binding(1) var<storage, read> xs: array<f32>;
37@group(0) @binding(2) var<storage, read> ys: array<f32>;
38@group(0) @binding(3) var<storage, read> zs: array<f32>;
39@group(0) @binding(4) var<storage, read> sorted: array<u32>;
40@group(0) @binding(5) var<storage, read> cell_start: array<u32>;
41@group(0) @binding(6) var<storage, read_write> out_normals: array<vec4<f32>>;
42
43fn rotate(a: ptr<function, array<vec3<f32>, 3>>,
44          v: ptr<function, array<vec3<f32>, 3>>,
45          p: u32, q: u32) {
46    let apq = (*a)[p][q];
47    if (abs(apq) < 1e-20) {
48        return;
49    }
50    let app = (*a)[p][p];
51    let aqq = (*a)[q][q];
52    let phi = 0.5 * (aqq - app) / apq;
53    var t: f32;
54    if (phi >= 0.0) {
55        t = 1.0 / (phi + sqrt(1.0 + phi * phi));
56    } else {
57        t = -1.0 / (-phi + sqrt(1.0 + phi * phi));
58    }
59    let c = 1.0 / sqrt(1.0 + t * t);
60    let s = t * c;
61    for (var r: u32 = 0u; r < 3u; r = r + 1u) {
62        let arp = (*a)[r][p];
63        let arq = (*a)[r][q];
64        (*a)[r][p] = c * arp - s * arq;
65        (*a)[r][q] = s * arp + c * arq;
66    }
67    for (var r: u32 = 0u; r < 3u; r = r + 1u) {
68        let apr = (*a)[p][r];
69        let aqr = (*a)[q][r];
70        (*a)[p][r] = c * apr - s * aqr;
71        (*a)[q][r] = s * apr + c * aqr;
72    }
73    for (var r: u32 = 0u; r < 3u; r = r + 1u) {
74        let vrp = (*v)[r][p];
75        let vrq = (*v)[r][q];
76        (*v)[r][p] = c * vrp - s * vrq;
77        (*v)[r][q] = s * vrp + c * vrq;
78    }
79}
80
81fn cell_coord(value: f32, origin: f32, inv_cell: f32, dim: u32) -> i32 {
82    let c = i32(floor((value - origin) * inv_cell));
83    return clamp(c, 0, i32(dim) - 1);
84}
85
86@compute @workgroup_size(256)
87fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
88    let i = gid.x;
89    if (i >= params.dims.w) {
90        return;
91    }
92    let px = xs[i];
93    let py = ys[i];
94    let pz = zs[i];
95    let dimx = params.dims.x;
96    let dimy = params.dims.y;
97    let dimz = params.dims.z;
98
99    let cx = cell_coord(px, params.origin.x, params.inv_cell, dimx);
100    let cy = cell_coord(py, params.origin.y, params.inv_cell, dimy);
101    let cz = cell_coord(pz, params.origin.z, params.inv_cell, dimz);
102
103    // First pass: mean over radius neighbors across the 27 adjacent cells.
104    var mean = vec3<f32>(0.0, 0.0, 0.0);
105    var count = 0.0;
106    for (var dz = -1; dz <= 1; dz = dz + 1) {
107        let nz = cz + dz;
108        if (nz < 0 || nz >= i32(dimz)) { continue; }
109        for (var dy = -1; dy <= 1; dy = dy + 1) {
110            let ny = cy + dy;
111            if (ny < 0 || ny >= i32(dimy)) { continue; }
112            for (var dx = -1; dx <= 1; dx = dx + 1) {
113                let nx = cx + dx;
114                if (nx < 0 || nx >= i32(dimx)) { continue; }
115                let cid = (u32(nz) * dimy + u32(ny)) * dimx + u32(nx);
116                let begin = cell_start[cid];
117                let end = cell_start[cid + 1u];
118                for (var s = begin; s < end; s = s + 1u) {
119                    let j = sorted[s];
120                    let d = vec3<f32>(xs[j] - px, ys[j] - py, zs[j] - pz);
121                    if (dot(d, d) <= params.radius_sq) {
122                        mean = mean + vec3<f32>(xs[j], ys[j], zs[j]);
123                        count = count + 1.0;
124                    }
125                }
126            }
127        }
128    }
129
130    if (count < 3.0) {
131        out_normals[i] = vec4<f32>(0.0, 0.0, 1.0, 0.0);
132        return;
133    }
134    mean = mean / count;
135
136    var c00 = 0.0; var c11 = 0.0; var c22 = 0.0;
137    var c01 = 0.0; var c02 = 0.0; var c12 = 0.0;
138    for (var dz = -1; dz <= 1; dz = dz + 1) {
139        let nz = cz + dz;
140        if (nz < 0 || nz >= i32(dimz)) { continue; }
141        for (var dy = -1; dy <= 1; dy = dy + 1) {
142            let ny = cy + dy;
143            if (ny < 0 || ny >= i32(dimy)) { continue; }
144            for (var dx = -1; dx <= 1; dx = dx + 1) {
145                let nx = cx + dx;
146                if (nx < 0 || nx >= i32(dimx)) { continue; }
147                let cid = (u32(nz) * dimy + u32(ny)) * dimx + u32(nx);
148                let begin = cell_start[cid];
149                let end = cell_start[cid + 1u];
150                for (var s = begin; s < end; s = s + 1u) {
151                    let j = sorted[s];
152                    let p = vec3<f32>(xs[j], ys[j], zs[j]);
153                    let rel = p - vec3<f32>(px, py, pz);
154                    if (dot(rel, rel) <= params.radius_sq) {
155                        let dd = p - mean;
156                        c00 = c00 + dd.x * dd.x;
157                        c11 = c11 + dd.y * dd.y;
158                        c22 = c22 + dd.z * dd.z;
159                        c01 = c01 + dd.x * dd.y;
160                        c02 = c02 + dd.x * dd.z;
161                        c12 = c12 + dd.y * dd.z;
162                    }
163                }
164            }
165        }
166    }
167
168    var a = array<vec3<f32>, 3>(
169        vec3<f32>(c00, c01, c02),
170        vec3<f32>(c01, c11, c12),
171        vec3<f32>(c02, c12, c22),
172    );
173    var v = array<vec3<f32>, 3>(
174        vec3<f32>(1.0, 0.0, 0.0),
175        vec3<f32>(0.0, 1.0, 0.0),
176        vec3<f32>(0.0, 0.0, 1.0),
177    );
178    for (var sweep: u32 = 0u; sweep < 16u; sweep = sweep + 1u) {
179        rotate(&a, &v, 0u, 1u);
180        rotate(&a, &v, 0u, 2u);
181        rotate(&a, &v, 1u, 2u);
182    }
183
184    let eig = vec3<f32>(a[0][0], a[1][1], a[2][2]);
185    var min_idx = 0u;
186    if (eig[1] < eig[min_idx]) { min_idx = 1u; }
187    if (eig[2] < eig[min_idx]) { min_idx = 2u; }
188    let normal = vec3<f32>(v[0][min_idx], v[1][min_idx], v[2][min_idx]);
189    let len = max(sqrt(dot(normal, normal)), 1e-20);
190    let unit = normal / len;
191    let trace = eig[0] + eig[1] + eig[2];
192    var curvature = 0.0;
193    if (trace > 0.0) {
194        curvature = eig[min_idx] / trace;
195    }
196    out_normals[i] = vec4<f32>(unit.x, unit.y, unit.z, curvature);
197}
198"#;
199
200/// Estimates per-point normals and curvature with a fully GPU radius neighbor
201/// search over a uniform grid.
202///
203/// The grid (cell size = `radius`) is built on the CPU with a counting sort
204/// (O(n)); the per-point neighbor gather, covariance, and eigen-decomposition
205/// all run on the GPU. Returns `SpatialError::InvalidArgument` when the bounding
206/// grid would exceed an internal cell cap (caller should fall back to the CPU
207/// KD-tree path).
208pub fn estimate_normals_grid_gpu(
209    runtime: &WgpuRuntime,
210    x: &[f32],
211    y: &[f32],
212    z: &[f32],
213    radius: f32,
214) -> SpatialResult<Vec<GpuNormal>> {
215    let point_count = x.len();
216    if y.len() != point_count || z.len() != point_count {
217        return Err(SpatialError::BufferLengthMismatch { expected: point_count, found: y.len() });
218    }
219    if point_count == 0 {
220        return Ok(Vec::new());
221    }
222    if radius <= 0.0 || radius.is_nan() {
223        return Err(SpatialError::InvalidArgument("grid radius must be positive".to_owned()));
224    }
225
226    let (origin, dims) = grid_bounds(x, y, z, radius)?;
227    let (sorted, cell_start) = build_grid(x, y, z, origin, dims, radius);
228
229    let device = runtime.device();
230    let queue = runtime.queue();
231    let inv_cell = 1.0 / radius;
232
233    let storage = wgpu::BufferUsages::STORAGE;
234    let x_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
235        label: Some("ng-x"),
236        contents: bytemuck::cast_slice(x),
237        usage: storage,
238    });
239    let y_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
240        label: Some("ng-y"),
241        contents: bytemuck::cast_slice(y),
242        usage: storage,
243    });
244    let z_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
245        label: Some("ng-z"),
246        contents: bytemuck::cast_slice(z),
247        usage: storage,
248    });
249    let sorted_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
250        label: Some("ng-sorted"),
251        contents: bytemuck::cast_slice(&sorted),
252        usage: storage,
253    });
254    let cell_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
255        label: Some("ng-cell-start"),
256        contents: bytemuck::cast_slice(&cell_start),
257        usage: storage,
258    });
259    let uniform = GridUniform {
260        origin: [origin[0], origin[1], origin[2], 0.0],
261        dims: [dims[0], dims[1], dims[2], point_count as u32],
262        inv_cell,
263        radius_sq: radius * radius,
264        _pad0: 0.0,
265        _pad1: 0.0,
266    };
267    let uniform_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
268        label: Some("ng-uniform"),
269        contents: bytemuck::bytes_of(&uniform),
270        usage: wgpu::BufferUsages::UNIFORM,
271    });
272
273    let output_len = (point_count * std::mem::size_of::<[f32; 4]>()) as u64;
274    let output_buf = device.create_buffer(&wgpu::BufferDescriptor {
275        label: Some("ng-output"),
276        size: output_len,
277        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
278        mapped_at_creation: false,
279    });
280
281    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
282        label: Some("ng-shader"),
283        source: wgpu::ShaderSource::Wgsl(NORMALS_GRID_WGSL.into()),
284    });
285    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
286        label: Some("ng-pipeline"),
287        layout: None,
288        module: &module,
289        entry_point: Some("main"),
290        compilation_options: wgpu::PipelineCompilationOptions::default(),
291        cache: None,
292    });
293    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
294        label: Some("ng-bind-group"),
295        layout: &pipeline.get_bind_group_layout(0),
296        entries: &[
297            wgpu::BindGroupEntry { binding: 0, resource: uniform_buf.as_entire_binding() },
298            wgpu::BindGroupEntry { binding: 1, resource: x_buf.as_entire_binding() },
299            wgpu::BindGroupEntry { binding: 2, resource: y_buf.as_entire_binding() },
300            wgpu::BindGroupEntry { binding: 3, resource: z_buf.as_entire_binding() },
301            wgpu::BindGroupEntry { binding: 4, resource: sorted_buf.as_entire_binding() },
302            wgpu::BindGroupEntry { binding: 5, resource: cell_buf.as_entire_binding() },
303            wgpu::BindGroupEntry { binding: 6, resource: output_buf.as_entire_binding() },
304        ],
305    });
306
307    let mut encoder =
308        device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("ng") });
309    {
310        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
311            label: Some("ng-pass"),
312            timestamp_writes: None,
313        });
314        pass.set_pipeline(&pipeline);
315        pass.set_bind_group(0, &bind_group, &[]);
316        pass.dispatch_workgroups((point_count as u32).div_ceil(WORKGROUP_SIZE), 1, 1);
317    }
318    queue.submit(Some(encoder.finish()));
319
320    let staging = device.create_buffer(&wgpu::BufferDescriptor {
321        label: Some("ng-staging"),
322        size: output_len,
323        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
324        mapped_at_creation: false,
325    });
326    let mut encoder =
327        device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("ng-rb") });
328    encoder.copy_buffer_to_buffer(&output_buf, 0, &staging, 0, output_len);
329    queue.submit(Some(encoder.finish()));
330
331    let slice = staging.slice(..);
332    let (sender, receiver) = std::sync::mpsc::channel();
333    slice.map_async(wgpu::MapMode::Read, move |result| {
334        let _ = sender.send(result);
335    });
336    device.poll(wgpu::Maintain::Wait);
337    receiver
338        .recv()
339        .map_err(|_| SpatialError::InvalidArgument("failed to receive wgpu map result".to_owned()))?
340        .map_err(|error| {
341            SpatialError::InvalidArgument(format!("failed to map wgpu buffer: {error}"))
342        })?;
343    let data = slice.get_mapped_range();
344    let raw: &[[f32; 4]] = bytemuck::cast_slice(&data);
345    let normals =
346        raw.iter().map(|v| GpuNormal { normal: [v[0], v[1], v[2]], curvature: v[3] }).collect();
347    drop(data);
348    staging.unmap();
349
350    Ok(normals)
351}
352
353#[cfg(test)]
354mod tests {
355    use super::estimate_normals_grid_gpu;
356    use crate::runtime::WgpuRuntime;
357
358    #[test]
359    fn planar_patch_has_vertical_normal() {
360        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
361        let mut x: Vec<f32> = Vec::new();
362        let mut y: Vec<f32> = Vec::new();
363        let mut z: Vec<f32> = Vec::new();
364        for i in 0..12 {
365            for j in 0..12 {
366                x.push(i as f32 * 0.1);
367                y.push(j as f32 * 0.1);
368                z.push(0.0);
369            }
370        }
371        let normals = estimate_normals_grid_gpu(&runtime, &x, &y, &z, 0.25).expect("grid normals");
372        assert_eq!(normals.len(), x.len());
373        for normal in &normals {
374            assert!(normal.normal[2].abs() > 0.99, "normal not vertical: {:?}", normal.normal);
375            assert!(normal.curvature < 1e-3, "curvature too high: {}", normal.curvature);
376        }
377    }
378}
spatialrust_gpu/kernels/normals_grid.rs

spatialrust_gpu/kernels/
normals_grid.rs
