Reduce branching in TrackedRenderPass (#7053)

# Objective
Speed up the render phase for rendering.

## Solution
 - Follow up #6988 and make the internals of atomic IDs `NonZeroU32`. This niches the `Option`s of the IDs in draw state, which reduces the size and branching behavior when evaluating for equality.
 - Require `&RenderDevice` to get the device's `Limits` when initializing a `TrackedRenderPass` to preallocate the bind groups and vertex buffer state in `DrawState`, this removes the branch on needing to resize those `Vec`s.

## Performance
This produces a similar speed up akin to that of #6885. This shows an approximate 6% speed up in `main_opaque_pass_3d` on `many_foxes` (408.79 us -> 388us). This should be orthogonal to the gains seen there.

![image](https://user-images.githubusercontent.com/3137680/209906239-e430f026-63c2-4b95-957e-a2045b810d79.png)

---

## Changelog
Added: `RenderContext::begin_tracked_render_pass`.
Changed: `TrackedRenderPass` now requires a `&RenderDevice` on construction.
Removed: `bevy_render::render_phase::DrawState`. It was not usable in any form outside of `bevy_render`.

## Migration Guide
TODO

Author: james7132

crates/bevy_core_pipeline/src/bloom/mod.rs

@@ -17,7 +17,6 @@ use bevy_render::{
     },
     prelude::Camera,
     render_graph::{Node, NodeRunError, RenderGraph, RenderGraphContext, SlotInfo, SlotType},
-    render_phase::TrackedRenderPass,
     render_resource::*,
     renderer::{RenderContext, RenderDevice},
     texture::{CachedTexture, TextureCache},
@@ -232,17 +231,15 @@ impl Node for BloomNode {
         {
             let view = &BloomTextures::texture_view(&textures.texture_a, 0);
             let mut prefilter_pass =
-                TrackedRenderPass::new(render_context.command_encoder.begin_render_pass(
-                    &RenderPassDescriptor {
-                        label: Some("bloom_prefilter_pass"),
-                        color_attachments: &[Some(RenderPassColorAttachment {
-                            view,
-                            resolve_target: None,
-                            ops: Operations::default(),
-                        })],
-                        depth_stencil_attachment: None,
-                    },
-                ));
+                render_context.begin_tracked_render_pass(RenderPassDescriptor {
+                    label: Some("bloom_prefilter_pass"),
+                    color_attachments: &[Some(RenderPassColorAttachment {
+                        view,
+                        resolve_target: None,
+                        ops: Operations::default(),
+                    })],
+                    depth_stencil_attachment: None,
+                });
             prefilter_pass.set_render_pipeline(downsampling_prefilter_pipeline);
             prefilter_pass.set_bind_group(
                 0,
@@ -258,17 +255,15 @@ impl Node for BloomNode {
         for mip in 1..textures.mip_count {
             let view = &BloomTextures::texture_view(&textures.texture_a, mip);
             let mut downsampling_pass =
-                TrackedRenderPass::new(render_context.command_encoder.begin_render_pass(
-                    &RenderPassDescriptor {
-                        label: Some("bloom_downsampling_pass"),
-                        color_attachments: &[Some(RenderPassColorAttachment {
-                            view,
-                            resolve_target: None,
-                            ops: Operations::default(),
-                        })],
-                        depth_stencil_attachment: None,
-                    },
-                ));
+                render_context.begin_tracked_render_pass(RenderPassDescriptor {
+                    label: Some("bloom_downsampling_pass"),
+                    color_attachments: &[Some(RenderPassColorAttachment {
+                        view,
+                        resolve_target: None,
+                        ops: Operations::default(),
+                    })],
+                    depth_stencil_attachment: None,
+                });
             downsampling_pass.set_render_pipeline(downsampling_pipeline);
             downsampling_pass.set_bind_group(
                 0,
@@ -284,17 +279,15 @@ impl Node for BloomNode {
         for mip in (1..textures.mip_count).rev() {
             let view = &BloomTextures::texture_view(&textures.texture_b, mip - 1);
             let mut upsampling_pass =
-                TrackedRenderPass::new(render_context.command_encoder.begin_render_pass(
-                    &RenderPassDescriptor {
-                        label: Some("bloom_upsampling_pass"),
-                        color_attachments: &[Some(RenderPassColorAttachment {
-                            view,
-                            resolve_target: None,
-                            ops: Operations::default(),
-                        })],
-                        depth_stencil_attachment: None,
-                    },
-                ));
+                render_context.begin_tracked_render_pass(RenderPassDescriptor {
+                    label: Some("bloom_upsampling_pass"),
+                    color_attachments: &[Some(RenderPassColorAttachment {
+                        view,
+                        resolve_target: None,
+                        ops: Operations::default(),
+                    })],
+                    depth_stencil_attachment: None,
+                });
             upsampling_pass.set_render_pipeline(upsampling_pipeline);
             upsampling_pass.set_bind_group(
                 0,
@@ -309,18 +302,16 @@ impl Node for BloomNode {
 
         {
             let mut upsampling_final_pass =
-                TrackedRenderPass::new(render_context.command_encoder.begin_render_pass(
-                    &RenderPassDescriptor {
-                        label: Some("bloom_upsampling_final_pass"),
-                        color_attachments: &[Some(view_target.get_unsampled_color_attachment(
-                            Operations {
-                                load: LoadOp::Load,
-                                store: true,
-                            },
-                        ))],
-                        depth_stencil_attachment: None,
-                    },
-                ));
+                render_context.begin_tracked_render_pass(RenderPassDescriptor {
+                    label: Some("bloom_upsampling_final_pass"),
+                    color_attachments: &[Some(view_target.get_unsampled_color_attachment(
+                        Operations {
+                            load: LoadOp::Load,
+                            store: true,
+                        },
+                    ))],
+                    depth_stencil_attachment: None,
+                });
             upsampling_final_pass.set_render_pipeline(upsampling_final_pipeline);
             upsampling_final_pass.set_bind_group(
                 0,

crates/bevy_core_pipeline/src/core_2d/main_pass_2d_node.rs

@@ -3,7 +3,6 @@ use crate::{
     core_2d::{camera_2d::Camera2d, Transparent2d},
 };
 use bevy_ecs::prelude::*;
-use bevy_render::render_phase::TrackedRenderPass;
 use bevy_render::{
     camera::ExtractedCamera,
     render_graph::{Node, NodeRunError, RenderGraphContext, SlotInfo, SlotType},
@@ -63,7 +62,8 @@ impl Node for MainPass2dNode {
         {
             #[cfg(feature = "trace")]
             let _main_pass_2d = info_span!("main_pass_2d").entered();
-            let pass_descriptor = RenderPassDescriptor {
+
+            let mut render_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
                 label: Some("main_pass_2d"),
                 color_attachments: &[Some(target.get_color_attachment(Operations {
                     load: match camera_2d.clear_color {
@@ -76,12 +76,7 @@ impl Node for MainPass2dNode {
                     store: true,
                 }))],
                 depth_stencil_attachment: None,
-            };
-
-            let render_pass = render_context
-                .command_encoder
-                .begin_render_pass(&pass_descriptor);
-            let mut render_pass = TrackedRenderPass::new(render_pass);
+            });
 
             if let Some(viewport) = camera.viewport.as_ref() {
                 render_pass.set_camera_viewport(viewport);

crates/bevy_core_pipeline/src/core_3d/main_pass_3d_node.rs

@@ -3,7 +3,6 @@ use crate::{
     core_3d::{AlphaMask3d, Camera3d, Opaque3d, Transparent3d},
 };
 use bevy_ecs::prelude::*;
-use bevy_render::render_phase::TrackedRenderPass;
 use bevy_render::{
     camera::ExtractedCamera,
     render_graph::{Node, NodeRunError, RenderGraphContext, SlotInfo, SlotType},
@@ -70,7 +69,8 @@ impl Node for MainPass3dNode {
             // NOTE: Scoped to drop the mutable borrow of render_context
             #[cfg(feature = "trace")]
             let _main_opaque_pass_3d_span = info_span!("main_opaque_pass_3d").entered();
-            let pass_descriptor = RenderPassDescriptor {
+
+            let mut render_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
                 label: Some("main_opaque_pass_3d"),
                 // NOTE: The opaque pass loads the color
                 // buffer as well as writing to it.
@@ -94,12 +94,7 @@ impl Node for MainPass3dNode {
                     }),
                     stencil_ops: None,
                 }),
-            };
-
-            let render_pass = render_context
-                .command_encoder
-                .begin_render_pass(&pass_descriptor);
-            let mut render_pass = TrackedRenderPass::new(render_pass);
+            });
 
             if let Some(viewport) = camera.viewport.as_ref() {
                 render_pass.set_camera_viewport(viewport);
@@ -113,7 +108,8 @@ impl Node for MainPass3dNode {
             // NOTE: Scoped to drop the mutable borrow of render_context
             #[cfg(feature = "trace")]
             let _main_alpha_mask_pass_3d_span = info_span!("main_alpha_mask_pass_3d").entered();
-            let pass_descriptor = RenderPassDescriptor {
+
+            let mut render_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
                 label: Some("main_alpha_mask_pass_3d"),
                 // NOTE: The alpha_mask pass loads the color buffer as well as overwriting it where appropriate.
                 color_attachments: &[Some(target.get_color_attachment(Operations {
@@ -129,12 +125,7 @@ impl Node for MainPass3dNode {
                     }),
                     stencil_ops: None,
                 }),
-            };
-
-            let render_pass = render_context
-                .command_encoder
-                .begin_render_pass(&pass_descriptor);
-            let mut render_pass = TrackedRenderPass::new(render_pass);
+            });
 
             if let Some(viewport) = camera.viewport.as_ref() {
                 render_pass.set_camera_viewport(viewport);
@@ -148,7 +139,8 @@ impl Node for MainPass3dNode {
             // NOTE: Scoped to drop the mutable borrow of render_context
             #[cfg(feature = "trace")]
             let _main_transparent_pass_3d_span = info_span!("main_transparent_pass_3d").entered();
-            let pass_descriptor = RenderPassDescriptor {
+
+            let mut render_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
                 label: Some("main_transparent_pass_3d"),
                 // NOTE: The transparent pass loads the color buffer as well as overwriting it where appropriate.
                 color_attachments: &[Some(target.get_color_attachment(Operations {
@@ -169,12 +161,7 @@ impl Node for MainPass3dNode {
                     }),
                     stencil_ops: None,
                 }),
-            };
-
-            let render_pass = render_context
-                .command_encoder
-                .begin_render_pass(&pass_descriptor);
-            let mut render_pass = TrackedRenderPass::new(render_pass);
+            });
 
             if let Some(viewport) = camera.viewport.as_ref() {
                 render_pass.set_camera_viewport(viewport);

crates/bevy_pbr/src/render/light.rs

@@ -1770,23 +1770,19 @@ impl Node for ShadowPassNode {
                     continue;
                 }
 
-                let pass_descriptor = RenderPassDescriptor {
-                    label: Some(&view_light.pass_name),
-                    color_attachments: &[],
-                    depth_stencil_attachment: Some(RenderPassDepthStencilAttachment {
-                        view: &view_light.depth_texture_view,
-                        depth_ops: Some(Operations {
-                            load: LoadOp::Clear(0.0),
-                            store: true,
+                let mut render_pass =
+                    render_context.begin_tracked_render_pass(RenderPassDescriptor {
+                        label: Some(&view_light.pass_name),
+                        color_attachments: &[],
+                        depth_stencil_attachment: Some(RenderPassDepthStencilAttachment {
+                            view: &view_light.depth_texture_view,
+                            depth_ops: Some(Operations {
+                                load: LoadOp::Clear(0.0),
+                                store: true,
+                            }),
+                            stencil_ops: None,
                         }),
-                        stencil_ops: None,
-                    }),
-                };
-
-                let render_pass = render_context
-                    .command_encoder
-                    .begin_render_pass(&pass_descriptor);
-                let mut render_pass = TrackedRenderPass::new(render_pass);
+                    });
 
                 shadow_phase.render(&mut render_pass, world, view_light_entity);
             }

crates/bevy_render/Cargo.toml

@@ -50,6 +50,7 @@ image = { version = "0.24", default-features = false }
 
 # misc
 wgpu = { version = "0.14.0", features = ["spirv"] }
+wgpu-hal = "0.14.1"
 codespan-reporting = "0.11.0"
 naga = { version = "0.10.0", features = ["glsl-in", "spv-in", "spv-out", "wgsl-in", "wgsl-out"] }
 serde = { version = "1", features = ["derive"] }

crates/bevy_render/src/render_phase/draw_state.rs

@@ -5,14 +5,16 @@ use crate::{
         BindGroup, BindGroupId, Buffer, BufferId, BufferSlice, RenderPipeline, RenderPipelineId,
         ShaderStages,
     },
+    renderer::RenderDevice,
 };
-use bevy_utils::tracing::trace;
+use bevy_utils::{default, tracing::trace};
 use std::ops::Range;
 use wgpu::{IndexFormat, RenderPass};
+use wgpu_hal::{MAX_BIND_GROUPS, MAX_VERTEX_BUFFERS};
 
 /// Tracks the current [`TrackedRenderPass`] state to ensure draw calls are valid.
 #[derive(Debug, Default)]
-pub struct DrawState {
+struct DrawState {
     pipeline: Option<RenderPipelineId>,
     bind_groups: Vec<(Option<BindGroupId>, Vec<u32>)>,
     vertex_buffers: Vec<Option<(BufferId, u64)>>,
@@ -26,12 +28,10 @@ impl DrawState {
         bind_group: BindGroupId,
         dynamic_indices: &[u32],
     ) {
-        if index >= self.bind_groups.len() {
-            self.bind_groups.resize(index + 1, (None, Vec::new()));
-        }
-        self.bind_groups[index].0 = Some(bind_group);
-        self.bind_groups[index].1.clear();
-        self.bind_groups[index].1.extend(dynamic_indices);
+        let group = &mut self.bind_groups[index];
+        group.0 = Some(bind_group);
+        group.1.clear();
+        group.1.extend(dynamic_indices);
     }
 
     pub fn is_bind_group_set(
@@ -48,9 +48,6 @@ impl DrawState {
     }
 
     pub fn set_vertex_buffer(&mut self, index: usize, buffer: BufferId, offset: u64) {
-        if index >= self.vertex_buffers.len() {
-            self.vertex_buffers.resize(index + 1, None);
-        }
         self.vertex_buffers[index] = Some((buffer, offset));
     }
 
@@ -98,9 +95,16 @@ pub struct TrackedRenderPass<'a> {
 
 impl<'a> TrackedRenderPass<'a> {
     /// Tracks the supplied render pass.
-    pub fn new(pass: RenderPass<'a>) -> Self {
+    pub fn new(device: &RenderDevice, pass: RenderPass<'a>) -> Self {
+        let limits = device.limits();
+        let max_bind_groups = limits.max_bind_groups as usize;
+        let max_vertex_buffers = limits.max_vertex_buffers as usize;
         Self {
-            state: DrawState::default(),
+            state: DrawState {
+                bind_groups: vec![(None, Vec::new()); max_bind_groups.min(MAX_BIND_GROUPS)],
+                vertex_buffers: vec![None; max_vertex_buffers.min(MAX_VERTEX_BUFFERS)],
+                ..default()
+            },
             pass,
         }
     }

crates/bevy_render/src/render_resource/resource_macros.rs

@@ -124,7 +124,7 @@ macro_rules! render_resource_wrapper {
 macro_rules! define_atomic_id {
     ($atomic_id_type:ident) => {
         #[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
-        pub struct $atomic_id_type(u32);
+        pub struct $atomic_id_type(core::num::NonZeroU32);
 
         // We use new instead of default to indicate that each ID created will be unique.
         #[allow(clippy::new_without_default)]
@@ -134,15 +134,13 @@ macro_rules! define_atomic_id {
 
                 static COUNTER: AtomicU32 = AtomicU32::new(1);
 
-                match COUNTER.fetch_add(1, Ordering::Relaxed) {
-                    0 => {
-                        panic!(
-                            "The system ran out of unique `{}`s.",
-                            stringify!($atomic_id_type)
-                        );
-                    }
-                    id => Self(id),
-                }
+                let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
+                Self(core::num::NonZeroU32::new(counter).unwrap_or_else(|| {
+                    panic!(
+                        "The system ran out of unique `{}`s.",
+                        stringify!($atomic_id_type)
+                    );
+                }))
             }
         }
     };

crates/bevy_render/src/renderer/mod.rs

@@ -8,6 +8,8 @@ pub use render_device::*;
 
 use crate::{
     render_graph::RenderGraph,
+    render_phase::TrackedRenderPass,
+    render_resource::RenderPassDescriptor,
     settings::{WgpuSettings, WgpuSettingsPriority},
     view::{ExtractedWindows, ViewTarget},
 };
@@ -279,3 +281,17 @@ pub struct RenderContext {
     pub render_device: RenderDevice,
     pub command_encoder: CommandEncoder,
 }
+
+impl RenderContext {
+    /// Creates a new [`TrackedRenderPass`] for the context,
+    /// configured using the provided `descriptor`.
+    pub fn begin_tracked_render_pass<'a>(
+        &'a mut self,
+        descriptor: RenderPassDescriptor<'a, '_>,
+    ) -> TrackedRenderPass<'a> {
+        TrackedRenderPass::new(
+            &self.render_device,
+            self.command_encoder.begin_render_pass(&descriptor),
+        )
+    }
+}