Add mesh shading info to naga IR

naga-cli/src/bin/naga.rs

@@ -64,6 +64,12 @@ struct Args {
     #[argh(option)]
     shader_model: Option<ShaderModelArg>,
 
+    /// the SPIR-V version to use if targeting SPIR-V
+    ///
+    /// For example, 1.0, 1.4, etc
+    #[argh(option)]
+    spirv_version: Option<SpirvVersionArg>,
+
     /// the shader stage, for example 'frag', 'vert', or 'compute'.
     /// if the shader stage is unspecified it will be derived from
     /// the file extension.
@@ -189,6 +195,22 @@ impl FromStr for ShaderModelArg {
     }
 }
 
+#[derive(Debug, Clone)]
+struct SpirvVersionArg(u8, u8);
+
+impl FromStr for SpirvVersionArg {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let dot = s
+            .find(".")
+            .ok_or_else(|| "Missing dot separator".to_owned())?;
+        let major = s[..dot].parse::<u8>().map_err(|e| e.to_string())?;
+        let minor = s[dot + 1..].parse::<u8>().map_err(|e| e.to_string())?;
+        Ok(Self(major, minor))
+    }
+}
+
 /// Newtype so we can implement [`FromStr`] for `ShaderSource`.
 #[derive(Debug, Clone, Copy)]
 struct ShaderStage(naga::ShaderStage);
@@ -465,6 +487,9 @@ fn run() -> anyhow::Result<()> {
     if let Some(ref version) = args.metal_version {
         params.msl.lang_version = version.0;
     }
+    if let Some(ref version) = args.spirv_version {
+        params.spv_out.lang_version = (version.0, version.1);
+    }
     params.keep_coordinate_space = args.keep_coordinate_space;
 
     params.dot.cfg_only = args.dot_cfg_only;

naga/src/back/dot/mod.rs

@@ -307,6 +307,25 @@ impl StatementGraph {
                         crate::RayQueryFunction::Terminate => "RayQueryTerminate",
                     }
                 }
+                S::MeshFunction(crate::MeshFunction::SetMeshOutputs {
+                    vertex_count,
+                    primitive_count,
+                }) => {
+                    self.dependencies.push((id, vertex_count, "vertex_count"));
+                    self.dependencies
+                        .push((id, primitive_count, "primitive_count"));
+                    "SetMeshOutputs"
+                }
+                S::MeshFunction(crate::MeshFunction::SetVertex { index, value }) => {
+                    self.dependencies.push((id, index, "index"));
+                    self.dependencies.push((id, value, "value"));
+                    "SetVertex"
+                }
+                S::MeshFunction(crate::MeshFunction::SetPrimitive { index, value }) => {
+                    self.dependencies.push((id, index, "index"));
+                    self.dependencies.push((id, value, "value"));
+                    "SetPrimitive"
+                }
                 S::SubgroupBallot { result, predicate } => {
                     if let Some(predicate) = predicate {
                         self.dependencies.push((id, predicate, "predicate"));

naga/src/back/glsl/features.rs

@@ -610,6 +610,7 @@ impl<W> Writer<'_, W> {
                     interpolation,
                     sampling,
                     blend_src,
+                    per_primitive: _,
                 } => {
                     if interpolation == Some(Interpolation::Linear) {
                         self.features.request(Features::NOPERSPECTIVE_QUALIFIER);

naga/src/back/glsl/mod.rs

@@ -139,7 +139,8 @@ impl crate::AddressSpace {
             | crate::AddressSpace::Uniform
             | crate::AddressSpace::Storage { .. }
             | crate::AddressSpace::Handle
-            | crate::AddressSpace::PushConstant => false,
+            | crate::AddressSpace::PushConstant
+            | crate::AddressSpace::TaskPayload => false,
         }
     }
 }
@@ -1300,6 +1301,9 @@ impl<'a, W: Write> Writer<'a, W> {
             crate::AddressSpace::Storage { .. } => {
                 self.write_interface_block(handle, global)?;
             }
+            crate::AddressSpace::TaskPayload => {
+                self.write_interface_block(handle, global)?;
+            }
             // A global variable in the `Function` address space is a
             // contradiction in terms.
             crate::AddressSpace::Function => unreachable!(),
@@ -1614,6 +1618,7 @@ impl<'a, W: Write> Writer<'a, W> {
                 interpolation,
                 sampling,
                 blend_src,
+                per_primitive: _,
             } => (location, interpolation, sampling, blend_src),
             crate::Binding::BuiltIn(built_in) => {
                 match built_in {
@@ -1732,6 +1737,7 @@ impl<'a, W: Write> Writer<'a, W> {
                 interpolation: None,
                 sampling: None,
                 blend_src,
+                per_primitive: false,
             },
             stage: self.entry_point.stage,
             options: VaryingOptions::from_writer_options(self.options, output),
@@ -2669,6 +2675,11 @@ impl<'a, W: Write> Writer<'a, W> {
                 self.write_image_atomic(ctx, image, coordinate, array_index, fun, value)?
             }
             Statement::RayQuery { .. } => unreachable!(),
+            Statement::MeshFunction(
+                crate::MeshFunction::SetMeshOutputs { .. }
+                | crate::MeshFunction::SetVertex { .. }
+                | crate::MeshFunction::SetPrimitive { .. },
+            ) => unreachable!(),
             Statement::SubgroupBallot { result, predicate } => {
                 write!(self.out, "{level}")?;
                 let res_name = Baked(result).to_string();
@@ -5247,6 +5258,15 @@ const fn glsl_built_in(built_in: crate::BuiltIn, options: VaryingOptions) -> &'s
         Bi::SubgroupId => "gl_SubgroupID",
         Bi::SubgroupSize => "gl_SubgroupSize",
         Bi::SubgroupInvocationId => "gl_SubgroupInvocationID",
+        // mesh
+        // TODO: figure out how to map these to glsl things as glsl treats them as arrays
+        Bi::CullPrimitive
+        | Bi::PointIndex
+        | Bi::LineIndices
+        | Bi::TriangleIndices
+        | Bi::MeshTaskSize => {
+            unimplemented!()
+        }
     }
 }
 
@@ -5262,6 +5282,7 @@ const fn glsl_storage_qualifier(space: crate::AddressSpace) -> Option<&'static s
         As::Handle => Some("uniform"),
         As::WorkGroup => Some("shared"),
         As::PushConstant => Some("uniform"),
+        As::TaskPayload => unreachable!(),
     }
 }
 

naga/src/back/hlsl/conv.rs

@@ -183,6 +183,9 @@ impl crate::BuiltIn {
             Self::PointSize | Self::ViewIndex | Self::PointCoord | Self::DrawID => {
                 return Err(Error::Custom(format!("Unsupported builtin {self:?}")))
             }
+            Self::CullPrimitive => "SV_CullPrimitive",
+            Self::PointIndex | Self::LineIndices | Self::TriangleIndices => unimplemented!(),
+            Self::MeshTaskSize => unreachable!(),
         })
     }
 }

naga/src/back/hlsl/writer.rs

@@ -507,7 +507,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
 
             self.write_wrapped_functions(module, &ctx)?;
 
-            if ep.stage == ShaderStage::Compute {
+            if ep.stage.compute_like() {
                 // HLSL is calling workgroup size "num threads"
                 let num_threads = ep.workgroup_size;
                 writeln!(
@@ -967,6 +967,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 self.write_type(module, global.ty)?;
                 ""
             }
+            crate::AddressSpace::TaskPayload => unimplemented!(),
             crate::AddressSpace::Uniform => {
                 // constant buffer declarations are expected to be inlined, e.g.
                 // `cbuffer foo: register(b0) { field1: type1; }`
@@ -2599,6 +2600,19 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     writeln!(self.out, ".Abort();")?;
                 }
             },
+            Statement::MeshFunction(crate::MeshFunction::SetMeshOutputs {
+                vertex_count,
+                primitive_count,
+            }) => {
+                write!(self.out, "{level}SetMeshOutputCounts(")?;
+                self.write_expr(module, vertex_count, func_ctx)?;
+                write!(self.out, ", ")?;
+                self.write_expr(module, primitive_count, func_ctx)?;
+                write!(self.out, ");")?;
+            }
+            Statement::MeshFunction(
+                crate::MeshFunction::SetVertex { .. } | crate::MeshFunction::SetPrimitive { .. },
+            ) => unimplemented!(),
             Statement::SubgroupBallot { result, predicate } => {
                 write!(self.out, "{level}")?;
                 let name = Baked(result).to_string();
@@ -3076,7 +3090,8 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                                 crate::AddressSpace::Function
                                 | crate::AddressSpace::Private
                                 | crate::AddressSpace::WorkGroup
-                                | crate::AddressSpace::PushConstant,
+                                | crate::AddressSpace::PushConstant
+                                | crate::AddressSpace::TaskPayload,
                             )
                             | None => true,
                             Some(crate::AddressSpace::Uniform) => {

naga/src/back/msl/mod.rs

@@ -540,6 +540,7 @@ impl Options {
                 interpolation,
                 sampling,
                 blend_src,
+                per_primitive: _,
             } => match mode {
                 LocationMode::VertexInput => Ok(ResolvedBinding::Attribute(location)),
                 LocationMode::FragmentOutput => {
@@ -697,6 +698,10 @@ impl ResolvedBinding {
                     Bi::CullDistance | Bi::ViewIndex | Bi::DrawID => {
                         return Err(Error::UnsupportedBuiltIn(built_in))
                     }
+                    Bi::CullPrimitive => "primitive_culled",
+                    // TODO: figure out how to make this written as a function call
+                    Bi::PointIndex | Bi::LineIndices | Bi::TriangleIndices => unimplemented!(),
+                    Bi::MeshTaskSize => unreachable!(),
                 };
                 write!(out, "{name}")?;
             }

naga/src/back/msl/writer.rs

@@ -594,7 +594,8 @@ impl crate::AddressSpace {
             | Self::Private
             | Self::WorkGroup
             | Self::PushConstant
-            | Self::Handle => true,
+            | Self::Handle
+            | Self::TaskPayload => true,
             Self::Function => false,
         }
     }
@@ -607,6 +608,7 @@ impl crate::AddressSpace {
             // may end up with "const" even if the binding is read-write,
             // and that should be OK.
             Self::Storage { .. } => true,
+            Self::TaskPayload => unimplemented!(),
             // These should always be read-write.
             Self::Private | Self::WorkGroup => false,
             // These translate to `constant` address space, no need for qualifiers.
@@ -623,6 +625,7 @@ impl crate::AddressSpace {
             Self::Storage { .. } => Some("device"),
             Self::Private | Self::Function => Some("thread"),
             Self::WorkGroup => Some("threadgroup"),
+            Self::TaskPayload => Some("object_data"),
         }
     }
 }
@@ -4060,6 +4063,14 @@ impl<W: Write> Writer<W> {
                         }
                     }
                 }
+                // TODO: write emitters for these
+                crate::Statement::MeshFunction(crate::MeshFunction::SetMeshOutputs { .. }) => {
+                    unimplemented!()
+                }
+                crate::Statement::MeshFunction(
+                    crate::MeshFunction::SetVertex { .. }
+                    | crate::MeshFunction::SetPrimitive { .. },
+                ) => unimplemented!(),
                 crate::Statement::SubgroupBallot { result, predicate } => {
                     write!(self.out, "{level}")?;
                     let name = self.namer.call("");
@@ -6619,7 +6630,7 @@ template <typename A>
                     LocationMode::Uniform,
                     false,
                 ),
-                crate::ShaderStage::Task | crate::ShaderStage::Mesh => unreachable!(),
+                crate::ShaderStage::Task | crate::ShaderStage::Mesh => unimplemented!(),
             };
 
             // Should this entry point be modified to do vertex pulling?
@@ -6686,6 +6697,9 @@ template <typename A>
                                 break;
                             }
                         }
+                        crate::AddressSpace::TaskPayload => {
+                            unimplemented!()
+                        }
                         crate::AddressSpace::Function
                         | crate::AddressSpace::Private
                         | crate::AddressSpace::WorkGroup => {}
@@ -7702,7 +7716,7 @@ mod workgroup_mem_init {
             fun_info: &valid::FunctionInfo,
         ) -> bool {
             options.zero_initialize_workgroup_memory
-                && ep.stage == crate::ShaderStage::Compute
+                && ep.stage.compute_like()
                 && module.global_variables.iter().any(|(handle, var)| {
                     !fun_info[handle].is_empty() && var.space == crate::AddressSpace::WorkGroup
                 })

naga/src/back/pipeline_constants.rs

@@ -39,6 +39,8 @@ pub enum PipelineConstantError {
     ValidationError(#[from] WithSpan<ValidationError>),
     #[error("workgroup_size override isn't strictly positive")]
     NegativeWorkgroupSize,
+    #[error("max vertices or max primitives is negative")]
+    NegativeMeshOutputMax,
 }
 
 /// Compact `module` and replace all overrides with constants.
@@ -243,6 +245,7 @@ pub fn process_overrides<'a>(
     for ep in entry_points.iter_mut() {
         process_function(&mut module, &override_map, &mut layouter, &mut ep.function)?;
         process_workgroup_size_override(&mut module, &adjusted_global_expressions, ep)?;
+        process_mesh_shader_overrides(&mut module, &adjusted_global_expressions, ep)?;
     }
     module.entry_points = entry_points;
     module.overrides = overrides;
@@ -296,6 +299,28 @@ fn process_workgroup_size_override(
     Ok(())
 }
 
+fn process_mesh_shader_overrides(
+    module: &mut Module,
+    adjusted_global_expressions: &HandleVec<Expression, Handle<Expression>>,
+    ep: &mut crate::EntryPoint,
+) -> Result<(), PipelineConstantError> {
+    if let Some(ref mut mesh_info) = ep.mesh_info {
+        if let Some(r#override) = mesh_info.max_vertices_override {
+            mesh_info.max_vertices = module
+                .to_ctx()
+                .eval_expr_to_u32(adjusted_global_expressions[r#override])
+                .map_err(|_| PipelineConstantError::NegativeWorkgroupSize)?;
+        }
+        if let Some(r#override) = mesh_info.max_primitives_override {
+            mesh_info.max_primitives = module
+                .to_ctx()
+                .eval_expr_to_u32(adjusted_global_expressions[r#override])
+                .map_err(|_| PipelineConstantError::NegativeWorkgroupSize)?;
+        }
+    }
+    Ok(())
+}
+
 /// Add a [`Constant`] to `module` for the override `old_h`.
 ///
 /// Add the new `Constant` to `override_map` and `adjusted_constant_initializers`.
@@ -835,6 +860,26 @@ fn adjust_stmt(new_pos: &HandleVec<Expression, Handle<Expression>>, stmt: &mut S
                 crate::RayQueryFunction::Terminate => {}
             }
         }
+        Statement::MeshFunction(crate::MeshFunction::SetMeshOutputs {
+            ref mut vertex_count,
+            ref mut primitive_count,
+        }) => {
+            adjust(vertex_count);
+            adjust(primitive_count);
+        }
+        Statement::MeshFunction(
+            crate::MeshFunction::SetVertex {
+                ref mut index,
+                ref mut value,
+            }
+            | crate::MeshFunction::SetPrimitive {
+                ref mut index,
+                ref mut value,
+            },
+        ) => {
+            adjust(index);
+            adjust(value);
+        }
         Statement::Break
         | Statement::Continue
         | Statement::Kill

naga/src/back/spv/block.rs

@@ -3654,6 +3654,7 @@ impl BlockContext<'_> {
                 } => {
                     self.write_subgroup_gather(mode, argument, result, &mut block)?;
                 }
+                Statement::MeshFunction(_) => unreachable!(),
             }
         }
 

naga/src/back/spv/helpers.rs

@@ -54,6 +54,7 @@ pub(super) const fn map_storage_class(space: crate::AddressSpace) -> spirv::Stor
         crate::AddressSpace::Uniform => spirv::StorageClass::Uniform,
         crate::AddressSpace::WorkGroup => spirv::StorageClass::Workgroup,
         crate::AddressSpace::PushConstant => spirv::StorageClass::PushConstant,
+        crate::AddressSpace::TaskPayload => unreachable!(),
     }
 }