Enable MSAA with OIT

assets/shaders/oit_compatible_custom_material.wesl

@@ -14,6 +14,8 @@ import bevy_core_pipeline::oit::draw::oit_draw;
 @fragment
 fn fragment(
     in: VertexOutput,
+    @if(MATERIAL_OIT_ENABLED)
+        @builtin(sample_mask) sample_mask: u32,
 ) -> @location(0) vec4<f32> {
 
     // This produces a noise-like varying transparency.
@@ -24,7 +26,7 @@ fn fragment(
 
 @if(MATERIAL_OIT_ENABLED) {
     // The input color of `oit_draw` should be alpha-premultiplied.
-    oit_draw(in.position, vec4f(color.rgb * color.a, color.a));
+    oit_draw(in.position, vec4f(color.rgb * color.a, color.a), sample_mask);
     discard;
 }
 

assets/shaders/oit_compatible_extended_material.wesl

@@ -31,6 +31,8 @@ struct Colors {
 @fragment
 fn fragment(
     in: VertexOutput,
+    @if(MATERIAL_OIT_ENABLED)
+        @builtin(sample_mask) sample_mask: u32,
 ) -> FragmentOutput {
     var out: FragmentOutput;
 
@@ -43,14 +45,14 @@ fn fragment(
     let alpha_mode = pbr_input.material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_RESERVED_BITS;
     if alpha_mode == pbr_types::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_BLEND {
         // The fragments will only be drawn during the oit resolve pass.
-        oit_draw(in.position, vec4(out.color.rgb * out.color.a, out.color.a));
+        oit_draw(in.position, vec4(out.color.rgb * out.color.a, out.color.a), sample_mask);
         discard;
     }
     // Both `Premultiplied` and `Add` colors are premultiplied in `premultiply_alpha()`
     if alpha_mode == pbr_types::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_PREMULTIPLIED
         || alpha_mode == pbr_types::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_ADD {
         // The fragments will only be drawn during the oit resolve pass.
-        oit_draw(in.position, out.color);
+        oit_draw(in.position, out.color, sample_mask);
         discard;
     }
 }

crates/bevy_core_pipeline/src/oit/draw.wesl

@@ -2,11 +2,36 @@ import bevy_pbr::render::mesh_view_bindings::{view, oit_nodes_capacity, oit_node
 import bevy_pbr::render::mesh_view_types::OitFragmentNode;
 import bevy_pbr::prepass::utils as prepass_utils;
 
+@if(OIT_ENABLED && DEPTH_PREPASS)
+// Returns the subset of `sample_mask` whose samples are in front of the opaque geometry
+// recorded by the depth prepass.
+//
+// When MSAA is off, `sample_mask` is `0x1` and `prepass_depth` ignores the index, so this
+// collapses to a single depth test.
+fn prepass_visible_samples(position: vec4f, sample_mask: u32) -> u32 {
+    var remaining = sample_mask;
+    var visible = 0u;
+    while remaining != 0u {
+        let sample_index = firstTrailingBit(remaining);
+        // Clear the lowest set bit.
+        remaining &= remaining - 1u;
+        if position.z >= prepass_utils::prepass_depth(position, sample_index) {
+            visible |= 1u << sample_index;
+        }
+    }
+    return visible;
+}
+
 @if(OIT_ENABLED)
 // Add the fragment to the oit buffer
-fn oit_draw(position: vec4f, color: vec4f) {
+fn oit_draw(position: vec4f, color: vec4f, sample_mask: u32) {
+    @if(DEPTH_PREPASS)
+    let visible_samples = prepass_visible_samples(position, sample_mask);
+    @else
+    let visible_samples = sample_mask;
+
     @if(DEPTH_PREPASS)
-    if position.z < prepass_utils::prepass_depth(position, 0u) {
+    if visible_samples == 0u {
         return;
     }
     // Don't add fully transparent fragments to the list
@@ -33,6 +58,7 @@ fn oit_draw(position: vec4f, color: vec4f) {
     node.next = atomicExchange(&oit_heads[screen_index], new_node_index + 1u) - 1u;
     node.color = bevy_pbr::render::rgb9e5::vec3_to_rgb9e5_(color.rgb);
     node.depth_alpha = pack_24bit_depth_8bit_alpha(position.z, color.a);
+    node.sample_mask = visible_samples;
     oit_nodes[new_node_index] = node;
 }
 

crates/bevy_core_pipeline/src/oit/mod.rs

@@ -15,7 +15,6 @@ use bevy_render::{
         UninitBufferVec,
     },
     renderer::{RenderDevice, RenderQueue},
-    view::Msaa,
     Render, RenderApp, RenderStartup, RenderSystems,
 };
 use bevy_shader::load_shader_library;
@@ -71,7 +70,9 @@ impl Default for OrderIndependentTransparencySettings {
 /// To enable OIT you need to add the [`OrderIndependentTransparencySettings`] component to the camera and set `Material::enable_oit` to true.
 /// Currently the supported alpha modes are `AlphaMode::Blend`, `AlphaMode::Premultiplied` and `AlphaMode::Add`.
 ///
-/// If you want to use OIT for your custom material you need to call `oit_draw(position, color)` in your fragment shader.
+/// If you want to use OIT for your custom material you need to call `oit_draw(position, color, sample_mask)`
+/// in your fragment shader, where `sample_mask` is the fragment's `@builtin(sample_mask)` input
+/// (stored with the fragment so the resolve pass can composite each MSAA sample separately).
 /// You also need to make sure that your fragment shader doesn't output any colors.
 ///
 /// # Implementation details
@@ -91,8 +92,7 @@ impl Plugin for OrderIndependentTransparencyPlugin {
         app.add_plugins((
             ExtractComponentPlugin::<OrderIndependentTransparencySettings>::default(),
             OitResolvePlugin,
-        ))
-        .add_systems(Update, check_msaa);
+        ));
 
         let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
             return;
@@ -133,19 +133,15 @@ fn configure_camera_depth_usages(
     }
 }
 
-fn check_msaa(cameras: Query<&Msaa, With<OrderIndependentTransparencySettings>>) {
-    for msaa in &cameras {
-        if msaa.samples() > 1 {
-            panic!("MSAA is not supported when using OrderIndependentTransparency");
-        }
-    }
-}
-
 #[derive(Clone, Copy, ShaderType)]
 pub struct OitFragmentNode {
     pub color: u32,
     pub depth_alpha: u32,
     pub next: u32,
+    /// The MSAA sample coverage mask of the fragment. For example a triangle may
+    /// be covering only three of the four msaa samples. Used by the resolve pass
+    /// to composite each sample separately when MSAA is enabled.
+    pub sample_mask: u32,
 }
 
 /// Holds the buffers that contain the data of all OIT layers.

crates/bevy_core_pipeline/src/oit/resolve/mod.rs

@@ -14,15 +14,16 @@ use bevy_render::{
     camera::ExtractedCamera,
     render_resource::{
         binding_types::{
-            storage_buffer_read_only_sized, storage_buffer_sized, texture_depth_2d, uniform_buffer,
+            storage_buffer_read_only_sized, storage_buffer_sized, texture_depth_2d,
+            texture_depth_2d_multisampled, uniform_buffer,
         },
         BindGroup, BindGroupEntries, BindGroupLayoutDescriptor, BindGroupLayoutEntries,
         BlendComponent, BlendState, CachedRenderPipelineId, ColorTargetState, ColorWrites,
-        DownlevelFlags, FragmentState, PipelineCache, RenderPipelineDescriptor, ShaderStages,
-        TextureFormat,
+        DownlevelFlags, FragmentState, MultisampleState, PipelineCache, RenderPipelineDescriptor,
+        ShaderStages, TextureFormat,
     },
     renderer::{RenderAdapter, RenderDevice},
-    view::{ExtractedView, ViewUniform, ViewUniforms},
+    view::{ExtractedView, Msaa, ViewUniform, ViewUniforms},
     Render, RenderApp, RenderSystems,
 };
 use bevy_shader::ShaderDefVal;
@@ -105,6 +106,8 @@ pub struct OitResolvePipeline {
     pub view_bind_group_layout: BindGroupLayoutDescriptor,
     /// Depth bind group layout.
     pub oit_depth_bind_group_layout: BindGroupLayoutDescriptor,
+    /// Depth bind group layout for multisampled depth textures.
+    pub oit_depth_multisampled_bind_group_layout: BindGroupLayoutDescriptor,
 }
 
 impl OitResolvePipeline {
@@ -129,9 +132,17 @@ impl OitResolvePipeline {
             "oit_depth_bind_group_layout",
             &BindGroupLayoutEntries::single(ShaderStages::FRAGMENT, texture_depth_2d()),
         );
+        let oit_depth_multisampled_bind_group_layout = BindGroupLayoutDescriptor::new(
+            "oit_depth_multisampled_bind_group_layout",
+            &BindGroupLayoutEntries::single(
+                ShaderStages::FRAGMENT,
+                texture_depth_2d_multisampled(),
+            ),
+        );
         OitResolvePipeline {
             view_bind_group_layout,
             oit_depth_bind_group_layout,
+            oit_depth_multisampled_bind_group_layout,
         }
     }
 }
@@ -145,6 +156,7 @@ pub struct OitResolvePipelineKey {
     target_format: TextureFormat,
     sorted_fragment_max_count: u32,
     depth_prepass: bool,
+    msaa: Msaa,
 }
 
 pub fn queue_oit_resolve_pipeline(
@@ -157,6 +169,7 @@ pub fn queue_oit_resolve_pipeline(
             &ExtractedView,
             &OrderIndependentTransparencySettings,
             Has<DepthPrepass>,
+            &Msaa,
         ),
         (
             With<OrderIndependentTransparencySettings>,
@@ -170,12 +183,13 @@ pub fn queue_oit_resolve_pipeline(
     mut cached_pipeline_id: Local<EntityHashMap<(OitResolvePipelineKey, CachedRenderPipelineId)>>,
 ) {
     let mut current_view_entities = EntityHashSet::default();
-    for (e, view, oit_settings, depth_prepass) in &cameras {
+    for (e, view, oit_settings, depth_prepass, msaa) in &cameras {
         current_view_entities.insert(e);
         let key = OitResolvePipelineKey {
             target_format: view.target_format,
             sorted_fragment_max_count: oit_settings.sorted_fragment_max_count,
             depth_prepass,
+            msaa: *msaa,
         };
 
         if let Some((cached_key, id)) = cached_pipeline_id.get(&e)
@@ -216,15 +230,28 @@ fn specialize_oit_resolve_pipeline(
         "SORTED_FRAGMENT_MAX_COUNT".into(),
         key.sorted_fragment_max_count,
     )];
+    if key.msaa.samples() > 1 {
+        shader_defs.push(ShaderDefVal::Bool("MULTISAMPLED".into(), true));
+    }
     if key.depth_prepass {
         shader_defs.push(ShaderDefVal::Bool("DEPTH_PREPASS".into(), true));
+    } else if key.msaa.samples() > 1 {
+        layout.push(
+            resolve_pipeline
+                .oit_depth_multisampled_bind_group_layout
+                .clone(),
+        );
     } else {
         layout.push(resolve_pipeline.oit_depth_bind_group_layout.clone());
     }
 
     RenderPipelineDescriptor {
         label: Some("oit_resolve_pipeline".into()),
         layout,
+        multisample: MultisampleState {
+            count: key.msaa.samples(),
+            ..default()
+        },
         fragment: Some(FragmentState {
             shader: load_embedded_asset!(asset_server, "oit_resolve.wesl"),
             shader_defs,

crates/bevy_core_pipeline/src/oit/resolve/node.rs

@@ -5,7 +5,7 @@ use bevy_render::{
     diagnostic::RecordDiagnostics,
     render_resource::{BindGroupEntries, PipelineCache, RenderPassDescriptor},
     renderer::{RenderContext, ViewQuery},
-    view::{ViewDepthStencilTexture, ViewTarget, ViewUniformOffset},
+    view::{Msaa, ViewDepthStencilTexture, ViewTarget, ViewUniformOffset},
 };
 
 use crate::prepass::DepthPrepass;
@@ -21,6 +21,7 @@ pub fn oit_resolve(
         &ViewDepthStencilTexture,
         Option<&MainPassResolutionOverride>,
         Has<DepthPrepass>,
+        &Msaa,
     )>,
     resolve_pipeline: Option<Res<OitResolvePipeline>>,
     bind_group: Option<Res<OitResolveBindGroup>>,
@@ -35,6 +36,7 @@ pub fn oit_resolve(
         depth,
         resolution_override,
         depth_prepass,
+        msaa,
     ) = view.into_inner();
 
     // This *must* run after main_transparent_pass_3d to reset the `oit_atomic_counter` and `oit_heads` buffer
@@ -53,9 +55,14 @@ pub fn oit_resolve(
     };
 
     let depth_bind_group = if !depth_prepass {
+        let depth_layout = if msaa.samples() > 1 {
+            &resolve_pipeline.oit_depth_multisampled_bind_group_layout
+        } else {
+            &resolve_pipeline.oit_depth_bind_group_layout
+        };
         Some(ctx.render_device().create_bind_group(
             "oit_resolve_depth_bind_group",
-            &pipeline_cache.get_bind_group_layout(&resolve_pipeline.oit_depth_bind_group_layout),
+            &pipeline_cache.get_bind_group_layout(depth_layout),
             &BindGroupEntries::single(depth_view),
         ))
     } else {

crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wesl

@@ -9,7 +9,9 @@ import bevy_pbr::render::mesh_view_types::OitFragmentNode;
 @group(0) @binding(2) var<storage, read_write> heads: array<u32>; // No need to be atomic
 @group(0) @binding(3) var<storage, read_write> atomic_counter: u32; // No need to be atomic
 
-@if(!DEPTH_PREPASS)
+@if(!DEPTH_PREPASS && MULTISAMPLED)
+@group(1) @binding(0) var depth: texture_depth_multisampled_2d;
+@if(!DEPTH_PREPASS && !MULTISAMPLED)
 @group(1) @binding(0) var depth: texture_depth_2d;
 
 struct OitFragment {
@@ -26,7 +28,13 @@ const LINKED_LIST_END_SENTINEL: u32 = 0xFFFFFFFFu;
 const SORTED_FRAGMENT_MAX_COUNT: u32 = constants::SORTED_FRAGMENT_MAX_COUNT;
 
 @fragment
-fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
+fn fragment(
+    in: FullscreenVertexOutput,
+    @if(MULTISAMPLED)
+        @builtin(sample_index) sample_index: u32,
+) -> @location(0) vec4<f32> {
+    @if(!MULTISAMPLED)
+    let sample_index = 0u;
     atomic_counter = 0u;
     let screen_index = u32(floor(in.position.x) + floor(in.position.y) * view.viewport.z);
 
@@ -43,16 +51,16 @@ fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
         // This is necessary because early z doesn't seem to trigger in the transparent pass.
         // This should be done during the draw pass so those fragments simply don't exist in the list,
         // but this requires a bigger refactor
-        let d = textureLoad(depth, vec2<i32>(in.position.xy), 0);
+        let d = textureLoad(depth, vec2<i32>(in.position.xy), i32(sample_index));
         @else
         let d = 0.0;
-        let color = resolve(head, d);
+        let color = resolve(head, d, sample_index);
         heads[screen_index] = 0u; // LINKED_LIST_END_SENTINEL + 1u;
         return color;
     }
 }
 
-fn resolve(head: u32, opaque_depth: f32) -> vec4<f32> {
+fn resolve(head: u32, opaque_depth: f32, sample_index: u32) -> vec4<f32> {
     // Contains all the colors and depth for this specific fragment
     // Fragments are sorted from front to back (depth values are in descending order)
     // This should make insertion sort slightly faster
@@ -69,6 +77,12 @@ fn resolve(head: u32, opaque_depth: f32) -> vec4<f32> {
         let fragment_node = nodes[current_node];
         current_node = fragment_node.next;
 
+        @if(MULTISAMPLED)
+        // Only composite fragments that cover this sample.
+        if (fragment_node.sample_mask & (1u << sample_index)) == 0u {
+            continue;
+        }
+
         @if(!DEPTH_PREPASS)
         // depth testing
         if fragment_node.depth_alpha < packed_opaque_depth {

crates/bevy_pbr/src/render/mesh_view_types.wesl

@@ -225,6 +225,7 @@ struct OitFragmentNode {
     color: u32,
     depth_alpha: u32,
     next: u32,
+    sample_mask: u32,
 };
 
 struct ClusteredDecal {

crates/bevy_pbr/src/render/pbr.wesl

@@ -39,6 +39,8 @@ fn fragment(
     vertex_output: VertexOutput,
 @if(!MESHLET_MESH_MATERIAL_PASS)
     @builtin(front_facing) is_front: bool,
+@if(MATERIAL_OIT_ENABLED)
+    @builtin(sample_mask) sample_mask: u32,
 ) -> FragmentOutput {
 @if(MESHLET_MESH_MATERIAL_PASS)
     let vertex_output = resolve_vertex_output(frag_coord);
@@ -95,14 +97,14 @@ fn fragment(
     let alpha_mode = pbr_input.material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_RESERVED_BITS;
     if alpha_mode == pbr_types::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_BLEND {
         // The fragments will only be drawn during the oit resolve pass.
-        oit_draw(in.position, vec4(out.color.rgb * out.color.a, out.color.a));
+        oit_draw(in.position, vec4(out.color.rgb * out.color.a, out.color.a), sample_mask);
         discard;
     }
     // Both `Premultiplied` and `Add` colors are premultiplied in `premultiply_alpha()`
     if alpha_mode == pbr_types::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_PREMULTIPLIED
         || alpha_mode == pbr_types::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_ADD {
         // The fragments will only be drawn during the oit resolve pass.
-        oit_draw(in.position, out.color);
+        oit_draw(in.position, out.color, sample_mask);
         discard;
     }
 }