bevy_render: delegate buffer aligning to render_resource_context (#842)

Author: mrk-its

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

@@ -15,12 +15,6 @@ use bevy_utils::HashMap;
 use renderer::{AssetRenderResourceBindings, BufferId, RenderResourceType, RenderResources};
 use std::{hash::Hash, marker::PhantomData, ops::DerefMut};
 
-pub const BIND_BUFFER_ALIGNMENT: usize = 256;
-
-fn get_aligned_dynamic_uniform_size(data_size: usize) -> usize {
-    BIND_BUFFER_ALIGNMENT * ((data_size as f32 / BIND_BUFFER_ALIGNMENT as f32).ceil() as usize)
-}
-
 #[derive(Debug)]
 struct QueuedBufferWrite {
     buffer: BufferId,
@@ -42,13 +36,9 @@ struct BufferArray<I> {
 }
 
 impl<I: Hash + Eq> BufferArray<I> {
-    pub fn new(item_size: usize, min_capacity: usize, align: bool) -> Self {
+    pub fn new(item_size: usize, min_capacity: usize) -> Self {
         BufferArray {
-            item_size: if align {
-                get_aligned_dynamic_uniform_size(item_size)
-            } else {
-                item_size
-            },
+            item_size,
             len: 0,
             buffer_capacity: 0,
             min_capacity,
@@ -158,13 +148,18 @@ where
     T: renderer::RenderResources,
 {
     /// Initialize this UniformBufferArrays using information from a RenderResources value.
-    fn initialize(&mut self, render_resources: &T) {
+    fn initialize(
+        &mut self,
+        render_resources: &T,
+        render_resource_context: &dyn RenderResourceContext,
+    ) {
         if self.buffer_arrays.len() != render_resources.render_resources_len() {
             let mut buffer_arrays = Vec::with_capacity(render_resources.render_resources_len());
             for render_resource in render_resources.iter() {
                 if let Some(RenderResourceType::Buffer) = render_resource.resource_type() {
                     let size = render_resource.buffer_byte_len().unwrap();
-                    buffer_arrays.push(Some(BufferArray::new(size, 10, true)));
+                    let aligned_size = render_resource_context.get_aligned_uniform_size(size, true);
+                    buffer_arrays.push(Some(BufferArray::new(aligned_size, 10)));
                 } else {
                     buffer_arrays.push(None);
                 }
@@ -247,8 +242,10 @@ where
                 Some(RenderResourceType::Buffer) => {
                     let size = render_resource.buffer_byte_len().unwrap();
                     let render_resource_name = uniforms.get_render_resource_name(i).unwrap();
+                    let aligned_size =
+                        render_resource_context.get_aligned_uniform_size(size, false);
                     let buffer_array = self.buffer_arrays[i].as_mut().unwrap();
-                    let range = 0..size as u64;
+                    let range = 0..aligned_size as u64;
                     let (target_buffer, target_offset) = if dynamic_uniforms {
                         let binding = buffer_array.get_binding(id).unwrap();
                         let dynamic_index = if let RenderResourceBinding::Buffer {
@@ -270,7 +267,7 @@ where
                                 size: current_size, ..
                             }) = render_resource_context.get_buffer_info(buffer_id)
                             {
-                                if size == current_size {
+                                if aligned_size == current_size {
                                     matching_buffer = Some(buffer_id);
                                 } else {
                                     render_resource_context.remove_buffer(buffer_id);
@@ -291,7 +288,7 @@ where
                             }
 
                             let buffer = render_resource_context.create_buffer(BufferInfo {
-                                size,
+                                size: aligned_size,
                                 buffer_usage: BufferUsage::COPY_DST | usage,
                                 ..Default::default()
                             });
@@ -432,7 +429,7 @@ fn render_resources_node_system<T: RenderResources>(
     uniform_buffer_arrays.begin_update();
     // initialize uniform buffer arrays using the first RenderResources
     if let Some((_, first, _, _)) = query.iter_mut().next() {
-        uniform_buffer_arrays.initialize(first);
+        uniform_buffer_arrays.initialize(first, render_resource_context);
     }
 
     for entity in query.removed::<T>() {
@@ -482,23 +479,6 @@ fn render_resources_node_system<T: RenderResources>(
         state
             .uniform_buffer_arrays
             .copy_staging_buffer_to_final_buffers(&mut state.command_queue, staging_buffer);
-    } else {
-        // TODO: can we just remove this?
-        let mut staging_buffer: [u8; 0] = [];
-        for (entity, uniforms, draw, mut render_pipelines) in query.iter_mut() {
-            if !draw.is_visible {
-                continue;
-            }
-
-            state.uniform_buffer_arrays.write_uniform_buffers(
-                entity,
-                &uniforms,
-                state.dynamic_uniforms,
-                render_resource_context,
-                &mut render_pipelines.bindings,
-                &mut staging_buffer,
-            );
-        }
     }
 }
 
@@ -579,7 +559,7 @@ fn asset_render_resources_node_system<T: RenderResources + Asset>(
     // 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);
-        uniform_buffer_arrays.initialize(asset);
+        uniform_buffer_arrays.initialize(asset, render_resource_context);
     }
 
     for asset_handle in modified_assets.iter() {
@@ -620,22 +600,6 @@ fn asset_render_resources_node_system<T: RenderResources + Asset>(
         state
             .uniform_buffer_arrays
             .copy_staging_buffer_to_final_buffers(&mut state.command_queue, staging_buffer);
-    } else {
-        let mut staging_buffer: [u8; 0] = [];
-        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,
-            );
-        }
     }
 
     for (asset_handle, draw, mut render_pipelines) in query.iter_mut() {

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

@@ -12,11 +12,6 @@ pub struct TextureCopyNode {
     pub texture_event_reader: EventReader<AssetEvent<Texture>>,
 }
 
-pub const ALIGNMENT: usize = 256;
-fn get_aligned(data_size: f32) -> usize {
-    ALIGNMENT * ((data_size / ALIGNMENT as f32).ceil() as usize)
-}
-
 impl Node for TextureCopyNode {
     fn update(
         &mut self,
@@ -34,7 +29,9 @@ impl Node for TextureCopyNode {
                     if let Some(texture) = textures.get(handle) {
                         let texture_descriptor: TextureDescriptor = texture.into();
                         let width = texture.size.x() as usize;
-                        let aligned_width = get_aligned(texture.size.x());
+                        let aligned_width = render_context
+                            .resources()
+                            .get_aligned_texture_size(texture.size.x() as usize);
                         let format_size = texture.format.pixel_size();
                         let mut aligned_data =
                             vec![0; format_size * aligned_width * texture.size.y() as usize];

crates/bevy_render/src/renderer/headless_render_resource_context.rs

@@ -140,4 +140,12 @@ impl RenderResourceContext for HeadlessRenderResourceContext {
     ) -> bool {
         false
     }
+
+    fn get_aligned_uniform_size(&self, size: usize, _dynamic: bool) -> usize {
+        size
+    }
+
+    fn get_aligned_texture_size(&self, size: usize) -> usize {
+        size
+    }
 }

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

@@ -31,9 +31,12 @@ impl SharedBuffers {
         buffer_usage: BufferUsage,
     ) -> Option<RenderResourceBinding> {
         if let Some(size) = render_resource.buffer_byte_len() {
+            let aligned_size = self
+                .render_resource_context
+                .get_aligned_uniform_size(size, false);
             // PERF: this buffer will be slow
             let staging_buffer = self.render_resource_context.create_buffer(BufferInfo {
-                size,
+                size: aligned_size,
                 buffer_usage: BufferUsage::COPY_SRC | BufferUsage::MAP_WRITE,
                 mapped_at_creation: true,
             });
@@ -49,7 +52,7 @@ impl SharedBuffers {
             self.render_resource_context.unmap_buffer(staging_buffer);
 
             let destination_buffer = self.render_resource_context.create_buffer(BufferInfo {
-                size,
+                size: aligned_size,
                 buffer_usage: BufferUsage::COPY_DST | buffer_usage,
                 ..Default::default()
             });
@@ -68,7 +71,7 @@ impl SharedBuffers {
             buffers.push(destination_buffer);
             Some(RenderResourceBinding::Buffer {
                 buffer: destination_buffer,
-                range: 0..size as u64,
+                range: 0..aligned_size as u64,
                 dynamic_index: None,
             })
         } else {

crates/bevy_render/src/renderer/render_resource_context.rs

@@ -33,7 +33,8 @@ pub trait RenderResourceContext: Downcast + Send + Sync + 'static {
     fn remove_texture(&self, texture: TextureId);
     fn remove_sampler(&self, sampler: SamplerId);
     fn get_buffer_info(&self, buffer: BufferId) -> Option<BufferInfo>;
-
+    fn get_aligned_uniform_size(&self, size: usize, dynamic: bool) -> usize;
+    fn get_aligned_texture_size(&self, data_size: usize) -> usize;
     fn set_asset_resource_untyped(
         &self,
         handle: HandleUntyped,

crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs

@@ -26,6 +26,9 @@ pub struct WgpuRenderResourceContext {
     pub resources: WgpuResources,
 }
 
+pub const BIND_BUFFER_ALIGNMENT: usize = 256;
+pub const TEXTURE_ALIGNMENT: usize = 256;
+
 impl WgpuRenderResourceContext {
     pub fn new(device: Arc<wgpu::Device>) -> Self {
         WgpuRenderResourceContext {
@@ -553,4 +556,16 @@ impl RenderResourceContext for WgpuRenderResourceContext {
         let buffer = buffers.get(&id).unwrap();
         buffer.unmap();
     }
+
+    fn get_aligned_texture_size(&self, size: usize) -> usize {
+        (size + TEXTURE_ALIGNMENT - 1) & !(TEXTURE_ALIGNMENT - 1)
+    }
+
+    fn get_aligned_uniform_size(&self, size: usize, dynamic: bool) -> usize {
+        if dynamic {
+            (size + BIND_BUFFER_ALIGNMENT - 1) & !(BIND_BUFFER_ALIGNMENT - 1)
+        } else {
+            size
+        }
+    }
 }