spatialrust_gpu/
gpu_frame.rs

1//! GPU-resident spatial frame ownership and chained execution.
2
3use spatialrust_core::{
4    PointSchema, SpatialError, SpatialResult, SpatialTensor, TransferDirection, TransferStats,
5};
6
7use crate::aoso_staging::runtime_device_key;
8use crate::{
9    build_radius_grid_aoso_gpu, downsample_voxel_centroid_aoso_chunks,
10    estimate_normals_radius_grid_aoso_gpu, reduce_voxel_attributes_aoso_chunks,
11    upload_spatial_tensor_xyz_chunks, AoSoAAttributeAggregation, AoSoAAttributeReduction,
12    AoSoAVoxelCentroidResult, GpuAoSoAttributeChunk, GpuAoSoNormals, GpuAoSoRadiusGrid,
13    GpuAoSoXyzBuffer, GpuVoxelSegments, WgpuRuntime,
14};
15
16/// GPU frame capabilities available to downstream algorithms.
17#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18pub enum GpuFrameCapability {
19    /// Interleaved XYZ positions.
20    Positions,
21    /// Per-point normals and curvature.
22    Normals,
23    /// Interleaved point attributes.
24    Attributes,
25    /// Voxel partition metadata.
26    VoxelSegments,
27    /// Sparse uniform radius grid.
28    RadiusGrid,
29}
30
31/// Transfer and logical-stage receipt for one GPU frame.
32#[derive(Clone, Debug, Default, PartialEq, Eq)]
33pub struct GpuExecutionReceipt {
34    transfers: TransferStats,
35    stages: Vec<&'static str>,
36}
37
38impl GpuExecutionReceipt {
39    /// Returns the common transfer accounting for this execution.
40    #[must_use]
41    pub const fn transfer_stats(&self) -> TransferStats {
42        self.transfers
43    }
44
45    /// Returns bytes uploaded from host memory.
46    #[must_use]
47    pub const fn host_to_device_bytes(&self) -> u64 {
48        self.transfers.host_to_device_bytes()
49    }
50
51    /// Returns bytes copied between GPU buffers.
52    #[must_use]
53    pub const fn gpu_to_gpu_bytes(&self) -> u64 {
54        self.transfers.device_to_device_bytes()
55    }
56
57    /// Returns bytes explicitly read back to host memory.
58    #[must_use]
59    pub const fn device_to_host_bytes(&self) -> u64 {
60        self.transfers.device_to_host_bytes()
61    }
62
63    /// Returns the logical GPU stages recorded by the high-level pipeline.
64    #[must_use]
65    pub fn stages(&self) -> &[&'static str] {
66        &self.stages
67    }
68}
69
70/// Owned GPU-resident point frame with explicit schema and device identity.
71pub struct GpuSpatialFrame {
72    schema: PointSchema,
73    device_key: usize,
74    positions: GpuAoSoXyzBuffer,
75    normals: Option<GpuAoSoNormals>,
76    attributes: Vec<GpuAoSoAttributeChunk>,
77    voxel_segments: Option<GpuVoxelSegments>,
78    radius_grid: Option<GpuAoSoRadiusGrid>,
79    receipt: GpuExecutionReceipt,
80}
81
82impl GpuSpatialFrame {
83    /// Creates a frame owning interleaved positions on `runtime`.
84    pub fn new(
85        runtime: &WgpuRuntime,
86        schema: PointSchema,
87        positions: GpuAoSoXyzBuffer,
88    ) -> SpatialResult<Self> {
89        let device_key = runtime_device_key(runtime);
90        if positions.device_key() != device_key {
91            return Err(SpatialError::InvalidArgument(
92                "position buffer belongs to a different runtime device".to_owned(),
93            ));
94        }
95        Ok(Self {
96            schema,
97            device_key,
98            positions,
99            normals: None,
100            attributes: Vec::new(),
101            voxel_segments: None,
102            radius_grid: None,
103            receipt: GpuExecutionReceipt::default(),
104        })
105    }
106
107    /// Returns the source point schema.
108    #[must_use]
109    pub const fn schema(&self) -> &PointSchema {
110        &self.schema
111    }
112
113    /// Returns the number of source points.
114    #[must_use]
115    pub const fn point_count(&self) -> usize {
116        self.positions.point_count()
117    }
118
119    /// Returns retained interleaved positions.
120    #[must_use]
121    pub const fn positions(&self) -> &GpuAoSoXyzBuffer {
122        &self.positions
123    }
124
125    /// Returns retained normals when attached.
126    #[must_use]
127    pub const fn normals(&self) -> Option<&GpuAoSoNormals> {
128        self.normals.as_ref()
129    }
130
131    /// Returns retained voxel segments when attached.
132    #[must_use]
133    pub const fn voxel_segments(&self) -> Option<&GpuVoxelSegments> {
134        self.voxel_segments.as_ref()
135    }
136
137    /// Returns retained radius grid when attached.
138    #[must_use]
139    pub const fn radius_grid(&self) -> Option<&GpuAoSoRadiusGrid> {
140        self.radius_grid.as_ref()
141    }
142
143    /// Returns execution and transfer accounting.
144    #[must_use]
145    pub const fn receipt(&self) -> &GpuExecutionReceipt {
146        &self.receipt
147    }
148
149    /// Returns whether a capability is currently attached.
150    #[must_use]
151    pub fn has_capability(&self, capability: GpuFrameCapability) -> bool {
152        match capability {
153            GpuFrameCapability::Positions => true,
154            GpuFrameCapability::Normals => self.normals.is_some(),
155            GpuFrameCapability::Attributes => !self.attributes.is_empty(),
156            GpuFrameCapability::VoxelSegments => self.voxel_segments.is_some(),
157            GpuFrameCapability::RadiusGrid => self.radius_grid.is_some(),
158        }
159    }
160
161    /// Verifies that `runtime` owns the frame's buffers.
162    pub fn validate_runtime(&self, runtime: &WgpuRuntime) -> SpatialResult<()> {
163        if self.device_key != runtime_device_key(runtime) {
164            return Err(SpatialError::InvalidArgument(
165                "GPU frame belongs to a different runtime device".to_owned(),
166            ));
167        }
168        Ok(())
169    }
170
171    /// Attaches per-point normal output after validating length and device.
172    pub fn attach_normals(
173        &mut self,
174        runtime: &WgpuRuntime,
175        normals: GpuAoSoNormals,
176    ) -> SpatialResult<()> {
177        self.validate_runtime(runtime)?;
178        if normals.device_key() != self.device_key {
179            return Err(SpatialError::InvalidArgument(
180                "normal buffer belongs to a different runtime device".to_owned(),
181            ));
182        }
183        if normals.point_count() != self.point_count() {
184            return Err(SpatialError::BufferLengthMismatch {
185                expected: self.point_count(),
186                found: normals.point_count(),
187            });
188        }
189        if let Some(previous) = self.normals.replace(normals) {
190            previous.recycle(runtime);
191        }
192        Ok(())
193    }
194
195    /// Attaches voxel segments after validating their source point count.
196    pub fn attach_voxel_segments(&mut self, segments: GpuVoxelSegments) -> SpatialResult<()> {
197        if segments.point_count() as usize != self.point_count() {
198            return Err(SpatialError::BufferLengthMismatch {
199                expected: self.point_count(),
200                found: segments.point_count() as usize,
201            });
202        }
203        self.voxel_segments = Some(segments);
204        Ok(())
205    }
206
207    /// Attaches a sparse radius grid after validating its point count.
208    pub fn attach_radius_grid(&mut self, grid: GpuAoSoRadiusGrid) -> SpatialResult<()> {
209        if grid.segments().point_count() as usize != self.point_count() {
210            return Err(SpatialError::BufferLengthMismatch {
211                expected: self.point_count(),
212                found: grid.segments().point_count() as usize,
213            });
214        }
215        self.radius_grid = Some(grid);
216        Ok(())
217    }
218
219    /// Attaches interleaved attribute chunks after validating total length.
220    pub fn attach_attributes(
221        &mut self,
222        runtime: &WgpuRuntime,
223        attributes: Vec<GpuAoSoAttributeChunk>,
224    ) -> SpatialResult<()> {
225        self.validate_runtime(runtime)?;
226        if attributes.iter().any(|attribute| attribute.device_key() != self.device_key) {
227            return Err(SpatialError::InvalidArgument(
228                "attribute buffer belongs to a different runtime device".to_owned(),
229            ));
230        }
231        let found = attributes.iter().map(GpuAoSoAttributeChunk::point_count).sum();
232        if found != self.point_count() {
233            return Err(SpatialError::BufferLengthMismatch { expected: self.point_count(), found });
234        }
235        for previous in std::mem::replace(&mut self.attributes, attributes) {
236            previous.recycle(runtime);
237        }
238        Ok(())
239    }
240
241    /// Rebuilds the sparse radius grid from retained positions.
242    pub fn rebuild_radius_grid(&mut self, runtime: &WgpuRuntime, radius: f32) -> SpatialResult<()> {
243        self.validate_runtime(runtime)?;
244        let grid = build_radius_grid_aoso_gpu(runtime, &self.positions, radius)?;
245        self.radius_grid = Some(grid);
246        self.receipt.stages.push("radius-grid");
247        Ok(())
248    }
249
250    /// Estimates normals using a cached matching grid or rebuilds it first.
251    pub fn estimate_normals(&mut self, runtime: &WgpuRuntime, radius: f32) -> SpatialResult<()> {
252        self.validate_runtime(runtime)?;
253        let matches = self.radius_grid.as_ref().is_some_and(|grid| grid.radius() == radius);
254        if !matches {
255            self.rebuild_radius_grid(runtime, radius)?;
256        }
257        let grid = self.radius_grid.as_ref().ok_or_else(|| {
258            SpatialError::InvalidArgument("radius grid is unavailable".to_owned())
259        })?;
260        let normals = estimate_normals_radius_grid_aoso_gpu(runtime, &self.positions, grid)?;
261        self.attach_normals(runtime, normals)?;
262        self.receipt.stages.push("radius-normals");
263        Ok(())
264    }
265
266    /// Reduces attached interleaved attributes using retained voxel segments.
267    pub fn reduce_attributes(
268        &mut self,
269        runtime: &WgpuRuntime,
270        aggregation: AoSoAAttributeAggregation,
271    ) -> SpatialResult<AoSoAAttributeReduction> {
272        self.validate_runtime(runtime)?;
273        if self.attributes.is_empty() {
274            return Err(SpatialError::InvalidArgument(
275                "GPU frame has no attached attributes".to_owned(),
276            ));
277        }
278        let segments = self.voxel_segments.as_ref().ok_or_else(|| {
279            SpatialError::InvalidArgument("GPU frame has no voxel segments".to_owned())
280        })?;
281        let reduced =
282            reduce_voxel_attributes_aoso_chunks(runtime, &self.attributes, segments, aggregation)?;
283        self.receipt.transfers.record(
284            TransferDirection::DeviceToHost,
285            (reduced.len() * reduced.layout().stride_f32() * std::mem::size_of::<f32>()) as u64,
286        );
287        self.receipt.stages.push("attribute-reduce");
288        Ok(reduced)
289    }
290
291    /// Explicitly reads interleaved positions back to CPU memory.
292    pub fn readback_positions(&mut self, runtime: &WgpuRuntime) -> SpatialResult<Vec<[f32; 3]>> {
293        self.validate_runtime(runtime)?;
294        let byte_len = self.positions.point_count() * 3 * std::mem::size_of::<f32>();
295        if byte_len == 0 {
296            return Ok(Vec::new());
297        }
298        let device = runtime.device();
299        let staging = device.create_buffer(&wgpu::BufferDescriptor {
300            label: Some("gpu-frame-position-readback"),
301            size: byte_len as u64,
302            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
303            mapped_at_creation: false,
304        });
305        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
306            label: Some("gpu-frame-position-readback-encoder"),
307        });
308        encoder.copy_buffer_to_buffer(self.positions.buffer(), 0, &staging, 0, byte_len as u64);
309        runtime.queue().submit(Some(encoder.finish()));
310        let slice = staging.slice(..);
311        let (sender, receiver) = std::sync::mpsc::channel();
312        slice.map_async(wgpu::MapMode::Read, move |result| {
313            let _ = sender.send(result);
314        });
315        device.poll(wgpu::Maintain::Wait);
316        receiver
317            .recv()
318            .map_err(|_| SpatialError::InvalidArgument("failed to receive map result".to_owned()))?
319            .map_err(|error| SpatialError::InvalidArgument(format!("map failed: {error}")))?;
320        let mapped = slice.get_mapped_range();
321        let values: &[f32] = bytemuck::cast_slice(&mapped);
322        let positions = values.chunks_exact(3).map(|p| [p[0], p[1], p[2]]).collect();
323        drop(mapped);
324        staging.unmap();
325        self.receipt.transfers.record(TransferDirection::DeviceToHost, byte_len as u64);
326        Ok(positions)
327    }
328
329    /// Recycles pooled frame buffers after validating the runtime device.
330    pub fn recycle(self, runtime: &WgpuRuntime) -> SpatialResult<()> {
331        self.validate_runtime(runtime)?;
332        self.positions.recycle(runtime);
333        if let Some(normals) = self.normals {
334            normals.recycle(runtime);
335        }
336        for attribute in self.attributes {
337            attribute.recycle(runtime);
338        }
339        Ok(())
340    }
341}
342
343/// Uploads a tensor and chains global voxel partitioning and radius normals.
344pub fn run_aoso_voxel_normal_frame(
345    runtime: &WgpuRuntime,
346    tensor: &SpatialTensor<'_>,
347    origin: [f32; 3],
348    inv_leaf: f32,
349    normal_radius: f32,
350) -> SpatialResult<GpuSpatialFrame> {
351    let chunks = upload_spatial_tensor_xyz_chunks(runtime, tensor)?;
352    let upload_bytes = chunks.iter().map(|chunk| chunk.byte_len()).sum();
353    let voxel = downsample_voxel_centroid_aoso_chunks(runtime, &chunks, origin, inv_leaf)?;
354    for chunk in chunks {
355        chunk.recycle(runtime);
356    }
357    let AoSoAVoxelCentroidResult { segments, positions, .. } = voxel;
358    let grid = build_radius_grid_aoso_gpu(runtime, &positions, normal_radius)?;
359    let normals = estimate_normals_radius_grid_aoso_gpu(runtime, &positions, &grid)?;
360    let mut frame = GpuSpatialFrame::new(runtime, tensor.schema().clone(), positions)?;
361    frame.attach_voxel_segments(segments)?;
362    frame.attach_radius_grid(grid)?;
363    frame.attach_normals(runtime, normals)?;
364    frame.receipt.transfers.record(TransferDirection::HostToDevice, upload_bytes);
365    frame.receipt.transfers.record(TransferDirection::DeviceToDevice, upload_bytes);
366    frame.receipt.stages = vec!["upload", "voxel-segments", "radius-grid", "radius-normals"];
367    Ok(frame)
368}
369
370#[cfg(test)]
371mod tests {
372    use super::{run_aoso_voxel_normal_frame, GpuFrameCapability};
373    use crate::{upload_spatial_tensor_attribute_chunks, AoSoAAttributeAggregation, WgpuRuntime};
374    use spatialrust_core::{
375        AoSoAAttributeLayout, PointCloudBuilder, SpatialTensor, StandardSchemas,
376    };
377
378    #[test]
379    fn chained_frame_owns_capabilities_and_tracks_readback() {
380        let mut builder = PointCloudBuilder::xyz();
381        let mut expected = Vec::new();
382        for row in 0..8 {
383            for column in 0..8 {
384                let point = [column as f32 * 0.1, row as f32 * 0.1, 0.0];
385                builder.push_point(point).unwrap();
386                expected.push(point);
387            }
388        }
389        let cloud = builder.build().unwrap();
390        let tensor = SpatialTensor::new(&cloud, 13).unwrap();
391        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
392
393        let mut frame =
394            run_aoso_voxel_normal_frame(&runtime, &tensor, [0.0; 3], 20.0, 0.25).unwrap();
395        assert_eq!(frame.point_count(), expected.len());
396        assert!(frame.has_capability(GpuFrameCapability::Positions));
397        assert!(frame.has_capability(GpuFrameCapability::Normals));
398        assert!(frame.has_capability(GpuFrameCapability::VoxelSegments));
399        assert!(frame.has_capability(GpuFrameCapability::RadiusGrid));
400        assert!(!frame.has_capability(GpuFrameCapability::Attributes));
401        assert_eq!(frame.normals().unwrap().point_count(), expected.len());
402        assert_eq!(
403            frame.receipt().stages(),
404            &["upload", "voxel-segments", "radius-grid", "radius-normals"]
405        );
406        assert_eq!(frame.receipt().host_to_device_bytes(), (expected.len() * 3 * 4) as u64);
407
408        let actual = frame.readback_positions(&runtime).unwrap();
409        assert_eq!(actual, expected);
410        assert_eq!(frame.receipt().device_to_host_bytes(), (expected.len() * 3 * 4) as u64);
411        assert!(frame.reduce_attributes(&runtime, AoSoAAttributeAggregation::Average).is_err());
412        frame.recycle(&runtime).unwrap();
413    }
414
415    #[test]
416    fn frame_rejects_a_different_runtime() {
417        let mut builder = PointCloudBuilder::xyz();
418        builder.push_point([0.0, 0.0, 0.0]).unwrap();
419        let cloud = builder.build().unwrap();
420        let tensor = SpatialTensor::new(&cloud, 1).unwrap();
421        let runtime = WgpuRuntime::new_headless().expect("first runtime");
422        let other = WgpuRuntime::new_headless().expect("second runtime");
423        let frame = run_aoso_voxel_normal_frame(&runtime, &tensor, [0.0; 3], 1.0, 1.0).unwrap();
424
425        assert!(frame.validate_runtime(&other).is_err());
426        frame.recycle(&runtime).unwrap();
427    }
428
429    #[test]
430    fn frame_native_normals_and_attributes_replace_safely() {
431        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyzinormal());
432        for point in [
433            [0.1, 0.0, 0.0, 2.0, 0.0, 0.0, 1.0],
434            [1.1, 0.0, 0.0, 4.0, 0.0, 0.0, 1.0],
435            [1.3, 0.0, 0.0, 8.0, 0.0, 0.0, 1.0],
436            [2.1, 0.0, 0.0, 6.0, 0.0, 0.0, 1.0],
437        ] {
438            builder.push_point(point).unwrap();
439        }
440        let cloud = builder.build().unwrap();
441        let tensor = SpatialTensor::new(&cloud, 2).unwrap();
442        let runtime = WgpuRuntime::new_headless().expect("wgpu runtime");
443        let attributes = upload_spatial_tensor_attribute_chunks(
444            &runtime,
445            &tensor,
446            AoSoAAttributeLayout::XYZ_INTENSITY_NORMALS,
447        )
448        .unwrap();
449        let mut frame = run_aoso_voxel_normal_frame(&runtime, &tensor, [0.0; 3], 1.0, 0.5).unwrap();
450        frame.attach_attributes(&runtime, attributes).unwrap();
451
452        let reduced =
453            frame.reduce_attributes(&runtime, AoSoAAttributeAggregation::Average).unwrap();
454        assert_eq!(reduced.len(), 3);
455        assert_eq!(reduced.as_slice()[10], 6.0);
456        assert!(frame.has_capability(GpuFrameCapability::Attributes));
457
458        frame.estimate_normals(&runtime, 0.5).unwrap();
459        assert_eq!(frame.radius_grid().unwrap().radius(), 0.5);
460        frame.estimate_normals(&runtime, 1.0).unwrap();
461        assert_eq!(frame.radius_grid().unwrap().radius(), 1.0);
462        assert_eq!(frame.normals().unwrap().point_count(), 4);
463        assert!(frame.receipt().stages().ends_with(&[
464            "radius-normals",
465            "radius-grid",
466            "radius-normals",
467        ]));
468        frame.recycle(&runtime).unwrap();
469    }
470}
spatialrust_gpu/gpu_frame.rs

spatialrust_gpu/
gpu_frame.rs
