[wgsl] gate primitive_index built-in behind its extension & [glsl-out] require extensions for gl_PrimitiveID if needed (gfx-rs#8879)

Author: inner-daemons

CHANGELOG.md

@@ -68,6 +68,7 @@ By @cwfitzgerald in [#8999](https://github.com/gfx-rs/wgpu/pull/8999).
         - `ir::Block::len`
 - Changed passthrough shaders to not require an entry point parameter, so that the same shader module may be used in multiple entry points. Also added support for metallib passthrough. By @inner-daemons in #8886.
 - Added support for `insert_debug_marker`, `push_debug_group` and `pop_debug_group` on WebGPU. By @evilpie in [#9017](https://github.com/gfx-rs/wgpu/pull/9017).
+- Added support for `enable primitive_index` and `@builtin(primitive_index)` with support on all platforms. By @inner-daemons in #8879.
 
 #### naga
 

naga/src/back/glsl/features.rs

@@ -57,6 +57,8 @@ bitflags::bitflags! {
         const TEXTURE_ATOMICS = 1 << 25;
         /// Image atomics
         const SHADER_BARYCENTRICS = 1 << 26;
+        /// Primitive index builtin
+        const PRIMITIVE_INDEX = 1 << 27;
     }
 }
 
@@ -298,6 +300,18 @@ impl FeaturesManager {
             )?;
         }
 
+        if self.0.contains(Features::PRIMITIVE_INDEX) {
+            match options.version {
+                Version::Embedded { version, .. } if version < 320 => {
+                    writeln!(out, "#extension GL_OES_geometry_shader : require")?;
+                }
+                Version::Desktop(version) if version < 150 => {
+                    writeln!(out, "#extension GL_ARB_geometry_shader4 : require")?;
+                }
+                _ => (),
+            }
+        }
+
         Ok(())
     }
 }
@@ -616,6 +630,9 @@ impl<W> Writer<'_, W> {
                     crate::BuiltIn::Barycentric { .. } => {
                         self.features.request(Features::SHADER_BARYCENTRICS)
                     }
+                    crate::BuiltIn::PrimitiveIndex => {
+                        self.features.request(Features::PRIMITIVE_INDEX)
+                    }
                     _ => {}
                 },
                 Binding::Location {

naga/src/back/spv/writer.rs

@@ -3083,6 +3083,7 @@ impl Writer {
                     Bi::PointCoord => BuiltIn::PointCoord,
                     Bi::FrontFacing => BuiltIn::FrontFacing,
                     Bi::PrimitiveIndex => {
+                        // Geometry shader capability is required for primitive index
                         self.require_any(
                             "`primitive_index` built-in",
                             &[spirv::Capability::Geometry],

naga/src/back/wgsl/writer.rs

@@ -278,96 +278,120 @@ impl<W: Write> Writer<W> {
     /// Helper method which writes all the `enable` declarations
     /// needed for a module.
     fn write_enable_declarations(&mut self, module: &Module) -> BackendResult {
-        let mut needs_f16 = false;
-        let mut needs_dual_source_blending = false;
-        let mut needs_clip_distances = false;
-        let mut needs_mesh_shaders = false;
-        let mut needs_cooperative_matrix = false;
+        #[derive(Default)]
+        struct RequiredEnabled {
+            f16: bool,
+            dual_source_blending: bool,
+            clip_distances: bool,
+            mesh_shaders: bool,
+            primitive_index: bool,
+            cooperative_matrix: bool,
+        }
+        let mut needed = RequiredEnabled::default();
+
+        let check_binding = |binding: &crate::Binding, needed: &mut RequiredEnabled| match *binding
+        {
+            crate::Binding::Location {
+                blend_src: Some(_), ..
+            } => {
+                needed.dual_source_blending = true;
+            }
+            crate::Binding::BuiltIn(crate::BuiltIn::ClipDistance) => {
+                needed.clip_distances = true;
+            }
+            crate::Binding::BuiltIn(crate::BuiltIn::PrimitiveIndex) => {
+                needed.primitive_index = true;
+            }
+            crate::Binding::Location {
+                per_primitive: true,
+                ..
+            } => {
+                needed.mesh_shaders = true;
+            }
+            crate::Binding::BuiltIn(
+                crate::BuiltIn::MeshTaskSize
+                | crate::BuiltIn::CullPrimitive
+                | crate::BuiltIn::PointIndex
+                | crate::BuiltIn::LineIndices
+                | crate::BuiltIn::TriangleIndices
+                | crate::BuiltIn::VertexCount
+                | crate::BuiltIn::Vertices
+                | crate::BuiltIn::PrimitiveCount
+                | crate::BuiltIn::Primitives,
+            ) => {
+                needed.mesh_shaders = true;
+            }
+            _ => {}
+        };
 
         // Determine which `enable` declarations are needed
         for (_, ty) in module.types.iter() {
             match ty.inner {
                 TypeInner::Scalar(scalar)
                 | TypeInner::Vector { scalar, .. }
                 | TypeInner::Matrix { scalar, .. } => {
-                    needs_f16 |= scalar == crate::Scalar::F16;
+                    needed.f16 |= scalar == crate::Scalar::F16;
                 }
                 TypeInner::Struct { ref members, .. } => {
                     for binding in members.iter().filter_map(|m| m.binding.as_ref()) {
-                        match *binding {
-                            crate::Binding::Location {
-                                blend_src: Some(_), ..
-                            } => {
-                                needs_dual_source_blending = true;
-                            }
-                            crate::Binding::BuiltIn(crate::BuiltIn::ClipDistance) => {
-                                needs_clip_distances = true;
-                            }
-                            crate::Binding::Location {
-                                per_primitive: true,
-                                ..
-                            } => {
-                                needs_mesh_shaders = true;
-                            }
-                            crate::Binding::BuiltIn(
-                                crate::BuiltIn::MeshTaskSize
-                                | crate::BuiltIn::CullPrimitive
-                                | crate::BuiltIn::PointIndex
-                                | crate::BuiltIn::LineIndices
-                                | crate::BuiltIn::TriangleIndices
-                                | crate::BuiltIn::VertexCount
-                                | crate::BuiltIn::Vertices
-                                | crate::BuiltIn::PrimitiveCount
-                                | crate::BuiltIn::Primitives,
-                            ) => {
-                                needs_mesh_shaders = true;
-                            }
-                            _ => {}
-                        }
+                        check_binding(binding, &mut needed);
                     }
                 }
                 TypeInner::CooperativeMatrix { .. } => {
-                    needs_cooperative_matrix = true;
+                    needed.cooperative_matrix = true;
                 }
                 _ => {}
             }
         }
 
-        if module
-            .entry_points
-            .iter()
-            .any(|ep| matches!(ep.stage, ShaderStage::Mesh | ShaderStage::Task))
-        {
-            needs_mesh_shaders = true;
+        for ep in &module.entry_points {
+            if matches!(ep.stage, ShaderStage::Mesh | ShaderStage::Task) {
+                needed.mesh_shaders = true;
+            }
+            if let Some(res) = ep.function.result.as_ref().and_then(|a| a.binding.as_ref()) {
+                check_binding(res, &mut needed);
+            }
+            for arg in ep
+                .function
+                .arguments
+                .iter()
+                .filter_map(|a| a.binding.as_ref())
+            {
+                check_binding(arg, &mut needed);
+            }
         }
 
         if module
             .global_variables
             .iter()
             .any(|gv| gv.1.space == crate::AddressSpace::TaskPayload)
         {
-            needs_mesh_shaders = true;
+            needed.mesh_shaders = true;
         }
 
         // Write required declarations
         let mut any_written = false;
-        if needs_f16 {
+        if needed.f16 {
             writeln!(self.out, "enable f16;")?;
             any_written = true;
         }
-        if needs_dual_source_blending {
+        if needed.dual_source_blending {
             writeln!(self.out, "enable dual_source_blending;")?;
             any_written = true;
         }
-        if needs_clip_distances {
+        if needed.clip_distances {
             writeln!(self.out, "enable clip_distances;")?;
             any_written = true;
         }
-        if needs_mesh_shaders {
+        if needed.mesh_shaders {
             writeln!(self.out, "enable wgpu_mesh_shader;")?;
             any_written = true;
         }
-        if needs_cooperative_matrix {
+        if needed.primitive_index {
+            writeln!(self.out, "enable primitive_index;")?;
+            any_written = true;
+        }
+        if needed.cooperative_matrix {
             writeln!(self.out, "enable wgpu_cooperative_matrix;")?;
             any_written = true;
         }

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

@@ -21,14 +21,8 @@ pub fn map_address_space<'a>(
         "immediate" => Ok(crate::AddressSpace::Immediate),
         "function" => Ok(crate::AddressSpace::Function),
         "task_payload" => {
-            if enable_extensions.contains(ImplementedEnableExtension::WgpuMeshShader) {
-                Ok(crate::AddressSpace::TaskPayload)
-            } else {
-                Err(Box::new(Error::EnableExtensionNotEnabled {
-                    span,
-                    kind: ImplementedEnableExtension::WgpuMeshShader.into(),
-                }))
-            }
+            enable_extensions.require(ImplementedEnableExtension::WgpuMeshShader, span)?;
+            Ok(crate::AddressSpace::TaskPayload)
         }
         "ray_payload" => {
             if enable_extensions.contains(ImplementedEnableExtension::WgpuRayTracingPipeline) {
@@ -154,12 +148,11 @@ pub fn map_built_in(
     };
     match built_in {
         crate::BuiltIn::ClipDistance => {
-            if !enable_extensions.contains(ImplementedEnableExtension::ClipDistances) {
-                return Err(Box::new(Error::EnableExtensionNotEnabled {
-                    span,
-                    kind: ImplementedEnableExtension::ClipDistances.into(),
-                }));
-            }
+            enable_extensions.require(ImplementedEnableExtension::ClipDistances, span)?
+        }
+
+        crate::BuiltIn::PrimitiveIndex => {
+            enable_extensions.require(ImplementedEnableExtension::PrimitiveIndex, span)?
         }
         crate::BuiltIn::CullPrimitive
         | crate::BuiltIn::PointIndex
@@ -169,12 +162,7 @@ pub fn map_built_in(
         | crate::BuiltIn::Vertices
         | crate::BuiltIn::PrimitiveCount
         | crate::BuiltIn::Primitives => {
-            if !enable_extensions.contains(ImplementedEnableExtension::WgpuMeshShader) {
-                return Err(Box::new(Error::EnableExtensionNotEnabled {
-                    span,
-                    kind: ImplementedEnableExtension::WgpuMeshShader.into(),
-                }));
-            }
+            enable_extensions.require(ImplementedEnableExtension::WgpuMeshShader, span)?
         }
         _ => {}
     }

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

@@ -19,6 +19,7 @@ pub(crate) struct EnableExtensions {
     f16: bool,
     clip_distances: bool,
     wgpu_cooperative_matrix: bool,
+    primitive_index: bool,
 }
 
 impl EnableExtensions {
@@ -32,6 +33,7 @@ impl EnableExtensions {
             dual_source_blending: false,
             clip_distances: false,
             wgpu_cooperative_matrix: false,
+            primitive_index: false,
         }
     }
 
@@ -50,6 +52,7 @@ impl EnableExtensions {
             ImplementedEnableExtension::F16 => &mut self.f16,
             ImplementedEnableExtension::ClipDistances => &mut self.clip_distances,
             ImplementedEnableExtension::WgpuCooperativeMatrix => &mut self.wgpu_cooperative_matrix,
+            ImplementedEnableExtension::PrimitiveIndex => &mut self.primitive_index,
         };
         *field = true;
     }
@@ -67,6 +70,22 @@ impl EnableExtensions {
             ImplementedEnableExtension::F16 => self.f16,
             ImplementedEnableExtension::ClipDistances => self.clip_distances,
             ImplementedEnableExtension::WgpuCooperativeMatrix => self.wgpu_cooperative_matrix,
+            ImplementedEnableExtension::PrimitiveIndex => self.primitive_index,
+        }
+    }
+
+    pub(crate) fn require(
+        &self,
+        ext: ImplementedEnableExtension,
+        span: Span,
+    ) -> Result<'static, ()> {
+        if !self.contains(ext) {
+            Err(Box::new(Error::EnableExtensionNotEnabled {
+                span,
+                kind: ext.into(),
+            }))
+        } else {
+            Ok(())
         }
     }
 }
@@ -124,9 +143,7 @@ impl EnableExtension {
                 Self::Implemented(ImplementedEnableExtension::WgpuCooperativeMatrix)
             }
             Self::SUBGROUPS => Self::Unimplemented(UnimplementedEnableExtension::Subgroups),
-            Self::PRIMITIVE_INDEX => {
-                Self::Unimplemented(UnimplementedEnableExtension::PrimitiveIndex)
-            }
+            Self::PRIMITIVE_INDEX => Self::Implemented(ImplementedEnableExtension::PrimitiveIndex),
             _ => return Err(Box::new(Error::UnknownEnableExtension(span, word))),
         })
     }
@@ -144,11 +161,11 @@ impl EnableExtension {
                 ImplementedEnableExtension::DualSourceBlending => Self::DUAL_SOURCE_BLENDING,
                 ImplementedEnableExtension::F16 => Self::F16,
                 ImplementedEnableExtension::ClipDistances => Self::CLIP_DISTANCES,
+                ImplementedEnableExtension::PrimitiveIndex => Self::PRIMITIVE_INDEX,
                 ImplementedEnableExtension::WgpuRayTracingPipeline => Self::RAY_TRACING_PIPELINE,
             },
             Self::Unimplemented(kind) => match kind {
                 UnimplementedEnableExtension::Subgroups => Self::SUBGROUPS,
-                UnimplementedEnableExtension::PrimitiveIndex => Self::PRIMITIVE_INDEX,
             },
         }
     }
@@ -186,6 +203,12 @@ pub enum ImplementedEnableExtension {
     WgpuRayTracingPipeline,
     /// Enables the `wgpu_cooperative_matrix` extension, native only.
     WgpuCooperativeMatrix,
+    /// Enables the `@builtin(primitive_index)` attribute in WGSL.
+    ///
+    /// In the WGSL standard, this corresponds to [`enable primitive-index;`].
+    ///
+    /// [`enable primitive-index;`]: https://www.w3.org/TR/WGSL/#extension-primitive_index
+    PrimitiveIndex,
 }
 
 impl ImplementedEnableExtension {
@@ -199,6 +222,7 @@ impl ImplementedEnableExtension {
         Self::WgpuRayQueryVertexReturn,
         Self::WgpuRayTracingPipeline,
         Self::WgpuCooperativeMatrix,
+        Self::PrimitiveIndex,
     ];
 
     /// Returns slice of all variants of [`ImplementedEnableExtension`].
@@ -218,6 +242,7 @@ impl ImplementedEnableExtension {
             Self::WgpuRayQueryVertexReturn => C::RAY_HIT_VERTEX_POSITION,
             Self::WgpuCooperativeMatrix => C::COOPERATIVE_MATRIX,
             Self::WgpuRayTracingPipeline => C::RAY_TRACING_PIPELINE,
+            Self::PrimitiveIndex => C::PRIMITIVE_INDEX,
         }
     }
 }
@@ -241,19 +266,12 @@ pub enum UnimplementedEnableExtension {
     ///
     /// [`enable subgroups;`]: https://www.w3.org/TR/WGSL/#extension-subgroups
     Subgroups,
-    /// Enables the `@builtin(primitive_index)` attribute in WGSL.
-    ///
-    /// In the WGSL standard, this corresponds to [`enable primitive-index;`].
-    ///
-    /// [`enable primitive-index;`]: https://www.w3.org/TR/WGSL/#extension-primitive_index
-    PrimitiveIndex,
 }
 
 impl UnimplementedEnableExtension {
     pub(crate) const fn tracking_issue_num(self) -> u16 {
         match self {
             Self::Subgroups => 5555,
-            Self::PrimitiveIndex => 8236,
         }
     }
 }

naga/tests/in/wgsl/extra.wgsl

@@ -1,3 +1,5 @@
+enable primitive_index;
+
 struct ImmediateData {
     index: u32,
     double: vec2<f32>,