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