tessera_ui_basic_components/pipelines/mean/
pipeline.rs

1use tessera_ui::{
2    compute::resource::ComputeResourceManager,
3    renderer::compute::{ComputablePipeline, ComputeBatchItem},
4    wgpu,
5};
6
7use super::command::MeanCommand;
8
9#[repr(C)]
10#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
11struct AreaUniform {
12    area_x: u32,
13    area_y: u32,
14    area_width: u32,
15    area_height: u32,
16}
17
18/// Pipeline for calculating mean luminance using a compute shader.
19pub struct MeanPipeline {
20    pipeline: wgpu::ComputePipeline,
21    bind_group_layout: wgpu::BindGroupLayout,
22}
23
24impl MeanPipeline {
25    pub fn new(device: &wgpu::Device) -> Self {
26        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
27            label: Some("Mean Shader"),
28            source: wgpu::ShaderSource::Wgsl(include_str!("mean.wgsl").into()),
29        });
30
31        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
32            entries: &[
33                // 0: Area Uniform
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(16).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: Result Buffer (Storage)
56                wgpu::BindGroupLayoutEntry {
57                    binding: 2,
58                    visibility: wgpu::ShaderStages::COMPUTE,
59                    ty: wgpu::BindingType::Buffer {
60                        ty: wgpu::BufferBindingType::Storage { read_only: false },
61                        has_dynamic_offset: false,
62                        min_binding_size: Some(std::num::NonZeroU64::new(8).unwrap()),
63                    },
64                    count: None,
65                },
66                // 3: Destination Texture (Storage)
67                wgpu::BindGroupLayoutEntry {
68                    binding: 3,
69                    visibility: wgpu::ShaderStages::COMPUTE,
70                    ty: wgpu::BindingType::StorageTexture {
71                        access: wgpu::StorageTextureAccess::WriteOnly,
72                        format: wgpu::TextureFormat::Rgba8Unorm,
73                        view_dimension: wgpu::TextureViewDimension::D2,
74                    },
75                    count: None,
76                },
77            ],
78            label: Some("mean_bind_group_layout"),
79        });
80
81        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
82            label: Some("Mean 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("Mean Pipeline"),
89            layout: Some(&pipeline_layout),
90            module: &shader,
91            entry_point: Some("main"),
92            compilation_options: Default::default(),
93            cache: None,
94        });
95
96        Self {
97            pipeline,
98            bind_group_layout,
99        }
100    }
101}
102
103impl ComputablePipeline<MeanCommand> for MeanPipeline {
104    /// Dispatches one or more mean luminance compute commands.
105    fn dispatch(
106        &mut self,
107        device: &wgpu::Device,
108        queue: &wgpu::Queue,
109        config: &wgpu::SurfaceConfiguration,
110        compute_pass: &mut wgpu::ComputePass<'_>,
111        items: &[ComputeBatchItem<'_, MeanCommand>],
112        resource_manager: &mut ComputeResourceManager,
113        input_view: &wgpu::TextureView,
114        output_view: &wgpu::TextureView,
115    ) {
116        for item in items {
117            let buffer_ref = item.command.result_buffer_ref();
118            let Some(result_buffer) = resource_manager.get(&buffer_ref) else {
119                continue;
120            };
121            queue.write_buffer(result_buffer, 0, bytemuck::cast_slice(&[0u32, 0u32]));
122            let target_area = item.target_area;
123            let area_uniform = AreaUniform {
124                area_x: target_area.x.0 as u32,
125                area_y: target_area.y.0 as u32,
126                area_width: target_area.width.0 as u32,
127                area_height: target_area.height.0 as u32,
128            };
129            let area_bytes = bytemuck::bytes_of(&area_uniform);
130            let area_buffer = device.create_buffer(&wgpu::BufferDescriptor {
131                label: Some("Mean Area Uniform Buffer"),
132                size: area_bytes.len() as u64,
133                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
134                mapped_at_creation: false,
135            });
136            queue.write_buffer(&area_buffer, 0, area_bytes);
137            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
138                layout: &self.bind_group_layout,
139                entries: &[
140                    wgpu::BindGroupEntry {
141                        binding: 0,
142                        resource: area_buffer.as_entire_binding(),
143                    },
144                    wgpu::BindGroupEntry {
145                        binding: 1,
146                        resource: wgpu::BindingResource::TextureView(input_view),
147                    },
148                    wgpu::BindGroupEntry {
149                        binding: 2,
150                        resource: result_buffer.as_entire_binding(),
151                    },
152                    wgpu::BindGroupEntry {
153                        binding: 3,
154                        resource: wgpu::BindingResource::TextureView(output_view),
155                    },
156                ],
157                label: Some("mean_bind_group"),
158            });
159
160            compute_pass.set_pipeline(&self.pipeline);
161            compute_pass.set_bind_group(0, &bind_group, &[]);
162            compute_pass.dispatch_workgroups(
163                config.width.div_ceil(8),
164                config.height.div_ceil(8),
165                1,
166            );
167        }
168    }
169}