comfy_wgpu/
pipelines.rs

1use std::sync::atomic::AtomicU32;
2
3use crate::*;
4
5static GENERATED_RENDER_TARGET_IDS: AtomicU32 = AtomicU32::new(1);
6
7#[derive(Clone, Debug)]
8pub struct RenderTargetParams {
9    pub label: String,
10    pub size: UVec2,
11    pub filter_mode: wgpu::FilterMode,
12}
13
14/// Creates a new render target with given dimensions. Among other parameters a label is
15/// required so that graphic debuggers like RenderDoc can display its name properly.
16pub fn create_render_target(
17    renderer: &mut WgpuRenderer,
18    params: &RenderTargetParams,
19) -> RenderTargetId {
20    let id = gen_render_target();
21
22    let c = &renderer.context;
23
24    let size = wgpu::Extent3d {
25        width: params.size.x,
26        height: params.size.y,
27        depth_or_array_layers: 1,
28    };
29
30    let texture = c.device.create_texture(&wgpu::TextureDescriptor {
31        label: Some(&params.label),
32        size,
33        mip_level_count: 1,
34        sample_count: 1,
35        dimension: wgpu::TextureDimension::D2,
36        format: wgpu::TextureFormat::Rgba16Float,
37        usage: wgpu::TextureUsages::TEXTURE_BINDING |
38            wgpu::TextureUsages::RENDER_ATTACHMENT,
39        view_formats: &[],
40    });
41
42    let view = texture.create_view(&wgpu::TextureViewDescriptor {
43        label: Some(&format!("{} View", params.label)),
44        format: None,
45        dimension: None,
46        aspect: wgpu::TextureAspect::All,
47        base_mip_level: 0,
48        mip_level_count: None,
49        base_array_layer: 0,
50        array_layer_count: None,
51    });
52
53    let sampler = c.device.create_sampler(&wgpu::SamplerDescriptor {
54        label: Some(&format!("{} Sampler", params.label)),
55        address_mode_u: wgpu::AddressMode::ClampToEdge,
56        address_mode_v: wgpu::AddressMode::ClampToEdge,
57        address_mode_w: wgpu::AddressMode::ClampToEdge,
58        mag_filter: params.filter_mode,
59        min_filter: params.filter_mode,
60        mipmap_filter: params.filter_mode,
61        ..Default::default()
62    });
63
64    let bind_group = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
65        label: Some(&format!("{} Bind Group", params.label)),
66        layout: &c.texture_layout,
67        entries: &[
68            wgpu::BindGroupEntry {
69                binding: 0,
70                resource: wgpu::BindingResource::TextureView(&view),
71            },
72            wgpu::BindGroupEntry {
73                binding: 1,
74                resource: wgpu::BindingResource::Sampler(&sampler),
75            },
76        ],
77    });
78
79    renderer.render_targets.borrow_mut().insert(id, UserRenderTarget {
80        creation_params: params.clone(),
81        texture,
82        view,
83        sampler,
84        bind_group,
85    });
86
87    id
88}
89
90pub struct UserRenderTarget {
91    pub creation_params: RenderTargetParams,
92    pub texture: wgpu::Texture,
93    pub view: wgpu::TextureView,
94    pub sampler: wgpu::Sampler,
95    pub bind_group: wgpu::BindGroup,
96}
97
98/// Allocates a new render target id
99fn gen_render_target() -> RenderTargetId {
100    let id = GENERATED_RENDER_TARGET_IDS
101        .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
102
103    RenderTargetId(id)
104}
105
106pub fn ensure_pipeline_exists(
107    c: &mut WgpuRenderer,
108    pass_data: &MeshDrawData,
109    sprite_shader_id: ShaderId,
110) -> String {
111    let shaders = c.shaders.borrow();
112
113    let maybe_shader_instance_id = pass_data.shader;
114
115    let maybe_shader = {
116        if maybe_shader_instance_id.0 > 0 {
117            let instance = get_shader_instance(maybe_shader_instance_id);
118            shaders.get(instance.id)
119        } else {
120            None
121        }
122    };
123
124    let name = format!(
125        "{} {:?} {:?} {:?}",
126        if maybe_shader_instance_id.0 > 0 {
127            "USER(Mesh)"
128        } else {
129            "BUILTIN(Mesh)"
130        },
131        pass_data.blend_mode,
132        maybe_shader,
133        c.enable_z_buffer
134    );
135
136    let mesh_pipeline = if let Some(shader) = maybe_shader {
137        RenderPipeline::User(
138            c.user_pipelines.entry(name.clone()).or_insert_with(|| {
139                create_user_pipeline(
140                    &name,
141                    pass_data,
142                    shader,
143                    &c.context,
144                    &c.texture_layout,
145                    &c.camera_bind_group_layout,
146                    c.enable_z_buffer,
147                )
148            }),
149        )
150    } else {
151        RenderPipeline::Wgpu(c.pipelines.entry(name.clone()).or_insert_with(
152            || {
153                create_render_pipeline_with_layout(
154                    &name,
155                    &c.context.device,
156                    wgpu::TextureFormat::Rgba16Float,
157                    &[&c.texture_layout, &c.camera_bind_group_layout],
158                    &[SpriteVertex::desc()],
159                    shaders.get(sprite_shader_id).unwrap(),
160                    pass_data.blend_mode,
161                    c.enable_z_buffer,
162                )
163                .unwrap()
164            },
165        ))
166    };
167
168    if let RenderPipeline::User(user_pipeline) = mesh_pipeline {
169        if maybe_shader_instance_id.0 > 0 {
170            let shader_instance = get_shader_instance(maybe_shader_instance_id);
171            let shader = shaders.get(shader_instance.id).unwrap();
172
173            for (buffer_name, buffer) in
174                user_pipeline.buffers.iter().sorted_by_key(|x| x.0)
175            {
176                if let Some(Uniform::F32(OrderedFloat(value))) =
177                    shader_instance.uniforms.get(buffer_name)
178                {
179                    c.context.queue.write_buffer(
180                        buffer,
181                        0,
182                        bytemuck::cast_slice(&[*value]),
183                    );
184                } else if let UniformDef::F32(Some(default_value)) =
185                    shader.uniform_defs.get(buffer_name).unwrap()
186                {
187                    c.context.queue.write_buffer(
188                        buffer,
189                        0,
190                        bytemuck::cast_slice(&[*default_value]),
191                    );
192                } else {
193                    panic!("No uniform value or default for {buffer_name}");
194                }
195            }
196        }
197    }
198
199    name
200}
201
202pub fn create_user_pipeline(
203    name: &str,
204    pass_data: &MeshDrawData,
205    shader: &Shader,
206    context: &GraphicsContext,
207    texture_layout: &Arc<wgpu::BindGroupLayout>,
208    camera_bind_group_layout: &wgpu::BindGroupLayout,
209    enable_z_buffer: bool,
210) -> UserRenderPipeline {
211    info!("Creating pipeline for shader: {:?}", shader.id);
212
213    let mut layout_entries = Vec::new();
214    let mut bind_group_entries = Vec::new();
215    let mut buffers = HashMap::new();
216
217    for (uniform_name, binding) in shader.bindings.iter() {
218        let uniform_def = shader.uniform_defs.get(uniform_name).unwrap();
219
220        layout_entries.push(wgpu::BindGroupLayoutEntry {
221            binding: *binding,
222            visibility: wgpu::ShaderStages::FRAGMENT,
223            ty: wgpu::BindingType::Buffer {
224                ty: wgpu::BufferBindingType::Uniform,
225                has_dynamic_offset: false,
226                min_binding_size: None,
227            },
228            count: None,
229        });
230
231        let uniform_buffer_usage =
232            wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST;
233
234        match uniform_def {
235            UniformDef::F32(maybe_default) => {
236                if let Some(value) = maybe_default {
237                    let buffer = context.device.create_buffer_init(
238                        &wgpu::util::BufferInitDescriptor {
239                            label: Some(&format!(
240                                "User UB: {} (default={})",
241                                uniform_name, value
242                            )),
243                            contents: bytemuck::cast_slice(&[*value]),
244                            usage: uniform_buffer_usage,
245                        },
246                    );
247
248                    buffers.insert(uniform_name.to_string(), buffer);
249                } else {
250                    let buffer =
251                        context.device.create_buffer(&wgpu::BufferDescriptor {
252                            label: Some(&format!(
253                                "User UB: {} (no-default)",
254                                uniform_name
255                            )),
256                            size: std::mem::size_of::<f32>() as u64,
257                            usage: uniform_buffer_usage,
258                            mapped_at_creation: false,
259                        });
260
261                    buffers.insert(uniform_name.to_string(), buffer);
262                }
263            }
264            UniformDef::Custom { .. } => {
265                unimplemented!("custom uniforms aren't available yet");
266            }
267        };
268    }
269
270    for (name, binding) in shader.bindings.iter() {
271        bind_group_entries.push(wgpu::BindGroupEntry {
272            binding: *binding,
273            resource: buffers.get(name).unwrap().as_entire_binding(),
274        });
275    }
276
277    let user_layout = context.device.create_bind_group_layout(
278        &wgpu::BindGroupLayoutDescriptor {
279            label: Some(&format!("User Layout: {}", name)),
280            entries: &layout_entries,
281        },
282    );
283
284    let pipeline = create_render_pipeline_with_layout(
285        name,
286        &context.device,
287        wgpu::TextureFormat::Rgba16Float,
288        &[&texture_layout, &camera_bind_group_layout, &user_layout],
289        &[SpriteVertex::desc()],
290        shader,
291        pass_data.blend_mode,
292        enable_z_buffer,
293    )
294    .unwrap();
295
296    let bind_group =
297        context.device.create_bind_group(&wgpu::BindGroupDescriptor {
298            label: Some("User Bind Group"),
299            layout: &user_layout,
300            entries: &bind_group_entries,
301        });
302
303    UserRenderPipeline { pipeline, layout: user_layout, bind_group, buffers }
304}