spatialrust_gpu/kernels/
euclidean_cluster.rs

1use spatialrust_core::{SpatialError, SpatialResult, TransferDirection, TransferStats};
2
3use crate::kernels::build_voxel_segments_from_positions_gpu;
4
5pub use spatialrust_search::euclidean_cluster_roots as euclidean_cluster_roots_grid;
6
7/// Runs the GPU sparse-grid stage and returns its transfer accounting together
8/// with deterministic host-side component roots.
9pub fn euclidean_cluster_roots_gpu_with_receipt(
10    runtime: &crate::runtime::WgpuRuntime,
11    x: &[f32],
12    y: &[f32],
13    z: &[f32],
14    cluster_tolerance: f32,
15) -> SpatialResult<(Vec<u32>, TransferStats)> {
16    if x.len() != y.len() || x.len() != z.len() {
17        return Err(SpatialError::InvalidArgument("xyz arrays must have equal length".to_owned()));
18    }
19    if cluster_tolerance <= 0.0 || cluster_tolerance.is_nan() {
20        return Err(SpatialError::InvalidArgument("cluster_tolerance must be positive".to_owned()));
21    }
22    if x.is_empty() {
23        return Ok((Vec::new(), TransferStats::default()));
24    }
25
26    // Key generation, sorting, and sparse-cell compaction are GPU kernels. The
27    // final connected-component union-find intentionally runs on the host so
28    // the public GPU path has deterministic minimum-root semantics and does not
29    // pretend that a CPU fallback is an accelerator kernel.
30    let mut origin = [x[0], y[0], z[0]];
31    for index in 1..x.len() {
32        origin[0] = origin[0].min(x[index]);
33        origin[1] = origin[1].min(y[index]);
34        origin[2] = origin[2].min(z[index]);
35    }
36    let segments =
37        build_voxel_segments_from_positions_gpu(runtime, x, y, z, origin, 1.0 / cluster_tolerance)?;
38    let mut transfers = TransferStats::default();
39    transfers
40        .record(TransferDirection::HostToDevice, (x.len() * 3 * std::mem::size_of::<f32>()) as u64);
41    let metadata_bytes = (segments.keys.len() * 4 * std::mem::size_of::<i32>()
42        + segments.cell_starts.len() * std::mem::size_of::<u32>()
43        + segments.point_indices.len() * std::mem::size_of::<u32>()
44        + std::mem::size_of::<u32>()) as u64;
45    transfers.record(TransferDirection::DeviceToHost, metadata_bytes);
46
47    let roots = spatialrust_search::euclidean_cluster_roots_from_segments(
48        x,
49        y,
50        z,
51        cluster_tolerance,
52        &segments.keys,
53        &segments.point_indices,
54        &segments.cell_starts,
55        &segments.cell_counts,
56    )?;
57    Ok((roots, transfers))
58}
59
60/// Connected-component roots via uniform-grid union-find.
61/// Computes Euclidean component roots using GPU sparse-grid construction.
62pub fn euclidean_cluster_roots_gpu(
63    runtime: &crate::runtime::WgpuRuntime,
64    x: &[f32],
65    y: &[f32],
66    z: &[f32],
67    cluster_tolerance: f32,
68) -> SpatialResult<Vec<u32>> {
69    euclidean_cluster_roots_gpu_with_receipt(runtime, x, y, z, cluster_tolerance)
70        .map(|(roots, _)| roots)
71}
spatialrust_gpu/kernels/euclidean_cluster.rs

spatialrust_gpu/kernels/
euclidean_cluster.rs
