WGSL parsing for mesh shaders

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),
@@ -1873,7 +1879,7 @@ impl<'a, W: Write> Writer<'a, W> {
         writeln!(self.out, ") {{")?;
 
         if self.options.zero_initialize_workgroup_memory
-            && ctx.ty.is_compute_entry_point(self.module)
+            && ctx.ty.is_compute_like_entry_point(self.module)
         {
             self.write_workgroup_variables_initialization(&ctx)?;
         }
@@ -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; }`
@@ -1764,7 +1765,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         module: &Module,
     ) -> bool {
         self.options.zero_initialize_workgroup_memory
-            && func_ctx.ty.is_compute_entry_point(module)
+            && func_ctx.ty.is_compute_like_entry_point(module)
             && module.global_variables.iter().any(|(handle, var)| {
                 !func_ctx.info[handle].is_empty() && var.space == crate::AddressSpace::WorkGroup
             })
@@ -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/mod.rs

@@ -139,11 +139,11 @@ pub enum FunctionType {
 }
 
 impl FunctionType {
-    /// Returns true if the function is an entry point for a compute shader.
-    pub fn is_compute_entry_point(&self, module: &crate::Module) -> bool {
+    /// Returns true if the function is an entry point for a compute-like shader.
+    pub fn is_compute_like_entry_point(&self, module: &crate::Module) -> bool {
         match *self {
             FunctionType::EntryPoint(index) => {
-                module.entry_points[index as usize].stage == crate::ShaderStage::Compute
+                module.entry_points[index as usize].stage.compute_like()
             }
             FunctionType::Function(_) => false,
         }

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 => {}
@@ -7683,7 +7697,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
                 })