use buffers + dynamic offsets for shaderuniforms

examples/simple_new_graph.rs

@@ -14,6 +14,14 @@ fn setup(world: &mut World) {
         (mesh_storage.add(cube), mesh_storage.add(plane))
     };
 
+    let mut indices = std::collections::HashMap::new();
+    indices.insert("StandardMaterial".to_string(), 0);
+    indices.insert("Object".to_string(), 0);
+
+    let mut indices_2 = std::collections::HashMap::new();
+    indices_2.insert("StandardMaterial".to_string(), 16);
+    indices_2.insert("Object".to_string(), 64);
+
     world.build()
         // plane
         // .add_archetype(MeshEntity {
@@ -32,25 +40,27 @@ fn setup(world: &mut World) {
                 uniform_selectors: vec![
                    uniform_selector::<StandardMaterial>, 
                    uniform_selector::<LocalToWorld>, 
-                ]
+                ],
+                dynamic_uniform_indices: indices,
             },
             local_to_world: LocalToWorld::identity(),
             translation: Translation::new(0.0, 0.0, 1.0),
         })
-        // .add_archetype(NewMeshEntity {
+        .add_archetype(NewMeshEntity {
-        //     mesh: cube_handle.clone(),
+            mesh: cube_handle.clone(),
-        //     material: StandardMaterial {
+            material: StandardMaterial {
-        //         albedo: math::vec4(0.0, 1.0, 0.0, 1.0),
+                albedo: math::vec4(0.0, 1.0, 0.0, 1.0),
-        //     },
+            },
-        //     shader_uniforms: ShaderUniforms {
+            shader_uniforms: ShaderUniforms {
-        //         uniform_selectors: vec![
+                uniform_selectors: vec![
-        //            uniform_selector::<StandardMaterial>, 
+                   uniform_selector::<StandardMaterial>, 
-        //            uniform_selector::<LocalToWorld>, 
+                   uniform_selector::<LocalToWorld>, 
-        //         ]
+                ],
-        //     },
+                dynamic_uniform_indices: indices_2,
-        //     local_to_world: LocalToWorld::identity(),
+            },
-        //     translation: Translation::new(-2.0, 0.0, 1.0),
+            local_to_world: LocalToWorld::identity(),
-        // })
+            translation: Translation::new(-2.0, 0.0, 1.0),
+        })
         // light
         // .add_archetype(LightEntity {
         //     light: Light {

src/render/render_graph_2/draw_target.rs

@@ -19,7 +19,7 @@ pub fn mesh_draw_target(world: &World, render_pass: &mut dyn RenderPass) {
     let mut last_mesh_id = None;
     let mesh_query =
         <(Read<ShaderUniforms>, Read<Handle<Mesh>>)>::query().filter(!component::<Instanced>());
-    for (_shader_uniforms, mesh) in mesh_query.iter(world) {
+    for (shader_uniforms, mesh) in mesh_query.iter(world) {
         let current_mesh_id = mesh.id;
 
         let mut should_load_mesh = last_mesh_id == None;
@@ -42,7 +42,7 @@ pub fn mesh_draw_target(world: &World, render_pass: &mut dyn RenderPass) {
         // TODO: re-getting the mesh isn't necessary. just store the index count
         if let Some(mesh_asset) = mesh_storage.get(mesh.id) {
             // TODO: validate bind group properties against shader uniform properties at least once
-            render_pass.setup_bind_groups();
+            render_pass.setup_bind_groups(&&*shader_uniforms);
             render_pass.draw_indexed(0..mesh_asset.indices.len() as u32, 0, 0..1);
         };
 

src/render/render_graph_2/pipelines/forward/mod.rs

@@ -45,7 +45,7 @@ impl ForwardPipelineBuilder for RenderGraphBuilder {
                     Binding {
                         name: "Object".to_string(),
                         bind_type: BindType::Uniform {
-                            dynamic: false,
+                            dynamic: true,
                             properties: vec![
                                 UniformProperty {
                                     name: "Model".to_string(),
@@ -57,7 +57,7 @@ impl ForwardPipelineBuilder for RenderGraphBuilder {
                     Binding {
                         name: "StandardMaterial".to_string(),
                         bind_type: BindType::Uniform {
-                            dynamic: false,
+                            dynamic: true,
                             properties: vec![
                                 UniformProperty {
                                     name: "Albedo".to_string(),

src/render/render_graph_2/renderer.rs

@@ -1,4 +1,4 @@
-use crate::{legion::prelude::*, render::render_graph_2::{RenderGraph, ResourceInfo, PipelineDescriptor}};
+use crate::{legion::prelude::*, render::render_graph_2::{RenderGraph, ResourceInfo, PipelineDescriptor, ShaderUniforms}};
 use std::ops::Range;
 
 pub trait Renderer {
@@ -19,5 +19,5 @@ pub trait RenderPass {
     fn set_index_buffer(&mut self, name: &str, offset: u64);
     fn set_vertex_buffer(&mut self, start_slot: u32, name: &str, offset: u64);
     fn draw_indexed(&mut self, indices: Range<u32>, base_vertex: i32, instances: Range<u32>);
-    fn setup_bind_groups(&mut self);
+    fn setup_bind_groups(&mut self, shader_uniforms: &ShaderUniforms);
 }

src/render/render_graph_2/resource_provider.rs

@@ -27,7 +27,6 @@ impl ResourceProvider for CameraResourceProvider {
             camera.update(width, height);
             let camera_matrix: [[f32; 4]; 4] =
                 (camera.view_matrix * local_to_world.0).to_cols_array_2d();
-            // TODO: use staging buffer here
             renderer.create_buffer_with_data(
                 resource_name::uniform::CAMERA,
                 camera_matrix.as_bytes(),

src/render/render_graph_2/shader.rs

@@ -3,28 +3,77 @@ use crate::{
         borrow::RefMap,
         prelude::{Entity, World},
     },
-    render::render_graph_2::{UniformPropertyType, BindType},
     math::Vec4,
+    render::render_graph_2::{BindType, UniformPropertyType},
 };
-use zerocopy::AsBytes;
+use std::collections::HashMap;
 use legion::storage::Component;
+use zerocopy::AsBytes;
 
 pub type ShaderUniformSelector = fn(Entity, &World) -> Option<RefMap<&dyn AsUniforms>>;
 pub struct ShaderUniforms {
     // used for distinguishing
     pub uniform_selectors: Vec<ShaderUniformSelector>,
+    pub dynamic_uniform_indices: HashMap<String, u64>,
 }
 
-impl<'a> ShaderUniforms {
+impl ShaderUniforms {
     pub fn new() -> Self {
         ShaderUniforms {
             uniform_selectors: Vec::new(),
+            dynamic_uniform_indices: HashMap::new(),
         }
     }
 
     pub fn add(&mut self, selector: ShaderUniformSelector) {
         self.uniform_selectors.push(selector);
     }
+
+    pub fn get_uniform_info<'a>(
+        &'a self,
+        world: &'a World,
+        entity: Entity,
+        uniform_name: &str,
+    ) -> Option<&'a UniformInfo> {
+        for uniform_selector in self.uniform_selectors.iter().rev() {
+            let uniforms = uniform_selector(entity, world).unwrap_or_else(|| {
+                panic!(
+                    "ShaderUniform selector points to a missing component. Uniform: {}",
+                    uniform_name
+                )
+            });
+
+            let info = uniforms.get_uniform_info(uniform_name);
+            if let Some(_) = info {
+              return info;
+            }
+        }
+
+        None
+    }
+
+    pub fn get_uniform_bytes<'a>(
+        &'a self,
+        world: &'a World,
+        entity: Entity,
+        uniform_name: &str,
+    ) -> Option<Vec<u8>> {
+        for uniform_selector in self.uniform_selectors.iter().rev() {
+            let uniforms = uniform_selector(entity, world).unwrap_or_else(|| {
+                panic!(
+                    "ShaderUniform selector points to a missing component. Uniform: {}",
+                    uniform_name
+                )
+            });
+
+            let bytes = uniforms.get_uniform_bytes(uniform_name);
+            if let Some(_) = bytes {
+              return bytes;
+            }
+        }
+
+        None
+    }
 }
 
 pub struct StandardMaterial {
@@ -59,50 +108,49 @@ impl GetBytes for Vec4 {
 }
 
 pub trait AsUniforms {
-    fn get_uniform_info(&self) -> &[UniformInfo];
+    fn get_uniform_infos(&self) -> &[UniformInfo];
+    fn get_uniform_info(&self, name: &str) -> Option<&UniformInfo>;
     fn get_uniform_layouts(&self) -> &[&[UniformPropertyType]];
     fn get_uniform_bytes(&self, name: &str) -> Option<Vec<u8>>;
     // TODO: support zero-copy uniforms
-    // fn get_uniform_value_ref(&self, index: usize) -> &[u8];
+    // fn get_uniform_bytes_ref(&self, name: &str) -> Option<&[u8]>;
 }
 
 // pub struct UniformInfo<'a> {
 //   pub name: &'a str,
-//   pub 
+//   pub
 // }
 
 pub struct UniformInfo<'a> {
-  pub name: &'a str,
+    pub name: &'a str,
-  pub bind_type: BindType,
+    pub bind_type: BindType,
 }
 
-pub fn uniform_selector<T>(entity: Entity, world: &World) -> Option<RefMap<&dyn AsUniforms>> where T: AsUniforms + Component {
+pub fn uniform_selector<T>(entity: Entity, world: &World) -> Option<RefMap<&dyn AsUniforms>>
-    world.get_component::<T>(entity).map(
+where
-        |c| {
+    T: AsUniforms + Component,
-        c.map_into(|s| {
+{
-            s as &dyn AsUniforms
+    world
-        })
+        .get_component::<T>(entity)
-    })
+        .map(|c| c.map_into(|s| s as &dyn AsUniforms))
 }
 
 // create this from a derive macro
-const STANDARD_MATERIAL_UNIFORM_INFO: &[UniformInfo] = &[
+const STANDARD_MATERIAL_UNIFORM_INFO: &[UniformInfo] = &[UniformInfo {
-  UniformInfo {
     name: "StandardMaterial",
     bind_type: BindType::Uniform {
-      dynamic: false,
+        dynamic: false,
-      // TODO: fill this in with properties
+        // TODO: fill this in with properties
-      properties: Vec::new()
+        properties: Vec::new(),
     },
-  }
+}];
-];
 
 // these are separate from BindType::Uniform{properties} because they need to be const
 const STANDARD_MATERIAL_UNIFORM_LAYOUTS: &[&[UniformPropertyType]] = &[&[]];
 
 // const ST
 impl AsUniforms for StandardMaterial {
-    fn get_uniform_info(&self) -> &[UniformInfo] {
+    fn get_uniform_infos(&self) -> &[UniformInfo] {
         STANDARD_MATERIAL_UNIFORM_INFO
     }
 
@@ -111,11 +159,18 @@ impl AsUniforms for StandardMaterial {
     }
 
     fn get_uniform_bytes(&self, name: &str) -> Option<Vec<u8>> {
-      match name {
+        match name {
-        "StandardMaterial" => Some(self.albedo.get_bytes()),
+            "StandardMaterial" => Some(self.albedo.get_bytes()),
-        _ => None,
+            _ => None,
-      }
+        }
     }
+    fn get_uniform_info(&self, name: &str) -> Option<&UniformInfo> {
+        match name {
+            "StandardMaterial" => Some(&STANDARD_MATERIAL_UNIFORM_INFO[0]),
+            _ => None,
+        }
+    }
+
     // fn iter_properties(&self) -> std::slice::Iter<&'static str>  {
     //   STANDARD_MATERIAL_PROPERTIES.iter()
     // }
@@ -128,37 +183,24 @@ impl AsUniforms for StandardMaterial {
     //     _ => None,
     //   }
     // }
-    // fn get_selector(&self) -> ShaderMaterialSelector {
-    //   |entity, world| {
-    //     world.get_component::<Self>(entity).map(
-    //       |c: Ref<StandardMaterial>| {
-    //         c.map_into(|s| {
-    //           s as &dyn ShaderMaterial
-    //         })
-    //       }
-    //     )
-    //   }
-    // }
 }
 
 // create this from a derive macro
-const LOCAL_TO_WORLD_UNIFORM_INFO: &[UniformInfo] = &[
+const LOCAL_TO_WORLD_UNIFORM_INFO: &[UniformInfo] = &[UniformInfo {
-  UniformInfo {
     name: "Object",
     bind_type: BindType::Uniform {
-      dynamic: false,
+        dynamic: false,
-      // TODO: fill this in with properties
+        // TODO: fill this in with properties
-      properties: Vec::new()
+        properties: Vec::new(),
     },
-  }
+}];
-];
 
 // these are separate from BindType::Uniform{properties} because they need to be const
 const LOCAL_TO_WORLD_UNIFORM_LAYOUTS: &[&[UniformPropertyType]] = &[&[]];
 
 // const ST
 impl AsUniforms for bevy_transform::prelude::LocalToWorld {
-    fn get_uniform_info(&self) -> &[UniformInfo] {
+    fn get_uniform_infos(&self) -> &[UniformInfo] {
         LOCAL_TO_WORLD_UNIFORM_INFO
     }
 
@@ -167,10 +209,16 @@ impl AsUniforms for bevy_transform::prelude::LocalToWorld {
     }
 
     fn get_uniform_bytes(&self, name: &str) -> Option<Vec<u8>> {
-      match name {
+        match name {
-        "Object" => Some(self.0.to_cols_array_2d().as_bytes().into()),
+            "Object" => Some(self.0.to_cols_array_2d().as_bytes().into()),
-        _ => None,
+            _ => None,
-      }
+        }
+    }
+    fn get_uniform_info(&self, name: &str) -> Option<&UniformInfo> {
+        match name {
+            "Object" => Some(&LOCAL_TO_WORLD_UNIFORM_INFO[0]),
+            _ => None,
+        }
     }
     // fn iter_properties(&self) -> std::slice::Iter<&'static str>  {
     //   STANDARD_MATERIAL_PROPERTIES.iter()
@@ -184,15 +232,4 @@ impl AsUniforms for bevy_transform::prelude::LocalToWorld {
     //     _ => None,
     //   }
     // }
-    // fn get_selector(&self) -> ShaderMaterialSelector {
-    //   |entity, world| {
-    //     world.get_component::<Self>(entity).map(
-    //       |c: Ref<StandardMaterial>| {
-    //         c.map_into(|s| {
-    //           s as &dyn ShaderMaterial
-    //         })
-    //       }
-    //     )
-    //   }
-    // }
 }

src/render/render_graph_2/wgpu_renderer.rs

@@ -350,6 +350,71 @@ impl WgpuRenderer {
             }
         }
     }
+
+    fn setup_dynamic_entity_shader_uniforms(&mut self, world: &World, render_graph: &RenderGraph, encoder: &mut wgpu::CommandEncoder) {
+        let mut dynamic_uniform_info = HashMap::new();
+
+        // retrieve all uniforms buffers that aren't aleady set. these are "dynamic" uniforms, which are set by the user in ShaderUniforms
+        // TODO: this breaks down in multiple ways:
+        // (1) resource_info will be set after the first run so this won't update.
+        // (2) if we create new buffers, the old bind groups will be invalid
+        for pipeline in render_graph.pipeline_descriptors.values() {
+            for bind_group in pipeline.pipeline_layout.bind_groups.iter() {
+                for binding in bind_group.bindings.iter() {
+                    if let None = self.resource_info.get(&binding.name) {
+                        if let BindType::Uniform { .. } = &binding.bind_type {
+                            if dynamic_uniform_info.contains_key(&binding.name) {
+                                continue;
+                            }
+
+                            dynamic_uniform_info.insert(binding.name.to_string(), UniformInfo {
+                                size: binding.bind_type.get_uniform_size().unwrap(),
+                                count: 0,
+                            });
+                        }
+                    }
+                }
+            }
+        }
+
+        // count the number of entities providing each uniform 
+        for (name, info) in dynamic_uniform_info.iter_mut() {
+            for (entity, shader_uniforms) in <Read<ShaderUniforms>>::query().iter_entities(world) {
+                if let Some(_) = shader_uniforms.get_uniform_info(world, entity, name) {
+                    info.count += 1;
+                }
+            }
+        }
+        
+        // allocate uniform buffers
+        for (name, info) in dynamic_uniform_info.iter() {
+            // TODO: maybe align to device
+            let size = info.size * info.count;
+            println!("{} {} {}", name, info.size, info.count);
+            self.create_buffer(name, size, wgpu::BufferUsage::COPY_DST | wgpu::BufferUsage::UNIFORM);
+        }
+
+        // copy entity uniform data to buffers
+        for (name, info) in dynamic_uniform_info.iter_mut() {
+            let size = info.size * info.count;
+            let mut current_index = 0;
+            let mapped = self.device.create_buffer_mapped(size as usize, wgpu::BufferUsage::COPY_SRC);
+            for ((entity, shader_uniforms), slot) in <Read<ShaderUniforms>>::query().iter_entities(world).zip(mapped.data.chunks_exact_mut(info.size as usize)) {
+                if let Some(bytes) = shader_uniforms.get_uniform_bytes(world, entity, name) {
+                    slot.copy_from_slice(bytes.as_slice());
+                }
+            }
+
+            let temp_buffer = mapped.finish();
+            let uniform_buffer = self.buffers.get(name);
+            encoder.copy_buffer_to_buffer(&temp_buffer, 0, uniform_buffer.unwrap(), 0, size);
+        }
+    }
+}
+
+pub struct UniformInfo {
+    pub size: u64,
+    pub count: u64,
 }
 
 impl Renderer for WgpuRenderer {
@@ -401,6 +466,8 @@ impl Renderer for WgpuRenderer {
         let mut encoder = self
             .device
             .create_command_encoder(&wgpu::CommandEncoderDescriptor { todo: 0 });
+        
+        self.setup_dynamic_entity_shader_uniforms(world, render_graph, &mut encoder);
 
         // setup, pipelines, bind groups, and resources
         for (pipeline_name, pipeline_descriptor) in render_graph.pipeline_descriptors.iter() {
@@ -420,16 +487,16 @@ impl Renderer for WgpuRenderer {
                 self.setup_bind_group(bind_group);
                 // TODO: Move this out of the for loop
                 // copy entity ShaderUniforms to buffers
-                let shader_uniform_query = <Read<ShaderUniforms>>::query();
+                // let shader_uniform_query = <Read<ShaderUniforms>>::query();
-                for (entity, shader_uniforms) in shader_uniform_query.iter_entities(world) {
+                // for (entity, shader_uniforms) in shader_uniform_query.iter_entities(world) {
-                    self.setup_entity_shader_uniforms(
+                //     self.setup_entity_shader_uniforms(
-                        bind_group,
+                //         bind_group,
-                        world,
+                //         world,
-                        entity,
+                //         entity,
-                        &&*shader_uniforms,
+                //         &&*shader_uniforms,
-                        &mut encoder,
+                //         &mut encoder,
-                    );
+                //     );
-                }
+                // }
             }
         }
 
@@ -543,7 +610,7 @@ impl<'a, 'b, 'c, 'd> RenderPass for WgpuRenderPass<'a, 'b, 'c, 'd> {
     }
 
     // TODO: maybe move setup to renderer.setup_bind_groups(&pipeline_desc);
-    fn setup_bind_groups(&mut self) {
+    fn setup_bind_groups(&mut self, shader_uniforms: &ShaderUniforms) {
         for (i, bind_group) in self
             .pipeline_descriptor
             .pipeline_layout
@@ -556,8 +623,22 @@ impl<'a, 'b, 'c, 'd> RenderPass for WgpuRenderPass<'a, 'b, 'c, 'd> {
             bind_group.hash(&mut hasher);
             let bind_group_id = hasher.finish();
             let bind_group_info = self.renderer.bind_groups.get(&bind_group_id).unwrap();
+
+            let mut dynamic_uniform_indices = Vec::new();
+            for binding in bind_group.bindings.iter() {
+                if let BindType::Uniform { dynamic, ..} = binding.bind_type {
+                    if !dynamic {
+                        continue;
+                    }
+
+                    if let Some(index) = shader_uniforms.dynamic_uniform_indices.get(&binding.name) {
+                        dynamic_uniform_indices.push(*index);
+                    }
+                }
+            } 
+
             self.render_pass
-                .set_bind_group(i as u32, &bind_group_info.bind_group, &[]);
+                .set_bind_group(i as u32, &bind_group_info.bind_group, dynamic_uniform_indices.as_slice());
         }
     }
 }