tessera_ui_basic_components/pipelines/mean/
pipeline.rs1use 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
18pub 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 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 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 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 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 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}