spatialrust_gpu/image/kernels/
box_blur.rs

1//! Gray box blur on the GPU.
2
3use bytemuck::{Pod, Zeroable};
4use spatialrust_core::{SpatialError, SpatialResult};
5use wgpu::util::DeviceExt;
6
7use crate::image::gpu_image::{create_texture, GpuImage, GpuImageReceipt};
8use crate::WgpuRuntime;
9
10const WORKGROUP: u32 = 256;
11
12/// Border extrapolation for GPU image filters.
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub enum GpuImageBorder {
15    /// Clamp coordinates to the nearest valid edge pixel (replicate).
16    Replicate,
17    /// Treat out-of-bounds samples as zero.
18    ConstantZero,
19}
20
21#[repr(C)]
22#[derive(Clone, Copy, Pod, Zeroable)]
23struct BlurParams {
24    width: u32,
25    height: u32,
26    kernel_width: u32,
27    kernel_height: u32,
28    border_mode: u32,
29    _pad0: u32,
30    _pad1: u32,
31    _pad2: u32,
32}
33
34const BLUR_WGSL: &str = r#"
35struct Params {
36    width: u32,
37    height: u32,
38    kernel_width: u32,
39    kernel_height: u32,
40    border_mode: u32,
41    pad0: u32,
42    pad1: u32,
43    pad2: u32,
44};
45
46@group(0) @binding(0) var<uniform> params: Params;
47@group(0) @binding(1) var input_px: texture_2d<u32>;
48@group(0) @binding(2) var output_px: texture_storage_2d<rgba8uint, write>;
49
50fn sample_gray(x: i32, y: i32) -> u32 {
51    var sx = x;
52    var sy = y;
53    if (params.border_mode == 0u) {
54        sx = clamp(sx, 0, i32(params.width) - 1);
55        sy = clamp(sy, 0, i32(params.height) - 1);
56    } else if (sx < 0 || sy < 0 || sx >= i32(params.width) || sy >= i32(params.height)) {
57        return 0u;
58    }
59    return textureLoad(input_px, vec2<i32>(sx, sy), 0).r;
60}
61
62@compute @workgroup_size(256)
63fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
64    let index = gid.x;
65    let pixel_count = params.width * params.height;
66    if (index >= pixel_count) {
67        return;
68    }
69    let x = i32(index % params.width);
70    let y = i32(index / params.width);
71    let radius_x = i32(params.kernel_width / 2u);
72    let radius_y = i32(params.kernel_height / 2u);
73    var sum: u32 = 0u;
74    var count: u32 = 0u;
75    for (var dy: i32 = -radius_y; dy <= radius_y; dy = dy + 1) {
76        for (var dx: i32 = -radius_x; dx <= radius_x; dx = dx + 1) {
77            sum = sum + sample_gray(x + dx, y + dy);
78            count = count + 1u;
79        }
80    }
81    let value = (sum + count / 2u) / count;
82    textureStore(output_px, vec2<i32>(x, y), vec4<u32>(value, 0u, 0u, 0u));
83}
84"#;
85
86pub(crate) struct BlurPipeline {
87    bind_group_layout: wgpu::BindGroupLayout,
88    pipeline: wgpu::ComputePipeline,
89}
90
91pub(crate) fn create_blur_pipeline(device: &wgpu::Device) -> BlurPipeline {
92    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
93        label: Some("gpu-image-blur-bgl"),
94        entries: &[
95            wgpu::BindGroupLayoutEntry {
96                binding: 0,
97                visibility: wgpu::ShaderStages::COMPUTE,
98                ty: wgpu::BindingType::Buffer {
99                    ty: wgpu::BufferBindingType::Uniform,
100                    has_dynamic_offset: false,
101                    min_binding_size: None,
102                },
103                count: None,
104            },
105            wgpu::BindGroupLayoutEntry {
106                binding: 1,
107                visibility: wgpu::ShaderStages::COMPUTE,
108                ty: wgpu::BindingType::Texture {
109                    sample_type: wgpu::TextureSampleType::Uint,
110                    view_dimension: wgpu::TextureViewDimension::D2,
111                    multisampled: false,
112                },
113                count: None,
114            },
115            wgpu::BindGroupLayoutEntry {
116                binding: 2,
117                visibility: wgpu::ShaderStages::COMPUTE,
118                ty: wgpu::BindingType::StorageTexture {
119                    access: wgpu::StorageTextureAccess::WriteOnly,
120                    format: wgpu::TextureFormat::Rgba8Uint,
121                    view_dimension: wgpu::TextureViewDimension::D2,
122                },
123                count: None,
124            },
125        ],
126    });
127    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
128        label: Some("gpu-image-blur-shader"),
129        source: wgpu::ShaderSource::Wgsl(BLUR_WGSL.into()),
130    });
131    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
132        label: Some("gpu-image-blur-pl"),
133        bind_group_layouts: &[&bind_group_layout],
134        push_constant_ranges: &[],
135    });
136    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
137        label: Some("gpu-image-blur-pipeline"),
138        layout: Some(&pipeline_layout),
139        module: &shader,
140        entry_point: Some("main"),
141        compilation_options: wgpu::PipelineCompilationOptions::default(),
142        cache: None,
143    });
144    BlurPipeline { bind_group_layout, pipeline }
145}
146
147/// Applies a square/rectangular mean box blur to a single-channel `GpuImage`.
148pub fn box_blur_gpu(
149    runtime: &WgpuRuntime,
150    source: &GpuImage,
151    kernel_width: u32,
152    kernel_height: u32,
153    border: GpuImageBorder,
154) -> SpatialResult<GpuImage> {
155    source.validate_runtime(runtime)?;
156    if source.channels() != 1 {
157        return Err(SpatialError::InvalidArgument(
158            "box_blur_gpu currently supports single-channel GpuImage inputs".to_owned(),
159        ));
160    }
161    if kernel_width == 0 || kernel_height == 0 || kernel_width % 2 == 0 || kernel_height % 2 == 0 {
162        return Err(SpatialError::InvalidArgument(
163            "box_blur_gpu requires positive odd kernel dimensions".to_owned(),
164        ));
165    }
166    let pixel_count = (source.width() as usize).saturating_mul(source.height() as usize);
167    let out_bytes = u64::from(source.width()) * u64::from(source.height()) * 4;
168    let output = create_texture(runtime, source.width(), source.height(), "gpu-image-blur-out");
169    let params = BlurParams {
170        width: source.width(),
171        height: source.height(),
172        kernel_width,
173        kernel_height,
174        border_mode: match border {
175            GpuImageBorder::Replicate => 0,
176            GpuImageBorder::ConstantZero => 1,
177        },
178        _pad0: 0,
179        _pad1: 0,
180        _pad2: 0,
181    };
182    let uniform = runtime.device().create_buffer_init(&wgpu::util::BufferInitDescriptor {
183        label: Some("gpu-image-blur-params"),
184        contents: bytemuck::bytes_of(&params),
185        usage: wgpu::BufferUsages::UNIFORM,
186    });
187    let pipeline =
188        runtime.image_blur_pipeline.get_or_init(|| create_blur_pipeline(runtime.device()));
189    let source_view = source.view();
190    let output_view = output.create_view(&wgpu::TextureViewDescriptor::default());
191    let bind_group = runtime.device().create_bind_group(&wgpu::BindGroupDescriptor {
192        label: Some("gpu-image-blur-bg"),
193        layout: &pipeline.bind_group_layout,
194        entries: &[
195            wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() },
196            wgpu::BindGroupEntry {
197                binding: 1,
198                resource: wgpu::BindingResource::TextureView(&source_view),
199            },
200            wgpu::BindGroupEntry {
201                binding: 2,
202                resource: wgpu::BindingResource::TextureView(&output_view),
203            },
204        ],
205    });
206    let mut encoder = runtime.device().create_command_encoder(&wgpu::CommandEncoderDescriptor {
207        label: Some("gpu-image-blur-encoder"),
208    });
209    {
210        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
211            label: Some("gpu-image-blur-pass"),
212            timestamp_writes: None,
213        });
214        pass.set_pipeline(&pipeline.pipeline);
215        pass.set_bind_group(0, &bind_group, &[]);
216        pass.dispatch_workgroups(pixel_count.div_ceil(WORKGROUP as usize) as u32, 1, 1);
217    }
218    runtime.queue().submit(Some(encoder.finish()));
219    let mut receipt = GpuImageReceipt::default();
220    receipt.merge_from(source.receipt());
221    receipt.record_gpu_to_gpu(out_bytes, "box_blur_gpu");
222    GpuImage::from_parts(
223        runtime,
224        source.width(),
225        source.height(),
226        1,
227        output,
228        source.metadata(),
229        receipt,
230    )
231}
spatialrust_gpu/image/kernels/box_blur.rs

spatialrust_gpu/image/kernels/
box_blur.rs
