Skip to main content

spatialrust_gpu/kernels/
gpu_segments.rs

1//! GPU-resident voxel segment buffers used across chained kernels.
2
3use spatialrust_core::SpatialResult;
4
5use crate::kernels::voxel_compact::{finalize_segments_from_readback, read_segment_metadata};
6use crate::kernels::voxel_segments::VoxelSegments;
7use crate::runtime::WgpuRuntime;
8
9/// Voxel segment metadata kept on the GPU between sort/compact and reduce passes.
10pub struct GpuVoxelSegments {
11    cell_count: u32,
12    point_count: u32,
13    keys: wgpu::Buffer,
14    point_indices: wgpu::Buffer,
15    cell_starts: wgpu::Buffer,
16}
17
18impl GpuVoxelSegments {
19    /// Creates GPU segment buffers.
20    pub(crate) fn new(
21        cell_count: u32,
22        point_count: u32,
23        keys: wgpu::Buffer,
24        point_indices: wgpu::Buffer,
25        cell_starts: wgpu::Buffer,
26    ) -> Self {
27        Self { cell_count, point_count, keys, point_indices, cell_starts }
28    }
29
30    /// Returns the number of occupied voxel cells.
31    #[must_use]
32    pub fn cell_count(&self) -> u32 {
33        self.cell_count
34    }
35
36    /// Returns the source point count represented by these segments.
37    #[must_use]
38    pub fn point_count(&self) -> u32 {
39        self.point_count
40    }
41
42    /// Returns the GPU buffer of unique voxel keys.
43    #[must_use]
44    pub fn keys_buffer(&self) -> &wgpu::Buffer {
45        &self.keys
46    }
47
48    /// Returns the GPU buffer of sorted point indices grouped by cell.
49    #[must_use]
50    pub fn point_indices_buffer(&self) -> &wgpu::Buffer {
51        &self.point_indices
52    }
53
54    /// Returns the GPU buffer of per-cell start offsets into `point_indices`.
55    #[must_use]
56    pub fn cell_starts_buffer(&self) -> &wgpu::Buffer {
57        &self.cell_starts
58    }
59
60    /// Readbacks segment metadata into the CPU-side representation.
61    pub fn to_voxel_segments(&self, runtime: &WgpuRuntime) -> SpatialResult<VoxelSegments> {
62        if self.cell_count == 0 {
63            return Ok(VoxelSegments {
64                keys: Vec::new(),
65                point_indices: Vec::new(),
66                cell_starts: Vec::new(),
67                cell_counts: Vec::new(),
68            });
69        }
70
71        let (keys, cell_starts, cell_counts, point_indices) =
72            read_segment_metadata(runtime, self, self.cell_count as usize)?;
73
74        Ok(finalize_segments_from_readback(keys, point_indices, cell_starts, cell_counts))
75    }
76}