wgsl: add spirv-unknown-naga-wgsl target, transpiling with naga 27

Author: Firestar99

Cargo.lock

@@ -61,6 +61,8 @@ itertools = "0.14.0"
 tracing.workspace = true
 tracing-subscriber.workspace = true
 tracing-tree = "0.4.0"
+naga = { version = "27.0.3", features = ["spv-in", "wgsl-out"] }
+strum = { version = "0.27.2", features = ["derive"] }
 
 [dev-dependencies]
 pretty_assertions = "1.0"

crates/rustc_codegen_spirv/Cargo.toml

@@ -133,6 +133,7 @@ mod custom_decorations;
 mod custom_insts;
 mod link;
 mod linker;
+mod naga_transpile;
 mod spirv_type;
 mod spirv_type_constraints;
 mod symbols;

crates/rustc_codegen_spirv/src/lib.rs

@@ -2,6 +2,7 @@
 use crate::maybe_pqp_cg_ssa as rustc_codegen_ssa;
 
 use crate::codegen_cx::{CodegenArgs, SpirvMetadata};
+use crate::naga_transpile::should_transpile;
 use crate::target::{SpirvTarget, SpirvTargetVariant};
 use crate::{SpirvCodegenBackend, SpirvModuleBuffer, linker};
 use ar::{Archive, GnuBuilder, Header};
@@ -323,6 +324,10 @@ fn post_link_single_module(
 
         drop(save_modules_timer);
     }
+
+    if let Some(transpile) = should_transpile(sess) {
+        transpile(sess, cg_args, &spv_binary, out_filename).ok();
+    }
 }
 
 fn do_spirv_opt(

crates/rustc_codegen_spirv/src/link.rs

@@ -0,0 +1,79 @@
+use crate::codegen_cx::CodegenArgs;
+use crate::target::{NagaTarget, SpirvTarget};
+use rustc_session::Session;
+use rustc_span::ErrorGuaranteed;
+use std::path::Path;
+
+pub type NagaTranspile = fn(
+    sess: &Session,
+    cg_args: &CodegenArgs,
+    spv_binary: &[u32],
+    out_filename: &Path,
+) -> Result<(), ErrorGuaranteed>;
+
+pub fn should_transpile(sess: &Session) -> Option<NagaTranspile> {
+    let target = SpirvTarget::parse_target(sess.opts.target_triple.tuple())
+        .expect("parsing should fail earlier");
+    match target {
+        SpirvTarget::Naga(NagaTarget::NAGA_WGSL) => Some(transpile::wgsl_transpile),
+        _ => None,
+    }
+}
+
+mod transpile {
+    use crate::codegen_cx::CodegenArgs;
+    use naga::error::ShaderError;
+    use naga::valid::Capabilities;
+    use rustc_session::Session;
+    use rustc_span::ErrorGuaranteed;
+    use std::path::Path;
+
+    pub fn wgsl_transpile(
+        sess: &Session,
+        _cg_args: &CodegenArgs,
+        spv_binary: &[u32],
+        out_filename: &Path,
+    ) -> Result<(), ErrorGuaranteed> {
+        // these should be params via spirv-builder
+        let opts = naga::front::spv::Options::default();
+        let capabilities = Capabilities::default();
+        let writer_flags = naga::back::wgsl::WriterFlags::empty();
+
+        let module = naga::front::spv::parse_u8_slice(bytemuck::cast_slice(spv_binary), &opts)
+            .map_err(|err| {
+                sess.dcx().err(format!(
+                    "Naga failed to parse spv: \n{}",
+                    ShaderError {
+                        source: String::new(),
+                        label: None,
+                        inner: Box::new(err),
+                    }
+                ))
+            })?;
+        let mut validator =
+            naga::valid::Validator::new(naga::valid::ValidationFlags::default(), capabilities);
+        let info = validator.validate(&module).map_err(|err| {
+            sess.dcx().err(format!(
+                "Naga validation failed: \n{}",
+                ShaderError {
+                    source: String::new(),
+                    label: None,
+                    inner: Box::new(err),
+                }
+            ))
+        })?;
+
+        let wgsl_dst = out_filename.with_extension("wgsl");
+        let wgsl = naga::back::wgsl::write_string(&module, &info, writer_flags).map_err(|err| {
+            sess.dcx()
+                .err(format!("Naga failed to write wgsl : \n{err}"))
+        })?;
+
+        std::fs::write(&wgsl_dst, wgsl).map_err(|err| {
+            sess.dcx()
+                .err(format!("failed to write wgsl to file: {err}"))
+        })?;
+
+        Ok(())
+    }
+}

crates/rustc_codegen_spirv/src/naga_transpile.rs

@@ -5,11 +5,16 @@ use std::cmp::Ordering;
 use std::fmt::{Debug, Display, Formatter};
 use std::ops::{Deref, DerefMut};
 use std::str::FromStr;
+use strum::{Display, EnumString, IntoStaticStr};
 
 #[derive(Clone, Eq, PartialEq)]
 pub enum TargetError {
+    /// If during parsing a target variant returns `UnknownTarget`, further variants will attempt to parse the string.
+    /// Returning another error means that you have recognized the target but something else is invalid, and we should
+    /// abort the parsing with your error.
     UnknownTarget(String),
     InvalidTargetVersion(SpirvTarget),
+    InvalidNagaVariant(String),
 }
 
 impl Display for TargetError {
@@ -21,6 +26,9 @@ impl Display for TargetError {
             TargetError::InvalidTargetVersion(target) => {
                 write!(f, "Invalid version in target `{}`", target.env())
             }
+            TargetError::InvalidNagaVariant(target) => {
+                write!(f, "Unknown naga out variant `{target}`")
+            }
         }
     }
 }
@@ -439,13 +447,71 @@ impl Display for OpenGLTarget {
     }
 }
 
+/// A naga target
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct NagaTarget {
+    pub out: NagaOut,
+}
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq, IntoStaticStr, Display, EnumString)]
+#[allow(clippy::upper_case_acronyms)]
+pub enum NagaOut {
+    #[strum(to_string = "wgsl")]
+    WGSL,
+}
+
+impl NagaTarget {
+    pub const NAGA_WGSL: Self = NagaTarget::new(NagaOut::WGSL);
+    pub const ALL_NAGA_TARGETS: &'static [Self] = &[Self::NAGA_WGSL];
+    /// emit spirv like naga targets were this target
+    pub const EMIT_SPIRV_LIKE: SpirvTarget = SpirvTarget::VULKAN_1_3;
+
+    pub const fn new(out: NagaOut) -> Self {
+        Self { out }
+    }
+}
+
+impl SpirvTargetVariant for NagaTarget {
+    fn validate(&self) -> Result<(), TargetError> {
+        Ok(())
+    }
+
+    fn to_spirv_tools(&self) -> spirv_tools::TargetEnv {
+        Self::EMIT_SPIRV_LIKE.to_spirv_tools()
+    }
+
+    fn spirv_version(&self) -> SpirvVersion {
+        Self::EMIT_SPIRV_LIKE.spirv_version()
+    }
+}
+
+impl FromStr for NagaTarget {
+    type Err = TargetError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let s = s
+            .strip_prefix("naga-")
+            .ok_or_else(|| TargetError::UnknownTarget(s.to_owned()))?;
+        Ok(Self::new(FromStr::from_str(s).map_err(|_e| {
+            TargetError::InvalidNagaVariant(s.to_owned())
+        })?))
+    }
+}
+
+impl Display for NagaTarget {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "naga-{}", self.out)
+    }
+}
+
 /// A rust-gpu target
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
 #[non_exhaustive]
 pub enum SpirvTarget {
     Universal(UniversalTarget),
     Vulkan(VulkanTarget),
     OpenGL(OpenGLTarget),
+    Naga(NagaTarget),
 }
 
 impl SpirvTarget {
@@ -467,12 +533,15 @@ impl SpirvTarget {
     pub const OPENGL_4_2: Self = Self::OpenGL(OpenGLTarget::OPENGL_4_2);
     pub const OPENGL_4_3: Self = Self::OpenGL(OpenGLTarget::OPENGL_4_3);
     pub const OPENGL_4_5: Self = Self::OpenGL(OpenGLTarget::OPENGL_4_5);
+    pub const NAGA_WGSL: Self = Self::Naga(NagaTarget::NAGA_WGSL);
 
+    #[allow(clippy::match_same_arms)]
     pub const fn memory_model(&self) -> MemoryModel {
         match self {
             SpirvTarget::Universal(_) => MemoryModel::Simple,
             SpirvTarget::Vulkan(_) => MemoryModel::Vulkan,
             SpirvTarget::OpenGL(_) => MemoryModel::GLSL450,
+            SpirvTarget::Naga(_) => MemoryModel::Vulkan,
         }
     }
 }
@@ -483,6 +552,7 @@ impl SpirvTargetVariant for SpirvTarget {
             SpirvTarget::Universal(t) => t.validate(),
             SpirvTarget::Vulkan(t) => t.validate(),
             SpirvTarget::OpenGL(t) => t.validate(),
+            SpirvTarget::Naga(t) => t.validate(),
         }
     }
 
@@ -491,6 +561,7 @@ impl SpirvTargetVariant for SpirvTarget {
             SpirvTarget::Universal(t) => t.to_spirv_tools(),
             SpirvTarget::Vulkan(t) => t.to_spirv_tools(),
             SpirvTarget::OpenGL(t) => t.to_spirv_tools(),
+            SpirvTarget::Naga(t) => t.to_spirv_tools(),
         }
     }
 
@@ -499,6 +570,7 @@ impl SpirvTargetVariant for SpirvTarget {
             SpirvTarget::Universal(t) => t.spirv_version(),
             SpirvTarget::Vulkan(t) => t.spirv_version(),
             SpirvTarget::OpenGL(t) => t.spirv_version(),
+            SpirvTarget::Naga(t) => t.spirv_version(),
         }
     }
 }
@@ -513,6 +585,9 @@ impl SpirvTarget {
         if matches!(result, Err(TargetError::UnknownTarget(..))) {
             result = OpenGLTarget::from_str(s).map(Self::OpenGL);
         }
+        if matches!(result, Err(TargetError::UnknownTarget(..))) {
+            result = NagaTarget::from_str(s).map(Self::Naga);
+        }
         result
     }
 
@@ -533,6 +608,7 @@ impl SpirvTarget {
             SpirvTarget::Universal(t) => t.to_string(),
             SpirvTarget::Vulkan(t) => t.to_string(),
             SpirvTarget::OpenGL(t) => t.to_string(),
+            SpirvTarget::Naga(t) => t.to_string(),
         }
     }
 
@@ -555,6 +631,7 @@ impl SpirvTarget {
                     .iter()
                     .map(|t| Self::OpenGL(*t)),
             )
+            .chain(NagaTarget::ALL_NAGA_TARGETS.iter().map(|t| Self::Naga(*t)))
     }
 }