tessera_ui_basic_components/pipelines/mean/
pipeline.rs

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