tessera_ui_basic_components/pipelines/
mean.rs

1use tessera_ui::{
2    compute::{ComputeResourceRef, resource::ComputeResourceManager},
3    renderer::compute::{ComputablePipeline, command::ComputeCommand},
4    wgpu,
5};
6
7// --- Command ---
8
9/// A command to calculate the mean luminance of the input texture.
10#[derive(Debug, Clone, Copy)]
11pub struct MeanCommand {
12    result_buffer_ref: ComputeResourceRef,
13}
14
15impl MeanCommand {
16    pub fn new(gpu: &wgpu::Device, compute_resource_manager: &mut ComputeResourceManager) -> Self {
17        let result_buffer = gpu.create_buffer(&wgpu::BufferDescriptor {
18            label: Some("Mean Result Buffer"),
19            size: 8, // two u32s: total_luminance, total_pixels
20            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
21            mapped_at_creation: false,
22        });
23
24        let result_buffer_ref = compute_resource_manager.push(result_buffer);
25        MeanCommand { result_buffer_ref }
26    }
27
28    pub fn result_buffer_ref(&self) -> ComputeResourceRef {
29        self.result_buffer_ref
30    }
31}
32
33impl ComputeCommand for MeanCommand {}
34
35// --- Pipeline ---
36
37pub struct MeanPipeline {
38    pipeline: wgpu::ComputePipeline,
39    bind_group_layout: wgpu::BindGroupLayout,
40}
41
42impl MeanPipeline {
43    pub fn new(device: &wgpu::Device) -> Self {
44        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
45            label: Some("Mean Shader"),
46            source: wgpu::ShaderSource::Wgsl(include_str!("mean/mean.wgsl").into()),
47        });
48
49        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
50            entries: &[
51                // 0: Source Texture
52                wgpu::BindGroupLayoutEntry {
53                    binding: 0,
54                    visibility: wgpu::ShaderStages::COMPUTE,
55                    ty: wgpu::BindingType::Texture {
56                        sample_type: wgpu::TextureSampleType::Float { filterable: false },
57                        view_dimension: wgpu::TextureViewDimension::D2,
58                        multisampled: false,
59                    },
60                    count: None,
61                },
62                // 1: Result Buffer (Storage)
63                wgpu::BindGroupLayoutEntry {
64                    binding: 1,
65                    visibility: wgpu::ShaderStages::COMPUTE,
66                    ty: wgpu::BindingType::Buffer {
67                        ty: wgpu::BufferBindingType::Storage { read_only: false },
68                        has_dynamic_offset: false,
69                        min_binding_size: Some(std::num::NonZeroU64::new(8).unwrap()),
70                    },
71                    count: None,
72                },
73                // 2: Destination Texture (Storage)
74                wgpu::BindGroupLayoutEntry {
75                    binding: 2,
76                    visibility: wgpu::ShaderStages::COMPUTE,
77                    ty: wgpu::BindingType::StorageTexture {
78                        access: wgpu::StorageTextureAccess::WriteOnly,
79                        format: wgpu::TextureFormat::Rgba8Unorm,
80                        view_dimension: wgpu::TextureViewDimension::D2,
81                    },
82                    count: None,
83                },
84            ],
85            label: Some("mean_bind_group_layout"),
86        });
87
88        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
89            label: Some("Mean Pipeline Layout"),
90            bind_group_layouts: &[&bind_group_layout],
91            push_constant_ranges: &[],
92        });
93
94        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
95            label: Some("Mean Pipeline"),
96            layout: Some(&pipeline_layout),
97            module: &shader,
98            entry_point: Some("main"),
99            compilation_options: Default::default(),
100            cache: None,
101        });
102
103        Self {
104            pipeline,
105            bind_group_layout,
106        }
107    }
108}
109
110impl ComputablePipeline<MeanCommand> for MeanPipeline {
111    fn dispatch(
112        &mut self,
113        device: &wgpu::Device,
114        _queue: &wgpu::Queue,
115        config: &wgpu::SurfaceConfiguration,
116        compute_pass: &mut wgpu::ComputePass<'_>,
117        command: &MeanCommand,
118        resource_manager: &mut ComputeResourceManager,
119        input_view: &wgpu::TextureView,
120        output_view: &wgpu::TextureView,
121    ) {
122        let result_buffer = resource_manager.get(&command.result_buffer_ref).unwrap();
123        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
124            layout: &self.bind_group_layout,
125            entries: &[
126                wgpu::BindGroupEntry {
127                    binding: 0,
128                    resource: wgpu::BindingResource::TextureView(input_view),
129                },
130                wgpu::BindGroupEntry {
131                    binding: 1,
132                    resource: result_buffer.as_entire_binding(),
133                },
134                wgpu::BindGroupEntry {
135                    binding: 2,
136                    resource: wgpu::BindingResource::TextureView(output_view),
137                },
138            ],
139            label: Some("mean_bind_group"),
140        });
141
142        compute_pass.set_pipeline(&self.pipeline);
143        compute_pass.set_bind_group(0, &bind_group, &[]);
144        compute_pass.dispatch_workgroups(config.width.div_ceil(8), config.height.div_ceil(8), 1);
145    }
146}