fix(naga): Reject unsupported enable extensions (#8913)

Previously we required that the enable extension be declared and that
the capability be present if the functionality was used, but allowed the
extension to be declared if the functionality was not used.

Author: andyleiserson

CHANGELOG.md

@@ -115,6 +115,7 @@ Bottom level categories:
 
 #### naga
 
+- Naga and `wgpu` now reject shaders with an `enable` directive for functionality that is not available, even if that functionality is not used by the shader. By @andyleiserson in [#8913](https://github.com/gfx-rs/wgpu/pull/8913).
 - Prevent UB from incorrectly using ray queries on HLSL. By @Vecvec in [#8763](https://github.com/gfx-rs/wgpu/pull/8763).
 
 ### deno\_webgpu

naga-cli/src/bin/naga.rs

@@ -144,6 +144,14 @@ struct Args {
     /// defines to be passed to the parser (only glsl is supported)
     #[argh(option, short = 'D')]
     defines: Vec<Defines>,
+
+    /// capabilities for parsing and validation.
+    ///
+    /// Can be a comma-separated list of capability names (e.g.,
+    /// "shader_float16,dual_source_blending"), a numeric bitflags value (e.g.,
+    /// "67108864"), the string "none", or the string "all".
+    #[argh(option, default = "CapabilitiesArg(naga::valid::Capabilities::all())")]
+    capabilities: CapabilitiesArg,
 }
 
 /// Newtype so we can implement [`FromStr`] for `BoundsCheckPolicy`.
@@ -334,6 +342,37 @@ impl FromStr for Defines {
     }
 }
 
+#[derive(Debug, Clone, Copy)]
+struct CapabilitiesArg(naga::valid::Capabilities);
+
+impl FromStr for CapabilitiesArg {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        use naga::valid::Capabilities;
+
+        let s = s.to_uppercase();
+
+        if s == "NONE" {
+            Ok(Self(Capabilities::empty()))
+        } else if s == "ALL" {
+            Ok(Self(Capabilities::all()))
+        } else if let Ok(bits) = s.parse::<u64>() {
+            Capabilities::from_bits(bits)
+                .map(Self)
+                .ok_or_else(|| format!("Invalid capabilities bitflags value: {bits}"))
+        } else {
+            s.split(',')
+                .try_fold(Capabilities::empty(), |acc, s| {
+                    Capabilities::from_name(s.trim())
+                        .map(|cap| acc | cap)
+                        .ok_or(format!("Unknown capability {}", s.trim()))
+                })
+                .map(Self)
+        }
+    }
+}
+
 #[derive(Default)]
 struct Parameters<'a> {
     validation_flags: naga::valid::ValidationFlags,
@@ -350,6 +389,7 @@ struct Parameters<'a> {
     input_kind: Option<InputKind>,
     shader_stage: Option<ShaderStage>,
     defines: FastHashMap<String, String>,
+    capabilities: naga::valid::Capabilities,
 
     /// We use this copy of `args.compact` to know whether we should pass the
     /// entrypoint to `process_overrides`, which will result in removal from
@@ -503,6 +543,7 @@ fn run() -> anyhow::Result<()> {
     );
 
     params.compact = args.compact;
+    params.capabilities = args.capabilities.0;
 
     if args.bulk_validate {
         return bulk_validate(&args, &params);
@@ -680,7 +721,12 @@ fn parse_input(input_path: &Path, input: Vec<u8>, params: &Parameters) -> anyhow
         },
         InputKind::Wgsl => {
             let input = String::from_utf8(input)?;
-            let result = naga::front::wgsl::parse_str(&input);
+            let options = naga::front::wgsl::Options {
+                parse_doc_comments: false,
+                capabilities: params.capabilities,
+            };
+            let mut frontend = naga::front::wgsl::Frontend::new_with_options(options);
+            let result = frontend.parse(&input);
             match result {
                 Ok(v) => Parsed {
                     module: v,
@@ -959,7 +1005,7 @@ fn bulk_validate(args: &Args, params: &Parameters) -> anyhow::Result<()> {
         };
 
         let mut validator =
-            naga::valid::Validator::new(params.validation_flags, naga::valid::Capabilities::all());
+            naga::valid::Validator::new(params.validation_flags, params.capabilities);
         validator.subgroup_stages(naga::valid::ShaderStages::all());
         validator.subgroup_operations(naga::valid::SubgroupOperationSet::all());
 

naga-test/src/lib.rs

@@ -84,6 +84,7 @@ impl From<&WgslInParameters> for naga::front::wgsl::Options {
     fn from(value: &WgslInParameters) -> Self {
         Self {
             parse_doc_comments: value.parse_doc_comments,
+            capabilities: naga::valid::Capabilities::all(),
         }
     }
 }

naga/src/front/wgsl/error.rs

@@ -387,6 +387,10 @@ pub(crate) enum Error<'a> {
         kind: EnableExtension,
         span: Span,
     },
+    EnableExtensionNotSupported {
+        kind: EnableExtension,
+        span: Span,
+    },
     LanguageExtensionNotYetImplemented {
         kind: UnimplementedLanguageExtension,
         span: Span,
@@ -1258,6 +1262,17 @@ impl<'a> Error<'a> {
                     ]
                 },
             },
+            Error::EnableExtensionNotSupported { kind, span } => ParseError {
+                message: format!(
+                    "the `{}` extension is not supported in the current environment",
+                    kind.to_ident()
+                ),
+                labels: vec![(
+                    span,
+                    "unsupported enable-extension".into(),
+                )],
+                notes: vec![],
+            },
             Error::LanguageExtensionNotYetImplemented { kind, span } => ParseError {
                 message: format!(
                     "the `{}` language extension is not yet supported",

naga/src/front/wgsl/parse/directive/enable_extension.rs

@@ -156,6 +156,7 @@ impl EnableExtension {
 
 /// A variant of [`EnableExtension::Implemented`].
 #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
+#[cfg_attr(test, derive(strum::VariantArray))]
 pub enum ImplementedEnableExtension {
     /// Enables `f16`/`half` primitive support in all shader languages.
     ///
@@ -187,6 +188,50 @@ pub enum ImplementedEnableExtension {
     WgpuCooperativeMatrix,
 }
 
+impl ImplementedEnableExtension {
+    /// A slice of all variants of [`ImplementedEnableExtension`].
+    pub const VARIANTS: &'static [Self] = &[
+        Self::F16,
+        Self::DualSourceBlending,
+        Self::ClipDistances,
+        Self::WgpuMeshShader,
+        Self::WgpuRayQuery,
+        Self::WgpuRayQueryVertexReturn,
+        Self::WgpuRayTracingPipeline,
+        Self::WgpuCooperativeMatrix,
+    ];
+
+    /// Returns slice of all variants of [`ImplementedEnableExtension`].
+    pub const fn all() -> &'static [Self] {
+        Self::VARIANTS
+    }
+
+    /// Returns the capability required for this enable extension.
+    pub const fn capability(self) -> crate::valid::Capabilities {
+        use crate::valid::Capabilities as C;
+        match self {
+            Self::F16 => C::SHADER_FLOAT16,
+            Self::DualSourceBlending => C::DUAL_SOURCE_BLENDING,
+            Self::ClipDistances => C::CLIP_DISTANCE,
+            Self::WgpuMeshShader => C::MESH_SHADER,
+            Self::WgpuRayQuery => C::RAY_QUERY,
+            Self::WgpuRayQueryVertexReturn => C::RAY_HIT_VERTEX_POSITION,
+            Self::WgpuCooperativeMatrix => C::COOPERATIVE_MATRIX,
+            Self::WgpuRayTracingPipeline => C::RAY_TRACING_PIPELINE,
+        }
+    }
+}
+
+#[test]
+/// Asserts that the manual implementation of VARIANTS is the same as the derived strum version would be
+/// while still allowing strum to be a dev-only dependency
+fn test_manual_variants_array_is_correct() {
+    assert_eq!(
+        <ImplementedEnableExtension as strum::VariantArray>::VARIANTS,
+        ImplementedEnableExtension::VARIANTS
+    );
+}
+
 /// A variant of [`EnableExtension::Unimplemented`].
 #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
 pub enum UnimplementedEnableExtension {

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

@@ -275,13 +275,16 @@ impl<'a> BindingParser<'a> {
 pub struct Options {
     /// Controls whether the parser should parse doc comments.
     pub parse_doc_comments: bool,
+    /// Capabilities to enable during parsing.
+    pub capabilities: crate::valid::Capabilities,
 }
 
 impl Options {
-    /// Creates a new [`Options`] without doc comments parsing.
+    /// Creates a new default [`Options`].
     pub const fn new() -> Self {
         Options {
             parse_doc_comments: false,
+            capabilities: crate::valid::Capabilities::all(),
         }
     }
 }
@@ -2231,6 +2234,14 @@ impl Parser {
                                     }))
                                 }
                             };
+                            // Check if the required capability is supported
+                            let required_capability = extension.capability();
+                            if !options.capabilities.contains(required_capability) {
+                                return Err(Box::new(Error::EnableExtensionNotSupported {
+                                    kind,
+                                    span,
+                                }));
+                            }
                             enable_extensions.add(extension);
                             Ok(())
                         })?;

naga/tests/naga/wgsl_errors.rs

@@ -13,6 +13,7 @@
 
 use naga::{
     compact::KeepUnused,
+    front::wgsl::{EnableExtension, ImplementedEnableExtension},
     valid::{self, Capabilities},
 };
 
@@ -4897,7 +4898,7 @@ fn check_ray_tracing_pipeline_payload_disallowed() {
                 "enable wgpu_ray_tracing_pipeline;
             @group(0) @binding(0) var acc_struct: acceleration_structure;
             var<ray_payload> payload: u32;
-            
+
             {stage} fn main() {output} {{_ = payload; {stmt}}}"
             ),
             Err(naga::valid::ValidationError::EntryPoint {
@@ -4908,3 +4909,45 @@ fn check_ray_tracing_pipeline_payload_disallowed() {
         );
     }
 }
+
+#[track_caller]
+fn check_with_capabilities(input: &str, snapshot: &str, capabilities: Capabilities) {
+    let mut options = naga::front::wgsl::Options::new();
+    options.capabilities = capabilities;
+    let mut frontend = naga::front::wgsl::Frontend::new_with_options(options);
+    let output = match frontend.parse(input) {
+        Ok(_) => panic!("expected parser error, but parsing succeeded!"),
+        Err(err) => err.emit_to_string(input),
+    };
+    if output != snapshot {
+        for diff in diff::lines(snapshot, &output) {
+            match diff {
+                diff::Result::Left(l) => println!("-{l}"),
+                diff::Result::Both(l, _) => println!(" {l}"),
+                diff::Result::Right(r) => println!("+{r}"),
+            }
+        }
+        panic!("Error snapshot failed");
+    }
+}
+
+#[test]
+fn enable_without_capability() {
+    for extension in ImplementedEnableExtension::all() {
+        let ident = EnableExtension::from(*extension).to_ident();
+        let carets = "^".repeat(ident.len());
+        check_with_capabilities(
+            &format!("enable {ident};"),
+            &format!(
+                "error: the `{ident}` extension is not supported in the current environment
+  ┌─ wgsl:1:8
+  │
+1 │ enable {ident};
+  │        {carets} unsupported enable-extension
+
+"
+            ),
+            !extension.capability(),
+        );
+    }
+}

tests/tests/wgpu-gpu/dual_source_blending.rs

@@ -18,7 +18,7 @@ fn vs_main() -> @builtin(position) vec4f {
 "#;
 
 const FRAGMENT_SHADER_WITHOUT_DUAL_SOURCE_BLENDING: &str = r#"
-@fragment 
+@fragment
 fn fs_main() -> @location(0) vec4f {
     return vec4f(1.0);
 }
@@ -31,7 +31,7 @@ struct FragmentOutput {
     @location(0) @blend_src(1) output1_: vec4<f32>,
 }
 
-@fragment 
+@fragment
 fn fs_main() -> FragmentOutput {
     return FragmentOutput(vec4<f32>(0.4f, 0.3f, 0.2f, 0.1f), vec4<f32>(0.9f, 0.8f, 0.7f, 0.6f));
 }
@@ -112,7 +112,7 @@ async fn dual_source_blending_disabled(ctx: TestingContext) {
                 source: ShaderSource::Wgsl(FRAGMENT_SHADER_WITH_DUAL_SOURCE_BLENDING.into()),
             });
         },
-        Some("Capability Capabilities(DUAL_SOURCE_BLENDING) is not supported"),
+        Some("the `dual_source_blending` extension is not supported in the current environment"),
     );
 }
 

tests/tests/wgpu-gpu/shader/compilation_messages/mod.rs

@@ -1,9 +1,14 @@
 use wgpu::include_wgsl;
 
-use wgpu_test::{gpu_test, GpuTestConfiguration, GpuTestInitializer, TestParameters};
+use wgpu_test::{fail, gpu_test, valid, GpuTestConfiguration, GpuTestInitializer, TestParameters};
 
 pub fn all_tests(vec: &mut Vec<GpuTestInitializer>) {
-    vec.extend([SHADER_COMPILE_SUCCESS, SHADER_COMPILE_ERROR]);
+    vec.extend([
+        SHADER_COMPILE_SUCCESS,
+        SHADER_COMPILE_ERROR,
+        ENABLE_EXTENSION_AVAILABLE,
+        ENABLE_EXTENSION_UNAVAILABLE,
+    ]);
 }
 
 #[gpu_test]
@@ -51,3 +56,55 @@ static SHADER_COMPILE_ERROR: GpuTestConfiguration = GpuTestConfiguration::new()
             "Expected the column number to be 33, because we're counting lines from 1"
         );
     });
+
+const ENABLE_EXTENSION_SHADER_SOURCE: &str = r#"
+    enable f16;
+
+    @compute @workgroup_size(1)
+    fn main() {}
+"#;
+
+#[gpu_test]
+static ENABLE_EXTENSION_AVAILABLE: GpuTestConfiguration = GpuTestConfiguration::new()
+    .parameters(
+        TestParameters::default()
+            .features(wgpu::Features::SHADER_F16)
+            .downlevel_flags(wgpu::DownlevelFlags::COMPUTE_SHADERS)
+            .limits(wgpu::Limits::downlevel_defaults()),
+    )
+    .run_async(|ctx| async move {
+        valid(&ctx.device, || {
+            let _ = ctx
+                .device
+                .create_shader_module(wgpu::ShaderModuleDescriptor {
+                    label: Some("shader declaring enable extension"),
+                    source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(
+                        ENABLE_EXTENSION_SHADER_SOURCE,
+                    )),
+                });
+        });
+    });
+
+#[gpu_test]
+static ENABLE_EXTENSION_UNAVAILABLE: GpuTestConfiguration = GpuTestConfiguration::new()
+    .parameters(
+        TestParameters::default()
+            // SHADER_F16 feature not requested
+            .downlevel_flags(wgpu::DownlevelFlags::COMPUTE_SHADERS)
+            .limits(wgpu::Limits::downlevel_defaults()),
+    )
+    .run_async(|ctx| async move {
+        fail(
+            &ctx.device,
+            || {
+                ctx.device
+                    .create_shader_module(wgpu::ShaderModuleDescriptor {
+                        label: Some("shader declaring enable extension"),
+                        source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(
+                            ENABLE_EXTENSION_SHADER_SOURCE,
+                        )),
+                    })
+            },
+            Some("the `f16` extension is not supported in the current environment"),
+        );
+    });