spatialrust_gpu/aoso_staging/
voxel.rs

1use super::*;
2
3/// Computes voxel keys directly from uploaded interleaved XYZ chunks.
4///
5/// The returned outer vector preserves chunk boundaries and order. Positions
6/// stay in their existing GPU buffers; only the computed keys are read back.
7/// All chunks must use the same `origin` and `inv_leaf` so keys are globally
8/// comparable across chunk boundaries.
9pub fn compute_voxel_keys_aoso_chunks(
10    runtime: &WgpuRuntime,
11    chunks: &[GpuAoSoXyzChunk],
12    origin: [f32; 3],
13    inv_leaf: f32,
14) -> SpatialResult<Vec<Vec<(i64, i64, i64)>>> {
15    if !inv_leaf.is_finite() || inv_leaf <= 0.0 {
16        return Err(SpatialError::InvalidArgument(
17            "inverse voxel leaf size must be finite and positive".to_owned(),
18        ));
19    }
20
21    chunks
22        .iter()
23        .map(|chunk| compute_voxel_keys_aoso_chunk(runtime, chunk, origin, inv_leaf))
24        .collect()
25}
26
27/// Runs global voxel centroid downsampling from uploaded AoSoA chunks.
28///
29/// Chunk buffers are concatenated with GPU-to-GPU copies. Key generation,
30/// sorting, segmentation, and centroid reduction then operate on that combined
31/// GPU buffer, so voxels spanning chunk boundaries are merged correctly. Only
32/// the final centroids are read back.
33pub fn downsample_voxel_centroid_aoso_chunks(
34    runtime: &WgpuRuntime,
35    chunks: &[GpuAoSoXyzChunk],
36    origin: [f32; 3],
37    inv_leaf: f32,
38) -> SpatialResult<AoSoAVoxelCentroidResult> {
39    if !inv_leaf.is_finite() || inv_leaf <= 0.0 {
40        return Err(SpatialError::InvalidArgument(
41            "inverse voxel leaf size must be finite and positive".to_owned(),
42        ));
43    }
44    let total_points = chunks.iter().try_fold(0usize, |total, chunk| {
45        total
46            .checked_add(chunk.point_count)
47            .ok_or_else(|| SpatialError::InvalidArgument("AoSoA point count overflow".to_owned()))
48    })?;
49    if total_points == 0 {
50        return Ok(AoSoAVoxelCentroidResult {
51            out_x: Vec::new(),
52            out_y: Vec::new(),
53            out_z: Vec::new(),
54            segments: build_voxel_segments_gpu_from_keys_buffer(
55                runtime,
56                &empty_storage_buffer(runtime),
57                0,
58                1,
59            )?,
60            positions: empty_aoso_buffer(runtime),
61        });
62    }
63    let point_count = u32::try_from(total_points).map_err(|_| {
64        SpatialError::InvalidArgument("AoSoA chunks exceed the GPU point limit".to_owned())
65    })?;
66    let positions = combine_aoso_chunks(runtime, chunks, total_points);
67    let keys =
68        dispatch_voxel_keys_aoso(runtime, positions.buffer(), point_count, origin, inv_leaf)?;
69    let segments = build_voxel_segments_gpu_from_keys_buffer(
70        runtime,
71        &keys,
72        point_count,
73        point_count.next_power_of_two(),
74    )?;
75    let (out_x, out_y, out_z) =
76        reduce_voxel_centroids_aoso(runtime, positions.buffer(), &segments)?;
77    Ok(AoSoAVoxelCentroidResult { out_x, out_y, out_z, segments, positions })
78}
79
80fn empty_aoso_buffer(runtime: &WgpuRuntime) -> GpuAoSoXyzBuffer {
81    GpuAoSoXyzBuffer {
82        buffer: runtime.device().create_buffer(&wgpu::BufferDescriptor {
83            label: Some("aoso-empty-positions"),
84            size: 4,
85            usage: wgpu::BufferUsages::STORAGE
86                | wgpu::BufferUsages::COPY_DST
87                | wgpu::BufferUsages::COPY_SRC,
88            mapped_at_creation: false,
89        }),
90        point_count: 0,
91        device_key: runtime_device_key(runtime),
92    }
93}
94
95pub(super) fn empty_storage_buffer(runtime: &WgpuRuntime) -> Buffer {
96    runtime.device().create_buffer(&wgpu::BufferDescriptor {
97        label: Some("aoso-empty-storage"),
98        size: 4,
99        usage: wgpu::BufferUsages::STORAGE,
100        mapped_at_creation: false,
101    })
102}
103
104fn combine_aoso_chunks(
105    runtime: &WgpuRuntime,
106    chunks: &[GpuAoSoXyzChunk],
107    total_points: usize,
108) -> GpuAoSoXyzBuffer {
109    let device = runtime.device();
110    let combined = device.create_buffer(&wgpu::BufferDescriptor {
111        label: Some("aoso-xyz-combined"),
112        size: (total_points * 3 * std::mem::size_of::<f32>()) as u64,
113        usage: wgpu::BufferUsages::STORAGE
114            | wgpu::BufferUsages::COPY_DST
115            | wgpu::BufferUsages::COPY_SRC,
116        mapped_at_creation: false,
117    });
118    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
119        label: Some("aoso-combine-encoder"),
120    });
121    let mut offset = 0;
122    for chunk in chunks {
123        let bytes = (chunk.point_count * 3 * std::mem::size_of::<f32>()) as u64;
124        if bytes > 0 {
125            encoder.copy_buffer_to_buffer(&chunk.buffer, 0, &combined, offset, bytes);
126            offset += bytes;
127        }
128    }
129    runtime.queue().submit(Some(encoder.finish()));
130    GpuAoSoXyzBuffer {
131        buffer: combined,
132        point_count: total_points,
133        device_key: runtime_device_key(runtime),
134    }
135}
136
137pub(super) fn dispatch_voxel_keys_aoso(
138    runtime: &WgpuRuntime,
139    positions: &Buffer,
140    point_count: u32,
141    origin: [f32; 3],
142    inv_leaf: f32,
143) -> SpatialResult<Buffer> {
144    let device = runtime.device();
145    let uniform = VoxelKeyUniform {
146        origin: [origin[0], origin[1], origin[2], 0.0],
147        inv_leaf,
148        point_count,
149        _pad0: 0,
150        _pad1: 0,
151    };
152    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
153        label: Some("voxel-key-aoso-uniform"),
154        contents: bytemuck::bytes_of(&uniform),
155        usage: wgpu::BufferUsages::UNIFORM,
156    });
157    let output = device.create_buffer(&wgpu::BufferDescriptor {
158        label: Some("voxel-key-aoso-output"),
159        size: u64::from(point_count) * std::mem::size_of::<VoxelKeyOutput>() as u64,
160        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
161        mapped_at_creation: false,
162    });
163    let pipelines = runtime.pipelines();
164    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
165        label: Some("voxel-key-aoso-bind-group"),
166        layout: &pipelines.voxel_keys_aoso.bind_group_layout,
167        entries: &[
168            wgpu::BindGroupEntry { binding: 0, resource: uniform_buffer.as_entire_binding() },
169            wgpu::BindGroupEntry { binding: 1, resource: positions.as_entire_binding() },
170            wgpu::BindGroupEntry { binding: 2, resource: output.as_entire_binding() },
171        ],
172    });
173    let mut encoder = device
174        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("voxel-key-aoso") });
175    {
176        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
177            label: Some("voxel-key-aoso-pass"),
178            timestamp_writes: None,
179        });
180        pass.set_pipeline(&pipelines.voxel_keys_aoso.pipeline);
181        pass.set_bind_group(0, &bind_group, &[]);
182        pass.dispatch_workgroups(point_count.div_ceil(WORKGROUP_SIZE), 1, 1);
183    }
184    runtime.queue().submit(Some(encoder.finish()));
185    Ok(output)
186}
187
188fn reduce_voxel_centroids_aoso(
189    runtime: &WgpuRuntime,
190    positions: &Buffer,
191    segments: &GpuVoxelSegments,
192) -> SpatialResult<(Vec<f32>, Vec<f32>, Vec<f32>)> {
193    let cell_count = segments.cell_count();
194    if cell_count == 0 {
195        return Ok((Vec::new(), Vec::new(), Vec::new()));
196    }
197    let device = runtime.device();
198    let output_len = u64::from(cell_count) * std::mem::size_of::<[f32; 4]>() as u64;
199    let output = device.create_buffer(&wgpu::BufferDescriptor {
200        label: Some("voxel-reduce-aoso-output"),
201        size: output_len,
202        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
203        mapped_at_creation: false,
204    });
205    let staging = device.create_buffer(&wgpu::BufferDescriptor {
206        label: Some("voxel-reduce-aoso-staging"),
207        size: output_len,
208        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
209        mapped_at_creation: false,
210    });
211    let uniform =
212        VoxelReduceUniform { cell_count, point_count: segments.point_count(), _pad0: 0, _pad1: 0 };
213    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
214        label: Some("voxel-reduce-aoso-uniform"),
215        contents: bytemuck::bytes_of(&uniform),
216        usage: wgpu::BufferUsages::UNIFORM,
217    });
218    let pipelines = runtime.pipelines();
219    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
220        label: Some("voxel-reduce-aoso-bind-group"),
221        layout: &pipelines.voxel_reduce_aoso.bind_group_layout,
222        entries: &[
223            wgpu::BindGroupEntry { binding: 0, resource: uniform_buffer.as_entire_binding() },
224            wgpu::BindGroupEntry {
225                binding: 1,
226                resource: segments.point_indices_buffer().as_entire_binding(),
227            },
228            wgpu::BindGroupEntry {
229                binding: 2,
230                resource: segments.cell_starts_buffer().as_entire_binding(),
231            },
232            wgpu::BindGroupEntry { binding: 3, resource: positions.as_entire_binding() },
233            wgpu::BindGroupEntry { binding: 4, resource: output.as_entire_binding() },
234        ],
235    });
236    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
237        label: Some("voxel-reduce-aoso-encoder"),
238    });
239    {
240        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
241            label: Some("voxel-reduce-aoso-pass"),
242            timestamp_writes: None,
243        });
244        pass.set_pipeline(&pipelines.voxel_reduce_aoso.pipeline);
245        pass.set_bind_group(0, &bind_group, &[]);
246        pass.dispatch_workgroups(cell_count.div_ceil(WORKGROUP_SIZE), 1, 1);
247    }
248    encoder.copy_buffer_to_buffer(&output, 0, &staging, 0, output_len);
249    runtime.queue().submit(Some(encoder.finish()));
250    let slice = staging.slice(..);
251    let (sender, receiver) = std::sync::mpsc::channel();
252    slice.map_async(wgpu::MapMode::Read, move |result| {
253        let _ = sender.send(result);
254    });
255    device.poll(wgpu::Maintain::Wait);
256    receiver
257        .recv()
258        .map_err(|_| SpatialError::InvalidArgument("failed to receive wgpu map result".to_owned()))?
259        .map_err(|error| {
260            SpatialError::InvalidArgument(format!("failed to map wgpu buffer: {error}"))
261        })?;
262    let data = slice.get_mapped_range();
263    let centroids: &[[f32; 4]] = bytemuck::cast_slice(&data);
264    let mut out_x = Vec::with_capacity(centroids.len());
265    let mut out_y = Vec::with_capacity(centroids.len());
266    let mut out_z = Vec::with_capacity(centroids.len());
267    for centroid in centroids {
268        out_x.push(centroid[0]);
269        out_y.push(centroid[1]);
270        out_z.push(centroid[2]);
271    }
272    drop(data);
273    staging.unmap();
274    Ok((out_x, out_y, out_z))
275}
276
277fn compute_voxel_keys_aoso_chunk(
278    runtime: &WgpuRuntime,
279    chunk: &GpuAoSoXyzChunk,
280    origin: [f32; 3],
281    inv_leaf: f32,
282) -> SpatialResult<Vec<(i64, i64, i64)>> {
283    if chunk.point_count == 0 {
284        return Ok(Vec::new());
285    }
286    let point_count = u32::try_from(chunk.point_count).map_err(|_| {
287        SpatialError::InvalidArgument("AoSoA chunk exceeds the GPU point limit".to_owned())
288    })?;
289    let device = runtime.device();
290    let queue = runtime.queue();
291    let uniform = VoxelKeyUniform {
292        origin: [origin[0], origin[1], origin[2], 0.0],
293        inv_leaf,
294        point_count,
295        _pad0: 0,
296        _pad1: 0,
297    };
298    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
299        label: Some("voxel-key-aoso-uniform"),
300        contents: bytemuck::bytes_of(&uniform),
301        usage: wgpu::BufferUsages::UNIFORM,
302    });
303    let output_len = chunk.point_count * std::mem::size_of::<VoxelKeyOutput>();
304    let output_buffer = device.create_buffer(&wgpu::BufferDescriptor {
305        label: Some("voxel-key-aoso-output"),
306        size: output_len as u64,
307        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
308        mapped_at_creation: false,
309    });
310    let staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
311        label: Some("voxel-key-aoso-staging"),
312        size: output_len as u64,
313        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
314        mapped_at_creation: false,
315    });
316    let pipelines = runtime.pipelines();
317    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
318        label: Some("voxel-key-aoso-bind-group"),
319        layout: &pipelines.voxel_keys_aoso.bind_group_layout,
320        entries: &[
321            wgpu::BindGroupEntry { binding: 0, resource: uniform_buffer.as_entire_binding() },
322            wgpu::BindGroupEntry { binding: 1, resource: chunk.buffer.as_entire_binding() },
323            wgpu::BindGroupEntry { binding: 2, resource: output_buffer.as_entire_binding() },
324        ],
325    });
326    let mut encoder = device
327        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("voxel-key-aoso") });
328    {
329        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
330            label: Some("voxel-key-aoso-pass"),
331            timestamp_writes: None,
332        });
333        pass.set_pipeline(&pipelines.voxel_keys_aoso.pipeline);
334        pass.set_bind_group(0, &bind_group, &[]);
335        pass.dispatch_workgroups(point_count.div_ceil(WORKGROUP_SIZE), 1, 1);
336    }
337    encoder.copy_buffer_to_buffer(&output_buffer, 0, &staging_buffer, 0, output_len as u64);
338    queue.submit(Some(encoder.finish()));
339
340    let slice = staging_buffer.slice(..);
341    let (sender, receiver) = std::sync::mpsc::channel();
342    slice.map_async(wgpu::MapMode::Read, move |result| {
343        let _ = sender.send(result);
344    });
345    device.poll(wgpu::Maintain::Wait);
346    receiver
347        .recv()
348        .map_err(|_| SpatialError::InvalidArgument("failed to receive wgpu map result".to_owned()))?
349        .map_err(|error| {
350            SpatialError::InvalidArgument(format!("failed to map wgpu buffer: {error}"))
351        })?;
352    let data = slice.get_mapped_range();
353    let keys = bytemuck::cast_slice::<u8, VoxelKeyOutput>(&data)
354        .iter()
355        .map(|key| (i64::from(key.ix), i64::from(key.iy), i64::from(key.iz)))
356        .collect();
357    drop(data);
358    staging_buffer.unmap();
359    Ok(keys)
360}
spatialrust_gpu/aoso_staging/voxel.rs

spatialrust_gpu/aoso_staging/
voxel.rs
