tessera_ui_basic_components/pipelines/contrast/
pipeline.rs

1use tessera_ui::{
2    compute::pipeline::{ComputablePipeline, ComputeContext},
3    wgpu,
4};
5
6use super::command::ContrastCommand;
7
8#[repr(C)]
9#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
10struct Uniforms {
11    contrast: f32,
12    area_x: u32,
13    area_y: u32,
14    area_width: u32,
15    area_height: u32,
16}
17
18/// Pipeline for applying contrast adjustment to an image using a compute shader.
19pub struct ContrastPipeline {
20    pipeline: wgpu::ComputePipeline,
21    bind_group_layout: wgpu::BindGroupLayout,
22}
23
24impl ContrastPipeline {
25    pub fn new(device: &wgpu::Device, pipeline_cache: Option<&wgpu::PipelineCache>) -> Self {
26        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
27            label: Some("Contrast Shader"),
28            source: wgpu::ShaderSource::Wgsl(include_str!("contrast.wgsl").into()),
29        });
30
31        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
32            entries: &[
33                // 0: Uniforms
34                wgpu::BindGroupLayoutEntry {
35                    binding: 0,
36                    visibility: wgpu::ShaderStages::COMPUTE,
37                    ty: wgpu::BindingType::Buffer {
38                        ty: wgpu::BufferBindingType::Uniform,
39                        has_dynamic_offset: false,
40                        min_binding_size: Some(std::num::NonZeroU64::new(20).unwrap()),
41                    },
42                    count: None,
43                },
44                // 1: Source Texture
45                wgpu::BindGroupLayoutEntry {
46                    binding: 1,
47                    visibility: wgpu::ShaderStages::COMPUTE,
48                    ty: wgpu::BindingType::Texture {
49                        sample_type: wgpu::TextureSampleType::Float { filterable: false },
50                        view_dimension: wgpu::TextureViewDimension::D2,
51                        multisampled: false,
52                    },
53                    count: None,
54                },
55                // 2: Destination Texture (Storage)
56                wgpu::BindGroupLayoutEntry {
57                    binding: 2,
58                    visibility: wgpu::ShaderStages::COMPUTE,
59                    ty: wgpu::BindingType::StorageTexture {
60                        access: wgpu::StorageTextureAccess::WriteOnly,
61                        format: wgpu::TextureFormat::Rgba8Unorm,
62                        view_dimension: wgpu::TextureViewDimension::D2,
63                    },
64                    count: None,
65                },
66                // 3: Mean Result Buffer (Storage)
67                wgpu::BindGroupLayoutEntry {
68                    binding: 3,
69                    visibility: wgpu::ShaderStages::COMPUTE,
70                    ty: wgpu::BindingType::Buffer {
71                        ty: wgpu::BufferBindingType::Storage { read_only: true },
72                        has_dynamic_offset: false,
73                        min_binding_size: Some(std::num::NonZeroU64::new(8).unwrap()),
74                    },
75                    count: None,
76                },
77            ],
78            label: Some("contrast_bind_group_layout"),
79        });
80
81        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
82            label: Some("Contrast Pipeline Layout"),
83            bind_group_layouts: &[&bind_group_layout],
84            push_constant_ranges: &[],
85        });
86
87        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
88            label: Some("Contrast Pipeline"),
89            layout: Some(&pipeline_layout),
90            module: &shader,
91            entry_point: Some("main"),
92            compilation_options: Default::default(),
93            cache: pipeline_cache,
94        });
95
96        Self {
97            pipeline,
98            bind_group_layout,
99        }
100    }
101}
102
103impl ComputablePipeline<ContrastCommand> for ContrastPipeline {
104    /// Dispatches one or more contrast adjustment compute commands.
105    fn dispatch(&mut self, context: &mut ComputeContext<ContrastCommand>) {
106        for item in context.items {
107            let Some(mean_buffer) = context
108                .resource_manager
109                .get(&item.command.mean_result_handle)
110            else {
111                continue;
112            };
113
114            let target_area = item.target_area;
115            let uniforms = Uniforms {
116                contrast: item.command.contrast,
117                area_x: target_area.x.0 as u32,
118                area_y: target_area.y.0 as u32,
119                area_width: target_area.width.0 as u32,
120                area_height: target_area.height.0 as u32,
121            };
122
123            let uniform_array = [uniforms];
124            let uniform_bytes = bytemuck::cast_slice(&uniform_array);
125            let uniform_buffer = context.device.create_buffer(&wgpu::BufferDescriptor {
126                label: Some("Contrast Uniform Buffer"),
127                size: uniform_bytes.len() as u64,
128                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
129                mapped_at_creation: false,
130            });
131            context
132                .queue
133                .write_buffer(&uniform_buffer, 0, uniform_bytes);
134
135            let bind_group = context
136                .device
137                .create_bind_group(&wgpu::BindGroupDescriptor {
138                    layout: &self.bind_group_layout,
139                    entries: &[
140                        wgpu::BindGroupEntry {
141                            binding: 0,
142                            resource: uniform_buffer.as_entire_binding(),
143                        },
144                        wgpu::BindGroupEntry {
145                            binding: 1,
146                            resource: wgpu::BindingResource::TextureView(context.input_view),
147                        },
148                        wgpu::BindGroupEntry {
149                            binding: 2,
150                            resource: wgpu::BindingResource::TextureView(context.output_view),
151                        },
152                        wgpu::BindGroupEntry {
153                            binding: 3,
154                            resource: mean_buffer.as_entire_binding(),
155                        },
156                    ],
157                    label: Some("contrast_bind_group"),
158                });
159
160            context.compute_pass.set_pipeline(&self.pipeline);
161            context.compute_pass.set_bind_group(0, &bind_group, &[]);
162            context.compute_pass.dispatch_workgroups(
163                context.config.width.div_ceil(8),
164                context.config.height.div_ceil(8),
165                1,
166            );
167        }
168    }
169}