spatialrust_gpu/aoso_staging/
attributes.rs

1use super::*;
2
3/// Aggregates uploaded interleaved records using global GPU voxel segments.
4pub fn reduce_voxel_attributes_aoso_chunks(
5    runtime: &WgpuRuntime,
6    chunks: &[GpuAoSoAttributeChunk],
7    segments: &GpuVoxelSegments,
8    aggregation: AoSoAAttributeAggregation,
9) -> SpatialResult<AoSoAAttributeReduction> {
10    let first = chunks.first().ok_or_else(|| {
11        SpatialError::InvalidArgument("attribute chunks must not be empty".to_owned())
12    })?;
13    let layout = first.layout;
14    if chunks.iter().any(|chunk| chunk.layout != layout) {
15        return Err(SpatialError::InvalidArgument(
16            "attribute chunks must use one AoSoA layout".to_owned(),
17        ));
18    }
19    let point_count: usize = chunks.iter().map(|chunk| chunk.point_count).sum();
20    if point_count != segments.point_count() as usize {
21        return Err(SpatialError::BufferLengthMismatch {
22            expected: segments.point_count() as usize,
23            found: point_count,
24        });
25    }
26    let cell_count = segments.cell_count();
27    if cell_count == 0 {
28        return Ok(AoSoAAttributeReduction { data: Vec::new(), point_count: 0, layout });
29    }
30
31    let stride = layout.stride_f32();
32    let device = runtime.device();
33    let combined_len = point_count * stride * std::mem::size_of::<f32>();
34    let combined = device.create_buffer(&wgpu::BufferDescriptor {
35        label: Some("aoso-attributes-combined"),
36        size: combined_len as u64,
37        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
38        mapped_at_creation: false,
39    });
40    let output_count = cell_count as usize * stride;
41    let output_len = output_count * std::mem::size_of::<f32>();
42    let output = device.create_buffer(&wgpu::BufferDescriptor {
43        label: Some("voxel-reduce-attributes-aoso-output"),
44        size: output_len as u64,
45        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
46        mapped_at_creation: false,
47    });
48    let staging = device.create_buffer(&wgpu::BufferDescriptor {
49        label: Some("voxel-reduce-attributes-aoso-staging"),
50        size: output_len as u64,
51        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
52        mapped_at_creation: false,
53    });
54    let uniform = AttributeReduceUniform {
55        cell_count,
56        point_count: segments.point_count(),
57        stride: stride as u32,
58        first_mode: u32::from(aggregation == AoSoAAttributeAggregation::First),
59    };
60    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
61        label: Some("voxel-reduce-attributes-aoso-uniform"),
62        contents: bytemuck::bytes_of(&uniform),
63        usage: wgpu::BufferUsages::UNIFORM,
64    });
65    let pipelines = runtime.pipelines();
66    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
67        label: Some("voxel-reduce-attributes-aoso-bind-group"),
68        layout: &pipelines.voxel_reduce_attributes_aoso.bind_group_layout,
69        entries: &[
70            wgpu::BindGroupEntry { binding: 0, resource: uniform_buffer.as_entire_binding() },
71            wgpu::BindGroupEntry {
72                binding: 1,
73                resource: segments.point_indices_buffer().as_entire_binding(),
74            },
75            wgpu::BindGroupEntry {
76                binding: 2,
77                resource: segments.cell_starts_buffer().as_entire_binding(),
78            },
79            wgpu::BindGroupEntry { binding: 3, resource: combined.as_entire_binding() },
80            wgpu::BindGroupEntry { binding: 4, resource: output.as_entire_binding() },
81        ],
82    });
83    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
84        label: Some("voxel-reduce-attributes-aoso-encoder"),
85    });
86    let mut offset = 0;
87    for chunk in chunks {
88        let bytes = chunk.buffer.size();
89        encoder.copy_buffer_to_buffer(&chunk.buffer, 0, &combined, offset, bytes);
90        offset += bytes;
91    }
92    {
93        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
94            label: Some("voxel-reduce-attributes-aoso-pass"),
95            timestamp_writes: None,
96        });
97        pass.set_pipeline(&pipelines.voxel_reduce_attributes_aoso.pipeline);
98        pass.set_bind_group(0, &bind_group, &[]);
99        pass.dispatch_workgroups((output_count as u32).div_ceil(WORKGROUP_SIZE), 1, 1);
100    }
101    encoder.copy_buffer_to_buffer(&output, 0, &staging, 0, output_len as u64);
102    runtime.queue().submit(Some(encoder.finish()));
103
104    let slice = staging.slice(..);
105    let (sender, receiver) = std::sync::mpsc::channel();
106    slice.map_async(wgpu::MapMode::Read, move |result| {
107        let _ = sender.send(result);
108    });
109    device.poll(wgpu::Maintain::Wait);
110    receiver
111        .recv()
112        .map_err(|_| SpatialError::InvalidArgument("failed to receive wgpu map result".to_owned()))?
113        .map_err(|error| {
114            SpatialError::InvalidArgument(format!("failed to map wgpu buffer: {error}"))
115        })?;
116    let mapped = slice.get_mapped_range();
117    let data = bytemuck::cast_slice::<u8, f32>(&mapped).to_vec();
118    drop(mapped);
119    staging.unmap();
120    Ok(AoSoAAttributeReduction { data, point_count: cell_count as usize, layout })
121}
spatialrust_gpu/aoso_staging/attributes.rs

spatialrust_gpu/aoso_staging/
attributes.rs
