Improved UI render batching (#8793)

# Objective

`prepare_uinodes` creates a new `UiBatch` whenever the texture changes,
when most often it's just queuing untextured quads. Instead of switching
textures, we can reduce the number of batches generated significantly by
adding a condition to the fragment shader so that it only multiplies by
the `textureSample` value when drawing a textured quad.

# Solution

Add a `mode` field to `UiVertex`.
In `prepare_uinodes` set `mode` to 0 if the quad is textured or 1 if
untextured.
Add a condition to the fragment shader that only multiplies by the
`color` value from `textureSample` if `mode` is set to 1.

---

## Changelog
* Added a `mode` field to `UiVertex`, and added an extra `u32` vertex
attribute to the shader and vertex buffer layout.
* In `prepare_uinodes` mode is set to 0 for the vertices of textured
quads, and 1 if untextured.
* Added a condition to the fragment shader in `ui.wgsl` that only
multiplies by the `color` value from `textureSample` if the mode is
equal to 0.

crates/bevy_ui/src/render/mod.rs

@@ -552,6 +552,7 @@ struct UiVertex {
     pub position: [f32; 3],
     pub uv: [f32; 2],
     pub color: [f32; 4],
+    pub mode: u32,
 }
 
 #[derive(Resource)]
@@ -585,6 +586,9 @@ pub struct UiBatch {
     pub z: f32,
 }
 
+const TEXTURED_QUAD: u32 = 0;
+const UNTEXTURED_QUAD: u32 = 1;
+
 pub fn prepare_uinodes(
     mut commands: Commands,
     render_device: Res<RenderDevice>,
@@ -601,20 +605,33 @@ pub fn prepare_uinodes(
 
     let mut start = 0;
     let mut end = 0;
-    let mut current_batch_handle = Default::default();
+    let mut current_batch_image = DEFAULT_IMAGE_HANDLE.typed();
     let mut last_z = 0.0;
+
+    #[inline]
+    fn is_textured(image: &Handle<Image>) -> bool {
+        image.id() != DEFAULT_IMAGE_HANDLE.id()
+    }
+
     for extracted_uinode in &extracted_uinodes.uinodes {
-        if current_batch_handle != extracted_uinode.image {
-            if start != end {
+        let mode = if is_textured(&extracted_uinode.image) {
+            if current_batch_image.id() != extracted_uinode.image.id() {
+                if is_textured(&current_batch_image) && start != end {
                     commands.spawn(UiBatch {
                         range: start..end,
-                    image: current_batch_handle,
+                        image: current_batch_image,
                         z: last_z,
                     });
                     start = end;
                 }
-            current_batch_handle = extracted_uinode.image.clone_weak();
+                current_batch_image = extracted_uinode.image.clone_weak();
             }
+            TEXTURED_QUAD
+        } else {
+            // Untextured `UiBatch`es are never spawned within the loop.
+            // If all the `extracted_uinodes` are untextured a single untextured UiBatch will be spawned after the loop terminates.
+            UNTEXTURED_QUAD
+        };
 
         let mut uinode_rect = extracted_uinode.rect;
 
@@ -672,7 +689,7 @@ pub fn prepare_uinodes(
                 continue;
             }
         }
-        let uvs = if current_batch_handle.id() == DEFAULT_IMAGE_HANDLE.id() {
+        let uvs = if mode == UNTEXTURED_QUAD {
             [Vec2::ZERO, Vec2::X, Vec2::ONE, Vec2::Y]
         } else {
             let atlas_extent = extracted_uinode.atlas_size.unwrap_or(uinode_rect.max);
@@ -717,6 +734,7 @@ pub fn prepare_uinodes(
                 position: positions_clipped[i].into(),
                 uv: uvs[i].into(),
                 color,
+                mode,
             });
         }
 
@@ -728,7 +746,7 @@ pub fn prepare_uinodes(
     if start != end {
         commands.spawn(UiBatch {
             range: start..end,
-            image: current_batch_handle,
+            image: current_batch_image,
             z: last_z,
         });
     }

crates/bevy_ui/src/render/pipeline.rs

@@ -77,6 +77,8 @@ impl SpecializedRenderPipeline for UiPipeline {
                 VertexFormat::Float32x2,
                 // color
                 VertexFormat::Float32x4,
+                // mode
+                VertexFormat::Uint32,
             ],
         );
         let shader_defs = Vec::new();

crates/bevy_ui/src/render/ui.wgsl

@@ -1,11 +1,14 @@
 #import bevy_render::view
 
+const TEXTURED_QUAD: u32 = 0u;
+
 @group(0) @binding(0)
 var<uniform> view: View;
 
 struct VertexOutput {
     @location(0) uv: vec2<f32>,
     @location(1) color: vec4<f32>,
+    @location(3) mode: u32,
     @builtin(position) position: vec4<f32>,
 };
 
@@ -14,11 +17,13 @@ fn vertex(
     @location(0) vertex_position: vec3<f32>,
     @location(1) vertex_uv: vec2<f32>,
     @location(2) vertex_color: vec4<f32>,
+    @location(3) mode: u32,
 ) -> VertexOutput {
     var out: VertexOutput;
     out.uv = vertex_uv;
     out.position = view.view_proj * vec4<f32>(vertex_position, 1.0);
     out.color = vertex_color;
+    out.mode = mode;
     return out;
 }
 
@@ -29,7 +34,12 @@ var sprite_sampler: sampler;
 
 @fragment
 fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {
+    // textureSample can only be called in unform control flow, not inside an if branch.
     var color = textureSample(sprite_texture, sprite_sampler, in.uv);
+    if in.mode == TEXTURED_QUAD {
         color = in.color * color;
+    } else {
+        color = in.color;   
+    }
     return color;
 }