uzor_vision/
gpu.rs

1use std::sync::OnceLock;
2
3use crate::cpu;
4
5#[derive(Clone, Copy, Debug, PartialEq, Eq)]
6pub enum GpuMode {
7    /// All stages on CPU.
8    Cpu,
9    /// GPU denoise + GPU SLIC assign; CPU center update / merge.
10    Hybrid,
11    /// SLIC loop stays on the device; one labels readback.
12    Gpu,
13}
14
15pub struct GpuCtx {
16    pub device: wgpu::Device,
17    pub queue: wgpu::Queue,
18    pub name: String,
19    bilateral: wgpu::ComputePipeline,
20    bilateral_bgl: wgpu::BindGroupLayout,
21    slic: wgpu::ComputePipeline,
22    slic_bgl: wgpu::BindGroupLayout,
23    slic_zero: wgpu::ComputePipeline,
24    slic_zero_bgl: wgpu::BindGroupLayout,
25    slic_accum: wgpu::ComputePipeline,
26    slic_accum_bgl: wgpu::BindGroupLayout,
27    slic_div: wgpu::ComputePipeline,
28    slic_div_bgl: wgpu::BindGroupLayout,
29}
30
31static GPU: OnceLock<Option<GpuCtx>> = OnceLock::new();
32
33pub fn ctx() -> Result<&'static GpuCtx, String> {
34    GPU.get_or_init(init_gpu)
35        .as_ref()
36        .ok_or_else(|| "no gpu adapter".into())
37}
38
39fn init_gpu() -> Option<GpuCtx> {
40    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
41    let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
42        power_preference: wgpu::PowerPreference::HighPerformance,
43        force_fallback_adapter: false,
44        compatible_surface: None,
45    }))
46    .ok()?;
47    let info = adapter.get_info();
48    let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
49        label: Some("uzor-vision"),
50        required_features: wgpu::Features::empty(),
51        required_limits: wgpu::Limits::default(),
52        memory_hints: wgpu::MemoryHints::default(),
53        trace: wgpu::Trace::Off,
54        experimental_features: wgpu::ExperimentalFeatures::default(),
55    }))
56    .ok()?;
57    let (bilateral, bilateral_bgl) = make_compute(
58        &device,
59        "bilateral",
60        BILATERAL_WGSL,
61        &[
62            wgpu::BufferBindingType::Storage { read_only: true },
63            wgpu::BufferBindingType::Storage { read_only: false },
64            wgpu::BufferBindingType::Uniform,
65        ],
66    );
67    let (slic, slic_bgl) = make_compute(
68        &device,
69        "slic",
70        SLIC_WGSL,
71        &[
72            wgpu::BufferBindingType::Storage { read_only: true },
73            wgpu::BufferBindingType::Storage { read_only: false },
74            wgpu::BufferBindingType::Storage { read_only: true },
75            wgpu::BufferBindingType::Uniform,
76        ],
77    );
78    let (slic_zero, slic_zero_bgl) = make_compute(
79        &device,
80        "slic-zero",
81        SLIC_ZERO_WGSL,
82        &[
83            wgpu::BufferBindingType::Storage { read_only: false },
84            wgpu::BufferBindingType::Uniform,
85        ],
86    );
87    let (slic_accum, slic_accum_bgl) = make_compute(
88        &device,
89        "slic-accum",
90        SLIC_ACCUM_WGSL,
91        &[
92            wgpu::BufferBindingType::Storage { read_only: true },
93            wgpu::BufferBindingType::Storage { read_only: true },
94            wgpu::BufferBindingType::Storage { read_only: false },
95            wgpu::BufferBindingType::Uniform,
96        ],
97    );
98    let (slic_div, slic_div_bgl) = make_compute(
99        &device,
100        "slic-div",
101        SLIC_DIV_WGSL,
102        &[
103            wgpu::BufferBindingType::Storage { read_only: false },
104            wgpu::BufferBindingType::Storage { read_only: false },
105            wgpu::BufferBindingType::Uniform,
106        ],
107    );
108    Some(GpuCtx {
109        device,
110        queue,
111        name: info.name,
112        bilateral,
113        bilateral_bgl,
114        slic,
115        slic_bgl,
116        slic_zero,
117        slic_zero_bgl,
118        slic_accum,
119        slic_accum_bgl,
120        slic_div,
121        slic_div_bgl,
122    })
123}
124
125fn make_compute(
126    device: &wgpu::Device,
127    label: &'static str,
128    wgsl: &'static str,
129    bindings: &[wgpu::BufferBindingType],
130) -> (wgpu::ComputePipeline, wgpu::BindGroupLayout) {
131    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
132        label: Some(label),
133        source: wgpu::ShaderSource::Wgsl(wgsl.into()),
134    });
135    let mut entries = Vec::new();
136    for (i, ty) in bindings.iter().enumerate() {
137        entries.push(wgpu::BindGroupLayoutEntry {
138            binding: i as u32,
139            visibility: wgpu::ShaderStages::COMPUTE,
140            ty: wgpu::BindingType::Buffer {
141                ty: *ty,
142                has_dynamic_offset: false,
143                min_binding_size: None,
144            },
145            count: None,
146        });
147    }
148    let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
149        label: Some(label),
150        entries: &entries,
151    });
152    let pl = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
153        label: Some(label),
154        bind_group_layouts: &[Some(&bgl)],
155        immediate_size: 0,
156    });
157    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
158        label: Some(label),
159        layout: Some(&pl),
160        module: &shader,
161        entry_point: Some("main"),
162        compilation_options: wgpu::PipelineCompilationOptions::default(),
163        cache: None,
164    });
165    (pipeline, bgl)
166}
167
168pub fn denoise(rgb: &[u8], w: usize, h: usize, mode: GpuMode) -> (Vec<u8>, String) {
169    let rgb = cpu::median3(rgb, w, h);
170    match mode {
171        GpuMode::Cpu => (cpu::bilateral(&rgb, w, h, 3, 2.0, 16.0), "cpu".into()),
172        GpuMode::Hybrid | GpuMode::Gpu => match gpu_bilateral(&rgb, w, h) {
173            Ok((out, name)) => {
174                let tag = if mode == GpuMode::Gpu { "gpu" } else { "hybrid" };
175                (out, format!("{tag}:{name}"))
176            }
177            Err(e) => {
178                eprintln!("gpu denoise failed ({e}); cpu fallback");
179                (
180                    cpu::bilateral(&rgb, w, h, 3, 2.0, 16.0),
181                    format!("cpu(fallback:{e})"),
182                )
183            }
184        },
185    }
186}
187
188fn gpu_bilateral(rgb: &[u8], w: usize, h: usize) -> Result<(Vec<u8>, String), String> {
189    let g = ctx()?;
190    let n = w * h;
191    let packed = cpu::pack_rgb(rgb, n);
192    let src = buf_storage_dst(&g.device, "b-src", bytemuck::cast_slice(&packed));
193    let dst = buf_rw(&g.device, "b-dst", (n * 4) as u64);
194    let params = BiParams {
195        width: w as u32,
196        height: h as u32,
197        radius: 3,
198        _pad: 0,
199        inv_2s: 0.5 / 4.0,
200        inv_2r: 0.5 / 256.0,
201        _pad2: 0.0,
202        _pad3: 0.0,
203    };
204    let ub = buf_uniform(&g.device, "b-p", bytemuck::bytes_of(&params));
205    dispatch_readback(
206        g,
207        &g.bilateral,
208        &g.bilateral_bgl,
209        &[&src, &dst, &ub],
210        &dst,
211        n * 4,
212        w,
213        h,
214    )
215    .map(|bytes| {
216        let packed_out: &[u32] = bytemuck::cast_slice(&bytes);
217        (cpu::unpack_rgb(packed_out), g.name.clone())
218    })
219}
220
221pub fn gpu_slic_assign(
222    rgb: &[u8],
223    w: usize,
224    h: usize,
225    centers: &[Center],
226    step: f32,
227    compact: f32,
228) -> Result<Vec<u32>, String> {
229    let g = ctx()?;
230    let n = w * h;
231    let packed = cpu::pack_rgb(rgb, n);
232    let src = buf_storage_dst(&g.device, "s-src", bytemuck::cast_slice(&packed));
233    let labels = buf_rw(&g.device, "s-lab", (n * 4) as u64);
234    let ctr = buf_storage_dst(&g.device, "s-ctr", bytemuck::cast_slice(centers));
235    let params = SlicParams {
236        width: w as u32,
237        height: h as u32,
238        k: centers.len() as u32,
239        step,
240        compact,
241        _pad: 0.0,
242        _pad2: 0.0,
243        _pad3: 0.0,
244    };
245    let ub = buf_uniform(&g.device, "s-p", bytemuck::bytes_of(&params));
246    let bytes = dispatch_readback(
247        g,
248        &g.slic,
249        &g.slic_bgl,
250        &[&src, &labels, &ctr, &ub],
251        &labels,
252        n * 4,
253        w,
254        h,
255    )?;
256    Ok(bytemuck::cast_slice(&bytes).to_vec())
257}
258
259/// Assign + accumulate + divide on the device. Labels read back once.
260pub fn gpu_slic_full(
261    rgb: &[u8],
262    w: usize,
263    h: usize,
264    centers: &[Center],
265    step: f32,
266    compact: f32,
267    iters: u32,
268) -> Result<Vec<u32>, String> {
269    let g = ctx()?;
270    let n = w * h;
271    let k = centers.len();
272    let packed = cpu::pack_rgb(rgb, n);
273    let src = buf_storage_dst(&g.device, "sf-src", bytemuck::cast_slice(&packed));
274    let labels = buf_rw(&g.device, "sf-lab", (n * 4) as u64);
275    let ctr = buf_storage_dst(&g.device, "sf-ctr", bytemuck::cast_slice(centers));
276    let acc = buf_rw(&g.device, "sf-acc", (k * 6 * 4) as u64);
277    let params = SlicParams {
278        width: w as u32,
279        height: h as u32,
280        k: k as u32,
281        step,
282        compact,
283        _pad: 0.0,
284        _pad2: 0.0,
285        _pad3: 0.0,
286    };
287    let ub = buf_uniform(&g.device, "sf-p", bytemuck::bytes_of(&params));
288    let bg_assign = bind(
289        &g.device,
290        &g.slic_bgl,
291        &[&src, &labels, &ctr, &ub],
292    );
293    let bg_zero = bind(&g.device, &g.slic_zero_bgl, &[&acc, &ub]);
294    let bg_accum = bind(
295        &g.device,
296        &g.slic_accum_bgl,
297        &[&src, &labels, &acc, &ub],
298    );
299    let bg_div = bind(&g.device, &g.slic_div_bgl, &[&acc, &ctr, &ub]);
300    let mut enc = g.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
301        label: Some("slic-full"),
302    });
303    let gx = ((w as u32) + 15) / 16;
304    let gy = ((h as u32) + 15) / 16;
305    let gz = ((k as u32) * 6 + 63) / 64;
306    let gd = (k as u32 + 63) / 64;
307    for _ in 0..iters.max(1) {
308        record(&mut enc, &g.slic, &bg_assign, gx, gy);
309        record(&mut enc, &g.slic_zero, &bg_zero, gz, 1);
310        record(&mut enc, &g.slic_accum, &bg_accum, gx, gy);
311        record(&mut enc, &g.slic_div, &bg_div, gd, 1);
312    }
313    let staging = g.device.create_buffer(&wgpu::BufferDescriptor {
314        label: Some("sf-read"),
315        size: (n * 4) as u64,
316        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
317        mapped_at_creation: false,
318    });
319    enc.copy_buffer_to_buffer(&labels, 0, &staging, 0, (n * 4) as u64);
320    g.queue.submit(Some(enc.finish()));
321    let slice = staging.slice(..);
322    let (tx, rx) = std::sync::mpsc::channel();
323    slice.map_async(wgpu::MapMode::Read, move |r| {
324        let _ = tx.send(r);
325    });
326    let _ = g.device.poll(wgpu::PollType::Wait {
327        submission_index: None,
328        timeout: None,
329    });
330    rx.recv()
331        .map_err(|e| e.to_string())?
332        .map_err(|e| format!("map {e}"))?;
333    let raw = slice.get_mapped_range();
334    let out: Vec<u32> = bytemuck::cast_slice(&raw).to_vec();
335    drop(raw);
336    staging.unmap();
337    Ok(out)
338}
339
340fn bind(device: &wgpu::Device, bgl: &wgpu::BindGroupLayout, bufs: &[&wgpu::Buffer]) -> wgpu::BindGroup {
341    let entries: Vec<wgpu::BindGroupEntry> = bufs
342        .iter()
343        .enumerate()
344        .map(|(i, b)| wgpu::BindGroupEntry {
345            binding: i as u32,
346            resource: b.as_entire_binding(),
347        })
348        .collect();
349    device.create_bind_group(&wgpu::BindGroupDescriptor {
350        label: Some("slic-bg"),
351        layout: bgl,
352        entries: &entries,
353    })
354}
355
356fn record(
357    enc: &mut wgpu::CommandEncoder,
358    pipe: &wgpu::ComputePipeline,
359    bg: &wgpu::BindGroup,
360    gx: u32,
361    gy: u32,
362) {
363    let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
364        label: Some("slic-pass"),
365        timestamp_writes: None,
366    });
367    pass.set_pipeline(pipe);
368    pass.set_bind_group(0, bg, &[]);
369    pass.dispatch_workgroups(gx, gy, 1);
370}
371
372fn buf_storage_dst(device: &wgpu::Device, label: &str, data: &[u8]) -> wgpu::Buffer {
373    let b = device.create_buffer(&wgpu::BufferDescriptor {
374        label: Some(label),
375        size: data.len() as u64,
376        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
377        mapped_at_creation: true,
378    });
379    b.slice(..).get_mapped_range_mut().copy_from_slice(data);
380    b.unmap();
381    b
382}
383
384fn buf_rw(device: &wgpu::Device, label: &str, size: u64) -> wgpu::Buffer {
385    device.create_buffer(&wgpu::BufferDescriptor {
386        label: Some(label),
387        size,
388        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
389        mapped_at_creation: false,
390    })
391}
392
393fn buf_uniform(device: &wgpu::Device, label: &str, data: &[u8]) -> wgpu::Buffer {
394    let b = device.create_buffer(&wgpu::BufferDescriptor {
395        label: Some(label),
396        size: data.len() as u64,
397        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
398        mapped_at_creation: true,
399    });
400    b.slice(..).get_mapped_range_mut().copy_from_slice(data);
401    b.unmap();
402    b
403}
404
405fn dispatch_readback(
406    g: &GpuCtx,
407    pipe: &wgpu::ComputePipeline,
408    bgl: &wgpu::BindGroupLayout,
409    bufs: &[&wgpu::Buffer],
410    copy_src: &wgpu::Buffer,
411    nbytes: usize,
412    w: usize,
413    h: usize,
414) -> Result<Vec<u8>, String> {
415    let mut entries = Vec::new();
416    for (i, b) in bufs.iter().enumerate() {
417        entries.push(wgpu::BindGroupEntry {
418            binding: i as u32,
419            resource: b.as_entire_binding(),
420        });
421    }
422    let bg = g.device.create_bind_group(&wgpu::BindGroupDescriptor {
423        label: Some("vision-bg"),
424        layout: bgl,
425        entries: &entries,
426    });
427    let staging = g.device.create_buffer(&wgpu::BufferDescriptor {
428        label: Some("vision-read"),
429        size: nbytes as u64,
430        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
431        mapped_at_creation: false,
432    });
433    let mut enc = g.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
434        label: Some("vision-enc"),
435    });
436    {
437        let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
438            label: Some("vision-pass"),
439            timestamp_writes: None,
440        });
441        pass.set_pipeline(pipe);
442        pass.set_bind_group(0, &bg, &[]);
443        pass.dispatch_workgroups(((w as u32) + 15) / 16, ((h as u32) + 15) / 16, 1);
444    }
445    enc.copy_buffer_to_buffer(copy_src, 0, &staging, 0, nbytes as u64);
446    g.queue.submit(Some(enc.finish()));
447    let slice = staging.slice(..);
448    let (tx, rx) = std::sync::mpsc::channel();
449    slice.map_async(wgpu::MapMode::Read, move |r| {
450        let _ = tx.send(r);
451    });
452    let _ = g.device.poll(wgpu::PollType::Wait {
453        submission_index: None,
454        timeout: None,
455    });
456    rx.recv()
457        .map_err(|e| e.to_string())?
458        .map_err(|e| format!("map {e}"))?;
459    let raw = slice.get_mapped_range();
460    let out = raw.to_vec();
461    drop(raw);
462    staging.unmap();
463    Ok(out)
464}
465
466#[repr(C)]
467#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
468struct BiParams {
469    width: u32,
470    height: u32,
471    radius: i32,
472    _pad: i32,
473    inv_2s: f32,
474    inv_2r: f32,
475    _pad2: f32,
476    _pad3: f32,
477}
478
479#[repr(C)]
480#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
481pub struct Center {
482    pub x: f32,
483    pub y: f32,
484    pub r: f32,
485    pub g: f32,
486    pub b: f32,
487    pub _p0: f32,
488    pub _p1: f32,
489    pub _p2: f32,
490}
491
492#[repr(C)]
493#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
494struct SlicParams {
495    width: u32,
496    height: u32,
497    k: u32,
498    step: f32,
499    compact: f32,
500    _pad: f32,
501    _pad2: f32,
502    _pad3: f32,
503}
504
505const BILATERAL_WGSL: &str = r#"
506struct Params {
507    width: u32,
508    height: u32,
509    radius: i32,
510    _pad: i32,
511    inv_2s: f32,
512    inv_2r: f32,
513    _pad2: f32,
514    _pad3: f32,
515}
516@group(0) @binding(0) var<storage, read> src: array<u32>;
517@group(0) @binding(1) var<storage, read_write> dst: array<u32>;
518@group(0) @binding(2) var<uniform> params: Params;
519fn unpack(p: u32) -> vec3<f32> {
520    return vec3<f32>(f32((p >> 16u) & 255u), f32((p >> 8u) & 255u), f32(p & 255u));
521}
522@compute @workgroup_size(16, 16)
523fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
524    let x = gid.x;
525    let y = gid.y;
526    if x >= params.width || y >= params.height { return; }
527    let i = y * params.width + x;
528    let c = unpack(src[i]);
529    var acc = vec3<f32>(0.0);
530    var wsum = 0.0;
531    let r = params.radius;
532    var dy = -r;
533    loop {
534        if dy > r { break; }
535        let ny = i32(y) + dy;
536        if ny >= 0 && ny < i32(params.height) {
537            var dx = -r;
538            loop {
539                if dx > r { break; }
540                let nx = i32(x) + dx;
541                if nx >= 0 && nx < i32(params.width) {
542                    let n = unpack(src[u32(ny) * params.width + u32(nx)]);
543                    let ds2 = f32(dx * dx + dy * dy);
544                    let d = n - c;
545                    let wt = exp(-ds2 * params.inv_2s - dot(d, d) * params.inv_2r);
546                    acc += n * wt;
547                    wsum += wt;
548                }
549                dx = dx + 1;
550            }
551        }
552        dy = dy + 1;
553    }
554    let o = acc / max(wsum, 1e-6);
555    dst[i] = (u32(clamp(o.r, 0.0, 255.0)) << 16u)
556        | (u32(clamp(o.g, 0.0, 255.0)) << 8u)
557        | u32(clamp(o.b, 0.0, 255.0));
558}
559"#;
560
561const SLIC_WGSL: &str = r#"
562struct Center { x: f32, y: f32, r: f32, g: f32, b: f32, _p0: f32, _p1: f32, _p2: f32 }
563struct Params { width: u32, height: u32, k: u32, step: f32, compact: f32, _pad: f32, _pad2: f32, _pad3: f32 }
564@group(0) @binding(0) var<storage, read> src: array<u32>;
565@group(0) @binding(1) var<storage, read_write> labels: array<u32>;
566@group(0) @binding(2) var<storage, read> centers: array<Center>;
567@group(0) @binding(3) var<uniform> params: Params;
568fn unpack(p: u32) -> vec3<f32> {
569    return vec3<f32>(f32((p >> 16u) & 255u), f32((p >> 8u) & 255u), f32(p & 255u));
570}
571@compute @workgroup_size(16, 16)
572fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
573    let x = gid.x;
574    let y = gid.y;
575    if x >= params.width || y >= params.height { return; }
576    let i = y * params.width + x;
577    let pix = unpack(src[i]);
578    let gx = i32(floor(f32(x) / params.step));
579    let gy = i32(floor(f32(y) / params.step));
580    var best = 0u;
581    var best_d = 1e20;
582    let gw = i32(ceil(f32(params.width) / params.step));
583    for (var oy = -1; oy <= 1; oy++) {
584        for (var ox = -1; ox <= 1; ox++) {
585            let cx = gx + ox;
586            let cy = gy + oy;
587            if cx < 0 || cy < 0 || cx >= gw { continue; }
588            let ci = u32(cy * gw + cx);
589            if ci >= params.k { continue; }
590            let c = centers[ci];
591            let dc = pix - vec3<f32>(c.r, c.g, c.b);
592            let dx = f32(x) - c.x;
593            let dy = f32(y) - c.y;
594            let d = dot(dc, dc) + params.compact * (dx * dx + dy * dy);
595            if d < best_d {
596                best_d = d;
597                best = ci;
598            }
599        }
600    }
601    labels[i] = best;
602}
603"#;
604
605const SLIC_ZERO_WGSL: &str = r#"
606struct Params { width: u32, height: u32, k: u32, step: f32, compact: f32, _pad: f32, _pad2: f32, _pad3: f32 }
607@group(0) @binding(0) var<storage, read_write> acc: array<atomic<u32>>;
608@group(0) @binding(1) var<uniform> params: Params;
609@compute @workgroup_size(64)
610fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
611    let i = gid.x;
612    if i >= params.k * 6u { return; }
613    atomicStore(&acc[i], 0u);
614}
615"#;
616
617const SLIC_ACCUM_WGSL: &str = r#"
618struct Params { width: u32, height: u32, k: u32, step: f32, compact: f32, _pad: f32, _pad2: f32, _pad3: f32 }
619@group(0) @binding(0) var<storage, read> src: array<u32>;
620@group(0) @binding(1) var<storage, read> labels: array<u32>;
621@group(0) @binding(2) var<storage, read_write> acc: array<atomic<u32>>;
622@group(0) @binding(3) var<uniform> params: Params;
623@compute @workgroup_size(16, 16)
624fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
625    let x = gid.x;
626    let y = gid.y;
627    if x >= params.width || y >= params.height { return; }
628    let i = y * params.width + x;
629    let lab = labels[i];
630    if lab >= params.k { return; }
631    let p = src[i];
632    let base = lab * 6u;
633    atomicAdd(&acc[base + 0u], x);
634    atomicAdd(&acc[base + 1u], y);
635    atomicAdd(&acc[base + 2u], (p >> 16u) & 255u);
636    atomicAdd(&acc[base + 3u], (p >> 8u) & 255u);
637    atomicAdd(&acc[base + 4u], p & 255u);
638    atomicAdd(&acc[base + 5u], 1u);
639}
640"#;
641
642const SLIC_DIV_WGSL: &str = r#"
643struct Center { x: f32, y: f32, r: f32, g: f32, b: f32, _p0: f32, _p1: f32, _p2: f32 }
644struct Params { width: u32, height: u32, k: u32, step: f32, compact: f32, _pad: f32, _pad2: f32, _pad3: f32 }
645@group(0) @binding(0) var<storage, read_write> acc: array<atomic<u32>>;
646@group(0) @binding(1) var<storage, read_write> centers: array<Center>;
647@group(0) @binding(2) var<uniform> params: Params;
648@compute @workgroup_size(64)
649fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
650    let i = gid.x;
651    if i >= params.k { return; }
652    let base = i * 6u;
653    let n = max(atomicLoad(&acc[base + 5u]), 1u);
654    let nf = f32(n);
655    centers[i].x = f32(atomicLoad(&acc[base + 0u])) / nf;
656    centers[i].y = f32(atomicLoad(&acc[base + 1u])) / nf;
657    centers[i].r = f32(atomicLoad(&acc[base + 2u])) / nf;
658    centers[i].g = f32(atomicLoad(&acc[base + 3u])) / nf;
659    centers[i].b = f32(atomicLoad(&acc[base + 4u])) / nf;
660}
661"#;
uzor_vision/gpu.rs

uzor_vision/
gpu.rs
