spatialrust_gpu/image/kernels/
gray.rs

1//! BT.601 RGB to gray on the GPU.
2
3use bytemuck::{Pod, Zeroable};
4use spatialrust_core::{SpatialError, SpatialResult};
5use spatialrust_image::{ColorSpace, ImageMetadata};
6use wgpu::util::DeviceExt;
7
8use crate::image::gpu_image::{create_texture, GpuImage, GpuImageReceipt};
9use crate::WgpuRuntime;
10
11const WORKGROUP: u32 = 256;
12
13#[repr(C)]
14#[derive(Clone, Copy, Pod, Zeroable)]
15struct GrayParams {
16    width: u32,
17    height: u32,
18    _pad0: u32,
19    _pad1: u32,
20}
21
22const GRAY_WGSL: &str = r#"
23struct Params { width: u32, height: u32, pad0: u32, pad1: u32, };
24
25@group(0) @binding(0) var<uniform> params: Params;
26@group(0) @binding(1) var input_px: texture_2d<u32>;
27@group(0) @binding(2) var output_px: texture_storage_2d<rgba8uint, write>;
28
29@compute @workgroup_size(256)
30fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
31    let index = gid.x;
32    let pixel_count = params.width * params.height;
33    if (index >= pixel_count) {
34        return;
35    }
36    let xy = vec2<i32>(i32(index % params.width), i32(index / params.width));
37    let rgb = textureLoad(input_px, xy, 0).rgb;
38    // Match CPU BT.601 fixed-point: (77*R + 150*G + 29*B + 128) >> 8
39    let gray = (77u * rgb.r + 150u * rgb.g + 29u * rgb.b + 128u) >> 8u;
40    textureStore(output_px, xy, vec4<u32>(gray, 0u, 0u, 0u));
41}
42"#;
43
44pub(crate) struct GrayPipeline {
45    bind_group_layout: wgpu::BindGroupLayout,
46    pipeline: wgpu::ComputePipeline,
47}
48
49pub(crate) fn create_gray_pipeline(device: &wgpu::Device) -> GrayPipeline {
50    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
51        label: Some("gpu-image-gray-bgl"),
52        entries: &[
53            wgpu::BindGroupLayoutEntry {
54                binding: 0,
55                visibility: wgpu::ShaderStages::COMPUTE,
56                ty: wgpu::BindingType::Buffer {
57                    ty: wgpu::BufferBindingType::Uniform,
58                    has_dynamic_offset: false,
59                    min_binding_size: None,
60                },
61                count: None,
62            },
63            wgpu::BindGroupLayoutEntry {
64                binding: 1,
65                visibility: wgpu::ShaderStages::COMPUTE,
66                ty: wgpu::BindingType::Texture {
67                    sample_type: wgpu::TextureSampleType::Uint,
68                    view_dimension: wgpu::TextureViewDimension::D2,
69                    multisampled: false,
70                },
71                count: None,
72            },
73            wgpu::BindGroupLayoutEntry {
74                binding: 2,
75                visibility: wgpu::ShaderStages::COMPUTE,
76                ty: wgpu::BindingType::StorageTexture {
77                    access: wgpu::StorageTextureAccess::WriteOnly,
78                    format: wgpu::TextureFormat::Rgba8Uint,
79                    view_dimension: wgpu::TextureViewDimension::D2,
80                },
81                count: None,
82            },
83        ],
84    });
85    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
86        label: Some("gpu-image-gray-shader"),
87        source: wgpu::ShaderSource::Wgsl(GRAY_WGSL.into()),
88    });
89    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
90        label: Some("gpu-image-gray-pl"),
91        bind_group_layouts: &[&bind_group_layout],
92        push_constant_ranges: &[],
93    });
94    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
95        label: Some("gpu-image-gray-pipeline"),
96        layout: Some(&pipeline_layout),
97        module: &shader,
98        entry_point: Some("main"),
99        compilation_options: wgpu::PipelineCompilationOptions::default(),
100        cache: None,
101    });
102    GrayPipeline { bind_group_layout, pipeline }
103}
104
105/// Converts an RGB `GpuImage` to a gray `GpuImage` using BT.601 fixed-point luma.
106pub fn rgb_to_gray_gpu(runtime: &WgpuRuntime, source: &GpuImage) -> SpatialResult<GpuImage> {
107    source.validate_runtime(runtime)?;
108    if source.channels() != 3 {
109        return Err(SpatialError::InvalidArgument(
110            "rgb_to_gray_gpu requires a 3-channel RGB GpuImage".to_owned(),
111        ));
112    }
113    let pixel_count = (source.width() as usize).saturating_mul(source.height() as usize);
114    let output = create_texture(runtime, source.width(), source.height(), "gpu-image-gray-out");
115    let out_bytes = u64::from(source.width()) * u64::from(source.height()) * 4;
116    let params = GrayParams { width: source.width(), height: source.height(), _pad0: 0, _pad1: 0 };
117    let uniform = runtime.device().create_buffer_init(&wgpu::util::BufferInitDescriptor {
118        label: Some("gpu-image-gray-params"),
119        contents: bytemuck::bytes_of(&params),
120        usage: wgpu::BufferUsages::UNIFORM,
121    });
122    let pipeline =
123        runtime.image_gray_pipeline.get_or_init(|| create_gray_pipeline(runtime.device()));
124    let source_view = source.view();
125    let output_view = output.create_view(&wgpu::TextureViewDescriptor::default());
126    let bind_group = runtime.device().create_bind_group(&wgpu::BindGroupDescriptor {
127        label: Some("gpu-image-gray-bg"),
128        layout: &pipeline.bind_group_layout,
129        entries: &[
130            wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() },
131            wgpu::BindGroupEntry {
132                binding: 1,
133                resource: wgpu::BindingResource::TextureView(&source_view),
134            },
135            wgpu::BindGroupEntry {
136                binding: 2,
137                resource: wgpu::BindingResource::TextureView(&output_view),
138            },
139        ],
140    });
141    let mut encoder = runtime.device().create_command_encoder(&wgpu::CommandEncoderDescriptor {
142        label: Some("gpu-image-gray-encoder"),
143    });
144    {
145        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
146            label: Some("gpu-image-gray-pass"),
147            timestamp_writes: None,
148        });
149        pass.set_pipeline(&pipeline.pipeline);
150        pass.set_bind_group(0, &bind_group, &[]);
151        pass.dispatch_workgroups(pixel_count.div_ceil(WORKGROUP as usize) as u32, 1, 1);
152    }
153    runtime.queue().submit(Some(encoder.finish()));
154    let mut receipt = GpuImageReceipt::default();
155    receipt.merge_from(source.receipt());
156    receipt.record_gpu_to_gpu(out_bytes, "rgb_to_gray_gpu");
157    let metadata = ImageMetadata { color_space: ColorSpace::Gray, ..source.metadata() };
158    GpuImage::from_parts(runtime, source.width(), source.height(), 1, output, metadata, receipt)
159}
spatialrust_gpu/image/kernels/gray.rs

spatialrust_gpu/image/kernels/
gray.rs
