shader reflection for dynamic uniforms

examples/custom_shader.rs

@@ -0,0 +1,111 @@
+use bevy::{
+    prelude::*,
+    render::{
+        render_graph_2::{PipelineDescriptor, StandardMaterial, resource_name, resource_providers::UniformResourceProvider},
+        Shader, ShaderStage, Vertex,
+    },
+};
+
+use bevy_derive::Uniforms;
+
+// #[derive(Uniforms)]
+struct MyMaterial {
+    pub color: Vec4
+}
+
+fn main() {
+    AppBuilder::new()
+        .add_defaults()
+        .setup_world(setup)
+        .setup_render_graph(|builder, pipeline_storage, shader_storage| {
+            builder
+                // .add_resource_provider(UniformResourceProvider::<MyMaterial>::new())
+                .add_pipeline_to_pass(
+                    resource_name::pass::MAIN,
+                    pipeline_storage,
+                    PipelineDescriptor::build(
+                        shader_storage, Shader::from_glsl(
+                            ShaderStage::Vertex,r#"
+                                #version 450
+                                layout(location = 0) in vec4 a_Pos;
+                                layout(location = 0) out vec4 v_Position;
+                                layout(set = 0, binding = 0) uniform Camera {
+                                    mat4 ViewProj;
+                                };
+                                layout(set = 1, binding = 0) uniform Object {
+                                    mat4 Model;
+                                };
+                                void main() {
+                                    v_Position = Model * vec4(a_Pos);
+                                    gl_Position = ViewProj * v_Position;
+                                }
+                            "#),
+                    )
+                    .with_fragment_shader(
+                        Shader::from_glsl(
+                            ShaderStage::Fragment, r#"
+                                #version 450
+                                layout(location = 0) in vec4 v_Position;
+                                layout(location = 0) out vec4 o_Target;
+                                layout(set = 1, binding = 1) uniform MyMaterial_color {
+                                    vec4 color;
+                                };
+                                void main() {
+                                    o_Target = color;
+                                }
+                        "#)
+                    )
+                    .with_depth_stencil_state(wgpu::DepthStencilStateDescriptor {
+                        format: wgpu::TextureFormat::Depth32Float,
+                        depth_write_enabled: true,
+                        depth_compare: wgpu::CompareFunction::Less,
+                        stencil_front: wgpu::StencilStateFaceDescriptor::IGNORE,
+                        stencil_back: wgpu::StencilStateFaceDescriptor::IGNORE,
+                        stencil_read_mask: 0,
+                        stencil_write_mask: 0,
+                    })
+                    .add_color_state(wgpu::ColorStateDescriptor {
+                        format: wgpu::TextureFormat::Bgra8UnormSrgb,
+                        color_blend: wgpu::BlendDescriptor::REPLACE,
+                        alpha_blend: wgpu::BlendDescriptor::REPLACE,
+                        write_mask: wgpu::ColorWrite::ALL,
+                    })
+                    .add_vertex_buffer_descriptor(Vertex::get_vertex_buffer_descriptor())
+                    .add_draw_target(resource_name::draw_target::ASSIGNED_MESHES)
+                    .build(),
+                )
+        })
+        .run();
+}
+
+fn setup(world: &mut World) {
+    let cube_handle = {
+        let mut mesh_storage = world.resources.get_mut::<AssetStorage<Mesh>>().unwrap();
+        mesh_storage.add(Mesh::load(MeshType::Cube))
+    };
+
+    world
+        .build()
+        // red cube
+        .add_archetype(NewMeshEntity {
+            mesh: cube_handle,
+            translation: Translation::new(0.0, 0.0, 1.0),
+            ..NewMeshEntity::default()
+        })
+        // camera
+        .add_archetype(CameraEntity {
+            camera: Camera::new(CameraType::Projection {
+                fov: std::f32::consts::PI / 4.0,
+                near: 1.0,
+                far: 1000.0,
+                aspect_ratio: 1.0,
+            }),
+            active_camera: ActiveCamera,
+            local_to_world: LocalToWorld(Mat4::look_at_rh(
+                Vec3::new(3.0, 8.0, 5.0),
+                Vec3::new(0.0, 0.0, 0.0),
+                Vec3::new(0.0, 0.0, 1.0),
+            )),
+        })
+        .build();
+}

src/app/app_builder.rs

@@ -8,7 +8,7 @@ use crate::{
         passes::*,
         render_graph_2::{
             passes::*, pipelines::*, renderers::wgpu_renderer::WgpuRenderer, resource_providers::*,
-            ShaderPipelineAssignments, StandardMaterial,
+            ShaderPipelineAssignments, StandardMaterial, RenderGraphBuilder
         },
         *,
     },
@@ -181,7 +181,25 @@ impl AppBuilder {
         self
     }
 
-    pub fn add_render_graph_defaults(mut self) -> Self {
+    pub fn add_render_graph_defaults(self) -> Self {
+        self.setup_render_graph(|builder, pipeline_storage, shader_storage| {
+            builder
+                .add_draw_target(resource_name::draw_target::MESHES, meshes_draw_target)
+                .add_draw_target(resource_name::draw_target::ASSIGNED_MESHES, assigned_meshes_draw_target)
+                .add_draw_target(resource_name::draw_target::UI, ui_draw_target)
+                .add_resource_provider(Box::new(CameraResourceProvider))
+                .add_resource_provider(Box::new(Camera2dResourceProvider))
+                .add_resource_provider(Box::new(LightResourceProvider::new(10)))
+                .add_resource_provider(Box::new(UiResourceProvider::new()))
+                .add_resource_provider(Box::new(UniformResourceProvider::<StandardMaterial>::new()))
+                .add_resource_provider(Box::new(UniformResourceProvider::<LocalToWorld>::new()))
+                .add_forward_pass()
+                .add_forward_pipeline(pipeline_storage, shader_storage)
+                .add_ui_pipeline(pipeline_storage, shader_storage)
+        })
+    }
+
+    pub fn setup_render_graph(mut self, setup: impl Fn(RenderGraphBuilder, &mut AssetStorage<PipelineDescriptor>, &mut AssetStorage<Shader>) -> RenderGraphBuilder) -> Self {
         {
             let mut pipeline_storage = self
                 .world
@@ -193,20 +211,7 @@ impl AppBuilder {
                 .resources
                 .get_mut::<AssetStorage<Shader>>()
                 .unwrap();
-            self.render_graph_builder = self
+            self.render_graph_builder = setup(self.render_graph_builder, &mut pipeline_storage, &mut shader_storage);
-                .render_graph_builder
-                .add_draw_target(resource_name::draw_target::MESHES, meshes_draw_target)
-                .add_draw_target(resource_name::draw_target::ASSIGNED_MESHES, assigned_meshes_draw_target)
-                .add_draw_target(resource_name::draw_target::UI, ui_draw_target)
-                .add_resource_provider(Box::new(CameraResourceProvider))
-                .add_resource_provider(Box::new(Camera2dResourceProvider))
-                .add_resource_provider(Box::new(LightResourceProvider::new(10)))
-                .add_resource_provider(Box::new(UiResourceProvider::new()))
-                .add_resource_provider(Box::new(UniformResourceProvider::<StandardMaterial>::new()))
-                .add_resource_provider(Box::new(UniformResourceProvider::<LocalToWorld>::new()))
-                .add_forward_pass()
-                .add_forward_pipeline(&mut pipeline_storage, &mut shader_storage)
-                .add_ui_pipeline(&mut pipeline_storage, &mut shader_storage);
         }
 
         self

src/asset/mod.rs

@@ -12,6 +12,7 @@ pub use texture::*;
 
 use std::{collections::HashMap, marker::PhantomData};
 
+#[derive(Copy)]
 pub struct Handle<T> {
     pub id: usize,
     marker: PhantomData<T>,

src/render/render_graph_2/passes/forward.rs

@@ -27,7 +27,7 @@ impl ForwardPassBuilder for RenderGraphBuilder {
             },
         )))
         .add_pass(
-            "main",
+            resource_name::pass::MAIN,
             PassDescriptor {
                 color_attachments: vec![RenderPassColorAttachmentDescriptor {
                     attachment: resource_name::texture::SWAP_CHAIN.to_string(),

src/render/render_graph_2/pipeline_layout.rs

@@ -38,9 +38,12 @@ impl PipelineLayout {
                }
             }
         }
+        let mut bind_groups_result = bind_groups.drain().map(|(_, value)| value).collect::<Vec<BindGroup>>();
 
+        // NOTE: for some reason bind groups need to be sorted by index. this is likely an issue with bevy and not with wgpu
+        bind_groups_result.sort_by(|a, b| a.index.partial_cmp(&b.index).unwrap());
         PipelineLayout {
-           bind_groups: bind_groups.drain().map(|(_, value)| value).collect()
+           bind_groups: bind_groups_result
         }
     }
 }

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

@@ -1,96 +1,35 @@
-use crate::{asset::AssetStorage, render::{
+use crate::{
-    render_graph_2::{
+    asset::AssetStorage,
-        pipeline_layout::*, PipelineDescriptor, RenderGraphBuilder, resource_name,
+    render::{
+        render_graph_2::{resource_name, PipelineDescriptor, RenderGraphBuilder},
+        shader::{Shader, ShaderStage},
+        Vertex,
     },
-    shader::{Shader, ShaderStage},
+};
-    Vertex,
-}};
 pub trait ForwardPipelineBuilder {
-    fn add_forward_pipeline(self, pipeline_descriptor_storage: &mut AssetStorage<PipelineDescriptor>, shader_storage: &mut AssetStorage<Shader>) -> Self;
+    fn add_forward_pipeline(
+        self,
+        pipeline_descriptor_storage: &mut AssetStorage<PipelineDescriptor>,
+        shader_storage: &mut AssetStorage<Shader>,
+    ) -> Self;
 }
 
 impl ForwardPipelineBuilder for RenderGraphBuilder {
-    fn add_forward_pipeline(self, pipeline_descriptor_storage: &mut AssetStorage<PipelineDescriptor>, shader_storage: &mut AssetStorage<Shader>) -> Self {
+    fn add_forward_pipeline(
+        self,
+        pipeline_descriptor_storage: &mut AssetStorage<PipelineDescriptor>,
+        shader_storage: &mut AssetStorage<Shader>,
+    ) -> Self {
         self.add_pipeline(
             pipeline_descriptor_storage,
             PipelineDescriptor::build(
                 shader_storage,
-                Shader::from_glsl(include_str!("forward.vert"), ShaderStage::Vertex),
+                Shader::from_glsl(ShaderStage::Vertex, include_str!("forward.vert")),
             )
             .with_fragment_shader(Shader::from_glsl(
-                include_str!("forward.frag"),
                 ShaderStage::Fragment,
+                include_str!("forward.frag"),
             ))
-            // .add_bind_group(BindGroup::new(0, vec![
-            //     Binding {
-            //         index: 0,
-            //         name: "Camera".to_string(),
-            //         bind_type: BindType::Uniform {
-            //             dynamic: false,
-            //             properties: vec![UniformProperty {
-            //                 name: "ViewProj".to_string(),
-            //                 property_type: UniformPropertyType::Mat4,
-            //             }],
-            //         },
-            //     },
-            //     Binding {
-            //         index: 1,
-            //         name: "Lights".to_string(),
-            //         bind_type: BindType::Uniform {
-            //             dynamic: false,
-            //             properties: vec![
-            //                 UniformProperty {
-            //                     name: "NumLights".to_string(),
-            //                     property_type: UniformPropertyType::UVec4,
-            //                 },
-            //                 UniformProperty {
-            //                     name: "SceneLights".to_string(),
-            //                     property_type: UniformPropertyType::Array(
-            //                         Box::new(UniformPropertyType::Struct(vec![
-            //                             UniformProperty {
-            //                                 name: "proj".to_string(),
-            //                                 property_type: UniformPropertyType::Mat4,
-            //                             },
-            //                             UniformProperty {
-            //                                 name: "pos".to_string(),
-            //                                 property_type: UniformPropertyType::Vec4,
-            //                             },
-            //                             UniformProperty {
-            //                                 name: "color".to_string(),
-            //                                 property_type: UniformPropertyType::Vec4,
-            //                             },
-            //                         ])),
-            //                         10, // max lights
-            //                     ),
-            //                 },
-            //             ],
-            //         },
-            //     },
-            // ]))
-            // .add_bind_group(BindGroup::new(1, vec![
-            //     Binding {
-            //         index: 0,
-            //         name: "Object".to_string(),
-            //         bind_type: BindType::Uniform {
-            //             dynamic: true,
-            //             properties: vec![UniformProperty {
-            //                 name: "Model".to_string(),
-            //                 property_type: UniformPropertyType::Mat4,
-            //             }],
-            //         },
-            //     },
-            //     Binding {
-            //         index: 1,
-            //         name: "StandardMaterial_albedo".to_string(),
-            //         bind_type: BindType::Uniform {
-            //             dynamic: true,
-            //             properties: vec![UniformProperty {
-            //                 name: "Albedo".to_string(),
-            //                 property_type: UniformPropertyType::Vec4,
-            //             }],
-            //         },
-            //     },
-            // ]))
             .with_rasterization_state(wgpu::RasterizationStateDescriptor {
                 front_face: wgpu::FrontFace::Ccw,
                 cull_mode: wgpu::CullMode::Back,

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

@@ -28,47 +28,12 @@ impl ForwardFlatPipelineBuilder for RenderGraphBuilder {
             pipeline_descriptor_storage,
             PipelineDescriptor::build(
                 shader_storage,
-                Shader::from_glsl(include_str!("forward_flat.vert"), ShaderStage::Vertex),
+                Shader::from_glsl(ShaderStage::Vertex, include_str!("forward_flat.vert")),
             )
             .with_fragment_shader(Shader::from_glsl(
-                include_str!("forward_flat.frag"),
                 ShaderStage::Fragment,
+                include_str!("forward_flat.frag")
             ))
-            .add_bind_group(BindGroup::new(0, vec![Binding {
-                index: 0,
-                name: "Camera".to_string(),
-                bind_type: BindType::Uniform {
-                    dynamic: false,
-                    properties: vec![UniformProperty {
-                        name: "ViewProj".to_string(),
-                        property_type: UniformPropertyType::Mat4,
-                    }],
-                },
-            }]))
-            .add_bind_group(BindGroup::new(1, vec![
-                Binding {
-                    index: 0,
-                    name: "Object".to_string(),
-                    bind_type: BindType::Uniform {
-                        dynamic: true,
-                        properties: vec![UniformProperty {
-                            name: "Model".to_string(),
-                            property_type: UniformPropertyType::Mat4,
-                        }],
-                    },
-                },
-                Binding {
-                    index: 1,
-                    name: "StandardMaterial_albedo".to_string(),
-                    bind_type: BindType::Uniform {
-                        dynamic: true,
-                        properties: vec![UniformProperty {
-                            name: "Albedo".to_string(),
-                            property_type: UniformPropertyType::Vec4,
-                        }],
-                    },
-                },
-            ]))
             .with_rasterization_state(wgpu::RasterizationStateDescriptor {
                 front_face: wgpu::FrontFace::Ccw,
                 cull_mode: wgpu::CullMode::Back,

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

@@ -27,23 +27,12 @@ impl UiPipelineBuilder for RenderGraphBuilder {
             pipeline_descriptor_storage,
             PipelineDescriptor::build(
                 shader_storage,
-                Shader::from_glsl(include_str!("ui.vert"), ShaderStage::Vertex),
+                Shader::from_glsl(ShaderStage::Vertex, include_str!("ui.vert")),
             )
             .with_fragment_shader(Shader::from_glsl(
-                include_str!("ui.frag"),
                 ShaderStage::Fragment,
+                include_str!("ui.frag")
             ))
-            .add_bind_group(BindGroup::new(0, vec![Binding {
-                index: 0,
-                name: "Camera2d".to_string(),
-                bind_type: BindType::Uniform {
-                    dynamic: false,
-                    properties: vec![UniformProperty {
-                        name: "ViewProj".to_string(),
-                        property_type: UniformPropertyType::Mat4,
-                    }],
-                },
-            }]))
             .with_rasterization_state(wgpu::RasterizationStateDescriptor {
                 front_face: wgpu::FrontFace::Ccw,
                 cull_mode: wgpu::CullMode::None,

src/render/render_graph_2/render_graph.rs

@@ -1,7 +1,7 @@
 use crate::{
     asset::{AssetStorage, Handle},
-    render::{
+    render::render_graph_2::{
-        render_graph_2::{PassDescriptor, PipelineDescriptor, ResourceProvider, TextureDescriptor, DrawTarget},
+        DrawTarget, PassDescriptor, PipelineDescriptor, ResourceProvider, TextureDescriptor,
     },
 };
 use std::collections::{HashMap, HashSet};
@@ -45,7 +45,6 @@ impl RenderGraph {
 pub struct RenderGraphBuilder {
     render_graph: RenderGraph,
     current_pass: Option<String>,
-
 }
 
 impl RenderGraphBuilder {
@@ -64,15 +63,33 @@ impl RenderGraphBuilder {
         self
     }
 
-    pub fn add_pipeline(mut self, pipeline_descriptor_storage: &mut AssetStorage<PipelineDescriptor>, pipeline: PipelineDescriptor) -> Self {
+    pub fn add_pipeline(
+        mut self,
+        pipeline_descriptor_storage: &mut AssetStorage<PipelineDescriptor>,
+        pipeline: PipelineDescriptor,
+    ) -> Self {
         if let Some(ref pass) = self.current_pass {
             let pipeline_descriptor_handle = pipeline_descriptor_storage.add(pipeline);
-            self.render_graph.add_pipeline(&pass, pipeline_descriptor_handle);
+            self.render_graph
+                .add_pipeline(&pass, pipeline_descriptor_handle);
         }
 
         self
     }
 
+    pub fn add_pipeline_to_pass(
+        mut self,
+        pass: &str,
+        pipeline_descriptor_storage: &mut AssetStorage<PipelineDescriptor>,
+        pipeline: PipelineDescriptor,
+    ) -> Self {
+        let pipeline_descriptor_handle = pipeline_descriptor_storage.add(pipeline);
+        self.render_graph
+            .add_pipeline(pass, pipeline_descriptor_handle);
+
+        self
+    }
+
     pub fn add_resource_provider(mut self, resource_provider: Box<dyn ResourceProvider>) -> Self {
         self.render_graph.resource_providers.push(resource_provider);
         self
@@ -86,7 +103,9 @@ impl RenderGraphBuilder {
     }
 
     pub fn add_draw_target(mut self, name: &str, draw_target: DrawTarget) -> Self {
-        self.render_graph.draw_targets.insert(name.to_string(), draw_target);
+        self.render_graph
+            .draw_targets
+            .insert(name.to_string(), draw_target);
         self
     }
 

src/render/render_graph_2/renderers/wgpu_renderer.rs

@@ -70,6 +70,7 @@ impl WgpuRenderer {
     }
 
     pub fn create_render_pipeline(
+        dynamic_uniform_buffer_info: &HashMap<String, DynamicUniformBufferInfo>,
         pipeline_descriptor: &mut PipelineDescriptor,
         bind_group_layouts: &mut HashMap<u64, wgpu::BindGroupLayout>,
         device: &wgpu::Device,
@@ -92,11 +93,35 @@ impl WgpuRenderer {
                 layouts.push(fragment_spirv.reflect_layout().unwrap());
             }
 
-            pipeline_descriptor.layout =
+            let mut layout = PipelineLayout::from_shader_layouts(&mut layouts);
-                PipelineLayoutType::Reflected(Some(PipelineLayout::from_shader_layouts(&mut layouts)));
+
+            // set each uniform binding to dynamic if there is a matching dynamic uniform buffer info
+            for mut bind_group in layout.bind_groups.iter_mut() {
+                bind_group.bindings = bind_group
+                    .bindings
+                    .iter()
+                    .cloned()
+                    .map(|mut binding| {
+                        if let BindType::Uniform {
+                            ref mut dynamic, ..
+                        } = binding.bind_type
+                        {
+                            if dynamic_uniform_buffer_info.contains_key(&binding.name) {
+                                *dynamic = true;
+                            }
+                        }
+
+                        binding
+                    })
+                    .collect();
+            }
+
+            pipeline_descriptor.layout = PipelineLayoutType::Reflected(Some(layout));
         }
 
         let layout = pipeline_descriptor.get_layout_mut().unwrap();
+        // println!("{:#?}", layout);
+        // println!();
 
         // setup new bind group layouts
         for bind_group in layout.bind_groups.iter_mut() {
@@ -326,6 +351,24 @@ impl WgpuRenderer {
     ) -> wgpu::ShaderModule {
         device.create_shader_module(&shader.get_spirv(macros))
     }
+
+    pub fn initialize_resource_providers(
+        &mut self,
+        world: &mut World,
+        render_graph: &mut RenderGraph,
+    ) {
+        self.encoder = Some(
+            self.device
+                .create_command_encoder(&wgpu::CommandEncoderDescriptor { todo: 0 }),
+        );
+        for resource_provider in render_graph.resource_providers.iter_mut() {
+            resource_provider.initialize(self, world);
+        }
+
+        // consume current encoder
+        let command_buffer = self.encoder.take().unwrap().finish();
+        self.queue.submit(&[command_buffer]);
+    }
 }
 
 impl Renderer for WgpuRenderer {
@@ -338,9 +381,8 @@ impl Renderer for WgpuRenderer {
         };
 
         self.surface = Some(surface);
-        for resource_provider in render_graph.resource_providers.iter_mut() {
+
-            resource_provider.initialize(self, world);
+        self.initialize_resource_providers(world, render_graph);
-        }
 
         self.resize(world, render_graph, window_size.width, window_size.height);
     }
@@ -425,6 +467,7 @@ impl Renderer for WgpuRenderer {
                     .as_ref()
                     .map(|handle| &*shader_storage.get(&handle).unwrap());
                 let render_pipeline = WgpuRenderer::create_render_pipeline(
+                    &self.dynamic_uniform_buffer_info,
                     pipeline_descriptor,
                     &mut self.bind_group_layouts,
                     &self.device,

src/render/render_graph_2/resource_name.rs

@@ -19,4 +19,8 @@ pub mod draw_target {
     pub const MESHES: &str = "Meshes";
     pub const ASSIGNED_MESHES: &str = "AssignedMeshes";
     pub const UI: &str = "Ui";
+}
+
+pub mod pass {
+    pub const MAIN: &str = "Main";
 }

src/render/render_graph_2/resource_providers/uniform_resource_provider.rs

@@ -31,7 +31,9 @@ impl<T> ResourceProvider for UniformResourceProvider<T>
 where
     T: AsUniforms + Send + Sync + 'static,
 {
-    fn initialize(&mut self, _renderer: &mut dyn Renderer, _world: &mut World) {}
+    fn initialize(&mut self, renderer: &mut dyn Renderer, world: &mut World) {
+        self.update(renderer, world);
+    }
 
     fn update(&mut self, renderer: &mut dyn Renderer, world: &mut World) {
         let query = <(Read<T>, Read<Renderable>)>::query();

src/render/shader.rs

@@ -61,7 +61,7 @@ pub struct Shader {
 }
 
 impl Shader {
-    pub fn from_glsl(glsl: &str, stage: ShaderStage) -> Shader {
+    pub fn from_glsl(stage: ShaderStage, glsl: &str) -> Shader {
         Shader {
             source: ShaderSource::Glsl(glsl.to_string()),
             stage,