optimize asset gpu data transfer (#987)

crates/bevy_render/src/draw.rs

@@ -1,12 +1,12 @@
 use crate::{
     pipeline::{PipelineCompiler, PipelineDescriptor, PipelineLayout, PipelineSpecialization},
     renderer::{
-        BindGroup, BindGroupId, BufferId, RenderResource, RenderResourceBinding,
+        AssetRenderResourceBindings, BindGroup, BindGroupId, BufferId, RenderResource,
-        RenderResourceBindings, RenderResourceContext, SharedBuffers,
+        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
+        Self::set_bind_groups_from_bindings_internal(
-            .current_pipeline
+            &self.current_pipeline,
-            .as_ref()
+            &self.pipelines,
-            .ok_or(DrawError::NoPipelineSet)?;
+            &**self.render_resource_context,
-        let pipeline_descriptor = self
+            Some(&mut self.asset_render_resource_bindings),
-            .pipelines
+            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() {
+        'bind_group_descriptors: for bind_group_descriptor in layout.bind_groups.iter() {
-            bindings.update_bind_groups(pipeline_descriptor, &**self.render_resource_context);
-        }
-        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) {
+    if let Some(asset) = changed_assets.values().next() {
-        let asset = assets.get(*first_handle).expect(EXPECT_ASSET_MESSAGE);
         uniform_buffer_arrays.initialize(asset, render_resource_context);
     }
 
-    for asset_handle in modified_assets.iter() {
+    for (asset_handle, asset) in changed_assets.iter() {
-        let asset = assets.get(*asset_handle).expect(EXPECT_ASSET_MESSAGE);
         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() {
+                if resized {
-                    let asset = assets.get(*asset_handle).expect(EXPECT_ASSET_MESSAGE);
+                    for (asset_handle, asset) in assets.iter() {
-                    let mut render_resource_bindings = asset_render_resource_bindings
+                        let mut render_resource_bindings = asset_render_resource_bindings
-                        .get_or_insert_mut(&Handle::<T>::weak(*asset_handle));
+                            .get_or_insert_mut(&Handle::<T>::weak(asset_handle));
-                    // TODO: only setup buffer if we haven't seen this handle before
+                        // TODO: only setup buffer if we haven't seen this handle before
-                    state.uniform_buffer_arrays.write_uniform_buffers(
+                        state.uniform_buffer_arrays.write_uniform_buffers(
-                        *asset_handle,
+                            asset_handle,
-                        &asset,
+                            &asset,
-                        state.dynamic_uniforms,
+                            state.dynamic_uniforms,
-                        render_resource_context,
+                            render_resource_context,
-                        &mut render_resource_bindings,
+                            &mut render_resource_bindings,
-                        &mut staging_buffer,
+                            &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() {
+    // update removed entity asset mapping
-        if !draw.is_visible {
+    for entity in entity_query.removed::<Handle<T>>() {
-            continue;
+        if let Ok(mut render_pipelines) = queries.q1_mut().get_mut(*entity) {
-        }
+            render_pipelines
-        if let Some(asset_bindings) = asset_render_resource_bindings.get(asset_handle) {
+                .bindings
-            render_pipelines.bindings.extend(asset_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>(

crates/bevy_render/src/renderer/render_resource/render_resource_bindings.rs

@@ -5,7 +5,7 @@ use crate::{
 };
 use bevy_asset::{Asset, Handle, HandleUntyped};
 use bevy_utils::{HashMap, HashSet};
-use std::ops::Range;
+use std::{any::TypeId, ops::Range};
 
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub enum RenderResourceBinding {
@@ -67,6 +67,7 @@ pub struct RenderResourceBindings {
     /// A Buffer that is filled with zeros that will be used for attributes required by the shader, but undefined by the mesh.
     pub vertex_fallback_buffer: Option<BufferId>,
     pub index_buffer: Option<BufferId>,
+    assets: HashSet<(HandleUntyped, TypeId)>,
     bind_groups: HashMap<BindGroupId, BindGroup>,
     bind_group_descriptors: HashMap<BindGroupDescriptorId, Option<BindGroupId>>,
     dirty_bind_groups: HashSet<BindGroupId>,
@@ -129,7 +130,7 @@ impl RenderResourceBindings {
         }
     }
 
-    pub fn update_bind_group(
+    fn update_bind_group_status(
         &mut self,
         bind_group_descriptor: &BindGroupDescriptor,
     ) -> BindGroupStatus {
@@ -149,6 +150,50 @@ impl RenderResourceBindings {
         }
     }
 
+    pub fn add_asset(&mut self, handle: HandleUntyped, type_id: TypeId) {
+        self.dynamic_bindings_generation += 1;
+        self.assets.insert((handle, type_id));
+    }
+
+    pub fn remove_asset_with_type(&mut self, type_id: TypeId) {
+        self.dynamic_bindings_generation += 1;
+        self.assets.retain(|(_, current_id)| *current_id != type_id);
+    }
+
+    pub fn iter_assets(&self) -> impl Iterator<Item = &(HandleUntyped, TypeId)> {
+        self.assets.iter()
+    }
+
+    pub fn update_bind_group(
+        &mut self,
+        bind_group_descriptor: &BindGroupDescriptor,
+        render_resource_context: &dyn RenderResourceContext,
+    ) -> Option<&BindGroup> {
+        let status = self.update_bind_group_status(bind_group_descriptor);
+        match status {
+            BindGroupStatus::Changed(id) => {
+                let bind_group = self
+                    .get_bind_group(id)
+                    .expect("`RenderResourceSet` was just changed, so it should exist.");
+                render_resource_context.create_bind_group(bind_group_descriptor.id, bind_group);
+                Some(bind_group)
+            }
+            BindGroupStatus::Unchanged(id) => {
+                // PERF: this is only required because RenderResourceContext::remove_stale_bind_groups doesn't inform RenderResourceBindings
+                // when a stale bind group has been removed
+                let bind_group = self
+                    .get_bind_group(id)
+                    .expect("`RenderResourceSet` was just changed, so it should exist.");
+                render_resource_context.create_bind_group(bind_group_descriptor.id, bind_group);
+                Some(bind_group)
+            }
+            BindGroupStatus::NoMatch => {
+                // ignore unchanged / unmatched render resource sets
+                None
+            }
+        }
+    }
+
     pub fn update_bind_groups(
         &mut self,
         pipeline: &PipelineDescriptor,
@@ -156,25 +201,7 @@ impl RenderResourceBindings {
     ) {
         let layout = pipeline.get_layout().unwrap();
         for bind_group_descriptor in layout.bind_groups.iter() {
-            match self.update_bind_group(bind_group_descriptor) {
+            self.update_bind_group(bind_group_descriptor, render_resource_context);
-                BindGroupStatus::Changed(id) => {
-                    let bind_group = self
-                        .get_bind_group(id)
-                        .expect("`RenderResourceSet` was just changed, so it should exist.");
-                    render_resource_context.create_bind_group(bind_group_descriptor.id, bind_group);
-                }
-                BindGroupStatus::Unchanged(id) => {
-                    // PERF: this is only required because RenderResourceContext::remove_stale_bind_groups doesn't inform RenderResourceBindings
-                    // when a stale bind group has been removed
-                    let bind_group = self
-                        .get_bind_group(id)
-                        .expect("`RenderResourceSet` was just changed, so it should exist.");
-                    render_resource_context.create_bind_group(bind_group_descriptor.id, bind_group);
-                }
-                BindGroupStatus::NoMatch => {
-                    // ignore unchanged / unmatched render resource sets
-                }
-            }
         }
     }
 
@@ -228,7 +255,11 @@ pub struct AssetRenderResourceBindings {
 
 impl AssetRenderResourceBindings {
     pub fn get<T: Asset>(&self, handle: &Handle<T>) -> Option<&RenderResourceBindings> {
-        self.bindings.get(&handle.clone_weak_untyped())
+        self.get_untyped(&handle.clone_weak_untyped())
+    }
+
+    pub fn get_untyped(&self, handle: &HandleUntyped) -> Option<&RenderResourceBindings> {
+        self.bindings.get(handle)
     }
 
     pub fn get_or_insert_mut<T: Asset>(
@@ -241,7 +272,14 @@ impl AssetRenderResourceBindings {
     }
 
     pub fn get_mut<T: Asset>(&mut self, handle: &Handle<T>) -> Option<&mut RenderResourceBindings> {
-        self.bindings.get_mut(&handle.clone_weak_untyped())
+        self.get_mut_untyped(&handle.clone_weak_untyped())
+    }
+
+    pub fn get_mut_untyped(
+        &mut self,
+        handle: &HandleUntyped,
+    ) -> Option<&mut RenderResourceBindings> {
+        self.bindings.get_mut(handle)
     }
 }
 
@@ -293,7 +331,7 @@ mod tests {
         equal_bindings.set("a", resource1.clone());
         equal_bindings.set("b", resource2.clone());
 
-        let status = bindings.update_bind_group(&bind_group_descriptor);
+        let status = bindings.update_bind_group_status(&bind_group_descriptor);
         let id = if let BindGroupStatus::Changed(id) = status {
             id
         } else {
@@ -301,7 +339,7 @@ mod tests {
         };
 
         let different_bind_group_status =
-            different_bindings.update_bind_group(&bind_group_descriptor);
+            different_bindings.update_bind_group_status(&bind_group_descriptor);
         if let BindGroupStatus::Changed(different_bind_group_id) = different_bind_group_status {
             assert_ne!(
                 id, different_bind_group_id,
@@ -312,7 +350,8 @@ mod tests {
             panic!("Expected a changed bind group.");
         };
 
-        let equal_bind_group_status = equal_bindings.update_bind_group(&bind_group_descriptor);
+        let equal_bind_group_status =
+            equal_bindings.update_bind_group_status(&bind_group_descriptor);
         if let BindGroupStatus::Changed(equal_bind_group_id) = equal_bind_group_status {
             assert_eq!(
                 id, equal_bind_group_id,
@@ -325,7 +364,7 @@ mod tests {
         let mut unmatched_bindings = RenderResourceBindings::default();
         unmatched_bindings.set("a", resource1.clone());
         let unmatched_bind_group_status =
-            unmatched_bindings.update_bind_group(&bind_group_descriptor);
+            unmatched_bindings.update_bind_group_status(&bind_group_descriptor);
         assert_eq!(unmatched_bind_group_status, BindGroupStatus::NoMatch);
     }
 }

crates/bevy_text/src/draw.rs

@@ -5,7 +5,7 @@ use bevy_render::{
     mesh,
     pipeline::{PipelineSpecialization, VertexBufferDescriptor},
     prelude::Msaa,
-    renderer::{AssetRenderResourceBindings, BindGroup, RenderResourceBindings, RenderResourceId},
+    renderer::{BindGroup, RenderResourceBindings, RenderResourceId},
 };
 use bevy_sprite::TextureAtlasSprite;
 use glyph_brush_layout::{HorizontalAlign, VerticalAlign};
@@ -46,7 +46,6 @@ impl Default for TextStyle {
 
 pub struct DrawableText<'a> {
     pub render_resource_bindings: &'a mut RenderResourceBindings,
-    pub asset_render_resource_bindings: &'a mut AssetRenderResourceBindings,
     pub position: Vec3,
     pub style: &'a TextStyle,
     pub text_glyphs: &'a Vec<PositionedGlyph>,
@@ -92,11 +91,7 @@ impl<'a> Drawable for DrawableText<'a> {
         context.set_bind_groups_from_bindings(draw, &mut [self.render_resource_bindings])?;
 
         for tv in self.text_glyphs {
-            let atlas_render_resource_bindings = self
+            context.set_asset_bind_groups(draw, &tv.atlas_info.texture_atlas)?;
-                .asset_render_resource_bindings
-                .get_mut(&tv.atlas_info.texture_atlas)
-                .unwrap();
-            context.set_bind_groups_from_bindings(draw, &mut [atlas_render_resource_bindings])?;
 
             let sprite = TextureAtlasSprite {
                 index: tv.atlas_info.glyph_index,

crates/bevy_ui/src/widget/text.rs

@@ -6,7 +6,7 @@ use bevy_render::{
     draw::{Draw, DrawContext, Drawable},
     mesh::Mesh,
     prelude::Msaa,
-    renderer::{AssetRenderResourceBindings, RenderResourceBindings},
+    renderer::RenderResourceBindings,
     texture::Texture,
 };
 use bevy_sprite::{TextureAtlas, QUAD_HANDLE};
@@ -145,7 +145,6 @@ pub fn draw_text_system(
     msaa: Res<Msaa>,
     meshes: Res<Assets<Mesh>>,
     mut render_resource_bindings: ResMut<RenderResourceBindings>,
-    mut asset_render_resource_bindings: ResMut<AssetRenderResourceBindings>,
     text_pipeline: Res<DefaultTextPipeline>,
     mut query: Query<(Entity, &mut Draw, &Text, &Node, &GlobalTransform)>,
 ) {
@@ -162,7 +161,6 @@ pub fn draw_text_system(
 
             let mut drawable_text = DrawableText {
                 render_resource_bindings: &mut render_resource_bindings,
-                asset_render_resource_bindings: &mut asset_render_resource_bindings,
                 position,
                 msaa: &msaa,
                 text_glyphs: &text_glyphs.glyphs,