est_render/gpu/pipeline/
render.rs

1use std::{
2    collections::HashMap,
3    hash::{DefaultHasher, Hash, Hasher},
4};
5
6use crate::utils::ArcRef;
7
8use super::{
9    manager::{GraphicsPipelineDesc, VertexAttributeLayout},
10    super::{
11        GPUInner,
12        texture::{Texture, TextureSampler, BlendState},
13        shader::{
14            bind_group_manager::BindGroupCreateInfo,
15            GraphicsShader,
16            IndexBufferSize,
17            ShaderBindingType,
18            ShaderCullMode,
19            ShaderFrontFace,
20            ShaderPollygonMode,
21            ShaderTopology,
22            graphics::GraphicsShaderType,
23            types::ShaderReflect
24        },
25        buffer::Buffer,
26        command::{
27            BindGroupAttachment,
28            utils::BindGroupType,
29            renderpass::IntermediateRenderPipeline,
30        },
31    }
32};
33
34#[derive(Debug, Clone, Hash)]
35pub struct RenderPipeline {
36    pub(crate) bind_group: Vec<(u32, wgpu::BindGroup)>,
37    pub(crate) pipeline_desc: GraphicsPipelineDesc,
38    pub(crate) index_format: Option<IndexBufferSize>,
39}
40
41#[derive(Debug, Clone)]
42pub struct RenderPipelineBuilder {
43    pub(crate) gpu: ArcRef<GPUInner>,
44    pub(crate) attachments: Vec<BindGroupAttachment>,
45    pub(crate) shader: Option<IntermediateRenderPipeline>,
46    pub(crate) blend: Option<wgpu::BlendState>,
47    pub(crate) color_write_mask: Option<wgpu::ColorWrites>,
48    pub(crate) shader_reflection: Option<Vec<ShaderReflect>>,
49}
50
51impl RenderPipelineBuilder {
52    pub(crate) fn new(gpu: ArcRef<GPUInner>) -> Self {
53        Self {
54            gpu,
55            attachments: Vec::new(),
56            shader: None,
57            blend: None,
58            color_write_mask: None,
59            shader_reflection: None,
60        }
61    }
62
63    #[inline]
64    pub fn set_blend(mut self, blend: Option<&BlendState>) -> Self {
65        match blend {
66            Some(blend) => {
67                self.blend = Some(blend.clone().into());
68                self.color_write_mask = Some(blend.clone().into());
69            }
70            None => {
71                self.blend = None;
72                self.color_write_mask = None;
73            }
74        }
75
76        self
77    }
78
79    #[inline]
80    pub fn set_shader(self, shader: Option<&GraphicsShader>) -> Self {
81        self.set_shader_with_options(shader, None, None, None, None, None)
82    }
83
84    #[inline]
85    pub fn set_shader_with_options(
86        mut self,
87        shader: Option<&GraphicsShader>,
88        topology: Option<ShaderTopology>,
89        cull_mode: Option<ShaderCullMode>,
90        front_face: Option<ShaderFrontFace>,
91        polygon_mode: Option<ShaderPollygonMode>,
92        index_format: Option<IndexBufferSize>,
93    ) -> Self {
94        match shader {
95            Some(shader) => {
96                let shader_inner = shader.inner.borrow();
97                let (vertex_shader, fragment_shader) = match &shader_inner.ty {
98                    GraphicsShaderType::GraphicsSplit {
99                        vertex_module,
100                        fragment_module,
101                    } => (vertex_module.clone(), fragment_module.clone()),
102                    GraphicsShaderType::GraphicsSingle { module } => (module.clone(), module.clone()),
103                };
104
105                let layout = shader_inner.bind_group_layouts.clone();
106
107                let vertex_reflect = shader_inner.reflection.iter().find(|r| {
108                    matches!(r, ShaderReflect::Vertex { .. })
109                        || matches!(r, ShaderReflect::VertexFragment { .. })
110                });
111
112                let fragment_reflect = shader_inner.reflection.iter().find(|r| {
113                    matches!(r, ShaderReflect::Fragment { .. })
114                        || matches!(r, ShaderReflect::VertexFragment { .. })
115                });
116
117                let vertex_entry_point = match vertex_reflect {
118                    Some(ShaderReflect::Vertex { entry_point, .. }) => Some(entry_point),
119                    Some(ShaderReflect::VertexFragment {
120                        vertex_entry_point, ..
121                    }) => Some(vertex_entry_point),
122                    _ => None,
123                };
124
125                let fragment_entry_point = match fragment_reflect {
126                    Some(ShaderReflect::Fragment { entry_point, .. }) => Some(entry_point),
127                    Some(ShaderReflect::VertexFragment {
128                        fragment_entry_point,
129                        ..
130                    }) => Some(fragment_entry_point),
131                    _ => None,
132                };
133
134                #[cfg(any(debug_assertions, feature = "enable-release-validation"))]
135                {
136                    if vertex_entry_point.is_none() {
137                        panic!("Vertex shader entry point is not found in shader reflection");
138                    }
139
140                    if fragment_entry_point.is_none() {
141                        panic!("Fragment shader entry point is not found in shader reflection");
142                    }
143                }
144
145                let vertex_entry_point = vertex_entry_point.unwrap();
146                let fragment_entry_point = fragment_entry_point.unwrap();
147
148                let attrib_inner = shader.attrib.borrow();
149                let shader_binding = IntermediateRenderPipeline {
150                    shader: (vertex_shader, fragment_shader),
151                    vertex_attribute: (attrib_inner.stride, attrib_inner.attributes.clone()),
152                    shader_entry: (vertex_entry_point.clone(), fragment_entry_point.clone()),
153                    layout: layout,
154                    topology: topology.unwrap_or(attrib_inner.topology),
155                    cull_mode: cull_mode.into(),
156                    front_face: front_face.unwrap_or(attrib_inner.front_face),
157                    polygon_mode: polygon_mode.unwrap_or(attrib_inner.polygon_mode),
158                    index_format: index_format.or_else(|| attrib_inner.index.clone()),
159                };
160
161                self.shader = Some(shader_binding);
162                self.shader_reflection = Some(shader_inner.reflection.clone());
163            }
164            None => {
165                self.shader = None;
166                self.shader_reflection = None;
167            }
168        }
169
170        self
171    }
172
173    #[inline]
174    pub fn set_attachment_sampler(
175        mut self,
176        group: u32,
177        binding: u32,
178        sampler: Option<&TextureSampler>,
179    ) -> Self {
180        match sampler {
181            Some(sampler) => {
182                let attachment = {
183                    let gpu_inner = self.gpu.borrow();
184
185                    BindGroupAttachment {
186                        group,
187                        binding,
188                        attachment: BindGroupType::Sampler(
189                            sampler.make_wgpu(gpu_inner.device()),
190                        ),
191                    }
192                };
193
194                self.insert_or_replace_attachment(group, binding, attachment);
195            }
196            None => {
197                self.remove_attachment(group, binding);
198            }
199        }
200
201        self
202    }
203
204    #[inline]
205    pub fn set_attachment_texture(
206        mut self,
207        group: u32,
208        binding: u32,
209        texture: Option<&Texture>,
210    ) -> Self {
211        match texture {
212            Some(texture) => {
213                let attachment = {
214                    BindGroupAttachment {
215                        group,
216                        binding,
217                        attachment: BindGroupType::Texture(
218                            texture.inner.borrow().wgpu_view.clone(),
219                        ),
220                    }
221                };
222
223                self.insert_or_replace_attachment(group, binding, attachment);
224            }
225            None => {
226                self.remove_attachment(group, binding);
227            }
228        }
229
230        self
231    }
232
233    #[inline]
234    pub fn set_attachment_texture_storage(
235        mut self,
236        group: u32,
237        binding: u32,
238        texture: Option<&Texture>,
239    ) -> Self {
240        match texture {
241            Some(texture) => {
242                let inner = texture.inner.borrow();
243                let attachment = BindGroupAttachment {
244                    group,
245                    binding,
246                    attachment: BindGroupType::TextureStorage(inner.wgpu_view.clone()),
247                };
248
249                self.insert_or_replace_attachment(group, binding, attachment);
250            }
251            None => {
252                self.remove_attachment(group, binding);
253            }
254        }
255
256        self
257    }
258
259    #[inline]
260    pub fn set_attachment_uniform(
261        mut self,
262        group: u32,
263        binding: u32,
264        buffer: Option<&Buffer>,
265    ) -> Self {
266        match buffer {
267            Some(buffer) => {
268                let inner = buffer.inner.borrow();
269                let attachment = BindGroupAttachment {
270                    group,
271                    binding,
272                    attachment: BindGroupType::Uniform(inner.buffer.clone()),
273                };
274
275                self.insert_or_replace_attachment(group, binding, attachment);
276            }
277            None => {
278                self.remove_attachment(group, binding);
279            }
280        }
281
282        self
283    }
284
285    #[inline]
286    pub fn set_attachment_uniform_vec<T>(
287        mut self,
288        group: u32,
289        binding: u32,
290        buffer: Option<Vec<T>>,
291    ) -> Self
292    where
293        T: bytemuck::Pod + bytemuck::Zeroable,
294    {
295        match buffer {
296            Some(buffer) => {
297                let attachment = {
298                    let mut inner = self.gpu.borrow_mut();
299
300                    let buffer = inner.create_buffer_with(&buffer, wgpu::BufferUsages::COPY_DST);
301                    BindGroupAttachment {
302                        group,
303                        binding,
304                        attachment: BindGroupType::Uniform(buffer),
305                    }
306                };
307
308                self.insert_or_replace_attachment(group, binding, attachment);
309            }
310            None => {
311                self.remove_attachment(group, binding);
312            }
313        }
314
315        self
316    }
317
318    #[inline]
319    pub fn set_attachment_uniform_raw<T>(
320        mut self,
321        group: u32,
322        binding: u32,
323        buffer: Option<&[T]>,
324    ) -> Self
325    where
326        T: bytemuck::Pod + bytemuck::Zeroable,
327    {
328        match buffer {
329            Some(buffer) => {
330                let mut inner = self.gpu.borrow_mut();
331
332                let buffer = inner.create_buffer_with(&buffer, wgpu::BufferUsages::COPY_DST);
333                let attachment = BindGroupAttachment {
334                    group,
335                    binding,
336                    attachment: BindGroupType::Uniform(buffer),
337                };
338
339                drop(inner);
340
341                self.insert_or_replace_attachment(group, binding, attachment);
342            }
343            None => {
344                self.remove_attachment(group, binding);
345            }
346        }
347
348        self
349    }
350
351    #[inline]
352    pub fn set_attachment_storage(
353        mut self,
354        group: u32,
355        binding: u32,
356        buffer: Option<&Buffer>,
357    ) -> Self {
358        match buffer {
359            Some(buffer) => {
360                let inner = buffer.inner.borrow();
361                let attachment = BindGroupAttachment {
362                    group,
363                    binding,
364                    attachment: BindGroupType::Storage(inner.buffer.clone()),
365                };
366
367                self.insert_or_replace_attachment(group, binding, attachment);
368            }
369            None => {
370                self.remove_attachment(group, binding);
371            }
372        }
373
374        self
375    }
376
377    #[inline]
378    pub fn set_attachment_storage_raw<T>(
379        mut self,
380        group: u32,
381        binding: u32,
382        buffer: Option<&[T]>,
383    ) -> Self
384    where
385        T: bytemuck::Pod + bytemuck::Zeroable,
386    {
387        match buffer {
388            Some(buffer) => {
389                let mut inner = self.gpu.borrow_mut();
390
391                let buffer = inner.create_buffer_with(&buffer, wgpu::BufferUsages::COPY_DST);
392                let attachment = BindGroupAttachment {
393                    group,
394                    binding,
395                    attachment: BindGroupType::Storage(buffer),
396                };
397
398                drop(inner);
399
400                self.insert_or_replace_attachment(group, binding, attachment);
401            }
402            None => {
403                self.remove_attachment(group, binding);
404            }
405        }
406
407        self
408    }
409
410    #[inline]
411    pub fn set_attachment_storage_vec<T>(
412        mut self,
413        group: u32,
414        binding: u32,
415        buffer: Option<Vec<T>>,
416    ) -> Self
417    where
418        T: bytemuck::Pod + bytemuck::Zeroable,
419    {
420        match buffer {
421            Some(buffer) => {
422                let mut inner = self.gpu.borrow_mut();
423
424                let buffer = inner.create_buffer_with(&buffer, wgpu::BufferUsages::COPY_DST);
425                let attachment = BindGroupAttachment {
426                    group,
427                    binding,
428                    attachment: BindGroupType::Storage(buffer),
429                };
430
431                drop(inner);
432
433                self.insert_or_replace_attachment(group, binding, attachment);
434            }
435            None => {
436                self.remove_attachment(group, binding);
437            }
438        }
439
440        self
441    }
442
443    #[inline]
444    pub(crate) fn remove_attachment(&mut self, group: u32, binding: u32) {
445        self.attachments
446            .retain(|a| a.group != group || a.binding != binding);
447    }
448
449    pub(crate) fn insert_or_replace_attachment(
450        &mut self,
451        group: u32,
452        binding: u32,
453        attachment: BindGroupAttachment,
454    ) {
455        let index = self
456            .attachments
457            .iter()
458            .position(|a| a.group == group && a.binding == binding);
459
460        if let Some(index) = index {
461            self.attachments[index] = attachment;
462        } else {
463            self.attachments.push(attachment);
464        }
465    }
466
467    pub fn build(self) -> Result<RenderPipeline, RenderPipelineError> {
468        if self.shader.is_none() {
469            return Err(RenderPipelineError::ShaderNotSet);
470        }
471
472        let shader_binding = self.shader.unwrap();
473        for attachment in &self.attachments {
474            let r#type = self
475                .shader_reflection
476                .as_ref()
477                .unwrap()
478                .iter()
479                .find_map(|b| {
480                    let bindings = match b {
481                        ShaderReflect::Vertex { bindings, .. }
482                        | ShaderReflect::Fragment { bindings, .. }
483                        | ShaderReflect::VertexFragment { bindings, .. } => bindings,
484                        _ => return None,
485                    };
486
487                    bindings.iter().find_map(|shaderbinding| {
488                        if shaderbinding.group == attachment.group
489                            && shaderbinding.binding == attachment.binding
490                        {
491                            Some(shaderbinding)
492                        } else {
493                            None
494                        }
495                    })
496                });
497
498            if r#type.is_none() {
499                return Err(RenderPipelineError::AttachmentNotSet(
500                    attachment.group,
501                    attachment.binding,
502                ));
503            }
504
505            let r#type = r#type.unwrap();
506
507            if !match r#type.ty {
508                ShaderBindingType::UniformBuffer(_) => {
509                    matches!(attachment.attachment, BindGroupType::Uniform(_))
510                }
511                ShaderBindingType::StorageBuffer(_, _) => {
512                    matches!(attachment.attachment, BindGroupType::Storage(_))
513                }
514                ShaderBindingType::StorageTexture(_) => {
515                    matches!(attachment.attachment, BindGroupType::TextureStorage(_))
516                }
517                ShaderBindingType::Sampler(_) => {
518                    matches!(attachment.attachment, BindGroupType::Sampler(_))
519                }
520                ShaderBindingType::Texture(_) => {
521                    matches!(attachment.attachment, BindGroupType::Texture(_))
522                }
523                ShaderBindingType::PushConstant(_) => {
524                    matches!(attachment.attachment, BindGroupType::Uniform(_))
525                }
526            } {
527                return Err(RenderPipelineError::InvalidAttachmentType(
528                    attachment.group,
529                    attachment.binding,
530                    r#type.ty,
531                ));
532            }
533        }
534
535        let bind_group_hash_key = {
536            let mut hasher = DefaultHasher::new();
537            hasher.write_u64(0u64); // Graphics shader hash id
538
539            for attachment in &self.attachments {
540                attachment.group.hash(&mut hasher);
541                attachment.binding.hash(&mut hasher);
542                match &attachment.attachment {
543                    BindGroupType::Uniform(uniform) => {
544                        uniform.hash(&mut hasher);
545                    }
546                    BindGroupType::Texture(texture) => {
547                        texture.hash(&mut hasher);
548                    }
549                    BindGroupType::TextureStorage(texture) => texture.hash(&mut hasher),
550                    BindGroupType::Sampler(sampler) => sampler.hash(&mut hasher),
551                    BindGroupType::Storage(storage) => storage.hash(&mut hasher),
552                }
553            }
554
555            hasher.finish()
556        };
557
558        let bind_group_attachments = {
559            let mut gpu_inner = self.gpu.borrow_mut();
560
561            match gpu_inner.get_bind_group(bind_group_hash_key) {
562                Some(bind_group) => bind_group,
563                None => {
564                    let mut bind_group_attachments: HashMap<u32, Vec<wgpu::BindGroupEntry>> =
565                        self.attachments.iter().fold(HashMap::new(), |mut map, e| {
566                            let (group, binding, attachment) = (e.group, e.binding, &e.attachment);
567                            let entry = match attachment {
568                                BindGroupType::Uniform(buffer) => wgpu::BindGroupEntry {
569                                    binding,
570                                    resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
571                                        buffer,
572                                        offset: 0,
573                                        size: None,
574                                    }),
575                                },
576                                BindGroupType::Texture(texture) => wgpu::BindGroupEntry {
577                                    binding,
578                                    resource: wgpu::BindingResource::TextureView(texture),
579                                },
580                                BindGroupType::Sampler(sampler) => wgpu::BindGroupEntry {
581                                    binding,
582                                    resource: wgpu::BindingResource::Sampler(sampler),
583                                },
584                                BindGroupType::Storage(buffer) => wgpu::BindGroupEntry {
585                                    binding,
586                                    resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
587                                        buffer,
588                                        offset: 0,
589                                        size: None,
590                                    }),
591                                },
592                                BindGroupType::TextureStorage(texture) => wgpu::BindGroupEntry {
593                                    binding,
594                                    resource: wgpu::BindingResource::TextureView(texture),
595                                },
596                            };
597
598                            map.entry(group).or_insert_with(Vec::new).push(entry);
599                            map
600                        });
601
602                    // sort each group attachments
603                    // group, binding
604                    // this is important for the bind group to be created in the correct order
605                    for entries in bind_group_attachments.values_mut() {
606                        entries.sort_by_key(|e| e.binding);
607                    }
608
609                    let bind_group = bind_group_attachments
610                        .iter()
611                        .map(|(group, entries)| {
612                            let layout = shader_binding
613                                .layout
614                                .iter()
615                                .find(|l| l.group == *group)
616                                .unwrap();
617
618                            (layout, entries.as_slice())
619                        })
620                        .collect::<Vec<_>>();
621
622                    let create_info = BindGroupCreateInfo {
623                        entries: bind_group,
624                    };
625
626                    gpu_inner.create_bind_group(bind_group_hash_key, create_info)
627                }
628            }
629        };
630
631        let attribute = &shader_binding.vertex_attribute;
632        let vertex_desc = VertexAttributeLayout {
633            stride: attribute.0 as wgpu::BufferAddress,
634            step_mode: wgpu::VertexStepMode::Vertex,
635            attributes: attribute.1.clone(),
636        };
637
638        let primitive_state = wgpu::PrimitiveState {
639            topology: shader_binding.topology.into(),
640            strip_index_format: None,
641            front_face: shader_binding.front_face.into(),
642            cull_mode: shader_binding.cull_mode.map(|c| c.into()),
643            polygon_mode: shader_binding.polygon_mode.into(),
644            unclipped_depth: false,
645            conservative: false,
646        };
647
648        let layout = shader_binding
649            .layout
650            .iter()
651            .map(|l| l.layout.clone())
652            .collect::<Vec<_>>();
653
654        let pipeline_desc = GraphicsPipelineDesc {
655            shaders: shader_binding.shader.clone(),
656            entry_point: shader_binding.shader_entry.clone(),
657            render_target: vec![(
658                wgpu::TextureFormat::Rgba8UnormSrgb,
659                self.blend.clone(),
660                self.color_write_mask.clone(),
661            )],
662            depth_stencil: None,
663            vertex_desc,
664            primitive_state,
665            bind_group_layout: layout,
666            msaa_count: 1,
667        };
668
669        Ok(RenderPipeline {
670            bind_group: bind_group_attachments,
671            pipeline_desc,
672            index_format: shader_binding.index_format,
673        })
674    }
675}
676
677#[derive(Debug, Clone, Copy)]
678pub enum RenderPipelineError {
679    ShaderNotSet,
680    InvalidShaderType,
681    AttachmentNotSet(u32, u32),
682    InvalidAttachmentType(u32, u32, ShaderBindingType),
683}