Optimize Asset Gpu Data Transfer

crates/bevy_render/src/draw.rs

@@ -1,12 +1,12 @@
 use crate::{
     pipeline::{PipelineCompiler, PipelineDescriptor, PipelineLayout, PipelineSpecialization},
     renderer::{
-        BindGroup, BindGroupId, BufferId, RenderResource, RenderResourceBinding,
-        RenderResourceBindings, RenderResourceContext, SharedBuffers,
+        AssetRenderResourceBindings, BindGroup, BindGroupId, BufferId, RenderResource,
+        RenderResourceBinding, RenderResourceBindings, RenderResourceContext, SharedBuffers,
     },
     shader::Shader,
 };
-use bevy_asset::{Assets, Handle};
+use bevy_asset::{Asset, Assets, Handle};
 use bevy_ecs::{Query, Res, ResMut, SystemParam};
 use bevy_reflect::Reflect;
 use std::{ops::Range, sync::Arc};
@@ -117,12 +117,15 @@ pub enum DrawError {
     PipelineHasNoLayout,
     #[error("failed to get a buffer for the given `RenderResource`")]
     BufferAllocationFailure,
+    #[error("the given asset does not have any render resources")]
+    MissingAssetRenderResources,
 }
 
 #[derive(SystemParam)]
 pub struct DrawContext<'a> {
     pub pipelines: ResMut<'a, Assets<PipelineDescriptor>>,
     pub shaders: ResMut<'a, Assets<Shader>>,
+    pub asset_render_resource_bindings: ResMut<'a, AssetRenderResourceBindings>,
     pub pipeline_compiler: ResMut<'a, PipelineCompiler>,
     pub render_resource_context: Res<'a, Box<dyn RenderResourceContext>>,
     pub shared_buffers: ResMut<'a, SharedBuffers>,
@@ -184,32 +187,91 @@ impl<'a> DrawContext<'a> {
         })
     }
 
+    pub fn set_asset_bind_groups<T: Asset>(
+        &mut self,
+        draw: &mut Draw,
+        asset_handle: &Handle<T>,
+    ) -> Result<(), DrawError> {
+        if let Some(asset_bindings) = self
+            .asset_render_resource_bindings
+            .get_mut_untyped(&asset_handle.clone_weak_untyped())
+        {
+            Self::set_bind_groups_from_bindings_internal(
+                &self.current_pipeline,
+                &self.pipelines,
+                &**self.render_resource_context,
+                None,
+                draw,
+                &mut [asset_bindings],
+            )
+        } else {
+            Err(DrawError::MissingAssetRenderResources)
+        }
+    }
+
     pub fn set_bind_groups_from_bindings(
-        &self,
+        &mut self,
         draw: &mut Draw,
         render_resource_bindings: &mut [&mut RenderResourceBindings],
     ) -> Result<(), DrawError> {
-        let pipeline = self
-            .current_pipeline
-            .as_ref()
-            .ok_or(DrawError::NoPipelineSet)?;
-        let pipeline_descriptor = self
-            .pipelines
+        Self::set_bind_groups_from_bindings_internal(
+            &self.current_pipeline,
+            &self.pipelines,
+            &**self.render_resource_context,
+            Some(&mut self.asset_render_resource_bindings),
+            draw,
+            render_resource_bindings,
+        )
+    }
+
+    fn set_bind_groups_from_bindings_internal(
+        current_pipeline: &Option<Handle<PipelineDescriptor>>,
+        pipelines: &Assets<PipelineDescriptor>,
+        render_resource_context: &dyn RenderResourceContext,
+        mut asset_render_resource_bindings: Option<&mut AssetRenderResourceBindings>,
+        draw: &mut Draw,
+        render_resource_bindings: &mut [&mut RenderResourceBindings],
+    ) -> Result<(), DrawError> {
+        let pipeline = current_pipeline.as_ref().ok_or(DrawError::NoPipelineSet)?;
+        let pipeline_descriptor = pipelines
             .get(pipeline)
             .ok_or(DrawError::NonExistentPipeline)?;
         let layout = pipeline_descriptor
             .get_layout()
             .ok_or(DrawError::PipelineHasNoLayout)?;
-        for bindings in render_resource_bindings.iter_mut() {
-            bindings.update_bind_groups(pipeline_descriptor, &**self.render_resource_context);
-        }
-        for bind_group_descriptor in layout.bind_groups.iter() {
+        'bind_group_descriptors: for bind_group_descriptor in layout.bind_groups.iter() {
             for bindings in render_resource_bindings.iter_mut() {
                 if let Some(bind_group) =
-                    bindings.get_descriptor_bind_group(bind_group_descriptor.id)
+                    bindings.update_bind_group(bind_group_descriptor, render_resource_context)
                 {
                     draw.set_bind_group(bind_group_descriptor.index, bind_group);
-                    break;
+                    continue 'bind_group_descriptors;
+                }
+            }
+
+            // if none of the given RenderResourceBindings have the current bind group, try their assets
+            let asset_render_resource_bindings =
+                if let Some(value) = asset_render_resource_bindings.as_mut() {
+                    value
+                } else {
+                    continue 'bind_group_descriptors;
+                };
+            for bindings in render_resource_bindings.iter_mut() {
+                for (asset_handle, _) in bindings.iter_assets() {
+                    let asset_bindings = if let Some(asset_bindings) =
+                        asset_render_resource_bindings.get_mut_untyped(asset_handle)
+                    {
+                        asset_bindings
+                    } else {
+                        continue;
+                    };
+
+                    if let Some(bind_group) = asset_bindings
+                        .update_bind_group(bind_group_descriptor, render_resource_context)
+                    {
+                        draw.set_bind_group(bind_group_descriptor.index, bind_group);
+                        continue 'bind_group_descriptors;
+                    }
                 }
             }
         }

crates/bevy_render/src/pipeline/render_pipelines.rs

@@ -114,6 +114,19 @@ pub fn draw_render_pipelines_system(
                     .collect::<HashSet<String>>();
                 pipeline.dynamic_bindings_generation =
                     render_pipelines.bindings.dynamic_bindings_generation();
+                for (handle, _) in render_pipelines.bindings.iter_assets() {
+                    if let Some(bindings) = draw_context
+                        .asset_render_resource_bindings
+                        .get_untyped(handle)
+                    {
+                        for binding in bindings.iter_dynamic_bindings() {
+                            pipeline
+                                .specialization
+                                .dynamic_bindings
+                                .insert(binding.to_string());
+                        }
+                    }
+                }
             }
         }
 

crates/bevy_render/src/render_graph/nodes/render_resources_node.rs

@@ -9,14 +9,15 @@ use crate::{
     texture,
 };
 
-use bevy_asset::{Asset, Assets, Handle, HandleId};
+use bevy_app::{EventReader, Events};
+use bevy_asset::{Asset, AssetEvent, Assets, Handle, HandleId};
 use bevy_ecs::{
     Changed, Commands, Entity, IntoSystem, Local, Query, QuerySet, Res, ResMut, Resources, System,
-    World,
+    With, World,
 };
 use bevy_utils::HashMap;
 use renderer::{AssetRenderResourceBindings, BufferId, RenderResourceType, RenderResources};
-use std::{hash::Hash, marker::PhantomData, ops::DerefMut};
+use std::{any::TypeId, hash::Hash, marker::PhantomData, ops::DerefMut};
 
 #[derive(Debug)]
 struct QueuedBufferWrite {
@@ -562,8 +563,6 @@ where
     }
 }
 
-const EXPECT_ASSET_MESSAGE: &str = "Only assets that exist should be in the modified assets list";
-
 impl<T> SystemNode for AssetRenderResourcesNode<T>
 where
     T: renderer::RenderResources + Asset,
@@ -583,35 +582,75 @@ where
     }
 }
 
+struct AssetRenderNodeState<T: Asset> {
+    event_reader: EventReader<AssetEvent<T>>,
+}
+
+impl<T: Asset> Default for AssetRenderNodeState<T> {
+    fn default() -> Self {
+        Self {
+            event_reader: Default::default(),
+        }
+    }
+}
+
+#[allow(clippy::clippy::too_many_arguments)]
 fn asset_render_resources_node_system<T: RenderResources + Asset>(
     mut state: Local<RenderResourcesNodeState<HandleId, T>>,
+    mut asset_state: Local<AssetRenderNodeState<T>>,
     assets: Res<Assets<T>>,
+    asset_events: Res<Events<AssetEvent<T>>>,
     mut asset_render_resource_bindings: ResMut<AssetRenderResourceBindings>,
     render_resource_context: Res<Box<dyn RenderResourceContext>>,
-    mut query: Query<(&Handle<T>, &Draw, &mut RenderPipelines)>,
+    mut queries: QuerySet<(
+        Query<(&Handle<T>, &mut RenderPipelines), Changed<Handle<T>>>,
+        Query<&mut RenderPipelines, With<Handle<T>>>,
+    )>,
+    entity_query: Query<Entity>,
 ) {
     let state = state.deref_mut();
     let uniform_buffer_arrays = &mut state.uniform_buffer_arrays;
     let render_resource_context = &**render_resource_context;
 
-    let modified_assets = assets.ids().collect::<Vec<_>>();
+    let mut changed_assets = HashMap::default();
+    for event in asset_state.event_reader.iter(&asset_events) {
+        match event {
+            AssetEvent::Created { ref handle } => {
+                if let Some(asset) = assets.get(handle) {
+                    changed_assets.insert(handle.id, asset);
+                }
+            }
+            AssetEvent::Modified { ref handle } => {
+                if let Some(asset) = assets.get(handle) {
+                    changed_assets.insert(handle.id, asset);
+                }
+            }
+            AssetEvent::Removed { ref handle } => {
+                uniform_buffer_arrays.remove_bindings(handle.id);
+                // if asset was modified and removed in the same update, ignore the modification
+                // events are ordered so future modification events are ok
+                changed_assets.remove(&handle.id);
+            }
+        }
+    }
 
     uniform_buffer_arrays.begin_update();
     // initialize uniform buffer arrays using the first RenderResources
-    if let Some(first_handle) = modified_assets.get(0) {
-        let asset = assets.get(*first_handle).expect(EXPECT_ASSET_MESSAGE);
+    if let Some(asset) = changed_assets.values().next() {
         uniform_buffer_arrays.initialize(asset, render_resource_context);
     }
 
-    for asset_handle in modified_assets.iter() {
-        let asset = assets.get(*asset_handle).expect(EXPECT_ASSET_MESSAGE);
+    for (asset_handle, asset) in changed_assets.iter() {
         uniform_buffer_arrays.prepare_uniform_buffers(*asset_handle, asset);
         let mut bindings =
             asset_render_resource_bindings.get_or_insert_mut(&Handle::<T>::weak(*asset_handle));
         setup_uniform_texture_resources::<T>(&asset, render_resource_context, &mut bindings);
     }
 
-    uniform_buffer_arrays.resize_buffer_arrays(render_resource_context);
+    let resized = uniform_buffer_arrays.resize_buffer_arrays(render_resource_context);
+    if resized {
+        uniform_buffer_arrays.set_required_staging_buffer_size_to_max()
+    }
     uniform_buffer_arrays.resize_staging_buffer(render_resource_context);
 
     if let Some(staging_buffer) = state.uniform_buffer_arrays.staging_buffer {
@@ -620,19 +659,34 @@ fn asset_render_resources_node_system<T: RenderResources + Asset>(
             staging_buffer,
             0..state.uniform_buffer_arrays.staging_buffer_size as u64,
             &mut |mut staging_buffer, _render_resource_context| {
-                for asset_handle in modified_assets.iter() {
-                    let asset = assets.get(*asset_handle).expect(EXPECT_ASSET_MESSAGE);
-                    let mut render_resource_bindings = asset_render_resource_bindings
-                        .get_or_insert_mut(&Handle::<T>::weak(*asset_handle));
-                    // TODO: only setup buffer if we haven't seen this handle before
-                    state.uniform_buffer_arrays.write_uniform_buffers(
-                        *asset_handle,
-                        &asset,
-                        state.dynamic_uniforms,
-                        render_resource_context,
-                        &mut render_resource_bindings,
-                        &mut staging_buffer,
-                    );
+                if resized {
+                    for (asset_handle, asset) in assets.iter() {
+                        let mut render_resource_bindings = asset_render_resource_bindings
+                            .get_or_insert_mut(&Handle::<T>::weak(asset_handle));
+                        // TODO: only setup buffer if we haven't seen this handle before
+                        state.uniform_buffer_arrays.write_uniform_buffers(
+                            asset_handle,
+                            &asset,
+                            state.dynamic_uniforms,
+                            render_resource_context,
+                            &mut render_resource_bindings,
+                            &mut staging_buffer,
+                        );
+                    }
+                } else {
+                    for (asset_handle, asset) in changed_assets.iter() {
+                        let mut render_resource_bindings = asset_render_resource_bindings
+                            .get_or_insert_mut(&Handle::<T>::weak(*asset_handle));
+                        // TODO: only setup buffer if we haven't seen this handle before
+                        state.uniform_buffer_arrays.write_uniform_buffers(
+                            *asset_handle,
+                            &asset,
+                            state.dynamic_uniforms,
+                            render_resource_context,
+                            &mut render_resource_bindings,
+                            &mut staging_buffer,
+                        );
+                    }
                 }
             },
         );
@@ -643,14 +697,24 @@ fn asset_render_resources_node_system<T: RenderResources + Asset>(
             .copy_staging_buffer_to_final_buffers(&mut state.command_queue, staging_buffer);
     }
 
-    for (asset_handle, draw, mut render_pipelines) in query.iter_mut() {
-        if !draw.is_visible {
-            continue;
-        }
-        if let Some(asset_bindings) = asset_render_resource_bindings.get(asset_handle) {
-            render_pipelines.bindings.extend(asset_bindings);
+    // update removed entity asset mapping
+    for entity in entity_query.removed::<Handle<T>>() {
+        if let Ok(mut render_pipelines) = queries.q1_mut().get_mut(*entity) {
+            render_pipelines
+                .bindings
+                .remove_asset_with_type(TypeId::of::<T>())
         }
     }
+
+    // update changed entity asset mapping
+    for (asset_handle, mut render_pipelines) in queries.q0_mut().iter_mut() {
+        render_pipelines
+            .bindings
+            .remove_asset_with_type(TypeId::of::<T>());
+        render_pipelines
+            .bindings
+            .add_asset(asset_handle.clone_weak_untyped(), TypeId::of::<T>());
+    }
 }
 
 fn setup_uniform_texture_resources<T>(