From eeb12a54d374f71e40c0508add30da5dd52030d1 Mon Sep 17 00:00:00 2001 From: rsasaki0109 Date: Mon, 12 Oct 2026 22:41:46 +0900 Subject: [PATCH] feat(gpu): add reusable query batches and resident bounded neighborhoods --- CHANGELOG.md | 7 + README.md | 7 + crates/spatialrust-gpu/src/kernels/mod.rs | 5 +- crates/spatialrust-gpu/src/kernels/nearest.rs | 8 + .../src/kernels/neighborhood_tests.rs | 394 +++++++++++++++ .../src/kernels/neighborhoods.rs | 470 ++++++++++++++++++ crates/spatialrust-gpu/src/lib.rs | 6 + crates/spatialrust-gpu/src/runtime.rs | 2 + .../src/shaders/neighborhoods.wgsl | 118 +++++ .../tests/neighborhood_benchmark.rs | 100 ++++ .../tests/gpu_neighborhood_smoke.rs | 33 ++ docs/ARCHITECTURE.md | 9 + docs/FEATURE_MATRIX.md | 15 + docs/ROADMAP.md | 2 +- notes/2026-10-12_gpu_bounded_neighborhoods.md | 187 +++++++ 15 files changed, 1361 insertions(+), 2 deletions(-) create mode 100644 crates/spatialrust-gpu/src/kernels/neighborhood_tests.rs create mode 100644 crates/spatialrust-gpu/src/kernels/neighborhoods.rs create mode 100644 crates/spatialrust-gpu/src/shaders/neighborhoods.wgsl create mode 100644 crates/spatialrust-gpu/tests/neighborhood_benchmark.rs create mode 100644 crates/spatialrust/tests/gpu_neighborhood_smoke.rs create mode 100644 notes/2026-10-12_gpu_bounded_neighborhoods.md diff --git a/CHANGELOG.md b/CHANGELOG.md index fe81321..05a2b93 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,13 @@ removed no sooner than the next major (see `docs/API_STABILITY.md`). ### Changed +- Add explicitly uploaded reusable GPU query batches and resident KD-tree + nearest-k/radius neighborhoods under `gpu-wgpu`. k=0..=64 retains CPU + selection and stable distance order; radius rows retain traversal order and + count all matches even when capacity is exceeded. Complete readback access + rejects overflow; named truncated access and count-only search are explicit. + Owned result buffers can feed subsequent GPU kernels without readback. + - Extend explicit GPU correspondence reuse to GICP under `register-gicp-gpu`, retaining unbounded matching followed by the original host distance gate, ordered Hessian/gradient reduction and CPU pose solving. Optional radius diff --git a/README.md b/README.md index c7b4649..e4314ca 100644 --- a/README.md +++ b/README.md @@ -649,6 +649,13 @@ GICP exposes the same explicit correspondence entry point under and share the transfer receipt. Its Hessian/gradient accumulation and pose solver stay on CPU. See [GICP measurements](notes/2026-10-12_gpu_gicp_correspondences.md). +GPU k-NN/radius search can retain its results for later device passes with +`GpuQueryBatch::upload` and +`GpuKdTree::{nearest_k_resident,radius_neighbors_resident}` (`gpu-wgpu`). Radius +capacity overflow is counted explicitly and complete readback rejects it. +This provides the neighborhood foundation for feature compute; FPFH/ISS/MLS +remain CPU algorithms. See [contracts and measurements](notes/2026-10-12_gpu_bounded_neighborhoods.md). + ### Library Load or save by file extension: diff --git a/crates/spatialrust-gpu/src/kernels/mod.rs b/crates/spatialrust-gpu/src/kernels/mod.rs index 1092bd4..5d35415 100644 --- a/crates/spatialrust-gpu/src/kernels/mod.rs +++ b/crates/spatialrust-gpu/src/kernels/mod.rs @@ -51,7 +51,10 @@ mod voxel_gather; pub(crate) mod cluster_union; pub(crate) mod nearest; -pub use nearest::GpuKdTree; +pub use nearest::{ + GpuKdTree, GpuNeighborhoodKind, GpuNeighborhoodReadback, GpuNeighborhoods, GpuQueryBatch, + MAX_GPU_K_NEIGHBORS, +}; #[cfg(feature = "gpu-wgpu")] mod euclidean_cluster; diff --git a/crates/spatialrust-gpu/src/kernels/nearest.rs b/crates/spatialrust-gpu/src/kernels/nearest.rs index d860ee1..393bfcb 100644 --- a/crates/spatialrust-gpu/src/kernels/nearest.rs +++ b/crates/spatialrust-gpu/src/kernels/nearest.rs @@ -7,6 +7,14 @@ use wgpu::util::DeviceExt; use crate::{scratch::ScratchBuffers, WgpuRuntime}; +#[path = "neighborhoods.rs"] +mod neighborhoods; +pub(crate) use neighborhoods::NeighborhoodPipelines; +pub use neighborhoods::{ + GpuNeighborhoodKind, GpuNeighborhoodReadback, GpuNeighborhoods, GpuQueryBatch, + MAX_GPU_K_NEIGHBORS, +}; + #[repr(C)] #[derive(Clone, Copy, Pod, Zeroable)] struct Node { diff --git a/crates/spatialrust-gpu/src/kernels/neighborhood_tests.rs b/crates/spatialrust-gpu/src/kernels/neighborhood_tests.rs new file mode 100644 index 0000000..b3c2780 --- /dev/null +++ b/crates/spatialrust-gpu/src/kernels/neighborhood_tests.rs @@ -0,0 +1,394 @@ +use super::*; +use spatialrust_search::{KdTree, NearestNeighborIndex, RadiusSearchIndex}; + +fn fixture(n: usize, kind: usize) -> (KdTree, Vec<[f32; 3]>) { + let points: Vec<_> = (0..n) + .map(|i| { + let j = n - i - 1; + match kind { + 0 => [(j % 19) as f32 * 0.5, (j / 19 % 19) as f32 * 0.5, (j / 361) as f32 * 0.5], + 1 => [j as f32 * 0.5, 0.0, 0.0], + 2 => [0.0; 3], + _ => [((j % 5) as f32 - 2.0) * 0.5, ((j / 5 % 5) as f32 - 2.0) * 0.5, 0.0], + } + }) + .collect(); + let xyz: [Vec<_>; 3] = std::array::from_fn(|a| points.iter().map(|p| p[a]).collect()); + (KdTree::from_slices(&xyz[0], &xyz[1], &xyz[2]), points) +} + +#[test] +fn neighborhood_plan_bounds_capacity_buffers_host_and_dispatch() { + let limits = wgpu::Limits::default(); + assert_eq!(plan(0, 64, GpuNeighborhoodKind::NearestK, &limits).unwrap(), 0); + assert_eq!(plan(1, 64, GpuNeighborhoodKind::NearestK, &limits).unwrap(), 520); + assert!(plan(1, 65, GpuNeighborhoodKind::NearestK, &limits).is_err()); + assert_eq!(plan(1, 65, GpuNeighborhoodKind::Radius, &limits).unwrap(), 528); + assert!(plan(0, u64::from(u32::MAX), GpuNeighborhoodKind::Radius, &limits).is_err()); + assert!(plan(u64::MAX, 1, GpuNeighborhoodKind::Radius, &limits).is_err()); + let max_queries = u64::from(limits.max_compute_workgroups_per_dimension) * 64; + assert!(plan(max_queries, 0, GpuNeighborhoodKind::Radius, &limits).is_ok()); + assert!(plan(max_queries + 1, 0, GpuNeighborhoodKind::Radius, &limits).is_err()); + let rows = u64::from(limits.max_storage_buffer_binding_size) / 520; + assert!(plan(rows, 64, GpuNeighborhoodKind::NearestK, &limits).is_ok()); + assert!(plan(rows + 1, 64, GpuNeighborhoodKind::NearestK, &limits).is_err()); + let mut restricted = limits.clone(); + restricted.max_compute_workgroup_size_x = 63; + assert!(plan(1, 1, GpuNeighborhoodKind::NearestK, &restricted).is_err()); + restricted = limits.clone(); + restricted.max_storage_buffers_per_shader_stage = 3; + assert!(plan(1, 1, GpuNeighborhoodKind::NearestK, &restricted).is_err()); +} + +#[test] +fn gpu_knn_retains_cpu_membership_stable_ties_and_independent_distances() { + let runtime = WgpuRuntime::new_headless().unwrap(); + let mut verified = 0; + for n in [1, 16, 17, 63, 64, 65, 255, 257, 513] { + for kind in 0..4 { + let (tree, points) = fixture(n, kind); + let target = GpuKdTree::upload(&runtime, &tree).unwrap(); + let queries: Vec<_> = points + .iter() + .step_by((n / 33).max(1)) + .flat_map(|&p| [p, [p[0] + 0.125, p[1] - 0.0625, p[2]], [p[0] + 0.25, p[1], p[2]]]) + .collect(); + let batch = GpuQueryBatch::upload(&runtime, &queries).unwrap(); + assert_eq!(batch.upload_transfers().host_to_device_bytes(), queries.len() as u64 * 16); + assert_eq!(batch.upload_transfers().device_to_host_bytes(), 0); + for k in [0, 1, 2, 8, 16, 32, 64] { + let expected: Vec<_> = + queries.iter().map(|q| tree.nearest_k(q[0], q[1], q[2], k)).collect(); + let output = target.nearest_k_resident(&runtime, &batch, k as u32).unwrap(); + assert_eq!(output.execution_transfers(), TransferStats::default()); + assert_eq!(output.kind(), GpuNeighborhoodKind::NearestK); + let downloaded = output.download(&runtime).unwrap(); + assert_eq!( + downloaded.complete_rows().unwrap(), + expected, + "n={n}, kind={kind}, k={k}" + ); + assert_eq!( + downloaded.transfers().device_to_host_bytes(), + queries.len() as u64 * (8 + 8 * k as u64) + ); + for (q, row) in queries.iter().zip(downloaded.complete_rows().unwrap()) { + let mut distances: Vec<_> = points + .iter() + .map(|p| { + let dx = p[0] - q[0]; + let dy = p[1] - q[1]; + let dz = p[2] - q[2]; + dx * dx + dy * dy + dz * dz + }) + .collect(); + distances.sort_by(f32::total_cmp); + distances.truncate(k); + assert_eq!( + row.iter().map(|r| r.distance_squared).collect::>(), + distances + ); + verified += row.len(); + } + } + } + } + eprintln!("GPU_NEIGHBORHOOD verified {verified} selected neighbors against CPU ordering and brute distances"); +} + +#[test] +fn gpu_radius_retains_cpu_order_exact_counts_and_explicit_overflow() { + let runtime = WgpuRuntime::new_headless().unwrap(); + let mut verified = 0; + for n in [1, 17, 65, 257, 513] { + for kind in 0..4 { + let (tree, points) = fixture(n, kind); + let target = GpuKdTree::upload(&runtime, &tree).unwrap(); + let queries: Vec<_> = + points.iter().step_by((n / 17).max(1)).copied().chain([[500.0; 3]]).collect(); + let batch = GpuQueryBatch::upload(&runtime, &queries).unwrap(); + for radius in [ + -1.0, + f32::NAN, + 0.0, + f32::from_bits(1.0_f32.to_bits() - 1), + 1.0, + f32::from_bits(1.0_f32.to_bits() + 1), + f32::INFINITY, + ] { + let expected: Vec<_> = + queries.iter().map(|q| tree.radius_search(q[0], q[1], q[2], radius)).collect(); + for capacity in [0, 3, 64, n] { + let output = target + .radius_neighbors_resident(&runtime, &batch, radius, capacity as u32) + .unwrap(); + let downloaded = output.download(&runtime).unwrap(); + let truncated: Vec<_> = expected + .iter() + .map(|row| row[..row.len().min(capacity)].to_vec()) + .collect(); + assert_eq!( + downloaded.truncated_rows(), + truncated, + "n={n}, kind={kind}, radius={radius}, cap={capacity}" + ); + assert_eq!( + downloaded.total_counts(), + expected.iter().map(|r| r.len() as u32).collect::>() + ); + let complete = expected.iter().all(|r| r.len() <= capacity); + assert_eq!(downloaded.is_complete(), complete); + assert_eq!(downloaded.complete_rows().is_ok(), complete); + assert_eq!(downloaded.into_complete_rows().is_ok(), complete); + for (q, row) in queries.iter().zip(&expected) { + let count = points + .iter() + .filter(|p| { + let dx = p[0] - q[0]; + let dy = p[1] - q[1]; + let dz = p[2] - q[2]; + radius >= 0.0 && dx * dx + dy * dy + dz * dz <= radius * radius + }) + .count(); + assert_eq!(row.len(), count); + verified += row.len().min(capacity); + } + } + } + } + } + eprintln!( + "GPU_NEIGHBORHOOD verified {verified} radius entries with exact counts and overflow guards" + ); +} + +#[test] +fn gpu_neighborhood_owns_live_outputs_and_reuses_queries_across_callers() { + let runtime = WgpuRuntime::new_headless().unwrap(); + let (tree, points) = fixture(4097, 0); + let target = GpuKdTree::upload(&runtime, &tree).unwrap(); + let batch = GpuQueryBatch::upload(&runtime, &points).unwrap(); + let expected: Vec<_> = points.iter().map(|q| tree.nearest_k(q[0], q[1], q[2], 16)).collect(); + let live = target.nearest_k_resident(&runtime, &batch, 16).unwrap(); + drop(tree); + std::thread::scope(|scope| { + for _ in 0..4 { + let (runtime, target, batch, expected) = (&runtime, &target, &batch, &expected); + scope.spawn(move || { + for _ in 0..2 { + let out = target.nearest_k_resident(runtime, batch, 16).unwrap(); + assert_eq!(out.download(runtime).unwrap().complete_rows().unwrap(), expected); + } + }); + } + }); + let other = target.radius_neighbors_resident(&runtime, &batch, 0.0, 1).unwrap(); + assert_ne!(live.buffer(), other.buffer()); + drop(target); + drop(batch); + assert_eq!(live.download(&runtime).unwrap().complete_rows().unwrap(), expected); + assert_eq!(other.download(&runtime).unwrap().total_counts(), vec![1; points.len()]); +} + +#[test] +fn gpu_neighborhood_validates_empty_inputs_runtime_and_metadata() { + let runtime = WgpuRuntime::new_headless().unwrap(); + let tree = KdTree::from_slices(&[f32::MAX, -f32::MAX], &[0.0; 2], &[0.0; 2]); + let target = GpuKdTree::upload(&runtime, &tree).unwrap(); + let batch = GpuQueryBatch::upload(&runtime, &[[0.0; 3]]).unwrap(); + for kind in [GpuNeighborhoodKind::NearestK, GpuNeighborhoodKind::Radius] { + let output = if kind == GpuNeighborhoodKind::NearestK { + target.nearest_k_resident(&runtime, &batch, 64) + } else { + target.radius_neighbors_resident(&runtime, &batch, f32::INFINITY, 64) + } + .unwrap(); + let rows = output.download(&runtime).unwrap().into_complete_rows().unwrap(); + let expected = if kind == GpuNeighborhoodKind::NearestK { + tree.nearest_k(0.0, 0.0, 0.0, 64) + } else { + tree.radius_search(0.0, 0.0, 0.0, f32::INFINITY) + }; + assert_eq!(rows, vec![expected]); + } + let empty = GpuQueryBatch::upload(&runtime, &[]).unwrap(); + let output = target.nearest_k_resident(&runtime, &empty, 64).unwrap(); + assert_eq!(output.readback_bytes(), 0); + assert_eq!(output.download(&runtime).unwrap().transfers(), TransferStats::default()); + let empty_target = GpuKdTree::upload(&runtime, &KdTree::from_slices(&[], &[], &[])).unwrap(); + let empty_rows = empty_target + .radius_neighbors_resident(&runtime, &batch, 1.0, 0) + .unwrap() + .download(&runtime) + .unwrap(); + assert!(empty_rows.is_complete()); + assert_eq!(empty_rows.total_counts(), [0]); + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + assert!(GpuQueryBatch::upload(&runtime, &[[value, 0.0, 0.0]]).is_err()); + } + assert!(target.nearest_k_resident(&runtime, &empty, 65).is_err()); + let other = WgpuRuntime::new_headless().unwrap(); + let other_queries = GpuQueryBatch::upload(&other, &[]).unwrap(); + assert!(target.nearest_k_resident(&runtime, &other_queries, 0).is_err()); + assert!(target.radius_neighbors_resident(&other, &empty, 0.0, 0).is_err()); + assert!(output.download(&other).is_err()); + // A downstream kernel may corrupt externally exposed storage. Reject bad + // headers instead of trusting them for host allocation/indexing. + let malformed = GpuNeighborhoods { + buffer: runtime.device().create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&[9_u32, 1, 0, 0]), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + }), + queries: 1, + capacity: 1, + bytes: 16, + kind: GpuNeighborhoodKind::NearestK, + identity: runtime.profiling_identity(), + }; + assert!(malformed.download(&runtime).is_err()); +} + +#[test] +fn downstream_gpu_kernel_consumes_resident_counts_and_indices_without_readback() { + let runtime = WgpuRuntime::new_headless().unwrap(); + let (tree, points) = fixture(257, 2); + let target = GpuKdTree::upload(&runtime, &tree).unwrap(); + let queries = GpuQueryBatch::upload(&runtime, &points[..65]).unwrap(); + let neighborhoods = target.radius_neighbors_resident(&runtime, &queries, 0.0, 3).unwrap(); + assert_eq!(neighborhoods.execution_transfers(), TransferStats::default()); + let shader = runtime.device().create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("neighborhood-consumer"), + source: wgpu::ShaderSource::Wgsl( + r#" + struct Params { count: u32, capacity: u32, pad0: u32, pad1: u32 } + @group(0) @binding(0) var params: Params; + @group(0) @binding(1) var neighbors: array; + @group(0) @binding(2) var output: array; + @compute @workgroup_size(64) + fn consume(@builtin(global_invocation_id) id: vec3) { + if id.x >= params.count { return; } + let base = id.x * (2u + 2u * params.capacity); + output[3u * id.x] = neighbors[base]; + output[3u * id.x + 1u] = neighbors[base + 1u]; + output[3u * id.x + 2u] = neighbors[base + 2u]; + } + "# + .into(), + ), + }); + let pipeline = runtime.device().create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("neighborhood-consumer"), + layout: None, + module: &shader, + entry_point: Some("consume"), + compilation_options: Default::default(), + cache: None, + }); + let uniform = runtime.device().create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&[65_u32, 3, 0, 0]), + usage: wgpu::BufferUsages::UNIFORM, + }); + let output = runtime.device().create_buffer(&wgpu::BufferDescriptor { + label: None, + size: 65 * 12, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let group = runtime.device().create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &pipeline.get_bind_group_layout(0), + entries: &[ + wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() }, + wgpu::BindGroupEntry { + binding: 1, + resource: neighborhoods.buffer().as_entire_binding(), + }, + wgpu::BindGroupEntry { binding: 2, resource: output.as_entire_binding() }, + ], + }); + let staging = runtime.device().create_buffer(&wgpu::BufferDescriptor { + label: None, + size: 65 * 12, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let mut encoder = runtime.device().create_command_encoder(&Default::default()); + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: None, + timestamp_writes: None, + }); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &group, &[]); + pass.dispatch_workgroups(2, 1, 1); + } + encoder.copy_buffer_to_buffer(&output, 0, &staging, 0, 65 * 12); + runtime.queue().submit(Some(encoder.finish())); + drop(neighborhoods); + drop(queries); + drop(target); + let slice = staging.slice(..); + let (sender, receiver) = std::sync::mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |r| { + let _ = sender.send(r); + }); + runtime.device().poll(wgpu::Maintain::Wait); + receiver.recv().unwrap().unwrap(); + let mapped = slice.get_mapped_range(); + let words: &[u32] = bytemuck::cast_slice(&mapped); + let first = tree.radius_search(0.0, 0.0, 0.0, 0.0)[0].index as u32; + for row in words.chunks_exact(3) { + assert_eq!(row, [3, 257, first]); + } + drop(mapped); + staging.unmap(); +} + +#[test] +fn large_partial_dispatch_preserves_original_indices_and_zero_padding() { + let runtime = WgpuRuntime::new_headless().unwrap(); + let (tree, points) = fixture(65537, 0); + let target = GpuKdTree::upload(&runtime, &tree).unwrap(); + let queries = GpuQueryBatch::upload(&runtime, &points).unwrap(); + for kind in [GpuNeighborhoodKind::NearestK, GpuNeighborhoodKind::Radius] { + let output = if kind == GpuNeighborhoodKind::NearestK { + target.nearest_k_resident(&runtime, &queries, 1) + } else { + target.radius_neighbors_resident(&runtime, &queries, 0.0, 1) + } + .unwrap(); + let readback = output.download(&runtime).unwrap(); + assert_eq!(readback.total_counts(), vec![1; points.len()]); + for (i, row) in readback.complete_rows().unwrap().iter().enumerate() { + assert_eq!(row, &[Neighbor { index: i, distance_squared: 0.0 }]); + } + } + // Target smaller than k: unused result slots must remain initialized zero. + let tiny = GpuKdTree::upload(&runtime, &KdTree::from_slices(&[0.0], &[0.0], &[0.0])).unwrap(); + let query = GpuQueryBatch::upload(&runtime, &[[0.0; 3]]).unwrap(); + let output = tiny.nearest_k_resident(&runtime, &query, 64).unwrap(); + let staging = runtime.device().create_buffer(&wgpu::BufferDescriptor { + label: None, + size: output.bytes, + usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut encoder = runtime.device().create_command_encoder(&Default::default()); + encoder.copy_buffer_to_buffer(output.buffer(), 0, &staging, 0, output.bytes); + runtime.queue().submit(Some(encoder.finish())); + let slice = staging.slice(..); + let (sender, receiver) = std::sync::mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |r| { + let _ = sender.send(r); + }); + runtime.device().poll(wgpu::Maintain::Wait); + receiver.recv().unwrap().unwrap(); + let mapped = slice.get_mapped_range(); + let words: &[u32] = bytemuck::cast_slice(&mapped); + assert_eq!(&words[..4], &[1, 1, 0, 0]); + assert!(words[4..].iter().all(|&word| word == 0)); + drop(mapped); + staging.unmap(); +} diff --git a/crates/spatialrust-gpu/src/kernels/neighborhoods.rs b/crates/spatialrust-gpu/src/kernels/neighborhoods.rs new file mode 100644 index 0000000..3226c3c --- /dev/null +++ b/crates/spatialrust-gpu/src/kernels/neighborhoods.rs @@ -0,0 +1,470 @@ +use super::{buffer_bytes, check_compute, invalid, Arc, GpuKdTree, Neighbor, Pod, Zeroable}; +use spatialrust_core::{SpatialError, SpatialResult, TransferDirection, TransferStats}; +use wgpu::util::DeviceExt; + +use crate::WgpuRuntime; + +/// Maximum k for the bounded GPU nearest-neighbor accumulator. +pub const MAX_GPU_K_NEIGHBORS: u32 = 64; + +/// Membership and ordering of a resident neighborhood result. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GpuNeighborhoodKind { + /// Up to k points, stably sorted by distance with CPU accumulator ties. + NearestK, + /// Radius membership in CPU near-first leaf traversal order. + Radius, +} + +/// Immutable, explicitly uploaded finite query positions, reusable across trees +/// and neighborhood passes on the same runtime. No readback occurs at upload. +pub struct GpuQueryBatch { + buffer: Option, + count: u32, + identity: Arc<()>, + transfers: TransferStats, +} + +impl GpuQueryBatch { + /// Uploads XYZ queries as padded 16-byte records. Finite XYZ and supported + /// buffer/dispatch sizes are required; empty batches allocate no buffer. + pub fn upload(runtime: &WgpuRuntime, queries: &[[f32; 3]]) -> SpatialResult { + let limits = runtime.device().limits(); + let bytes = buffer_bytes(queries.len() as u64, 16, &limits)?; + check_compute(&limits, queries.len() as u64)?; + if queries.iter().flatten().any(|v| !v.is_finite()) { + return Err(invalid()); + } + let mut batch = Self { + buffer: None, + count: queries.len() as u32, + identity: runtime.profiling_identity(), + transfers: TransferStats::default(), + }; + if !queries.is_empty() { + let packed: Vec<_> = queries.iter().map(|p| [p[0], p[1], p[2], 0.0]).collect(); + batch.buffer = + Some(runtime.device().create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("neighborhood-queries"), + contents: bytemuck::cast_slice(&packed), + usage: wgpu::BufferUsages::STORAGE, + })); + batch.transfers.record(TransferDirection::HostToDevice, bytes); + } + Ok(batch) + } + + /// Number of XYZ query records. + #[must_use] + pub const fn query_count(&self) -> u32 { + self.count + } + + /// Payload bytes transferred during explicit upload. + #[must_use] + pub const fn upload_transfers(&self) -> TransferStats { + self.transfers + } +} + +/// Immutable GPU-resident bounded neighborhoods. Creating this result submits +/// GPU search without payload uploads/readbacks or a completion wait. Query/target +/// buffers are explicitly uploaded beforehand. Outputs are owned, never pooled. +pub struct GpuNeighborhoods { + buffer: wgpu::Buffer, + queries: u32, + capacity: u32, + bytes: u64, + kind: GpuNeighborhoodKind, + identity: Arc<()>, +} + +impl GpuNeighborhoods { + /// Output storage for subsequent GPU kernels on the same runtime. Each row + /// has `2 + 2*capacity` u32 words: stored count, total count, then pairs of + /// original target index and f32 squared-distance bits. Nearest rows are + /// sorted by distance; radius rows retain traversal order. Respect stored + /// counts; unused slots are zero. Empty results have an unused 8-byte buffer. + /// Radius consumers must handle total count > capacity before using a row. + #[must_use] + pub fn buffer(&self) -> &wgpu::Buffer { + &self.buffer + } + + /// Number of result rows, in query input order. + #[must_use] + pub const fn query_count(&self) -> u32 { + self.queries + } + + /// Maximum stored neighbors per row; radius counts may exceed this. + #[must_use] + pub const fn capacity(&self) -> u32 { + self.capacity + } + + /// Membership and ordering contract. + #[must_use] + pub const fn kind(&self) -> GpuNeighborhoodKind { + self.kind + } + + /// Payload bytes in a complete explicit readback, including row padding. + #[must_use] + pub const fn readback_bytes(&self) -> u64 { + self.bytes + } + + /// Host/device payload transfers performed by search itself. Target/query + /// uploads are separate; uniform control data is excluded. + #[must_use] + pub fn execution_transfers(&self) -> TransferStats { + TransferStats::default() + } + + /// Explicitly downloads rows and exact counts, waiting for queue completion. + /// Use complete_rows/into_complete_rows to reject radius overflow, or + /// truncated_rows to intentionally inspect a bounded traversal prefix. + pub fn download(&self, runtime: &WgpuRuntime) -> SpatialResult { + if !Arc::ptr_eq(&self.identity, &runtime.profiling_identity()) { + return Err(invalid()); + } + if self.queries == 0 { + return Ok(GpuNeighborhoodReadback { + rows: Vec::new(), + totals: Vec::new(), + transfers: TransferStats::default(), + }); + } + let staging = runtime.device().create_buffer(&wgpu::BufferDescriptor { + label: Some("neighborhood-readback"), + size: self.bytes, + usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut encoder = runtime.device().create_command_encoder(&Default::default()); + encoder.copy_buffer_to_buffer(&self.buffer, 0, &staging, 0, self.bytes); + runtime.queue().submit(Some(encoder.finish())); + let slice = staging.slice(..); + let (sender, receiver) = std::sync::mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |r| { + let _ = sender.send(r); + }); + runtime.device().poll(wgpu::Maintain::Wait); + receiver + .recv() + .map_err(|e| SpatialError::InvalidArgument(e.to_string()))? + .map_err(|e| SpatialError::InvalidArgument(e.to_string()))?; + let mapped = slice.get_mapped_range(); + let words: &[u32] = bytemuck::cast_slice(&mapped); + let stride = 2 + 2 * self.capacity as usize; + let mut rows = Vec::with_capacity(self.queries as usize); + let mut totals = Vec::with_capacity(self.queries as usize); + for row in words.chunks_exact(stride) { + if row[0] > self.capacity + || row[1] < row[0] + || (self.kind == GpuNeighborhoodKind::NearestK && row[0] != row[1]) + { + drop(mapped); + staging.unmap(); + return Err(SpatialError::InvalidArgument( + "GPU neighborhood metadata is invalid".to_owned(), + )); + } + rows.push( + (0..row[0] as usize) + .map(|j| Neighbor { + index: row[2 + 2 * j] as usize, + distance_squared: f32::from_bits(row[3 + 2 * j]), + }) + .collect(), + ); + totals.push(row[1]); + } + drop(mapped); + staging.unmap(); + let mut transfers = TransferStats::default(); + transfers.record(TransferDirection::DeviceToHost, self.bytes); + Ok(GpuNeighborhoodReadback { rows, totals, transfers }) + } +} + +/// CPU rows from an explicit readback, with radius-overflow information kept +/// separate from stored neighbors. Complete access rejects any incomplete row. +pub struct GpuNeighborhoodReadback { + rows: Vec>, + totals: Vec, + transfers: TransferStats, +} + +impl GpuNeighborhoodReadback { + /// True if every radius match fit its row (always true for nearest-k). + #[must_use] + pub fn is_complete(&self) -> bool { + self.rows.iter().zip(&self.totals).all(|(row, &total)| row.len() == total as usize) + } + + /// Exact radius match counts, or selected counts for nearest-k. + #[must_use] + pub fn total_counts(&self) -> &[u32] { + &self.totals + } + + /// Explicit bounded row access, allowing radius truncation. This is a + /// traversal prefix, not necessarily the closest points inside the radius. + #[must_use] + pub fn truncated_rows(&self) -> &[Vec] { + &self.rows + } + + /// Complete rows, or an error identifying the first overflowed radius row. + pub fn complete_rows(&self) -> SpatialResult<&[Vec]> { + self.check_complete()?; + Ok(&self.rows) + } + + /// Consumes this readback, rejecting radius overflow before returning rows. + pub fn into_complete_rows(self) -> SpatialResult>> { + self.check_complete()?; + Ok(self.rows) + } + + /// Payload bytes transferred during this explicit readback. + #[must_use] + pub const fn transfers(&self) -> TransferStats { + self.transfers + } + + fn check_complete(&self) -> SpatialResult<()> { + if let Some((i, (row, &total))) = self + .rows + .iter() + .zip(&self.totals) + .enumerate() + .find(|(_, (row, total))| row.len() < **total as usize) + { + return Err(SpatialError::InvalidArgument(format!( + "radius neighborhood row {i} contains {total} matches but stores only {}; increase capacity before consuming complete neighborhoods", row.len() + ))); + } + Ok(()) + } +} + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct Params { + count: u32, + root: u32, + capacity: u32, + bound: f32, +} + +pub(crate) struct NeighborhoodPipelines { + layout: wgpu::BindGroupLayout, + nearest: wgpu::ComputePipeline, + radius: wgpu::ComputePipeline, +} + +impl NeighborhoodPipelines { + fn new(device: &wgpu::Device) -> Self { + let entries: Vec<_> = (0..5) + .map(|binding| wgpu::BindGroupLayoutEntry { + binding, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: if binding == 0 { + wgpu::BufferBindingType::Uniform + } else { + wgpu::BufferBindingType::Storage { read_only: binding != 4 } + }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }) + .collect(); + let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("neighborhood-layout"), + entries: &entries, + }); + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("neighborhood-pipeline-layout"), + bind_group_layouts: &[&layout], + push_constant_ranges: &[], + }); + let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("neighborhood-shader"), + source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/neighborhoods.wgsl").into()), + }); + let pipeline = |entry| { + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("neighborhood-pipeline"), + layout: Some(&pipeline_layout), + module: &module, + entry_point: Some(entry), + compilation_options: Default::default(), + cache: None, + }) + }; + Self { nearest: pipeline("nearest_k"), radius: pipeline("radius_neighbors"), layout } + } +} + +fn plan( + queries: u64, + capacity: u64, + kind: GpuNeighborhoodKind, + limits: &wgpu::Limits, +) -> SpatialResult { + if capacity > u64::from((u32::MAX - 2) / 2) + || (kind == GpuNeighborhoodKind::NearestK && capacity > u64::from(MAX_GPU_K_NEIGHBORS)) + { + return Err(invalid()); + } + check_compute(limits, queries)?; + buffer_bytes(queries, 16, limits)?; + let stride = capacity.checked_mul(8).and_then(|v| v.checked_add(8)).ok_or_else(invalid)?; + let bytes = buffer_bytes(queries, stride, limits)?; + // Also bound shader word offsets and CPU decoding capacities before work. + let host_slots = queries + .checked_mul(capacity) + .and_then(|n| n.checked_mul(std::mem::size_of::() as u64)) + .ok_or_else(invalid)?; + let host_rows = + queries.checked_mul(std::mem::size_of::>() as u64).ok_or_else(invalid)?; + if bytes / 4 > u64::from(u32::MAX) + || host_slots > isize::MAX as u64 + || host_rows > isize::MAX as u64 + { + return Err(invalid()); + } + Ok(bytes) +} + +impl GpuKdTree { + /// Submits up to k nearest neighbors for each explicitly uploaded query. + /// Results stay on GPU, sorted by distance with the CPU KD-tree selection + /// and stable tie order. Supports k=0..=64; includes self/duplicate points. + /// No payload upload/readback or host wait occurs. Unsupported sizes error + /// before allocation. Finite targets/queries are required at upload; GPU + /// f32 rounding/underflow may differ across backends near search boundaries. + pub fn nearest_k_resident( + &self, + runtime: &WgpuRuntime, + queries: &GpuQueryBatch, + k: u32, + ) -> SpatialResult { + self.neighborhoods(runtime, queries, k, f32::INFINITY, GpuNeighborhoodKind::NearestK) + } + + /// Submits radius search (radius is not squared) with bounded row capacity. + /// All matches are counted on GPU; the first capacity matches are stored + /// in CPU traversal order. Capacity=0 performs count-only search. Negative/ + /// NaN radii produce empty rows; infinite/overflowing radius squares are + /// unbounded. Self/duplicate points are included. No silent truncation: + /// readback complete access rejects overflow, and device consumers must + /// inspect total counts. Target/query payloads are reused without copies. + pub fn radius_neighbors_resident( + &self, + runtime: &WgpuRuntime, + queries: &GpuQueryBatch, + radius: f32, + capacity: u32, + ) -> SpatialResult { + self.neighborhoods( + runtime, + queries, + capacity, + if radius < 0.0 { f32::NAN } else { radius * radius }, + GpuNeighborhoodKind::Radius, + ) + } + + fn neighborhoods( + &self, + runtime: &WgpuRuntime, + queries: &GpuQueryBatch, + capacity: u32, + bound: f32, + kind: GpuNeighborhoodKind, + ) -> SpatialResult { + let identity = runtime.profiling_identity(); + if !Arc::ptr_eq(&self.identity, &identity) || !Arc::ptr_eq(&queries.identity, &identity) { + return Err(invalid()); + } + let bytes = + plan(u64::from(queries.count), u64::from(capacity), kind, &runtime.device().limits())?; + let output = GpuNeighborhoods { + buffer: runtime.device().create_buffer(&wgpu::BufferDescriptor { + label: Some("neighborhood-results"), + size: bytes.max(8), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }), + queries: queries.count, + capacity, + bytes, + kind, + identity, + }; + // Owned wgpu buffers start at zero, including unused slots and headers + // on empty/invalid-gate paths. They never alias recycled scratch. + if queries.count == 0 + || self.root.is_none() + || bound.is_nan() + || (kind == GpuNeighborhoodKind::NearestK && capacity == 0) + { + return Ok(output); + } + let cache = runtime + .neighborhood_pipelines + .get_or_init(|| NeighborhoodPipelines::new(runtime.device())); + let params = Params { count: queries.count, root: self.root.unwrap(), capacity, bound }; + let uniform = runtime.device().create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("neighborhood-params"), + contents: bytemuck::bytes_of(¶ms), + usage: wgpu::BufferUsages::UNIFORM, + }); + let buffers = [ + &uniform, + self.nodes.as_ref().unwrap(), + self.points.as_ref().unwrap(), + queries.buffer.as_ref().unwrap(), + &output.buffer, + ]; + let entries: Vec<_> = buffers + .iter() + .enumerate() + .map(|(i, b)| wgpu::BindGroupEntry { + binding: i as u32, + resource: b.as_entire_binding(), + }) + .collect(); + let group = runtime.device().create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("neighborhood-bindings"), + layout: &cache.layout, + entries: &entries, + }); + let mut encoder = runtime.device().create_command_encoder(&Default::default()); + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("neighborhood-search"), + timestamp_writes: None, + }); + pass.set_pipeline(if kind == GpuNeighborhoodKind::NearestK { + &cache.nearest + } else { + &cache.radius + }); + pass.set_bind_group(0, &group, &[]); + pass.dispatch_workgroups(queries.count.div_ceil(64), 1, 1); + } + runtime.queue().submit(Some(encoder.finish())); + Ok(output) + } +} + +#[cfg(all(test, not(target_arch = "wasm32")))] +#[path = "neighborhood_tests.rs"] +mod tests; diff --git a/crates/spatialrust-gpu/src/lib.rs b/crates/spatialrust-gpu/src/lib.rs index abe6342..10ae695 100644 --- a/crates/spatialrust-gpu/src/lib.rs +++ b/crates/spatialrust-gpu/src/lib.rs @@ -50,6 +50,12 @@ mod image; pub use buffer::DeviceBuffer; pub use device::{GpuDevice, WgpuDevice}; +#[cfg(feature = "gpu-wgpu")] +pub use kernels::{ + GpuNeighborhoodKind, GpuNeighborhoodReadback, GpuNeighborhoods, GpuQueryBatch, + MAX_GPU_K_NEIGHBORS, +}; + #[cfg(feature = "gpu-wgpu")] pub use kernels::{ build_voxel_segments, build_voxel_segments_from_positions_gpu, diff --git a/crates/spatialrust-gpu/src/runtime.rs b/crates/spatialrust-gpu/src/runtime.rs index 24054b7..4de0a6b 100644 --- a/crates/spatialrust-gpu/src/runtime.rs +++ b/crates/spatialrust-gpu/src/runtime.rs @@ -24,6 +24,7 @@ pub struct WgpuRuntime { pub(crate) ransac_plane_pipelines: OnceLock, pub(crate) cluster_pipelines: OnceLock, pub(crate) nearest_pipeline: OnceLock, + pub(crate) neighborhood_pipelines: OnceLock, pub(crate) point_estimator_pipelines: OnceLock, max_gather_channels: u32, @@ -386,6 +387,7 @@ impl WgpuRuntime { ransac_plane_pipelines: OnceLock::new(), cluster_pipelines: OnceLock::new(), nearest_pipeline: OnceLock::new(), + neighborhood_pipelines: OnceLock::new(), point_estimator_pipelines: OnceLock::new(), max_gather_channels, upload_pool: GpuBufferPool::default(), diff --git a/crates/spatialrust-gpu/src/shaders/neighborhoods.wgsl b/crates/spatialrust-gpu/src/shaders/neighborhoods.wgsl new file mode 100644 index 0000000..9d87647 --- /dev/null +++ b/crates/spatialrust-gpu/src/shaders/neighborhoods.wgsl @@ -0,0 +1,118 @@ +struct Params { count: u32, root: u32, capacity: u32, bound: f32 } +struct Node { + split: f32, axis: u32, left: u32, right: u32, + start: u32, end: u32, pad0: u32, pad1: u32, +} +struct Point { xyz: vec3, index: u32 } +struct Neighbor { index: u32, distance: f32 } +@group(0) @binding(0) var params: Params; +@group(0) @binding(1) var nodes: array; +@group(0) @binding(2) var points: array; +@group(0) @binding(3) var queries: array>; +// Row: stored_count, total_count, then capacity pairs of (index, distance bits). +@group(0) @binding(4) var results: array; + +fn search(i: u32, radius: bool) { + let q = queries[i].xyz; + let base = i * (2u + 2u * params.capacity); + var best: array; + var count = 0u; + var total = 0u; + var worst = 0.0; + var worst_index = 0u; + var far_nodes: array; + var planes: array; + var depth = 0u; + var current = params.root; + loop { + let node = nodes[current]; + if node.axis != 0xffffffffu { + let diff = q[node.axis] - node.split; + far_nodes[depth] = select(node.left, node.right, diff <= 0.0); + planes[depth] = diff * diff; + depth += 1u; + current = select(node.right, node.left, diff <= 0.0); + continue; + } + for (var j = node.start; j < node.end; j += 1u) { + let point = points[j]; + let dx = point.xyz.x - q.x; + let dy = point.xyz.y - q.y; + let dz = point.xyz.z - q.z; + let distance = dx * dx + dy * dy + dz * dz; + if radius { + if distance <= params.bound { + total += 1u; + if count < params.capacity { + results[base + 2u + count * 2u] = point.index; + results[base + 3u + count * 2u] = bitcast(distance); + count += 1u; + } + } + } else { + if count < params.capacity { + best[count] = Neighbor(point.index, distance); + count += 1u; + if count == 1u || distance > worst { + worst = distance; + worst_index = count - 1u; + } + } else if distance < worst { + best[worst_index] = Neighbor(point.index, distance); + worst = best[0].distance; + worst_index = 0u; + for (var k = 1u; k < count; k += 1u) { + if best[k].distance > worst { + worst = best[k].distance; + worst_index = k; + } + } + } + } + } + var found = false; + loop { + if depth == 0u { break; } + depth -= 1u; + let visit = select(count < params.capacity || planes[depth] < worst, + planes[depth] <= params.bound, radius); + if visit { + current = far_nodes[depth]; + found = true; + break; + } + } + if !found { break; } + } + if !radius { + // Stable distance sort matches the CPU accumulator's final sort_by, + // including ties between retained/replaced accumulator slots. + for (var k = 1u; k < count; k += 1u) { + let candidate = best[k]; + var j = k; + loop { + if j == 0u { break; } + if best[j - 1u].distance <= candidate.distance { break; } + best[j] = best[j - 1u]; + j -= 1u; + } + best[j] = candidate; + } + for (var k = 0u; k < count; k += 1u) { + results[base + 2u + k * 2u] = best[k].index; + results[base + 3u + k * 2u] = bitcast(best[k].distance); + } + total = count; + } + results[base] = count; + results[base + 1u] = total; +} + +@compute @workgroup_size(64) +fn nearest_k(@builtin(global_invocation_id) id: vec3) { + if id.x < params.count { search(id.x, false); } +} +@compute @workgroup_size(64) +fn radius_neighbors(@builtin(global_invocation_id) id: vec3) { + if id.x < params.count { search(id.x, true); } +} diff --git a/crates/spatialrust-gpu/tests/neighborhood_benchmark.rs b/crates/spatialrust-gpu/tests/neighborhood_benchmark.rs new file mode 100644 index 0000000..a46a7a7 --- /dev/null +++ b/crates/spatialrust-gpu/tests/neighborhood_benchmark.rs @@ -0,0 +1,100 @@ +#![cfg(all(feature = "gpu-wgpu", not(target_arch = "wasm32")))] + +use spatialrust_gpu::{GpuKdTree, GpuQueryBatch, WgpuRuntime}; +use spatialrust_search::{KdTree, NearestNeighborIndex, RadiusSearchIndex}; +use std::time::Instant; + +#[test] +#[ignore = "completed and resident neighborhood benchmark; run alone in release mode"] +fn completed_and_resident_neighborhoods() { + let runtime = WgpuRuntime::new_headless().unwrap(); + let info = runtime.adapter_info(); + println!( + "GPU_NEIGHBORHOOD_ENV backend={} device_type={} cpu_workers=1", + info.backend, info.device_type + ); + for n in [257, 4097, 16384, 65537] { + for geometry in 0..2 { + let points: Vec<_> = (0..n) + .map(|i| { + let j = n - i - 1; + if geometry == 0 { + [(j % 41) as f32 * 0.5, (j / 41 % 41) as f32 * 0.5, (j / 1681) as f32 * 0.5] + } else { + [j as f32 * 0.5, 0.0, 0.0] + } + }) + .collect(); + let xyz: [Vec<_>; 3] = std::array::from_fn(|a| points.iter().map(|p| p[a]).collect()); + let tree = KdTree::from_slices(&xyz[0], &xyz[1], &xyz[2]); + let target = GpuKdTree::upload(&runtime, &tree).unwrap(); + let queries: Vec<_> = + points.iter().map(|p| [p[0] + 0.0625, p[1] - 0.03125, p[2]]).collect(); + let resident_queries = GpuQueryBatch::upload(&runtime, &queries).unwrap(); + for kind in 0..2 { + let cpu = || { + queries + .iter() + .map(|q| { + if kind == 0 { + tree.nearest_k(q[0], q[1], q[2], 16) + } else { + tree.radius_search(q[0], q[1], q[2], 0.75) + } + }) + .collect::>() + }; + let run = |batch: &GpuQueryBatch| { + if kind == 0 { + target.nearest_k_resident(&runtime, batch, 16).unwrap() + } else { + target.radius_neighbors_resident(&runtime, batch, 0.75, 32).unwrap() + } + }; + let gpu = || { + let batch = GpuQueryBatch::upload(&runtime, &queries).unwrap(); + run(&batch).download(&runtime).unwrap().into_complete_rows().unwrap() + }; + let expected = cpu(); + for _ in 0..3 { + assert_eq!(gpu(), expected); + assert_eq!( + run(&resident_queries) + .download(&runtime) + .unwrap() + .into_complete_rows() + .unwrap(), + expected + ); + } + let (mut cpu_ms, mut gpu_ms, mut resident_ms) = + (Vec::new(), Vec::new(), Vec::new()); + for round in 0..20 { + for slot in 0..3 { + let stage = (round + slot) % 3; + let start = Instant::now(); + if stage == 2 { + let output = run(&resident_queries); + runtime.wait_idle(); + resident_ms.push(start.elapsed().as_secs_f64() * 1000.0); + assert_eq!( + output.download(&runtime).unwrap().into_complete_rows().unwrap(), + expected + ); + } else { + let rows = if stage == 0 { cpu() } else { gpu() }; + let elapsed = start.elapsed().as_secs_f64() * 1000.0; + if stage == 0 { + cpu_ms.push(elapsed) + } else { + gpu_ms.push(elapsed) + } + assert_eq!(rows, expected); + } + } + } + println!("GPU_NEIGHBORHOOD_BENCH n={n} geometry={geometry} kind={kind} cpu_ms={cpu_ms:?} gpu_ms={gpu_ms:?} resident_ms={resident_ms:?}"); + } + } + } +} diff --git a/crates/spatialrust/tests/gpu_neighborhood_smoke.rs b/crates/spatialrust/tests/gpu_neighborhood_smoke.rs new file mode 100644 index 0000000..24c6d39 --- /dev/null +++ b/crates/spatialrust/tests/gpu_neighborhood_smoke.rs @@ -0,0 +1,33 @@ +#![cfg(all(feature = "gpu-wgpu", not(target_arch = "wasm32")))] + +use spatialrust::gpu::{GpuKdTree, GpuQueryBatch, WgpuRuntime}; +use spatialrust::search::{KdTree, NearestNeighborIndex, RadiusSearchIndex}; + +#[test] +#[ignore = "requires a native wgpu adapter; run explicitly with --include-ignored"] +fn facade_resident_neighborhoods_reuse_queries_and_reject_radius_overflow() { + let runtime = WgpuRuntime::new_headless().unwrap(); + let x: Vec<_> = (0..257).map(|i| i as f32 * 0.5).collect(); + let yz = vec![0.0; x.len()]; + let tree = KdTree::from_slices(&x, &yz, &yz); + let target = GpuKdTree::upload(&runtime, &tree).unwrap(); + let queries = GpuQueryBatch::upload(&runtime, &[[0.0; 3], [64.0, 0.0, 0.0]]).unwrap(); + let knn = target.nearest_k_resident(&runtime, &queries, 16).unwrap(); + assert_eq!(knn.execution_transfers().device_to_host_bytes(), 0); + let rows = knn.download(&runtime).unwrap().into_complete_rows().unwrap(); + assert_eq!(rows, [tree.nearest_k(0.0, 0.0, 0.0, 16), tree.nearest_k(64.0, 0.0, 0.0, 16)]); + let radius = target.radius_neighbors_resident(&runtime, &queries, 1.0, 1).unwrap(); + let downloaded = radius.download(&runtime).unwrap(); + assert_eq!(downloaded.total_counts(), [3, 5]); + assert!(downloaded.complete_rows().is_err()); + let full = target + .radius_neighbors_resident(&runtime, &queries, 1.0, 5) + .unwrap() + .download(&runtime) + .unwrap(); + assert_eq!( + full.complete_rows().unwrap(), + [tree.radius_search(0.0, 0.0, 0.0, 1.0), tree.radius_search(64.0, 0.0, 0.0, 1.0)] + ); + assert_eq!(full.transfers().device_to_host_bytes(), 2 * (8 + 5 * 8)); +} diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 06bfd76..6bc167b 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -23,6 +23,15 @@ caller-provided runtime and reports all input/grid payload and padded output transfers; errors propagate. Without a radius, covariance remains CPU k-NN. The legacy native shared-runtime covariance route keeps its fallback contract. +Reusable finite `GpuQueryBatch` uploads and owned `GpuNeighborhoods` make the +next feature-compute boundary explicit. GPU KD-tree k-NN/radius search writes +fixed-capacity rows with stored/total counts and original index/distance pairs. +Nearest-k rows retain CPU stable ordering; radius rows retain traversal order +and expose overflow. Complete CPU readback rejects incomplete neighborhoods; +GPU consumers must inspect counts before using rows. Target/query payloads +remain resident between passes. This adds no core/search dependency and does +not change feature algorithm placement; FPFH/ISS/MLS device compute follows. + | Area | Decision | | --- | --- | | Core model | `SpatialTensor` + `SpatialIndex` + `SpatialAlgorithm` + `SpatialRuntime` | diff --git a/docs/FEATURE_MATRIX.md b/docs/FEATURE_MATRIX.md index f3bda8d..451d67b 100644 --- a/docs/FEATURE_MATRIX.md +++ b/docs/FEATURE_MATRIX.md @@ -6,6 +6,21 @@ should normally depend on `spatialrust` and select only the profiles they need. ## User-facing profiles +`gpu-wgpu` exposes `GpuQueryBatch::upload` and +`GpuKdTree::{nearest_k_resident,radius_neighbors_resident}` for shared GPU +neighborhoods. Target/query uploads are explicit and reusable; search submits +without payload copies or a host wait and returns owned `GpuNeighborhoods`. +Each row stores two u32 counts plus bounded index/distance pairs. k=0..=64 is +stably sorted by distance with CPU selection ties. Radius membership retains +CPU traversal order, counts all matches, and supports capacity=0 count-only +search. `download` accounts for padded readback bytes; complete row access +rejects overflow while `truncated_rows` explicitly returns a traversal prefix. +Self/duplicate points are included. GPU feature consumers must check counts +and preserve query/target semantics. Feature algorithms remain separate CPU +APIs until their device compute is added. Native correctness and completed/ +resident measurements are in +[the neighborhood note](../notes/2026-10-12_gpu_bounded_neighborhoods.md). + `register-icp-gpu` enables the explicit `IcpRegistration::align_with_gpu_correspondences(runtime, source, target)` API. The target KD-tree is built on CPU and uploaded once; each update and fitness diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index d1db1f1..eb9f49e 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -1128,6 +1128,6 @@ Dependencies and explicit transfers come before broader algorithm coverage. | --- | --- | --- | --- | | ICP correspondences | Complete | Reusable runtime-owned GPU KD-tree and explicit batched nearest uploads/readback; opt-in ICP update/fitness integration with CPU solver | CPU/brute nearest ties and gates, ownership/reuse/capacity, pose/trace/receipt checks and completed query/registration comparisons; `notes/2026-10-12_gpu_icp_correspondences.md` | | GICP correspondences | Complete | Reuse GPU target for unbounded matching followed by host gating; explicit radius covariance runtime and full receipt with CPU Hessian/pose solver | Pose/stopping/gating/receipt/unequal-cloud checks, error propagation, completed comparisons for CPU/radius-GPU covariance; `notes/2026-10-12_gpu_gicp_correspondences.md` | -| Feature neighborhoods | Planned | Bounded resident radius/k-NN results shared by descriptors, keypoints and smoothing | Exact membership/order/capacity, degenerate/tied supports, explicit transfer receipts and completed-work measurements | +| Feature neighborhoods | Resident API complete; feature consumers planned | Reused target/query buffers, GPU k-NN and bounded radius rows with exact counts/overflow, explicit readback and raw storage for device consumers | CPU/brute membership/order, degenerate/tied/partial/ownership/capacity cases, device consumer test and completed/resident measurements; `notes/2026-10-12_gpu_bounded_neighborhoods.md` | | FPFH/ISS/MLS compute | Planned | Device feature/smoothing phases on the shared neighborhood representation | CPU/math references, attribute/schema preservation, numeric tolerances where reductions change, transfer-inclusive benefit | | Registration updates | Planned | Explicit correspondence reduction/pose-update boundaries after GPU search is established | Pose/convergence/fitness references, numerical conditioning, source/target residency and end-to-end registration timing | diff --git a/notes/2026-10-12_gpu_bounded_neighborhoods.md b/notes/2026-10-12_gpu_bounded_neighborhoods.md new file mode 100644 index 0000000..5a5cc69 --- /dev/null +++ b/notes/2026-10-12_gpu_bounded_neighborhoods.md @@ -0,0 +1,187 @@ +# Bounded resident GPU neighborhoods + +Date: 2026-10-12. Named slice: reusable query uploads and resident k-NN/radius +results, building on the GPU KD-tree from PR #164 and preserving PR #165 APIs. + +## Data and execution boundary + +`GpuQueryBatch::upload(runtime, xyz)` validates finite XYZ, u32 counts, device +buffer/dispatch limits and host capacity, then explicitly uploads padded +16-byte query records. Empty batches allocate no query buffer. The owned, +immutable batch can outlive its host input and be reused across target trees +and passes on the same runtime. The target's CPU-built KD-tree upload remains +explicit and separately accounted, as before. + +`GpuKdTree::nearest_k_resident` supports k=0..=64. One GPU invocation traverses +the balanced target tree with a 32-level deferred far-child stack, reproduces +the CPU accumulator's first-worst-slot replacement and strict tie rejection, +then stably sorts the retained slots by distance. This preserves the CPU +`nearest_k` membership AND ordered ties on the tested fixtures. Self and +duplicate points are included; feature consumers may request k+1 and apply +their own self-exclusion policy. Finite XYZ can still produce infinite squared +distances, which retain the CPU selection/pruning behavior. + +`radius_neighbors_resident(runtime, queries, radius, capacity)` uses a radius, +not a squared gate. It preserves CPU near-first leaf traversal order, counts +ALL radius matches even after the row is full, and stores only the first +capacity matches. Capacity=0 performs count-only search. Negative/NaN radii +produce empty rows. Positive infinity and overflowing positive radius squares +are unbounded, as on CPU. A truncated row is a traversal prefix, not a closest- +point selection. Radius capacity is not restricted to 64, but buffers, shader +word offsets and decoding capacities must fit validated limits. + +Search submits device work and returns owned `GpuNeighborhoods` without a +payload upload, readback or completion wait. Uniform controls are excluded from +payload receipts. Target/query buffers are reused; result buffers are separately +owned and never returned to scratch pools. Each row contains `2 + 2*capacity` +u32 words: stored count, total count, then original target index / f32 squared- +distance bit pairs. k-NN rows are sorted; radius rows retain traversal order. +Unused slots are zero, including fast-path rows. Empty results expose an unused +8-byte physical buffer, with zero logical rows/readback bytes. + +`GpuNeighborhoods::buffer` lets subsequent GPU kernels consume that layout on +the same runtime. Consumers must respect stored counts and handle radius total +counts larger than capacity before relying on complete neighborhoods. Inputs +and outputs are immutable application data; custom kernels must not overwrite +them. Public search/download methods validate runtime ownership even for empty +or invalid-radius paths. The shader needs four storage bindings, one uniform +and 64 invocations per workgroup. Row size, total storage, u32 shader offsets, +CPU row/neighbor allocation bounds and dispatch are validated before allocation. + +`download` explicitly waits for completion and copies the FULL padded row +buffer: `queries * (8 + 8*capacity)` bytes. It returns exact counts and stored +rows; malformed GPU counts reject before host neighbor allocation/indexing. +`complete_rows` and `into_complete_rows` reject the first overflowed radius row. +`truncated_rows` is the explicitly named opt-in access to bounded prefixes. +`is_complete` and `total_counts` let callers choose a larger capacity or bounded +query batches rather than silently lose neighbors. + +Target upload is `16*T + 32*tree_nodes` H2D bytes. Query upload is `16*Q` H2D. +Resident search has zero payload transfer bytes. Download is `Q*(8+8*C)` D2H, +including padding. No global CPU nearest search or automatic CPU fallback is +used by these new GPU methods. Existing nearest-one ICP/GICP and feature APIs +are unchanged. FPFH/ISS/MLS remain CPU algorithms; their device compute and +feature-specific neighborhood integration are the next slices, not completed +by merely exposing this resident data model. + +GPU f32 rounding/underflow may differ across backends near gate/tie boundaries; +exact parity below is for selected native fixtures, not every floating-point +input. Native synchronous download/consumer execution is verified. WASM build +success alone does not establish browser execution. + +## Correctness + +- 467,364 selected nearest neighbors match exact CPU membership/ordered ties + and independent exhaustive sorted distances, across k=0/1/2/8/16/32/64 and + target sizes 1/16/17/63/64/65/255/257/513. Fixtures include regular lattice, + reversed line, coincident points and repeated symmetric coordinates. +- 199,846 stored radius entries match CPU traversal prefixes, with exact total + counts also checked by exhaustive membership scans. Zero/one-ULP/negative/ + NaN/infinite radii, capacities 0/3/64/full-target and missing matches cover + complete and overflow paths. +- Four concurrent callers twice each reuse a 4,097-query batch and target; + 524,416 selected neighbors remain exact. Earlier outputs survive later + k-NN/radius calls and CPU tree, query and target drops. +- Two 65,537-query partial dispatches preserve all original self indices + (131,074 neighbors). A raw readback confirms unused slots remain zero when + target count is smaller than k. Empty/invalid/finite-overflow data and wrong + target/query/output runtime ownership, unsupported counts/capacities and + malformed result metadata reject as specified. +- A downstream native GPU test consumes radius counts and original indices + directly from resident storage, before any neighborhood readback; it detects + 257 total / 3 stored neighbors per row. The queue retains resources even + after the owning target/query/result handles are dropped. One native facade + test checks reuse, complete/overflow access and padded transfer receipts. + +Final validation: all GPU features passed 196 tests (23 ignored), registration +passed 35 unit tests (1 ignored) plus all six explicit GPU ICP/GICP regression +tests, and the explicit native facade test passed. GPU all-feature/all-target +clippy passed with warnings denied. Minimal GPU, GPU image/AoSoA WASM, +formatting and diff checks passed. Source comparison against parent 8b1226f +confirms the nearest-one shader/upload/search and search/feature/registration +algorithm files are unchanged; only the new neighborhood API/cache is added. + +## Measurement method + +Native wgpu Vulkan discrete-class adapter, release build, CPU baseline one +worker. No adapter model is recorded. Target CPU construction and target GPU +upload occur once before timing for BOTH reused-index comparisons. Query XYZ +are supplied by the caller and offset x by 1/16, y by -1/32. Fixtures are a +half-unit lattice and a reversed half-unit line, with equal target/query counts. +Nearest-k uses k=16; radius uses 0.75 and capacity=32 (all selected fixture rows +fit). Two isolated surveys each contain 20 cyclic-order CPU / completed GPU / +resident GPU triples for 16 conditions. A CPU reference and three checked GPU +warmups of each path precede timing. No other local builds/tests run during +measurements; output checks and result drop are outside measured intervals. + +Completed GPU time includes query packing/upload, owned result allocation, +search, queue completion, padded download, metadata/overflow checking and CPU +row decoding. CPU time includes all scalar KD-tree queries and returned row +vectors. Resident time separately includes result allocation, submission and +queue completion with target/queries already uploaded, and ends before CPU +readback/decoding. Resident results are downloaded/checked AFTER their timer. +Resident columns therefore describe a different transfer boundary and are NOT +used to claim a CPU speedup or an end-to-end feature speedup. No GPU-only pass +timestamps are claimed. Target uploads and CPU construction are excluded from +these reused-index query timings and must be added for cold single-use calls. + +Tables give medians in ms and the geometric mean paired CPU/completed-GPU ratio; +above one favors the completed GPU API. Brackets are separate 95% paired log- +ratio bootstrap intervals with 10,000 resamples and deterministic seed 20261012 +plus record index. Geometry, query count, k/capacity, input reuse, CPU worker +count and transfer boundary matter. Synthetic fixtures do not establish a +universal crossover or performance on an entire feature/registration pipeline. + +| Survey | Queries | Geometry | Query | CPU ms | Complete GPU ms | Resident completion ms | CPU/complete GPU [95% interval] | +| --- | ---: | --- | --- | ---: | ---: | ---: | ---: | +| 1 | 257 | lattice | k=16 | 0.192 | 1.129 | 0.218 | 0.167 [0.157, 0.176] | +| 1 | 257 | lattice | radius=0.75, C=32 | 0.103 | 1.096 | 0.165 | 0.096 [0.091, 0.102] | +| 1 | 257 | reversed line | k=16 | 0.146 | 1.167 | 0.200 | 0.128 [0.124, 0.133] | +| 1 | 257 | reversed line | radius=0.75, C=32 | 0.044 | 1.162 | 0.153 | 0.038 [0.037, 0.039] | +| 1 | 4,097 | lattice | k=16 | 4.731 | 1.833 | 0.563 | 2.598 [2.525, 2.674] | +| 1 | 4,097 | lattice | radius=0.75, C=32 | 2.721 | 1.577 | 0.354 | 1.739 [1.655, 1.816] | +| 1 | 4,097 | reversed line | k=16 | 2.421 | 1.538 | 0.339 | 1.548 [1.479, 1.613] | +| 1 | 4,097 | reversed line | radius=0.75, C=32 | 0.711 | 1.251 | 0.158 | 0.554 [0.531, 0.571] | +| 1 | 16,384 | lattice | k=16 | 25.935 | 4.124 | 1.653 | 6.236 [5.894, 6.574] | +| 1 | 16,384 | lattice | radius=0.75, C=32 | 17.572 | 5.656 | 1.769 | 3.162 [3.024, 3.303] | +| 1 | 16,384 | reversed line | k=16 | 9.875 | 5.373 | 2.449 | 1.857 [1.794, 1.920] | +| 1 | 16,384 | reversed line | radius=0.75, C=32 | 3.309 | 6.393 | 2.796 | 0.545 [0.520, 0.576] | +| 1 | 65,537 | lattice | k=16 | 125.638 | 34.424 | 8.370 | 4.124 [3.615, 4.693] | +| 1 | 65,537 | lattice | radius=0.75, C=32 | 75.407 | 31.469 | 12.887 | 2.574 [2.262, 2.925] | +| 1 | 65,537 | reversed line | k=16 | 43.134 | 28.232 | 14.003 | 1.316 [1.183, 1.449] | +| 1 | 65,537 | reversed line | radius=0.75, C=32 | 13.534 | 18.782 | 6.446 | 0.771 [0.733, 0.815] | +| 2 | 257 | lattice | k=16 | 0.220 | 1.087 | 0.214 | 0.196 [0.191, 0.200] | +| 2 | 257 | lattice | radius=0.75, C=32 | 0.102 | 1.047 | 0.146 | 0.098 [0.096, 0.099] | +| 2 | 257 | reversed line | k=16 | 0.153 | 1.162 | 0.203 | 0.130 [0.126, 0.135] | +| 2 | 257 | reversed line | radius=0.75, C=32 | 0.044 | 1.063 | 0.139 | 0.042 [0.041, 0.043] | +| 2 | 4,097 | lattice | k=16 | 4.537 | 1.659 | 0.481 | 2.740 [2.702, 2.780] | +| 2 | 4,097 | lattice | radius=0.75, C=32 | 2.704 | 1.402 | 0.279 | 1.898 [1.837, 1.957] | +| 2 | 4,097 | reversed line | k=16 | 2.367 | 1.387 | 0.259 | 1.648 [1.555, 1.741] | +| 2 | 4,097 | reversed line | radius=0.75, C=32 | 0.714 | 1.240 | 0.155 | 0.575 [0.565, 0.583] | +| 2 | 16,384 | lattice | k=16 | 26.155 | 4.102 | 1.648 | 6.409 [6.080, 6.760] | +| 2 | 16,384 | lattice | radius=0.75, C=32 | 16.635 | 5.458 | 1.775 | 2.848 [2.649, 3.037] | +| 2 | 16,384 | reversed line | k=16 | 9.965 | 5.896 | 2.597 | 1.418 [1.259, 1.594] | +| 2 | 16,384 | reversed line | radius=0.75, C=32 | 3.331 | 8.049 | 3.498 | 0.425 [0.412, 0.439] | +| 2 | 65,537 | lattice | k=16 | 144.535 | 23.909 | 8.979 | 5.041 [4.044, 6.106] | +| 2 | 65,537 | lattice | radius=0.75, C=32 | 101.913 | 34.310 | 14.801 | 3.046 [2.889, 3.209] | +| 2 | 65,537 | reversed line | k=16 | 58.495 | 42.144 | 15.069 | 1.449 [1.265, 1.666] | +| 2 | 65,537 | reversed line | radius=0.75, C=32 | 22.245 | 25.248 | 9.141 | 0.829 [0.749, 0.908] | + +At 65,537 lattice queries, completed k-NN improves by 4.124x / 5.041x and radius +by 2.574x / 3.046x across the two surveys. Both retain transfer-inclusive gains. +257-query cases favor CPU. Reversed-line radius queries favor CPU at every size +in both surveys, even though the resident-only boundary is cheaper. Larger +line k-NN cases have smaller and more variable gains than the lattice. All +adverse cases and separate intervals are retained. These differences justify +explicit GPU APIs without changing any automatic feature placement thresholds. + +Reproduce each survey alone: + +```powershell +cargo test -p spatialrust-gpu --features gpu-wgpu --release --test neighborhood_benchmark -- --ignored --nocapture +``` + +Local raw logs/statistical records stay under +`C:/Users/rsasa/Workspace/SpatialRust/target/automation`. Benchmark output prints +backend, device class and CPU worker count only; published tables contain no +adapter model or private sensor/customer data.