spatialrust_gpu/aoso_staging/
normals.rs

1use super::voxel::{dispatch_voxel_keys_aoso, empty_storage_buffer};
2use super::*;
3
4/// Estimates normals directly from a retained interleaved XYZ GPU buffer.
5///
6/// `neighbors` is a global flattened `point_count * k` index array. Position
7/// data remains GPU-resident; only neighbor indices are uploaded. The returned
8/// normal buffer stays on the GPU until [`GpuAoSoNormals::readback`] is called.
9pub fn estimate_normals_aoso_gpu(
10    runtime: &WgpuRuntime,
11    positions: &GpuAoSoXyzBuffer,
12    neighbors: &[u32],
13    k: u32,
14) -> SpatialResult<GpuAoSoNormals> {
15    let point_count = positions.point_count;
16    if k == 0 || neighbors.len() != point_count * k as usize {
17        return Err(SpatialError::InvalidArgument(format!(
18            "neighbors must have point_count*k = {} entries, got {}",
19            point_count * k as usize,
20            neighbors.len()
21        )));
22    }
23    let device = runtime.device();
24    if point_count == 0 {
25        return Ok(GpuAoSoNormals {
26            buffer: device.create_buffer(&wgpu::BufferDescriptor {
27                label: Some("normals-aoso-empty"),
28                size: 4,
29                usage: wgpu::BufferUsages::STORAGE
30                    | wgpu::BufferUsages::COPY_SRC
31                    | wgpu::BufferUsages::COPY_DST,
32                mapped_at_creation: false,
33            }),
34            point_count: 0,
35            device_key: runtime_device_key(runtime),
36        });
37    }
38    let point_count_u32 = u32::try_from(point_count).map_err(|_| {
39        SpatialError::InvalidArgument("AoSoA positions exceed the GPU point limit".to_owned())
40    })?;
41    let neighbor_buffer = runtime.upload_u32_storage("normals-aoso-neighbors", neighbors)?;
42    let uniform = AoSoANormalsUniform { point_count: point_count_u32, k, _pad0: 0, _pad1: 0 };
43    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
44        label: Some("normals-aoso-uniform"),
45        contents: bytemuck::bytes_of(&uniform),
46        usage: wgpu::BufferUsages::UNIFORM,
47    });
48    let output = device.create_buffer(&wgpu::BufferDescriptor {
49        label: Some("normals-aoso-output"),
50        size: (point_count * std::mem::size_of::<[f32; 4]>()) as u64,
51        usage: wgpu::BufferUsages::STORAGE
52            | wgpu::BufferUsages::COPY_SRC
53            | wgpu::BufferUsages::COPY_DST,
54        mapped_at_creation: false,
55    });
56    let shader_source = crate::kernels::NORMALS_WGSL
57        .replace(
58            "@group(0) @binding(1) var<storage, read> xs: array<f32>;\n@group(0) @binding(2) var<storage, read> ys: array<f32>;\n@group(0) @binding(3) var<storage, read> zs: array<f32>;",
59            "@group(0) @binding(1) var<storage, read> positions: array<f32>;",
60        )
61        .replace("xs[idx]", "positions[idx * 3u]")
62        .replace("ys[idx]", "positions[idx * 3u + 1u]")
63        .replace("zs[idx]", "positions[idx * 3u + 2u]");
64    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
65        label: Some("normals-aoso-shader"),
66        source: wgpu::ShaderSource::Wgsl(shader_source.into()),
67    });
68    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
69        label: Some("normals-aoso-pipeline"),
70        layout: None,
71        module: &module,
72        entry_point: Some("main"),
73        compilation_options: wgpu::PipelineCompilationOptions::default(),
74        cache: None,
75    });
76    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
77        label: Some("normals-aoso-bind-group"),
78        layout: &pipeline.get_bind_group_layout(0),
79        entries: &[
80            wgpu::BindGroupEntry { binding: 0, resource: uniform_buffer.as_entire_binding() },
81            wgpu::BindGroupEntry { binding: 1, resource: positions.buffer.as_entire_binding() },
82            wgpu::BindGroupEntry { binding: 4, resource: neighbor_buffer.as_entire_binding() },
83            wgpu::BindGroupEntry { binding: 5, resource: output.as_entire_binding() },
84        ],
85    });
86    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
87        label: Some("normals-aoso-encoder"),
88    });
89    {
90        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
91            label: Some("normals-aoso-pass"),
92            timestamp_writes: None,
93        });
94        pass.set_pipeline(&pipeline);
95        pass.set_bind_group(0, &bind_group, &[]);
96        pass.dispatch_workgroups(point_count_u32.div_ceil(WORKGROUP_SIZE), 1, 1);
97    }
98    runtime.queue().submit(Some(encoder.finish()));
99    Ok(GpuAoSoNormals { buffer: output, point_count, device_key: runtime_device_key(runtime) })
100}
101
102/// Builds a sparse uniform radius grid directly from retained AoSoA positions.
103///
104/// Cell keys, sorting, and segment compaction stay on the GPU. Grid keys use a
105/// zero origin and `floor(position / radius)`; negative coordinates therefore
106/// remain valid without requiring CPU-side bounds discovery.
107pub fn build_radius_grid_aoso_gpu(
108    runtime: &WgpuRuntime,
109    positions: &GpuAoSoXyzBuffer,
110    radius: f32,
111) -> SpatialResult<GpuAoSoRadiusGrid> {
112    if !radius.is_finite() || radius <= 0.0 {
113        return Err(SpatialError::InvalidArgument(
114            "grid radius must be finite and positive".to_owned(),
115        ));
116    }
117    let point_count = u32::try_from(positions.point_count).map_err(|_| {
118        SpatialError::InvalidArgument("AoSoA positions exceed the GPU point limit".to_owned())
119    })?;
120    let empty = empty_storage_buffer(runtime);
121    let segments = if point_count == 0 {
122        build_voxel_segments_gpu_from_keys_buffer(runtime, &empty, 0, 1)?
123    } else {
124        let keys = dispatch_voxel_keys_aoso(
125            runtime,
126            positions.buffer(),
127            point_count,
128            [0.0; 3],
129            1.0 / radius,
130        )?;
131        build_voxel_segments_gpu_from_keys_buffer(
132            runtime,
133            &keys,
134            point_count,
135            point_count.next_power_of_two(),
136        )?
137    };
138    Ok(GpuAoSoRadiusGrid { radius, segments })
139}
140
141/// Estimates radius normals directly from a GPU-resident sparse AoSoA grid.
142pub fn estimate_normals_radius_grid_aoso_gpu(
143    runtime: &WgpuRuntime,
144    positions: &GpuAoSoXyzBuffer,
145    grid: &GpuAoSoRadiusGrid,
146) -> SpatialResult<GpuAoSoNormals> {
147    let point_count = u32::try_from(positions.point_count).map_err(|_| {
148        SpatialError::InvalidArgument("AoSoA positions exceed the GPU point limit".to_owned())
149    })?;
150    if grid.segments.point_count() != point_count {
151        return Err(SpatialError::BufferLengthMismatch {
152            expected: point_count as usize,
153            found: grid.segments.point_count() as usize,
154        });
155    }
156    if point_count == 0 {
157        return estimate_normals_aoso_gpu(runtime, positions, &[], 1);
158    }
159    let device = runtime.device();
160    let uniform = SparseGridNormalsUniform {
161        origin: [0.0; 4],
162        dims: [grid.segments.cell_count(), 0, 0, point_count],
163        inv_cell: 1.0 / grid.radius,
164        radius_sq: grid.radius * grid.radius,
165        _pad0: 0.0,
166        _pad1: 0.0,
167    };
168    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
169        label: Some("normals-radius-aoso-uniform"),
170        contents: bytemuck::bytes_of(&uniform),
171        usage: wgpu::BufferUsages::UNIFORM,
172    });
173    let output = device.create_buffer(&wgpu::BufferDescriptor {
174        label: Some("normals-radius-aoso-output"),
175        size: u64::from(point_count) * std::mem::size_of::<[f32; 4]>() as u64,
176        usage: wgpu::BufferUsages::STORAGE
177            | wgpu::BufferUsages::COPY_SRC
178            | wgpu::BufferUsages::COPY_DST,
179        mapped_at_creation: false,
180    });
181    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
182        label: Some("normals-radius-aoso-shader"),
183        source: wgpu::ShaderSource::Wgsl(sparse_grid_normals_shader().into()),
184    });
185    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
186        label: Some("normals-radius-aoso-pipeline"),
187        layout: None,
188        module: &module,
189        entry_point: Some("main"),
190        compilation_options: wgpu::PipelineCompilationOptions::default(),
191        cache: None,
192    });
193    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
194        label: Some("normals-radius-aoso-bind-group"),
195        layout: &pipeline.get_bind_group_layout(0),
196        entries: &[
197            wgpu::BindGroupEntry { binding: 0, resource: uniform_buffer.as_entire_binding() },
198            wgpu::BindGroupEntry { binding: 1, resource: positions.buffer.as_entire_binding() },
199            wgpu::BindGroupEntry {
200                binding: 2,
201                resource: grid.segments.keys_buffer().as_entire_binding(),
202            },
203            wgpu::BindGroupEntry {
204                binding: 3,
205                resource: grid.segments.point_indices_buffer().as_entire_binding(),
206            },
207            wgpu::BindGroupEntry {
208                binding: 4,
209                resource: grid.segments.cell_starts_buffer().as_entire_binding(),
210            },
211            wgpu::BindGroupEntry { binding: 5, resource: output.as_entire_binding() },
212        ],
213    });
214    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
215        label: Some("normals-radius-aoso-encoder"),
216    });
217    {
218        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
219            label: Some("normals-radius-aoso-pass"),
220            timestamp_writes: None,
221        });
222        pass.set_pipeline(&pipeline);
223        pass.set_bind_group(0, &bind_group, &[]);
224        pass.dispatch_workgroups(point_count.div_ceil(WORKGROUP_SIZE), 1, 1);
225    }
226    runtime.queue().submit(Some(encoder.finish()));
227    Ok(GpuAoSoNormals {
228        buffer: output,
229        point_count: point_count as usize,
230        device_key: runtime_device_key(runtime),
231    })
232}
233
234fn sparse_grid_normals_shader() -> String {
235    let declarations = "@group(0) @binding(1) var<storage, read> xs: array<f32>;\n@group(0) @binding(2) var<storage, read> ys: array<f32>;\n@group(0) @binding(3) var<storage, read> zs: array<f32>;\n@group(0) @binding(4) var<storage, read> sorted: array<u32>;\n@group(0) @binding(5) var<storage, read> cell_start: array<u32>;\n@group(0) @binding(6) var<storage, read_write> out_normals: array<vec4<f32>>;";
236    let sparse_declarations = r#"struct SparseKey { ix: i32, iy: i32, iz: i32, pad: i32, }
237@group(0) @binding(1) var<storage, read> positions: array<f32>;
238@group(0) @binding(2) var<storage, read> keys: array<SparseKey>;
239@group(0) @binding(3) var<storage, read> sorted: array<u32>;
240@group(0) @binding(4) var<storage, read> cell_start: array<u32>;
241@group(0) @binding(5) var<storage, read_write> out_normals: array<vec4<f32>>;"#;
242    let cell_coord = r#"fn cell_coord(value: f32, origin: f32, inv_cell: f32, dim: u32) -> i32 {
243    let c = i32(floor((value - origin) * inv_cell));
244    return clamp(c, 0, i32(dim) - 1);
245}"#;
246    let lookup = r#"fn cell_coord(value: f32, origin: f32, inv_cell: f32, dim: u32) -> i32 { return i32(floor((value - origin) * inv_cell)); }
247fn key_less(key: SparseKey, x: i32, y: i32, z: i32) -> bool {
248    if (key.ix != x) { return key.ix < x; }
249    if (key.iy != y) { return key.iy < y; }
250    return key.iz < z;
251}
252fn find_cell(x: i32, y: i32, z: i32, count: u32) -> i32 {
253    var low = 0u; var high = count;
254    loop { if (low >= high) { break; } let mid = low + (high-low)/2u; if (key_less(keys[mid],x,y,z)) { low=mid+1u; } else { high=mid; } }
255    if (low < count) { let key=keys[low]; if (key.ix==x && key.iy==y && key.iz==z) { return i32(low); } }
256    return -1;
257}"#;
258    let dense = "let cid = (u32(nz) * dimy + u32(ny)) * dimx + u32(nx);\n                let begin = cell_start[cid];\n                let end = cell_start[cid + 1u];";
259    let sparse = "let found = find_cell(nx, ny, nz, cell_count);\n                if (found < 0) { continue; }\n                let cid = u32(found);\n                let begin = cell_start[cid];\n                let end = select(params.dims.w, cell_start[cid + 1u], cid + 1u < cell_count);";
260    crate::kernels::NORMALS_GRID_WGSL
261        .replace(declarations, sparse_declarations)
262        .replace(cell_coord, lookup)
263        .replace("let px = xs[i];", "let px = positions[i * 3u];")
264        .replace("let py = ys[i];", "let py = positions[i * 3u + 1u];")
265        .replace("let pz = zs[i];", "let pz = positions[i * 3u + 2u];")
266        .replace("let dimx = params.dims.x;\n    let dimy = params.dims.y;\n    let dimz = params.dims.z;", "let cell_count = params.dims.x;")
267        .replace(", dimx);", ", cell_count);")
268        .replace(", dimy);", ", cell_count);")
269        .replace(", dimz);", ", cell_count);")
270        .replace("if (nz < 0 || nz >= i32(dimz)) { continue; }", "")
271        .replace("if (ny < 0 || ny >= i32(dimy)) { continue; }", "")
272        .replace("if (nx < 0 || nx >= i32(dimx)) { continue; }", "")
273        .replace(dense, sparse)
274        .replace("xs[j]", "positions[j * 3u]")
275        .replace("ys[j]", "positions[j * 3u + 1u]")
276        .replace("zs[j]", "positions[j * 3u + 2u]")
277}
spatialrust_gpu/aoso_staging/normals.rs

spatialrust_gpu/aoso_staging/
normals.rs
