tessera_ui_basic_components/pipelines/contrast/
pipeline.rs

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