spatialrust_gpu/kernels/
icp_trim.rs

1//! Private resident stable distance/index selection for ICP updates.
2use super::*;
3
4pub(super) const SHADER: &str = include_str!("../shaders/icp_trim.wgsl");
5
6/// Arithmetic used for retained ICP moments; transform/search remain f32.
7#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
8pub enum GpuIcpTrimPrecision {
9    /// Canonical f32 normalized coordinates/products with compensated sums.
10    /// Lower arithmetic cost, but coordinate/product rounding can bias moments.
11    CompensatedF32,
12    /// Two f32 words through offsets, normalization, mean division and products.
13    /// Carries quotient/product residuals to reduce moment bias; hardware f32
14    /// transforms/search and their gate/tie precision remain unchanged.
15    #[default]
16    DoubleSingle,
17}
18
19#[repr(C)]
20#[derive(Clone, Copy, Pod, Zeroable)]
21pub(super) struct TrimParams {
22    count: u32,
23    groups: u32,
24    step: u32,
25    mantissa_lo: u32,
26    mantissa_hi: u32,
27    shift: u32,
28    bound: f32,
29    targets: u32,
30}
31
32pub(super) fn plan_trim(count: u64, limits: &wgpu::Limits) -> SpatialResult<u32> {
33    let groups = plan(count, limits)?;
34    let bound = limits.max_buffer_size.min(u64::from(limits.max_storage_buffer_binding_size));
35    if u64::from(groups) * 1024 > bound
36        || limits.max_storage_buffers_per_shader_stage < 7
37        || limits.max_bindings_per_bind_group < 8
38        || limits.max_compute_workgroup_storage_size < 2048
39    {
40        return Err(invalid());
41    }
42    Ok(groups)
43}
44
45impl TrimParams {
46    fn new(count: u32, groups: u32, fraction: f64, bound: f32, targets: u32) -> Self {
47        let bits = fraction.to_bits();
48        let exponent = ((bits >> 52) & 2047) as u32;
49        let mantissa = (bits & ((1u64 << 52) - 1)) | if exponent == 0 { 0 } else { 1u64 << 52 };
50        Self {
51            count,
52            groups,
53            step: 0,
54            mantissa_lo: mantissa as u32,
55            mantissa_hi: (mantissa >> 32) as u32,
56            shift: if exponent == 0 { 1074 } else { 1075 - exponent },
57            bound,
58            targets,
59        }
60    }
61}
62
63pub(crate) struct TrimPipelines {
64    reduction: std::sync::OnceLock<IcpReductionPipelines>,
65    precise: std::sync::OnceLock<IcpReductionPipelines>,
66    layout: wgpu::BindGroupLayout,
67    histogram: wgpu::ComputePipeline,
68    select: wgpu::ComputePipeline,
69    wide_select: std::sync::OnceLock<wgpu::ComputePipeline>,
70}
71
72fn use_wide_selector(groups: u32, limits: &wgpu::Limits) -> bool {
73    groups >= 64
74        && limits.max_compute_invocations_per_workgroup >= 256
75        && limits.max_compute_workgroup_size_x >= 256
76}
77
78impl TrimPipelines {
79    pub(super) fn new(device: &wgpu::Device) -> Self {
80        let entries: Vec<_> = (0..4)
81            .map(|binding| wgpu::BindGroupLayoutEntry {
82                binding,
83                visibility: wgpu::ShaderStages::COMPUTE,
84                ty: wgpu::BindingType::Buffer {
85                    ty: if binding == 0 {
86                        wgpu::BufferBindingType::Uniform
87                    } else {
88                        wgpu::BufferBindingType::Storage { read_only: binding == 1 }
89                    },
90                    has_dynamic_offset: false,
91                    min_binding_size: None,
92                },
93                count: None,
94            })
95            .collect();
96        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
97            label: Some("icp-trim-layout"),
98            entries: &entries,
99        });
100        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
101            label: Some("icp-trim-layout"),
102            bind_group_layouts: &[&layout],
103            push_constant_ranges: &[],
104        });
105        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
106            label: Some("icp-trim-shader"),
107            source: wgpu::ShaderSource::Wgsl(SHADER.into()),
108        });
109        let [histogram, select] = ["histogram_prefix", "select_digit"].map(|entry| {
110            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
111                label: Some("icp-trim-pipeline"),
112                layout: Some(&pipeline_layout),
113                module: &shader,
114                entry_point: Some(entry),
115                compilation_options: Default::default(),
116                cache: None,
117            })
118        });
119        Self {
120            reduction: std::sync::OnceLock::new(),
121            precise: std::sync::OnceLock::new(),
122            layout,
123            histogram,
124            select,
125            wide_select: std::sync::OnceLock::new(),
126        }
127    }
128
129    fn selector(&self, device: &wgpu::Device, groups: u32) -> &wgpu::ComputePipeline {
130        if !use_wide_selector(groups, &device.limits()) {
131            return &self.select;
132        }
133        self.wide_select.get_or_init(|| {
134            let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
135                label: Some("icp-trim-wide-selector"),
136                source: wgpu::ShaderSource::Wgsl(SHADER.into()),
137            });
138            let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
139                label: Some("icp-trim-wide-selector-layout"),
140                bind_group_layouts: &[&self.layout],
141                push_constant_ranges: &[],
142            });
143            let constants = std::collections::HashMap::from([
144                ("select_lanes".to_owned(), 256.0),
145                ("select_tree".to_owned(), 1.0),
146            ]);
147            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
148                label: Some("icp-trim-wide-selector"),
149                layout: Some(&layout),
150                module: &shader,
151                entry_point: Some("select_digit"),
152                compilation_options: wgpu::PipelineCompilationOptions {
153                    constants: &constants,
154                    ..Default::default()
155                },
156                cache: None,
157            })
158        })
159    }
160
161    pub(super) fn reduction(
162        &self,
163        device: &wgpu::Device,
164        precision: GpuIcpTrimPrecision,
165    ) -> &IcpReductionPipelines {
166        match precision {
167            GpuIcpTrimPrecision::CompensatedF32 => self.reduction.get_or_init(|| {
168                IcpReductionPipelines::with_shader(device, &trimmed_shader(), true)
169            }),
170            GpuIcpTrimPrecision::DoubleSingle => self.precise.get_or_init(|| {
171                IcpReductionPipelines::with_shader(device, &precise_shader(), true)
172            }),
173        }
174    }
175
176    pub(super) fn encode(
177        &self,
178        runtime: &WgpuRuntime,
179        encoder: &mut wgpu::CommandEncoder,
180        scratch: &mut ScratchBuffers<'_>,
181        params: TrimParams,
182        matches: &wgpu::Buffer,
183        selection: &wgpu::Buffer,
184    ) {
185        let histograms = scratch.acquire("icp-trim-histograms", u64::from(params.groups) * 1024);
186        let select = self.selector(runtime.device(), params.groups);
187        encoder.clear_buffer(selection, 0, None);
188        // Source indices have known leading zero bytes; do not dispatch those
189        // digits. Distance still uses all four bytes, without a quantized cutoff.
190        let first_index_step = 4 + params.count.saturating_sub(1).leading_zeros() / 8;
191        for step in (0..4).chain(first_index_step..8) {
192            let uniform = runtime.device().create_buffer_init(&wgpu::util::BufferInitDescriptor {
193                label: Some("icp-trim-params"),
194                contents: bytemuck::bytes_of(&TrimParams { step, ..params }),
195                usage: wgpu::BufferUsages::UNIFORM,
196            });
197            let entries: Vec<_> = [&uniform, matches, &histograms, selection]
198                .iter()
199                .enumerate()
200                .map(|(binding, buffer)| wgpu::BindGroupEntry {
201                    binding: binding as u32,
202                    resource: buffer.as_entire_binding(),
203                })
204                .collect();
205            let group = runtime.device().create_bind_group(&wgpu::BindGroupDescriptor {
206                label: Some("icp-trim-bindings"),
207                layout: &self.layout,
208                entries: &entries,
209            });
210            dispatch(encoder, &self.histogram, &group, params.groups, 1);
211            dispatch(encoder, select, &group, 1, 1);
212        }
213    }
214}
215
216pub(super) fn encode_selection(
217    runtime: &WgpuRuntime,
218    encoder: &mut wgpu::CommandEncoder,
219    scratch: &mut ScratchBuffers<'_>,
220    params: &Params,
221    fraction: f64,
222    matches: &wgpu::Buffer,
223    selection: &wgpu::Buffer,
224) {
225    runtime.icp_trim_pipelines.get().unwrap().encode(
226        runtime,
227        encoder,
228        scratch,
229        TrimParams::new(params.count, params.groups, fraction, params.bound, params.targets),
230        matches,
231        selection,
232    );
233}
234
235fn replace_once(shader: &mut String, from: &str, to: &str) {
236    assert_eq!(shader.matches(from).count(), 1, "ICP trim shader anchor changed: {from}");
237    *shader = shader.replacen(from, to, 1);
238}
239
240// Keep the untrimmed numerical kernels canonical; change only selection/count
241// semantics. Assertions make later edits to the shared shader fail explicitly.
242fn trimmed_shader() -> String {
243    let mut shader = include_str!("../shaders/icp_reduction.wgsl").replace("\r\n", "\n");
244    shader.push_str("\nstruct Selection { gated:u32, retained:u32, rank:u32, distance:u32, source_index:u32, unused0:u32, unused1:u32, unused2:u32 }\n@group(0) @binding(7) var<storage,read> selection: Selection;\nfn retained_match(i:u32)->bool { let distance=bitcast<u32>(matches[i].distance)&0x7fffffffu; return selection.retained>0u && (distance<selection.distance || (distance==selection.distance && i<=selection.source_index)); }\n");
245    replace_once(&mut shader, "if !valid_match(i) { continue; }\n        let p=", "if !valid_match(i) { continue; }\n        if term!=15u && !retained_match(i) { continue; }\n        let p=");
246    replace_once(
247        &mut shader,
248        "if !valid_match(i) { continue; }\n            let s=",
249        "if !valid_match(i) || !retained_match(i) { continue; }\n            let s=",
250    );
251    replace_once(&mut shader, "result.count=count;", "result.count=count; result.reserved0=selection.gated;\n    if count!=selection.retained { atomicOr(&result.status,2u); }");
252    replace_once(
253        &mut shader,
254        "if count>0u {\n                if term<6u",
255        "if count>0u || term==15u {\n                if term<6u",
256    );
257    replace_once(&mut shader, "if atomicLoad(&result.status)!=0u || result.count==0u {", "if result.count==0u { result.source_anchor=vec4<f32>(0.0); result.target_anchor=vec4<f32>(0.0); }\n    if atomicLoad(&result.status)!=0u || result.reserved0==0u {");
258    replace_once(
259        &mut shader,
260        "result.count=0u; result.source_anchor=",
261        "result.count=0u; result.reserved0=0u; result.source_anchor=",
262    );
263    shader
264}
265
266fn precise_shader() -> String {
267    let mut shader = trimmed_shader();
268    shader.push_str(include_str!("../shaders/icp_trim_moments.wgsl"));
269    replace_once(&mut shader, "var value=0.0;", "var value=vec2<f32>(0.0);");
270    replace_once(
271        &mut shader,
272        "value=((p-anchor)/params.scale)[term];",
273        "value=ds_offset(p[term],anchor[term],params.scale);",
274    );
275    replace_once(
276        &mut shader,
277        "value=((q-params.target_anchor.xyz)/params.scale)[term-3u];",
278        "value=ds_offset(q[term-3u],params.target_anchor[term-3u],params.scale);",
279    );
280    replace_once(&mut shader, "value=dot3(d,d);", "value=vec2<f32>(dot3(d,d),0.0);");
281    replace_once(&mut shader, "if !finite(value) {", "if !finite(value.x) || !finite(value.y) {");
282    replace_once(
283        &mut shader,
284        "total=add(total,value);",
285        "if term==15u { total=add(total,value.x); } else { total=precise_merge(total,value); }",
286    );
287    replace_once(&mut shader, "result.values[term]=merge(vec2<f32>(total.x/f32(count),0.0),vec2<f32>(total.y/f32(count),0.0));", "result.values[term]=ds_divide(total,f32(count));");
288    replace_once(
289        &mut shader,
290        "let mean_s=result.values[row].x+result.values[row].y;",
291        "let mean_s=result.values[row];",
292    );
293    replace_once(
294        &mut shader,
295        "let mean_t=result.values[col+3u].x+result.values[col+3u].y;",
296        "let mean_t=result.values[col+3u];",
297    );
298    replace_once(
299        &mut shader,
300        "let s=(queries[i].xyz-result.source_anchor.xyz)/params.scale;",
301        "let s=ds_offset(queries[i][row],result.source_anchor[row],params.scale);",
302    );
303    replace_once(&mut shader, "let t=(target_points[matches[i].index].xyz-result.target_anchor.xyz)/params.scale;", "let t=ds_offset(target_points[matches[i].index][col],result.target_anchor[col],params.scale);");
304    replace_once(
305        &mut shader,
306        "total=add(total,(s[row]-mean_s)*(t[col]-mean_t));",
307        "total=precise_merge(total,ds_product(precise_merge(s,-mean_s),precise_merge(t,-mean_t)));",
308    );
309    replace_once(&mut shader, "if lane<stride { sums[lane]=merge(sums[lane],sums[lane+stride]); counts[lane]+=counts[lane+stride]; }", "if lane<stride { if term==15u { sums[lane]=merge(sums[lane],sums[lane+stride]); } else { sums[lane]=precise_merge(sums[lane],sums[lane+stride]); } counts[lane]+=counts[lane+stride]; }");
310    let group_merge =
311        "for (var g=0u;g<params.groups;g+=1u) { total=merge(total,partials[g].values[term]); }";
312    assert_eq!(shader.matches(group_merge).count(), 2);
313    shader=shader.replacen(group_merge,"for (var g=0u;g<params.groups;g+=1u) { if term==15u { total=merge(total,partials[g].values[term]); } else { total=precise_merge(total,partials[g].values[term]); } }",1);
314    replace_once(
315        &mut shader,
316        "if lane<stride { sums[lane]=merge(sums[lane],sums[lane+stride]); }",
317        "if lane<stride { sums[lane]=precise_merge(sums[lane],sums[lane+stride]); }",
318    );
319    // The remaining group merge belongs to covariance finalization.
320    replace_once(&mut shader, "for (var g=0u;g<params.groups;g+=1u) { total=merge(total,partials[g].values[term]); }", "for (var g=0u;g<params.groups;g+=1u) { total=precise_merge(total,partials[g].values[term]); }");
321    shader
322}
323
324#[cfg(all(test, not(target_arch = "wasm32")))]
325#[path = "icp_trim_tests.rs"]
326mod tests;
spatialrust_gpu/kernels/icp_trim.rs

spatialrust_gpu/kernels/
icp_trim.rs
