Vulkan support for SHADER_EARLY_DEPTH_TEST and fix to conservative depth optimizations

CHANGELOG.md

@@ -81,6 +81,7 @@ Naga now infers the correct binding layout when a resource appears only in an as
 
 - Mark `readonly_and_readwrite_storage_textures` & `packed_4x8_integer_dot_product` language extensions as implemented. By @teoxoy in [#7543](https://github.com/gfx-rs/wgpu/pull/7543)
 - `naga::back::hlsl::Writer::new` has a new `pipeline_options` argument. `hlsl::PipelineOptions::default()` can be passed as a default. The `shader_stage` and `entry_point` members of `pipeline_options` can be used to write only a single entry point when using the HLSL and MSL backends (GLSL and SPIR-V already had this functionality). The Metal and DX12 HALs now write only a single entry point when loading shaders. By @andyleiserson in [#7626](https://github.com/gfx-rs/wgpu/pull/7626).
+- Implemented `early_depth_test` for SPIR-V backend, enabling `SHADER_EARLY_DEPTH_TEST` for Vulkan. Additionally, fixed conservative depth optimizations when using `early_depth_test`. The syntax for forcing early depth tests is now `@early_depth_test(force)` instead of `@early_depth_test`. By @dzamkov in [#7676](https://github.com/gfx-rs/wgpu/pull/7676).
 
 #### D3D12
 

naga/src/back/glsl/features.rs

@@ -300,14 +300,16 @@ impl<W> Writer<'_, W> {
     pub(super) fn collect_required_features(&mut self) -> BackendResult {
         let ep_info = self.info.get_entry_point(self.entry_point_idx as usize);
 
-        if let Some(depth_test) = self.entry_point.early_depth_test {
-            // If IMAGE_LOAD_STORE is supported for this version of GLSL
-            if self.options.version.supports_early_depth_test() {
-                self.features.request(Features::IMAGE_LOAD_STORE);
-            }
-
-            if depth_test.conservative.is_some() {
-                self.features.request(Features::CONSERVATIVE_DEPTH);
+        if let Some(early_depth_test) = self.entry_point.early_depth_test {
+            match early_depth_test {
+                crate::EarlyDepthTest::Force => {
+                    if self.options.version.supports_early_depth_test() {
+                        self.features.request(Features::IMAGE_LOAD_STORE);
+                    }
+                }
+                crate::EarlyDepthTest::Allow { .. } => {
+                    self.features.request(Features::CONSERVATIVE_DEPTH);
+                }
             }
         }
 

naga/src/back/glsl/mod.rs

@@ -750,22 +750,23 @@ impl<'a, W: Write> Writer<'a, W> {
         }
 
         // Enable early depth tests if needed
-        if let Some(depth_test) = self.entry_point.early_depth_test {
+        if let Some(early_depth_test) = self.entry_point.early_depth_test {
             // If early depth test is supported for this version of GLSL
             if self.options.version.supports_early_depth_test() {
-                writeln!(self.out, "layout(early_fragment_tests) in;")?;
-
-                if let Some(conservative) = depth_test.conservative {
-                    use crate::ConservativeDepth as Cd;
-
-                    let depth = match conservative {
-                        Cd::GreaterEqual => "greater",
-                        Cd::LessEqual => "less",
-                        Cd::Unchanged => "unchanged",
-                    };
-                    writeln!(self.out, "layout (depth_{depth}) out float gl_FragDepth;")?;
+                match early_depth_test {
+                    crate::EarlyDepthTest::Force => {
+                        writeln!(self.out, "layout(early_fragment_tests) in;")?;
+                    }
+                    crate::EarlyDepthTest::Allow { conservative, .. } => {
+                        use crate::ConservativeDepth as Cd;
+                        let depth = match conservative {
+                            Cd::GreaterEqual => "greater",
+                            Cd::LessEqual => "less",
+                            Cd::Unchanged => "unchanged",
+                        };
+                        writeln!(self.out, "layout (depth_{depth}) out float gl_FragDepth;")?;
+                    }
                 }
-                writeln!(self.out)?;
             } else {
                 log::warn!(
                     "Early depth testing is not supported for this version of GLSL: {}",

naga/src/back/spv/writer.rs

@@ -1133,6 +1133,35 @@ impl Writer {
             crate::ShaderStage::Vertex => spirv::ExecutionModel::Vertex,
             crate::ShaderStage::Fragment => {
                 self.write_execution_mode(function_id, spirv::ExecutionMode::OriginUpperLeft)?;
+                match entry_point.early_depth_test {
+                    Some(crate::EarlyDepthTest::Force) => {
+                        self.write_execution_mode(
+                            function_id,
+                            spirv::ExecutionMode::EarlyFragmentTests,
+                        )?;
+                    }
+                    Some(crate::EarlyDepthTest::Allow { conservative }) => {
+                        // TODO: Consider emitting EarlyAndLateFragmentTestsAMD here, if available.
+                        // https://github.khronos.org/SPIRV-Registry/extensions/AMD/SPV_AMD_shader_early_and_late_fragment_tests.html
+                        // This permits early depth tests even if the shader writes to a storage
+                        // binding
+                        match conservative {
+                            crate::ConservativeDepth::GreaterEqual => self.write_execution_mode(
+                                function_id,
+                                spirv::ExecutionMode::DepthGreater,
+                            )?,
+                            crate::ConservativeDepth::LessEqual => self.write_execution_mode(
+                                function_id,
+                                spirv::ExecutionMode::DepthLess,
+                            )?,
+                            crate::ConservativeDepth::Unchanged => self.write_execution_mode(
+                                function_id,
+                                spirv::ExecutionMode::DepthUnchanged,
+                            )?,
+                        }
+                    }
+                    None => {}
+                }
                 if let Some(ref result) = entry_point.function.result {
                     if contains_builtin(
                         result.binding.as_ref(),

naga/src/front/glsl/functions.rs

@@ -1370,7 +1370,7 @@ impl Frontend {
         ctx.module.entry_points.push(EntryPoint {
             name: "main".to_string(),
             stage: self.meta.stage,
-            early_depth_test: Some(crate::EarlyDepthTest { conservative: None })
+            early_depth_test: Some(crate::EarlyDepthTest::Force)
                 .filter(|_| self.meta.early_fragment_tests),
             workgroup_size: self.meta.workgroup_size,
             workgroup_size_overrides: None,

naga/src/front/spv/mod.rs

@@ -4825,24 +4825,49 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
 
         match mode {
             ExecutionMode::EarlyFragmentTests => {
-                if ep.early_depth_test.is_none() {
-                    ep.early_depth_test = Some(crate::EarlyDepthTest { conservative: None });
-                }
+                ep.early_depth_test = Some(crate::EarlyDepthTest::Force);
             }
             ExecutionMode::DepthUnchanged => {
-                ep.early_depth_test = Some(crate::EarlyDepthTest {
-                    conservative: Some(crate::ConservativeDepth::Unchanged),
-                });
+                if let &mut Some(ref mut early_depth_test) = &mut ep.early_depth_test {
+                    if let &mut crate::EarlyDepthTest::Allow {
+                        ref mut conservative,
+                    } = early_depth_test
+                    {
+                        *conservative = crate::ConservativeDepth::Unchanged;
+                    }
+                } else {
+                    ep.early_depth_test = Some(crate::EarlyDepthTest::Allow {
+                        conservative: crate::ConservativeDepth::Unchanged,
+                    });
+                }
             }
             ExecutionMode::DepthGreater => {
-                ep.early_depth_test = Some(crate::EarlyDepthTest {
-                    conservative: Some(crate::ConservativeDepth::GreaterEqual),
-                });
+                if let &mut Some(ref mut early_depth_test) = &mut ep.early_depth_test {
+                    if let &mut crate::EarlyDepthTest::Allow {
+                        ref mut conservative,
+                    } = early_depth_test
+                    {
+                        *conservative = crate::ConservativeDepth::GreaterEqual;
+                    }
+                } else {
+                    ep.early_depth_test = Some(crate::EarlyDepthTest::Allow {
+                        conservative: crate::ConservativeDepth::GreaterEqual,
+                    });
+                }
             }
             ExecutionMode::DepthLess => {
-                ep.early_depth_test = Some(crate::EarlyDepthTest {
-                    conservative: Some(crate::ConservativeDepth::LessEqual),
-                });
+                if let &mut Some(ref mut early_depth_test) = &mut ep.early_depth_test {
+                    if let &mut crate::EarlyDepthTest::Allow {
+                        ref mut conservative,
+                    } = early_depth_test
+                    {
+                        *conservative = crate::ConservativeDepth::LessEqual;
+                    }
+                } else {
+                    ep.early_depth_test = Some(crate::EarlyDepthTest::Allow {
+                        conservative: crate::ConservativeDepth::LessEqual,
+                    });
+                }
             }
             ExecutionMode::DepthReplacing => {
                 // Ignored because it can be deduced from the IR.

naga/src/front/wgsl/parse/mod.rs

@@ -2838,15 +2838,17 @@ impl Parser {
                     workgroup_size.set(new_workgroup_size, name_span)?;
                 }
                 "early_depth_test" => {
-                    let conservative = if lexer.skip(Token::Paren('(')) {
-                        let (ident, ident_span) = lexer.next_ident_with_span()?;
-                        let value = conv::map_conservative_depth(ident, ident_span)?;
-                        lexer.expect(Token::Paren(')'))?;
-                        Some(value)
+                    lexer.expect(Token::Paren('('))?;
+                    let (ident, ident_span) = lexer.next_ident_with_span()?;
+                    let value = if ident == "force" {
+                        crate::EarlyDepthTest::Force
                     } else {
-                        None
+                        crate::EarlyDepthTest::Allow {
+                            conservative: conv::map_conservative_depth(ident, ident_span)?,
+                        }
                     };
-                    early_depth_test.set(crate::EarlyDepthTest { conservative }, name_span)?;
+                    lexer.expect(Token::Paren(')'))?;
+                    early_depth_test.set(value, name_span)?;
                 }
                 "must_use" => {
                     must_use.set(name_span, name_span)?;

naga/src/ir/mod.rs

@@ -237,19 +237,27 @@ use crate::{FastIndexMap, NamedExpressions};
 
 pub use block::Block;
 
-/// Early fragment tests.
+/// Explicitly allows early depth/stencil tests.
 ///
-/// In a standard situation, if a driver determines that it is possible to switch on early depth test, it will.
+/// Normally, depth/stencil tests are performed after fragment shading. However, as an optimization,
+/// most drivers will move the depth/stencil tests before fragment shading if this does not
+/// have any observable consequences. This optimization is disabled under the following
+/// circumstances:
+///   - `discard` is called in the fragment shader.
+///   - The fragment shader writes to the depth buffer.
+///   - The fragment shader writes to any storage bindings.
 ///
-/// Typical situations when early depth test is switched off:
-///   - Calling `discard` in a shader.
-///   - Writing to the depth buffer, unless ConservativeDepth is enabled.
+/// When `EarlyDepthTest` is set, it is allowed to perform an early depth/stencil test even if the
+/// above conditions are not met. When [`EarlyDepthTest::Force`] is used, depth/stencil tests
+/// **must** be performed before fragment shading.
 ///
-/// To use in a shader:
+/// To force early depth/stencil tests in a shader:
 ///   - GLSL: `layout(early_fragment_tests) in;`
 ///   - HLSL: `Attribute earlydepthstencil`
 ///   - SPIR-V: `ExecutionMode EarlyFragmentTests`
-///   - WGSL: `@early_depth_test`
+///   - WGSL: `@early_depth_test(force)`
+///
+/// This may also be enabled in a shader by specifying a [`ConservativeDepth`].
 ///
 /// For more, see:
 ///   - <https://www.khronos.org/opengl/wiki/Early_Fragment_Test#Explicit_specification>
@@ -259,8 +267,24 @@ pub use block::Block;
 #[cfg_attr(feature = "serialize", derive(Serialize))]
 #[cfg_attr(feature = "deserialize", derive(Deserialize))]
 #[cfg_attr(feature = "arbitrary", derive(Arbitrary))]
-pub struct EarlyDepthTest {
-    pub conservative: Option<ConservativeDepth>,
+pub enum EarlyDepthTest {
+    /// Requires depth/stencil tests to be performed before fragment shading.
+    ///
+    /// This will disable depth/stencil tests after fragment shading, so discarding the fragment
+    /// or overwriting the fragment depth will have no effect.
+    Force,
+
+    /// Allows an additional depth/stencil test to be performed before fragment shading.
+    ///
+    /// It is up to the driver to decide whether early tests are performed. Unlike `Force`, this
+    /// does not disable depth/stencil tests after fragment shading.
+    Allow {
+        /// Specifies restrictions on how the depth value can be modified within the fragment
+        /// shader.
+        ///
+        /// This may be taken into account when deciding whether to perform early tests.
+        conservative: ConservativeDepth,
+    },
 }
 
 /// Enables adjusting depth without disabling early Z.

naga/tests/in/wgsl/early-depth-test-conservative.toml

@@ -0,0 +1,5 @@
+god_mode = true
+targets = "SPIRV | GLSL"
+
+[glsl]
+version.Desktop = 420

naga/tests/in/wgsl/early-depth-test-conservative.wgsl

@@ -0,0 +1,5 @@
+@fragment
+@early_depth_test(less_equal)
+fn main(@builtin(position) pos: vec4<f32>) -> @builtin(frag_depth) f32 {
+    return pos.z - 0.1;
+}

naga/tests/in/wgsl/early-depth-test-force.toml

@@ -0,0 +1,2 @@
+god_mode = true
+targets = "SPIRV | GLSL"

naga/tests/in/wgsl/early-depth-test-force.wgsl

@@ -0,0 +1,5 @@
+@fragment
+@early_depth_test(force)
+fn main() -> @location(0) vec4<f32> {
+    return vec4<f32>(0.4, 0.3, 0.2, 0.1);
+}

naga/tests/out/glsl/wgsl-early-depth-test-conservative.main.Fragment.glsl

@@ -0,0 +1,9 @@
+#version 420 core
+layout (depth_less) out float gl_FragDepth;
+
+void main() {
+    vec4 pos = gl_FragCoord;
+    gl_FragDepth = (pos.z - 0.1);
+    return;
+}
+

naga/tests/out/glsl/wgsl-early-depth-test-force.main.Fragment.glsl

@@ -0,0 +1,13 @@
+#version 310 es
+
+precision highp float;
+precision highp int;
+
+layout(early_fragment_tests) in;
+layout(location = 0) out vec4 _fs2p_location0;
+
+void main() {
+    _fs2p_location0 = vec4(0.4, 0.3, 0.2, 0.1);
+    return;
+}
+

naga/tests/out/spv/wgsl-early-depth-test-conservative.spvasm

@@ -0,0 +1,32 @@
+; SPIR-V
+; Version: 1.1
+; Generator: rspirv
+; Bound: 17
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %11 "main" %6 %9
+OpExecutionMode %11 OriginUpperLeft
+OpExecutionMode %11 DepthLess
+OpExecutionMode %11 DepthReplacing
+OpDecorate %6 BuiltIn FragCoord
+OpDecorate %9 BuiltIn FragDepth
+%2 = OpTypeVoid
+%3 = OpTypeFloat 32
+%4 = OpTypeVector %3 4
+%7 = OpTypePointer Input %4
+%6 = OpVariable  %7  Input
+%10 = OpTypePointer Output %3
+%9 = OpVariable  %10  Output
+%12 = OpTypeFunction %2
+%13 = OpConstant  %3  0.1
+%11 = OpFunction  %2  None %12
+%5 = OpLabel
+%8 = OpLoad  %4  %6
+OpBranch %14
+%14 = OpLabel
+%15 = OpCompositeExtract  %3  %8 2
+%16 = OpFSub  %3  %15 %13
+OpStore %9 %16
+OpReturn
+OpFunctionEnd

naga/tests/out/spv/wgsl-early-depth-test-force.spvasm

@@ -0,0 +1,29 @@
+; SPIR-V
+; Version: 1.1
+; Generator: rspirv
+; Bound: 16
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %8 "main" %6
+OpExecutionMode %8 OriginUpperLeft
+OpExecutionMode %8 EarlyFragmentTests
+OpDecorate %6 Location 0
+%2 = OpTypeVoid
+%4 = OpTypeFloat 32
+%3 = OpTypeVector %4 4
+%7 = OpTypePointer Output %3
+%6 = OpVariable  %7  Output
+%9 = OpTypeFunction %2
+%10 = OpConstant  %4  0.4
+%11 = OpConstant  %4  0.3
+%12 = OpConstant  %4  0.2
+%13 = OpConstant  %4  0.1
+%14 = OpConstantComposite  %3  %10 %11 %12 %13
+%8 = OpFunction  %2  None %9
+%5 = OpLabel
+OpBranch %15
+%15 = OpLabel
+OpStore %6 %14
+OpReturn
+OpFunctionEnd

wgpu-hal/src/vulkan/adapter.rs

@@ -546,6 +546,7 @@ impl PhysicalDeviceFeatures {
             | F::TEXTURE_ADAPTER_SPECIFIC_FORMAT_FEATURES
             | F::CLEAR_TEXTURE
             | F::PIPELINE_CACHE
+            | F::SHADER_EARLY_DEPTH_TEST
             | F::TEXTURE_ATOMIC;
 
         let mut dl_flags = Df::COMPUTE_SHADERS