spatialrust_gpu/image/kernels/
spatial.rs

1//! Device-resident resize, edge, and morphology texture kernels.
2
3use bytemuck::{Pod, Zeroable};
4use spatialrust_core::{SpatialError, SpatialResult};
5
6use crate::image::gpu_image::{create_texture, GpuImage, GpuImageReceipt};
7use crate::WgpuRuntime;
8
9const WORKGROUP_X: u32 = 16;
10const WORKGROUP_Y: u32 = 16;
11
12/// Single-channel morphology operation executed on a GPU texture.
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub enum GpuMorphology {
15    /// Replaces each pixel with the neighborhood minimum.
16    Erode,
17    /// Replaces each pixel with the neighborhood maximum.
18    Dilate,
19}
20
21#[repr(C)]
22#[derive(Clone, Copy, Pod, Zeroable)]
23struct SpatialParams {
24    source_width: u32,
25    source_height: u32,
26    output_width: u32,
27    output_height: u32,
28    kernel_width: u32,
29    kernel_height: u32,
30    operation: u32,
31    _pad: u32,
32}
33
34const SPATIAL_WGSL: &str = r#"
35struct Params {
36    source_width: u32,
37    source_height: u32,
38    output_width: u32,
39    output_height: u32,
40    kernel_width: u32,
41    kernel_height: u32,
42    operation: u32,
43    pad: 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 in_output(gid: vec3<u32>) -> bool {
51    return gid.x < params.output_width && gid.y < params.output_height;
52}
53
54fn gray_at(x: i32, y: i32) -> i32 {
55    let sx = clamp(x, 0, i32(params.source_width) - 1);
56    let sy = clamp(y, 0, i32(params.source_height) - 1);
57    return i32(textureLoad(input_px, vec2<i32>(sx, sy), 0).r);
58}
59
60@compute @workgroup_size(16, 16)
61fn resize_nearest(@builtin(global_invocation_id) gid: vec3<u32>) {
62    if (!in_output(gid)) { return; }
63    let sx = min((gid.x * params.source_width) / params.output_width, params.source_width - 1u);
64    let sy = min((gid.y * params.source_height) / params.output_height, params.source_height - 1u);
65    textureStore(output_px, vec2<i32>(gid.xy), textureLoad(input_px, vec2<i32>(i32(sx), i32(sy)), 0));
66}
67
68@compute @workgroup_size(16, 16)
69fn sobel(@builtin(global_invocation_id) gid: vec3<u32>) {
70    if (!in_output(gid)) { return; }
71    let x = i32(gid.x);
72    let y = i32(gid.y);
73    let gx = -gray_at(x - 1, y - 1) + gray_at(x + 1, y - 1)
74        - 2 * gray_at(x - 1, y) + 2 * gray_at(x + 1, y)
75        - gray_at(x - 1, y + 1) + gray_at(x + 1, y + 1);
76    let gy = -gray_at(x - 1, y - 1) - 2 * gray_at(x, y - 1) - gray_at(x + 1, y - 1)
77        + gray_at(x - 1, y + 1) + 2 * gray_at(x, y + 1) + gray_at(x + 1, y + 1);
78    let magnitude = u32(clamp(abs(gx) + abs(gy), 0, 255));
79    textureStore(output_px, vec2<i32>(x, y), vec4<u32>(magnitude, 0u, 0u, 0u));
80}
81
82@compute @workgroup_size(16, 16)
83fn morphology(@builtin(global_invocation_id) gid: vec3<u32>) {
84    if (!in_output(gid)) { return; }
85    let x = i32(gid.x);
86    let y = i32(gid.y);
87    let rx = i32(params.kernel_width / 2u);
88    let ry = i32(params.kernel_height / 2u);
89    var value = select(255, 0, params.operation == 1u);
90    for (var dy = -ry; dy <= ry; dy = dy + 1) {
91        for (var dx = -rx; dx <= rx; dx = dx + 1) {
92            let sample = gray_at(x + dx, y + dy);
93            if (params.operation == 0u) { value = min(value, sample); }
94            else { value = max(value, sample); }
95        }
96    }
97    textureStore(output_px, vec2<i32>(x, y), vec4<u32>(u32(value), 0u, 0u, 0u));
98}
99"#;
100
101pub(crate) struct SpatialPipelines {
102    layout: wgpu::BindGroupLayout,
103    resize: wgpu::ComputePipeline,
104    sobel: wgpu::ComputePipeline,
105    morphology: wgpu::ComputePipeline,
106}
107
108pub(crate) fn create_spatial_pipelines(device: &wgpu::Device) -> SpatialPipelines {
109    let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
110        label: Some("gpu-image-spatial-bgl"),
111        entries: &[
112            wgpu::BindGroupLayoutEntry {
113                binding: 0,
114                visibility: wgpu::ShaderStages::COMPUTE,
115                ty: wgpu::BindingType::Buffer {
116                    ty: wgpu::BufferBindingType::Uniform,
117                    has_dynamic_offset: false,
118                    min_binding_size: None,
119                },
120                count: None,
121            },
122            wgpu::BindGroupLayoutEntry {
123                binding: 1,
124                visibility: wgpu::ShaderStages::COMPUTE,
125                ty: wgpu::BindingType::Texture {
126                    sample_type: wgpu::TextureSampleType::Uint,
127                    view_dimension: wgpu::TextureViewDimension::D2,
128                    multisampled: false,
129                },
130                count: None,
131            },
132            wgpu::BindGroupLayoutEntry {
133                binding: 2,
134                visibility: wgpu::ShaderStages::COMPUTE,
135                ty: wgpu::BindingType::StorageTexture {
136                    access: wgpu::StorageTextureAccess::WriteOnly,
137                    format: wgpu::TextureFormat::Rgba8Uint,
138                    view_dimension: wgpu::TextureViewDimension::D2,
139                },
140                count: None,
141            },
142        ],
143    });
144    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
145        label: Some("gpu-image-spatial-shader"),
146        source: wgpu::ShaderSource::Wgsl(SPATIAL_WGSL.into()),
147    });
148    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
149        label: Some("gpu-image-spatial-pl"),
150        bind_group_layouts: &[&layout],
151        push_constant_ranges: &[],
152    });
153    let create = |entry_point| {
154        device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
155            label: Some(entry_point),
156            layout: Some(&pipeline_layout),
157            module: &shader,
158            entry_point: Some(entry_point),
159            compilation_options: wgpu::PipelineCompilationOptions::default(),
160            cache: None,
161        })
162    };
163    SpatialPipelines {
164        layout,
165        resize: create("resize_nearest"),
166        sobel: create("sobel"),
167        morphology: create("morphology"),
168    }
169}
170
171/// Resizes a device image with nearest-neighbor sampling without host transfer.
172pub fn resize_nearest_gpu(
173    runtime: &WgpuRuntime,
174    source: &GpuImage,
175    width: u32,
176    height: u32,
177) -> SpatialResult<GpuImage> {
178    if width == 0 || height == 0 {
179        return Err(SpatialError::InvalidArgument(
180            "GPU resize output dimensions must be positive".to_owned(),
181        ));
182    }
183    dispatch(
184        runtime,
185        source,
186        width,
187        height,
188        1,
189        1,
190        0,
191        &runtime
192            .image_spatial_pipelines
193            .get_or_init(|| create_spatial_pipelines(runtime.device()))
194            .resize,
195        "resize_nearest_gpu",
196        source.channels(),
197    )
198}
199
200/// Computes clamped L1 Sobel magnitude on a single-channel device image.
201pub fn sobel_gpu(runtime: &WgpuRuntime, source: &GpuImage) -> SpatialResult<GpuImage> {
202    require_gray(source, "sobel_gpu")?;
203    dispatch(
204        runtime,
205        source,
206        source.width(),
207        source.height(),
208        3,
209        3,
210        0,
211        &runtime
212            .image_spatial_pipelines
213            .get_or_init(|| create_spatial_pipelines(runtime.device()))
214            .sobel,
215        "sobel_gpu",
216        1,
217    )
218}
219
220/// Applies rectangular erosion or dilation without leaving device memory.
221pub fn morphology_gpu(
222    runtime: &WgpuRuntime,
223    source: &GpuImage,
224    kernel_width: u32,
225    kernel_height: u32,
226    operation: GpuMorphology,
227) -> SpatialResult<GpuImage> {
228    require_gray(source, "morphology_gpu")?;
229    if kernel_width == 0
230        || kernel_height == 0
231        || kernel_width % 2 == 0
232        || kernel_height % 2 == 0
233        || kernel_width > 31
234        || kernel_height > 31
235    {
236        return Err(SpatialError::InvalidArgument(
237            "GPU morphology kernels must be odd and in 1..=31".to_owned(),
238        ));
239    }
240    dispatch(
241        runtime,
242        source,
243        source.width(),
244        source.height(),
245        kernel_width,
246        kernel_height,
247        match operation {
248            GpuMorphology::Erode => 0,
249            GpuMorphology::Dilate => 1,
250        },
251        &runtime
252            .image_spatial_pipelines
253            .get_or_init(|| create_spatial_pipelines(runtime.device()))
254            .morphology,
255        "morphology_gpu",
256        1,
257    )
258}
259
260#[allow(clippy::too_many_arguments)]
261fn dispatch(
262    runtime: &WgpuRuntime,
263    source: &GpuImage,
264    output_width: u32,
265    output_height: u32,
266    kernel_width: u32,
267    kernel_height: u32,
268    operation: u32,
269    pipeline: &wgpu::ComputePipeline,
270    stage: &'static str,
271    output_channels: u32,
272) -> SpatialResult<GpuImage> {
273    source.validate_runtime(runtime)?;
274    let output = create_texture(runtime, output_width, output_height, stage);
275    let params = SpatialParams {
276        source_width: source.width(),
277        source_height: source.height(),
278        output_width,
279        output_height,
280        kernel_width,
281        kernel_height,
282        operation,
283        _pad: 0,
284    };
285    let uniform = runtime.device().create_buffer(&wgpu::BufferDescriptor {
286        label: Some("gpu-image-spatial-params"),
287        size: std::mem::size_of::<SpatialParams>() as u64,
288        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
289        mapped_at_creation: false,
290    });
291    runtime.queue().write_buffer(&uniform, 0, bytemuck::bytes_of(&params));
292    let source_view = source.view();
293    let output_view = output.create_view(&wgpu::TextureViewDescriptor::default());
294    let bind_group = runtime.device().create_bind_group(&wgpu::BindGroupDescriptor {
295        label: Some("gpu-image-spatial-bg"),
296        layout: &runtime
297            .image_spatial_pipelines
298            .get_or_init(|| create_spatial_pipelines(runtime.device()))
299            .layout,
300        entries: &[
301            wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() },
302            wgpu::BindGroupEntry {
303                binding: 1,
304                resource: wgpu::BindingResource::TextureView(&source_view),
305            },
306            wgpu::BindGroupEntry {
307                binding: 2,
308                resource: wgpu::BindingResource::TextureView(&output_view),
309            },
310        ],
311    });
312    let mut encoder = runtime.device().create_command_encoder(&wgpu::CommandEncoderDescriptor {
313        label: Some("gpu-image-spatial-encoder"),
314    });
315    {
316        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
317            label: Some(stage),
318            timestamp_writes: None,
319        });
320        pass.set_pipeline(pipeline);
321        pass.set_bind_group(0, &bind_group, &[]);
322        pass.dispatch_workgroups(
323            output_width.div_ceil(WORKGROUP_X),
324            output_height.div_ceil(WORKGROUP_Y),
325            1,
326        );
327    }
328    runtime.queue().submit(Some(encoder.finish()));
329    let output_bytes = u64::from(output_width) * u64::from(output_height) * 4;
330    let mut receipt = GpuImageReceipt::default();
331    receipt.merge_from(source.receipt());
332    receipt.record_gpu_to_gpu(output_bytes, stage);
333    GpuImage::from_parts(
334        runtime,
335        output_width,
336        output_height,
337        output_channels,
338        output,
339        source.metadata(),
340        receipt,
341    )
342}
343
344fn require_gray(source: &GpuImage, operation: &str) -> SpatialResult<()> {
345    if source.channels() != 1 {
346        return Err(SpatialError::InvalidArgument(format!(
347            "{operation} requires a single-channel GpuImage"
348        )));
349    }
350    Ok(())
351}
spatialrust_gpu/image/kernels/spatial.rs

spatialrust_gpu/image/kernels/
spatial.rs
