render: make it easier to share bind groups and bind group layouts between passes

lyra-game/src/render/graph/mod.rs

@@ -2,7 +2,6 @@ mod pass;
 use std::{
     cell::RefCell,
     collections::{HashMap, VecDeque},
-    ptr::NonNull,
     rc::Rc,
     sync::Arc,
 };
@@ -19,15 +18,11 @@ pub use slot_desc::*;
 mod execution_path;
 
 use rustc_hash::{FxHashMap, FxHashSet};
-use tracing::{debug, debug_span, instrument, trace};
-use wgpu::{util::DeviceExt, RenderPass};
+use tracing::{debug_span, instrument, trace};
 
 use self::execution_path::GraphExecutionPath;
 
-use super::{
-    renderer::{BasicRenderer, Renderer},
-    resource::{Pipeline, RenderPipeline},
-};
+use super::resource::{Pipeline, RenderPipeline};
 
 //#[derive(Clone)]
 struct PassEntry {
@@ -35,6 +30,14 @@ struct PassEntry {
     desc: Arc<RenderGraphPassDesc>,
 }
 
+pub struct BindGroupEntry {
+    pub name: String,
+    /// BindGroup
+    pub bg: Rc<wgpu::BindGroup>,
+    /// BindGroupLayout
+    pub layout: Option<Rc<wgpu::BindGroupLayout>>,
+}
+
 struct ResourcedSlot {
     name: String,
     //slot: RenderPassSlot,
@@ -58,7 +61,8 @@ pub struct RenderGraph {
     slot_names: HashMap<String, u64>,
     passes: FxHashMap<u64, PassEntry>,
     // TODO: Use a SlotMap
-    bind_groups: FxHashMap<u64, wgpu::BindGroup>,
+    bind_groups: FxHashMap<u64, BindGroupEntry>,
+    bind_group_names: FxHashMap<String, u64>,
     // TODO: make pipelines a `type` parameter in RenderPasses,
     // then the pipelines can be retrieved via TypeId to the pass.
     ///
@@ -70,7 +74,11 @@ pub struct RenderGraph {
 }
 
 impl RenderGraph {
-    pub fn new(device: Rc<wgpu::Device>, queue: Rc<wgpu::Queue>, surface_config: wgpu::SurfaceConfiguration) -> Self {
+    pub fn new(
+        device: Rc<wgpu::Device>,
+        queue: Rc<wgpu::Queue>,
+        surface_config: wgpu::SurfaceConfiguration,
+    ) -> Self {
         let mut slots = FxHashMap::default();
         let mut slot_names = HashMap::default();
 
@@ -92,6 +100,7 @@ impl RenderGraph {
             slot_names,
             passes: Default::default(),
             bind_groups: Default::default(),
+            bind_group_names: Default::default(),
             pipelines: Default::default(),
             current_id: 1,
             exec_path: None,
@@ -129,19 +138,19 @@ impl RenderGraph {
                 .get(&slot.name)
                 .and_then(|id| self.slots.get_mut(id).map(|s| (id, s)))
             {
-                // if there is a slot of the same name
-                slot.id = *id;
-
                 debug_assert_eq!(
                     slot.ty, other.ty,
                     "slot {} in pass {} does not match existing slot of same name",
                     slot.name, desc.name
                 );
+
+                // if there is a slot of the same name
+                slot.id = *id;
             } else {
                 let res_slot = ResourcedSlot {
                     name: slot.name.clone(),
                     ty: slot.ty,
-                    value: SlotValue::None,
+                    value: slot.value.clone().unwrap_or(SlotValue::None),
                 };
 
                 self.slots.insert(slot.id, res_slot);
@@ -149,6 +158,19 @@ impl RenderGraph {
             }
         }
 
+        for (name, bg, bgl) in &desc.bind_groups {
+            let bg_id = self.next_id();
+            self.bind_groups.insert(
+                bg_id,
+                BindGroupEntry {
+                    name: name.clone(),
+                    bg: bg.clone(),
+                    layout: bgl.clone(),
+                },
+            );
+            self.bind_group_names.insert(name.clone(), bg_id);
+        }
+
         self.passes.insert(
             desc.id,
             PassEntry {
@@ -189,6 +211,28 @@ impl RenderGraph {
             inner.prepare(world, &mut context);
         }
 
+        // Queue all buffer writes to the gpu
+        {
+            let s = debug_span!("queue_buffer_writes");
+            let _e = s.enter();
+
+            while let Some(bufwr) = context.buffer_writes.pop_front() {
+                let slot = self
+                    .slots
+                    .get(&self.slot_id(&bufwr.target_slot).unwrap())
+                    .expect(&format!(
+                        "Failed to find slot '{}' for buffer write",
+                        bufwr.target_slot
+                    ));
+                let buf = slot
+                    .value
+                    .as_buffer()
+                    .expect(&format!("Slot '{}' is not a buffer", bufwr.target_slot));
+
+                self.queue.write_buffer(buf, bufwr.offset, &bufwr.bytes);
+            }
+        }
+
         // create the execution path for the graph. This will be executed in `RenderGraph::render`
         let builtin = {
             let mut h = FxHashSet::default();
@@ -197,7 +241,10 @@ impl RenderGraph {
         };
         let descs = self.passes.values().map(|p| &*p.desc).collect();
         let path = GraphExecutionPath::new(builtin, descs);
-        trace!("Found {} steps in the rendergraph to execute", path.queue.len());
+        trace!(
+            "Found {} steps in the rendergraph to execute",
+            path.queue.len()
+        );
         self.exec_path = Some(path);
     }
 
@@ -216,7 +263,7 @@ impl RenderGraph {
             "window_texture_view",
             "unexpected slot where 'window_texture_view' should be"
         );
-        window_tv_slot.value = SlotValue::TextureView(view);
+        window_tv_slot.value = SlotValue::TextureView(Rc::new(view));
 
         let mut encoders = vec![];
         while let Some(pass_id) = path.queue.pop_front() {
@@ -246,37 +293,45 @@ impl RenderGraph {
     }
 
     #[inline(always)]
-    pub fn try_bind_group(&self, id: u64) -> Option<&wgpu::BindGroup> {
-        self.bind_groups.get(&id)
+    pub fn try_bind_group(&self, id: u64) -> Option<&Rc<wgpu::BindGroup>> {
+        self.bind_groups.get(&id).map(|e| &e.bg)
     }
 
     #[inline(always)]
-    pub fn bind_group(&self, id: u64) -> &wgpu::BindGroup {
+    pub fn bind_group(&self, id: u64) -> &Rc<wgpu::BindGroup> {
         self.try_bind_group(id).expect("Unknown id for bind group")
     }
 
     #[inline(always)]
-    pub fn try_slot_bind_group(&self, id: u64) -> Option<&wgpu::BindGroup> {
-        todo!()
+    pub fn try_bind_group_layout(&self, id: u64) -> Option<&Rc<wgpu::BindGroupLayout>> {
+        self.bind_groups.get(&id).and_then(|e| e.layout.as_ref())
     }
 
     #[inline(always)]
-    pub fn slot_bind_group(&self, id: u64) -> &wgpu::BindGroup {
-        todo!()
+    pub fn bind_group_layout(&self, id: u64) -> &Rc<wgpu::BindGroupLayout> {
+        self.try_bind_group_layout(id)
+            .expect("Unknown id for bind group layout")
+    }
+
+    #[inline(always)]
+    pub fn bind_group_id(&self, name: &str) -> Option<u64> {
+        self.bind_group_names.get(name).copied()
     }
 }
 
 pub(crate) struct BufferWrite {
     /// The name of the slot that has the resource that will be written
     target_slot: String,
+    offset: u64,
     bytes: Vec<u8>,
 }
 
+#[allow(dead_code)]
 pub struct RenderGraphContext<'a> {
     /// Becomes None when the encoder is submitted
     pub(crate) encoder: Option<wgpu::CommandEncoder>,
     pub(crate) queue: &'a wgpu::Queue,
-    pub(crate) buffer_writes: Vec<BufferWrite>,
+    pub(crate) buffer_writes: VecDeque<BufferWrite>,
     renderpass_desc: Vec<wgpu::RenderPassDescriptor<'a, 'a>>,
 }
 
@@ -285,7 +340,7 @@ impl<'a> RenderGraphContext<'a> {
         Self {
             encoder,
             queue,
-            buffer_writes: vec![],
+            buffer_writes: Default::default(),
             renderpass_desc: vec![],
         }
     }
@@ -296,28 +351,43 @@ impl<'a> RenderGraphContext<'a> {
     ) -> wgpu::RenderPass {
         self.encoder
             .as_mut()
-            .expect("RenderGraphContext is missing a command encoder. This is likely \
-                because you are trying to run render commands in the prepare stage.")
+            .expect(
+                "RenderGraphContext is missing a command encoder. This is likely \
+                because you are trying to run render commands in the prepare stage.",
+            )
             .begin_render_pass(&desc)
     }
 
     pub fn begin_compute_pass(&mut self, desc: &wgpu::ComputePassDescriptor) -> wgpu::ComputePass {
         self.encoder
             .as_mut()
-            .expect("RenderGraphContext is missing a command encoder. This is likely \
-                because you are trying to run render commands in the prepare stage.")
+            .expect(
+                "RenderGraphContext is missing a command encoder. This is likely \
+                because you are trying to run render commands in the prepare stage.",
+            )
             .begin_compute_pass(desc)
     }
 
-    pub fn write_buffer(&mut self, target_slot: &str, bytes: &[u8]) {
-        //self.queue.write_buffer(buffer, offset, data)
-        self.buffer_writes.push(BufferWrite {
+    /// Queue a data write to a buffer at that is contained in `target_slot`.
+    ///
+    /// This does not submit the data to the GPU immediately, or add it to the `wgpu::Queue`. The
+    /// data will be submitted to the GPU queue right after the prepare stage for all passes
+    /// is ran.
+    pub fn queue_buffer_write(&mut self, target_slot: &str, offset: u64, bytes: &[u8]) {
+        self.buffer_writes.push_back(BufferWrite {
             target_slot: target_slot.to_string(),
+            offset,
             bytes: bytes.to_vec(),
         })
     }
 
-    pub fn write_buffer_muck<T: bytemuck::NoUninit>(&mut self, target_slot: &str, bytes: T) {
-        self.write_buffer(target_slot, bytemuck::bytes_of(&bytes));
+    /// Write
+    pub fn queue_buffer_write_with<T: bytemuck::NoUninit>(
+        &mut self,
+        target_slot: &str,
+        offset: u64,
+        bytes: T,
+    ) {
+        self.queue_buffer_write(target_slot, offset, bytemuck::bytes_of(&bytes));
     }
 }

lyra-game/src/render/graph/pass.rs

@@ -4,10 +4,7 @@ use lyra_ecs::World;
 
 use crate::render::resource::RenderPipelineDescriptor;
 
-use super::{
-    BufferDescriptor, BufferInitDescriptor, RenderGraph, RenderGraphContext, SamplerDescriptor,
-    TextureDescriptor, TextureViewDescriptor,
-};
+use super::{RenderGraph, RenderGraphContext};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
 pub enum RenderPassType {
@@ -24,12 +21,12 @@ pub enum SlotType {
     Buffer,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum SlotValue {
     None,
-    TextureView(wgpu::TextureView),
-    Sampler(wgpu::Sampler),
-    Texture(wgpu::Texture),
+    TextureView(Rc<wgpu::TextureView>),
+    Sampler(Rc<wgpu::Sampler>),
+    Texture(Rc<wgpu::Texture>),
     Buffer(Rc<wgpu::Buffer>),
 }
 
@@ -37,17 +34,24 @@ impl SlotValue {
     /// Gets `self` as a texture, panics if self is not a instance of [`SlotValue::Texture`].
     pub fn as_texture(&self) -> &wgpu::Texture {
         match self {
-            Self::Texture(t) => t,
+            Self::Texture(v) => v,
             _ => panic!("self is not an instance of SlotValue::Texture"),
         }
     }
 
     pub fn as_texture_view(&self) -> &wgpu::TextureView {
         match self {
-            Self::TextureView(t) => t,
+            Self::TextureView(v) => v,
             _ => panic!("self is not an instance of SlotValue::TextureView"),
         }
     }
+
+    pub fn as_buffer(&self) -> Option<&wgpu::Buffer> {
+        match self {
+            Self::Buffer(v) => Some(v),
+            _ => None,
+        }
+    }
 }
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 pub enum SlotAttribute {
@@ -63,7 +67,7 @@ pub struct RenderPassSlot {
     pub name: String,
     /// The descriptor of the slot value.
     /// This will be `None` if this slot is an input.
-    pub value: Option<Rc<SlotValue>>,
+    pub value: Option<SlotValue>,
 }
 
 #[derive(Clone)]
@@ -128,6 +132,11 @@ pub struct RenderGraphPassDesc {
     pub slots: Vec<RenderPassSlot>,
     slot_names: HashMap<String, u64>,
     pub pipeline_desc: Option<RenderPipelineDescriptor>,
+    pub bind_groups: Vec<(
+        String,
+        Rc<wgpu::BindGroup>,
+        Option<Rc<wgpu::BindGroupLayout>>,
+    )>,
 }
 
 impl RenderGraphPassDesc {
@@ -136,6 +145,7 @@ impl RenderGraphPassDesc {
         name: &str,
         pass_type: RenderPassType,
         pipeline_desc: Option<RenderPipelineDescriptor>,
+        bind_groups: Vec<(&str, Rc<wgpu::BindGroup>, Option<Rc<wgpu::BindGroupLayout>>)>,
     ) -> Self {
         Self {
             id,
@@ -144,6 +154,10 @@ impl RenderGraphPassDesc {
             slots: vec![],
             slot_names: HashMap::default(),
             pipeline_desc,
+            bind_groups: bind_groups
+                .into_iter()
+                .map(|bg| (bg.0.to_string(), bg.1, bg.2))
+                .collect(),
         }
     }
 
@@ -161,10 +175,7 @@ impl RenderGraphPassDesc {
         value: Option<SlotValue>,
     ) {
         debug_assert!(
-            matches!(
-                value,
-                None | Some(SlotValue::Buffer(_))
-            ),
+            matches!(value, None | Some(SlotValue::Buffer(_))),
             "slot value is not a buffer"
         );
 
@@ -173,7 +184,7 @@ impl RenderGraphPassDesc {
             name: name.to_string(),
             ty: SlotType::Buffer,
             attribute,
-            value: value.map(|v| Rc::new(v)),
+            value,
         };
         self.add_slot(slot);
     }
@@ -187,10 +198,7 @@ impl RenderGraphPassDesc {
         value: Option<SlotValue>,
     ) {
         debug_assert!(
-            matches!(
-                value,
-                None | Some(SlotValue::Texture(_))
-            ),
+            matches!(value, None | Some(SlotValue::Texture(_))),
             "slot value is not a texture"
         );
 
@@ -199,7 +207,7 @@ impl RenderGraphPassDesc {
             name: name.to_string(),
             ty: SlotType::Texture,
             attribute,
-            value: value.map(|v| Rc::new(v)),
+            value,
         };
         self.add_slot(slot);
     }
@@ -213,10 +221,7 @@ impl RenderGraphPassDesc {
         value: Option<SlotValue>,
     ) {
         debug_assert!(
-            matches!(
-                value,
-                None | Some(SlotValue::TextureView(_))
-            ),
+            matches!(value, None | Some(SlotValue::TextureView(_))),
             "slot value is not a texture view"
         );
 
@@ -225,7 +230,7 @@ impl RenderGraphPassDesc {
             name: name.to_string(),
             ty: SlotType::TextureView,
             attribute,
-            value: value.map(|v| Rc::new(v)),
+            value,
         };
         self.add_slot(slot);
     }
@@ -239,10 +244,7 @@ impl RenderGraphPassDesc {
         value: Option<SlotValue>,
     ) {
         debug_assert!(
-            matches!(
-                value,
-                None | Some(SlotValue::Sampler(_))
-            ),
+            matches!(value, None | Some(SlotValue::Sampler(_))),
             "slot value is not a sampler"
         );
 
@@ -251,7 +253,7 @@ impl RenderGraphPassDesc {
             name: name.to_string(),
             ty: SlotType::Sampler,
             attribute,
-            value: value.map(|v| Rc::new(v)),
+            value,
         };
         self.add_slot(slot);
     }

lyra-game/src/render/graph/passes/triangle.rs

@@ -1,22 +1,15 @@
 use std::rc::Rc;
 
-use glam::UVec2;
-use tracing::warn;
-use wgpu::include_wgsl;
-
 use crate::{
     render::{
-        camera::{CameraUniform, RenderCamera},
         graph::{
-            BufferDescriptor, BufferInitDescriptor, PipelineShaderDesc, RenderGraphContext,
-            RenderGraphPass, RenderGraphPassDesc, RenderGraphPipelineInfo, RenderPassType,
+            RenderGraphContext, RenderGraphPass, RenderGraphPassDesc, RenderPassType,
             SlotAttribute, SlotValue,
         },
         render_buffer::BufferWrapper,
-        renderer::ScreenSize,
         resource::{FragmentState, RenderPipelineDescriptor, Shader, VertexState},
     },
-    scene::CameraComponent, DeltaTime,
+    DeltaTime,
 };
 
 /// Supplies some basic things other passes needs.
@@ -24,8 +17,8 @@ use crate::{
 /// screen size buffer, camera buffer,
 #[derive(Default)]
 pub struct TrianglePass {
-    color_bg: Option<wgpu::BindGroup>,
-    color_buf: Option<Rc<wgpu::Buffer>>,
+    //color_bg: Option<Rc<wgpu::BindGroup>>,
+    //color_buf: Option<Rc<wgpu::Buffer>>,
     acc: f32,
 }
 
@@ -46,16 +39,16 @@ impl RenderGraphPass for TrianglePass {
         });
 
         let device = graph.device();
-
         let (color_bgl, color_bg, color_buf, _) = BufferWrapper::builder()
             .label_prefix("color")
             .visibility(wgpu::ShaderStages::FRAGMENT)
             .buffer_usage(wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST)
             .contents(&[glam::Vec4::new(0.902, 0.639, 0.451, 1.0)])
             .finish_parts(device);
-        let color_buf = Rc::new(color_buf);
-        self.color_bg = Some(color_bg);
-        self.color_buf = Some(color_buf.clone());
+        let color_bgl = Rc::new(color_bgl);
+        let color_bg = Rc::new(color_bg);
+
+        //let color_buf = Rc::new(color_buf);
 
         let mut desc = RenderGraphPassDesc::new(
             graph.next_id(),
@@ -63,7 +56,7 @@ impl RenderGraphPass for TrianglePass {
             RenderPassType::Render,
             Some(RenderPipelineDescriptor {
                 label: Some("triangle_pipeline".into()),
-                layouts: vec![color_bgl],
+                layouts: vec![color_bgl.clone()],
                 push_constant_ranges: vec![],
                 vertex: VertexState {
                     module: shader.clone(),
@@ -84,6 +77,7 @@ impl RenderGraphPass for TrianglePass {
                 multisample: wgpu::MultisampleState::default(),
                 multiview: None,
             }),
+            vec![("color_bg", color_bg, Some(color_bgl))],
         );
 
         desc.add_texture_view_slot(
@@ -97,7 +91,7 @@ impl RenderGraphPass for TrianglePass {
             graph.next_id(),
             "color_buffer",
             SlotAttribute::Output,
-            Some(SlotValue::Buffer(color_buf)),
+            Some(SlotValue::Buffer(Rc::new(color_buf))),
         );
 
         desc
@@ -105,7 +99,7 @@ impl RenderGraphPass for TrianglePass {
 
     fn prepare(&mut self, world: &mut lyra_ecs::World, context: &mut RenderGraphContext) {
         const SPEED: f32 = 1.5;
-        
+
         let dt = **world.get_resource::<DeltaTime>();
         self.acc += dt;
         let x = (self.acc * SPEED).sin();
@@ -113,7 +107,7 @@ impl RenderGraphPass for TrianglePass {
         let z = ((self.acc + 5.35) * SPEED).sin();
 
         let color = glam::Vec4::new(x, y, z, 1.0);
-        context.queue.write_buffer(self.color_buf.as_ref().unwrap(), 0, bytemuck::bytes_of(&color));
+        context.queue_buffer_write_with("color_buffer", 0, color);
     }
 
     fn execute(
@@ -131,7 +125,8 @@ impl RenderGraphPass for TrianglePass {
         //context.queue.write_buffer(buffer, offset, data)
 
         //let color_bg = graph.bind_group(graph.slot_id("color_buffer").unwrap());
-        let color_bg = self.color_bg.as_ref().unwrap();
+        //let color_bg = self.color_bg.as_ref().unwrap();
+        let color_bg = graph.bind_group(graph.bind_group_id("color_bg").unwrap());
 
         let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
             label: Some("TrianglePass"),

lyra-game/src/render/resource/render_pipeline.rs

@@ -1,6 +1,6 @@
 use std::{num::NonZeroU32, ops::Deref, rc::Rc};
 
-use wgpu::{BindGroupLayout, PipelineLayout};
+use wgpu::PipelineLayout;
 
 #[derive(Debug, Default, Clone)]
 pub struct VertexBufferLayout {
@@ -44,7 +44,7 @@ pub struct FragmentState {
 //#[derive(Debug, Clone)]
 pub struct RenderPipelineDescriptor {
     pub label: Option<String>,
-    pub layouts: Vec<wgpu::BindGroupLayout>,
+    pub layouts: Vec<Rc<wgpu::BindGroupLayout>>,
     pub push_constant_ranges: Vec<wgpu::PushConstantRange>,
     pub vertex: VertexState,
     pub fragment: Option<FragmentState>,
@@ -57,7 +57,11 @@ pub struct RenderPipelineDescriptor {
 impl RenderPipelineDescriptor {
     /// Create the [`wgpu::PipelineLayout`] for this pipeline
     pub(crate) fn create_layout(&self, device: &wgpu::Device) -> wgpu::PipelineLayout {
-        let bgs = self.layouts.iter().collect::<Vec<_>>();
+        let bgs = self
+            .layouts
+            .iter()
+            .map(|bg| bg.as_ref())
+            .collect::<Vec<_>>();
 
         device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
             label: None, //self.label.as_ref().map(|s| format!("{}Layout", s)),
@@ -65,51 +69,6 @@ impl RenderPipelineDescriptor {
             push_constant_ranges: &self.push_constant_ranges,
         })
     }
-
-    /* fn as_wgpu<'a>(&'a self, device: &wgpu::Device, layout: Option<&'a wgpu::PipelineLayout>) -> wgpu::RenderPipelineDescriptor<'a> {
-        let vbuffers = self.vertex.buffers.iter().map(|vbl| {
-            wgpu::VertexBufferLayout {
-                array_stride: vbl.array_stride,
-                step_mode: vbl.step_mode,
-                attributes: &vbl.attributes
-            }
-        }).collect::<Vec<_>>();
-        let vmodule = device.create_shader_module(wgpu::ShaderModuleDescriptor {
-            label: self.vertex.module.label.as_ref().map(|s| s.as_str()),
-            source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(&self.vertex.module.source)),
-        });
-        let vstate = wgpu::VertexState {
-            module: &vmodule,
-            entry_point: &self.vertex.entry_point,
-            buffers: &vbuffers,
-        };
-
-        let fmodule = self.fragment.as_ref().map(|f| {
-            device.create_shader_module(wgpu::ShaderModuleDescriptor {
-                label: f.module.label.as_ref().map(|s| s.as_str()),
-                source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(&f.module.source)),
-            })
-        });
-
-        let fstate = self.fragment.as_ref().map(move |f| {
-            wgpu::FragmentState {
-                module: fmodule.as_ref().unwrap(),
-                entry_point: &f.entry_point,
-                targets: &f.targets,
-            }
-        });
-
-        wgpu::RenderPipelineDescriptor {
-            label: self.label.as_deref(),
-            layout,
-            vertex: vstate,
-            primitive: self.primitive,
-            depth_stencil: self.depth_stencil,
-            multisample: self.multisample,
-            fragment: fstate,
-            multiview: self.multiview,
-        }
-    } */
 }
 
 pub struct RenderPipeline {