Add AlphaMask2d phase (#14724)

# Objective

- Bevy now supports an opaque phase for mesh2d, but it's very common for
2d textures to have a transparent alpha channel.

## Solution

- Add an alpha mask phase identical to the one in 3d. It will do the
alpha masking in the shader before drawing the mesh.
- Uses the BinnedRenderPhase
- Since it's an opaque draw it also correctly writes to depth

## Testing

- Tested the mes2d_alpha_mode example and the bevymark example with
alpha mask mode enabled

---

## Showcase


![image](https://github.com/user-attachments/assets/9e5e4561-d0a7-4aa3-b049-d4b1247d5ed4)

The white logo on the right is rendered with alpha mask enabled.

Running the bevymark example I can get 65fps for 120k mesh2d all using
alpha mask.

## Notes

This is one more step for mesh2d improvements
https://github.com/bevyengine/bevy/issues/13265

---------

Co-authored-by: Kristoffer Søholm <k.soeholm@gmail.com>

crates/bevy_core_pipeline/src/core_2d/main_opaque_pass_2d_node.rs

@@ -13,7 +13,10 @@ use bevy_utils::tracing::error;
 #[cfg(feature = "trace")]
 use bevy_utils::tracing::info_span;
 
-/// A [`bevy_render::render_graph::Node`] that runs the [`Opaque2d`] [`ViewBinnedRenderPhases`]
+use super::AlphaMask2d;
+
+/// A [`bevy_render::render_graph::Node`] that runs the
+/// [`Opaque2d`] [`ViewBinnedRenderPhases`] and [`AlphaMask2d`] [`ViewBinnedRenderPhases`]
 #[derive(Default)]
 pub struct MainOpaquePass2dNode;
 impl ViewNode for MainOpaquePass2dNode {
@@ -30,7 +33,10 @@ impl ViewNode for MainOpaquePass2dNode {
         (camera, target, depth): QueryItem<'w, Self::ViewQuery>,
         world: &'w World,
     ) -> Result<(), NodeRunError> {
-        let Some(opaque_phases) = world.get_resource::<ViewBinnedRenderPhases<Opaque2d>>() else {
+        let (Some(opaque_phases), Some(alpha_mask_phases)) = (
+            world.get_resource::<ViewBinnedRenderPhases<Opaque2d>>(),
+            world.get_resource::<ViewBinnedRenderPhases<AlphaMask2d>>(),
+        ) else {
             return Ok(());
         };
 
@@ -40,7 +46,10 @@ impl ViewNode for MainOpaquePass2dNode {
         let depth_stencil_attachment = Some(depth.get_attachment(StoreOp::Store));
 
         let view_entity = graph.view_entity();
-        let Some(opaque_phase) = opaque_phases.get(&view_entity) else {
+        let (Some(opaque_phase), Some(alpha_mask_phase)) = (
+            opaque_phases.get(&view_entity),
+            alpha_mask_phases.get(&view_entity),
+        ) else {
             return Ok(());
         };
         render_context.add_command_buffer_generation_task(move |render_device| {
@@ -77,6 +86,15 @@ impl ViewNode for MainOpaquePass2dNode {
                 }
             }
 
+            // Alpha mask draws
+            if !alpha_mask_phase.is_empty() {
+                #[cfg(feature = "trace")]
+                let _alpha_mask_main_pass_2d_span = info_span!("alpha_mask_main_pass_2d").entered();
+                if let Err(err) = alpha_mask_phase.render(&mut render_pass, world, view_entity) {
+                    error!("Error encountered while rendering the 2d alpha mask phase {err:?}");
+                }
+            }
+
             pass_span.end(&mut render_pass);
             drop(render_pass);
             command_encoder.finish()

crates/bevy_core_pipeline/src/core_2d/mod.rs

@@ -78,9 +78,11 @@ impl Plugin for Core2dPlugin {
         };
         render_app
             .init_resource::<DrawFunctions<Opaque2d>>()
+            .init_resource::<DrawFunctions<AlphaMask2d>>()
             .init_resource::<DrawFunctions<Transparent2d>>()
             .init_resource::<ViewSortedRenderPhases<Transparent2d>>()
             .init_resource::<ViewBinnedRenderPhases<Opaque2d>>()
+            .init_resource::<ViewBinnedRenderPhases<AlphaMask2d>>()
             .add_systems(ExtractSchedule, extract_core_2d_camera_phases)
             .add_systems(
                 Render,
@@ -147,8 +149,6 @@ pub struct Opaque2dBinKey {
     /// the ID of another type of asset.
     pub asset_id: UntypedAssetId,
     /// The ID of a bind group specific to the material.
-    ///
-    /// In the case of PBR, this is the `MaterialBindGroupId`.
     pub material_bind_group_id: Option<BindGroupId>,
 }
 
@@ -207,6 +207,92 @@ impl CachedRenderPipelinePhaseItem for Opaque2d {
     }
 }
 
+/// Alpha mask 2D [`BinnedPhaseItem`]s.
+pub struct AlphaMask2d {
+    /// The key, which determines which can be batched.
+    pub key: AlphaMask2dBinKey,
+    /// An entity from which data will be fetched, including the mesh if
+    /// applicable.
+    pub representative_entity: Entity,
+    /// The ranges of instances.
+    pub batch_range: Range<u32>,
+    /// An extra index, which is either a dynamic offset or an index in the
+    /// indirect parameters list.
+    pub extra_index: PhaseItemExtraIndex,
+}
+
+/// Data that must be identical in order to batch phase items together.
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct AlphaMask2dBinKey {
+    /// The identifier of the render pipeline.
+    pub pipeline: CachedRenderPipelineId,
+    /// The function used to draw.
+    pub draw_function: DrawFunctionId,
+    /// The asset that this phase item is associated with.
+    ///
+    /// Normally, this is the ID of the mesh, but for non-mesh items it might be
+    /// the ID of another type of asset.
+    pub asset_id: UntypedAssetId,
+    /// The ID of a bind group specific to the material.
+    pub material_bind_group_id: Option<BindGroupId>,
+}
+
+impl PhaseItem for AlphaMask2d {
+    #[inline]
+    fn entity(&self) -> Entity {
+        self.representative_entity
+    }
+
+    #[inline]
+    fn draw_function(&self) -> DrawFunctionId {
+        self.key.draw_function
+    }
+
+    #[inline]
+    fn batch_range(&self) -> &Range<u32> {
+        &self.batch_range
+    }
+
+    #[inline]
+    fn batch_range_mut(&mut self) -> &mut Range<u32> {
+        &mut self.batch_range
+    }
+
+    fn extra_index(&self) -> PhaseItemExtraIndex {
+        self.extra_index
+    }
+
+    fn batch_range_and_extra_index_mut(&mut self) -> (&mut Range<u32>, &mut PhaseItemExtraIndex) {
+        (&mut self.batch_range, &mut self.extra_index)
+    }
+}
+
+impl BinnedPhaseItem for AlphaMask2d {
+    type BinKey = AlphaMask2dBinKey;
+
+    fn new(
+        key: Self::BinKey,
+        representative_entity: Entity,
+        batch_range: Range<u32>,
+        extra_index: PhaseItemExtraIndex,
+    ) -> Self {
+        AlphaMask2d {
+            key,
+            representative_entity,
+            batch_range,
+            extra_index,
+        }
+    }
+}
+
+impl CachedRenderPipelinePhaseItem for AlphaMask2d {
+    #[inline]
+    fn cached_pipeline(&self) -> CachedRenderPipelineId {
+        self.key.pipeline
+    }
+}
+
+/// Transparent 2D [`SortedPhaseItem`]s.
 pub struct Transparent2d {
     pub sort_key: FloatOrd,
     pub entity: Entity,
@@ -274,6 +360,7 @@ pub fn extract_core_2d_camera_phases(
     mut commands: Commands,
     mut transparent_2d_phases: ResMut<ViewSortedRenderPhases<Transparent2d>>,
     mut opaque_2d_phases: ResMut<ViewBinnedRenderPhases<Opaque2d>>,
+    mut alpha_mask_2d_phases: ResMut<ViewBinnedRenderPhases<AlphaMask2d>>,
     cameras_2d: Extract<Query<(Entity, &Camera), With<Camera2d>>>,
     mut live_entities: Local<EntityHashSet>,
 ) {
@@ -287,6 +374,7 @@ pub fn extract_core_2d_camera_phases(
         commands.get_or_spawn(entity);
         transparent_2d_phases.insert_or_clear(entity);
         opaque_2d_phases.insert_or_clear(entity);
+        alpha_mask_2d_phases.insert_or_clear(entity);
 
         live_entities.insert(entity);
     }
@@ -294,6 +382,7 @@ pub fn extract_core_2d_camera_phases(
     // Clear out all dead views.
     transparent_2d_phases.retain(|camera_entity, _| live_entities.contains(camera_entity));
     opaque_2d_phases.retain(|camera_entity, _| live_entities.contains(camera_entity));
+    alpha_mask_2d_phases.retain(|camera_entity, _| live_entities.contains(camera_entity));
 }
 
 pub fn prepare_core_2d_depth_textures(

crates/bevy_sprite/src/mesh2d/color_material.rs

@@ -98,16 +98,29 @@ bitflags::bitflags! {
     #[repr(transparent)]
     pub struct ColorMaterialFlags: u32 {
         const TEXTURE                    = 1 << 0;
+        /// Bitmask reserving bits for the [`AlphaMode2d`]
+        /// Values are just sequential values bitshifted into
+        /// the bitmask, and can range from 0 to 3.
+        const ALPHA_MODE_RESERVED_BITS   = Self::ALPHA_MODE_MASK_BITS << Self::ALPHA_MODE_SHIFT_BITS;
+        const ALPHA_MODE_OPAQUE          = 0 << Self::ALPHA_MODE_SHIFT_BITS;
+        const ALPHA_MODE_MASK            = 1 << Self::ALPHA_MODE_SHIFT_BITS;
+        const ALPHA_MODE_BLEND           = 2 << Self::ALPHA_MODE_SHIFT_BITS;
         const NONE                       = 0;
         const UNINITIALIZED              = 0xFFFF;
     }
 }
 
+impl ColorMaterialFlags {
+    const ALPHA_MODE_MASK_BITS: u32 = 0b11;
+    const ALPHA_MODE_SHIFT_BITS: u32 = 32 - Self::ALPHA_MODE_MASK_BITS.count_ones();
+}
+
 /// The GPU representation of the uniform data of a [`ColorMaterial`].
 #[derive(Clone, Default, ShaderType)]
 pub struct ColorMaterialUniform {
     pub color: Vec4,
     pub flags: u32,
+    pub alpha_cutoff: f32,
 }
 
 impl AsBindGroupShaderType<ColorMaterialUniform> for ColorMaterial {
@@ -117,9 +130,20 @@ impl AsBindGroupShaderType<ColorMaterialUniform> for ColorMaterial {
             flags |= ColorMaterialFlags::TEXTURE;
         }
 
+        // Defaults to 0.5 like in 3d
+        let mut alpha_cutoff = 0.5;
+        match self.alpha_mode {
+            AlphaMode2d::Opaque => flags |= ColorMaterialFlags::ALPHA_MODE_OPAQUE,
+            AlphaMode2d::Mask(c) => {
+                alpha_cutoff = c;
+                flags |= ColorMaterialFlags::ALPHA_MODE_MASK;
+            }
+            AlphaMode2d::Blend => flags |= ColorMaterialFlags::ALPHA_MODE_BLEND,
+        };
         ColorMaterialUniform {
             color: LinearRgba::from(self.color).to_f32_array().into(),
             flags: flags.bits(),
+            alpha_cutoff,
         }
     }
 }

crates/bevy_sprite/src/mesh2d/color_material.wgsl

@@ -11,8 +11,14 @@ struct ColorMaterial {
     color: vec4<f32>,
     // 'flags' is a bit field indicating various options. u32 is 32 bits so we have up to 32 options.
     flags: u32,
+    alpha_cutoff: f32,
 };
+
 const COLOR_MATERIAL_FLAGS_TEXTURE_BIT: u32              = 1u;
+const COLOR_MATERIAL_FLAGS_ALPHA_MODE_RESERVED_BITS: u32 = 3221225472u; // (0b11u32 << 30)
+const COLOR_MATERIAL_FLAGS_ALPHA_MODE_OPAQUE: u32        = 0u;          // (0u32 << 30)
+const COLOR_MATERIAL_FLAGS_ALPHA_MODE_MASK: u32          = 1073741824u; // (1u32 << 30)
+const COLOR_MATERIAL_FLAGS_ALPHA_MODE_BLEND: u32         = 2147483648u; // (2u32 << 30)
 
 @group(2) @binding(0) var<uniform> material: ColorMaterial;
 @group(2) @binding(1) var texture: texture_2d<f32>;
@@ -23,14 +29,41 @@ fn fragment(
     mesh: VertexOutput,
 ) -> @location(0) vec4<f32> {
     var output_color: vec4<f32> = material.color;
+
 #ifdef VERTEX_COLORS
     output_color = output_color * mesh.color;
 #endif
+
     if ((material.flags & COLOR_MATERIAL_FLAGS_TEXTURE_BIT) != 0u) {
         output_color = output_color * textureSample(texture, texture_sampler, mesh.uv);
     }
+
+    output_color = alpha_discard(material, output_color);
+
 #ifdef TONEMAP_IN_SHADER
     output_color = tonemapping::tone_mapping(output_color, view.color_grading);
 #endif
     return output_color;
 }
+
+fn alpha_discard(material: ColorMaterial, output_color: vec4<f32>) -> vec4<f32> {
+    var color = output_color;
+    let alpha_mode = material.flags & COLOR_MATERIAL_FLAGS_ALPHA_MODE_RESERVED_BITS;
+    if alpha_mode == COLOR_MATERIAL_FLAGS_ALPHA_MODE_OPAQUE {
+        // NOTE: If rendering as opaque, alpha should be ignored so set to 1.0
+        color.a = 1.0;
+    }
+#ifdef MAY_DISCARD
+    else if alpha_mode == COLOR_MATERIAL_FLAGS_ALPHA_MODE_MASK {
+       if color.a >= material.alpha_cutoff {
+            // NOTE: If rendering as masked alpha and >= the cutoff, render as fully opaque
+            color.a = 1.0;
+        } else {
+            // NOTE: output_color.a < in.material.alpha_cutoff should not be rendered
+            discard;
+        }
+    }
+#endif // MAY_DISCARD
+
+    return color;
+}

crates/bevy_sprite/src/mesh2d/material.rs

@@ -1,7 +1,7 @@
 use bevy_app::{App, Plugin};
 use bevy_asset::{Asset, AssetApp, AssetId, AssetServer, Handle};
 use bevy_core_pipeline::{
-    core_2d::{Opaque2d, Opaque2dBinKey, Transparent2d},
+    core_2d::{AlphaMask2d, AlphaMask2dBinKey, Opaque2d, Opaque2dBinKey, Transparent2d},
     tonemapping::{DebandDither, Tonemapping},
 };
 use bevy_derive::{Deref, DerefMut};
@@ -149,6 +149,15 @@ pub enum AlphaMode2d {
     /// Base color alpha values are overridden to be fully opaque (1.0).
     #[default]
     Opaque,
+    /// Reduce transparency to fully opaque or fully transparent
+    /// based on a threshold.
+    ///
+    /// Compares the base color alpha value to the specified threshold.
+    /// If the value is below the threshold,
+    /// considers the color to be fully transparent (alpha is set to 0.0).
+    /// If it is equal to or above the threshold,
+    /// considers the color to be fully opaque (alpha is set to 1.0).
+    Mask(f32),
     /// The base color alpha value defines the opacity of the color.
     /// Standard alpha-blending is used to blend the fragment's color
     /// with the color behind it.
@@ -176,6 +185,7 @@ where
         if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
             render_app
                 .add_render_command::<Opaque2d, DrawMaterial2d<M>>()
+                .add_render_command::<AlphaMask2d, DrawMaterial2d<M>>()
                 .add_render_command::<Transparent2d, DrawMaterial2d<M>>()
                 .init_resource::<RenderMaterial2dInstances<M>>()
                 .init_resource::<SpecializedMeshPipelines<Material2dPipeline<M>>>()
@@ -374,6 +384,7 @@ impl<P: PhaseItem, M: Material2d, const I: usize> RenderCommand<P>
 pub const fn alpha_mode_pipeline_key(alpha_mode: AlphaMode2d) -> Mesh2dPipelineKey {
     match alpha_mode {
         AlphaMode2d::Blend => Mesh2dPipelineKey::BLEND_ALPHA,
+        AlphaMode2d::Mask(_) => Mesh2dPipelineKey::MAY_DISCARD,
         _ => Mesh2dPipelineKey::NONE,
     }
 }
@@ -396,6 +407,7 @@ pub const fn tonemapping_pipeline_key(tonemapping: Tonemapping) -> Mesh2dPipelin
 #[allow(clippy::too_many_arguments)]
 pub fn queue_material2d_meshes<M: Material2d>(
     opaque_draw_functions: Res<DrawFunctions<Opaque2d>>,
+    alpha_mask_draw_functions: Res<DrawFunctions<AlphaMask2d>>,
     transparent_draw_functions: Res<DrawFunctions<Transparent2d>>,
     material2d_pipeline: Res<Material2dPipeline<M>>,
     mut pipelines: ResMut<SpecializedMeshPipelines<Material2dPipeline<M>>>,
@@ -406,6 +418,7 @@ pub fn queue_material2d_meshes<M: Material2d>(
     render_material_instances: Res<RenderMaterial2dInstances<M>>,
     mut transparent_render_phases: ResMut<ViewSortedRenderPhases<Transparent2d>>,
     mut opaque_render_phases: ResMut<ViewBinnedRenderPhases<Opaque2d>>,
+    mut alpha_mask_render_phases: ResMut<ViewBinnedRenderPhases<AlphaMask2d>>,
     mut views: Query<(
         Entity,
         &ExtractedView,
@@ -425,13 +438,16 @@ pub fn queue_material2d_meshes<M: Material2d>(
         let Some(transparent_phase) = transparent_render_phases.get_mut(&view_entity) else {
             continue;
         };
-
         let Some(opaque_phase) = opaque_render_phases.get_mut(&view_entity) else {
             continue;
         };
+        let Some(alpha_mask_phase) = alpha_mask_render_phases.get_mut(&view_entity) else {
+            continue;
+        };
 
         let draw_transparent_2d = transparent_draw_functions.read().id::<DrawMaterial2d<M>>();
         let draw_opaque_2d = opaque_draw_functions.read().id::<DrawMaterial2d<M>>();
+        let draw_alpha_mask_2d = alpha_mask_draw_functions.read().id::<DrawMaterial2d<M>>();
 
         let mut view_key = Mesh2dPipelineKey::from_msaa_samples(msaa.samples())
             | Mesh2dPipelineKey::from_hdr(view.hdr);
@@ -497,6 +513,19 @@ pub fn queue_material2d_meshes<M: Material2d>(
                         BinnedRenderPhaseType::mesh(mesh_instance.automatic_batching),
                     );
                 }
+                AlphaMode2d::Mask(_) => {
+                    let bin_key = AlphaMask2dBinKey {
+                        pipeline: pipeline_id,
+                        draw_function: draw_alpha_mask_2d,
+                        asset_id: mesh_instance.mesh_asset_id.into(),
+                        material_bind_group_id: material_2d.get_bind_group_id().0,
+                    };
+                    alpha_mask_phase.add(
+                        bin_key,
+                        *visible_entity,
+                        BinnedRenderPhaseType::mesh(mesh_instance.automatic_batching),
+                    );
+                }
                 AlphaMode2d::Blend => {
                     transparent_phase.add(Transparent2d {
                         entity: *visible_entity,

crates/bevy_sprite/src/mesh2d/mesh.rs

@@ -1,9 +1,11 @@
 use bevy_app::Plugin;
 use bevy_asset::{load_internal_asset, AssetId, Handle};
 
-use bevy_core_pipeline::core_2d::{Camera2d, Opaque2d, Transparent2d, CORE_2D_DEPTH_FORMAT};
+use bevy_core_pipeline::{
-use bevy_core_pipeline::tonemapping::{
+    core_2d::{AlphaMask2d, Camera2d, Opaque2d, Transparent2d, CORE_2D_DEPTH_FORMAT},
+    tonemapping::{
         get_lut_bind_group_layout_entries, get_lut_bindings, Tonemapping, TonemappingLuts,
+    },
 };
 use bevy_derive::{Deref, DerefMut};
 use bevy_ecs::{
@@ -114,6 +116,8 @@ impl Plugin for Mesh2dRenderPlugin {
                     (
                         batch_and_prepare_binned_render_phase::<Opaque2d, Mesh2dPipeline>
                             .in_set(RenderSet::PrepareResources),
+                        batch_and_prepare_binned_render_phase::<AlphaMask2d, Mesh2dPipeline>
+                            .in_set(RenderSet::PrepareResources),
                         batch_and_prepare_sorted_render_phase::<Transparent2d, Mesh2dPipeline>
                             .in_set(RenderSet::PrepareResources),
                         write_batched_instance_buffer::<Mesh2dPipeline>
@@ -475,6 +479,7 @@ bitflags::bitflags! {
         const TONEMAP_IN_SHADER                 = 1 << 1;
         const DEBAND_DITHER                     = 1 << 2;
         const BLEND_ALPHA                       = 1 << 3;
+        const MAY_DISCARD                       = 1 << 4;
         const MSAA_RESERVED_BITS                = Self::MSAA_MASK_BITS << Self::MSAA_SHIFT_BITS;
         const PRIMITIVE_TOPOLOGY_RESERVED_BITS  = Self::PRIMITIVE_TOPOLOGY_MASK_BITS << Self::PRIMITIVE_TOPOLOGY_SHIFT_BITS;
         const TONEMAP_METHOD_RESERVED_BITS      = Self::TONEMAP_METHOD_MASK_BITS << Self::TONEMAP_METHOD_SHIFT_BITS;
@@ -619,6 +624,10 @@ impl SpecializedMeshPipeline for Mesh2dPipeline {
             }
         }
 
+        if key.contains(Mesh2dPipelineKey::MAY_DISCARD) {
+            shader_defs.push("MAY_DISCARD".into());
+        }
+
         let vertex_buffer_layout = layout.0.get_layout(&vertex_attributes)?;
 
         let format = match key.contains(Mesh2dPipelineKey::HDR) {

examples/2d/mesh2d_alpha_mode.rs

@@ -59,16 +59,16 @@ fn setup(
         ..default()
     });
 
-    // Test the interaction between opaque and transparent meshes
+    // Test the interaction between opaque/mask and transparent meshes
     // The white sprite should be:
-    // - fully opaque
+    // - only the icon is opaque but background is transparent
     // - on top of the green sprite
     // - behind the blue sprite
     commands.spawn(MaterialMesh2dBundle {
         mesh: mesh_handle.clone().into(),
         material: materials.add(ColorMaterial {
             color: WHITE.into(),
-            alpha_mode: AlphaMode2d::Opaque,
+            alpha_mode: AlphaMode2d::Mask(0.5),
             texture: Some(texture_handle.clone()),
         }),
         transform: Transform::from_xyz(200.0, 0.0, 0.0),

examples/stress_tests/bevymark.rs

@@ -103,6 +103,7 @@ enum AlphaMode {
     Opaque,
     #[default]
     Blend,
+    AlphaMask,
 }
 
 impl FromStr for AlphaMode {
@@ -112,8 +113,9 @@ impl FromStr for AlphaMode {
         match s {
             "opaque" => Ok(Self::Opaque),
             "blend" => Ok(Self::Blend),
+            "alpha_mask" => Ok(Self::AlphaMask),
             _ => Err(format!(
-                "Unknown alpha mode: '{s}', valid modes: 'opaque', 'blend'"
+                "Unknown alpha mode: '{s}', valid modes: 'opaque', 'blend', 'alpha_mask'"
             )),
         }
     }
@@ -601,6 +603,7 @@ fn init_materials(
     let alpha_mode = match args.alpha_mode {
         AlphaMode::Opaque => AlphaMode2d::Opaque,
         AlphaMode::Blend => AlphaMode2d::Blend,
+        AlphaMode::AlphaMask => AlphaMode2d::Mask(0.5),
     };
 
     let mut materials = Vec::with_capacity(capacity);