Match entrypoints by both stage and name

Author: andyleiserson

CHANGELOG.md

@@ -68,7 +68,7 @@ Naga now infers the correct binding layout when a resource appears only in an as
 #### Naga
 
 - 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 `entry_point` member 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).
+- `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).
 
 #### D3D12
 

benches/benches/wgpu-benchmark/shader.rs

@@ -349,7 +349,7 @@ fn backends(c: &mut Criterion) {
             let options = naga::back::hlsl::Options::default();
             let mut string = String::new();
             for input in &inputs.inner {
-                let pipeline_options = naga::back::hlsl::PipelineOptions { entry_point: None };
+                let pipeline_options = Default::default();
                 let mut writer =
                     naga::back::hlsl::Writer::new(&mut string, &options, &pipeline_options);
                 let _ = writer.write(

naga-cli/src/bin/naga.rs

@@ -824,7 +824,7 @@ fn write_output(
                     .unwrap_pretty();
 
             let mut buffer = String::new();
-            let pipeline_options = hlsl::PipelineOptions { entry_point: None };
+            let pipeline_options = Default::default();
             let mut writer = hlsl::Writer::new(&mut buffer, &params.hlsl, &pipeline_options);
             writer.write(&module, &info, None).unwrap_pretty();
             fs::write(output_path, buffer)?;

naga/src/back/hlsl/mod.rs

@@ -119,7 +119,7 @@ use core::fmt::Error as FmtError;
 
 use thiserror::Error;
 
-use crate::{back, proc};
+use crate::{back, ir, proc};
 
 #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serialize", derive(serde::Serialize))]
@@ -442,9 +442,12 @@ pub struct ReflectionInfo {
 pub struct PipelineOptions {
     /// The entry point to write.
     ///
+    /// Entry points are identified by a shader stage specification,
+    /// and a name.
+    ///
     /// If `None`, all entry points will be written. If `Some` and the entry
     /// point is not found, an error will be thrown while writing.
-    pub entry_point: Option<String>,
+    pub entry_point: Option<(ir::ShaderStage, String)>,
 }
 
 #[derive(Error, Debug)]
@@ -461,8 +464,8 @@ pub enum Error {
     Override,
     #[error(transparent)]
     ResolveArraySizeError(#[from] proc::ResolveArraySizeError),
-    #[error("entry point with name '{0}' not found")]
-    EntryPointNotFound(String),
+    #[error("entry point with stage {0:?} and name '{1}' not found")]
+    EntryPointNotFound(ir::ShaderStage, String),
 }
 
 #[derive(PartialEq, Eq, Hash)]

naga/src/back/hlsl/writer.rs

@@ -1,5 +1,4 @@
 use alloc::{
-    borrow::ToOwned,
     format,
     string::{String, ToString},
     vec::Vec,
@@ -389,8 +388,8 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             writeln!(self.out)?;
         }
 
-        let ep_range = get_entry_points(module, self.pipeline_options.entry_point.as_deref())
-            .map_err(|name| Error::EntryPointNotFound(name.to_owned()))?;
+        let ep_range = get_entry_points(module, self.pipeline_options.entry_point.as_ref())
+            .map_err(|(stage, name)| Error::EntryPointNotFound(stage, name))?;
 
         // Write all entry points wrapped structs
         for index in ep_range.clone() {

naga/src/back/mod.rs

@@ -9,7 +9,7 @@ Backend functions that export shader [`Module`](super::Module)s into binary and
     )
 )]
 
-use alloc::string::String;
+use alloc::{borrow::ToOwned, string::String};
 
 #[cfg(dot_out)]
 pub mod dot;
@@ -86,17 +86,17 @@ impl core::fmt::Display for Level {
 /// `entry_point` is given, returns the complete range of entry point indices.
 /// If `entry_point` is given but does not exist, returns an error.
 #[cfg(any(hlsl_out, msl_out))]
-fn get_entry_points<'a>(
+fn get_entry_points(
     module: &crate::ir::Module,
-    entry_point: Option<&'a str>,
-) -> Result<core::ops::Range<usize>, &'a str> {
-    if let Some(entry_point) = entry_point {
+    entry_point: Option<&(crate::ir::ShaderStage, String)>,
+) -> Result<core::ops::Range<usize>, (crate::ir::ShaderStage, String)> {
+    if let Some(&(stage, ref name)) = entry_point {
         let Some(ep_index) = module
             .entry_points
             .iter()
-            .position(|ep| ep.name == *entry_point)
+            .position(|ep| ep.stage == stage && ep.name == *name)
         else {
-            return Err(entry_point);
+            return Err((stage, name.to_owned()));
         };
         Ok(ep_index..ep_index + 1)
     } else {

naga/src/back/msl/mod.rs

@@ -52,7 +52,7 @@ use alloc::{
 };
 use core::fmt::{Error as FmtError, Write};
 
-use crate::{arena::Handle, proc::index, valid::ModuleInfo};
+use crate::{arena::Handle, ir, proc::index, valid::ModuleInfo};
 
 mod keywords;
 pub mod sampler;
@@ -184,7 +184,7 @@ pub enum Error {
     #[error("can not use writeable storage buffers in fragment stage prior to MSL 1.2")]
     UnsupportedWriteableStorageBuffer,
     #[error("can not use writeable storage textures in {0:?} stage prior to MSL 1.2")]
-    UnsupportedWriteableStorageTexture(crate::ShaderStage),
+    UnsupportedWriteableStorageTexture(ir::ShaderStage),
     #[error("can not use read-write storage textures prior to MSL 1.2")]
     UnsupportedRWStorageTexture,
     #[error("array of '{0}' is not supported for target MSL version")]
@@ -199,8 +199,8 @@ pub enum Error {
     UnsupportedBitCast(crate::TypeInner),
     #[error(transparent)]
     ResolveArraySizeError(#[from] crate::proc::ResolveArraySizeError),
-    #[error("entry point with name '{0}' not found")]
-    EntryPointNotFound(String),
+    #[error("entry point with stage {0:?} and name '{1}' not found")]
+    EntryPointNotFound(ir::ShaderStage, String),
 }
 
 #[derive(Clone, Debug, PartialEq, thiserror::Error)]
@@ -424,9 +424,12 @@ pub struct VertexBufferMapping {
 pub struct PipelineOptions {
     /// The entry point to write.
     ///
+    /// Entry points are identified by a shader stage specification,
+    /// and a name.
+    ///
     /// If `None`, all entry points will be written. If `Some` and the entry
     /// point is not found, an error will be thrown while writing.
-    pub entry_point: Option<String>,
+    pub entry_point: Option<(ir::ShaderStage, String)>,
 
     /// Allow `BuiltIn::PointSize` and inject it if doesn't exist.
     ///
@@ -745,5 +748,5 @@ pub fn write_string(
 
 #[test]
 fn test_error_size() {
-    assert_eq!(size_of::<Error>(), 32);
+    assert_eq!(size_of::<Error>(), 40);
 }

naga/src/back/msl/writer.rs

@@ -1,5 +1,4 @@
 use alloc::{
-    borrow::ToOwned,
     format,
     string::{String, ToString},
     vec,
@@ -5873,8 +5872,8 @@ template <typename A>
             self.named_expressions.clear();
         }
 
-        let ep_range = get_entry_points(module, pipeline_options.entry_point.as_deref())
-            .map_err(|name| Error::EntryPointNotFound(name.to_owned()))?;
+        let ep_range = get_entry_points(module, pipeline_options.entry_point.as_ref())
+            .map_err(|(stage, name)| Error::EntryPointNotFound(stage, name))?;
 
         let mut info = TranslationInfo {
             entry_point_names: Vec::with_capacity(ep_range.len()),

naga/tests/naga/snapshots.rs

@@ -741,7 +741,7 @@ fn write_output_hlsl(
             .expect("override evaluation failed");
 
     let mut buffer = String::new();
-    let pipeline_options = hlsl::PipelineOptions { entry_point: None };
+    let pipeline_options = Default::default();
     let mut writer = hlsl::Writer::new(&mut buffer, options, &pipeline_options);
     let reflection_info = writer
         .write(&module, &info, frag_ep.as_ref())

wgpu-hal/src/dx12/device.rs

@@ -300,7 +300,7 @@ impl super::Device {
         };
 
         let pipeline_options = hlsl::PipelineOptions {
-            entry_point: Some(stage.entry_point.to_string()),
+            entry_point: Some((naga_stage, stage.entry_point.to_string())),
         };
 
         //TODO: reuse the writer